Sync with upstream v0.6.18

[demandal25]

Summary

Syncs the port from upstream v0.5.3 to v0.6.18 — 1125 upstream commits and 2224 files since the merge base 2628bebcf (2025-11-20). vLLM and SGLang both pin flashinfer-python==0.6.18 exactly on main, so this is the surface their call sites are written against.

Scope is deliberately the sync itself: land the merge, restore the ROCm halves, get the port green. No new upstream features are adopted, the forked headers are not refreshed, ROCm's public API is not brought to parity with 0.6.18's reworked wrappers, and the AITER pin stays at 0.1.20. The two follow-ups that scoping implies are written down at the bottom rather than left implicit.

Read it commit by commit: ad239a313 is the merge with every shared code file taken from upstream verbatim, and each commit after it restores one subsystem.

What changed

The merge — ad239a313

git merge-tree raised 24 conflicts, 15 of them code — modest for a jump this size, and the payoff of the evacuation work in #305, #314, #318, #325 and #334. Resolution was upstream-verbatim for every shared code file, so git diff v0.6.18 over those paths is empty at that commit and the ROCm side comes back in reviewable pieces afterwards.

• flashinfer/norm.py → flashinfer/norm/ — upstream turned the module into a package, so the deletion is accepted and the dispatch is re-homed in fc730dd08. This is the one conflict that fails silently if fumbled: leaving norm.py beside the package resolves to the package and the ROCm routing just disappears, with no import error.
• Deleted — CHANGELOG.md and tests/utils/test_block_sparse_indices_to_vector_sparse_offsets.py (both gone upstream); version.txt stays deleted, the port versions from setuptools-scm.
• Ours — README.md, CONTRIBUTING.md, benchmarks/README.md, CLAUDE.md, .claude/skills/benchmark-kernel/SKILL.md, pyproject.toml. .gitignore is the union.
• pytest.ini is not a conflict — it predates the merge base and upstream has not touched it since, so the port's deletion carries through and [tool.pytest.ini_options] still governs test selection.

Restoring the ROCm side

• flashinfer/jit/{env,core}.py (365533b3e) — the two files that gate everything else. env.py is unchanged in shape; two things shrank, because upstream deleted get_nvshmem_{include,lib}_dirs outright and now does the extra_include_paths → Path conversion itself. core.py needed real work: upstream refactored JitSpec into an ABC with a JitSpecNvcc subclass and moved build_and_load() onto the base, so the port's old class_name= plumbing no longer fits — it is dropped rather than ported, since no caller has ever passed it.
• flashinfer/jit/rocm/cpp_ext.py (same commit) — emit absolute output paths. Upstream moved its own generator to absolute paths and switched build() to run ninja with -C self.build_dir instead of the JIT root; the ROCm generator still emitted $name/<obj>, which under the new workdir would have nested every artifact one level deeper.
• The four __init__ files (275f64be8) — mechanical. Upstream's export list re-indented under if IS_CUDA:, the port's elif IS_HIP: / else: arms re-attached unchanged. No symbol the port binds was renamed or dropped.
• flashinfer/{activation,page,pod,rope}.py (2d1228dca) — the backend="auto"|"native"|"aiter" parameter and its maybe_* short-circuit. pod.py absorbs a fifth CUDA-only plan() argument (uniform_q_len) as one more element of the tail the port already splits off; rope.py grew a seventh _fake_* op.
• flashinfer/norm/__init__.py (fc730dd08) — plus one addition, below.
• The benchmark harness (3b742fa08, 4b20e8011) — against a routines/attention.py upstream changed by +1991/−177.

Defects this surfaced

• import flashinfer was dead on ROCm (a5fc63be4). flashinfer/quantization/ is a package in 0.6.18 and its __init__ imports fp4/fp8 eagerly, which reach CUDA-only jit exports: ImportError: cannot import name 'sm121a_nvcc_flags'. Importing the submodule directly does not help — the package __init__ runs either way — so the gate goes where the coupling is, matching comm/__init__.py. packbits is the only member ROCm compiles.
• The wheel would have been indistinguishable from upstream's (a5fc63be4). scripts/git_describe_rocm.py took the closest ancestor tag regardless of shape; after the merge that is upstream's own v0.6.18 at distance 364 against v0.5.3+amd.2 at 1249, and a bare tag yields a version with no +amd local segment and no error to say so. Fork tags now win at any distance. Restricting the git tag -l glob to v*+amd.* instead would have broken six of the thirteen existing cases.
• Norm silently stops defaulting to the native kernel if nvidia-cutlass-dsl is ever installed (fc730dd08). Upstream routes norm through CuTe DSL by default and ROCm escaped only via an except ImportError — i.e. because the package happens to be absent. _USE_CUDA_NORM is now forced under IS_HIP, so the default the port chose in fix(rocm): default the RMSNorm ops to the native kernel, with the benchmark that decided it #332 does not depend on what is in site-packages.
• The benchmark harness aborted before running anything (4b20e8011). Three module-level imports added in 0.6.18 fail on ROCm, so parse_args died on ImportError: cannot import name 'autotune' from 'flashinfer' — which made every ROCm hook restored in 3b742fa08 unreachable. All three uses are CUDA-only paths and move to their call sites.
• Every sampling call raised at dispatch (fe3c4b4c0). flashinfer/sampling.py is shared verbatim and 0.6.18 widened 10 of its 11 entry points, so the first call died with sampling::sampling_from_probs() expected at most 6 argument(s) but received 9. Three unrelated widenings landed at once: the scalar philox pair became optional per-request seed/offset tensors (7 ops), the five *_from_probs ops gained a valid output, and the renorm ops gained multi-CTA scratch plus an is_deterministic switch for AIR top-p. None maps onto a ROCm kernel — a per-request seed is rejected rather than silently collapsed to seed_val, valid is uniformly true because ROCm's kernels carry the last-valid-index fallback but not upstream's reject flag, and the scratch is unused because the single-CTA ternary-search kernels are deterministic already. top_k_top_p_sampling_from_{logits,probs} also grew a radix top-k fast path that builds csrc/topk.cu, which csrc/rocm does not have, so it failed inside ninja rather than at import; gated off on HIP. tests/rocm/test_sampling_binding_abi.py compares the two sides statically — no GPU, and it would have caught all three widenings at merge time.
• tests/attention/test_page.py stopped collecting, so a bare pytest died before running anything (d023ac824). 0.6.18 gave it a module-level nvfp4 import. It is dropped from testpaths — its one HIP-relevant case is covered further by tests/rocm/test_append_paged_kv_cache_aiter.py, and a native-only 5-D case was added there because every existing one is @requires_aiter. Only a full-testpath run sees this; a tests/rocm/-only run cannot.
• The CUDA-only import gate was never installed on ROCm (a8523a350). gate_cuda_only_modules() had one call site, flashinfer/comm/__init__.py, and import flashinfer does not reach flashinfer.comm on HIP — so the gate was up only for callers that imported comm first. Every existing test hid it by importing comm itself. Called from flashinfer/rocm/api.py now.
• The release artifacts carried NVIDIA-licensed sources (7bded9b66). csrc/fmha_v2/ and flashinfer/jit/attention/fmha_v2/ are new in 0.6.18 and 61 of their files carry the NVIDIA TensorRT Source Code License, which forbids redistribution. The sdist shipped 77 of them and the wheel four Python modules, including the licensed generator_utils.py. The wheel case needed exclude-package-data, not the packages.find exclude: the directory has no __init__.py, so the files arrive as package data — include-package-data defaults true for a pyproject project and setuptools-scm's file finder then sweeps every tracked file under flashinfer/ in.
• pre-commit CI was red on the new Validate CUDA configuration step (7bded9b66). Not noise: with pyproject.toml reshaped to suit its hand-rolled TOML parser it then demands nvidia-cutlass-dsl[cu13]>=4.6.2a0 in a cu13 extra, which a ROCm wheel cannot declare. The step is dropped.
• Two tests were asserting against the wrong file (4b20e8011). test_hip_gqa_group_sizes_match_the_kernel_dispatch read upstream's include/flashinfer/utils.cuh, which gained a group_size == 6 arm; the HIP decode kernel compiles against the forked rocm/dispatch.cuh, still {1,2,3,4,8}, so the constant was right and the test was reading the wrong header — it has pointed at the upstream original since the headers were forked, and the merge only exposed it. test_nvshmem_helpers_stay_absent_on_rocm named gen_nvshmem_module(), which 0.6.18 deletes.

What v0.6.18 adds without conflicting

A merge only reports files both sides changed. 1767 files arrive that the port did not have, and several change behaviour with nothing to announce them (b9cbb433e, 07ec7c8f1):

