refactor: move the ROCm C++ to root csrc/rocm and apply the directory/rocm convention

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

@@ -6,24 +6,24 @@ 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`](../../../csrc/rocm/norm.cu) +
+[`flashinfer_norm_binding.cu`](../../../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`](../../../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 | `csrc/rocm/<op>.cu` | PyTorch launcher: `at::Tensor` in, `at::hip::getCurrentHIPStream()`, `TORCH_CHECK`, `DISPATCH_PYTORCH_DTYPE_*`. |
+| 3 | `csrc/rocm/flashinfer_<op>_binding.cu` | `TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { m.def("<op>", <op>); }`. |
+| 4 (opt) | `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. |
+| 7 | `tests/rocm/test_<op>.py` | Correctness tests; FP32 reference math, loose BF16 tolerances. |
 | 8 | `flashinfer/jit/__init__.py` (`IS_HIP` branch) | `from .<op> import gen_<op>_module as gen_<op>_module`. |
 | 9 | `flashinfer/__init__.py` (`IS_HIP` branch) | `from .<op> import <op> as <op>`. |
 | 10 (opt) | `flashinfer/aot_hip.py` | Register `gen_<op>_module` for pre-compiled wheels. |
@@ -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` | `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); }` |
@@ -46,18 +46,18 @@ When porting an upstream kernel, mechanically rewrite:
 | `c10::cuda::OptionalCUDAGuard` | `c10::hip::OptionalHIPGuardMasqueradingAsCUDA` |
 | `nvcc` flags via `extra_cuda_cflags=[...]` | **Same kwarg name** (`extra_cuda_cflags`) — internally routed to `hipcc`. |
 | `flashinfer/aot.py` registration | `flashinfer/aot_hip.py` |
-| `tests/test_op.py` | `tests/rocm_tests/test_op_hip.py` |
+| `tests/test_op.py` | `tests/rocm/test_op.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`. Sources live at `csrc/rocm/`; the build materializes them into the package, so the resolved path is `flashinfer/csrc/rocm/`. **Never hardcode either.** |
 | `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).
-- **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.
+- **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 intrinsic header in [`include/flashinfer/rocm/`](../../../include/flashinfer/rocm) rather than forking the kernel into `csrc/rocm/`. Existing ones: `mma.h`, `memory_ops.h`, `math.h`, `vec_dtypes.h`. Symbols go in `flashinfer::`, grouped by area (`flashinfer::math`, `flashinfer::memory`, `flashinfer::mma`), keeping the upstream namesake's names; each such header carries an `#error` tripwire on upstream's include guard.
 - **`-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`](../../../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`)
@@ -74,5 +74,5 @@ When porting an upstream kernel, mechanically rewrite:
 - [ ] Binding registered via `TORCH_LIBRARY_FRAGMENT`.
 - [ ] JIT generator uses `jit_env.FLASHINFER_CSRC_DIR` (not hardcoded `csrc/`).
 - [ ] Both `flashinfer/jit/__init__.py` and `flashinfer/__init__.py` IS_HIP branches updated.
-- [ ] Test file under `tests/rocm_tests/` named `test_*_hip.py`.
+- [ ] Test file under `tests/rocm/` named `test_*.py`.
 - [ ] `pre-commit run -a` clean.

.claude/skills/benchmark-kernel/SKILL.md

@@ -6,8 +6,8 @@ description: Guide for benchmarking FlashInfer+ROCm kernels on AMD Instinct (CDN
 # Benchmarking FlashInfer+ROCm Kernels
 
 For a real driver script to copy, see
-[`benchmarks/rocm_benchmarks/bench_fa2_prefill.py`](../../../benchmarks/rocm_benchmarks/bench_fa2_prefill.py) and [`benchmarks/rocm_benchmarks/bench_aiter_prefill.py`](../../../benchmarks/rocm_benchmarks/bench_aiter_prefill.py)
-For the in-repo profiler wrapper, see [`rocm_profiler/rocm_profiler.py`](../../../rocm_profiler/rocm_profiler.py).
+[`benchmarks/rocm/bench_fa2_prefill.py`](../../../benchmarks/rocm/bench_fa2_prefill.py) and [`benchmarks/rocm/bench_aiter_prefill.py`](../../../benchmarks/rocm/bench_aiter_prefill.py)
+For the in-repo profiler wrapper, see [`profiler/rocm/rocm_profiler.py`](../../../profiler/rocm/rocm_profiler.py).
 
 ## Timing method matrix
 
@@ -53,7 +53,7 @@ Or pass a path to a `rocprofv3`-native YAML for a custom counter set.
 
 Driver script flags: `--timing-only` (skip rocprofv3), `--skip-roofline`, `--replot` (regen PNG from existing CSVs, no GPU), `--list-presets`.
 
-Output (under `benchmarks/rocm_benchmarks/`, gitignored):
+Output (under `benchmarks/rocm/`, gitignored):
 
 ```text
 <label>_timing.csv             # median + std per config

.claude/skills/code-coverage/SKILL.md

@@ -90,7 +90,7 @@ docker exec <container> git config --global --add safe.directory '*'
 
 There is **no GPU CI** in this repo, so no automation produces this number —
 every workflow runs on `ubuntu-latest` or a CPU runner. The classifier's own
-tests (`tests/rocm_tests/test_amd_coverage.py`) do run there, via
+tests (`tests/rocm/test_amd_coverage.py`) do run there, via
 `arch-caps-conformance.yml`.
 
 ## Failure modes worth recognising

.claude/skills/pr-workflow/SKILL.md

@@ -132,26 +132,34 @@ literal path fails with `not a directory`.
 
 ### A fresh worktree is source-only
 
-Two gitignored, generated files must be recreated or the JIT will not build:
+Three gitignored, generated paths must be recreated or the JIT will not build:
 
 ```bash
 cd tmp/worktrees/<branch-name>
-rm -rf flashinfer/include                    # -f alone will not clear a real dir
+rm -rf flashinfer/include flashinfer/csrc    # -f alone will not clear a real dir
 ln -s ../include flashinfer/include          # MUST be relative
+mkdir -p flashinfer/csrc && ln -s ../../csrc/rocm flashinfer/csrc/rocm
 cp <main-checkout>/flashinfer/_version.py flashinfer/_version.py   # see below if absent
 ```
 
 Clear the path first. `-f` replaces a dangling or stale *symlink*, but against
 a real directory `ln` silently creates `flashinfer/include/include` **inside**
-it and exits 0 — leaving a broken tree with no error to go on. Deleting
-`flashinfer/include` is safe: it is gitignored and generated, and the real
-headers live in `include/` at the repo root.
+it and exits 0 — leaving a broken tree with no error to go on. Deleting both is
+safe: they are gitignored and generated, and the real sources live in
+`include/` and `csrc/rocm/` at the repo root.
 
+Link `flashinfer/csrc/rocm`, never `flashinfer/csrc` — the latter would expose
+upstream's 850-file CUDA `csrc/` tree inside the package.
+
+- `flashinfer/csrc/rocm` — `get_include_paths.get_csrc_dir()` returns
+  `<pkg>/csrc/rocm`, which becomes `FLASHINFER_CSRC_DIR`. Missing, the JIT
+  fails on a missing `.cu` rather than on the directory, because
+  `jit/core.py` creates it with `exist_ok=True` at import.
 - `flashinfer/include` — `get_include_paths.get_include()` returns
   `<pkg>/include`, and the JIT passes that through `.resolve()` into
   `-isystem`. Both failure modes emit a well-formed flag pointing at a
   directory that isn't there, so the compile fails with
-  `'flashinfer/attention/aiter/batch_prefill.cuh' file not found` rather than
+  `'flashinfer/rocm/attention/aiter/batch_prefill.cuh' file not found` rather than
   anything naming the include path:
   - missing entirely → `-isystem <pkg>/include`, a path that does not exist.
   - copied as an **absolute** symlink into a container mount point

.github/workflows/arch-caps-conformance.yml

@@ -1,7 +1,7 @@
 # Hardware-less conformance check for the per-op/per-arch capability table.
 #
 # Every other ROCm test in this repo needs an MI300X/MI355X and a ROCm toolchain,
-# so nothing in tests/rocm_tests/ runs on a pull request today. The capability
+# so nothing in tests/rocm/ runs on a pull request today. The capability
 # table is the one piece that can: it is pure data plus routing logic, and its
 # suite is written to load flashinfer/arch_caps.py and flashinfer/hip_utils.py
 # directly rather than importing the package (flashinfer/__init__.py raises on a
@@ -45,40 +45,40 @@ jobs:
           python -m pip install --upgrade pip pytest build
           python -m pip install 'setuptools>=80' 'setuptools-scm>=9.2' 'packaging>=24'
 
-      # --noconftest: tests/conftest.py and tests/rocm_tests/conftest.py both
+      # --noconftest: tests/conftest.py and tests/rocm/conftest.py both
       # import torch and flashinfer at module scope, so collecting under
-      # tests/rocm_tests/ would fail here long before this file is reached. The
+      # tests/rocm/ would fail here long before this file is reached. The
       # suite defines its own fixtures and needs nothing from either conftest;
       # test_suite_loads_without_torch is what keeps that true.
       - name: Capability-table conformance (no GPU, no torch)
-        run: pytest --noconftest tests/rocm_tests/test_arch_caps_hip.py
+        run: pytest --noconftest tests/rocm/test_arch_caps.py
 
       # Flag plumbing in the HIP JIT ninja generator. It stubs torch to load
       # cpp_ext_hip.py, which is what lets it run here; without this step that
-      # stubbing buys nothing, since no other lane collects tests/rocm_tests/.
+      # stubbing buys nothing, since no other lane collects tests/rocm/.
       - name: HIP JIT flag plumbing (no GPU, no torch)
-        run: pytest --noconftest tests/rocm_tests/test_jit_flag_hooks.py
+        run: pytest --noconftest tests/rocm/test_jit_flag_hooks.py
 
       # The other hardware-less suite: ownership classification for
       # scripts/amd_coverage.py. Mostly throwaway git repositories, but one case
       # resolves the base against this checkout and skips when it cannot, so -rs
       # keeps that from vanishing silently.
       - name: Coverage-ownership classification (no GPU, no torch)
-        run: pytest -rs --noconftest tests/rocm_tests/test_amd_coverage.py
+        run: pytest -rs --noconftest tests/rocm/test_amd_coverage.py
 
       # The repo-maintenance scripts. Git and filesystem only, so this is the
-      # only lane that runs them -- tests/rocm_tests/ is otherwise GPU-gated.
+      # only lane that runs them -- tests/rocm/ is otherwise GPU-gated.
       # The matrix suite doubles as the README staleness check for anyone who
       # committed with pre-commit disabled.
       # -rs because test_build_backend.py importorskips `build`; without it a
       # missing dependency reads as a pass rather than as an unrun wheel check.
       - name: Repo-maintenance scripts (no GPU, no torch)
         run: |
           pytest -rs --noconftest \
-            tests/rocm_tests/test_upstream_canary.py \
-            tests/rocm_tests/test_git_describe_rocm.py \
-            tests/rocm_tests/test_gen_arch_support_matrix.py \
-            tests/rocm_tests/test_build_backend.py
+            tests/rocm/test_upstream_canary.py \
+            tests/rocm/test_git_describe.py \
+            tests/rocm/test_gen_arch_support_matrix.py \
+            tests/rocm/test_build_backend.py
 
       # The unit tests once stayed green while the tool itself raised on every
       # invocation, so CI has to actually invoke it. This also fails loudly if

.gitignore

@@ -12,6 +12,7 @@ docs/generated/
 flashinfer/_build_meta.py
 flashinfer/data/
 flashinfer/include
+flashinfer/csrc/
 flashinfer/_version.py
 flashinfer/__config__.py
 flashinfer/jit/aot_config.py
@@ -197,7 +198,7 @@ cython_debug/
 .nfs*
 *.png
 # rocm_profiler generated outputs (timing CSVs, counter YAMLs, counter CSVs, roofline PNGs)
-# These live in benchmarks/rocm_benchmarks/ at runtime and must not be committed.
+# These live in benchmarks/rocm/ at runtime and must not be committed.
 *_counters.yml
 *_timing.csv
 *_counter_collection.csv

.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|attention/aiter)/.*\.(cuh|hpp|h)$
+        files: ^include/flashinfer/rocm/.*\.(cuh|hpp|h)$
         pass_filenames: true
 
   # The README's per-architecture support matrix is generated from

