Skip to content

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

Merged
demandal25 merged 13 commits into
amd-integrationfrom
rocm-evacuate-per-op
Aug 29, 2026
Merged

refactor(rocm): remove the gpu_iface layer and flatten the ROCm headers#329
demandal25 merged 13 commits into
amd-integrationfrom
rocm-evacuate-per-op

Conversation

@demandal25

@demandal25 demandal25 commented Aug 29, 2026

Copy link
Copy Markdown
Collaborator

Summary

Final PR of the upstream-conflict-surface plan. Evacuates the last large ROCm-only block out of an upstream file, moves the ROCm C++ sources under flashinfer/csrc/rocm/, then removes the gpu_iface layer entirely — flattening 19 headers into include/flashinfer/rocm/ and dropping the namespace. Along the way it fixes two live defects the investigation turned up.

Ten commits, each self-contained and reviewable on its own.

Two defects fixed

A macro redefinition. gpu_iface/dispatch.cuh and attention/dispatch.cuh both defined DISPATCH_NUM_MMA_KV with different bodies, both used #pragma once, and three headers reached both:

include/flashinfer/rocm/attention/dispatch.cuh:31:9: warning: 'DISPATCH_NUM_MMA_KV' macro redefined [-Wmacro-redefined]

The build does not pass -Wmacro-redefined, so nothing reported it. Merging is behaviour-preserving, not a judgement call: all seven expansion sites are in prefill.cuh, pod.cuh and batch_pod.cuh, and each includes cascade.cuh (which pulls the other copy) immediately before the sibling "dispatch.cuh" — so the fallback body already won everywhere the macro is actually expanded.

A silent header shadow. rocm/sampling.cuh included upstream <flashinfer/allocator.h>, which pulls upstream exception.h — but our exception.h had already claimed the shared guard FLASHINFER_EXCEPTION_H_. Preprocessing showed our copy expanding and upstream's producing nothing; grep -c 'class Warning|FLASHINFER_WARN' on the output returned 0. Upstream's allocator was compiling against our exception header in sampling.cu and renorm.cu, and worked only because it happens not to use the suppressed symbols. Pointing it at the fork's allocator removes the last path from a ROCm TU to any upstream header.

What changed

  • flashinfer/rocm/torch_compile.py — new; the opt-in torch.library registration that was ~90 lines inside flashinfer/utils.py (114 → 49). utils.py keeps upstream's version split and both signatures, delegating only the bodies, because a plain re-import would enlarge the diff against upstream.
  • flashinfer/csrc_rocm/flashinfer/csrc/rocm/ — 65 files. Deliberately package-internal rather than repo-root csrc/rocm/: upstream stages root trees into flashinfer/data/ via machinery this fork does not run, upstream's csrc/ has no backend-named subdirectories, and it is a tree upstream actively grows.
  • include/gpu_iface/include/flashinfer/rocm/ — 19 headers flattened up two levels, ~140 include paths rewritten. _hip suffixes stay: with wrapper and backend in one directory, that suffix is what distinguishes them.
  • The gpu_iface namespace is gone. Levels survive only where upstream also groups (flashinfer::math, flashinfer::memory). Four headers turned out to be pure pass-throughs once the namespace went and are deleted. flashinfer::mma_hip is the one rename — three of its signatures match upstream's flashinfer::mma exactly.
  • 22 fork headers gained their own include guards plus an #error naming the upstream file if its guard is already set, so the shadow class of bug above fails loudly instead of silently.

Deliberate calls

  • include/flashinfer/rocm/utils.cuh is excluded from the canary's drift pairing. It shares upstream's basename but 411 of 443 lines differ — a HIP rewrite that never tracked upstream. Everything else pairs by default, which newly puts layout.cuh and fastdiv.cuh into the drift report.
  • backend/hip/mma_debug_utils_hip.h (411 lines of MMA fragment-layout debug printers) is deleted, along with error.hpp. Both had zero includers. The debug printers are recoverable from git history if the next fragment-layout bug wants them.
  • The anonymous namespace in vec_dtypes_hip.h shadows HIP's __float2bfloat162_rn, and the flatten widened that shadow from gpu_iface::vec_dtypes to flashinfer. Its callers moved with it so in-file resolution is unchanged, and no other header calls either helper.
  • docs/rocm/coverage-gfx942.json is untouched. It is a snapshot taken at 0e68a2feb, already 340 commits behind, and is regenerated by a gfx942 coverage run rather than edited.

Test plan

gfx950 (MI350X), ROCm 7.2, torch 2.9.1+rocm7.2, aiter 0.1.10, in a container against this worktree.

  • Negative control first: #error poisoned into a header the tests must reach — the suite fails. Without that, every green below is meaningless.
  • All 42 fork headers compile -fsyntax-only standalone, zero errors and zero macro redefinitions, for both --offload-arch=gfx942 and gfx950
  • Preprocessor proof that no fork header's closure reaches any upstream include/flashinfer/*.h
  • Tripwire fires with the exact pair named: include/flashinfer/exception.h and include/flashinfer/rocm/exception.h both define FLASHINFER_EXCEPTION_H_
  • -Wmacro-redefined on the dispatch pair: 1 → 0
  • Wheel inventory diffed in full: 437 → 431 entries, headers 222 → 215 (−2 dead, −2 deduped, −4 folded shims, +1 memory_types.hpp), csrc 65 = 65, nothing lost
  • Merge-conflict surface against v0.6.18rc3: 26 → 25
  • Canary tracked forked headers 23 → 25; pytest tests/rocm_tests/test_upstream_canary.py green
  • pre-commit run -a