• Ten upstream workflows. Seven fire on this fork's own activity and cannot work here — pr-labeler.yml (pull_request_target, no branch filter, needs FLASHINFER_BOT_TOKEN) would be red on every PR and issue-claim.yml (cron */30) roughly 48 times a day. They are dropped. pr-test.yml, pr-api-doc-checks.yml and pr-api-doc-comment.yml are branches: [main] and stay dormant, so they are left alone — deleting an upstream file buys a permanent conflict on every future sync.
• pre-commit run -a wanted to rewrite 24 upstream files. The port kept its own pyproject.toml through the merge and so never inherited the [tool.ruff] extend-exclude block 0.6.18 added for generated and vendored kernels; without it, 16 upstream files reformat. markdownlint is fork-added — upstream does not run it — so aiming it at the whole tree rewrote docs/design_docs/, .claude/skills/add-cuda-kernel and flashinfer/moe_ep/**/*.md. Both fixes are in the config, not the files.
• flashinfer.comm.mixed_comm joins CUDA_ONLY_MODULES: it is where the NVSHMEM build landed, and it renders .cu templates into FLASHINFER_GEN_SRC_DIR before reaching import nvidia.nvshmem, so without the gate a ROCm caller writes CUDA sources on its way to the failure.

Architecture / design notes

