feat(hip): cascade attention support on ROCm using HIP

[demandal25]

Summary

• Exposes the full cascade kernel family on HIP: merge_state, merge_state_in_place, merge_states, variable_length_merge_states, attention_sum, variable_length_attention_sum
• Splits cascade.cu into kernels + bindings files, mirroring the CUDA layout so gen_cascade_module() resolves both files unchanged
• Exports MultiLevelCascadeAttentionWrapper, BatchDecodeWithSharedPrefixPagedKVCacheWrapper, BatchPrefillWithSharedPrefixPagedKVCacheWrapper, and the merge_state* family from flashinfer.__init__ on HIP
• Adds optional partial_state=(out, lse) kwarg to BatchPrefillWithPagedKVCacheWrapper.run() for fused cascade epilogue (gated by FLASHINFER_HIP_FUSED_CASCADE=1)
• Tunes merge kernel launch geometry for CDNA3: wave-64 threadblocks for head_dim≤128, reduces num_smem_stages from 4→1 (pred_load is synchronous on CDNA3)
• Adds warp_sync_state to generic/cascade.cuh with __builtin_amdgcn_wave_barrier() on HIP and __syncwarp() on CUDA, enabling future persistent merge-states kernel support

Design choice: no cascade_rocm.py (unlike prefill_rocm.py / decode_rocm.py)

Other backends-divergent modules (prefill, decode, mla) are forked into separate *_rocm.py files and aliased via sys.modules in flashinfer/__init__.py. Cascade is shared in a single cascade.py with a small inline HIP branch. This is intentional. The heuristic used:

Fork the file when the method bodies diverge. Share the file when only the composition diverges.

prefill.py vs prefill_rocm.py are ~3000 lines each, and the method bodies are genuinely different on each backend (JIT compilation path, backend dispatch list — fa2/fa3/cudnn/trtllm on CUDA vs fa2/aiter/auto on HIP — kernel param packing layouts, ROCm-specific PDL handling, partial_state plumbing, AITER bootstrap, etc.). Sharing the file would mean nearly every method becoming a big if IS_HIP: ... else: ... block. Forking is correct.

cascade.py, by contrast, is mostly orchestration. MultiLevelCascadeAttentionWrapper.run loops over per-level prefill wrappers and merges results. Both the wrappers (BatchPrefillWithPagedKVCacheWrapper) and the merge helper (merge_state_in_place) are already backend-swapped via the sys.modules alias in __init__.py, so the cascade orchestration is correct on both backends automatically — the loop body needs no if IS_HIP: at all. The only HIP-specific divergence is one optional kernel-fusion optimization, gated behind an env var.

Concretely, if cascade were forked into cascade_rocm.py, this PR would duplicate ~500 lines of identical orchestration plumbing (the MultiLevelCascadeAttentionWrapper, BatchDecodeWithSharedPrefixPagedKVCacheWrapper, BatchPrefillWithSharedPrefixPagedKVCacheWrapper, and the merge_state* wrapper functions) for the sake of one 4-line behavioral difference. Every future cascade.py improvement would then need to be ported twice. The cost would dwarf the benefit.

Known consequence: the partial_state kwarg on the HIP-side BatchPrefillWithPagedKVCacheWrapper.run() does not exist on the CUDA variant, so static type-checking at the cascade.py call site sees only the CUDA overloads and flags a call-overload error. This will need a localized # type: ignore[call-overload] (or equivalent) at that call site. That is a meaningfully smaller maintenance cost than the ~500 LOC duplication a cascade_rocm.py fork would impose.

Architecture note: two cascade merge paths in generic/prefill.cuh

The ROCm prefill kernel uses generic/cascade.cuh differently from the CUDA prefill.cuh:

CUDA — cascade merging is always a post-kernel API call:

prefill kernel → (writes partial outputs) → MergeStates / VariableLengthMergeStates

ROCm — two separate mechanisms:

1. Non-split-KV path (fused epilogue, lines 2325–2370): When params.partial_o != nullptr && !partition_kv, the cascade merge is inlined inside the prefill kernel. The kernel reads the prior cascade level's partial_o/partial_lse from global memory and blends them with the current output registers before the output write — no separate kernel launch, and no call into generic/cascade.cuh. This is the path enabled by FLASHINFER_HIP_FUSED_CASCADE=1.

2. Split-KV path (post-kernel, lines 2636–2644): Follows the same pattern as CUDA — VariableLengthMergeStates merges the KV-split chunks, then MergeStateInPlace folds in the prior cascade level. Both are post-kernel calls into generic/cascade.cuh.

Test plan

• test_cascade_hip.py: covers merge_state, merge_state_in_place, merge_states (fp16/bf16, multiple shapes), merge_state_in_place with boolean mask (all-ones, all-zeros, random), and fused_cascade_epilogue correctness vs float32 reference
• test_shared_prefix_kernels_hip.py: ports test_batch_attention_with_shared_prefix_paged_kv_cache from the CUDA suite (MultiLevelCascadeAttentionWrapper vs legacy shared-prefix wrappers); adds test_multilevel_cascade_fused_vs_unfused comparing the FLASHINFER_HIP_FUSED_CASCADE path against the standalone merge_state_in_place path for head_dim in {128, 256} and shared_kv_len in {128, 512, 2048}

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR extends FlashInfer’s cascade attention feature set on ROCm/HIP by wiring up the cascade kernel family (and associated wrappers) in the HIP extension, aligning the ROCm cascade module layout with CUDA, and adding an optional fused “cascade epilogue” path for prefill that merges a prior cascade state directly in-kernel.

Changes:

• Adds/exports ROCm cascade bindings (merge/merge-in-place/merge-states + variable-length variants + attention-sum variants) and makes HIP flashinfer.__init__ expose cascade wrappers + merge_state* APIs.
• Introduces an optional fused cascade epilogue for HIP FA2 prefill via partial_state=(out, lse) (gated by FLASHINFER_HIP_FUSED_CASCADE=1) and corresponding kernel parameter plumbing.
• Adds ROCm test coverage for cascade ops and shared-prefix multi-level cascade behavior; tunes merge kernel launch geometry for CDNA3.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tests/rocm_tests/test_shared_prefix_kernels_hip.py	Adds ROCm tests comparing multilevel cascade vs shared-prefix wrappers and fused vs unfused paths.
tests/rocm_tests/test_cascade_hip.py	Adds ROCm tests for merge_state APIs, masks, and fused epilogue correctness.
tests/rocm_tests/test_activation_hip.py	Replaces large header docstring with SPDX header.
include/flashinfer/attention/generic/prefill.cuh	Implements HIP-only fused cascade epilogue merge and post-merge handling in dispatched path.
include/flashinfer/attention/generic/default_prefill_params.cuh	Adds partial_o / partial_lse pointers to prefill params.
include/flashinfer/attention/generic/cascade.cuh	Adds helper + CDNA3 tuning for merge/attention-sum kernels.
flashinfer/prefill_rocm.py	Adds partial_state plumbing/validation and passes optional partial tensors to ROCm FA2 op.
flashinfer/csrc_rocm/flashinfer_cascade_binding.cu	Registers ROCm cascade ops in the extension library.
flashinfer/csrc_rocm/cascade.cu	Implements ROCm C++ entrypoints for variable-length merge/attention-sum ops.
flashinfer/csrc_rocm/batch_prefill.cu	Extends ROCm paged prefill entrypoint with optional partial-state tensors.
flashinfer/csrc_rocm/batch_prefill_jit_pybind.cu	Updates the pybind declaration to match the new optional args.
flashinfer/csrc_rocm/batch_prefill_customize_config.jinja	Extends generated paged params with optional partial-state pointers.
flashinfer/cascade.py	Adds HIP fused-cascade toggle and routes multi-level cascade to the fused prefill path when enabled.
flashinfer/init.py	Exports cascade wrappers and merge_state* from the HIP entrypoint after ROCm module injection.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.