gfx942 is compile-verified only. Both architectures are codegen'd and compiled in every build, but the CDNA3 slurm node has no HIP dev headers (/opt/rocm/include/hip/hip_runtime.h absent), no flashinfer image, and a full root filesystem, so nothing could execute there. Stated rather than glossed.

Clear your JIT cache before first userm -rf ~/.cache/flashinfer/. JitSpec.build() only writes build.ninja when it is missing, so a warm cache keeps referencing deleted header paths.

test_mla_and_utils_guards_hip.py::TestRequireAiterMla::test_a_missing_package_names_the_install_command fails on this branch and on amd-integration: #294 dropped the github.com/ROCm/aiter URL from the message and #324's test still asserts it. This branch touches neither file.

Copilot AI lite review requested due to automatic review settings August 29, 2026 05:24

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Refactors the ROCm fork surface by moving ROCm-only torch.compile custom-op registration out of flashinfer/utils.py, relocating the ROCm PyTorch binding sources under flashinfer/csrc/rocm/, and moving the ROCm-only gpu_iface header layer under include/flashinfer/rocm/gpu_iface/ (plus updating references, packaging globs, and canary/coverage tooling to match).

Changes:

  • Extract torch.library custom-op registration/guards into flashinfer/rocm/torch_compile.py and delegate from flashinfer/utils.py.
  • Rename/move ROCm binding sources to flashinfer/csrc/rocm/ and update build/package/test references (pyproject.toml, MANIFEST.in, tests, scripts).
  • Move gpu_iface under include/flashinfer/rocm/gpu_iface/, update includes, and adjust scripts/upstream_canary.py pairing logic + tests to avoid basename collisions.

Reviewed changes

Copilot reviewed 67 out of 134 changed files in this pull request and generated no comments.

