Repository navigation
Bring release/0.6.18+amd.1 up to the v0.6.18+amd.1 tag - #396
Merged
demandal25 merged 3 commits intoOct 21, 2026
Merged
Conversation
…#358) ## Summary `CUDA_ONLY_MODULES` exists so that importing an unported upstream module on ROCm raises one catchable `ImportError` naming the module, rather than an obscure `cannot import name ...` from inside the JIT. An import sweep of all 69 top-level `flashinfer.*` modules found eight that reach ROCm through a different error, and five that import cleanly but can never run. Neither group was gated or documented, so a caller could not tell "unsupported" from "you hit a bug". This gates the eight, documents the rest, and adds the static guard that would have caught them during the v0.6.18 sync. Sweep environment: `docker/Dockerfile.rocm` (torch 2.12.0+rocm10.0, gfx942/MI325X), editable install of `0.6.18+amd.1`. ### What changed - **`flashinfer/rocm/__init__.py`** — eight entries added to `CUDA_ONLY_MODULES`, each with the error it replaces: `deep_gemm` and `green_ctx` (both reach `cuda_utils.py:31`, which re-raises without `cuda-python`), `gdn_prefill` and `mamba.ssd_combined` (both `No module named 'cutlass'` — CuTe DSL imported at module scope), `parallel_attention` (`cannot import name 'fmha_varlen'`), `aot` and `__main__` (`jit/cpp_ext.py` shells out to `nvcc` with no HIP branch), `tactics_blocklist_gen` (`cannot import name 'ActivationType'`). - **`tests/rocm/test_kernel_source_coverage.py`** — new. Every `gen_*` import an op makes must resolve to a source under `csrc/rocm`, or the module must be gated, shadowed, or in a named `_UNPORTED` allowlist. - **`.github/workflows/arch-caps-conformance.yml`** — runs the new test in the hardware-less lane. - **`docs/rocm/backends.md`** — the "Not available on ROCm" table grows, and gains two subsections: "Imports, but has no ROCm kernel" and "Unverified". ### Architecture / design notes **`aot` is gated, not shadowed.** The tempting fix is to alias `flashinfer.aot` to `flashinfer.rocm.aot` via `SHADOW_MODULES`. It does not work and would not be correct if it did. A `sys.modules` alias cannot serve `python -m flashinfer.aot`: `runpy._get_module_details` asks the aliased spec's loader for `flashinfer.aot`, and `SourceFileLoader._check_name_wrapper` raises `loader for flashinfer.rocm.aot cannot handle flashinfer.aot`. And the twin is not a drop-in — `gen_all_modules` takes 6 parameters against upstream's 14, `gen_attention` 4 against 10, `copy_built_kernels` adds a required positional `build_dir`, and `gen_fa3` / `gen_xqa` / `detect_sm_capabilities` are absent — so the alias would promise upstream's API and deliver a narrower one, which is what `test_api_parity.py` exists to prevent. `flashinfer.rocm.aot` is already the documented entry point (`CONTRIBUTING.md`, `amd-flashinfer-jit-cache/build_backend.py`); the gate just makes the wrong door say so. `scripts/verify_all_modules_compiled.py` is the only in-tree importer of `flashinfer.aot` and runs solely from the CUDA `pr-test.yml` matrix, where the gate is not installed. **Five unported modules are deliberately left importable.** `topk`, `concat_ops`, `mhc`, `xqa` and `nvfp4_attention_sm120` import cleanly and die in ninja on first call, because `csrc/rocm` ships no matching source. Gating them would break working code: `topk_varlen/topk_varlen.py:45` imports `topk` at module scope, and `tests/trace/template_registry.py` registers four of the five. (`sampling.py:1385` imports `topk` too, but behind a `not IS_HIP` guard, so it is not a reason.) The gate's message is "CUDA-only"; these are "not ported yet", and conflating the two costs more than it documents. The ninja error already names the missing file. **The guard is static, not a runtime sweep.** The obvious test — import every module and assert each either works or raises the gate error — cannot run where it is needed. `arch-caps-conformance.yml` is the only lane that sees a PR, and it installs no torch. A `pkgutil` sweep would also miss 11 of the 20 existing gate entries, which are submodules (`comm.*`, `quantization.fp4_quantization`/`fp8_quantization`) — including the v0.6.18 change this is meant to catch. Parsing `FLASHINFER_CSRC_DIR / "..."` out of the JIT specs needs neither torch nor a GPU, and catches the cause rather than the symptom. The test also asserts `scripts/rocm_api_parity.py`'s `PAIRS` covers every `SHADOW_MODULES` entry — nothing linked the two before, so a future shadow would enrol in no signature audit. **`cuda_utils` stays ungated** — it is private, and the two modules that reach it are now gated themselves. ### Known gaps left alone The allowlist surfaced four latent breaks, all lazily built and reached by no in-tree caller: `utils.set_log_level()` needs `logging.cc`, `api_logging`'s opt-in GPU stats counter needs `api_log_stats.cu`, `norm`'s fused rmsnorm+silu variant needs `rmsnorm_silu.cu`, and `mamba.selective_state_update` needs the SSM kernels — none present under `csrc/rocm`. `norm` itself is a supported op; only that one entry point is unbacked. Separately, `kda`/`kda_decode`/`kda_kernels` (`tvm_ffi`), `profiler` (`tg4perfetto`) and `artifacts` (`requests`) fail on optional dependencies the dev image does not install, so their ROCm status is still unknown — documented as unverified rather than guessed at. ## Test plan - [x] Import sweep over all 69 top-level modules, gfx942/MI325X: all eight report the uniform gate error; the only remaining non-`OK` rows are `cuda_utils` (intentional) and the optional-dependency misses. - [x] `topk_varlen.topk_varlen`, `sampling`, and the five ungated modules still import. - [x] `pytest tests/rocm/{test_kernel_source_coverage,test_comm_import_gate,test_api_parity,test_aot,test_rocm_dispatch}.py` — 138 passed, 2 skipped. - [x] `python -m flashinfer.aot --help` exits with `flashinfer.aot is CUDA-only and not available on ROCm`, pinned by a subprocess test — runpy asks the loader for code and never reaches `exec_module`, so `_CudaOnlyLoader` needs `get_code`. - [x] New test passes under `python:3.12-slim` with no torch and `--noconftest`, which is the lane it has to survive. - [x] `pre-commit run -a` Not run: the GPU suite, unchanged by this PR — no kernel, wrapper or routing code is touched. The sweep ran on gfx942 only; the gate is arch-invariant (`flashinfer/rocm/api.py:30` passes a constant frozenset), so gfx950 adds nothing here. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
## Summary Three follow-ups to #358, all documentation plus one comment fix. The support matrix listed what works and nothing else, so a reader arriving from upstream could not tell whether an op was missing because ROCm lacks it or because the table only enumerates the supported set. #358 established the three classes in `docs/rocm/backends.md`; this surfaces them in the README. Along the way the README was condensed, and the pr-workflow quality gate gained a documentation step so the next change of this kind is caught before the push rather than after. ## What changed - **`README.md`** — new "What upstream has that this does not" subsection under Support matrix: gated CUDA-only, imports but has no ROCm kernel, unverified, with a few names each and a pointer to `backends.md` for the full lists. It sits outside the generated block, so `gen_arch_support_matrix.py` and its staleness test are unaffected. - **`README.md`, `flashinfer/rocm/torch_compile.py`** — "Dynamo" now reads "TorchDynamo". - **`README.md`, `docs/rocm/backends.md`, `flashinfer/rocm/arch_caps.py`** — condensation, including two paragraphs relocated rather than shrunk. - **`.claude/skills/pr-workflow/SKILL.md`** — a documentation step in the pre-PR quality gate. ## Notes **Why "TorchDynamo".** The original statement was correct — Dynamo here is the `torch.compile` tracing frontend, a PyTorch component that works on ROCm and is covered by `tests/rocm/test_torch_compile.py`. But the name collides with NVIDIA's Dynamo inference server, which is a bad ambiguity for a ROCm port to leave in place. **Condensation: 2908 → 2720 words, while adding the new section**, so the inherited prose is about 15% shorter. Nothing was dropped. Two blocks moved to where the subject already lives: the soft-capped-prefill analysis is now in `backends.md` under Per-op notes alongside the other per-op constraints, and the `-n auto` mechanics were already spelled out in `CONTRIBUTING.md`, so the README keeps the rule and cites it. The `cascade` note was condensed at its source in `arch_caps.py` and the table regenerated — the table itself is generated and was not hand-edited. **The new gate step** lists which file owns which subject (README, `docs/rocm/backends.md`, CONTRIBUTING.md, `CLAUDE.md` / `.claude/skills/`, the per-package READMEs), and carries the two traps worth stating once: the support matrix is generated, so edit `arch_caps.py` and regenerate; and a fact stated in two files will drift, so cite the owner rather than restating it. ## Test plan - [x] `pre-commit run -a` - [x] `pytest tests/rocm/test_gen_arch_support_matrix.py tests/rocm/test_arch_caps.py` — the regenerated table matches `arch_caps.py`. - [x] Every `backends.md` anchor the README links to resolves, and every module named in the new subsection appears in `backends.md`. No GPU tests: no kernel, wrapper or routing code changes. The only non-documentation edits are a comment in `flashinfer/rocm/torch_compile.py` and a `note=` string in `arch_caps.py`, neither of which affects behaviour. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
) ## Summary The "What upstream has that this does not" subsection added in #359 classified by two different things at once: the first row was a cause ("CUDA-only") welded to a mechanism ("gated"), while the other two named only how the failure surfaced. Anything NVIDIA-only that happened to import cleanly therefore landed in the wrong row, and "Unverified" became a bucket for whatever had not been checked. This checks it, and splits the section by cause with the failure mode as its own column. ## What changed - **`README.md`** — two groups, NVIDIA-only and not ported yet, plus a sentence for the two modules that are simply uninstallable here. - **`docs/rocm/backends.md`** — matching sections: "Not ported yet", "NVIDIA-only, but not gated", "Blocked on a missing package, not on ROCm", with the evidence per entry and why each ungated NVIDIA-only module is not gated. ## Architecture / design notes **What checking moved.** Five modules were called unverified and are not: | Module | Evidence | | :--- | :--- | | `msa_ops`, `moe_ep` | import `cutlass.cute`; `moe_ep` also NVSHMEM | | `gdn_decode` | every kernel under `gdn_kernels/` is `@cute.jit` | | `cute_dsl`, `cutile` | `cute_dsl/__init__.py:37` hides the whole surface behind `is_cute_dsl_available()`, so on ROCm it imports as an empty shell | None of those "may well work", which is what the old row claimed. Going the other way, `diffusion_ops` is not unverified either — it re-exports four fused DiT entry points from `flashinfer.norm`, and `csrc/rocm/norm.cu` defines none of them, so it is unported. `nvfp4_attention_sm120` (`supported_compute_capability([120, 121])`) and `tllm_utils` (five `nv_internal` TensorRT-LLM sources) move up to NVIDIA-only. **Why the third group is gone rather than renamed.** After those moves, what remained was `profiler` and `artifacts`, which fail to import for want of `tg4perfetto` and `requests`. That is a statement about `docker/Dockerfile.rocm`, not about ROCm support, so it reads as one sentence rather than as a support category. **Cause and mechanism are now separate.** Being NVIDIA-only does not imply being gated: `autotuner` imports `tllm_utils` at module scope, so gating it would break a working module, and `gdn_decode` fails late where `gdn_prefill` fails at import. The table says which, per row. ## Test plan - [x] `pre-commit run -a` - [x] `pytest tests/rocm/test_kernel_source_coverage.py tests/rocm/test_gen_arch_support_matrix.py` - [x] All 26 modules the README names appear in `backends.md`; every anchor the README and the `backends.md` TOC use resolves. Documentation only — no code changes, so no GPU tests. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
🟢 Approval recommended
Only a non-blocking documentation clarification remains in this release-tag alignment.
1 open issue
What changed in this PR
Aligns the release branch with the shipped v0.6.18+amd.1 tag, restoring tag reachability for version conformance checks.
Changes:
- Adds uniform ROCm import gates and static kernel-source checks.
- Clarifies unsupported modules and condenses documentation.
- Requires a merge commit to preserve the tagged history.
| File | Description |
|---|---|
| tests/rocm/test_kernel_source_coverage.py | Adds static source and classification checks. |
| tests/rocm/test_comm_import_gate.py | Tests gated command-line module execution. |
| README.md | Condenses guidance and explains support gaps. |
| flashinfer/rocm/torch_compile.py | Clarifies TorchDynamo naming in a docstring. |
| flashinfer/rocm/arch_caps.py | Condenses the cascade support note. |
| flashinfer/rocm/__init__.py | Extends import gates and loader errors. |
| docs/rocm/backends.md | Documents unsupported modules and prefill constraints. |
| .github/workflows/arch-caps-conformance.yml | Runs the source-coverage guard. |
| .claude/skills/pr-workflow/SKILL.md | Adds documentation checks and formatting guidance. |
Review effort: Balanced
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
release/0.6.18+amd.1sits at3578e732c(#357), three commits behind thev0.6.18+amd.1tag (7e2913df7): #358, #359, #360. From the branch,git describe --match '*+amd.*'resolvesv0.5.3+amd.2, sotest_the_ref_agrees_with_the_forks_own_release_tagfails on every PR into it, including #389.This brings the branch up to the shipped tag. It adds no new commits; merge with Create a merge commit (the only method the branch allows) so the tag commit stays reachable.
Test plan
arch-caps-conformance.ymlfile list, CI's dependencies) on tag + [26.10] [DOC] FlashInfer: documentation updates for release #389: 189 passed; on the current branch + [26.10] [DOC] FlashInfer: documentation updates for release #389: 1 failed, 188 passed (the assertion above).