diff --git a/.claude/skills/add-rocm-kernel/SKILL.md b/.claude/skills/add-rocm-kernel/SKILL.md index d2e7c82882..a30f47da77 100644 --- a/.claude/skills/add-rocm-kernel/SKILL.md +++ b/.claude/skills/add-rocm-kernel/SKILL.md @@ -6,11 +6,11 @@ description: Step-by-step tutorial for adding new HIP kernels to FlashInfer+ROCm # Adding a New Kernel to FlashInfer+ROCm For a complete worked example to copy, read these together: -[`norm.cu`](../../../flashinfer/csrc_rocm/norm.cu) + -[`flashinfer_norm_binding.cu`](../../../flashinfer/csrc_rocm/flashinfer_norm_binding.cu) + +[`norm.cu`](../../../flashinfer/csrc/rocm/norm.cu) + +[`flashinfer_norm_binding.cu`](../../../flashinfer/csrc/rocm/flashinfer_norm_binding.cu) + [`jit/norm.py`](../../../flashinfer/jit/norm.py) + [`norm.py`](../../../flashinfer/norm.py). For plan-run / multi-backend / FP8 see -[`batch_prefill.cu`](../../../flashinfer/csrc_rocm/batch_prefill.cu) + +[`batch_prefill.cu`](../../../flashinfer/csrc/rocm/batch_prefill.cu) + [`prefill_rocm.py`](../../../flashinfer/prefill_rocm.py). ## File touchpoints (every new op needs each row, in order) @@ -18,9 +18,9 @@ For a complete worked example to copy, read these together: | Step | File | Purpose | | --- | --- | --- | | 1 | `include/flashinfer/.cuh` | Framework-agnostic kernel + launcher template. **No `` includes here.** | -| 2 | `flashinfer/csrc_rocm/.cu` | PyTorch launcher: `at::Tensor` in, `at::hip::getCurrentHIPStream()`, `TORCH_CHECK`, `DISPATCH_PYTORCH_DTYPE_*`. | -| 3 | `flashinfer/csrc_rocm/flashinfer__binding.cu` | `TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { m.def("", ); }`. | -| 4 (opt) | `flashinfer/csrc_rocm/_customize_config.jinja` | Compile-time type specialization. Skip if runtime dispatch is enough. | +| 2 | `flashinfer/csrc/rocm/.cu` | PyTorch launcher: `at::Tensor` in, `at::hip::getCurrentHIPStream()`, `TORCH_CHECK`, `DISPATCH_PYTORCH_DTYPE_*`. | +| 3 | `flashinfer/csrc/rocm/flashinfer__binding.cu` | `TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { m.def("", ); }`. | +| 4 (opt) | `flashinfer/csrc/rocm/_customize_config.jinja` | Compile-time type specialization. Skip if runtime dispatch is enough. | | 5 | `flashinfer/jit/.py` | `gen__module() -> JitSpec` via `gen_jit_spec(...)`. | | 6 | `flashinfer/.py` | Python API: `@functools.cache` module loader, destination-passing (`out=`). | | 7 | `tests/rocm_tests/test__hip.py` | Correctness tests; FP32 reference math, loose BF16 tolerances. | @@ -36,7 +36,7 @@ When porting an upstream kernel, mechanically rewrite: | Upstream CUDA | This fork | | --- | --- | -| `csrc/.cu` | `flashinfer/csrc_rocm/.cu` | +| `csrc/.cu` | `flashinfer/csrc/rocm/.cu` | | `#include "tvm_ffi_utils.h"` | `#include "pytorch_extension_utils.h"` | | `tvm::ffi::TensorView` | `at::Tensor` | | `TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, op)` | `TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { m.def("op", op); }` | @@ -48,16 +48,16 @@ When porting an upstream kernel, mechanically rewrite: | `flashinfer/aot.py` registration | `flashinfer/aot_hip.py` | | `tests/test_op.py` | `tests/rocm_tests/test_op_hip.py` | | `supported_major_versions=[9, 10]` | No analogue. Guard at Python layer via `FLASHINFER_SUPPORTED_ROCM_ARCHS`. | -| `csrc/` (hardcoded) | `jit_env.FLASHINFER_CSRC_DIR` resolves to `flashinfer/csrc_rocm/` on HIP. **Never hardcode `csrc/`.** | +| `csrc/` (hardcoded) | `jit_env.FLASHINFER_CSRC_DIR` resolves to `flashinfer/csrc/rocm/` on HIP. **Never hardcode `csrc/`.** | | `PYBIND11_MODULE(...)` | **Don't.** Use `TORCH_LIBRARY_FRAGMENT` (integrates with `torch.compile`). | ## Non-obvious gotchas - **PyTorch's ROCm masquerade.** `input.device.type == "cuda"` even on AMD. Never check for `"hip"`. PyTorch's HIP namespaces are reachable via `at::hip::...` and `c10::hip::OptionalHIPGuardMasqueradingAsCUDA` (literally the type name). -- **`gpu_iface` over duplication.** If a primitive (MMA intrinsic, cross-lane shuffle, dtype container, warp reduction) needs a HIP-specific implementation, add it under [`include/gpu_iface/backend/hip/`](../../../include/gpu_iface) and expose a common name from the top-level `gpu_iface/` header. Don't fork the kernel into `csrc_rocm/`. Existing HIP backends: `mma_hip.h`, `memory_ops_hip.h`, `math_hip.h`, `vec_dtypes_hip.h`. +- **Shared intrinsics over duplication.** If a primitive (MMA intrinsic, cross-lane shuffle, dtype container, warp reduction) needs a HIP-specific implementation, add it to the matching `_hip.h` in [`include/flashinfer/rocm/`](../../../include/flashinfer/rocm) rather than forking the kernel into `csrc/rocm/`. Existing ones: `mma_hip.h`, `memory_ops_hip.h`, `math_hip.h`, `vec_dtypes_hip.h`. Symbols go in `flashinfer::`, grouped by area (`flashinfer::math`, `flashinfer::memory`); `flashinfer::mma_hip` is renamed rather than tripwired because three of its signatures match upstream's `flashinfer::mma` exactly. - **`-ffast-math` adds `-ffinite-math-only` on clang/hipcc.** [`jit/core.py`](../../../flashinfer/jit/core.py) explicitly re-adds `-fno-finite-math-only` so kernels that use `-inf` as a sentinel (online-softmax Map+Reduce) keep working. CUDA's `-use_fast_math` does *not* enable finite-math-only — divergence to be aware of when porting. - **`gen_jit_spec` auto-injects `--offload-arch=gfxNNN`** for every target arch plus `COMMON_HIPCC_FLAGS` (`-DFLASHINFER_ENABLE_HIP`, FP8 enables, etc.). Don't add `--offload-arch` by hand. -- **Validation macros** live in [`pytorch_extension_utils.h`](../../../flashinfer/csrc_rocm/pytorch_extension_utils.h): `CHECK_INPUT` (GPU + contiguous), `CHECK_LAST_DIM_CONTIGUOUS_INPUT`, `CHECK_EQ`, `CHECK_DIM`, `CHECK_GE`, `CHECK_SHAPE`. Dispatch macros: `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16` (FP16+BF16), `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP8` (E4M3+E5M2, both `_fnuz` on CDNA3/4), and the unsuffixed `DISPATCH_PYTORCH_DTYPE_TO_CTYPE` (FP16+BF16+FP8 combined). There is **no** `_FP16_FP32` variant — if you need FP32, dispatch manually. +- **Validation macros** live in [`pytorch_extension_utils.h`](../../../flashinfer/csrc/rocm/pytorch_extension_utils.h): `CHECK_INPUT` (GPU + contiguous), `CHECK_LAST_DIM_CONTIGUOUS_INPUT`, `CHECK_EQ`, `CHECK_DIM`, `CHECK_GE`, `CHECK_SHAPE`. Dispatch macros: `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16` (FP16+BF16), `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP8` (E4M3+E5M2, both `_fnuz` on CDNA3/4), and the unsuffixed `DISPATCH_PYTORCH_DTYPE_TO_CTYPE` (FP16+BF16+FP8 combined). There is **no** `_FP16_FP32` variant — if you need FP32, dispatch manually. - **The `_jit_pybind.cu` naming pattern** (e.g. `batch_decode_jit_pybind.cu`) is used by newer AITER-integrated bindings; the older `flashinfer__binding.cu` pattern is used by everything else. Both work — match the neighbors. ## CDNA3 (`gfx942`) vs CDNA4 (`gfx950`) diff --git a/.claude/skills/code-coverage/SKILL.md b/.claude/skills/code-coverage/SKILL.md index 56bbab95ac..8ff033e218 100644 --- a/.claude/skills/code-coverage/SKILL.md +++ b/.claude/skills/code-coverage/SKILL.md @@ -55,7 +55,7 @@ though no ROCm box executes it. Detection is AST-based and matches only a bare `IS_CUDA` test; a compound condition stays in the denominator rather than being dropped on a guess. -**`csrc_rocm` reach is not coverage.** JIT-built HIP has no line data. The +**`csrc/rocm` reach is not coverage.** JIT-built HIP has no line data. The report says how many of its translation units a run built and loaded, via the `tests/jit_reach_plugin.py` hook on `JitSpec.load`. Do not quote it as a percentage or add it to the Python number. diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 78256ff194..94666af2cc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -77,7 +77,7 @@ repos: name: reuse (selective files) entry: ./scripts/reuse-check-selective.sh language: system - files: ^(include/flashinfer/rocm/|include/gpu_iface/).*\.(cuh|hpp|h)$ + files: ^include/flashinfer/(rocm|attention/aiter)/.*\.(cuh|hpp|h)$ pass_filenames: true # The README's per-architecture support matrix is generated from diff --git a/CLAUDE.md b/CLAUDE.md index f126c420cd..c105ef550d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -97,7 +97,7 @@ JIT generator (the HIP path injects `-O3` before `extra_cuda_cflags`, so trailin **Framework separation**: Torch headers **must not** be included in `include/` files. `include/` is framework-agnostic (raw pointers only); -`flashinfer/csrc_rocm/` is where PyTorch tensor handling lives. Violations +`flashinfer/csrc/rocm/` is where PyTorch tensor handling lives. Violations cause subtle build failures. **Test parallelism**: `pytest -n auto` automatically halves the physical GPU diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ba1cf94fe7..9cce871901 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -131,7 +131,7 @@ python3 scripts/amd_coverage.py # re-score an existing **What gets counted.** Files we added are scored whole. Upstream files we merely edited are scored **only on the lines our diff touched**, so upstream's untested code neither flatters nor penalises the number. A third tier covers files with a zero-line Python diff whose implementation is ours anyway through the `FLASHINFER_CSRC_DIR` redirect — `sampling.py` and friends — which no diff can discover; they are declared in `scripts/coverage_ownership.toml`, one entry per file with a reason. -**What is deliberately left out, and why the report says so.** Lines inside `if IS_CUDA:` are excluded and counted in the output: the port re-indented upstream code under those guards, so git attributes it to us even though no ROCm box can execute it — in `flashinfer/jit/env.py` that is about half the owned lines. Lines that run at `import flashinfer` are reported as their own bucket rather than in the headline, because `tests/conftest.py` imports the package at collection and would otherwise credit every module-level statement before a test body runs. C++ under `csrc_rocm/` is JIT-compiled and has no line data at all; instead the report counts how many of its translation units a run actually built and loaded, labelled as reach, not coverage. +**What is deliberately left out, and why the report says so.** Lines inside `if IS_CUDA:` are excluded and counted in the output: the port re-indented upstream code under those guards, so git attributes it to us even though no ROCm box can execute it — in `flashinfer/jit/env.py` that is about half the owned lines. Lines that run at `import flashinfer` are reported as their own bucket rather than in the headline, because `tests/conftest.py` imports the package at collection and would otherwise credit every module-level statement before a test body runs. C++ under `flashinfer/csrc/rocm/` is JIT-compiled and has no line data at all; instead the report counts how many of its translation units a run actually built and loaded, labelled as reach, not coverage. **The last measured run is committed** at [`docs/rocm/coverage-gfx942.json`](docs/rocm/coverage-gfx942.json), so a change that drops coverage shows up as a reviewable diff rather than going unnoticed until someone re-runs the suite by hand (about 75 minutes under `--cov` instrumentation, against the ~20 min in the table above uninstrumented). Refresh it in the same commit as any change that moves the number, and before each `+amd.N` tag, by adding `--json-out docs/rocm/coverage-gfx942.json` to the invocation above; relative paths are anchored to the repository root, so it does not matter where you run it from. There is no automated gate — a stale payload is invisible, and the artifact records no HEAD sha to check it against, so this is a convention rather than an enforcement. @@ -148,11 +148,11 @@ The classifier itself is covered by `tests/rocm_tests/test_amd_coverage.py`, whi ```text flashinfer/ ├── include/ # framework-agnostic kernel headers (raw pointers only) -│ ├── flashinfer/ # FlashInfer kernel implementations -│ └── gpu_iface/backend/hip/ # HIP intrinsics behind a common header surface +│ └── flashinfer/ # FlashInfer kernel implementations +│ └── rocm/ # fork-owned headers, incl. the HIP intrinsics ├── csrc/ # upstream CUDA op registration (PyTorch bindings) ├── flashinfer/ -│ ├── csrc_rocm/ # HIP op registration (PyTorch bindings) — the ROCm analog of csrc/ +│ ├── csrc/rocm/ # HIP op registration (PyTorch bindings) — the ROCm analog of csrc/ │ ├── jit/ # Python JIT compilation infra (cpp_ext_hip.py is the HIP entry) │ └── *.py # Python user-facing API (e.g. attention.py, mla_rocm.py) ├── tests/rocm_tests/ # HIP test suite (test_*_hip.py) @@ -162,23 +162,32 @@ flashinfer/ **Framework separation.** `include/` files must remain framework-agnostic — no PyTorch headers, raw pointers only. PyTorch tensor handling for HIP -ops lives in `flashinfer/csrc_rocm/`. Violating this causes subtle build +ops lives in `flashinfer/csrc/rocm/`. Violating this causes subtle build failures because the same headers are pulled into the JIT compilation pipeline that has no PyTorch on its include path. -**`csrc/` vs `flashinfer/csrc_rocm/`.** `csrc/` is the upstream CUDA op +**`csrc/` vs `flashinfer/csrc/rocm/`.** `csrc/` is the upstream CUDA op registration tree — keep it in sync with upstream where possible to reduce merge conflicts. New HIP-specific op bindings go in -`flashinfer/csrc_rocm/`, with a `_hip` or `_aiter` suffix when the file +`flashinfer/csrc/rocm/`, with a `_hip` or `_aiter` suffix when the file routes to a HIP-specific code path or to AITER. -**`include/gpu_iface/`.** A common header surface (`math_ops.hpp`, -`mma_ops.hpp`, `memory_ops.hpp`, …) over HIP intrinsics. It once spanned -CUDA too; that half is gone, so a non-HIP compiler now gets an `#error` -from `macros.hpp`. When you need a new intrinsic, add the abstraction in -`gpu_iface/` and implement it under `gpu_iface/backend/hip/`. Don't -reach for `hipcub`, `__hip_*`, or inline asm from inside -`include/flashinfer/` — go through `gpu_iface`. +**HIP intrinsics.** `include/flashinfer/rocm/*_hip.h` wrap the HIP intrinsics +(`math_hip.h`, `mma_hip.h`, `memory_ops_hip.h`, `vec_dtypes_hip.h`). These once +sat behind a `gpu_iface` abstraction spanning CUDA too; that half is gone, so a +non-HIP compiler now gets an `#error` from `macros.hpp`. Put a new intrinsic in +the matching `_hip.h` if it is a general primitive. A kernel that needs +`hipcub` or one inline-asm builtin inline is fine — several under +`rocm/attention/` do — but anything a second kernel would want belongs in the +shared header. + +Symbols live in `flashinfer::`, grouped by what they do — `flashinfer::math` +matches upstream's name, `flashinfer::memory` is ours (upstream calls the same +area `cp_async`). Where a fork header and its upstream namesake can coexist +they keep the same name and the fork header carries an `#error` tripwire on +upstream's include guard; `flashinfer::mma_hip` is renamed instead, because +three of its signatures match upstream's `flashinfer::mma` exactly and the two +are meant to be usable together. # Additive-Only: the rule that keeps upstream syncs cheap @@ -189,13 +198,16 @@ however large — are close to free at merge time, because upstream has nothing to merge them against. The exception is a path upstream later adds too: that conflicts as add/add, with no common ancestor to help resolve it, which is how `CLAUDE.md` and `.claude/skills/benchmark-kernel/SKILL.md` got onto the -conflict list. Prefer a `_rocm`/`_aiter`-suffixed name for anything upstream -might plausibly create. +conflict list. For anything upstream might plausibly create, prefer a +`rocm/` subdirectory over a sibling file: `flashinfer/csrc/rocm/` and +`include/flashinfer/rocm/` collide with nothing even as upstream grows those +trees. A `_rocm`/`_aiter` suffix is the fallback where a subdirectory does not +fit, as with the `_hip.h` intrinsic headers. **So: add files, don't edit them.** Concretely, prefer in this order: 1. **Source-path redirect.** `FLASHINFER_CSRC_DIR` already points at - `flashinfer/csrc_rocm/` on ROCm (see `flashinfer/jit/env.py` and + `flashinfer/csrc/rocm/` on ROCm (see `flashinfer/jit/env.py` and `flashinfer/get_include_paths.py`), so a shared JIT generator naming `sampling.cu` picks up the HIP source with **zero Python diff**. This is why `flashinfer/sampling.py` and `flashinfer/quantization.py` contain no HIP @@ -217,9 +229,11 @@ additions such as `prefill_rocm.py` that are not edits to anything, and an in-place edit under `csrc/` or `include/` never appears there at all. **Forked headers are exempt from conflicts and therefore from warnings.** -Everything under `include/flashinfer/rocm/` is a fork of an upstream header -re-expressed on `gpu_iface` — `rocm/attention/` for the attention headers, -plus `rocm/sampling.cuh` and `rocm/quantization.cuh`. Their upstream +Much of `include/flashinfer/rocm/` is a fork of an upstream header — the +`rocm/attention/` set, plus `sampling.cuh`, `quantization.cuh`, `layout.cuh`, +`fastdiv.cuh` and `exception.h`. The `_hip.h` intrinsics and their types headers have +no upstream counterpart, and `utils.cuh` shares a basename without forking +anything, so `upstream_canary.py` excludes it by exact path. Their upstream originals are byte-identical to the merge base and will merge cleanly forever, so a fix landing upstream reaches the original and *not* the fork, with nothing conflicting to tell you. The canary's drift report is the only signal, and a fix @@ -242,9 +256,9 @@ state; what matters is that your change does not lengthen the list. # Adding a Kernel 1. **Kernel implementation** — framework-agnostic header(s) in - `include/flashinfer/`, using `gpu_iface/` for any CUDA/HIP-divergent + `include/flashinfer/rocm/`, using the `_hip.h` headers for any HIP-specific intrinsic. -2. **PyTorch binding** — register the op in `flashinfer/csrc_rocm/`. +2. **PyTorch binding** — register the op in `flashinfer/csrc/rocm/`. The only layer that may include Torch headers. 3. **JIT generator** — add the op's JIT spec in `flashinfer/jit/*.py`. 4. **Python interface** — expose the user-facing API in `flashinfer/*.py`. diff --git a/MANIFEST.in b/MANIFEST.in index b256dde6b9..57682d8df0 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -28,9 +28,9 @@ include NOTICE include README.md recursive-include licenses *.txt -# csrc_rocm is declared as package-data, but graft it explicitly so the sdist is +# csrc/rocm is declared as package-data, but graft it explicitly so the sdist is # self-sufficient regardless of setuptools' package-data-in-sdist behavior. -recursive-include flashinfer/csrc_rocm *.cu *.cc *.h *.jinja +recursive-include flashinfer/csrc/rocm *.cu *.cc *.h *.jinja # setuptools-scm's file finder adds every git-tracked file to the sdist, so # development and CI machinery has to be pruned explicitly or it ships in the diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index a96e1a8db4..3dfec24036 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -45,7 +45,7 @@ specifier. **`aiter_utils.AITER_MIN_VERSION` (0.1.16) is a hard floor**, enforced before routing. FlashInfer links AITER's C++ symbols by mangled name -(`flashinfer/csrc_rocm/aiter_loader.cc`) and vendors its argument structs +(`flashinfer/csrc/rocm/aiter_loader.cc`) and vendors its argument structs (`include/flashinfer/attention/aiter/`) at the 0.1.16 layout, so an older release shifts field offsets instead of failing to load. Below the floor `auto` will not select AITER and an explicit `backend="aiter"` raises. diff --git a/docs/rocm/coverage-gfx942.json b/docs/rocm/coverage-gfx942.json index 39cfdbc278..0df0438f04 100644 --- a/docs/rocm/coverage-gfx942.json +++ b/docs/rocm/coverage-gfx942.json @@ -347,7 +347,7 @@ { "path": "flashinfer/jit/cascade.py", "tier": "C", - "reason": "builds csrc_rocm/{cascade,flashinfer_cascade_binding}.cu", + "reason": "builds flashinfer/csrc/rocm/{cascade,flashinfer_cascade_binding}.cu", "owned": 1, "covered": 1, "import_time": 3, @@ -429,7 +429,7 @@ { "path": "flashinfer/jit/quantization.py", "tier": "C", - "reason": "builds csrc_rocm/{quantization,flashinfer_quantization_binding}.cu", + "reason": "builds flashinfer/csrc/rocm/{quantization,flashinfer_quantization_binding}.cu", "owned": 1, "covered": 1, "import_time": 3, @@ -506,7 +506,7 @@ { "path": "flashinfer/jit/sampling.py", "tier": "C", - "reason": "builds csrc_rocm/{sampling,renorm,flashinfer_sampling_binding}.cu", + "reason": "builds flashinfer/csrc/rocm/{sampling,renorm,flashinfer_sampling_binding}.cu", "owned": 1, "covered": 1, "import_time": 3, @@ -727,7 +727,7 @@ { "path": "flashinfer/quantization.py", "tier": "C", - "reason": "wrapper over csrc_rocm/quantization.cu", + "reason": "wrapper over flashinfer/csrc/rocm/quantization.cu", "owned": 19, "covered": 19, "import_time": 13, @@ -822,7 +822,7 @@ { "path": "flashinfer/sampling.py", "tier": "C", - "reason": "wrapper over csrc_rocm/sampling.cu and renorm.cu", + "reason": "wrapper over flashinfer/csrc/rocm/sampling.cu and renorm.cu", "owned": 228, "covered": 187, "import_time": 25, diff --git a/flashinfer/csrc_rocm/activation.cu b/flashinfer/csrc/rocm/activation.cu similarity index 100% rename from flashinfer/csrc_rocm/activation.cu rename to flashinfer/csrc/rocm/activation.cu diff --git a/flashinfer/csrc_rocm/activation_aiter.cu b/flashinfer/csrc/rocm/activation_aiter.cu similarity index 100% rename from flashinfer/csrc_rocm/activation_aiter.cu rename to flashinfer/csrc/rocm/activation_aiter.cu diff --git a/flashinfer/csrc_rocm/activation_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/activation_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/activation_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/activation_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/aiter_loader.cc b/flashinfer/csrc/rocm/aiter_loader.cc similarity index 99% rename from flashinfer/csrc_rocm/aiter_loader.cc rename to flashinfer/csrc/rocm/aiter_loader.cc index e981a52fb8..47c31864d7 100644 --- a/flashinfer/csrc_rocm/aiter_loader.cc +++ b/flashinfer/csrc/rocm/aiter_loader.cc @@ -26,7 +26,7 @@ constexpr const char* kAbiPinNote = "\n C++ ABI. If AITER was upgraded, that is the likely cause. Reinstall the pin:" "\n pip install amd-aiter==0.1.10 \\" "\n --extra-index-url https://pypi.amd.com/rocm-7.1.1/simple" - "\n or re-pin the symbols in flashinfer/csrc_rocm/aiter_loader.cc."; + "\n or re-pin the symbols in flashinfer/csrc/rocm/aiter_loader.cc."; std::string get_jit_dir() { if (const char* env = std::getenv("AITER_JIT_DIR")) return env; diff --git a/flashinfer/csrc_rocm/aiter_tensor_compat.h b/flashinfer/csrc/rocm/aiter_tensor_compat.h similarity index 100% rename from flashinfer/csrc_rocm/aiter_tensor_compat.h rename to flashinfer/csrc/rocm/aiter_tensor_compat.h diff --git a/flashinfer/csrc_rocm/aot_extension_utils.h b/flashinfer/csrc/rocm/aot_extension_utils.h similarity index 100% rename from flashinfer/csrc_rocm/aot_extension_utils.h rename to flashinfer/csrc/rocm/aot_extension_utils.h diff --git a/flashinfer/csrc_rocm/batch_decode.cu b/flashinfer/csrc/rocm/batch_decode.cu similarity index 99% rename from flashinfer/csrc_rocm/batch_decode.cu rename to flashinfer/csrc/rocm/batch_decode.cu index 002fed2d76..3ffe438816 100644 --- a/flashinfer/csrc_rocm/batch_decode.cu +++ b/flashinfer/csrc/rocm/batch_decode.cu @@ -15,7 +15,7 @@ */ #include #include -#include +#include #include #include "batch_decode_config.inc" diff --git a/flashinfer/csrc_rocm/batch_decode_aiter.cu b/flashinfer/csrc/rocm/batch_decode_aiter.cu similarity index 100% rename from flashinfer/csrc_rocm/batch_decode_aiter.cu rename to flashinfer/csrc/rocm/batch_decode_aiter.cu diff --git a/flashinfer/csrc_rocm/batch_decode_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/batch_decode_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/batch_decode_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/batch_decode_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/batch_decode_customize_config.jinja b/flashinfer/csrc/rocm/batch_decode_customize_config.jinja similarity index 95% rename from flashinfer/csrc_rocm/batch_decode_customize_config.jinja rename to flashinfer/csrc/rocm/batch_decode_customize_config.jinja index c570a1a623..35974e162e 100644 --- a/flashinfer/csrc_rocm/batch_decode_customize_config.jinja +++ b/flashinfer/csrc/rocm/batch_decode_customize_config.jinja @@ -1,6 +1,6 @@ #pragma once -#include -#include +#include +#include #include #include diff --git a/flashinfer/csrc_rocm/batch_decode_jit_pybind.cu b/flashinfer/csrc/rocm/batch_decode_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/batch_decode_jit_pybind.cu rename to flashinfer/csrc/rocm/batch_decode_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/batch_decode_kernel_inst.jinja b/flashinfer/csrc/rocm/batch_decode_kernel_inst.jinja similarity index 100% rename from flashinfer/csrc_rocm/batch_decode_kernel_inst.jinja rename to flashinfer/csrc/rocm/batch_decode_kernel_inst.jinja diff --git a/flashinfer/csrc_rocm/batch_pod.cu b/flashinfer/csrc/rocm/batch_pod.cu similarity index 99% rename from flashinfer/csrc_rocm/batch_pod.cu rename to flashinfer/csrc/rocm/batch_pod.cu index bfd4e548c2..74fac8d7c2 100644 --- a/flashinfer/csrc_rocm/batch_pod.cu +++ b/flashinfer/csrc/rocm/batch_pod.cu @@ -7,9 +7,9 @@ #include #include -#include -#include -#include +#include +#include +#include #include #include "batch_pod_config.inc" diff --git a/flashinfer/csrc_rocm/batch_pod_customize_config.jinja b/flashinfer/csrc/rocm/batch_pod_customize_config.jinja similarity index 93% rename from flashinfer/csrc_rocm/batch_pod_customize_config.jinja rename to flashinfer/csrc/rocm/batch_pod_customize_config.jinja index f373e9b309..bc4882b27d 100644 --- a/flashinfer/csrc_rocm/batch_pod_customize_config.jinja +++ b/flashinfer/csrc/rocm/batch_pod_customize_config.jinja @@ -1,8 +1,8 @@ #pragma once -#include -#include -#include -#include +#include +#include +#include +#include #include #include diff --git a/flashinfer/csrc_rocm/batch_pod_jit_pybind.cu b/flashinfer/csrc/rocm/batch_pod_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/batch_pod_jit_pybind.cu rename to flashinfer/csrc/rocm/batch_pod_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/batch_pod_kernel_inst.jinja b/flashinfer/csrc/rocm/batch_pod_kernel_inst.jinja similarity index 100% rename from flashinfer/csrc_rocm/batch_pod_kernel_inst.jinja rename to flashinfer/csrc/rocm/batch_pod_kernel_inst.jinja diff --git a/flashinfer/csrc_rocm/batch_prefill.cu b/flashinfer/csrc/rocm/batch_prefill.cu similarity index 99% rename from flashinfer/csrc_rocm/batch_prefill.cu rename to flashinfer/csrc/rocm/batch_prefill.cu index 4059913352..3b47ce5a6f 100644 --- a/flashinfer/csrc_rocm/batch_prefill.cu +++ b/flashinfer/csrc/rocm/batch_prefill.cu @@ -17,7 +17,7 @@ #include #include -#include +#include #include #include "batch_prefill_config.inc" diff --git a/flashinfer/csrc_rocm/batch_prefill_aiter_customize_config.jinja b/flashinfer/csrc/rocm/batch_prefill_aiter_customize_config.jinja similarity index 100% rename from flashinfer/csrc_rocm/batch_prefill_aiter_customize_config.jinja rename to flashinfer/csrc/rocm/batch_prefill_aiter_customize_config.jinja diff --git a/flashinfer/csrc_rocm/batch_prefill_customize_config.jinja b/flashinfer/csrc/rocm/batch_prefill_customize_config.jinja similarity index 95% rename from flashinfer/csrc_rocm/batch_prefill_customize_config.jinja rename to flashinfer/csrc/rocm/batch_prefill_customize_config.jinja index c50f0211bb..b4b99552fa 100644 --- a/flashinfer/csrc_rocm/batch_prefill_customize_config.jinja +++ b/flashinfer/csrc/rocm/batch_prefill_customize_config.jinja @@ -1,8 +1,8 @@ #pragma once -#include -#include -#include -#include +#include +#include +#include +#include #include #include diff --git a/flashinfer/csrc_rocm/batch_prefill_jit_pybind.cu b/flashinfer/csrc/rocm/batch_prefill_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/batch_prefill_jit_pybind.cu rename to flashinfer/csrc/rocm/batch_prefill_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/batch_prefill_paged_aiter.cu b/flashinfer/csrc/rocm/batch_prefill_paged_aiter.cu similarity index 99% rename from flashinfer/csrc_rocm/batch_prefill_paged_aiter.cu rename to flashinfer/csrc/rocm/batch_prefill_paged_aiter.cu index 678b15ac58..5879443451 100644 --- a/flashinfer/csrc_rocm/batch_prefill_paged_aiter.cu +++ b/flashinfer/csrc/rocm/batch_prefill_paged_aiter.cu @@ -11,8 +11,8 @@ #include #include -#include -#include +#include +#include #include #include diff --git a/flashinfer/csrc_rocm/batch_prefill_paged_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/batch_prefill_paged_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/batch_prefill_paged_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/batch_prefill_paged_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/batch_prefill_paged_kernel_inst.jinja b/flashinfer/csrc/rocm/batch_prefill_paged_kernel_inst.jinja similarity index 100% rename from flashinfer/csrc_rocm/batch_prefill_paged_kernel_inst.jinja rename to flashinfer/csrc/rocm/batch_prefill_paged_kernel_inst.jinja diff --git a/flashinfer/csrc_rocm/batch_prefill_ragged_kernel_inst.jinja b/flashinfer/csrc/rocm/batch_prefill_ragged_kernel_inst.jinja similarity index 100% rename from flashinfer/csrc_rocm/batch_prefill_ragged_kernel_inst.jinja rename to flashinfer/csrc/rocm/batch_prefill_ragged_kernel_inst.jinja diff --git a/flashinfer/csrc_rocm/batch_ragged_prefill_aiter.cu b/flashinfer/csrc/rocm/batch_ragged_prefill_aiter.cu similarity index 98% rename from flashinfer/csrc_rocm/batch_ragged_prefill_aiter.cu rename to flashinfer/csrc/rocm/batch_ragged_prefill_aiter.cu index 3e751783a5..6a42b06ced 100644 --- a/flashinfer/csrc_rocm/batch_ragged_prefill_aiter.cu +++ b/flashinfer/csrc/rocm/batch_ragged_prefill_aiter.cu @@ -13,8 +13,8 @@ #include #include -#include -#include +#include +#include #include #include "batch_prefill_aiter_config.inc" diff --git a/flashinfer/csrc_rocm/batch_ragged_prefill_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/batch_ragged_prefill_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/batch_ragged_prefill_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/batch_ragged_prefill_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/cascade.cu b/flashinfer/csrc/rocm/cascade.cu similarity index 100% rename from flashinfer/csrc_rocm/cascade.cu rename to flashinfer/csrc/rocm/cascade.cu diff --git a/flashinfer/csrc_rocm/flashinfer_cascade_binding.cu b/flashinfer/csrc/rocm/flashinfer_cascade_binding.cu similarity index 100% rename from flashinfer/csrc_rocm/flashinfer_cascade_binding.cu rename to flashinfer/csrc/rocm/flashinfer_cascade_binding.cu diff --git a/flashinfer/csrc_rocm/flashinfer_norm_binding.cu b/flashinfer/csrc/rocm/flashinfer_norm_binding.cu similarity index 100% rename from flashinfer/csrc_rocm/flashinfer_norm_binding.cu rename to flashinfer/csrc/rocm/flashinfer_norm_binding.cu diff --git a/flashinfer/csrc_rocm/flashinfer_ops.cu b/flashinfer/csrc/rocm/flashinfer_ops.cu similarity index 100% rename from flashinfer/csrc_rocm/flashinfer_ops.cu rename to flashinfer/csrc/rocm/flashinfer_ops.cu diff --git a/flashinfer/csrc_rocm/flashinfer_page_binding.cu b/flashinfer/csrc/rocm/flashinfer_page_binding.cu similarity index 100% rename from flashinfer/csrc_rocm/flashinfer_page_binding.cu rename to flashinfer/csrc/rocm/flashinfer_page_binding.cu diff --git a/flashinfer/csrc_rocm/flashinfer_quantization_binding.cu b/flashinfer/csrc/rocm/flashinfer_quantization_binding.cu similarity index 100% rename from flashinfer/csrc_rocm/flashinfer_quantization_binding.cu rename to flashinfer/csrc/rocm/flashinfer_quantization_binding.cu diff --git a/flashinfer/csrc_rocm/flashinfer_rope_binding.cu b/flashinfer/csrc/rocm/flashinfer_rope_binding.cu similarity index 100% rename from flashinfer/csrc_rocm/flashinfer_rope_binding.cu rename to flashinfer/csrc/rocm/flashinfer_rope_binding.cu diff --git a/flashinfer/csrc_rocm/flashinfer_sampling_binding.cu b/flashinfer/csrc/rocm/flashinfer_sampling_binding.cu similarity index 100% rename from flashinfer/csrc_rocm/flashinfer_sampling_binding.cu rename to flashinfer/csrc/rocm/flashinfer_sampling_binding.cu diff --git a/flashinfer/csrc_rocm/fused_moe_aiter.cu b/flashinfer/csrc/rocm/fused_moe_aiter.cu similarity index 100% rename from flashinfer/csrc_rocm/fused_moe_aiter.cu rename to flashinfer/csrc/rocm/fused_moe_aiter.cu diff --git a/flashinfer/csrc_rocm/fused_moe_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/fused_moe_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/fused_moe_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/fused_moe_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/norm.cu b/flashinfer/csrc/rocm/norm.cu similarity index 100% rename from flashinfer/csrc_rocm/norm.cu rename to flashinfer/csrc/rocm/norm.cu diff --git a/flashinfer/csrc_rocm/norm_aiter.cu b/flashinfer/csrc/rocm/norm_aiter.cu similarity index 100% rename from flashinfer/csrc_rocm/norm_aiter.cu rename to flashinfer/csrc/rocm/norm_aiter.cu diff --git a/flashinfer/csrc_rocm/norm_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/norm_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/norm_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/norm_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/page.cu b/flashinfer/csrc/rocm/page.cu similarity index 100% rename from flashinfer/csrc_rocm/page.cu rename to flashinfer/csrc/rocm/page.cu diff --git a/flashinfer/csrc_rocm/page_aiter.cu b/flashinfer/csrc/rocm/page_aiter.cu similarity index 100% rename from flashinfer/csrc_rocm/page_aiter.cu rename to flashinfer/csrc/rocm/page_aiter.cu diff --git a/flashinfer/csrc_rocm/page_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/page_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/page_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/page_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/pod.cu b/flashinfer/csrc/rocm/pod.cu similarity index 99% rename from flashinfer/csrc_rocm/pod.cu rename to flashinfer/csrc/rocm/pod.cu index 32ea5b9b87..7a786adf5e 100644 --- a/flashinfer/csrc_rocm/pod.cu +++ b/flashinfer/csrc/rocm/pod.cu @@ -7,8 +7,8 @@ #include #include -#include -#include +#include +#include #include #include "pod_config.inc" diff --git a/flashinfer/csrc_rocm/pod_customize_config.jinja b/flashinfer/csrc/rocm/pod_customize_config.jinja similarity index 93% rename from flashinfer/csrc_rocm/pod_customize_config.jinja rename to flashinfer/csrc/rocm/pod_customize_config.jinja index ceb6cc3cd2..6149664d7b 100644 --- a/flashinfer/csrc_rocm/pod_customize_config.jinja +++ b/flashinfer/csrc/rocm/pod_customize_config.jinja @@ -1,8 +1,8 @@ #pragma once -#include -#include -#include -#include +#include +#include +#include +#include #include #include diff --git a/flashinfer/csrc_rocm/pod_jit_pybind.cu b/flashinfer/csrc/rocm/pod_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/pod_jit_pybind.cu rename to flashinfer/csrc/rocm/pod_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/pod_kernel_inst.jinja b/flashinfer/csrc/rocm/pod_kernel_inst.jinja similarity index 100% rename from flashinfer/csrc_rocm/pod_kernel_inst.jinja rename to flashinfer/csrc/rocm/pod_kernel_inst.jinja diff --git a/flashinfer/csrc_rocm/pytorch_conversion_utils.h b/flashinfer/csrc/rocm/pytorch_conversion_utils.h similarity index 100% rename from flashinfer/csrc_rocm/pytorch_conversion_utils.h rename to flashinfer/csrc/rocm/pytorch_conversion_utils.h diff --git a/flashinfer/csrc_rocm/pytorch_extension_utils.h b/flashinfer/csrc/rocm/pytorch_extension_utils.h similarity index 100% rename from flashinfer/csrc_rocm/pytorch_extension_utils.h rename to flashinfer/csrc/rocm/pytorch_extension_utils.h diff --git a/flashinfer/csrc_rocm/quantization.cu b/flashinfer/csrc/rocm/quantization.cu similarity index 98% rename from flashinfer/csrc_rocm/quantization.cu rename to flashinfer/csrc/rocm/quantization.cu index c6916bde53..7d0f81a3c4 100644 --- a/flashinfer/csrc_rocm/quantization.cu +++ b/flashinfer/csrc/rocm/quantization.cu @@ -13,8 +13,8 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include #include -#include #include "pytorch_extension_utils.h" diff --git a/flashinfer/csrc_rocm/renorm.cu b/flashinfer/csrc/rocm/renorm.cu similarity index 100% rename from flashinfer/csrc_rocm/renorm.cu rename to flashinfer/csrc/rocm/renorm.cu diff --git a/flashinfer/csrc_rocm/rope.cu b/flashinfer/csrc/rocm/rope.cu similarity index 100% rename from flashinfer/csrc_rocm/rope.cu rename to flashinfer/csrc/rocm/rope.cu diff --git a/flashinfer/csrc_rocm/rope_aiter.cu b/flashinfer/csrc/rocm/rope_aiter.cu similarity index 100% rename from flashinfer/csrc_rocm/rope_aiter.cu rename to flashinfer/csrc/rocm/rope_aiter.cu diff --git a/flashinfer/csrc_rocm/rope_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/rope_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/rope_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/rope_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/runtime_utils.h b/flashinfer/csrc/rocm/runtime_utils.h similarity index 100% rename from flashinfer/csrc_rocm/runtime_utils.h rename to flashinfer/csrc/rocm/runtime_utils.h diff --git a/flashinfer/csrc_rocm/sampling.cu b/flashinfer/csrc/rocm/sampling.cu similarity index 100% rename from flashinfer/csrc_rocm/sampling.cu rename to flashinfer/csrc/rocm/sampling.cu diff --git a/flashinfer/csrc_rocm/single_decode.cu b/flashinfer/csrc/rocm/single_decode.cu similarity index 100% rename from flashinfer/csrc_rocm/single_decode.cu rename to flashinfer/csrc/rocm/single_decode.cu diff --git a/flashinfer/csrc_rocm/single_decode_customize_config.jinja b/flashinfer/csrc/rocm/single_decode_customize_config.jinja similarity index 95% rename from flashinfer/csrc_rocm/single_decode_customize_config.jinja rename to flashinfer/csrc/rocm/single_decode_customize_config.jinja index 4ca5c9d2dc..f87555d705 100644 --- a/flashinfer/csrc_rocm/single_decode_customize_config.jinja +++ b/flashinfer/csrc/rocm/single_decode_customize_config.jinja @@ -1,6 +1,6 @@ #pragma once -#include -#include +#include +#include #include #include diff --git a/flashinfer/csrc_rocm/single_decode_jit_pybind.cu b/flashinfer/csrc/rocm/single_decode_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/single_decode_jit_pybind.cu rename to flashinfer/csrc/rocm/single_decode_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/single_decode_kernel_inst.jinja b/flashinfer/csrc/rocm/single_decode_kernel_inst.jinja similarity index 100% rename from flashinfer/csrc_rocm/single_decode_kernel_inst.jinja rename to flashinfer/csrc/rocm/single_decode_kernel_inst.jinja diff --git a/flashinfer/csrc_rocm/single_prefill.cu b/flashinfer/csrc/rocm/single_prefill.cu similarity index 98% rename from flashinfer/csrc_rocm/single_prefill.cu rename to flashinfer/csrc/rocm/single_prefill.cu index ee6b55a7bc..e7b7fc216b 100644 --- a/flashinfer/csrc_rocm/single_prefill.cu +++ b/flashinfer/csrc/rocm/single_prefill.cu @@ -16,8 +16,8 @@ #include #include -#include -#include +#include +#include #include #include "pytorch_extension_utils.h" diff --git a/flashinfer/csrc_rocm/single_prefill_aiter.cu b/flashinfer/csrc/rocm/single_prefill_aiter.cu similarity index 98% rename from flashinfer/csrc_rocm/single_prefill_aiter.cu rename to flashinfer/csrc/rocm/single_prefill_aiter.cu index a25c3a15b3..a6aafed27d 100644 --- a/flashinfer/csrc_rocm/single_prefill_aiter.cu +++ b/flashinfer/csrc/rocm/single_prefill_aiter.cu @@ -8,8 +8,8 @@ #include #include -#include -#include +#include +#include #include #include #include diff --git a/flashinfer/csrc_rocm/single_prefill_aiter_jit_pybind.cu b/flashinfer/csrc/rocm/single_prefill_aiter_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/single_prefill_aiter_jit_pybind.cu rename to flashinfer/csrc/rocm/single_prefill_aiter_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/single_prefill_customize_config.jinja b/flashinfer/csrc/rocm/single_prefill_customize_config.jinja similarity index 92% rename from flashinfer/csrc_rocm/single_prefill_customize_config.jinja rename to flashinfer/csrc/rocm/single_prefill_customize_config.jinja index c1c06bbf33..b5bbfdbe2f 100644 --- a/flashinfer/csrc_rocm/single_prefill_customize_config.jinja +++ b/flashinfer/csrc/rocm/single_prefill_customize_config.jinja @@ -1,8 +1,8 @@ #pragma once -#include -#include -#include -#include +#include +#include +#include +#include #include #include diff --git a/flashinfer/csrc_rocm/single_prefill_jit_pybind.cu b/flashinfer/csrc/rocm/single_prefill_jit_pybind.cu similarity index 100% rename from flashinfer/csrc_rocm/single_prefill_jit_pybind.cu rename to flashinfer/csrc/rocm/single_prefill_jit_pybind.cu diff --git a/flashinfer/csrc_rocm/single_prefill_kernel_inst.jinja b/flashinfer/csrc/rocm/single_prefill_kernel_inst.jinja similarity index 100% rename from flashinfer/csrc_rocm/single_prefill_kernel_inst.jinja rename to flashinfer/csrc/rocm/single_prefill_kernel_inst.jinja diff --git a/flashinfer/get_include_paths.py b/flashinfer/get_include_paths.py index ba7e29eee4..7e68596921 100644 --- a/flashinfer/get_include_paths.py +++ b/flashinfer/get_include_paths.py @@ -47,5 +47,5 @@ def get_csrc_dir(): csrc_dir : str Path to flashinfer's C++/ROCm source files. """ - csrc_dir = pathlib.Path(__file__).parent / "csrc_rocm" + csrc_dir = pathlib.Path(__file__).parent / "csrc/rocm" return str(csrc_dir) diff --git a/flashinfer/jit/aiter_source.py b/flashinfer/jit/aiter_source.py index 09a61ee376..e39d892a74 100644 --- a/flashinfer/jit/aiter_source.py +++ b/flashinfer/jit/aiter_source.py @@ -3,7 +3,7 @@ """ Shared plumbing for FlashInfer's C++-level AITER backends (ROCm). -FlashInfer wraps AITER kernels by compiling a small ``csrc_rocm/*_aiter.cu`` shim +FlashInfer wraps AITER kernels by compiling a small ``csrc/rocm/*_aiter.cu`` shim that calls AITER's C++ entry point directly and links the symbol-visible AITER ``.so``. Prefer ``#include``-ing AITER's real header, so a signature change is a compile error rather than a load-time ``undefined symbol``. Fall back to a diff --git a/flashinfer/jit/env.py b/flashinfer/jit/env.py index 6efa97d3bf..8e5ff669fa 100644 --- a/flashinfer/jit/env.py +++ b/flashinfer/jit/env.py @@ -144,7 +144,7 @@ def _get_workspace_dir_name() -> pathlib.Path: if IS_CUDA: # These must stay gated, unlike the helpers above: jit/comm.py imports on # ROCm, so gen_nvshmem_module() would reach them and fail on an absent - # nvidia.nvshmem -- or, with NVSHMEM_* set, on a source csrc_rocm lacks. + # nvidia.nvshmem -- or, with NVSHMEM_* set, on a source csrc/rocm lacks. def get_nvshmem_include_dirs(): paths = os.environ.get("NVSHMEM_INCLUDE_PATH") if paths is not None: diff --git a/flashinfer/jit/rocm/activation.py b/flashinfer/jit/rocm/activation.py index f600889648..7d6c61d00e 100644 --- a/flashinfer/jit/rocm/activation.py +++ b/flashinfer/jit/rocm/activation.py @@ -5,7 +5,7 @@ """The HIP source template for the act-and-mul kernels.""" activation_templ = r""" - #include + #include #include #include "pytorch_extension_utils.h" #include diff --git a/flashinfer/rocm/api.py b/flashinfer/rocm/api.py index 2049d5805a..618e120b60 100644 --- a/flashinfer/rocm/api.py +++ b/flashinfer/rocm/api.py @@ -104,4 +104,4 @@ from ..pod import BatchPODWithPagedKVCacheWrapper as BatchPODWithPagedKVCacheWrapper from ..utils import next_positive_power_of_2 as next_positive_power_of_2 -from ..utils import use_torch_custom_ops_enabled as use_torch_custom_ops_enabled +from .torch_compile import use_torch_custom_ops_enabled as use_torch_custom_ops_enabled diff --git a/flashinfer/rocm/torch_compile.py b/flashinfer/rocm/torch_compile.py new file mode 100644 index 0000000000..31ba97f59e --- /dev/null +++ b/flashinfer/rocm/torch_compile.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Opt-in ``torch.library`` custom-op registration. + +Real registration happens behind FLASHINFER_USE_TORCH_CUSTOM_OPS=1; otherwise +ops carry a guard that raises under torch.compile instead of letting Dynamo +trace into an extension. utils.py keeps upstream's unguarded no-op below +torch 2.4, so the guard only exists from 2.4 up. +""" + +import contextlib +import functools +import os +import warnings +from typing import Callable, Iterable, Optional, Sequence, Union + +import torch +from torch.torch_version import TorchVersion +from torch.torch_version import __version__ as torch_version + +# torch.library.custom_op adds dispatch overhead, which is why upstream leaves it +# off: https://github.com/vllm-project/vllm/blob/36e76700453924c8d421db99af70a88a1df835cd/vllm/utils.py#L1660-L1674 +_USE_TORCH_CUSTOM_OPS = TorchVersion(torch_version) >= TorchVersion( + "2.4" +) and os.environ.get("FLASHINFER_USE_TORCH_CUSTOM_OPS", "").strip().lower() in ( + "1", + "true", + "yes", + "on", +) + + +def use_torch_custom_ops_enabled() -> bool: + """Whether opaque ``torch.library`` custom ops are active.""" + return _USE_TORCH_CUSTOM_OPS + + +def _guard_compile(f: Callable, op_name: str) -> Callable: + """Wrap ``f`` to raise if traced by torch.compile while registration is off.""" + + @functools.wraps(f) + def wrapper(*args, **kwargs): + if torch.compiler.is_compiling(): + raise RuntimeError( + f"torch.compile traced into flashinfer op '{op_name}' but " + "custom ops are not enabled. Set the environment variable " + "FLASHINFER_USE_TORCH_CUSTOM_OPS=1 before importing " + "flashinfer to use torch.compile." + ) + return f(*args, **kwargs) + + return wrapper + + +def register_custom_op( + name: str, + fn: Optional[Callable] = None, + /, + *, + mutates_args: Union[str, Iterable[str]], + device_types: Optional[Union[str, Sequence[str]]] = None, + schema: Optional[str] = None, +) -> Callable: + def decorator(f: Callable) -> Callable: + if not _USE_TORCH_CUSTOM_OPS: + return _guard_compile(f, name) + try: + return torch.library.custom_op( + name, + f, + mutates_args=mutates_args, + device_types=device_types, + schema=schema, + ) + except (ValueError, TypeError): + # Schema inference rejects some parameter types, e.g. + # Optional[torch.Generator]. Fall back to the guard so tracing + # still fails loudly instead of entering the extension. + warnings.warn( + f"Could not register '{name}' as a torch.library custom op " + "(unsupported parameter type in schema inference); falling back " + "to compile guard. torch.compile will raise a RuntimeError if it " + "traces into this op.", + stacklevel=2, + ) + return _guard_compile(f, name) + + if fn is not None: + return decorator(fn) + return decorator + + +def register_fake_op( + name: str, + fn: Optional[Callable] = None, +) -> Callable: + def decorator(f: Callable) -> Callable: + if _USE_TORCH_CUSTOM_OPS: + with contextlib.suppress(Exception): + torch.library.register_fake(name, f) + return f + + if fn is not None: + return decorator(fn) + return decorator diff --git a/flashinfer/utils.py b/flashinfer/utils.py index 6f9fe7e8b9..59d2b0967d 100644 --- a/flashinfer/utils.py +++ b/flashinfer/utils.py @@ -14,11 +14,8 @@ limitations under the License. """ -import contextlib import functools -import warnings import math -import os from enum import Enum from typing import Callable, Dict, Iterable, Optional, Sequence, Tuple, Union @@ -29,6 +26,10 @@ import inspect from .jit.spdlog import gen_spdlog_module +from .rocm import torch_compile as _rocm_torch_compile +from .rocm.torch_compile import ( + use_torch_custom_ops_enabled as use_torch_custom_ops_enabled, +) def plan_info_vec_as_tensor( @@ -331,26 +332,6 @@ def _check_cached_qkv_data_type( ) -# When True, kernels are wrapped in ``torch.library.custom_op`` so ``torch.compile`` / Dynamo -# do not trace into extensions that touch tensor data pointers (see PyTorch custom ops docs). -# Set environment variable ``FLASHINFER_USE_TORCH_CUSTOM_OPS=1`` before importing ``flashinfer``. -# NOTE(Zihao): ``torch.library.custom_op`` adds dispatch overhead; see -# https://github.com/vllm-project/vllm/blob/36e76700453924c8d421db99af70a88a1df835cd/vllm/utils.py#L1660-L1674 -_USE_TORCH_CUSTOM_OPS = TorchVersion(torch_version) >= TorchVersion( - "2.4" -) and os.environ.get("FLASHINFER_USE_TORCH_CUSTOM_OPS", "").strip().lower() in ( - "1", - "true", - "yes", - "on", -) - - -def use_torch_custom_ops_enabled() -> bool: - """Return whether opaque ``torch.library`` custom ops are active (effective behavior).""" - return _USE_TORCH_CUSTOM_OPS - - if TorchVersion(torch_version) < TorchVersion("2.4"): def register_custom_op( @@ -372,23 +353,6 @@ def register_fake_op( else: - def _guard_compile(f: Callable, op_name: str) -> Callable: - """Wrap *f* so it raises a clear error when called under ``torch.compile`` - without ``FLASHINFER_USE_TORCH_CUSTOM_OPS=1``.""" - - @functools.wraps(f) - def wrapper(*args, **kwargs): - if torch.compiler.is_compiling(): - raise RuntimeError( - f"torch.compile traced into flashinfer op '{op_name}' but " - "custom ops are not enabled. Set the environment variable " - "FLASHINFER_USE_TORCH_CUSTOM_OPS=1 before importing " - "flashinfer to use torch.compile." - ) - return f(*args, **kwargs) - - return wrapper - def register_custom_op( name: str, fn: Optional[Callable] = None, @@ -398,48 +362,19 @@ def register_custom_op( device_types: Optional[Union[str, Sequence[str]]] = None, schema: Optional[str] = None, ) -> Callable: - def decorator(f: Callable) -> Callable: - if not _USE_TORCH_CUSTOM_OPS: - return _guard_compile(f, name) - try: - return torch.library.custom_op( - name, - f, - mutates_args=mutates_args, - device_types=device_types, - schema=schema, - ) - except (ValueError, TypeError): - # Some parameter types (e.g. Optional[torch.Generator]) are not - # supported by torch.library.custom_op's schema inference. Fall - # back to the compile guard so torch.compile still raises a - # clear error instead of tracing into the extension. - warnings.warn( - f"Could not register '{name}' as a torch.library custom op " - "(unsupported parameter type in schema inference); falling back " - "to compile guard. torch.compile will raise a RuntimeError if it " - "traces into this op.", - stacklevel=2, - ) - return _guard_compile(f, name) - - if fn is not None: - return decorator(fn) - return decorator + return _rocm_torch_compile.register_custom_op( + name, + fn, + mutates_args=mutates_args, + device_types=device_types, + schema=schema, + ) def register_fake_op( name: str, fn: Optional[Callable] = None, ) -> Callable: - def decorator(f: Callable) -> Callable: - if _USE_TORCH_CUSTOM_OPS: - with contextlib.suppress(Exception): - torch.library.register_fake(name, f) - return f - - if fn is not None: - return decorator(fn) - return decorator + return _rocm_torch_compile.register_fake_op(name, fn) def determine_gemm_backend(device: torch.device) -> str: diff --git a/include/flashinfer/rocm/attention/activation.cuh b/include/flashinfer/rocm/attention/activation.cuh index cbcc1ea718..ab73c19b3f 100644 --- a/include/flashinfer/rocm/attention/activation.cuh +++ b/include/flashinfer/rocm/attention/activation.cuh @@ -2,19 +2,23 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_ACTIVATION_CUH_ -#define FLASHINFER_ACTIVATION_CUH_ +#ifdef FLASHINFER_ACTIVATION_CUH_ +#error \ + "include/flashinfer/activation.cuh and include/flashinfer/rocm/attention/activation.cuh both define FLASHINFER_ACTIVATION_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_ACTIVATION_CUH_ +#define FLASHINFER_ROCM_ATTENTION_ACTIVATION_CUH_ #include -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/utils.cuh" -#include "gpu_iface/vec_dtypes.hpp" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/utils.cuh" +#include "flashinfer/rocm/vec_dtypes_hip.h" namespace flashinfer { -using namespace gpu_iface::vec_dtypes; namespace activation { // Adaptive launch config for act_and_mul_kernel. One block per token underfills @@ -22,7 +26,7 @@ namespace activation { // across blocks_per_row blocks on gridDim.y until the total block count covers // the CU array. For large num_tokens this resolves to blocks_per_row == 1, i.e. // the original one-block-per-token launch. Single definition shared by the AOT -// launcher (flashinfer/csrc_rocm/activation.cu) and the JIT template +// launcher (flashinfer/csrc/rocm/activation.cu) and the JIT template // (flashinfer/jit/activation.py) so the two paths cannot drift. inline void act_and_mul_launch_dims(int d, int64_t num_tokens, uint32_t vec_size, int dev_id, dim3& grid_dim, dim3& block_dim) { @@ -101,4 +105,4 @@ __global__ void act_and_mul_kernel(T* __restrict__ out, const T* __restrict__ in } // namespace activation } // namespace flashinfer -#endif // FLASHINFER_ACTIVATION_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_ACTIVATION_CUH_ diff --git a/include/flashinfer/rocm/attention/allocator.h b/include/flashinfer/rocm/attention/allocator.h index aef9a7acc5..6669cb00ac 100644 --- a/include/flashinfer/rocm/attention/allocator.h +++ b/include/flashinfer/rocm/attention/allocator.h @@ -2,13 +2,18 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_ALLOCATOR_H_ -#define FLASHINFER_ALLOCATOR_H_ +#ifdef FLASHINFER_ALLOCATOR_H_ +#error \ + "include/flashinfer/allocator.h and include/flashinfer/rocm/attention/allocator.h both define FLASHINFER_ALLOCATOR_H_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_ALLOCATOR_H_ +#define FLASHINFER_ROCM_ATTENTION_ALLOCATOR_H_ #include #include -#include "exception.h" +#include "flashinfer/rocm/exception.h" namespace flashinfer { @@ -32,8 +37,9 @@ struct AlignedAllocator { return result; } else { std::ostringstream oss; - oss << "Failed to allocate memory for " << name << " with size " << size << " and alignment " - << alignment << " in AlignedAllocator"; + oss << "Buffer overflow when allocating memory for " << name << " with size " << size + << " and alignment " << alignment << ", but only " << remaining_space + << " bytes available in AlignedAllocator. Increase the workspace buffer size."; FLASHINFER_ERROR(oss.str()); } return nullptr; @@ -48,4 +54,4 @@ struct AlignedAllocator { } // namespace flashinfer -#endif // FLASHINFER_ALLOCATOR_H_ +#endif // FLASHINFER_ROCM_ATTENTION_ALLOCATOR_H_ diff --git a/include/flashinfer/rocm/attention/attention_impl.cuh b/include/flashinfer/rocm/attention/attention_impl.cuh index 05c3d6cd35..dbaaabf031 100644 --- a/include/flashinfer/rocm/attention/attention_impl.cuh +++ b/include/flashinfer/rocm/attention/attention_impl.cuh @@ -3,8 +3,13 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once -#ifndef FLASHINFER_ATTENTION_IMPL_CUH_ -#define FLASHINFER_ATTENTION_IMPL_CUH_ +#ifdef FLASHINFER_ATTENTION_IMPL_CUH_ +#error \ + "include/flashinfer/attention_impl.cuh and include/flashinfer/rocm/attention/attention_impl.cuh both define FLASHINFER_ATTENTION_IMPL_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_ATTENTION_IMPL_CUH_ +#define FLASHINFER_ROCM_ATTENTION_ATTENTION_IMPL_CUH_ #include "cascade.cuh" #include "decode.cuh" @@ -13,4 +18,4 @@ #include "prefill.cuh" #include "variants.cuh" -#endif // FLASHINFER_ATTENTION_IMPL_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_ATTENTION_IMPL_CUH_ diff --git a/include/flashinfer/rocm/attention/batch_pod.cuh b/include/flashinfer/rocm/attention/batch_pod.cuh index ffa1979f66..a63f1ff680 100644 --- a/include/flashinfer/rocm/attention/batch_pod.cuh +++ b/include/flashinfer/rocm/attention/batch_pod.cuh @@ -3,14 +3,18 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#ifdef FLASHINFER_BATCH_POD_CUH_ +#error \ + "include/flashinfer/attention/batch_pod.cuh and include/flashinfer/rocm/attention/batch_pod.cuh both define FLASHINFER_BATCH_POD_CUH_; include only one" +#endif #include "cascade.cuh" -#include "dispatch.cuh" -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/sm_id.hpp" -#include "gpu_iface/utils.cuh" +#include "flashinfer/rocm/dispatch.cuh" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/sm_id.hpp" +#include "flashinfer/rocm/utils.cuh" #include "prefill.cuh" #include "variants.cuh" @@ -46,7 +50,7 @@ __global__ __launch_bounds__(std::max( constexpr int blk_factor_p = 1; constexpr int blk_factor_d = 1; - linear_bid = static_cast(gpu_iface::get_processor_id() % static_cast(num_SMs)); + linear_bid = static_cast(get_processor_id() % static_cast(num_SMs)); const int prefill_slots = (prefill_blocks + blk_factor_p - 1) / blk_factor_p; const int decode_slots = (decode_blocks + blk_factor_d - 1) / blk_factor_d; diff --git a/include/flashinfer/rocm/attention/cascade.cuh b/include/flashinfer/rocm/attention/cascade.cuh index 28b69a96c6..31b6766c1b 100644 --- a/include/flashinfer/rocm/attention/cascade.cuh +++ b/include/flashinfer/rocm/attention/cascade.cuh @@ -3,22 +3,26 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once -#include "gpu_iface/dispatch.cuh" -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/memory_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/utils.cuh" +#ifdef FLASHINFER_CASCADE_CUH_ +#error \ + "include/flashinfer/attention/cascade.cuh and include/flashinfer/rocm/attention/cascade.cuh both define FLASHINFER_CASCADE_CUH_; include only one" +#endif +#include "flashinfer/rocm/dispatch.cuh" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/memory_ops_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/utils.cuh" #include "state.cuh" #ifdef PLATFORM_HIP_DEVICE -#include "gpu_iface/conversion_utils.h" +#include "flashinfer/rocm/conversion_utils.h" #endif namespace flashinfer { -using PrefetchMode = gpu_iface::memory::PrefetchMode; -using SharedMemFillMode = gpu_iface::memory::SharedMemFillMode; +using PrefetchMode = memory::PrefetchMode; +using SharedMemFillMode = memory::SharedMemFillMode; /*! * \brief The kernel that merges the self-attention state of two index sets @@ -48,8 +52,8 @@ __global__ void MergeStateKernel(DTypeIn* __restrict__ v_a, float* __restrict__ float s_a_val = s_a[pos * num_heads + head_idx]; float s_b_val = s_b[pos * num_heads + head_idx]; float s_max = max(s_a_val, s_b_val); - s_a_val = gpu_iface::math::ptx_exp2(s_a_val - s_max); - s_b_val = gpu_iface::math::ptx_exp2(s_b_val - s_max); + s_a_val = math::ptx_exp2(s_a_val - s_max); + s_b_val = math::ptx_exp2(s_b_val - s_max); float a_scale = s_a_val / (s_a_val + s_b_val); float b_scale = s_b_val / (s_a_val + s_b_val); vec_t v_a_vec, v_b_vec, v_merged_vec; @@ -61,7 +65,7 @@ __global__ void MergeStateKernel(DTypeIn* __restrict__ v_a, float* __restrict__ } v_merged_vec.cast_store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); if (s_merged != nullptr) { - s_merged[pos * num_heads + head_idx] = gpu_iface::math::ptx_log2(s_a_val + s_b_val) + s_max; + s_merged[pos * num_heads + head_idx] = math::ptx_log2(s_a_val + s_b_val) + s_max; } } @@ -94,8 +98,8 @@ __global__ void MergeStateInPlaceKernel(DType* __restrict__ v, float* __restrict float s_val = s[pos * num_heads + head_idx]; float s_other_val = s_other[pos * num_heads + head_idx]; float s_max = max(s_val, s_other_val); - s_val = gpu_iface::math::ptx_exp2(s_val - s_max); - s_other_val = gpu_iface::math::ptx_exp2(s_other_val - s_max); + s_val = math::ptx_exp2(s_val - s_max); + s_other_val = math::ptx_exp2(s_other_val - s_max); float scale = s_val / (s_val + s_other_val); float other_scale = s_other_val / (s_val + s_other_val); vec_t v_vec, v_other_vec; @@ -107,7 +111,7 @@ __global__ void MergeStateInPlaceKernel(DType* __restrict__ v, float* __restrict } v_vec.cast_store(v + (pos * num_heads + head_idx) * head_dim + tx * vec_size); if (s != nullptr) { - s[pos * num_heads + head_idx] = gpu_iface::math::ptx_log2(s_val + s_other_val) + s_max; + s[pos * num_heads + head_idx] = math::ptx_log2(s_val + s_other_val) + s_max; } } @@ -188,7 +192,7 @@ __global__ void AttentionSumKernel(DTypeIn* __restrict__ V, DTypeO* __restrict__ if (num_index_sets == 0) { vec_t v; #ifdef PLATFORM_HIP_DEVICE - v.fill(fi::con::explicit_casting(0.f)); + v.fill(explicit_casting(0.f)); #else v.fill(DTypeO(0.f)); #endif @@ -230,13 +234,13 @@ __global__ void MergeStatesKernel(DTypeIn* __restrict__ V, float* __restrict__ S if (num_index_sets == 0) { vec_t v; #ifdef PLATFORM_HIP_DEVICE - v.fill(fi::con::explicit_casting(0.f)); + v.fill(explicit_casting(0.f)); #else v.fill(DTypeO(0.f)); #endif v.store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); if (s_merged != nullptr) { - s_merged[pos * num_heads + head_idx] = -gpu_iface::math::inf; + s_merged[pos * num_heads + head_idx] = -math::inf; } return; } @@ -305,12 +309,12 @@ __global__ void MergeStatesLargeNumIndexSetsKernel(DTypeIn* __restrict__ V, floa #pragma unroll for (uint32_t iter = 0; iter < num_smem_stages; ++iter) { - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + (iter * bdy + ty) * head_dim + tx * vec_size, V + ((pos * num_index_sets + (iter * bdy + ty)) * num_heads + head_idx) * head_dim + tx * vec_size, (iter * bdy + ty) < num_index_sets); - gpu_iface::memory::commit_group(); + memory::commit_group(); } #pragma unroll 4 for (uint32_t iter = 0; iter < ceil_div(num_index_sets, bdy); ++iter) { @@ -321,7 +325,7 @@ __global__ void MergeStatesLargeNumIndexSetsKernel(DTypeIn* __restrict__ V, floa : 0.f; __syncthreads(); } - gpu_iface::memory::wait_group(); + memory::wait_group(); __syncthreads(); vec_t v; v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size); @@ -330,7 +334,7 @@ __global__ void MergeStatesLargeNumIndexSetsKernel(DTypeIn* __restrict__ V, floa st.merge(v, s, 1); } __syncthreads(); - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size, V + ((pos * num_index_sets + ((iter + num_smem_stages) * bdy + ty)) * num_heads + @@ -338,9 +342,9 @@ __global__ void MergeStatesLargeNumIndexSetsKernel(DTypeIn* __restrict__ V, floa head_dim + tx * vec_size, (iter + num_smem_stages) * bdy + ty < num_index_sets); - gpu_iface::memory::commit_group(); + memory::commit_group(); } - gpu_iface::memory::wait_group<0>(); + memory::wait_group<0>(); __syncthreads(); st.normalize(); @@ -408,13 +412,13 @@ __global__ void PersistentVariableLengthMergeStatesKernel( if (num_index_sets == 0) { vec_t v; #ifdef PLATFORM_HIP_DEVICE - v.fill(fi::con::explicit_casting(0.f)); + v.fill(explicit_casting(0.f)); #else v.fill(DTypeO(0.f)); #endif v.store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); if (s_merged != nullptr) { - s_merged[pos * num_heads + head_idx] = -gpu_iface::math::inf; + s_merged[pos * num_heads + head_idx] = -math::inf; } continue; } @@ -431,11 +435,11 @@ __global__ void PersistentVariableLengthMergeStatesKernel( #pragma unroll for (uint32_t iter = 0; iter < num_smem_stages; ++iter) { - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + (iter * bdy + ty) * head_dim + tx * vec_size, V + ((indptr[pos] + (iter * bdy + ty)) * num_heads + head_idx) * head_dim + tx * vec_size, (iter * bdy + ty) < num_index_sets); - gpu_iface::memory::commit_group(); + memory::commit_group(); } #pragma unroll 4 for (uint32_t iter = 0; iter < ceil_div(num_index_sets, bdy); ++iter) { @@ -446,7 +450,7 @@ __global__ void PersistentVariableLengthMergeStatesKernel( : 0.f; __syncthreads(); } - gpu_iface::memory::wait_group(); + memory::wait_group(); __syncthreads(); vec_t v; v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size); @@ -455,16 +459,16 @@ __global__ void PersistentVariableLengthMergeStatesKernel( st.merge(v, s, 1); } __syncthreads(); - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size, V + ((indptr[pos] + ((iter + num_smem_stages) * bdy + ty)) * num_heads + head_idx) * head_dim + tx * vec_size, (iter + num_smem_stages) * bdy + ty < num_index_sets); - gpu_iface::memory::commit_group(); + memory::commit_group(); } - gpu_iface::memory::wait_group<0>(); + memory::wait_group<0>(); __syncthreads(); st.normalize(); @@ -509,7 +513,7 @@ __global__ void PersistentVariableLengthAttentionSumKernel(DTypeIn* __restrict__ if (num_index_sets == 0) { vec_t v; #ifdef PLATFORM_HIP_DEVICE - v.fill(fi::con::explicit_casting(0.f)); + v.fill(explicit_casting(0.f)); #else v.fill(DTypeO(0.f)); #endif @@ -526,15 +530,15 @@ __global__ void PersistentVariableLengthAttentionSumKernel(DTypeIn* __restrict__ #pragma unroll for (uint32_t iter = 0; iter < num_smem_stages; ++iter) { - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + (iter * bdy + ty) * head_dim + tx * vec_size, V + ((indptr[pos] + (iter * bdy + ty)) * num_heads + head_idx) * head_dim + tx * vec_size, (iter * bdy + ty) < num_index_sets); - gpu_iface::memory::commit_group(); + memory::commit_group(); } #pragma unroll 4 for (uint32_t iter = 0; iter < ceil_div(num_index_sets, bdy); ++iter) { - gpu_iface::memory::wait_group(); + memory::wait_group(); __syncthreads(); vec_t v; v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size); @@ -545,16 +549,16 @@ __global__ void PersistentVariableLengthAttentionSumKernel(DTypeIn* __restrict__ } } __syncthreads(); - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size, V + ((indptr[pos] + ((iter + num_smem_stages) * bdy + ty)) * num_heads + head_idx) * head_dim + tx * vec_size, (iter + num_smem_stages) * bdy + ty < num_index_sets); - gpu_iface::memory::commit_group(); + memory::commit_group(); } - gpu_iface::memory::wait_group<0>(); + memory::wait_group<0>(); __syncthreads(); threadblock_sum(v_sum_vec, v_smem); diff --git a/include/flashinfer/rocm/attention/decode.cuh b/include/flashinfer/rocm/attention/decode.cuh index 13851e5a00..184612c063 100644 --- a/include/flashinfer/rocm/attention/decode.cuh +++ b/include/flashinfer/rocm/attention/decode.cuh @@ -2,19 +2,24 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_DECODE_CUH_ -#define FLASHINFER_DECODE_CUH_ +#ifdef FLASHINFER_DECODE_CUH_ +#error \ + "include/flashinfer/attention/decode.cuh and include/flashinfer/rocm/attention/decode.cuh both define FLASHINFER_DECODE_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_DECODE_CUH_ +#define FLASHINFER_ROCM_ATTENTION_DECODE_CUH_ #include #include "cascade.cuh" #include "decode_tuning.cuh" -#include "gpu_iface/cooperative_groups.h" -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/memory_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/utils.cuh" +#include "flashinfer/rocm/cooperative_groups.h" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/memory_ops_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/utils.cuh" #include "pos_enc.cuh" #include "state.cuh" @@ -35,9 +40,8 @@ inline void CheckSmemBudget(uint32_t smem_size, int dev_id) { } namespace cg = cooperative_groups; -using PrefetchMode = gpu_iface::memory::PrefetchMode; -using SharedMemFillMode = gpu_iface::memory::SharedMemFillMode; -using namespace gpu_iface::vec_dtypes; +using PrefetchMode = memory::PrefetchMode; +using SharedMemFillMode = memory::SharedMemFillMode; namespace { /*! @@ -279,23 +283,23 @@ __global__ void SingleDecodeWithKVCacheKernel(const Params params) { #pragma unroll for (uint32_t iter = 0; iter < num_stages_smem; ++iter) { for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { - gpu_iface::memory::pred_load( + memory::pred_load( k_smem + (((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + tx * vec_size, k + (producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j) * kv_stride_n + kv_head_idx * kv_stride_h + tx * vec_size, producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j < chunk_end); } - gpu_iface::memory::commit_group(); + memory::commit_group(); for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + (((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + tx * vec_size, v + (producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j) * kv_stride_n + kv_head_idx * kv_stride_h + tx * vec_size, producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j < chunk_end); } - gpu_iface::memory::commit_group(); + memory::commit_group(); producer_kv_idx_base += bdy * bdz * tile_size_per_bdx; } @@ -307,7 +311,7 @@ __global__ void SingleDecodeWithKVCacheKernel(const Params params) { #pragma unroll 2 for (uint32_t iter = 0; iter < ceil_div(kv_chunk_size, tile_size_per_bdx * bdy * bdz); ++iter) { // compute qk - gpu_iface::memory::wait_group<2 * num_stages_smem - 1>(); + memory::wait_group<2 * num_stages_smem - 1>(); block.sync(); compute_qk( params, variant, /*batch_idx=*/0, @@ -317,17 +321,17 @@ __global__ void SingleDecodeWithKVCacheKernel(const Params params) { block.sync(); // load k for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { - gpu_iface::memory::pred_load( + memory::pred_load( k_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + tx * vec_size, k + (producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j) * kv_stride_n + kv_head_idx * kv_stride_h + tx * vec_size, producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j < chunk_end); } - gpu_iface::memory::commit_group(); + memory::commit_group(); // update m/d/o state - gpu_iface::memory::wait_group<2 * num_stages_smem - 1>(); + memory::wait_group<2 * num_stages_smem - 1>(); block.sync(); update_local_state( v_smem + (stage_idx * bdz + tz) * bdy * tile_size_per_bdx * head_dim, s, stage_idx, @@ -336,20 +340,20 @@ __global__ void SingleDecodeWithKVCacheKernel(const Params params) { // load v for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + tx * vec_size, v + (producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j) * kv_stride_n + kv_head_idx * kv_stride_h + tx * vec_size, producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j < chunk_end); } - gpu_iface::memory::commit_group(); + memory::commit_group(); stage_idx = (stage_idx + 1) % num_stages_smem; producer_kv_idx_base += tile_size_per_bdx * bdy * bdz; consumer_kv_idx_base += tile_size_per_bdx * bdy * bdz; } - gpu_iface::memory::wait_group<0>(); + memory::wait_group<0>(); block.sync(); // sync local state of all warps inside a threadblock @@ -489,22 +493,22 @@ __device__ __inline__ void BatchDecodeWithPagedKVCacheDevice(const Params& param } #pragma unroll for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { - gpu_iface::memory::pred_load( + memory::pred_load( k_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + tx * vec_size, paged_kv.k_data + kv_offset[j], ((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j < chunk_size); } - gpu_iface::memory::commit_group(); + memory::commit_group(); #pragma unroll for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + tx * vec_size, paged_kv.v_data + kv_offset[j], ((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j < chunk_size); } - gpu_iface::memory::commit_group(); + memory::commit_group(); stage_idx = (stage_idx + 1) % num_stages_smem; } @@ -526,7 +530,7 @@ __device__ __inline__ void BatchDecodeWithPagedKVCacheDevice(const Params& param } } // compute qk - gpu_iface::memory::wait_group<2 * num_stages_smem - 1>(); + memory::wait_group<2 * num_stages_smem - 1>(); block.sync(); compute_qk( params, variant, batch_idx, @@ -548,16 +552,16 @@ __device__ __inline__ void BatchDecodeWithPagedKVCacheDevice(const Params& param // load k tiles #pragma unroll for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { - gpu_iface::memory::pred_load( + memory::pred_load( k_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + tx * vec_size, paged_kv.k_data + kv_offset[j], (((iter + num_stages_smem) * bdz + tz) * bdy + ty) * tile_size_per_bdx + j < chunk_size); } - gpu_iface::memory::commit_group(); + memory::commit_group(); // update m/d/o states - gpu_iface::memory::wait_group<2 * num_stages_smem - 1>(); + memory::wait_group<2 * num_stages_smem - 1>(); block.sync(); update_local_state( v_smem + (stage_idx * bdz + tz) * bdy * tile_size_per_bdx * head_dim, s, stage_idx, st, tx); @@ -566,16 +570,16 @@ __device__ __inline__ void BatchDecodeWithPagedKVCacheDevice(const Params& param // load v tiles #pragma unroll for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { - gpu_iface::memory::pred_load( + memory::pred_load( v_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + tx * vec_size, paged_kv.v_data + kv_offset[j], (((iter + num_stages_smem) * bdz + tz) * bdy + ty) * tile_size_per_bdx + j < chunk_size); } - gpu_iface::memory::commit_group(); + memory::commit_group(); stage_idx = (stage_idx + 1) % num_stages_smem; } - gpu_iface::memory::wait_group<0>(); + memory::wait_group<0>(); block.sync(); // sync local state of all warps inside a threadblock @@ -676,7 +680,7 @@ gpuError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTy // This has been hard coded to 2U. Previous implementation involved a macro redirection that // always resulted in 2U for H100 or CDNA3 architecture. Please take a look at - // gpu_iface/dispatch.cuh - DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM macro + // flashinfer/rocm/dispatch.cuh - DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM macro constexpr uint32_t NUM_STAGES_SMEM = 2U; const uint32_t smem_size = @@ -990,25 +994,25 @@ __global__ void BatchDecodeWithPagedKVCacheKernelMLA(Params params) { is_valid_range = (iter * kv_iter_len + dim2_offset(bdy, tz, ty)) < cur_chunk_len; offset_bytes = ckv_offset_smem[dim3_offset(bdz, bdy, iter, tz, ty)] + tx * vec_size_ckv; - gpu_iface::memory::pred_load( + memory::pred_load( ckv_smem + (stage_idx * kv_iter_len + dim2_offset(bdy, tz, ty)) * head_dim_ckv + tx * vec_size_ckv, paged_kv.ckv_data + offset_bytes, is_valid_range); offset_bytes = kpe_offset_smem[dim3_offset(bdz, bdy, iter, tz, ty)] + tx / tx_fold * vec_size_ckv; - gpu_iface::memory::pred_load( + memory::pred_load( kpe_smem + (stage_idx * kv_iter_len + dim2_offset(bdy, tz, ty)) * head_dim_kpe + tx / tx_fold * vec_size_ckv, paged_kv.kpe_data + offset_bytes, is_valid_range); - gpu_iface::memory::commit_group(); + memory::commit_group(); stage_idx = (stage_idx + 1) % num_stages_smem; } #pragma unroll for (uint32_t iter = 0; iter < ceil_div(cur_chunk_len, kv_iter_len); ++iter) { - gpu_iface::memory::wait_group<1 * num_stages_smem - 1>(); + memory::wait_group<1 * num_stages_smem - 1>(); block.sync(); const int32_t kv_idx_base = (paged_kv.rope_pos_offset == nullptr ? 0 : paged_kv.rope_pos_offset[mapped_batch_idx]) + @@ -1038,22 +1042,22 @@ __global__ void BatchDecodeWithPagedKVCacheKernelMLA(Params params) { ((iter + num_stages_smem) * kv_iter_len + dim2_offset(bdy, tz, ty)) < cur_chunk_len; offset_bytes = ckv_offset_smem[dim3_offset(bdz, bdy, (iter + num_stages_smem) % bdx, tz, ty)] + tx * vec_size_ckv; - gpu_iface::memory::pred_load( + memory::pred_load( ckv_smem + (stage_idx * kv_iter_len + dim2_offset(bdy, tz, ty)) * head_dim_ckv + tx * vec_size_ckv, paged_kv.ckv_data + offset_bytes, is_valid_range); offset_bytes = kpe_offset_smem[dim3_offset(bdz, bdy, (iter + num_stages_smem) % bdx, tz, ty)] + tx / tx_fold * vec_size_ckv; - gpu_iface::memory::pred_load( + memory::pred_load( kpe_smem + (stage_idx * kv_iter_len + dim2_offset(bdy, tz, ty)) * head_dim_kpe + tx / tx_fold * vec_size_ckv, paged_kv.kpe_data + offset_bytes, is_valid_range); - gpu_iface::memory::commit_group(); + memory::commit_group(); stage_idx = (stage_idx + 1) % num_stages_smem; } - gpu_iface::memory::wait_group<0>(); + memory::wait_group<0>(); block.sync(); if (bdz != 1) { @@ -1141,4 +1145,4 @@ gpuError_t BatchDecodeWithPagedKVCacheDispatchedMLA(Params params, typename Para } // namespace flashinfer -#endif // FLASHINFER_DECODE_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_DECODE_CUH_ diff --git a/include/flashinfer/rocm/attention/default_decode_params.cuh b/include/flashinfer/rocm/attention/default_decode_params.cuh index 78cbc2f31f..9f3af76665 100644 --- a/include/flashinfer/rocm/attention/default_decode_params.cuh +++ b/include/flashinfer/rocm/attention/default_decode_params.cuh @@ -3,13 +3,18 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once -#ifndef FLASHINFER_DECODE_PARAMS_CUH_ -#define FLASHINFER_DECODE_PARAMS_CUH_ +#ifdef FLASHINFER_DECODE_PARAMS_CUH_ +#error \ + "include/flashinfer/attention/default_decode_params.cuh and include/flashinfer/rocm/attention/default_decode_params.cuh both define FLASHINFER_DECODE_PARAMS_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_DEFAULT_DECODE_PARAMS_CUH_ +#define FLASHINFER_ROCM_ATTENTION_DEFAULT_DECODE_PARAMS_CUH_ #include -#include "gpu_iface/layout.cuh" -#include "gpu_iface/platform.hpp" +#include "flashinfer/rocm/layout.cuh" +#include "flashinfer/rocm/platform.hpp" #include "page.cuh" namespace flashinfer { @@ -265,4 +270,4 @@ struct BatchDecodeParamsMLA { } // namespace flashinfer -#endif // FLASHINFER_DECODE_PARAMS_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_DEFAULT_DECODE_PARAMS_CUH_ diff --git a/include/flashinfer/rocm/attention/default_prefill_params.cuh b/include/flashinfer/rocm/attention/default_prefill_params.cuh index 5f59fb3720..06a513f071 100644 --- a/include/flashinfer/rocm/attention/default_prefill_params.cuh +++ b/include/flashinfer/rocm/attention/default_prefill_params.cuh @@ -1,13 +1,18 @@ // SPDX-FileCopyrightText: 2023-2025 FlashInfer team. // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_PREFILL_PARAMS_CUH_ -#define FLASHINFER_PREFILL_PARAMS_CUH_ +#ifdef FLASHINFER_PREFILL_PARAMS_CUH_ +#error \ + "include/flashinfer/attention/default_prefill_params.cuh and include/flashinfer/rocm/attention/default_prefill_params.cuh both define FLASHINFER_PREFILL_PARAMS_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_DEFAULT_PREFILL_PARAMS_CUH_ +#define FLASHINFER_ROCM_ATTENTION_DEFAULT_PREFILL_PARAMS_CUH_ #include #include -#include "gpu_iface/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" #include "page.cuh" namespace flashinfer { @@ -377,4 +382,4 @@ struct BatchPrefillPagedParams { } // namespace flashinfer -#endif // FLASHINFER_PREFILL_PARAMS_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_DEFAULT_PREFILL_PARAMS_CUH_ diff --git a/include/flashinfer/rocm/attention/dispatch.cuh b/include/flashinfer/rocm/attention/dispatch.cuh deleted file mode 100644 index 8cfca26d28..0000000000 --- a/include/flashinfer/rocm/attention/dispatch.cuh +++ /dev/null @@ -1,216 +0,0 @@ -// SPDX-FileCopyrightText: 2023-2025 FlashInfer team. -// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once - -#include "gpu_iface/enums.hpp" -#include "gpu_iface/exception.h" - -#define DISPATCH_USE_FP16_QK_REDUCTION(use_fp16_qk_reduction, USE_FP16_QK_REDUCTION, ...) \ - if (use_fp16_qk_reduction) { \ - FLASHINFER_ERROR("FP16_QK_REDUCTION disabled at compile time"); \ - } else { \ - constexpr bool USE_FP16_QK_REDUCTION = false; \ - __VA_ARGS__ \ - } - -#define DISPATCH_NUM_MMA_Q(num_mma_q, NUM_MMA_Q, ...) \ - if (num_mma_q == 1) { \ - constexpr size_t NUM_MMA_Q = 1; \ - __VA_ARGS__ \ - } else if (num_mma_q == 2) { \ - constexpr size_t NUM_MMA_Q = 2; \ - __VA_ARGS__ \ - } else { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported num_mma_q: " << num_mma_q; \ - FLASHINFER_ERROR(err_msg.str()); \ - } - -#define DISPATCH_NUM_MMA_KV(max_mma_kv, NUM_MMA_KV, ...) \ - if (max_mma_kv >= 8) { \ - constexpr size_t NUM_MMA_KV = 8; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 4) { \ - constexpr size_t NUM_MMA_KV = 4; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 2) { \ - constexpr size_t NUM_MMA_KV = 2; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 1) { \ - constexpr size_t NUM_MMA_KV = 1; \ - __VA_ARGS__ \ - } else { \ - /* Fallback for AMD GPUs with tight shared memory constraints */ \ - /* Use NUM_MMA_KV=1 and issue warning */ \ - constexpr size_t NUM_MMA_KV = 1; \ - __VA_ARGS__ \ - } - -#define DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, ...) \ - switch (cta_tile_q) { \ - case 128: { \ - constexpr uint32_t CTA_TILE_Q = 128; \ - __VA_ARGS__ \ - break; \ - } \ - case 64: { \ - constexpr uint32_t CTA_TILE_Q = 64; \ - __VA_ARGS__ \ - break; \ - } \ - case 16: { \ - constexpr uint32_t CTA_TILE_Q = 16; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported cta_tile_q: " << cta_tile_q; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_GQA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \ - if (group_size == 1) { \ - constexpr size_t GROUP_SIZE = 1; \ - __VA_ARGS__ \ - } else if (group_size == 2) { \ - constexpr size_t GROUP_SIZE = 2; \ - __VA_ARGS__ \ - } else if (group_size == 3) { \ - constexpr size_t GROUP_SIZE = 3; \ - __VA_ARGS__ \ - } else if (group_size == 4) { \ - constexpr size_t GROUP_SIZE = 4; \ - __VA_ARGS__ \ - } else if (group_size == 8) { \ - constexpr size_t GROUP_SIZE = 8; \ - __VA_ARGS__ \ - } else { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported group_size: " << group_size; \ - FLASHINFER_ERROR(err_msg.str()); \ - } - -#define DISPATCH_MASK_MODE(mask_mode, MASK_MODE, ...) \ - switch (mask_mode) { \ - case MaskMode::kNone: { \ - constexpr MaskMode MASK_MODE = MaskMode::kNone; \ - __VA_ARGS__ \ - break; \ - } \ - case MaskMode::kCausal: { \ - constexpr MaskMode MASK_MODE = MaskMode::kCausal; \ - __VA_ARGS__ \ - break; \ - } \ - case MaskMode::kCustom: { \ - constexpr MaskMode MASK_MODE = MaskMode::kCustom; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported mask_mode: " << int(mask_mode); \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -// convert head_dim to compile-time constant -#define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...) \ - switch (head_dim) { \ - case 64: { \ - constexpr size_t HEAD_DIM = 64; \ - __VA_ARGS__ \ - break; \ - } \ - case 128: { \ - constexpr size_t HEAD_DIM = 128; \ - __VA_ARGS__ \ - break; \ - } \ - case 256: { \ - constexpr size_t HEAD_DIM = 256; \ - __VA_ARGS__ \ - break; \ - } \ - case 512: { \ - constexpr size_t HEAD_DIM = 512; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported head_dim: " << head_dim; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_POS_ENCODING_MODE(pos_encoding_mode, POS_ENCODING_MODE, ...) \ - switch (pos_encoding_mode) { \ - case PosEncodingMode::kNone: { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kNone; \ - __VA_ARGS__ \ - break; \ - } \ - case PosEncodingMode::kRoPELlama: { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kRoPELlama; \ - __VA_ARGS__ \ - break; \ - } \ - case PosEncodingMode::kALiBi: { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kALiBi; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported pos_encoding_mode: " << int(pos_encoding_mode); \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_ALIGNED_VEC_SIZE(aligned_vec_size, ALIGNED_VEC_SIZE, ...) \ - switch (aligned_vec_size) { \ - case 16: { \ - constexpr size_t ALIGNED_VEC_SIZE = 16; \ - __VA_ARGS__ \ - break; \ - } \ - case 8: { \ - constexpr size_t ALIGNED_VEC_SIZE = 8; \ - __VA_ARGS__ \ - break; \ - } \ - case 4: { \ - constexpr size_t ALIGNED_VEC_SIZE = 4; \ - __VA_ARGS__ \ - break; \ - } \ - case 2: { \ - constexpr size_t ALIGNED_VEC_SIZE = 2; \ - __VA_ARGS__ \ - break; \ - } \ - case 1: { \ - constexpr size_t ALIGNED_VEC_SIZE = 1; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported aligned_vec_size: " << aligned_vec_size; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, ...) \ - if (compute_capacity.first >= 8) { \ - constexpr uint32_t NUM_STAGES_SMEM = 2; \ - __VA_ARGS__ \ - } else { \ - constexpr uint32_t NUM_STAGES_SMEM = 1; \ - __VA_ARGS__ \ - } diff --git a/include/flashinfer/rocm/attention/frag_layout_swizzle.cuh b/include/flashinfer/rocm/attention/frag_layout_swizzle.cuh index 6f66c7ee65..cabdf00cb3 100644 --- a/include/flashinfer/rocm/attention/frag_layout_swizzle.cuh +++ b/include/flashinfer/rocm/attention/frag_layout_swizzle.cuh @@ -2,12 +2,17 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_FRAG_LAYOUT_SWIZZLE_CUH_ -#define FLASHINFER_FRAG_LAYOUT_SWIZZLE_CUH_ +#ifdef FLASHINFER_FRAG_LAYOUT_SWIZZLE_CUH_ +#error \ + "include/flashinfer/frag_layout_swizzle.cuh and include/flashinfer/rocm/attention/frag_layout_swizzle.cuh both define FLASHINFER_FRAG_LAYOUT_SWIZZLE_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_FRAG_LAYOUT_SWIZZLE_CUH_ +#define FLASHINFER_ROCM_ATTENTION_FRAG_LAYOUT_SWIZZLE_CUH_ #include -#include "gpu_iface/platform.hpp" +#include "flashinfer/rocm/platform.hpp" // Define platform-specific full mask for warp/wavefront operations constexpr uint64_t WARP_FULL_MASK = 0xffffffffffffffffULL; // 64-bit mask for HIP @@ -30,4 +35,4 @@ __device__ __forceinline__ uint32_t frag_layout_swizzle_16b_to_8b_trans(uint32_t return x; } -#endif // FLASHINFER_FRAG_LAYOUT_SWIZZLE_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_FRAG_LAYOUT_SWIZZLE_CUH_ diff --git a/include/flashinfer/rocm/attention/heap.h b/include/flashinfer/rocm/attention/heap.h index 1bdcf8aa13..373e14f352 100644 --- a/include/flashinfer/rocm/attention/heap.h +++ b/include/flashinfer/rocm/attention/heap.h @@ -2,8 +2,13 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_ATTENTION_HEAP_H -#define FLASHINFER_ATTENTION_HEAP_H +#ifdef FLASHINFER_ATTENTION_HEAP_H +#error \ + "include/flashinfer/attention/heap.h and include/flashinfer/rocm/attention/heap.h both define FLASHINFER_ATTENTION_HEAP_H; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_HEAP_H_ +#define FLASHINFER_ROCM_ATTENTION_HEAP_H_ #include #include @@ -52,4 +57,4 @@ class MinHeap { } // namespace flashinfer -#endif // FLASHINFER_ATTENTION_HEAP_H +#endif // FLASHINFER_ROCM_ATTENTION_HEAP_H_ diff --git a/include/flashinfer/rocm/attention/norm.cuh b/include/flashinfer/rocm/attention/norm.cuh index df2056ac07..484ac004eb 100644 --- a/include/flashinfer/rocm/attention/norm.cuh +++ b/include/flashinfer/rocm/attention/norm.cuh @@ -3,20 +3,24 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once -#ifndef FLASHINFER_NORM_CUH_ -#define FLASHINFER_NORM_CUH_ +#ifdef FLASHINFER_NORM_CUH_ +#error \ + "include/flashinfer/norm.cuh and include/flashinfer/rocm/attention/norm.cuh both define FLASHINFER_NORM_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_NORM_CUH_ +#define FLASHINFER_ROCM_ATTENTION_NORM_CUH_ #include -#include "gpu_iface/dispatch.cuh" -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/utils.cuh" -#include "gpu_iface/vec_dtypes.hpp" +#include "flashinfer/rocm/dispatch.cuh" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/utils.cuh" +#include "flashinfer/rocm/vec_dtypes_hip.h" namespace flashinfer { -using namespace gpu_iface::vec_dtypes; namespace norm { // Threads per lane group. Must divide the wavefront so a group's shuffle @@ -24,8 +28,8 @@ namespace norm { constexpr uint32_t kLaneGroupSize = 32; constexpr uint32_t kMaxBlockSize = 1024; -static_assert(kLaneGroupSize <= static_cast(gpu_iface::kWarpSize) && - static_cast(gpu_iface::kWarpSize) % kLaneGroupSize == 0); +static_assert(kLaneGroupSize <= static_cast(kWarpSize) && + static_cast(kWarpSize) % kLaneGroupSize == 0); // Stage 2 folds num_warps partial sums inside a single lane group. Ceiling // division, matching how the launchers derive num_warps. static_assert((kMaxBlockSize + kLaneGroupSize - 1) / kLaneGroupSize <= kLaneGroupSize); @@ -332,4 +336,4 @@ gpuError_t GemmaFusedAddRMSNorm(T* input, T* residual, T* weight, uint32_t batch } // namespace flashinfer -#endif // FLASHINFER_NORM_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_NORM_CUH_ diff --git a/include/flashinfer/rocm/attention/page.cuh b/include/flashinfer/rocm/attention/page.cuh index 98a6c166d6..7f439755fb 100644 --- a/include/flashinfer/rocm/attention/page.cuh +++ b/include/flashinfer/rocm/attention/page.cuh @@ -3,19 +3,22 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#ifdef FLASHINFER_PAGE_CUH_ +#error \ + "include/flashinfer/page.cuh and include/flashinfer/rocm/attention/page.cuh both define FLASHINFER_PAGE_CUH_; include only one" +#endif #include -#include "gpu_iface/dispatch.cuh" -#include "gpu_iface/exception.h" -#include "gpu_iface/fastdiv.cuh" -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/layout.cuh" -#include "gpu_iface/utils.cuh" -#include "gpu_iface/vec_dtypes.hpp" +#include "flashinfer/rocm/dispatch.cuh" +#include "flashinfer/rocm/exception.h" +#include "flashinfer/rocm/fastdiv.cuh" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/layout.cuh" +#include "flashinfer/rocm/utils.cuh" +#include "flashinfer/rocm/vec_dtypes_hip.h" namespace flashinfer { -using namespace gpu_iface::vec_dtypes; /*! * \brief Paged key-value cache * \tparam layout The layout of last 3 dimensions in KV-Cache. diff --git a/include/flashinfer/rocm/attention/permuted_smem.cuh b/include/flashinfer/rocm/attention/permuted_smem.cuh index 03f7e7493c..a21a4e6ea4 100644 --- a/include/flashinfer/rocm/attention/permuted_smem.cuh +++ b/include/flashinfer/rocm/attention/permuted_smem.cuh @@ -2,14 +2,19 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_PERMUTED_SMEM_CUH_ -#define FLASHINFER_PERMUTED_SMEM_CUH_ +#ifdef FLASHINFER_PERMUTED_SMEM_CUH_ +#error \ + "include/flashinfer/permuted_smem.cuh and include/flashinfer/rocm/attention/permuted_smem.cuh both define FLASHINFER_PERMUTED_SMEM_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_PERMUTED_SMEM_CUH_ +#define FLASHINFER_ROCM_ATTENTION_PERMUTED_SMEM_CUH_ -#include "gpu_iface/memory_ops.hpp" -#include "gpu_iface/mma_ops.hpp" -#include "gpu_iface/platform.hpp" +#include "flashinfer/rocm/memory_ops_hip.h" +#include "flashinfer/rocm/mma_hip.h" +#include "flashinfer/rocm/platform.hpp" -namespace gpu_mem = flashinfer::gpu_iface::memory; +namespace gpu_mem = flashinfer::memory; namespace flashinfer { @@ -239,7 +244,7 @@ struct smem_t { template __device__ __forceinline__ void load_matrix_m16n16_trans(uint32_t offset, T* frag) { load_fragment(offset, frag); - gpu_iface::mma::transpose_mma_tile(frag); + mma_hip::transpose_mma_tile(frag); } #endif @@ -322,4 +327,4 @@ struct smem_t { } // namespace flashinfer -#endif // FLASHINFER_PERMUTED_SMEM_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_PERMUTED_SMEM_CUH_ diff --git a/include/flashinfer/rocm/attention/pod.cuh b/include/flashinfer/rocm/attention/pod.cuh index b68d131768..4fb8395d5c 100644 --- a/include/flashinfer/rocm/attention/pod.cuh +++ b/include/flashinfer/rocm/attention/pod.cuh @@ -3,14 +3,18 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#ifdef FLASHINFER_POD_CUH_ +#error \ + "include/flashinfer/attention/pod.cuh and include/flashinfer/rocm/attention/pod.cuh both define FLASHINFER_POD_CUH_; include only one" +#endif #include "cascade.cuh" -#include "dispatch.cuh" -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/sm_id.hpp" -#include "gpu_iface/utils.cuh" +#include "flashinfer/rocm/dispatch.cuh" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/sm_id.hpp" +#include "flashinfer/rocm/utils.cuh" #include "prefill.cuh" #include "variants.cuh" @@ -48,7 +52,7 @@ __global__ __launch_bounds__(std::max( constexpr int blk_factor_p = 1; constexpr int blk_factor_d = 1; - linear_bid = static_cast(gpu_iface::get_processor_id() % static_cast(num_SMs)); + linear_bid = static_cast(get_processor_id() % static_cast(num_SMs)); const int prefill_slots = (prefill_blocks + blk_factor_p - 1) / blk_factor_p; const int decode_slots = (decode_blocks + blk_factor_d - 1) / blk_factor_d; diff --git a/include/flashinfer/rocm/attention/pos_enc.cuh b/include/flashinfer/rocm/attention/pos_enc.cuh index 4d7ed202de..9d9caec972 100644 --- a/include/flashinfer/rocm/attention/pos_enc.cuh +++ b/include/flashinfer/rocm/attention/pos_enc.cuh @@ -2,8 +2,13 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_POS_ENC_CUH_ -#define FLASHINFER_POS_ENC_CUH_ +#ifdef FLASHINFER_POS_ENC_CUH_ +#error \ + "include/flashinfer/pos_enc.cuh and include/flashinfer/rocm/attention/pos_enc.cuh both define FLASHINFER_POS_ENC_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_POS_ENC_CUH_ +#define FLASHINFER_ROCM_ATTENTION_POS_ENC_CUH_ #include #include @@ -11,20 +16,19 @@ #include #include -#include "gpu_iface/dispatch.cuh" -#include "gpu_iface/enums.hpp" -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/layout.cuh" -#include "gpu_iface/macros.hpp" -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/utils.cuh" -#include "gpu_iface/vec_dtypes.hpp" +#include "flashinfer/rocm/dispatch.cuh" +#include "flashinfer/rocm/enums.hpp" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/layout.cuh" +#include "flashinfer/rocm/macros.hpp" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/utils.cuh" +#include "flashinfer/rocm/vec_dtypes_hip.h" #include "page.cuh" namespace flashinfer { -using namespace gpu_iface::vec_dtypes; /*! * \brief Convert PosEncodingMode to string * \param pos_encoding_mode A PosEncodingMode value @@ -43,10 +47,9 @@ inline std::string PosEncodingModeToString(const PosEncodingMode& pos_encoding_m } __device__ __forceinline__ float get_alibi_slope(uint32_t head_idx, uint32_t num_heads) { - int n = (int)gpu_iface::math::ptx_exp2(gpu_iface::math::ptx_log2(float(num_heads))); - return head_idx < n - ? gpu_iface::math::ptx_exp2(-8.f * float(head_idx + 1) / float(n)) - : gpu_iface::math::ptx_exp2(-4.f * float((head_idx + 1 - n) * 2 - 1) / float(n)); + int n = (int)math::ptx_exp2(math::ptx_log2(float(num_heads))); + return head_idx < n ? math::ptx_exp2(-8.f * float(head_idx + 1) / float(n)) + : math::ptx_exp2(-4.f * float((head_idx + 1 - n) * 2 - 1) / float(n)); } /*! @@ -579,10 +582,10 @@ gpuError_t BatchQKApplyRotaryPosIdsCosSinCache( DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { // 16-byte vectorised loads; at least one element per thread per wavefront lane - constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / gpu_iface::kWarpSize); + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / kWarpSize); constexpr uint32_t bdx = HEAD_DIM / vec_size; // at least 2 wavefronts per block for occupancy on CDNA3 / 2 warps on CUDA - uint32_t num_threads = std::max(2U * gpu_iface::kWarpSize, bdx); + uint32_t num_threads = std::max(2U * kWarpSize, bdx); uint32_t bdy = num_threads / bdx; uint32_t nblks_x = (nnz + bdy - 1) / bdy; void* args[] = {(void*)&q, @@ -649,9 +652,9 @@ gpuError_t BatchQKApplyRotaryPosIds( DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { - constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / gpu_iface::kWarpSize); + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / kWarpSize); constexpr uint32_t bdx = HEAD_DIM / vec_size; - uint32_t num_threads = std::max(2U * gpu_iface::kWarpSize, bdx); + uint32_t num_threads = std::max(2U * kWarpSize, bdx); uint32_t bdy = num_threads / bdx; uint32_t nblks_x = (nnz + bdy - 1) / bdy; @@ -717,9 +720,9 @@ gpuError_t BatchQKApplyRotary(DType* q, DType* k, DType* q_rope, DType* k_rope, DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { - constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / gpu_iface::kWarpSize); + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / kWarpSize); constexpr uint32_t bdx = HEAD_DIM / vec_size; - uint32_t num_threads = std::max(2U * gpu_iface::kWarpSize, bdx); + uint32_t num_threads = std::max(2U * kWarpSize, bdx); uint32_t bdy = num_threads / bdx; dim3 nblks(batch_size * (num_qo_heads + num_kv_heads)); dim3 nthrs(bdx, bdy); @@ -783,9 +786,9 @@ gpuError_t BatchQKApplyLlama31Rotary( DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { - constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / gpu_iface::kWarpSize); + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / kWarpSize); constexpr uint32_t bdx = HEAD_DIM / vec_size; - uint32_t num_threads = std::max(2U * gpu_iface::kWarpSize, bdx); + uint32_t num_threads = std::max(2U * kWarpSize, bdx); uint32_t bdy = num_threads / bdx; dim3 nblks(batch_size * (num_qo_heads + num_kv_heads)); dim3 nthrs(bdx, bdy); @@ -834,9 +837,9 @@ gpuError_t BatchQKApplyLlama31RotaryPosIds( DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { - constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / gpu_iface::kWarpSize); + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / kWarpSize); constexpr uint32_t bdx = HEAD_DIM / vec_size; - uint32_t num_threads = std::max(2U * gpu_iface::kWarpSize, bdx); + uint32_t num_threads = std::max(2U * kWarpSize, bdx); uint32_t bdy = num_threads / bdx; dim3 nblks((nnz + bdy - 1) / bdy); dim3 nthrs(bdx, bdy); @@ -928,7 +931,7 @@ __global__ void RopeQuantizeKernel( uint32_t k_rope_end = q_rope_end + num_kv_heads * rope_chunks; uint32_t k_nope_end = k_rope_end + num_kv_heads * no_rope_chunks; - using vec_t = flashinfer::gpu_iface::vec_dtypes::vec_t; + using vec_t = flashinfer::vec_t; vec_t cos, sin; if (bx * bdy + ty < nnz) { const uint32_t idx = bx * bdy + ty; @@ -1090,7 +1093,7 @@ __global__ void RopeQuantizeAppendPagedKVCacheKernel( uint32_t k_rope_end = q_rope_end + num_kv_heads * rope_chunks; uint32_t k_nope_end = k_rope_end + num_kv_heads * no_rope_chunks; - using vec_t = flashinfer::gpu_iface::vec_dtypes::vec_t; + using vec_t = flashinfer::vec_t; vec_t cos, sin; if (bx * bdy + ty < nnz) { const uint32_t idx = bx * bdy + ty; @@ -1268,7 +1271,7 @@ gpuError_t RopeQuantize( // 16-byte vector loads — single global_load_dwordx4 on CDNA3 constexpr uint32_t vec_size = 16 / sizeof(DType); constexpr uint32_t bdx = ROPE_DIM / vec_size; - uint32_t num_threads = std::max(2U * gpu_iface::kWarpSize, bdx); + uint32_t num_threads = std::max(2U * kWarpSize, bdx); uint32_t bdy = num_threads / bdx; uint32_t nblks_x = (nnz + bdy - 1) / bdy; uint32_t rope_chunk_size = rope_dim; @@ -1338,7 +1341,7 @@ gpuError_t RopeQuantizeAppendPagedKVCache( // 16-byte vector loads — single global_load_dwordx4 on CDNA3 constexpr uint32_t vec_size = 16 / sizeof(DType); constexpr uint32_t bdx = ROPE_DIM / vec_size; - uint32_t num_threads = std::max(2U * gpu_iface::kWarpSize, bdx); + uint32_t num_threads = std::max(2U * kWarpSize, bdx); uint32_t bdy = num_threads / bdx; uint32_t nblks_x = (nnz + bdy - 1) / bdy; uint32_t rope_chunks = 1; @@ -1410,7 +1413,7 @@ gpuError_t RopeQuantizeAppendPagedMLACache( DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { constexpr uint32_t vec_size = 16 / sizeof(DType); constexpr uint32_t bdx = ROPE_DIM / vec_size; - uint32_t num_threads = std::max(2U * gpu_iface::kWarpSize, bdx); + uint32_t num_threads = std::max(2U * kWarpSize, bdx); uint32_t bdy = num_threads / bdx; uint32_t nblks_x = (nnz + bdy - 1) / bdy; uint32_t rope_chunks = 1; @@ -1463,4 +1466,4 @@ gpuError_t RopeQuantizeAppendPagedMLACache( } // namespace flashinfer -#endif // FLASHINFER_POS_ENC_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_POS_ENC_CUH_ diff --git a/include/flashinfer/rocm/attention/prefill.cuh b/include/flashinfer/rocm/attention/prefill.cuh index d3cfa73674..f390662fa1 100644 --- a/include/flashinfer/rocm/attention/prefill.cuh +++ b/include/flashinfer/rocm/attention/prefill.cuh @@ -3,14 +3,18 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#ifdef FLASHINFER_PREFILL_CUH_ +#error \ + "include/flashinfer/attention/prefill.cuh and include/flashinfer/rocm/attention/prefill.cuh both define FLASHINFER_PREFILL_CUH_; include only one" +#endif -#include "gpu_iface/cooperative_groups.h" -#include "gpu_iface/fastdiv.cuh" -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/memory_ops.hpp" -#include "gpu_iface/mma_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/utils.cuh" +#include "flashinfer/rocm/cooperative_groups.h" +#include "flashinfer/rocm/fastdiv.cuh" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/memory_ops_hip.h" +#include "flashinfer/rocm/mma_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/utils.cuh" #ifdef FP16_QK_REDUCTION_SUPPORTED #include "../../fp16.h" @@ -18,7 +22,7 @@ #include #include "cascade.cuh" -#include "dispatch.cuh" +#include "flashinfer/rocm/dispatch.cuh" #include "frag_layout_swizzle.cuh" #include "page.cuh" #include "permuted_smem.cuh" @@ -30,14 +34,9 @@ namespace flashinfer { DEFINE_HAS_MEMBER(maybe_q_rope_offset) DEFINE_HAS_MEMBER(maybe_k_rope_offset) -namespace cg = gpu_iface::cg; -namespace memory = gpu_iface::memory; -namespace mma = gpu_iface::mma; - -using gpu_iface::vec_dtypes::vec_cast; -using mma::MMAMode; +using mma_hip::MMAMode; -constexpr uint32_t WARP_SIZE = gpu_iface::kWarpSize; +constexpr uint32_t WARP_SIZE = kWarpSize; constexpr uint32_t get_num_warps_q(const uint32_t cta_tile_q) { if (cta_tile_q > 16) { @@ -161,9 +160,9 @@ struct KernelTraits { template static constexpr DT getNegInf() { if constexpr (std::is_same::value) { - return std::bit_cast(fp16_ieee_from_fp32_value(-gpu_iface::math::inf)); + return std::bit_cast(fp16_ieee_from_fp32_value(-math::inf)); } else { - return static_cast(-gpu_iface::math::inf); + return static_cast(-math::inf); } } @@ -174,7 +173,7 @@ struct KernelTraits { "Set -DFP16_QK_REDUCTION_SUPPORTED and install boost_math " "then recompile to support fp16 reduction"); static constexpr DTypeQKAccum MaskFillValue = - AttentionVariant::use_softmax ? DTypeQKAccum(-gpu_iface::math::inf) : DTypeQKAccum(0.f); + AttentionVariant::use_softmax ? DTypeQKAccum(-math::inf) : DTypeQKAccum(0.f); #endif }; @@ -472,7 +471,7 @@ __device__ __forceinline__ void init_states( for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { #pragma unroll for (uint32_t j = 0; j < NUM_ACCUM_ROWS_PER_THREAD; ++j) { - m[mma_q][j] = typename KTraits::DTypeQKAccum(-gpu_iface::math::inf); + m[mma_q][j] = typename KTraits::DTypeQKAccum(-math::inf); d[mma_q][j] = 1.f; } } @@ -767,10 +766,10 @@ __device__ __forceinline__ void compute_qk( for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { if constexpr (std::is_same_v) { if (mma_d == 0) { - mma::mma_sync_m16n16k16_row_col_f16f16f32( + mma_hip::mma_sync_m16n16k16_row_col_f16f16f32( s_frag[mma_q][mma_kv], a_frag[mma_q], b_frag); } else { - mma::mma_sync_m16n16k16_row_col_f16f16f32( + mma_hip::mma_sync_m16n16k16_row_col_f16f16f32( s_frag[mma_q][mma_kv], a_frag[mma_q], b_frag); } } else if (std::is_same_v) { @@ -928,11 +927,11 @@ __device__ __forceinline__ void update_mdo_states( m[mma_q][j] = max(m[mma_q][j], s_frag[mma_q][mma_kv][j]); } // Butterfly reduction across all threads in the band - m[mma_q][j] = max(m[mma_q][j], gpu_iface::math::shfl_xor_sync(m[mma_q][j], 0x8)); - m[mma_q][j] = max(m[mma_q][j], gpu_iface::math::shfl_xor_sync(m[mma_q][j], 0x4)); - m[mma_q][j] = max(m[mma_q][j], gpu_iface::math::shfl_xor_sync(m[mma_q][j], 0x2)); - m[mma_q][j] = max(m[mma_q][j], gpu_iface::math::shfl_xor_sync(m[mma_q][j], 0x1)); - float o_scale = gpu_iface::math::ptx_exp2(m_prev * sm_scale - m[mma_q][j] * sm_scale); + m[mma_q][j] = max(m[mma_q][j], math::shfl_xor_sync(m[mma_q][j], 0x8)); + m[mma_q][j] = max(m[mma_q][j], math::shfl_xor_sync(m[mma_q][j], 0x4)); + m[mma_q][j] = max(m[mma_q][j], math::shfl_xor_sync(m[mma_q][j], 0x2)); + m[mma_q][j] = max(m[mma_q][j], math::shfl_xor_sync(m[mma_q][j], 0x1)); + float o_scale = math::ptx_exp2(m_prev * sm_scale - m[mma_q][j] * sm_scale); d[mma_q][j] *= o_scale; #pragma unroll @@ -941,8 +940,8 @@ __device__ __forceinline__ void update_mdo_states( } #pragma unroll for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { - s_frag[mma_q][mma_kv][j] = gpu_iface::math::ptx_exp2( - s_frag[mma_q][mma_kv][j] * sm_scale - m[mma_q][j] * sm_scale); + s_frag[mma_q][mma_kv][j] = + math::ptx_exp2(s_frag[mma_q][mma_kv][j] * sm_scale - m[mma_q][j] * sm_scale); } } } @@ -982,7 +981,7 @@ __device__ __forceinline__ void compute_sfm_v( for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { #pragma unroll for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { - mma::transpose_mma_tile(reinterpret_cast(s_frag_f16[mma_q][mma_kv])); + mma_hip::transpose_mma_tile(reinterpret_cast(s_frag_f16[mma_q][mma_kv])); } } @@ -992,7 +991,7 @@ __device__ __forceinline__ void compute_sfm_v( #pragma unroll for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { if constexpr (std::is_same_v) { - mma::m16k16_rowsum_f16f16f32(d[mma_q], s_frag_f16[mma_q][mma_kv]); + mma_hip::m16k16_rowsum_f16f16f32(d[mma_q], s_frag_f16[mma_q][mma_kv]); } else { static_assert(!std::is_same_v, "FP16 reduction path not implemented for CDNA3"); @@ -1018,10 +1017,10 @@ __device__ __forceinline__ void compute_sfm_v( #pragma unroll for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { if constexpr (std::is_same_v) { - mma::mma_sync_m16n16k16_row_col_f16f16f32( + mma_hip::mma_sync_m16n16k16_row_col_f16f16f32( o_frag[mma_q][mma_d], (uint32_t*)s_frag_f16[mma_q][mma_kv], b_frag); } else { - mma::mma_sync_m16n16k16_row_col_f16f16f32( + mma_hip::mma_sync_m16n16k16_row_col_f16f16f32( o_frag[mma_q][mma_d], (uint32_t*)s_frag[mma_q][mma_kv], b_frag); } } @@ -1059,8 +1058,8 @@ __device__ __forceinline__ void normalize_d( for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { #pragma unroll for (uint32_t j = 0; j < KTraits::NUM_ACCUM_ROWS_PER_THREAD; ++j) { - d_rcp[mma_q][j] = (m[mma_q][j] != typename KTraits::DTypeQKAccum(-gpu_iface::math::inf)) - ? gpu_iface::math::ptx_rcp(d[mma_q][j]) + d_rcp[mma_q][j] = (m[mma_q][j] != typename KTraits::DTypeQKAccum(-math::inf)) + ? math::ptx_rcp(d[mma_q][j]) : 0.f; } } @@ -1088,7 +1087,7 @@ __device__ __forceinline__ void finalize_m( for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { #pragma unroll for (uint32_t j = 0; j < KTraits::NUM_ACCUM_ROWS_PER_THREAD; ++j) { - if (m[mma_q][j] != typename KTraits::DTypeQKAccum(-gpu_iface::math::inf)) { + if (m[mma_q][j] != typename KTraits::DTypeQKAccum(-math::inf)) { m[mma_q][j] *= variant.sm_scale_log2; } } @@ -1165,21 +1164,20 @@ __device__ __forceinline__ void threadblock_sync_mdo_states( float o_scale[NARPT][KTraits::NUM_WARPS_KV]; #pragma unroll for (uint32_t j = 0; j < NARPT; ++j) { - float m_new = -gpu_iface::math::inf, d_new = 1.f; + float m_new = -math::inf, d_new = 1.f; #pragma unroll for (uint32_t i = 0; i < KTraits::NUM_WARPS_KV; ++i) { float2 md = smem_md[i * KTraits::NUM_MMA_Q * 16 + mma_q * 16 + ln_grp_idx * NARPT + j]; float m_prev = m_new, d_prev = d_new; m_new = max(m_new, md.x); - d_new = d_prev * gpu_iface::math::ptx_exp2(m_prev - m_new) + - md.y * gpu_iface::math::ptx_exp2(md.x - m_new); + d_new = d_prev * math::ptx_exp2(m_prev - m_new) + md.y * math::ptx_exp2(md.x - m_new); } #pragma unroll for (uint32_t i = 0; i < KTraits::NUM_WARPS_KV; ++i) { float2 md = smem_md[i * KTraits::NUM_MMA_Q * 16 + mma_q * 16 + ln_grp_idx * NARPT + j]; float mi = md.x; - o_scale[j][i] = gpu_iface::math::ptx_exp2(float(mi - m_new)); + o_scale[j][i] = math::ptx_exp2(float(mi - m_new)); } m[mma_q][j] = typename KTraits::DTypeQKAccum(m_new); d[mma_q][j] = d_new; @@ -1603,10 +1601,10 @@ __device__ __forceinline__ void SinglePrefillWithKVCacheDevice( if (qo_idx < qo_len) { if (partition_kv) { lse[(qo_idx * num_chunks + chunk_idx) * num_qo_heads + qo_head_idx] = - gpu_iface::math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); } else { lse[qo_idx * num_qo_heads + qo_head_idx] = - gpu_iface::math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); } } } @@ -2054,10 +2052,10 @@ __global__ __launch_bounds__(KTraits::NUM_THREADS) void BatchPrefillWithRaggedKV if (qo_idx < qo_len) { if (partition_kv) { lse[(o_indptr[request_idx] + qo_idx * num_kv_chunks + kv_tile_idx) * num_qo_heads + - qo_head_idx] = gpu_iface::math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + qo_head_idx] = math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); } else { lse[(o_indptr[request_idx] + qo_idx) * num_qo_heads + qo_head_idx] = - gpu_iface::math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); } } } @@ -2342,7 +2340,7 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( const uint32_t qo_head_idx = kv_head_idx * group_size + r; const uint32_t qo_idx = q_idx; if (qo_idx < qo_upper_bound) { - const float s_cur = gpu_iface::math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + const float s_cur = math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); const float s_partial = params.partial_lse[(o_indptr[request_idx] + qo_idx) * num_qo_heads + qo_head_idx]; const float s_max = fmaxf(s_cur, s_partial); @@ -2396,10 +2394,10 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( if (qo_idx < qo_upper_bound) { if (partition_kv) { lse[(o_indptr[request_idx] + qo_idx * num_kv_chunks + kv_tile_idx) * num_qo_heads + - qo_head_idx] = gpu_iface::math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + qo_head_idx] = math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); } else { lse[(o_indptr[request_idx] + qo_idx) * num_qo_heads + qo_head_idx] = - gpu_iface::math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); } } } diff --git a/include/flashinfer/rocm/attention/scheduler.cuh b/include/flashinfer/rocm/attention/scheduler.cuh index ace27b345b..9a2342318f 100644 --- a/include/flashinfer/rocm/attention/scheduler.cuh +++ b/include/flashinfer/rocm/attention/scheduler.cuh @@ -2,8 +2,13 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_ATTENTION_SCHEDULER_CUH_ -#define FLASHINFER_ATTENTION_SCHEDULER_CUH_ +#ifdef FLASHINFER_ATTENTION_SCHEDULER_CUH_ +#error \ + "include/flashinfer/attention/scheduler.cuh and include/flashinfer/rocm/attention/scheduler.cuh both define FLASHINFER_ATTENTION_SCHEDULER_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_SCHEDULER_CUH_ +#define FLASHINFER_ROCM_ATTENTION_SCHEDULER_CUH_ #include #include @@ -13,11 +18,11 @@ #include "allocator.h" #include "decode_tuning.cuh" -#include "exception.h" -#include "gpu_iface/dispatch.cuh" -#include "gpu_iface/gpu_runtime_compat.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/utils.cuh" +#include "flashinfer/rocm/dispatch.cuh" +#include "flashinfer/rocm/exception.h" +#include "flashinfer/rocm/gpu_runtime_compat.hpp" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/utils.cuh" #include "heap.h" #include "pos_enc.cuh" @@ -1346,4 +1351,4 @@ inline gpuError_t MLAPlan(void* float_buffer, size_t float_workspace_size_in_byt } } // namespace flashinfer -#endif // FLASHINFER_ATTENTION_SCHEDULER_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_SCHEDULER_CUH_ diff --git a/include/flashinfer/rocm/attention/state.cuh b/include/flashinfer/rocm/attention/state.cuh index db851224e8..c58c6cc202 100644 --- a/include/flashinfer/rocm/attention/state.cuh +++ b/include/flashinfer/rocm/attention/state.cuh @@ -2,19 +2,23 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_STATE_CUH_ -#define FLASHINFER_STATE_CUH_ +#ifdef FLASHINFER_STATE_CUH_ +#error \ + "include/flashinfer/attention/state.cuh and include/flashinfer/rocm/attention/state.cuh both define FLASHINFER_STATE_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_STATE_CUH_ +#define FLASHINFER_ROCM_ATTENTION_STATE_CUH_ -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/vec_dtypes.hpp" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/vec_dtypes_hip.h" #if defined(PLATFORM_HIP_DEVICE) #define HIP_ENABLE_WARP_SYNC_BUILTINS 1 #endif namespace flashinfer { -using namespace gpu_iface::vec_dtypes; /*! * \brief The flashattention state. * \tparam vec_size The size of the vector used in o. @@ -75,4 +79,4 @@ struct state_t { } // namespace flashinfer -#endif // FLASHINFER_STATE_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_STATE_CUH_ diff --git a/include/flashinfer/rocm/attention/variant_helper.cuh b/include/flashinfer/rocm/attention/variant_helper.cuh index c37575e35d..2fae6cecb2 100644 --- a/include/flashinfer/rocm/attention/variant_helper.cuh +++ b/include/flashinfer/rocm/attention/variant_helper.cuh @@ -2,12 +2,17 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_ATTENTION_VARIANT_HELPER_H -#define FLASHINFER_ATTENTION_VARIANT_HELPER_H +#ifdef FLASHINFER_ATTENTION_VARIANT_HELPER_H +#error \ + "include/flashinfer/attention/variant_helper.cuh and include/flashinfer/rocm/attention/variant_helper.cuh both define FLASHINFER_ATTENTION_VARIANT_HELPER_H; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_VARIANT_HELPER_CUH_ +#define FLASHINFER_ROCM_ATTENTION_VARIANT_HELPER_CUH_ #include -#include "gpu_iface/platform.hpp" +#include "flashinfer/rocm/platform.hpp" namespace flashinfer { @@ -51,4 +56,4 @@ struct AttentionVariantBase { } // namespace flashinfer -#endif // FLASHINFER_ATTENTION_VARIANT_HELPER_H +#endif // FLASHINFER_ROCM_ATTENTION_VARIANT_HELPER_CUH_ diff --git a/include/flashinfer/rocm/attention/variants.cuh b/include/flashinfer/rocm/attention/variants.cuh index bab6143eaa..3784b8528e 100644 --- a/include/flashinfer/rocm/attention/variants.cuh +++ b/include/flashinfer/rocm/attention/variants.cuh @@ -2,21 +2,26 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_ATTENTION_VARIANTS_CUH_ -#define FLASHINFER_ATTENTION_VARIANTS_CUH_ +#ifdef FLASHINFER_ATTENTION_VARIANTS_CUH_ +#error \ + "include/flashinfer/attention/variants.cuh and include/flashinfer/rocm/attention/variants.cuh both define FLASHINFER_ATTENTION_VARIANTS_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_ATTENTION_VARIANTS_CUH_ +#define FLASHINFER_ROCM_ATTENTION_VARIANTS_CUH_ #include #include -#include "gpu_iface/math_ops.hpp" -#include "gpu_iface/memory_ops.hpp" -#include "gpu_iface/platform.hpp" -#include "gpu_iface/utils.cuh" +#include "flashinfer/rocm/math_hip.h" +#include "flashinfer/rocm/memory_ops_hip.h" +#include "flashinfer/rocm/platform.hpp" +#include "flashinfer/rocm/utils.cuh" #include "variant_helper.cuh" namespace flashinfer { -using gpu_iface::memory::SharedMemFillMode; +using memory::SharedMemFillMode; DEFINE_HAS_MEMBER(maybe_mask_indptr) @@ -37,13 +42,13 @@ struct DefaultAttention : AttentionVariantBase { qo_len = params.get_qo_len(batch_idx); kv_len = params.get_kv_len(batch_idx); if constexpr (use_logits_soft_cap) { - soft_cap_pre_tanh_scale = params.sm_scale * gpu_iface::math::ptx_rcp(params.logits_soft_cap); - sm_scale_log2 = gpu_iface::math::log2e * params.logits_soft_cap; + soft_cap_pre_tanh_scale = params.sm_scale * math::ptx_rcp(params.logits_soft_cap); + sm_scale_log2 = math::log2e * params.logits_soft_cap; } else { if constexpr (use_alibi) { - sm_scale_log2 = gpu_iface::math::log2e; + sm_scale_log2 = math::log2e; } else { - sm_scale_log2 = params.sm_scale * gpu_iface::math::log2e; + sm_scale_log2 = params.sm_scale * math::log2e; } } if constexpr (use_custom_mask) { @@ -64,7 +69,7 @@ struct DefaultAttention : AttentionVariantBase { params.maybe_alibi_slopes[qo_head_idx] * float(int(kv_idx) - int(qo_idx)); } if constexpr (use_logits_soft_cap) { - logits = float(gpu_iface::math::tanh(logits * soft_cap_pre_tanh_scale)); + logits = float(math::tanh(logits * soft_cap_pre_tanh_scale)); } return logits; }) @@ -88,4 +93,4 @@ struct DefaultAttention : AttentionVariantBase { }; // namespace flashinfer -#endif // FLASHINFER_ATTENTION_VARIANTS_CUH_ +#endif // FLASHINFER_ROCM_ATTENTION_VARIANTS_CUH_ diff --git a/include/gpu_iface/conversion_utils.h b/include/flashinfer/rocm/conversion_utils.h similarity index 87% rename from include/gpu_iface/conversion_utils.h rename to include/flashinfer/rocm/conversion_utils.h index 4fb683a51a..60bcf979dd 100644 --- a/include/gpu_iface/conversion_utils.h +++ b/include/flashinfer/rocm/conversion_utils.h @@ -3,13 +3,18 @@ #pragma once +// clang-format off +// macros.hpp first: its non-HIP #error must fire before a missing hip/ header. +#include "macros.hpp" +// clang-format on + #include #include #include #include #include -namespace fi::con { +namespace flashinfer { template __host__ __device__ __inline__ DTypeOut explicit_casting(DTypeIn value) { return DTypeOut(value); @@ -51,4 +56,4 @@ __host__ __device__ __inline__ __hip_bfloat16 explicit_casting<__hip_bfloat16, _ __hip_bfloat16 value) { return value; } -} // namespace fi::con +} // namespace flashinfer diff --git a/include/gpu_iface/cooperative_groups.h b/include/flashinfer/rocm/cooperative_groups.h similarity index 91% rename from include/gpu_iface/cooperative_groups.h rename to include/flashinfer/rocm/cooperative_groups.h index cb4d8c0f90..5d1ebf643a 100644 --- a/include/gpu_iface/cooperative_groups.h +++ b/include/flashinfer/rocm/cooperative_groups.h @@ -10,7 +10,5 @@ #include // clang-format on namespace flashinfer { -namespace gpu_iface { namespace cg = ::cooperative_groups; -} // namespace gpu_iface } // namespace flashinfer diff --git a/include/gpu_iface/dispatch.cuh b/include/flashinfer/rocm/dispatch.cuh similarity index 89% rename from include/gpu_iface/dispatch.cuh rename to include/flashinfer/rocm/dispatch.cuh index 135eb4e0d7..4d2628c10e 100644 --- a/include/gpu_iface/dispatch.cuh +++ b/include/flashinfer/rocm/dispatch.cuh @@ -3,9 +3,15 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +// Not a shared include guard but the same hazard: upstream's utils.cuh defines +// eleven of the DISPATCH_* macros below, and the last definition would win. +#ifdef FLASHINFER_UTILS_CUH_ +#error \ + "include/flashinfer/utils.cuh and include/flashinfer/rocm/dispatch.cuh both define the DISPATCH_* macros; include only one" +#endif -#include "enums.hpp" -#include "gpu_iface/exception.h" +#include "flashinfer/rocm/enums.hpp" +#include "flashinfer/rocm/exception.h" #define DISPATCH_USE_FP16_QK_REDUCTION(use_fp16_qk_reduction, USE_FP16_QK_REDUCTION, ...) \ if (use_fp16_qk_reduction) { \ @@ -28,23 +34,24 @@ FLASHINFER_ERROR(err_msg.str()); \ } -#define DISPATCH_NUM_MMA_KV(max_mma_kv, NUM_MMA_KV, ...) \ - if (max_mma_kv >= 8) { \ - constexpr size_t NUM_MMA_KV = 8; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 4) { \ - constexpr size_t NUM_MMA_KV = 4; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 2) { \ - constexpr size_t NUM_MMA_KV = 2; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 1) { \ - constexpr size_t NUM_MMA_KV = 1; \ - __VA_ARGS__ \ - } else { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported max_mma_kv: " << max_mma_kv; \ - FLASHINFER_ERROR(err_msg.str()); \ +#define DISPATCH_NUM_MMA_KV(max_mma_kv, NUM_MMA_KV, ...) \ + if (max_mma_kv >= 8) { \ + constexpr size_t NUM_MMA_KV = 8; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 4) { \ + constexpr size_t NUM_MMA_KV = 4; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 2) { \ + constexpr size_t NUM_MMA_KV = 2; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 1) { \ + constexpr size_t NUM_MMA_KV = 1; \ + __VA_ARGS__ \ + } else { \ + /* Fallback for AMD GPUs with tight shared memory constraints */ \ + /* Use NUM_MMA_KV=1 and issue warning */ \ + constexpr size_t NUM_MMA_KV = 1; \ + __VA_ARGS__ \ } #define DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, ...) \ @@ -205,6 +212,15 @@ } \ } +#define DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, ...) \ + if (compute_capacity.first >= 8) { \ + constexpr uint32_t NUM_STAGES_SMEM = 2; \ + __VA_ARGS__ \ + } else { \ + constexpr uint32_t NUM_STAGES_SMEM = 1; \ + __VA_ARGS__ \ + } + #define DISPATCH_ROPE_DIM(rope_dim, ROPE_DIM, ...) \ switch (rope_dim) { \ case 16: { \ @@ -239,12 +255,3 @@ FLASHINFER_ERROR(err_msg.str()); \ } \ } - -#define DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, ...) \ - if (compute_capacity.first >= 8) { \ - constexpr uint32_t NUM_STAGES_SMEM = 2; \ - __VA_ARGS__ \ - } else { \ - constexpr uint32_t NUM_STAGES_SMEM = 1; \ - __VA_ARGS__ \ - } diff --git a/include/gpu_iface/enums.hpp b/include/flashinfer/rocm/enums.hpp similarity index 100% rename from include/gpu_iface/enums.hpp rename to include/flashinfer/rocm/enums.hpp diff --git a/include/flashinfer/rocm/attention/exception.h b/include/flashinfer/rocm/exception.h similarity index 74% rename from include/flashinfer/rocm/attention/exception.h rename to include/flashinfer/rocm/exception.h index fe6e83b690..35b9c4f645 100644 --- a/include/flashinfer/rocm/attention/exception.h +++ b/include/flashinfer/rocm/exception.h @@ -2,11 +2,17 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_EXCEPTION_H_ -#define FLASHINFER_EXCEPTION_H_ +#ifdef FLASHINFER_EXCEPTION_H_ +#error \ + "include/flashinfer/exception.h and include/flashinfer/rocm/exception.h both define FLASHINFER_EXCEPTION_H_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_EXCEPTION_H_ +#define FLASHINFER_ROCM_EXCEPTION_H_ #include #include +#include namespace flashinfer { @@ -34,4 +40,4 @@ class Error : public std::exception { } // namespace flashinfer -#endif // FLASHINFER_EXCEPTION_H_ +#endif // FLASHINFER_ROCM_EXCEPTION_H_ diff --git a/include/gpu_iface/fastdiv.cuh b/include/flashinfer/rocm/fastdiv.cuh similarity index 88% rename from include/gpu_iface/fastdiv.cuh rename to include/flashinfer/rocm/fastdiv.cuh index a19d01def7..403c27b483 100644 --- a/include/gpu_iface/fastdiv.cuh +++ b/include/flashinfer/rocm/fastdiv.cuh @@ -3,8 +3,13 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_FASTDIV_CUH_ -#define FLASHINFER_FASTDIV_CUH_ +#ifdef FLASHINFER_FASTDIV_CUH_ +#error \ + "include/flashinfer/fastdiv.cuh and include/flashinfer/rocm/fastdiv.cuh both define FLASHINFER_FASTDIV_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_FASTDIV_CUH_ +#define FLASHINFER_ROCM_FASTDIV_CUH_ #include @@ -97,4 +102,4 @@ __host__ __device__ __forceinline__ uint32_t operator%(const uint32_t n, } // namespace flashinfer -#endif // FLASHINFER_FASTDIV_CUH_ +#endif // FLASHINFER_ROCM_FASTDIV_CUH_ diff --git a/include/gpu_iface/gpu_runtime_compat.hpp b/include/flashinfer/rocm/gpu_runtime_compat.hpp similarity index 100% rename from include/gpu_iface/gpu_runtime_compat.hpp rename to include/flashinfer/rocm/gpu_runtime_compat.hpp diff --git a/include/gpu_iface/layout.cuh b/include/flashinfer/rocm/layout.cuh similarity index 94% rename from include/gpu_iface/layout.cuh rename to include/flashinfer/rocm/layout.cuh index bb66739b10..fc81d775c2 100644 --- a/include/gpu_iface/layout.cuh +++ b/include/flashinfer/rocm/layout.cuh @@ -2,8 +2,13 @@ // SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 -#ifndef FLASHINFER_LAYOUT_CUH_ -#define FLASHINFER_LAYOUT_CUH_ +#ifdef FLASHINFER_LAYOUT_CUH_ +#error \ + "include/flashinfer/layout.cuh and include/flashinfer/rocm/layout.cuh both define FLASHINFER_LAYOUT_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_LAYOUT_CUH_ +#define FLASHINFER_ROCM_LAYOUT_CUH_ #include #include @@ -116,4 +121,4 @@ inline std::string QKVLayoutToString(const QKVLayout& layout) { } } // namespace flashinfer -#endif // FLASHINFER_LAYOUT_CUH_ +#endif // FLASHINFER_ROCM_LAYOUT_CUH_ diff --git a/include/gpu_iface/macros.hpp b/include/flashinfer/rocm/macros.hpp similarity index 92% rename from include/gpu_iface/macros.hpp rename to include/flashinfer/rocm/macros.hpp index 4eaecf349e..13b8bb3d15 100644 --- a/include/gpu_iface/macros.hpp +++ b/include/flashinfer/rocm/macros.hpp @@ -20,7 +20,7 @@ #endif #else -// The CUDA backend was removed; gpu_iface serves HIP only. Fail here with a +// The CUDA backend was removed; these headers serve HIP only. Fail here with a // named diagnostic rather than deeper in a missing hip/ header. #error "flashinfer ROCm requires a HIP compiler (__HIP__ / __HIPCC__)." #endif diff --git a/include/gpu_iface/backend/hip/math_hip.h b/include/flashinfer/rocm/math_hip.h similarity index 88% rename from include/gpu_iface/backend/hip/math_hip.h rename to include/flashinfer/rocm/math_hip.h index 6f98b19990..5a65471aa6 100644 --- a/include/gpu_iface/backend/hip/math_hip.h +++ b/include/flashinfer/rocm/math_hip.h @@ -3,11 +3,21 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once -#ifndef FLASHINFER_MATH_CUH_ -#define FLASHINFER_MATH_CUH_ +#ifdef FLASHINFER_MATH_CUH_ +#error \ + "include/flashinfer/math.cuh and include/flashinfer/rocm/math_hip.h both define FLASHINFER_MATH_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_MATH_HIP_H_ +#define FLASHINFER_ROCM_MATH_HIP_H_ #define HIP_ENABLE_WARP_SYNC_BUILTINS 1 +// clang-format off +// macros.hpp first: its non-HIP #error must fire before a missing hip/ header. +#include "macros.hpp" +// clang-format on + #include #include #include @@ -105,4 +115,4 @@ __forceinline__ __device__ __half2 tanh<__half2>(__half2 x) { } } // namespace flashinfer::math -#endif // FLASHINFER_MATH_CUH_ +#endif // FLASHINFER_ROCM_MATH_HIP_H_ diff --git a/include/gpu_iface/backend/hip/memory_ops_hip.h b/include/flashinfer/rocm/memory_ops_hip.h similarity index 94% rename from include/gpu_iface/backend/hip/memory_ops_hip.h rename to include/flashinfer/rocm/memory_ops_hip.h index 50e757effc..67a375e569 100644 --- a/include/gpu_iface/backend/hip/memory_ops_hip.h +++ b/include/flashinfer/rocm/memory_ops_hip.h @@ -2,14 +2,19 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +// clang-format off +// macros.hpp first: its non-HIP #error must fire before a missing hip/ header. +#include "macros.hpp" +// clang-format on + #include #include +#include "memory_types.hpp" +#include "platform.hpp" + namespace flashinfer { -namespace gpu_iface { namespace memory { -namespace detail { -namespace hip { __device__ __forceinline__ void commit_group() { // Currently a no-op for HIP @@ -83,8 +88,5 @@ __device__ __forceinline__ void pred_load(T* smem_ptr, const T* gmem_ptr, bool p } } -} // namespace hip -} // namespace detail } // namespace memory -} // namespace gpu_iface } // namespace flashinfer diff --git a/include/flashinfer/rocm/memory_types.hpp b/include/flashinfer/rocm/memory_types.hpp new file mode 100644 index 0000000000..a2ec307ab8 --- /dev/null +++ b/include/flashinfer/rocm/memory_types.hpp @@ -0,0 +1,27 @@ +// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once +#include "platform.hpp" + +namespace flashinfer { +namespace memory { + +/** + * @brief Control options for shared memory fill behavior + */ +enum class SharedMemFillMode { + kFillZero, // Fill zero to shared memory when predicate is false + kNoFill // Do not fill zero to shared memory when predicate is false +}; + +/** + * @brief Control options for memory prefetch behavior + */ +enum class PrefetchMode { + kNoPrefetch, // Do not fetch additional data from global memory to L2 + kPrefetch // Fetch additional data from global memory to L2 +}; + +} // namespace memory +} // namespace flashinfer diff --git a/include/gpu_iface/backend/hip/mma_hip.h b/include/flashinfer/rocm/mma_hip.h similarity index 97% rename from include/gpu_iface/backend/hip/mma_hip.h rename to include/flashinfer/rocm/mma_hip.h index b24dccd895..8de0f87fd6 100644 --- a/include/gpu_iface/backend/hip/mma_hip.h +++ b/include/flashinfer/rocm/mma_hip.h @@ -5,8 +5,8 @@ #include -#include "gpu_iface/mma_types.hpp" -#include "gpu_iface/platform.hpp" +#include "flashinfer/rocm/mma_types.hpp" +#include "flashinfer/rocm/platform.hpp" namespace { using f16 = _Float16; @@ -16,9 +16,7 @@ using f32x4 = float __attribute__((ext_vector_type(4))); } // namespace namespace flashinfer { -namespace gpu_iface { -namespace mma_impl { -namespace hip { +namespace mma_hip { #define FLASHINFER_RUNTIME_ASSERT(x) assert(0 && x) @@ -140,7 +138,7 @@ __device__ __forceinline__ void load_fragment(uint32_t* R, const T* smem_ptr) { } // MMA operation for FP16 inputs with FP32 accumulator -template +template __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) { #if defined(__HIP_DEVICE_COMPILE__) && (__gfx90a__ || __gfx908__ || __gfx942__ || __gfx950__) @@ -149,7 +147,7 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, u "T must be __half or __hip_bfloat16"); // Initialize C if requested - if constexpr (mma_mode == mma::MMAMode::kInit) { + if constexpr (mma_mode == mma_hip::MMAMode::kInit) { C[0] = 0.0f; C[1] = 0.0f; C[2] = 0.0f; @@ -246,7 +244,5 @@ __device__ __forceinline__ void m16k32_rowsum_f8f8f32(float* d_frag, DType* s_fr FLASHINFER_RUNTIME_ASSERT("FP8 rowsum not implemented for AMD"); } -} // namespace hip -} // namespace mma_impl -} // namespace gpu_iface +} // namespace mma_hip } // namespace flashinfer diff --git a/include/gpu_iface/mma_types.hpp b/include/flashinfer/rocm/mma_types.hpp similarity index 73% rename from include/gpu_iface/mma_types.hpp rename to include/flashinfer/rocm/mma_types.hpp index f675fdd7ec..9b4e166da6 100644 --- a/include/gpu_iface/mma_types.hpp +++ b/include/flashinfer/rocm/mma_types.hpp @@ -4,14 +4,12 @@ #pragma once namespace flashinfer { -namespace gpu_iface { -namespace mma { +namespace mma_hip { enum class MMAMode { kInit = 0U, kInplaceUpdate = 1U, }; -} // namespace mma -} // namespace gpu_iface +} // namespace mma_hip } // namespace flashinfer diff --git a/include/gpu_iface/platform.hpp b/include/flashinfer/rocm/platform.hpp similarity index 85% rename from include/gpu_iface/platform.hpp rename to include/flashinfer/rocm/platform.hpp index a1922eed67..cca4a4f8cb 100644 --- a/include/gpu_iface/platform.hpp +++ b/include/flashinfer/rocm/platform.hpp @@ -6,10 +6,8 @@ #include "macros.hpp" namespace flashinfer { -namespace gpu_iface { // Platform-agnostic stream type constexpr int kWarpSize = 64; -} // namespace gpu_iface } // namespace flashinfer diff --git a/include/flashinfer/rocm/quantization.cuh b/include/flashinfer/rocm/quantization.cuh index 8ea404508f..d1d8b3a281 100644 --- a/include/flashinfer/rocm/quantization.cuh +++ b/include/flashinfer/rocm/quantization.cuh @@ -17,10 +17,15 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#ifndef FLASHINFER_QUANTIZATION_CUH_ -#define FLASHINFER_QUANTIZATION_CUH_ -#include -#include +#ifdef FLASHINFER_QUANTIZATION_CUH_ +#error \ + "include/flashinfer/quantization.cuh and include/flashinfer/rocm/quantization.cuh both define FLASHINFER_QUANTIZATION_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_QUANTIZATION_CUH_ +#define FLASHINFER_ROCM_QUANTIZATION_CUH_ +#include +#include // CUB/hipCUB abstraction #ifdef PLATFORM_HIP_DEVICE @@ -31,7 +36,7 @@ namespace block_ops = hipcub; namespace block_ops = cub; #endif -#include +#include namespace flashinfer { namespace quantization { @@ -128,4 +133,4 @@ gpuError_t SegmentPackBits(bool* input, uint8_t* output, IdType* input_indptr, } // namespace quantization } // namespace flashinfer -#endif // FLASHINFER_QUANTIZATION_CUH_ +#endif // FLASHINFER_ROCM_QUANTIZATION_CUH_ diff --git a/include/flashinfer/rocm/sampling.cuh b/include/flashinfer/rocm/sampling.cuh index 415ae5fff0..6083756b62 100644 --- a/include/flashinfer/rocm/sampling.cuh +++ b/include/flashinfer/rocm/sampling.cuh @@ -18,45 +18,34 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#ifndef FLASHINFER_SAMPLING_CUH_ -#define FLASHINFER_SAMPLING_CUH_ +#ifdef FLASHINFER_SAMPLING_CUH_ +#error \ + "include/flashinfer/sampling.cuh and include/flashinfer/rocm/sampling.cuh both define FLASHINFER_SAMPLING_CUH_; include only one" +#endif -// gpu_iface portability layer provides FI_GPU_CALL, gpuError_t, gpuStream_t, -// gpuLaunchKernel, gpuFuncSetAttribute, etc. on both CUDA and HIP. -// macros.hpp must be first: it defines PLATFORM_HIP_DEVICE / PLATFORM_CUDA_DEVICE. -#include -#include -#include +#ifndef FLASHINFER_ROCM_SAMPLING_CUH_ +#define FLASHINFER_ROCM_SAMPLING_CUH_ -#ifdef PLATFORM_HIP_DEVICE -// --- HIP-specific: hiprand, hipcub, and gpu_iface math/utils/vec_dtypes --- +// clang-format off +// macros.hpp first: its non-HIP #error must fire before a missing hip/ header. +#include +// clang-format on #include #include +#include +#include #include namespace cub = hipcub; -#include -#include -#include -#else -// --- CUDA-specific --- -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#endif // PLATFORM_HIP_DEVICE +#include +#include -// allocator.h is a pure-C++ header; include it on both CUDA and HIP. -#include +#include + +// The fork's allocator, not upstream's: upstream's pulls upstream exception.h, +// which our exception.h now tripwires against. +#include #include #include @@ -90,9 +79,6 @@ namespace flashinfer { namespace sampling { using namespace cub; -#ifdef PLATFORM_HIP_DEVICE -using namespace gpu_iface::vec_dtypes; -#endif // Warp/wavefront size: 32 for NVIDIA, 64 for AMD #ifdef PLATFORM_HIP_DEVICE @@ -2399,4 +2385,4 @@ gpuError_t ChainSpeculativeSampling(DType* draft_probs, IdType* draft_token_ids, } // namespace flashinfer -#endif // FLASHINFER_SAMPLING_CUH_ +#endif // FLASHINFER_ROCM_SAMPLING_CUH_ diff --git a/include/gpu_iface/sm_id.hpp b/include/flashinfer/rocm/sm_id.hpp similarity index 93% rename from include/gpu_iface/sm_id.hpp rename to include/flashinfer/rocm/sm_id.hpp index 448ae726d5..f06a1ff0f8 100644 --- a/include/gpu_iface/sm_id.hpp +++ b/include/flashinfer/rocm/sm_id.hpp @@ -7,7 +7,6 @@ #include "macros.hpp" namespace flashinfer { -namespace gpu_iface { __device__ __forceinline__ uint32_t get_processor_id() { // Read HW_ID (id=4, offset=0, size=32) and extract bits [15:8], which pack @@ -18,5 +17,4 @@ __device__ __forceinline__ uint32_t get_processor_id() { return (hw_id >> 8) & 0xFF; } -} // namespace gpu_iface } // namespace flashinfer diff --git a/include/gpu_iface/utils.cuh b/include/flashinfer/rocm/utils.cuh similarity index 94% rename from include/gpu_iface/utils.cuh rename to include/flashinfer/rocm/utils.cuh index 980fafc659..ded40cf34e 100644 --- a/include/gpu_iface/utils.cuh +++ b/include/flashinfer/rocm/utils.cuh @@ -3,6 +3,10 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#ifdef FLASHINFER_UTILS_CUH_ +#error \ + "include/flashinfer/utils.cuh and include/flashinfer/rocm/utils.cuh both define FLASHINFER_UTILS_CUH_; include only one" +#endif #include #include @@ -73,9 +77,8 @@ inline uint32_t FA2ForcedCtaTileQ() { // would otherwise skip leading space and a '+' while still rejecting "64 ". const bool digits_only = (env[0] >= '0' && env[0] <= '9'); if (!digits_only || *end != '\0' || (value != 16ul && value != 64ul && value != 128ul)) { - // Throw std:: directly: gpu_iface/exception.h shares the FLASHINFER_EXCEPTION_H_ - // guard with flashinfer/exception.h, so including it here could suppress that - // header's variadic FLASHINFER_CHECK. + // Throw std:: directly rather than pull in exception.h: this header is included + // widely, and exception.h tripwires against upstream's FLASHINFER_EXCEPTION_H_. std::ostringstream err_msg; err_msg << "FLASHINFER_ROCM_FORCE_CTA_TILE_Q must be 16, 64, or 128, got \"" << env << "\""; throw std::invalid_argument(err_msg.str()); diff --git a/include/gpu_iface/backend/hip/vec_dtypes_hip.h b/include/flashinfer/rocm/vec_dtypes_hip.h similarity index 99% rename from include/gpu_iface/backend/hip/vec_dtypes_hip.h rename to include/flashinfer/rocm/vec_dtypes_hip.h index 8e4abdc668..6a12ac9b44 100644 --- a/include/gpu_iface/backend/hip/vec_dtypes_hip.h +++ b/include/flashinfer/rocm/vec_dtypes_hip.h @@ -3,11 +3,21 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once -#ifndef VEC_DTYPES_CUH_ -#define VEC_DTYPES_CUH_ +#ifdef VEC_DTYPES_CUH_ +#error \ + "include/flashinfer/vec_dtypes.cuh and include/flashinfer/rocm/vec_dtypes_hip.h both define VEC_DTYPES_CUH_; include only one" +#endif + +#ifndef FLASHINFER_ROCM_VEC_DTYPES_HIP_H_ +#define FLASHINFER_ROCM_VEC_DTYPES_HIP_H_ #define HIP_ENABLE_WARP_SYNC_BUILTINS 1 +// clang-format off +// macros.hpp first: its non-HIP #error must fire before a missing hip/ header. +#include "macros.hpp" +// clang-format on + #include #include #include @@ -17,8 +27,14 @@ #include +#include "platform.hpp" + #define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ +// Opened here, after the system includes: this header used to be included from +// inside these namespaces, so wrapping the includes too would nest them. +namespace flashinfer { + namespace { __host__ __device__ inline __hip_bfloat162 __float2bfloat162_rn(const float a) { return __hip_bfloat162{__float2bfloat16(a), __float2bfloat16(a)}; @@ -32,13 +48,8 @@ FLASHINFER_INLINE __hip_bfloat162 make_bfloat162(const __hip_bfloat16 x, const _ } } // namespace -namespace detail { -namespace hip { - #define FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED -#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ - #if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 < 120400) && \ (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)) // CUDA version < 12.4 and GPU architecture < 80 @@ -1906,5 +1917,5 @@ struct vec_t { } } }; -} // namespace hip -} // namespace detail + +} // namespace flashinfer diff --git a/include/gpu_iface/backend/hip/mma_debug_utils_hip.h b/include/gpu_iface/backend/hip/mma_debug_utils_hip.h deleted file mode 100644 index a992fa9f27..0000000000 --- a/include/gpu_iface/backend/hip/mma_debug_utils_hip.h +++ /dev/null @@ -1,411 +0,0 @@ -// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once - -#include "gpu_iface/backend/hip/mma_hip.h" -#include "gpu_iface/fastdiv.cuh" -#include "gpu_iface/gpu_runtime_compat.hpp" - -namespace { -constexpr uint32_t MMA_COLS = 16; -constexpr uint32_t MMA_ROWS_PER_THREAD = 4; -} // namespace - -namespace flashinfer::gpu_iface::debug_utils::hip { - -enum class MatrixLayout { A, B }; - -/// @brief Initializes a 2D LDS array with lexicographical values (0, 1, 2, ...). -/// @param lds_array Pointer to the shared memory array. -/// @param dimY The height of the 2D array. -/// @param dimX The width of the 2D array. -__device__ void lexicographic_init_lds_array(half* lds_array, uint32_t dimY, uint32_t dimX) { - const int tid = threadIdx.x; - if (tid == 0) { - for (int y = 0; y < dimY; ++y) { - for (int x = 0; x < dimX; ++x) { - lds_array[y * dimX + x] = __half(y * dimX + x); - } - } - } - __syncthreads(); -} - -/// @brief Loads a 16x16 tile from LDS into registers using the A-matrix layout pattern. -/// @details Each thread `T_(16*c + r)` loads a 1x4 horizontal fragment from `LDS[r, 4*c : 4*c+3]`. -/// @tparam T The data type of the LDS array, must be `__half`. -/// @param lds_array Pointer to the shared memory array. -/// @param R Pointer to the thread's registers (uint32_t[2]). -/// @param dimX The width of the LDS array. -template -__device__ void load_amatrix_layout(T* lds_array, uint32_t* R, uint32_t dimX) { - static_assert(std::is_same_v, "Only supported for __half types"); - const int lane_id = threadIdx.x % 64; - const int row = lane_id % MMA_COLS; - const int col_start = (lane_id / MMA_COLS) * MMA_ROWS_PER_THREAD; - - auto offset = lds_array + row * dimX + col_start; - mma_impl::hip::load_fragment(R, offset); -} - -/// @brief Loads a 16x16 tile from LDS into registers using the B-matrix layout pattern. -/// @details Uses an efficient load-and-transpose strategy. A 4x4 block of threads loads a -/// contiguous 4x4 tile from LDS and then performs an in-register transpose, -/// resulting in each thread holding a column fragment. -/// @tparam T The data type of the LDS array, must be `__half`. -/// @param arr Pointer to the shared memory array. -/// @param R Pointer to the thread's registers (uint32_t[2]). -/// @param dimY The height of the LDS array. -template -__device__ void load_bmatrix_layout(T* arr, uint32_t* R, uint32_t dimY) { - static_assert(std::is_same_v, "Only supported for __half types"); - const int lane_id = threadIdx.x % 64; - int b_idx = - ((lane_id % MMA_ROWS_PER_THREAD) + MMA_ROWS_PER_THREAD * (lane_id / MMA_COLS)) * dimY + - ((lane_id % MMA_COLS) / MMA_ROWS_PER_THREAD) * MMA_ROWS_PER_THREAD; - mma_impl::hip::load_quad_transposed_fragment<__half>(R, &arr[b_idx]); -} - -/// @brief Prints a single MMA fragment (typically 4 or 8 elements). -/// @details Simple low-level printer for a single [ELEMS_PER_FRAGMENT] array. -/// Works for both A-matrix layout (row strip) and B-matrix layout (column strip). -/// @tparam T The data type of the fragment (e.g., float, __half). -/// @tparam ELEMS_PER_FRAGMENT The number of elements per fragment (typically 4 or 8). -/// @param values Pointer to the fragment values. -template -__device__ void debug_print_frag(const T* values) { - printf("["); - for (uint32_t i = 0; i < ELEMS_PER_FRAGMENT; ++i) { - printf("%10.6f", float(values[i])); - if (i < ELEMS_PER_FRAGMENT - 1) printf(", "); - } - printf("]"); -} - -/// @brief Prints all MMA fragments from a thread's registers. -/// @details Loops over [NUM_MMA_ROW][NUM_MMA_COL][ELEMS_PER_FRAGMENT] array. -/// Works for all fragment types: Q, K (A-matrix), S, O (B-matrix), etc. -/// @tparam T The data type of the fragments (e.g., float, __half). -/// @tparam NUM_MMA_ROW Number of MMA tiles in the row dimension. -/// @tparam NUM_MMA_COL Number of MMA tiles in the column dimension. -/// @tparam ELEMS_PER_FRAGMENT The number of elements per fragment (typically 4 or 8). -/// @param frag The 3D fragment array from the thread's registers. -/// @param frag_name A string name to identify which fragment is being printed. -/// @param tidx The x component of the thread to print from. -/// @param tidy The y component of the thread to print from. -/// @param tidz The z component of the thread to print from. -template -__device__ void debug_print_frag_registers(const T (*frag)[NUM_MMA_COL][ELEMS_PER_FRAGMENT], - const char* frag_name = "frag", const uint32_t tidx = 0, - const uint32_t tidy = 0, const uint32_t tidz = 0) { - if (threadIdx.x == tidx && threadIdx.y == tidy && threadIdx.z == tidz) { - printf("Thread (%u,%u,%u) %s registers:\n", tidx, tidy, tidz, frag_name); - for (uint32_t mma_row = 0; mma_row < NUM_MMA_ROW; ++mma_row) { - for (uint32_t mma_col = 0; mma_col < NUM_MMA_COL; ++mma_col) { - printf(" %s[%u][%u]: ", frag_name, mma_row, mma_col); - debug_print_frag(frag[mma_row][mma_col]); - printf("\n"); - } - } - printf("\n"); - } -} - -/// @brief Prints a 2D LDS array to the console from a single thread. -/// @tparam T The data type of the LDS array, must be `__half`. -/// @param lds_array Pointer to the shared memory array. -/// @param dimY The height of the 2D array. -/// @param dimX The width of the 2D array. -template -__device__ void print_lds_array(T* lds_array, uint32_t dimY, uint32_t dimX, - const char* title = "LDS Array") { - static_assert(std::is_same_v, "Only supported for __half types"); - if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) { - printf("%s (%dx%d):\n", title, dimY, dimX); - for (int y = 0; y < dimY; ++y) { - for (int x = 0; x < dimX; ++x) { - if (x == dimX - 1) { - printf("%10.6f", (float)lds_array[y * dimX + x]); - } else { - printf("%10.6f ", float(lds_array[y * dimX + x])); - } - } - printf("\n"); - } - printf("\n"); - } - __syncthreads(); -} - -/// @brief Prints a 2D LDS array of floats to the console from a single thread. -__device__ void print_lds_array(float* lds_array, uint32_t dimY, uint32_t dimX, - const char* title = "LDS Array (float)") { - if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) { - printf("%s (%dx%d):\n", title, dimY, dimX); - for (int y = 0; y < dimY; ++y) { - for (int x = 0; x < dimX; ++x) { - if (x == dimX - 1) { - printf("%10.6f", lds_array[y * dimX + x]); - } else { - printf("%10.6f ", lds_array[y * dimX + x]); - } - } - printf("\n"); - } - printf("\n"); - } -} - -/// @brief Prints a 1D LDS array of floats to the console from a single thread. -/// @details Useful for printing row-wise statistics like m or d values. -__device__ void print_lds_array_1d(float* lds_array, uint32_t dim, - const char* title = "LDS Array 1D (float)") { - if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) { - printf("%s (%d elements):\n", title, dim); - for (int i = 0; i < dim; ++i) { - printf("%10.6f ", lds_array[i]); - if ((i + 1) % 16 == 0) printf("\n"); // Line break every 16 elements - } - if (dim % 16 != 0) printf("\n"); - printf("\n"); - } - __syncthreads(); -} - -/// @brief Writes an A-matrix fragment from registers to shared memory. -/// @details In the A-matrix layout, each thread owns a row slice of a 16x16 fragment. -/// Thread T_(16*c + r) owns row r, columns [4*c : 4*c+3]. -/// This function reconstructs the full logical tile from distributed row fragments. -/// @tparam T The data type of the fragments and LDS array (e.g., float or half). -/// @tparam NUM_MMA_ROW The number of fragments along the rows dimension per thread. -/// @tparam NUM_MMA_COL The number of fragments along the column dimension per thread. -/// @tparam ELEMS_PER_FRAGMENT The number of elements per fragment (typically 4). -/// @param frag The 3D fragment array from the thread's registers. -/// @param lds_scratchpad Pointer to the shared memory array. -/// @param lds_stride The width/stride of the lds_scratchpad. -/// @param tid The thread's index within the block (threadIdx). -template -__device__ void write_amatrix_frag_to_lds(const T (*frag)[NUM_MMA_COL][ELEMS_PER_FRAGMENT], - T* lds_scratchpad, const uint32_t lds_stride, - const dim3 tid = threadIdx) { - const int lane_id = tid.x % 64; - const int warp_idx_q = tid.y; - - // Calculate the starting row in the LDS tile for this entire warp. - const uint32_t warp_base_row = warp_idx_q * NUM_MMA_ROW * MMA_COLS; - -#pragma unroll - for (uint32_t mma_row = 0; mma_row < NUM_MMA_ROW; ++mma_row) { -#pragma unroll - for (uint32_t mma_col = 0; mma_col < NUM_MMA_COL; ++mma_col) { - // -- Calculate the top-left corner of the 16x16 fragment this thread contributes to -- - const uint32_t frag_row_offset = mma_row * MMA_COLS; - const uint32_t frag_col_offset = mma_col * MMA_COLS; - - // -- Calculate the specific 1x4 element strip this thread writes within that fragment -- - // A-matrix layout: each thread handles a row strip. - // Thread lane_id = 16*c + r owns row r, columns [4*c : 4*c+3] - const uint32_t thread_row_in_frag = lane_id % MMA_COLS; - const uint32_t thread_start_col_in_frag = (lane_id / MMA_COLS) * MMA_ROWS_PER_THREAD; - - // -- Combine all offsets and write the 1x4 row strip to LDS -- - const T* values = frag[mma_row][mma_col]; - - // The row is fixed for all 4 elements in the strip. - const uint32_t final_row = warp_base_row + frag_row_offset + thread_row_in_frag; - - for (int i = 0; i < MMA_ROWS_PER_THREAD; ++i) { - // The column for this element is the thread's starting column + the element's index. - const uint32_t final_col = frag_col_offset + thread_start_col_in_frag + i; - - // Calculate destination and write the value. - T* dest = lds_scratchpad + final_row * lds_stride + final_col; - *dest = values[i]; - } - } - } -} - -/// @brief Generic function to materialize 2D fragment arrays into shared memory. -/// @details Works for both s_frag (attention scores) and o_frag (output accumulator). -/// Reconstructs a logical tile from distributed register fragments. -/// @tparam T The data type of the fragments and LDS array (e.g., float or half). -/// @tparam NUM_MMA_ROW The number of fragments along the rows dimension per thread. -/// @tparam NUM_MMA_COL The number of fragments along the column dimension per thread. -/// For s_frag: NUM_MMA_KV (KV sequence length) -/// For o_frag: NUM_MMA_D_VO (head dimension) -/// @tparam ELEMS_PER_FRAGMENT The number of elements per fragment (typically 4). -/// @param frag The 3D fragment array from the thread's registers. -/// @param lds_scratchpad Pointer to the shared memory array. -/// @param lds_stride The width/stride of the lds_scratchpad. -/// @param tid The thread's index within the block (threadIdx). -template -__device__ void write_frag_to_lds(const T (*frag)[NUM_MMA_COL][ELEMS_PER_FRAGMENT], - T* lds_scratchpad, const uint32_t lds_stride, - const dim3 tid = threadIdx) { - const int lane_id = tid.x % 64; - const int warp_idx_q = tid.y; - - // Calculate the starting row in the LDS tile for this entire warp. - const uint32_t warp_base_row = warp_idx_q * NUM_MMA_ROW * MMA_COLS; - -#pragma unroll - for (uint32_t mma_q = 0; mma_q < NUM_MMA_ROW; ++mma_q) { -#pragma unroll - for (uint32_t mma_col = 0; mma_col < NUM_MMA_COL; ++mma_col) { - // -- Calculate the top-left corner of the 16x16 fragment this thread contributes to -- - const uint32_t frag_row_offset = mma_q * MMA_COLS; - const uint32_t frag_col_offset = mma_col * MMA_COLS; - - // -- Calculate the specific 4x1 element strip this thread writes within that fragment -- - // This logic correctly materializes a B-layout fragment (column strip). - // Each thread T_c handles column 'c' of the fragment. - // The 4 threads in a "column" of the warp (e.g., lanes 0, 16, 32, 48) - // handle the 4 rows of that column strip. - const uint32_t thread_start_row_in_frag = (lane_id / MMA_COLS) * MMA_ROWS_PER_THREAD; - const uint32_t thread_col_in_frag = (lane_id % MMA_COLS); - - // -- Combine all offsets and write the 4x1 column strip to LDS -- - const T* values = frag[mma_q][mma_col]; - for (int i = 0; i < MMA_ROWS_PER_THREAD; ++i) { - // The row for this element is the thread's starting row + the element's index in the strip. - const uint32_t final_row = warp_base_row + frag_row_offset + thread_start_row_in_frag + i; - // The column is fixed for all 4 elements in the strip. - const uint32_t final_col = frag_col_offset + thread_col_in_frag; - - // Calculate destination and write the value. - T* dest = lds_scratchpad + final_row * lds_stride + final_col; - *dest = values[i]; - } - } - } -} - -/// @brief Convenience wrapper for s_frag (attention scores). -template -__device__ void write_s_frag_to_lds(const T (*s_frag)[NUM_MMA_KV][ELEMS_PER_FRAGMENT], - T* lds_scratchpad, const uint32_t lds_stride, - const dim3 tid = threadIdx) { - write_frag_to_lds(s_frag, lds_scratchpad, - lds_stride, tid); -} - -/// @brief Convenience wrapper for o_frag (output accumulator). -template -__device__ void write_o_frag_to_lds(const T (*o_frag)[NUM_MMA_D_VO][ELEMS_PER_FRAGMENT], - T* lds_scratchpad, const uint32_t lds_stride, - const dim3 tid = threadIdx) { - write_frag_to_lds(o_frag, lds_scratchpad, - lds_stride, tid); -} - -/// @brief Generic function to materialize 1D row-wise values (m or d) into shared memory. -/// @details Writes row-wise statistics (like max or denominator) from register arrays -/// to a 1D shared memory array, with one value per row. -/// @tparam T The data type (typically float). -/// @tparam NUM_MMA_Q The number of fragments along the Q dimension per thread. -/// @tparam NUM_ACCUM_ROWS_PER_THREAD The number of accumulator rows per thread (typically 4). -/// @param values The 2D array from registers [NUM_MMA_Q][NUM_ACCUM_ROWS_PER_THREAD]. -/// @param lds_scratchpad Pointer to the 1D shared memory array. -/// @param tid The thread's index within the block (threadIdx). -template -__device__ void write_row_values_to_lds(const T (*values)[NUM_ACCUM_ROWS_PER_THREAD], - T* lds_scratchpad, const dim3 tid = threadIdx) { - const int lane_idx = tid.x; - const int warp_idx_q = tid.y; - - // Each group of 16 threads (a "row group") handles 4 rows. - // We only need one thread from each group to write the results. - if (lane_idx % MMA_COLS == 0) { - // Base row index for this warp's Q tile - const uint32_t warp_base_row = warp_idx_q * NUM_MMA_Q * MMA_COLS; - -#pragma unroll - for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { - // Base row for this specific MMA instruction within the warp's tile - const uint32_t mma_base_row = mma_q * MMA_COLS; - -#pragma unroll - for (uint32_t j = 0; j < NUM_ACCUM_ROWS_PER_THREAD; ++j) { - // The thread's lane_idx determines which group of 4 rows it is in. - // e.g., lane 0 is in group 0, lane 16 is in group 1, etc. - const uint32_t row_group_offset = (lane_idx / MMA_COLS) * NUM_ACCUM_ROWS_PER_THREAD; - - // The final row index in the logical matrix - const uint32_t final_row_idx = warp_base_row + mma_base_row + row_group_offset + j; - - lds_scratchpad[final_row_idx] = values[mma_q][j]; - } - } - } -} - -/// @brief Convenience wrapper for m (row-wise max) values. -template -__device__ void write_m_to_lds(const T (*m)[NUM_ACCUM_ROWS_PER_THREAD], T* lds_scratchpad, - const dim3 tid = threadIdx) { - write_row_values_to_lds(m, lds_scratchpad, tid); -} - -/// @brief Convenience wrapper for d (denominator) values. -template -__device__ void write_d_to_lds(const T (*d)[NUM_ACCUM_ROWS_PER_THREAD], T* lds_scratchpad, - const dim3 tid = threadIdx) { - write_row_values_to_lds(d, lds_scratchpad, tid); -} - -// Legacy alias for backward compatibility -template -__device__ void write_m_new_to_lds(const T (*m)[NUM_ACCUM_ROWS_PER_THREAD], T* lds_scratchpad, - const dim3 tid = threadIdx) { - write_m_to_lds(m, lds_scratchpad, tid); -} - -/// @brief Reads O matrix from global memory and prints it. -/// @details This function reads back the O matrix that was written to global memory -/// by write_o_reg_gmem and prints it for validation. -/// @tparam DTypeO The data type of the O matrix in global memory (typically __half). -/// @param o_ptr_base Pointer to the base of the O matrix in global memory. -/// @param o_stride_n Stride between consecutive queries (sequence dimension). -/// @param o_stride_h Stride between consecutive heads. -/// @param num_rows Number of rows to read (typically CTA_TILE_Q = 128). -/// @param num_cols Number of columns to read (typically HEAD_DIM = 64). -/// @param qo_packed_idx_base Base index for query packing (for GQA). -/// @param group_size Group size for grouped query attention. -/// @param kv_head_idx The KV head index. -/// @param header_text Optional header text to print before the matrix. -/// @param tid Thread index. -template -__device__ void debug_print_o_from_gmem(DTypeO* o_ptr_base, const uint32_t o_stride_n, - const uint32_t o_stride_h, const uint32_t num_rows, - const uint32_t num_cols, const uint32_t qo_packed_idx_base, - const uint_fastdiv group_size, const uint32_t kv_head_idx, - const char* header_text = "O from global memory", - const dim3 tid = threadIdx) { - if (tid.x == 0 && tid.y == 0 && tid.z == 0) { - printf("\n%s (%dx%d):\n", header_text, num_rows, num_cols); - - for (uint32_t row = 0; row < num_rows; ++row) { - // Compute the q and r indices for GQA - uint32_t q, r; - group_size.divmod(qo_packed_idx_base + row, q, r); - const uint32_t qo_head_idx = kv_head_idx * group_size + r; - - // Print row values - for (uint32_t col = 0; col < num_cols; ++col) { - DTypeO* ptr = o_ptr_base + q * o_stride_n + qo_head_idx * o_stride_h + col; - float val = float(*ptr); - printf("%10.6f", val); - if (col < num_cols - 1) { - printf(" "); - } - } - printf("\n"); - } - printf("\n"); - } - __syncthreads(); -} - -} // namespace flashinfer::gpu_iface::debug_utils::hip diff --git a/include/gpu_iface/error.hpp b/include/gpu_iface/error.hpp deleted file mode 100644 index 45bdbea595..0000000000 --- a/include/gpu_iface/error.hpp +++ /dev/null @@ -1,37 +0,0 @@ -// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once -#include -#include -#include - -#include "platform.hpp" - -namespace flashinfer { -namespace gpu_iface { - -// Platform-agnostic error type -class GpuError { - private: - int code_; - std::string message_; - - public: - GpuError() : code_(0) {} - GpuError(int code, std::string message) : code_(code), message_(std::move(message)) {} - - bool isSuccess() const { return code_ == 0; } - int code() const { return code_; } - const std::string& message() const { return message_; } - - hipError_t getNative() const { return static_cast(code_); } -}; - -// Create error from message -inline GpuError CreateError(std::string message) { - return GpuError(static_cast(hipErrorUnknown), std::move(message)); -} - -} // namespace gpu_iface -} // namespace flashinfer diff --git a/include/gpu_iface/exception.h b/include/gpu_iface/exception.h deleted file mode 100644 index 74071938ee..0000000000 --- a/include/gpu_iface/exception.h +++ /dev/null @@ -1,38 +0,0 @@ -// SPDX-FileCopyrightText: 2024-2025 FlashInfer team. -// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. -// SPDX-License-Identifier: Apache-2.0 - -#ifndef FLASHINFER_EXCEPTION_H_ -#define FLASHINFER_EXCEPTION_H_ - -#include -#include -#include - -namespace flashinfer { - -class Error : public std::exception { - private: - std::string message_; - - public: - Error(const std::string& func, const std::string& file, int line, const std::string& message) { - std::ostringstream oss; - oss << "Error in function '" << func << "' " - << "at " << file << ":" << line << ": " << message; - message_ = oss.str(); - } - - virtual const char* what() const noexcept override { return message_.c_str(); } -}; - -#define FLASHINFER_ERROR(message) throw Error(__FUNCTION__, __FILE__, __LINE__, message) - -#define FLASHINFER_CHECK(condition, message) \ - if (!(condition)) { \ - FLASHINFER_ERROR(message); \ - } - -} // namespace flashinfer - -#endif // FLASHINFER_EXCEPTION_H_ diff --git a/include/gpu_iface/math_ops.hpp b/include/gpu_iface/math_ops.hpp deleted file mode 100644 index a2119a1ca1..0000000000 --- a/include/gpu_iface/math_ops.hpp +++ /dev/null @@ -1,25 +0,0 @@ -// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once -#include "macros.hpp" - -// Include platform-specific implementations -#include "backend/hip/math_hip.h" - -namespace flashinfer { -namespace gpu_iface { -namespace math { -using flashinfer::math::inf; -using flashinfer::math::log2e; -using flashinfer::math::loge2; -using flashinfer::math::ptx_exp2; -using flashinfer::math::ptx_log2; -using flashinfer::math::ptx_rcp; -using flashinfer::math::rsqrt; -using flashinfer::math::shfl_xor_sync; -using flashinfer::math::tanh; - -} // namespace math -} // namespace gpu_iface -} // namespace flashinfer diff --git a/include/gpu_iface/memory_ops.hpp b/include/gpu_iface/memory_ops.hpp deleted file mode 100644 index c244344d6f..0000000000 --- a/include/gpu_iface/memory_ops.hpp +++ /dev/null @@ -1,124 +0,0 @@ -// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once -#include "platform.hpp" - -namespace flashinfer { -namespace gpu_iface { -namespace memory { - -/** - * @brief Control options for shared memory fill behavior - */ -enum class SharedMemFillMode { - kFillZero, // Fill zero to shared memory when predicate is false - kNoFill // Do not fill zero to shared memory when predicate is false -}; - -/** - * @brief Control options for memory prefetch behavior - */ -enum class PrefetchMode { - kNoPrefetch, // Do not fetch additional data from global memory to L2 - kPrefetch // Fetch additional data from global memory to L2 -}; - -// Include platform-specific implementations -#include "backend/hip/memory_ops_hip.h" -namespace mem_detail = flashinfer::gpu_iface::memory::detail::hip; - -/** - * @brief Commits pending asynchronous memory operations to a group - */ -__device__ __forceinline__ void commit_group() { mem_detail::commit_group(); } - -/** - * @brief Waits until N most recent groups of async operations are complete - * - * @tparam N Number of most recent groups to wait for (0-7) - */ -template -__device__ __forceinline__ void wait_group() { - mem_detail::wait_group(); -} - -/** - * @brief Asynchronously loads 128 bits from global to shared memory - * - * @tparam PrefetchOpt Prefetch option - * @tparam T Data type - * @param smem_ptr Destination shared memory pointer - * @param gmem_ptr Source global memory pointer - */ -template -__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) { - mem_detail::load_128b(smem_ptr, gmem_ptr); -} - -template -__device__ __forceinline__ void load_64b(T* smem_ptr, const T* gmem_ptr) { -#if defined(PLATFORM_HIP_DEVICE) - mem_detail::load_64b(smem_ptr, gmem_ptr); -#else -#error "load_64b not implemented for this platform" -#endif -} - -/** - * @brief Conditionally loads 128 bits from global to shared memory - * - * @tparam PrefetchOpt Prefetch option - * @tparam FillOpt Memory fill option - * @tparam T Data type - * @param smem_ptr Destination shared memory pointer - * @param gmem_ptr Source global memory pointer - * @param predicate Condition for executing the load - */ -template -__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) { - mem_detail::pred_load_128b(smem_ptr, gmem_ptr, predicate); -} - -template -__device__ __forceinline__ void pred_load_64b(T* smem_ptr, const T* gmem_ptr, bool predicate) { -#if defined(PLATFORM_HIP_DEVICE) - mem_detail::pred_load_64b(smem_ptr, gmem_ptr, predicate); -#else -#error "pred_load_64b not implemented for this platform" -#endif -} - -/** - * @brief Loads N bits (128 or 256) from global to shared memory - * - * @tparam NumBits Number of bits to load (128 or 256) - * @tparam PrefetchOpt Prefetch option - * @tparam T Data type - * @param smem_ptr Destination shared memory pointer - * @param gmem_ptr Source global memory pointer - */ -template -__device__ __forceinline__ void load(T* smem_ptr, const T* gmem_ptr) { - mem_detail::load(smem_ptr, gmem_ptr); -} - -/** - * @brief Conditionally loads N bits from global to shared memory - * - * @tparam NumBits Number of bits to load (128 or 256) - * @tparam PrefetchOpt Prefetch option - * @tparam FillOpt Memory fill option - * @tparam T Data type - * @param smem_ptr Destination shared memory pointer - * @param gmem_ptr Source global memory pointer - * @param predicate Condition for executing the load - */ -template -__device__ __forceinline__ void pred_load(T* smem_ptr, const T* gmem_ptr, bool predicate) { - mem_detail::pred_load(smem_ptr, gmem_ptr, predicate); -} - -} // namespace memory -} // namespace gpu_iface -} // namespace flashinfer diff --git a/include/gpu_iface/mma_ops.hpp b/include/gpu_iface/mma_ops.hpp deleted file mode 100644 index 98edeb09c4..0000000000 --- a/include/gpu_iface/mma_ops.hpp +++ /dev/null @@ -1,65 +0,0 @@ -// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once -#include "gpu_iface/mma_types.hpp" -#include "gpu_iface/platform.hpp" - -// Include platform-specific implementations -#include "backend/hip/mma_hip.h" -namespace mma_detail = flashinfer::gpu_iface::mma_impl::hip; - -namespace flashinfer { -namespace gpu_iface { -namespace mma { - -/*! - * \brief Loads data from shared memory to fragment - * \tparam T data type of the fragment - * \param R pointer to the fragment - * \param smem_ptr pointer to the shared memory - */ -// Call this load fragment -// inside mma there is impl of load - -template -__device__ __forceinline__ void load_fragment(uint32_t* R, const T* smem_ptr) { - mma_detail::load_fragment(R, smem_ptr); -} - -#if defined(PLATFORM_HIP_DEVICE) -/*! - * \brief Performs a full 16x16 in-register matrix transpose for CDNA3 MFMA tiles - * \details Converts between A-matrix layout (row-major) and B/C/D-matrix layout (column-major) - * by combining intra-quad and inter-quad fragment transpositions. - * \param R Pointer to 2 uint32_t registers containing the fragment data - */ -__device__ __forceinline__ void transpose_mma_tile(uint32_t* R) { - mma_detail::transpose_mma_tile(R); -} -#endif - -/*! - * \brief An m16n16k16 gemm kernel using MMA instructions for row - * major and column major f16 matrix multiplication, accumulated in f32. - * - * \tparam T data type of the fragment - * \tparam mma_mode whether we are initializing the accumulator or updating it - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, - uint32_t* B) { - mma_detail::mma_sync_m16n16k16_row_col_f16f16f32(C, A, B); -} - -template -__device__ __forceinline__ void m16k16_rowsum_f16f16f32(float* d, DType* s) { - mma_detail::m16k16_rowsum_f16f16f32(d, s); -} - -} // namespace mma -} // namespace gpu_iface -} // namespace flashinfer diff --git a/include/gpu_iface/vec_dtypes.hpp b/include/gpu_iface/vec_dtypes.hpp deleted file mode 100644 index daa1f1876d..0000000000 --- a/include/gpu_iface/vec_dtypes.hpp +++ /dev/null @@ -1,25 +0,0 @@ -// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once -#include -#include - -#include "platform.hpp" - -namespace flashinfer { -namespace gpu_iface { -namespace vec_dtypes { - -// Include the appropriate backend implementation -#include "backend/hip/vec_dtypes_hip.h" -namespace vec_t_detail = flashinfer::gpu_iface::vec_dtypes::detail::hip; - -// Re-export types and functions from the appropriate backend -// This allows code to use flashinfer::gpu_iface::vec_dtypes::vec_t -using vec_t_detail::vec_cast; -using vec_t_detail::vec_t; - -} // namespace vec_dtypes -} // namespace gpu_iface -} // namespace flashinfer diff --git a/pyproject.toml b/pyproject.toml index bdff3d2a29..66da8c8133 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,15 +82,15 @@ include = ["flashinfer", "flashinfer.*"] [tool.setuptools.package-data] # scikit-build-core's `wheel.packages = ["flashinfer"]` swept these in # implicitly; setuptools requires them to be named. Both matter at runtime: -# get_csrc_dir() resolves flashinfer/csrc_rocm and get_include() resolves +# get_csrc_dir() resolves flashinfer/csrc/rocm and get_include() resolves # flashinfer/include, so dropping either ships a wheel whose JIT cannot compile. # flashinfer/include itself is produced by build_backend_rocm.py, not by git. flashinfer = [ "py.typed", - "csrc_rocm/**/*.cu", - "csrc_rocm/**/*.cc", - "csrc_rocm/**/*.h", - "csrc_rocm/**/*.jinja", + "csrc/rocm/**/*.cu", + "csrc/rocm/**/*.cc", + "csrc/rocm/**/*.h", + "csrc/rocm/**/*.jinja", "include/**/*.cuh", "include/**/*.h", "include/**/*.hpp", diff --git a/scripts/amd_coverage.py b/scripts/amd_coverage.py index 904961a271..1855f6c55c 100644 --- a/scripts/amd_coverage.py +++ b/scripts/amd_coverage.py @@ -44,7 +44,7 @@ _HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@") -_CSRC_DIR = "flashinfer/csrc_rocm" +_CSRC_DIR = "flashinfer/csrc/rocm" _JIT_REACH_ENV = "FLASHINFER_JIT_REACH_DIR" _TIER_LABEL = { @@ -584,7 +584,7 @@ def _stamp_matches(out_dir: Path, data_file: Path) -> bool: def _jit_reach( repo: Path, out_dir: Path, data_file: Path ) -> Optional[Tuple[int, List[str]]]: - """(reached, unreached) csrc_rocm translation units, merged across xdist workers. + """(reached, unreached) csrc/rocm translation units, merged across xdist workers. Not a coverage figure: the HIP sources are JIT-compiled and have no line data. It says only which of them the Python tests caused to be built at all. @@ -699,7 +699,7 @@ def _report( if reach is not None: reached, unreached = reach total = reached + len(unreached) - print("== csrc_rocm reach ==") + print("== flashinfer/csrc/rocm reach ==") print( f" {reached} of {total} translation units were built and loaded by a test" ) @@ -722,7 +722,9 @@ def _report( f" {len(unowned)} upstream files on the measured surface, attributed to no tier" ) print(" tests/ and benchmarks/ are outside the measured surface") - print(" C++/HIP under csrc_rocm/ and include/ has no line coverage (JIT-built)") + print( + " C++/HIP under flashinfer/csrc/rocm/ and include/ has no line coverage (JIT-built)" + ) print() # None, not 0.0: with no executable statements outside the import baseline diff --git a/scripts/coverage_ownership.toml b/scripts/coverage_ownership.toml index c09d49e7b3..4d9b08b7e4 100644 --- a/scripts/coverage_ownership.toml +++ b/scripts/coverage_ownership.toml @@ -12,29 +12,29 @@ # generated artifacts, which a checkout that has never been built does not have. # Files with a zero-line Python diff whose implementation is nonetheless ours: -# jit_env.FLASHINFER_CSRC_DIR resolves to flashinfer/csrc_rocm/ on HIP, so these +# jit_env.FLASHINFER_CSRC_DIR resolves to flashinfer/csrc/rocm/ on HIP, so these # wrappers load AMD kernels through an unmodified upstream source path. A # diff-based classifier cannot see that, and scoring them as upstream would # credit none of the sampling, quantization or cascade surface to us. [[redirect_owned]] path = "flashinfer/sampling.py" -reason = "wrapper over csrc_rocm/sampling.cu and renorm.cu" +reason = "wrapper over flashinfer/csrc/rocm/sampling.cu and renorm.cu" [[redirect_owned]] path = "flashinfer/quantization.py" -reason = "wrapper over csrc_rocm/quantization.cu" +reason = "wrapper over flashinfer/csrc/rocm/quantization.cu" [[redirect_owned]] path = "flashinfer/jit/sampling.py" -reason = "builds csrc_rocm/{sampling,renorm,flashinfer_sampling_binding}.cu" +reason = "builds flashinfer/csrc/rocm/{sampling,renorm,flashinfer_sampling_binding}.cu" [[redirect_owned]] path = "flashinfer/jit/quantization.py" -reason = "builds csrc_rocm/{quantization,flashinfer_quantization_binding}.cu" +reason = "builds flashinfer/csrc/rocm/{quantization,flashinfer_quantization_binding}.cu" [[redirect_owned]] path = "flashinfer/jit/cascade.py" -reason = "builds csrc_rocm/{cascade,flashinfer_cascade_binding}.cu" +reason = "builds flashinfer/csrc/rocm/{cascade,flashinfer_cascade_binding}.cu" # Rulings that contradict what a heuristic would guess. Recorded so the answer # is reviewed once rather than re-derived per run. diff --git a/scripts/upstream_canary.py b/scripts/upstream_canary.py index 9391f50d29..5e77e2447f 100755 --- a/scripts/upstream_canary.py +++ b/scripts/upstream_canary.py @@ -40,6 +40,10 @@ _FORKED_DIR = "include/flashinfer/rocm" +# Shares a basename with an upstream header but forks nothing: a HIP rewrite +# that happens to reuse the names. Pairing it would report drift forever. +_FORKED_EXCLUDE = frozenset({f"{_FORKED_DIR}/utils.cuh"}) + # rocm/ mirrors upstream's layout but flattens one level: rocm/attention/ holds # headers upstream keeps both in attention/ and directly under flashinfer/. _UPSTREAM_HEADER_DIRS = ("include/flashinfer/attention", "include/flashinfer") @@ -254,7 +258,9 @@ def _report_drift(repo: str, ours: str, theirs: str, churn: Dict[str, Churn]) -> """ print("== forked-header drift ==") forked = sorted( - p for p in _ls_tree(repo, ours, _FORKED_DIR) if p.endswith(_HEADER_SUFFIXES) + p + for p in _ls_tree(repo, ours, _FORKED_DIR) + if p.endswith(_HEADER_SUFFIXES) and p not in _FORKED_EXCLUDE ) if not forked: print(f"no forked headers under {_FORKED_DIR} in {ours[:12]}") diff --git a/tests/jit_reach_plugin.py b/tests/jit_reach_plugin.py index 6f391903e0..1401cbc45c 100644 --- a/tests/jit_reach_plugin.py +++ b/tests/jit_reach_plugin.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: Apache-2.0 -"""Record which ``csrc_rocm/`` translation units a test run actually loaded. +"""Record which ``csrc/rocm/`` translation units a test run actually loaded. JIT-built HIP has no line coverage; this is the honest substitute, reported separately from the Python percentage. Enable with ``-p jit_reach_plugin`` @@ -27,7 +27,7 @@ def _record(sources) -> None: if name.endswith((".cu", ".cc")): # Match on basename: jit/attention/modules_hip.py copies each source # into FLASHINFER_GEN_SRC_DIR and registers the copy, so the - # recorded path points into the JIT cache, not into csrc_rocm/. + # recorded path points into the JIT cache, not into csrc/rocm/. _reached.add(name) diff --git a/tests/rocm_tests/test_amd_coverage.py b/tests/rocm_tests/test_amd_coverage.py index f232235121..7dc514c5c2 100644 --- a/tests/rocm_tests/test_amd_coverage.py +++ b/tests/rocm_tests/test_amd_coverage.py @@ -728,8 +728,8 @@ def test_reach_needs_a_stamp_tying_shards_to_this_data_file(self, tmp_path): A cutoff on mtime therefore rejects every genuine run; the stamp is what distinguishes "this run" from a previous one left in the same directory. """ - (tmp_path / "flashinfer" / "csrc_rocm").mkdir(parents=True) - (tmp_path / "flashinfer" / "csrc_rocm" / "a.cu").write_text("x") + (tmp_path / "flashinfer" / "csrc/rocm").mkdir(parents=True) + (tmp_path / "flashinfer" / "csrc/rocm" / "a.cu").write_text("x") data = tmp_path / ".coverage" shard = tmp_path / "jit-reach.gw0.json" shard.write_text('["a.cu"]', encoding="utf-8") @@ -753,9 +753,9 @@ def test_zero_reached_is_reported_not_treated_as_absent(self, tmp_path): The plugin writes a shard even with nothing loaded, so an empty shard must survive as a real result rather than collapsing to None. """ - (tmp_path / "flashinfer" / "csrc_rocm").mkdir(parents=True) + (tmp_path / "flashinfer" / "csrc/rocm").mkdir(parents=True) for name in ("a.cu", "b.cu"): - (tmp_path / "flashinfer" / "csrc_rocm" / name).write_text("x") + (tmp_path / "flashinfer" / "csrc/rocm" / name).write_text("x") data = tmp_path / ".coverage" (tmp_path / "jit-reach.gw0.json").write_text("[]", encoding="utf-8") data.write_text("x", encoding="utf-8") diff --git a/tests/rocm_tests/test_build_backend.py b/tests/rocm_tests/test_build_backend.py index 8cd462ccab..26ea1d7c3e 100644 --- a/tests/rocm_tests/test_build_backend.py +++ b/tests/rocm_tests/test_build_backend.py @@ -178,7 +178,7 @@ def _isolated_project(dest): def test_wheel_carries_the_paths_the_jit_resolves(tmp_path, monkeypatch): """Assert on the artifact: the helpers above do not prove what setuptools ships. - csrc_rocm reaches the wheel through both package-data and MANIFEST.in, so + csrc/rocm reaches the wheel through both package-data and MANIFEST.in, so the counts here catch losing it, not which of the two carried it. """ build = pytest.importorskip("build") @@ -201,10 +201,10 @@ def test_wheel_carries_the_paths_the_jit_resolves(tmp_path, monkeypatch): for f in (_REPO_ROOT / "include").rglob("*") if f.suffix in {".cuh", ".h", ".hpp"} ) - csrc = {f for f in names if f.startswith("flashinfer/csrc_rocm/")} + csrc = {f for f in names if f.startswith("flashinfer/csrc/rocm/")} assert len(csrc) == sum( 1 - for f in (_REPO_ROOT / "flashinfer" / "csrc_rocm").rglob("*") + for f in (_REPO_ROOT / "flashinfer" / "csrc/rocm").rglob("*") if f.suffix in {".cu", ".cc", ".h", ".jinja"} ) # The sibling projects leaked in once via an unanchored packages.find glob. diff --git a/tests/rocm_tests/test_customize_prefill_use_softmax_hip.py b/tests/rocm_tests/test_customize_prefill_use_softmax_hip.py index 05da35b7c5..a360a47e70 100644 --- a/tests/rocm_tests/test_customize_prefill_use_softmax_hip.py +++ b/tests/rocm_tests/test_customize_prefill_use_softmax_hip.py @@ -22,7 +22,7 @@ # ROCm's variant contract differs from the CUDA examples in # tests/utils/test_jit_example.py: window_left is read unconditionally by -# prefill.cuh, and the math helpers live in gpu_iface::math. +# prefill.cuh, and the math helpers live in math. FLASH_SIGMOID_DECL = r""" struct FlashSigmoid : AttentionVariantBase { static constexpr bool use_softmax = false; @@ -33,8 +33,8 @@ template __device__ __host__ FlashSigmoid(const Params& params, uint32_t batch_idx, uint8_t* smem_ptr) { - logits_scale_log2 = params.logits_scale * gpu_iface::math::log2e; - sigmoid_bias_log2e = params.sigmoid_bias * gpu_iface::math::log2e; + logits_scale_log2 = params.logits_scale * math::log2e; + sigmoid_bias_log2e = params.sigmoid_bias * math::log2e; qo_len = params.get_qo_len(batch_idx); kv_len = params.get_kv_len(batch_idx); window_left = kv_len; @@ -42,8 +42,8 @@ REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, { - return gpu_iface::math::ptx_rcp( - 1.f + gpu_iface::math::ptx_exp2(-float(logits) * logits_scale_log2 - + return math::ptx_rcp( + 1.f + math::ptx_exp2(-float(logits) * logits_scale_log2 - sigmoid_bias_log2e)); }) }; diff --git a/tests/rocm_tests/test_fused_moe_aiter_hip.py b/tests/rocm_tests/test_fused_moe_aiter_hip.py index ed8cd8da10..ac60aa1804 100644 --- a/tests/rocm_tests/test_fused_moe_aiter_hip.py +++ b/tests/rocm_tests/test_fused_moe_aiter_hip.py @@ -168,7 +168,7 @@ def test_supported_block_m_matches_the_shim(): shim = ( pathlib.Path(__file__).parents[2] / "flashinfer" - / "csrc_rocm" + / "csrc/rocm" / "fused_moe_aiter.cu" ).read_text() body = re.search(r"bool\s+is_supported_block_m\s*\([^)]*\)\s*\{(.*?)\}", shim, re.S) diff --git a/tests/rocm_tests/test_jit_env_hip.py b/tests/rocm_tests/test_jit_env_hip.py index 8cde3383f8..5f8d26b493 100644 --- a/tests/rocm_tests/test_jit_env_hip.py +++ b/tests/rocm_tests/test_jit_env_hip.py @@ -110,7 +110,7 @@ def test_nvshmem_helpers_stay_absent_on_rocm(): jit/comm.py imports fine on ROCm, so with these defined the failure moves to an absent nvidia.nvshmem -- or, with NVSHMEM_* set, to a build against - flashinfer/csrc_rocm/nvshmem_binding.cu, which does not exist. + flashinfer/csrc/rocm/nvshmem_binding.cu, which does not exist. """ from flashinfer.jit import comm as jit_comm from flashinfer.jit import env as e diff --git a/tests/rocm_tests/test_jit_flag_hooks.py b/tests/rocm_tests/test_jit_flag_hooks.py index 56973c080c..c8bf824791 100644 --- a/tests/rocm_tests/test_jit_flag_hooks.py +++ b/tests/rocm_tests/test_jit_flag_hooks.py @@ -47,7 +47,7 @@ def _load_module(): "flashinfer.jit": types.ModuleType("flashinfer.jit"), "flashinfer.jit.env": types.SimpleNamespace( FLASHINFER_INCLUDE_DIR=Path("/stub/include"), - FLASHINFER_CSRC_DIR=Path("/stub/csrc_rocm"), + FLASHINFER_CSRC_DIR=Path("/stub/csrc/rocm"), ), } saved = {k: sys.modules.get(k) for k in stubs} @@ -259,7 +259,7 @@ def test_non_system_mode_switches_own_headers_to_plain_include(self, tmp_path): text = _ninja(tmp_path, FLASHINFER_OWN_HEADERS_NON_SYSTEM="1") assert "-I/stub/include" in text - assert "-I/stub/csrc_rocm" in text + assert "-I/stub/csrc/rocm" in text assert "-isystem /stub/include" not in text # Third-party headers must stay -isystem regardless (ninja keeps the var). assert "-isystem $torch_home/include" in text diff --git a/tests/rocm_tests/test_torch_compile_hip.py b/tests/rocm_tests/test_torch_compile_hip.py index 9a7634387a..5711e15479 100644 --- a/tests/rocm_tests/test_torch_compile_hip.py +++ b/tests/rocm_tests/test_torch_compile_hip.py @@ -4,7 +4,7 @@ """Tests for ``FLASHINFER_USE_TORCH_CUSTOM_OPS`` and ``torch.compile`` on ROCm. -Because ``_USE_TORCH_CUSTOM_OPS`` is evaluated at import time, each test that +Because the flag is evaluated at import time, each test that needs a different env-var value runs in a subprocess so the module is freshly imported with the desired setting. """ diff --git a/tests/rocm_tests/test_upstream_canary.py b/tests/rocm_tests/test_upstream_canary.py index ad73ee800f..88a455d72d 100644 --- a/tests/rocm_tests/test_upstream_canary.py +++ b/tests/rocm_tests/test_upstream_canary.py @@ -432,6 +432,50 @@ def test_non_header_suffixes_are_ignored(self, repo, capsys): assert "no forked headers" in capsys.readouterr().out + def test_excluded_header_is_not_paired_against_upstream(self, repo, capsys): + """utils.cuh shares upstream's basename but forks nothing. + + Without the exclusion the pairing reports upstream churn against a HIP + rewrite that never tracked it. + """ + for path in uc._FORKED_EXCLUDE: + _write(repo, path, "// hip rewrite\n") + _write( + repo, f"include/flashinfer/{path.rsplit('/', 1)[1]}", "// upstream\n" + ) + _commit(repo, "headers") + + uc._report_drift( + str(repo), + "HEAD", + "HEAD", + { + f"include/flashinfer/{p.rsplit('/', 1)[1]}": uc.Churn(5, 5) + for p in uc._FORKED_EXCLUDE + }, + ) + out = capsys.readouterr().out + + assert "no forked headers" in out + assert "utils.cuh" not in out + + def test_a_sibling_of_an_excluded_header_is_still_paired(self, repo, capsys): + """The exclusion is exact paths, not a prefix or a basename match.""" + self._forked(repo, "layout.cuh") + _write(repo, "include/flashinfer/layout.cuh", "// upstream\n") + _commit(repo, "headers") + + uc._report_drift( + str(repo), + "HEAD", + "HEAD", + {"include/flashinfer/layout.cuh": uc.Churn(4, 4)}, + ) + + assert ( + "include/flashinfer/layout.cuh upstream +4/-4" in capsys.readouterr().out + ) + class TestResolve: def test_returns_the_full_oid(self, repo):