Show a summary per file
File Description
tests/rocm_tests/test_upstream_canary.py Adds coverage for excluding gpu_iface from upstream-canary fork pairing (with layout.cuh exempted back in).
tests/rocm_tests/test_torch_compile_hip.py Docstring wording update to match refactored compile-flag handling.
tests/rocm_tests/test_jit_flag_hooks.py Updates stubbed FLASHINFER_CSRC_DIR path to csrc/rocm.
tests/rocm_tests/test_jit_env_hip.py Updates doc/comment references from csrc_rocm to csrc/rocm.
tests/rocm_tests/test_fused_moe_aiter_hip.py Updates path lookup for shim source under flashinfer/csrc/rocm.
tests/rocm_tests/test_build_backend.py Updates wheel-contents assertions for flashinfer/csrc/rocm/.
tests/rocm_tests/test_amd_coverage.py Updates test fixtures to create flashinfer/csrc/rocm paths.
tests/jit_reach_plugin.py Updates docstrings/comments from csrc_rocm to csrc/rocm.
scripts/upstream_canary.py Excludes include/flashinfer/rocm/gpu_iface/ from basename pairing (with layout.cuh exempt).
scripts/coverage_ownership.toml Updates redirect-owned reasons to reference flashinfer/csrc/rocm paths.
scripts/amd_coverage.py Updates csrc reach path label and _CSRC_DIR to flashinfer/csrc/rocm.
pyproject.toml Updates setuptools package-data globs to csrc/rocm/**.
MANIFEST.in Updates sdist graft path to flashinfer/csrc/rocm.
include/flashinfer/rocm/sampling.cuh Updates include paths to new flashinfer/rocm/gpu_iface/* location.
include/flashinfer/rocm/quantization.cuh Updates include paths to new flashinfer/rocm/gpu_iface/* location.
include/flashinfer/rocm/gpu_iface/vec_dtypes.hpp New/moved gpu_iface header under include/flashinfer/rocm/gpu_iface/.
include/flashinfer/rocm/gpu_iface/utils.cuh Updates comments/paths for gpu_iface include relocation.
include/flashinfer/rocm/gpu_iface/sm_id.hpp New/moved gpu_iface header under include/flashinfer/rocm/gpu_iface/.
include/flashinfer/rocm/gpu_iface/platform.hpp New/moved gpu_iface header under include/flashinfer/rocm/gpu_iface/.
include/flashinfer/rocm/gpu_iface/mma_types.hpp New/moved gpu_iface header under include/flashinfer/rocm/gpu_iface/.
include/flashinfer/rocm/gpu_iface/mma_ops.hpp Updates internal includes to new flashinfer/rocm/gpu_iface/* paths.
include/flashinfer/rocm/gpu_iface/memory_ops.hpp New/moved gpu_iface header; wraps backend HIP impl.
include/flashinfer/rocm/gpu_iface/math_ops.hpp New/moved gpu_iface header; re-exports math intrinsics.
include/flashinfer/rocm/gpu_iface/macros.hpp New/moved gpu_iface header; HIP-only platform gating/macros.
include/flashinfer/rocm/gpu_iface/layout.cuh New/moved fork-tracking header (exempted from canary exclusion).
include/flashinfer/rocm/gpu_iface/gpu_runtime_compat.hpp New/moved gpu_iface runtime compat layer (HIP mappings + helpers).
include/flashinfer/rocm/gpu_iface/fastdiv.cuh New/moved gpu_iface header providing fast division helper.
include/flashinfer/rocm/gpu_iface/exception.h New/moved gpu_iface exception/check utilities.
include/flashinfer/rocm/gpu_iface/error.hpp New/moved gpu_iface error wrapper.
include/flashinfer/rocm/gpu_iface/enums.hpp New/moved enums used by dispatch paths.
include/flashinfer/rocm/gpu_iface/dispatch.cuh Updates include to new gpu_iface exception header path.
include/flashinfer/rocm/gpu_iface/cooperative_groups.h New/moved gpu_iface cooperative-groups shim.
include/flashinfer/rocm/gpu_iface/conversion_utils.h New/moved conversion helper header.
include/flashinfer/rocm/gpu_iface/backend/hip/mma_hip.h Updates include paths to relocated gpu_iface headers.
include/flashinfer/rocm/gpu_iface/backend/hip/mma_debug_utils_hip.h Updates include paths to relocated gpu_iface headers.
include/flashinfer/rocm/gpu_iface/backend/hip/memory_ops_hip.h New/moved HIP backend for memory ops.
include/flashinfer/rocm/gpu_iface/backend/hip/math_hip.h New/moved HIP backend math intrinsics implementation.
include/flashinfer/rocm/attention/variants.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/variant_helper.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/state.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/scheduler.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/prefill.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/pos_enc.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/pod.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/permuted_smem.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/page.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/norm.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/frag_layout_swizzle.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/dispatch.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/default_prefill_params.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/default_decode_params.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/decode.cuh Updates gpu_iface includes and fixes in-comment path reference.
include/flashinfer/rocm/attention/cascade.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/batch_pod.cuh Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
include/flashinfer/rocm/attention/activation.cuh Updates comment reference to moved activation source path.
flashinfer/utils.py Delegates torch.custom-op registration to flashinfer/rocm/torch_compile.py and re-exports the flag helper.
flashinfer/rocm/torch_compile.py New module holding compile guard + torch.library registration helpers.
flashinfer/rocm/api.py Imports use_torch_custom_ops_enabled directly from .torch_compile.
flashinfer/jit/rocm/activation.py Updates template include path for platform.hpp relocation.
flashinfer/jit/env.py Updates comment reference from csrc_rocm to csrc/rocm.
flashinfer/jit/aiter_source.py Updates module docstring reference from csrc_rocm to csrc/rocm.
flashinfer/get_include_paths.py Updates get_csrc_dir() to return flashinfer/csrc/rocm.
flashinfer/csrc/rocm/single_prefill.cu Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/single_prefill_kernel_inst.jinja New/moved instantiation template under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/single_prefill_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/single_prefill_customize_config.jinja Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/single_prefill_aiter.cu Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/single_prefill_aiter_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/single_decode.cu New/moved decode binding under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/single_decode_kernel_inst.jinja New/moved instantiation template under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/single_decode_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/single_decode_customize_config.jinja Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/sampling.cu New/moved sampling bindings under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/runtime_utils.h New/moved visibility macro header under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/rope_aiter.cu New/moved AITER RoPE binding under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/rope_aiter_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/renorm.cu New/moved renorm bindings under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/quantization.cu Updates include to relocated gpu_iface runtime compat header.
flashinfer/csrc/rocm/pytorch_conversion_utils.h New/moved vector↔tensor helpers for ROCm bindings.
flashinfer/csrc/rocm/pod.cu Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/pod_kernel_inst.jinja New/moved instantiation template under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/pod_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/pod_customize_config.jinja Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/page.cu New/moved page bindings under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/page_aiter.cu New/moved AITER paged-cache append binding under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/page_aiter_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/norm.cu New/moved norm bindings under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/norm_aiter.cu New/moved AITER norm bindings under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/norm_aiter_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/fused_moe_aiter_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/flashinfer_sampling_binding.cu New/moved torch-library fragment binding for sampling ops.
flashinfer/csrc/rocm/flashinfer_rope_binding.cu New/moved torch-library fragment binding for RoPE ops.
flashinfer/csrc/rocm/flashinfer_quantization_binding.cu New/moved torch-library fragment binding for quantization ops.
flashinfer/csrc/rocm/flashinfer_page_binding.cu New/moved torch-library fragment binding for page ops.
flashinfer/csrc/rocm/flashinfer_norm_binding.cu New/moved torch-library fragment binding for norm ops.
flashinfer/csrc/rocm/flashinfer_cascade_binding.cu New/moved torch-library fragment binding for cascade ops.
flashinfer/csrc/rocm/cascade.cu New/moved cascade bindings under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/batch_ragged_prefill_aiter.cu Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/batch_ragged_prefill_aiter_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/batch_prefill.cu Updates gpu_iface include to relocated enums.hpp.
flashinfer/csrc/rocm/batch_prefill_ragged_kernel_inst.jinja New/moved instantiation template under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/batch_prefill_paged_kernel_inst.jinja New/moved instantiation template under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/batch_prefill_paged_aiter.cu Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/batch_prefill_paged_aiter_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/batch_prefill_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/batch_prefill_customize_config.jinja Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/batch_prefill_aiter_customize_config.jinja New/moved AITER config template.
flashinfer/csrc/rocm/batch_pod.cu Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/batch_pod_kernel_inst.jinja New/moved instantiation template under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/batch_pod_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/batch_pod_customize_config.jinja Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/batch_decode.cu Updates gpu_iface includes to relocated utils.cuh and uses conversion utils.
flashinfer/csrc/rocm/batch_decode_kernel_inst.jinja New/moved instantiation template under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/batch_decode_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/batch_decode_customize_config.jinja Updates gpu_iface includes to flashinfer/rocm/gpu_iface/*.
flashinfer/csrc/rocm/batch_decode_aiter.cu New/moved AITER PA-v1 decode harness under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/batch_decode_aiter_jit_pybind.cu New/moved torch-library fragment binding.
flashinfer/csrc/rocm/aot_extension_utils.h New/moved AOT dispatch macro helpers under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/aiter_tensor_compat.h New/moved adapter from at::Tensor to aiter_tensor_t.
flashinfer/csrc/rocm/aiter_loader.cc Updates error-message path reference for moved file.
flashinfer/csrc/rocm/activation.cu New/moved activation kernels under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/activation_aiter.cu New/moved AITER activation binding under flashinfer/csrc/rocm.
flashinfer/csrc/rocm/activation_aiter_jit_pybind.cu New/moved torch-library fragment binding.
docs/rocm/backends.md Updates documented aiter_loader.cc path to flashinfer/csrc/rocm.
CONTRIBUTING.md Updates documentation of ROCm directory layout and canary behavior post-move.
CLAUDE.md Updates path reference for ROCm binding directory rename.
.pre-commit-config.yaml Widens REUSE selective hook path regex to include include/flashinfer/attention/aiter/.
.claude/skills/code-coverage/SKILL.md Updates skill text references from csrc_rocm to csrc/rocm.
.claude/skills/add-rocm-kernel/SKILL.md Updates skill text references from csrc_rocm to csrc/rocm and new gpu_iface path.

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

@demandal25
demandal25 force-pushed the rocm-evacuate-per-op branch from a766467 to d9c1ad1 Compare August 29, 2026 13:48
Copilot AI review requested due to automatic review settings August 29, 2026 13:48
@demandal25 demandal25 changed the title refactor(rocm): evacuate the torch.compile machinery and move csrc_rocm/gpu_iface under rocm/ refactor(rocm): remove the gpu_iface layer and flatten the ROCm headers Aug 29, 2026

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 89 out of 140 changed files in this pull request and generated no new comments.

Suppressed comments (2)

include/flashinfer/rocm/vec_dtypes_hip.h:1919

  • This file opens an include guard (FLASHINFER_ROCM_VEC_DTYPES_HIP_H_) but never closes it with #endif, which will cause a preprocessor error at end-of-file.
    include/flashinfer/rocm/vec_dtypes_hip.h:49
  • FLASHINFER_INLINE is defined twice in the same header without an intervening #undef, which can trigger -Wmacro-redefined warnings (especially since this PR is actively eliminating macro redefinitions).

…ls.py

The fork's opt-in torch.compile support (#210, #230) replaced upstream's
two no-op decorators with ~90 lines in flashinfer/utils.py. That is the
largest single block we add to an upstream file outside the JIT tree.

Moved to flashinfer/rocm/torch_compile.py. utils.py keeps upstream's version
split and two delegating bodies rather than a plain re-import: upstream
defines both functions with those exact signatures inside the same
`if TorchVersion < 2.4 / else`, so delegating shrinks the hunk to the
function bodies. The contextlib, warnings and os imports go with the block.

  flashinfer/utils.py   114 -> 49  (-57%)

utils.py still re-exports use_torch_custom_ops_enabled for
flashinfer/__init__.py:169, which is upstream-owned and sits in the IS_CUDA
arm. On ROCm the name reaches the package surface through rocm/api.py, which
this commit repoints at .torch_compile.

The flag keeps its leading underscore. It is read once per op at import when
the decorator body runs, so a public name would advertise a monkeypatch
target that cannot change any already-registered op.

No cycle: flashinfer/rocm/__init__.py imports only stdlib, so utils.py can
reach the new module at import time even though rocm/api.py imports utils.

Measured on gfx950 (MI350X), ROCm 7.2, torch 2.9.1+rocm7.2, aiter 0.1.10:
test_torch_compile_hip.py and test_mla_and_utils_guards_hip.py pass, including
the subprocess cases that set FLASHINFER_USE_TORCH_CUSTOM_OPS to 0 and 1 and
assert torch.compile respectively raises and succeeds. gfx942 unverified.
Applies the rocm/-subfolder convention to the C++ tree: the ROCm sources
become flashinfer/csrc/rocm/ rather than a _rocm suffix on the parent.

Deliberately package-internal rather than the repo-root csrc/rocm/ the plan
named. Root csrc/ is upstream's and now 568 files at v0.6.18rc3, so putting
fork-only sources inside it invites incidental conflict for no gain -- the
directory has no upstream counterpart either way, so neither location
changes the merge surface. It would also cost a second build-backend
materialisation step: root csrc/ is outside the package, so the wheel would
need include/'s copy-into-package treatment, adding a path where a wheel
ships without the sources its JIT needs.

flashinfer/csrc/ does not exist upstream today, so there is no add/add
conflict now. It is a plainer name than a _rocm suffix would be, so this
accepts a small future add/add risk against reading better; if upstream ever
adds flashinfer/csrc/, the resolution is to keep both subtrees.

get_csrc_dir(), the package-data globs and the MANIFEST graft move with it;
51 references across 16 files follow. CHANGELOG.md keeps the old name --
those entries describe releases that shipped it.
docs/rocm/coverage-gfx942.json also keeps it: that file is a measurement
snapshot taken at 0e68a2f, where csrc_rocm/ was the correct path, and it
is regenerated by a gfx942 coverage run rather than edited by hand.

Verified by diffing the full wheel inventory, base vs HEAD, rather than
spot-checking: 437 -> 438 entries, the 65 csrc_rocm/ paths replaced
one-for-one by csrc/rocm/, and the single addition is
flashinfer/rocm/torch_compile.py from the previous commit. Nothing dropped.
The 65 includes aiter_tensor_compat.h, which #294 added under the old path
after this branch started -- the rebase resolution is a rename, and `git
rebase --skip` would have dropped it and broken three *_aiter.cu compiles.

test_build_backend.py::test_wheel_carries_the_paths_the_jit_resolves asserts
the same property from inside the suite and passes once `build` is installed;
it is importorskip'd, so it is silently absent in an environment without it.
gpu_iface/ and rocm/attention/ each carried a dispatch.cuh and an
exception.h. Both dispatch copies define DISPATCH_NUM_MMA_KV with different
bodies and both use #pragma once, so a TU reaching both silently redefines
the macro:

  attention/dispatch.cuh:31:9: warning: 'DISPATCH_NUM_MMA_KV' macro redefined

Nothing surfaced it because the build does not pass -Wmacro-redefined.

Merged to include/flashinfer/rocm/{dispatch.cuh,exception.h}, above
attention/ because rocm/sampling.cuh needs them too.

dispatch.cuh keeps attention's DISPATCH_NUM_MMA_KV -- the one that falls
back to NUM_MMA_KV=1 rather than raising -- and gains gpu_iface's
DISPATCH_ROPE_DIM, which had no counterpart. That is behaviour-preserving,
not a preference: all seven expansion sites live in prefill.cuh (1694, 2485,
2600), pod.cuh (234, 246) and batch_pod.cuh (212, 223), and each of those
three includes cascade.cuh -- which pulls gpu_iface's copy -- immediately
before the sibling "dispatch.cuh". Attention's definition therefore already
won at every site. gpu_iface's raising body was defined in other TUs but
reached no expansion site, so deleting it changes nothing that ran.

exception.h takes gpu_iface's copy, which is the superset: it carries
#include <string>, which attention's omits while still using std::string.

The relative includes had to be qualified -- attention/{prefill,pod,
batch_pod}.cuh said "dispatch.cuh" and attention/{allocator.h,scheduler.cuh}
said "exception.h", both of which resolved to the now-deleted siblings. A
missed one hard-errors rather than silently picking up the wrong file.

Verified on gfx942 and gfx950: the redefinition probe goes from 1 warning to
0, and cascade, scheduler, norm, pos_enc, page, sampling and quantization
each compile -fsyntax-only with zero errors and zero redefinitions.
rocm/sampling.cuh included upstream's <flashinfer/allocator.h>, which
includes upstream's exception.h. Our exception.h claims the same guard,
FLASHINFER_EXCEPTION_H_, and wins the race -- so upstream's compiled to
nothing and its allocator was built against our header instead. Preprocessing
rocm/sampling.cuh showed our copy expanded first and upstream's producing no
output; grep for 'class Warning' and FLASHINFER_WARN over the result returned
0, confirming upstream's declarations were silently absent. It worked only
because allocator.h happens not to use them.

Point it at the fork's rocm/attention/allocator.h, which differs from
upstream's only in the overflow diagnostic and provides a superset of what
sampling.cuh uses (AlignedAllocator, aligned_alloc<T>). The fork copy's own
include closure is just itself plus rocm/exception.h, so nothing upstream
comes back in through the side door.

Also drops the #else arm that selected the CUDA headers. gpu_iface/macros.hpp
#errors on a non-HIP compiler, so that branch could never be taken; it was
the only place a ROCm header named <flashinfer/math.cuh> and
<flashinfer/vec_dtypes.cuh>, both of which share guards with our HIP copies.

The closure of rocm/sampling.cuh is now entirely fork-owned -- flashinfer/rocm/*
and gpu_iface/* only, zero include/flashinfer/*.h -- which is what makes the
tripwires in a later commit meaningful rather than a list of exemptions.
Compiles -fsyntax-only with 0 errors and 0 redefinitions on gfx942 and gfx950.
memory_ops_hip.h was not standalone-compilable. It names PrefetchMode and
SharedMemFillMode unqualified and got them only from the scope its includer
created: memory_ops.hpp included it from *inside*
flashinfer::gpu_iface::memory, so the backend's own namespace chain nested
inside that one and the real FQN was doubled --
flashinfer::gpu_iface::memory::flashinfer::gpu_iface::memory::detail::hip.
The alias on the next line resolved only because unqualified lookup finds the
outer `flashinfer` again from in there.

Extracting the two enums to memory_types.hpp lets the backend include what it
needs and the wrapper include it at global scope. Proof that this is a real
fix rather than a textual one: a TU naming
flashinfer::gpu_iface::memory::detail::hip::commit_group fails before with

  error: no member named 'detail' in namespace 'flashinfer::gpu_iface::memory'

and compiles after. Rewriting only the alias would have left it compiling via
the same lookup trick, with the doubling intact.

vec_dtypes_hip.h had the same shape but no namespace of its own at all -- it
opened a bare `detail::hip`, so lifting its include to global scope would have
put vec_t and vec_cast in the global namespace. It now declares
flashinfer::gpu_iface::vec_dtypes explicitly, opened *after* its system
includes: those were previously expanded inside three fork namespaces and
survived only because platform.hpp had already pulled them at global scope.
Its anonymous namespace, which shadows HIP's own __float2bfloat162_rn, stays
nested where it is today rather than becoming ambiguous at global scope.

Also deletes error.hpp and backend/hip/mma_debug_utils_hip.h -- zero
includers; the only references to gpu_iface::debug_utils were that header's
own namespace braces.

All 45 fork headers compile -fsyntax-only standalone with zero errors and
zero macro redefinitions on both gfx942 and gfx950.
Since #320 deleted the CUDA backend, include/gpu_iface/ has been HIP-only
headers sitting at the include root as if it were a shared abstraction, with a
backend/hip/ level implying a backend selection that does not exist --
macros.hpp #errors on a non-HIP compiler.

19 headers move up into include/flashinfer/rocm/, losing both directory
levels. No basename collisions: the four backend files all carry _hip
suffixes, and the two that would have clashed with rocm/attention/ were merged
in an earlier commit. The _hip suffixes stay -- with wrapper and backend now
in one directory, that suffix is what distinguishes them.

~140 include lines follow, including 22 across 14 .jinja templates and one in
a Python string at jit/rocm/activation.py, neither of which any C++ tool would
have caught.

upstream_canary pairs forked headers to upstream by basename, so the flatten
puts exception.h, fastdiv.cuh and layout.cuh into the drift report for the
first time -- all three are genuine forks and were previously invisible to it.
utils.cuh is excluded: it shares upstream's basename but 411 of 443 lines
differ, so it is a HIP rewrite that never tracked upstream and would report
drift forever. The exclusion is an exact-path set rather than a directory
prefix, with a test that a sibling is still paired.

Verified on gfx942 and gfx950: all 45 fork headers compile -fsyntax-only
standalone with zero errors and zero macro redefinitions. Negative control: a
TU including the old <gpu_iface/platform.hpp> now fails "file not found", so
nothing is still resolving the old layout from another include root.
gpu_iface described a backend-selection layer that no longer exists: #320
deleted the CUDA half, and macros.hpp #errors on a non-HIP compiler. What
remained was five levels of namespace -- flashinfer::gpu_iface::vec_dtypes::
detail::hip::vec_t -- for what upstream calls flashinfer::vec_t.

Levels are kept only where upstream also groups: math and memory survive
(upstream has flashinfer::math and flashinfer::cp_async). gpu_iface, the
single-backend vec_dtypes wrapper and the detail::hip impl split all go, since
there is no second backend to disambiguate against.

mma is the exception. Its MMAMode, mma_sync_m16n16k16_row_col_f16f16f32 and
m16k16_rowsum_f16f16f32 have exact-signature counterparts in upstream's
flashinfer::mma (mma.cuh:57,318,522), and both headers use bare #pragma once,
so no include guard would catch the collision. They land in
flashinfer::mma_hip instead.

Four headers turned out to be pure pass-throughs once the namespace went --
math_ops, vec_dtypes, memory_ops and mma_ops each re-exported the backend into
the namespace the backend now declares itself, which is a redefinition rather
than a shim. They are deleted and their includers point at the backend
headers, which absorbed the macros.hpp and platform.hpp includes the wrappers
had been contributing.

conversion_utils.h's `namespace fi::con` was the only namespace in the repo
not rooted at flashinfer; explicit_casting joins flashinfer.

Two things no C++ tool would have found: 22 include lines across 14 .jinja
templates plus jit/rocm/activation.py, and five gpu_iface::math uses inside
the JIT-compiled C++ string literal in
test_customize_prefill_use_softmax_hip.py, which fails at test time rather
than build time.

`grep -rn gpu_iface` now returns only CHANGELOG history plus the docs and the
pre-commit path, both retargeted in the next commit. All 42 fork headers
compile -fsyntax-only standalone with zero errors and zero macro redefinitions
on gfx942 and gfx950.
The tree diagram, the intrinsics paragraph and the add-rocm-kernel skill all
described a gpu_iface/backend/hip/ layout and a namespace that no longer
exist. The paragraph's instruction -- "add the abstraction in gpu_iface/ and
implement it under gpu_iface/backend/hip/" -- would have been actively
misleading.

Also rewrites the naming rule. It said to prefer a _rocm/_aiter suffix for
anything upstream might plausibly create; the port has since settled on a
rocm/ subdirectory instead, which is what flashinfer/csrc/rocm/ and
include/flashinfer/rocm/ are, and which collides with nothing even as upstream
grows those trees. The suffix stays as the fallback where a subdirectory does
not fit, which is now just the _hip.h intrinsic headers.

The forked-headers paragraph gains the three headers the flatten newly exposed
to the drift report (layout.cuh, fastdiv.cuh, exception.h) and names utils.cuh
as the one exclusion, so the canary's exclude set is documented where a reader
looks for it rather than only in the script.

The reuse pre-commit pattern loses its dead include/gpu_iface/ alternative and
gains include/flashinfer/attention/aiter/ -- six AMD-authored headers that no
licence check has ever reached.

CLAUDE.md needed no change; it never mentioned gpu_iface.
22 fork headers reused their upstream namesake's include guard, so whichever
arrived first silently suppressed the other. That is how the sampling path
came to compile upstream's allocator.h against our exception.h, fixed two
commits ago -- the failure mode is a wrong-but-compatible header vanishing,
which compiles.

Each now defines FLASHINFER_ROCM_<PATH>_ instead, plus an #error if the
upstream guard is already set. Both include orders are now loud:

  upstream first -> "include/flashinfer/exception.h and
    include/flashinfer/rocm/exception.h both define FLASHINFER_EXCEPTION_H_;
    include only one"
  fork first     -> upstream's header is no longer suppressed, so it reaches
                    the compiler and produces redefinition errors

Ugly in the second direction, but ugly-and-loud beats silent. Upstream's
headers are not touched; the additive-only policy rules that out, which is why
only one direction gets the sentence.

Safe to do now and not before: a preprocessor scan of all 42 fork headers
shows none reaches an upstream include/flashinfer/*.h any more, so nothing was
relying on the suppression. It matters more after the flatten, because
include/flashinfer/rocm/utils.cuh now reads like a peer of
include/flashinfer/utils.cuh, whose ten same-signature functions it shadows.

Note upstream math.cuh and vec_dtypes.cuh cannot be reached from a ROCm TU
regardless -- they include cuda_fp16.h -- so their guards were never the live
risk. exception.h, allocator.h, utils.cuh, layout.cuh and fastdiv.cuh are
pure C++ and were.

All 42 headers compile -fsyntax-only with zero errors and zero redefinitions
on gfx942 and gfx950; pre-commit clean.
The re-guard in the previous commit keyed on `^#ifndef`, so it silently passed
over every fork header that used a bare `#pragma once` -- including
rocm/utils.cuh, which that commit's own message names as the motivating case.
Six headers with an upstream namesake were left without a tripwire:
utils.cuh and attention/{prefill,cascade,page,pod,batch_pod}.cuh.

They never shared a guard, so neither include order was ever silent; they just
produced unnamed redefinition errors where the other 22 name the pair. They
now carry the same #ifdef/#error, so the scheme is uniform.

Also corrects three things this PR invalidated on its own:

- utils.cuh and sampling.cuh still explained an include choice by the
  FLASHINFER_EXCEPTION_H_ race, which the re-guard removed two commits later.
- sampling.cuh had an empty `#ifdef PLATFORM_HIP_DEVICE` / `#endif` left where
  a `using namespace` line was deleted.
- CONTRIBUTING and the add-rocm-kernel skill claimed upstream declares
  `flashinfer::memory` (it declares `cp_async`) and that upstream's mma.cuh has
  no include guard (it has FLASHINFER_MMA_CUH_). The real rule is that a fork
  header keeps its name and tripwires when the two are meant to coexist, and is
  renamed when they are meant to be usable together, which is why mma_hip is
  the one rename. The intrinsics rule also read as an absolute ban on hipcub
  and inline asm outside the shared headers, which six live fork headers break.

Known and accepted: the anonymous namespace in vec_dtypes_hip.h shadows HIP's
__float2bfloat162_rn, and the flatten widened that shadow from
flashinfer::gpu_iface::vec_dtypes to flashinfer. Its callers moved with it, so
in-file resolution is unchanged, and no other header calls either helper.
Three follow-ups from reviewing the flatten.

macros.hpp is what #errors on a non-HIP compiler, so it has to be included
before any hip/ header or the failure is "hip/hip_fp16.h: No such file or
directory" instead of the named diagnostic. cooperative_groups.h already
carries a clang-format-off block saying exactly that; math_hip.h and
sampling.cuh gained macros.hpp during the flatten and clang-format sorted it
after the system includes, silently defeating it. Both now use the same pin.
sampling.cuh's comment asserted the ordering it did not have.

rocm/dispatch.cuh defines eleven DISPATCH_* macros that upstream's utils.cuh
also defines. The guard tripwires added earlier only cover headers that share
an include guard, and these two share neither a guard nor a basename, so the
overlap had no cover at all -- a TU reaching both would silently take
whichever came last, flipping DISPATCH_NUM_MMA_KV between the AMD fallback and
upstream's throwing arm. It now tripwires on FLASHINFER_UTILS_CUH_.

The fork's AlignedAllocator reported only the name, size and alignment on
overflow. Pointing sampling.cuh at it (two commits ago) therefore downgraded
what a user sees when they under-size the workspace, since upstream's copy
names the shortfall and the remedy. Restores that wording, which improves
every ROCm Plan() path rather than just sampling.
vec_dtypes_hip.h defined FLASHINFER_INLINE twice, identically and 21 lines
apart. Pre-existing -- the pre-flatten copy has both -- but a PR whose subject
is removing macro redefinitions should not leave one behind.

docs/rocm/coverage-gfx942.json still named csrc_rocm/ in five `reason`
strings. Those are labels echoed from scripts/coverage_ownership.toml, which
this branch already updated, so leaving them meant the two files disagreed and
the repo carried paths that no longer exist. Only the strings change: the
execution_percent, the base sha and every per-file count are byte-identical,
asserted rather than eyeballed. The snapshot itself stays stale -- its base is
340 commits back and regenerating it needs a gfx942 coverage run -- which is a
pre-existing gap this branch does not close.
@demandal25
demandal25 force-pushed the rocm-evacuate-per-op branch from d9c1ad1 to 4b9ceea Compare August 29, 2026 14:52
Copilot AI review requested due to automatic review settings August 29, 2026 14:52
@demandal25

Copy link
Copy Markdown
Collaborator Author

Suppressed comments from review 5058247273 (against d9c1ad1aa):

vec_dtypes_hip.h:1919 unterminated FLASHINFER_ROCM_VEC_DTYPES_HIP_H_ — declined. The guard is balanced and the header compiles:

$ grep -cE '^[[:space:]]*#[[:space:]]*(if|ifdef|ifndef)\b' vec_dtypes_hip.h   -> 13
$ grep -cE '^[[:space:]]*#[[:space:]]*endif\b'             vec_dtypes_hip.h   -> 13
$ hipcc -fsyntax-only --offload-arch=gfx942 -I include tu.cu; echo $?        -> 0

An unterminated conditional is a hard clang error, so it cannot be present.

vec_dtypes_hip.h:49 duplicate FLASHINFER_INLINE — fixed in 4b9ceea. Pre-existing (the pre-flatten copy has both, identical), but out of place in a PR about macro redefinitions.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 90 out of 141 changed files in this pull request and generated no new comments.

Suppressed comments (2)

include/flashinfer/rocm/vec_dtypes_hip.h:18

  • This header includes <hip/...> headers before any project-level guard like macros.hpp. On a non-HIP compiler that can fail with missing HIP headers instead of the intended explicit diagnostic from macros.hpp. Include macros.hpp first (before any <hip/...> includes) for a consistent, clear error path.
    include/flashinfer/rocm/memory_ops_hip.h:6
  • This header includes <hip/...> headers before any project-level guard like macros.hpp. On a non-HIP compiler that can fail with missing HIP headers instead of the explicit diagnostic from macros.hpp. Include macros.hpp first (before any <hip/...> includes) for a consistent failure mode.

… headers

d14c973 did this for math_hip.h; a sweep of the 34 HIP-dependent
headers under include/flashinfer/rocm (41 files in total) found four
more that still failed with "fatal error: hip/hip_bf16.h: No such file
or directory" on a non-HIP compiler instead of the named diagnostic
from macros.hpp:25.

  memory_ops_hip.h, vec_dtypes_hip.h, conversion_utils.h  -- fixed here
  attention/permuted_smem.cuh                             -- transitively,
      via its first include (memory_ops_hip.h)

Copilot flagged two of the four (review 5058411285). The sweep is what
found the other two, so the failure mode is now gone tree-wide rather
than at the two reported sites:

  for h in $(find include/flashinfer/rocm \( -name '*.h' -o -name '*.hpp' \
                  -o -name '*.cuh' \)); do
    echo "#include \"${h#include/}\"" > t.cpp
    g++ -fsyntax-only -I include t.cpp
  done
  # before: 30 named-diagnostic, 4 hip-header-not-found
  # after:  34 named-diagnostic, 0

macros.hpp carries a conditional `#define __forceinline__ inline`, so
ordering it ahead of the hip/ includes has to be shown not to weaken
always_inline tree-wide. What happens: hipcc predefines no
__forceinline__, so the #ifndef *does* fire and the macro is briefly
`inline` --

  echo '#include "flashinfer/rocm/macros.hpp"' > b.hip
  hipcc -dM -E -I include --offload-arch=gfx942 b.hip | grep -w __forceinline__
  #define __forceinline__ inline

-- and hip/amd_detail/host_defines.h:233 then unconditionally redefines
it back. Nothing sits between the two, so the intermediate value is
never applied to a declaration:

  hipcc -E -I include m.hip | grep -v '^#'   # parent vs this commit
  # 14 diff lines, all blank; 0 non-blank token difference

The resulting redefinition is invisible to a default build because
host_defines.h is a system header; -Wsystem-headers surfaces it. It is
not new here -- math_hip.h, platform.hpp and fastdiv.cuh already reach
macros.hpp first -- but this commit widens it, so it is named rather
than left to be rediscovered. Unlike the DISPATCH_NUM_MMA_KV
redefinition C3 removed, both definitions are known and the later one
deterministically wins.

An earlier draft of this message claimed the #ifndef never fires and
that -Wmacro-redefined reported zero; both were artefacts of measuring
the post-host_defines.h value and of system-header warning suppression.
Copilot AI review requested due to automatic review settings August 29, 2026 23:07
@demandal25

Copy link
Copy Markdown
Collaborator Author

Suppressed comments from review 5058411285 (against 4b9ceea88) — both accepted, fixed in f6e460629.

vec_dtypes_hip.h:18 and memory_ops_hip.h:6, both the same claim: macros.hpp must precede the <hip/...> includes so its non-HIP #error fires first. Sweeping all 41 fork headers found four with that ordering, not two — conversion_utils.h and attention/permuted_smem.cuh as well (the latter transitively, via memory_ops_hip.h). Before: 30 named-diagnostic / 4 hip/... No such file. After: 34 / 0.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 90 out of 141 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

include/flashinfer/rocm/attention/allocator.h:16

  • This header uses std::string (including as a value parameter) but doesn't include , which can lead to compile errors when included from a minimal translation unit.

@demandal25

Copy link
Copy Markdown
Collaborator Author

Suppressed comment from review 5059414122 (against f6e460629) — declined.

attention/allocator.h:16 missing <string>: it compiles clean as a minimal TU, because <string> arrives from an explicit include in the header it pulls on the very line flagged.

$ echo '#include "flashinfer/rocm/attention/allocator.h"' > alloc.hip
$ hipcc -fsyntax-only -I include --offload-arch=gfx942 alloc.hip; echo $?
0
$ grep -n include include/flashinfer/rocm/exception.h
13:#include <exception>
14:#include <sstream>
15:#include <string>

Not an incidental libstdc++ transitive — exception.h:15 is an in-repo explicit include, and allocator.h:16 includes exception.h directly.

@demandal25
demandal25 merged commit 5064ec5 into amd-integration Aug 29, 2026
3 checks passed
@demandal25
demandal25 deleted the rocm-evacuate-per-op branch August 29, 2026 23:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants