refactor(rocm): remove the gpu_iface layer and flatten the ROCm headers

.claude/skills/add-rocm-kernel/SKILL.md

@@ -6,21 +6,21 @@ 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)
 
 | Step | File | Purpose |
 | --- | --- | --- |
 | 1 | `include/flashinfer/<op>.cuh` | Framework-agnostic kernel + launcher template. **No `<torch/...>` includes here.** |
-| 2 | `flashinfer/csrc_rocm/<op>.cu` | PyTorch launcher: `at::Tensor` in, `at::hip::getCurrentHIPStream()`, `TORCH_CHECK`, `DISPATCH_PYTORCH_DTYPE_*`. |
-| 3 | `flashinfer/csrc_rocm/flashinfer_<op>_binding.cu` | `TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { m.def("<op>", <op>); }`. |
-| 4 (opt) | `flashinfer/csrc_rocm/<op>_customize_config.jinja` | Compile-time type specialization. Skip if runtime dispatch is enough. |
+| 2 | `flashinfer/csrc/rocm/<op>.cu` | PyTorch launcher: `at::Tensor` in, `at::hip::getCurrentHIPStream()`, `TORCH_CHECK`, `DISPATCH_PYTORCH_DTYPE_*`. |
+| 3 | `flashinfer/csrc/rocm/flashinfer_<op>_binding.cu` | `TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { m.def("<op>", <op>); }`. |
+| 4 (opt) | `flashinfer/csrc/rocm/<op>_customize_config.jinja` | Compile-time type specialization. Skip if runtime dispatch is enough. |
 | 5 | `flashinfer/jit/<op>.py` | `gen_<op>_module() -> JitSpec` via `gen_jit_spec(...)`. |
 | 6 | `flashinfer/<op>.py` | Python API: `@functools.cache` module loader, destination-passing (`out=`). |
 | 7 | `tests/rocm_tests/test_<op>_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/<op>.cu` | `flashinfer/csrc_rocm/<op>.cu` |
+| `csrc/<op>.cu` | `flashinfer/csrc/rocm/<op>.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_<op>_binding.cu` pattern is used by everything else. Both work — match the neighbors.
 
 ## CDNA3 (`gfx942`) vs CDNA4 (`gfx950`)

.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.

.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

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

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`.

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

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.