Merge, not rebase. The two prior upgrades (#156 for v0.3.1, #173 for v0.5.3) were both git merge <upstream tag> into a topic branch, and the same reasoning holds: a rebase replays 353 fork commits over 1125 upstream commits and re-resolves the same files at every step, where a merge resolves each once. The catch-up with 128f9dfd5 mid-branch is a merge for the same reason — rebasing a branch that carries the v0.6.18 merge commit would re-raise all 24 conflicts against the same tag for no gain, which is why the v0.5.3 upgrade did the same thing in 61c11fdb8.

Both architectures. The merge changes zero files under include/flashinfer/rocm, csrc/rocm, flashinfer/rocm, tests/rocm, benchmarks/rocm or docs/rocm, and the ROCm jinja templates include only flashinfer/rocm/* headers, so no upstream header reaches the HIP build. The 196-file include/ replacement is inert for gfx942 and gfx950 alike and the risk is concentrated in Python — which is why the arch-specific treatment below is thin rather than absent.

Coverage needs no code change, but does need the tag. scripts/amd_coverage.py derives its base from the fork's own *+amd.* tag, so it self-corrects the moment v0.6.18+amd.1 is cut. Until then it derives v0.5.3, which is not an ancestor of the merge, and reports 827 owned .py files instead of 39 — green, because the arch-caps job only asserts owned is non-empty, and meaningless. Cutting the tag right after this merges is the fix; papering over a one-PR window in the tool would be worse. Verified rather than assumed, in a throwaway clone at f58311959: tagging v0.6.18+amd.1 moves the version from v0.5.3+amd.2.dev1266 to v0.6.18+amd.1, the derived base to 69ff11fc4 ("bump version to 0.6.18"), and the owned set from 827 files to 56.

A user holding a 0.5.3-line amd-flashinfer-jit-cache wheel will hit a RuntimeError from flashinfer/jit/rocm/env.py::get_aot_dir(), which checks cache_version.startswith(flashinfer_version). The jit-cache wheel has to be rebuilt alongside the tag.

Test plan

• import flashinfer and import flashinfer.comm on ROCm with no GPU visible — the cheapest signal, and the one that caught the quantization package defect.
• pre-commit run -a — clean, without reformatting a single upstream file.
• pytest -rs --noconftest tests/rocm/test_{upstream_canary,git_describe,gen_arch_support_matrix,build_backend}.py — 142 passed, no GPU.
• pytest tests/rocm/test_{benchmark_harness,jit_env,comm_import_gate}.py on gfx942 — 51 passed.
• A/B on the version fix: reverted it, confirmed the new case picks the bare v0.6.18 tag, restored.
• pytest -n 4 --reruns 2 -m "not slow" tests/rocm/ on gfx942 (MI300X) and gfx950 (MI350X), both from a detached pin-run worktree at acf125ecd — zero failures on both, exit 0, --check confirming neither tree moved mid-run. An earlier pass at fe3c4b4c0 returned an identical 11 failures on both arches, all in two host-only files, fixed in f1d69e483. Only 7bded9b66 (packaging metadata and a CI workflow) landed after the suite SHA; the two test files that parse pyproject.toml were re-run against it.
• A/B on each new test that asserts a behaviour this PR adds: the binding-ABI check (dropped valid from one declaration), the fp16/bf16 cast-back (dropped the copy_), valid (dropped mark_all_valid), and the import gate (insert → append). Every one fails with the fix reverted — the gate test did not, first time round, which is why it now gates a module that exists.
• Wheel and sdist built and their full inventories diffed: 986 -> 982 and 2101 -> 2027 entries, the delta being exactly the NVIDIA-licensed fmha_v2 files and nothing else. A scan of both for the licence string comes back empty, and the installed wheel imports with the module absent.
• scripts/upstream_canary.py on both refs. Against v0.6.18: clean merge -- nothing to do, and no upstream changes to any of the 25 forked headers — the port is exactly in sync with the tag it targets. Against upstream/main: 38 conflicted, 32 of them code, but see below before reading that as fork delta.

Follow-ups

• v0.6.18 is not on upstream/main. It is cut from a release branch that diverged at 61a6c6518 (2026-08-19), so git merge-base v0.6.18 upstream/main is that commit and not the tag. The canary's 38 conflicts against main therefore attribute the whole v0.6.18 release delta to "ours" — fused_moe/runners.py, cute_dsl/tuner.py and the moe tests dominate the list, and none of them is fork-edited. The honest number for this PR is the v0.6.18 one: clean. It does mean the next sync should pick a tag or a main commit deliberately rather than assume tags are ancestors.
• Forked-header drift is zero right now, and that is the moment it starts accruing. The 25 forked headers match their v0.6.18 counterparts, so the canary reports nothing; against main seven already show churn (attention/prefill.cuh +102/−42, norm.cuh +137/−0, and five smaller). They merge cleanly forever and go stale silently, so the drift report is the only signal.
• ROCm public-API parity with 0.6.18. Upstream reworked prefill.py (+3019/−195), decode.py (+1514/−437) and turned mla.py into a package. flashinfer/rocm/{decode,prefill,mla}.py are standalone and unaffected by the merge, but a vLLM or SGLang caller on the 0.6.18 pin will find gaps. block_sparse_indices_to_vector_sparse_offsets is one concrete instance: upstream removed the Python op, and only csrc/rocm/ still implements it, with nothing calling it.

[Copilot]

Pull request overview

Syncs the AMD ROCm port of FlashInfer to upstream v0.6.18, bringing the repo’s Python/C++ surface area and test suite into alignment with the version pinned by vLLM/SGLang, while re-applying ROCm-specific routing, packaging, and CI adjustments needed to keep the port functional.

Changes:

• Updates/extends multiple CUDA and ROCm-facing modules (JIT, comm, attention/norm/sampling/topk, versioning) to match upstream v0.6.18 structure and APIs.
• Adds/updates test coverage across trace/reference correctness, sharding infrastructure, ROCm utilities, and various CUDA-only paths.
• Refreshes CI/devcontainer and repository automation scripts/config to match the new upstream layout and workflows.

Reviewed changes

Copilot reviewed 65 out of 2221 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
flashinfer/version.py	Updates version metadata loading and backward-compat aliasing for git commit/version fields.
tests/utils/test_jit_warmup.py	Adds a CUDA-arch-specific warmup test for SM90-only fa3 backend.
csrc/bgmv_moe/moe_bgmv_ops.h	Introduces a public header for BGMV MoE TVM-FFI entry points.
tests/comm/test_nvshmem.py	Adjusts NVSHMEM smoke checks and import paths for basic API access.
tests/rocm/test_git_describe.py	Adds regression coverage for fork-tag precedence over nearer upstream tags.
tests/rocm/test_benchmark_harness.py	Fixes ROCm header reference for HIP group-size dispatch assertions.
flashinfer/rocm/api.py	Moves ROCm packbits exports to quantization.packbits submodule to avoid eager CUDA-only imports.
flashinfer/rocm/prefill.py	Updates ROCm imports to avoid quantization package side effects and keep ROCm importable.
.github/workflows/pre-commit.yml	Switches workflow Python selection to .python-version and adds CUDA config validation step.
ci/docker-tags.yml	Updates CI image tags to the newer CUDA CI baselines.
docs/api/*.rst	Adds/expands API docs for new/reshaped modules (topk, pod, mhc, mamba, kda, etc.).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

Review details

• Files reviewed: 65/2232 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Some newly-modified tests can raise at collection time on CPU-only/no-GPU environments due to unguarded torch.cuda.get_device_capability() calls.

Review details

• Files reviewed: 66/2234 changed files
• Comments generated: 0 new
• Review effort level: Lite