CLAUDE.md

@@ -35,12 +35,19 @@ stale-binary failures.
 git worktree add -b <branch-name> tmp/worktrees/<branch-name> origin/amd-integration
 ```
 
-A fresh worktree is source-only: recreate `flashinfer/include` as a **relative**
-symlink (`rm -rf flashinfer/include && ln -s ../include flashinfer/include` — the
-`rm` matters, since `ln -f` will not clear a real directory) and copy
-`flashinfer/_version.py` from the main checkout, or the JIT will not build.
-`_version.py` only exists there once that checkout has been installed or built.
-Details: `pr-workflow` skill.
+A fresh worktree is source-only. Recreate both generated trees as **relative**
+symlinks and copy `flashinfer/_version.py` from the main checkout, or the JIT
+will not build:
+
+```bash
+rm -rf flashinfer/include flashinfer/csrc     # ln -f will not clear a real dir
+ln -s ../include flashinfer/include
+mkdir -p flashinfer/csrc && ln -s ../../csrc/rocm flashinfer/csrc/rocm
+cp <main-checkout>/flashinfer/_version.py flashinfer/_version.py
+```
+
+`_version.py` only exists in the main checkout once it has been installed or
+built. Details: `pr-workflow` skill.
 
 ## Essential Commands
 
@@ -97,7 +104,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
+`csrc/rocm/` is where PyTorch tensor handling lives. Violations
 cause subtle build failures.
 
 **Test parallelism**: `pytest -n auto` automatically halves the physical GPU