feat(hip): port RoPE to ROCm

[demandal25]

Summary

Closes the remaining gaps between the CUDA and ROCm RoPE implementations: adds int64 position-id dispatch, enables MLA (DeepSeek-class) paths in the fused RoPE+FP8 quantize and paged-cache append kernels, and tunes launch geometry for CDNA3/Wave64.

What changed

DISPATCH_PYTORCH_IDTYPE_TO_CTYPE macro

• flashinfer/csrc_rocm/pytorch_extension_utils.h — New macro mirroring CUDA's DISPATCH_DLPACK_IDTYPE_TO_CTYPE, dispatching at::kInt→int32_t and at::kLong→int64_t. Used by all five apply_rope* / apply_llama31_rope* HIP bindings so callers can freely pass either dtype for position tensors.

MLA support in the fused RoPE kernel

• include/flashinfer/attention/generic/pos_enc.cuh — Adds typename CacheT = paged_kv_t<QuantType, IdType> as a 7th template parameter on RopeQuantizeAppendPagedKVCacheKernel. A constexpr bool IS_MLA = std::is_same<CacheT, paged_kv_mla_t<...>>::value branch selects at compile time:
  • K pointer arithmetic: 2D stride (no head dim) for MLA vs 3D for GQA/MHA
  • Cache writes: get_kpe_ptr / get_ckv_ptr (MLA) vs get_k_ptr / get_v_ptr (GQA/MHA)
  • V section: skipped entirely for MLA — total_blocks_y excludes the V blocks when IS_MLA
• New RopeQuantizeAppendPagedMLACache<DType, IdType, QuantType> host launcher constructs a paged_kv_mla_t and calls the shared kernel.

CDNA3 launch-geometry tuning

Replaces NVIDIA-ported warp32 heuristics with Wave64-aware geometry throughout pos_enc.cuh:

Kernel group	Old	New
BatchQKApplyRotary* / BatchQKApplyLlama31Rotary*	vec_size = max(16/sizeof(DType), HEAD_DIM/32), 128-thread block	vec_size = max(16/sizeof(DType), HEAD_DIM/kWarpSize), num_threads = max(2*kWarpSize, bdx)
RopeQuantize / RopeQuantizeAppendPagedKVCache	vec_size = 32/sizeof(DType), fixed 128 threads	vec_size = 16/sizeof(DType) (single global_load_dwordx4), num_threads = max(2*kWarpSize, bdx)

All changes use gpu_iface::kWarpSize so CUDA continues to see warp32 geometry unchanged.

HIP rope bindings

• flashinfer/csrc_rocm/rope.cu
  • apply_rope / apply_llama31_rope: nested DISPATCH_PYTORCH_IDTYPE_TO_CTYPE inside the existing dtype dispatch supports both int32 and int64 position tensors. Added TORCH_CHECK(offsets.scalar_type() == indptr.scalar_type()) to catch dtype mismatches between the two index tensors.
  • rope_quantize: 2D-K MLA branch — when k_rope_in.ndim() == 2, sets num_kv_heads=1 and aliases the head stride to the token stride. Added TORCH_CHECK(no_rope_dim == 0 || no_rope_dim % rope_dim == 0) to prevent OOB in the K-nope tiling loop.
  • rope_quantize_append_paged_kv_cache: replaces the TORCH_CHECK(!has_mla_caches) hard-error with a full paged_kv_mla_t dispatch via RopeQuantizeAppendPagedMLACache; adds TORCH_CHECK(!(has_gqa_caches && has_mla_caches)) so callers providing both sets of caches fail fast; moves v_in dtype and contiguity validation inside the GQA-only block.
• flashinfer/rope.py: adds v.dtype != q_rope.dtype check in the Python GQA path of rope_quantize_fp8_append_paged_kv_cache.

Test coverage

• tests/rocm_tests/test_rope_hip.py:
  • Enables idtype parametrize (torch.int32 / torch.int64) on test_rope_pos_ids
  • Un-skips test_rope_cos_sin_cache — uses 5× looser tolerances motivated by measured ~3e-2 bfloat16 rounding error (two cast_load/cast_store round-trips through bf16 in the interleaved rotation path)
  • Adds test_mla_rope_quantize (2D-K, DeepSeek-class config)
  • Adds MLA variants to test_generalized_rope_quantize_hip and test_generalized_rope_quantize_append_kv_cache_hip
  • Adds test_rope_quantize_fp8_append_paged_kv_cache_decode_hip (decode scenario)
  • Updates _rope_apply_interleave_f32 reference helper to handle 2D K natively
• tests/rocm_tests/test_activation_hip.py: removes stale if __name__ == "__main__": block.

Test plan

• pytest tests/rocm_tests/test_rope_hip.py -m "not slow" — 13,079 passed on MI300X (gfx942)
• No regressions in existing GQA/MHA rope paths
• apply_rope / apply_llama31_rope pass with int64 position IDs across all (head_dim, interleave, inplace) variants
• test_rope_cos_sin_cache passes on HIP (previously skipped) with bfloat16 inputs
• test_mla_rope_quantize and MLA variants in fused-quantize and cache-append tests pass
• pytest tests/rocm_tests/ -m "not slow" — all rope tests pass; full suite results tracked separately
• pre-commit run -a passes

Generated with Claude Code

[Copilot]

Pull request overview

This PR extends the ROCm/HIP RoPE implementation to match CUDA feature coverage by adding int64 position-id dispatch, enabling MLA (2D-K) support in fused RoPE+FP8 quantize and paged-cache append paths, and tuning kernel launch geometry (notably for CDNA3).

Changes:

• Add DISPATCH_PYTORCH_IDTYPE_TO_CTYPE and use it in HIP RoPE bindings to support int32/int64 position-id tensors.
• Unify the fused RoPE+quantize+append kernel to support both GQA/MHA (paged_kv_t) and MLA (paged_kv_mla_t) via a CacheT template + constexpr branching; add an MLA-specific host launcher.
• Expand ROCm test coverage for pos-id dtype variants and MLA paths; re-enable the cos/sin-cache test on HIP with looser tolerances.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
flashinfer/csrc_rocm/pytorch_extension_utils.h	Adds an idtype dispatch macro used by ROCm bindings to accept int32/int64 tensors.
flashinfer/csrc_rocm/rope.cu	Updates HIP bindings to dispatch idtype, adds MLA support in quantize/append paths, and routes MLA cache appends through paged_kv_mla_t.
include/flashinfer/attention/generic/pos_enc.cuh	Generalizes the fused kernel to support MLA cache layout and updates launch geometry using kWarpSize.
tests/rocm_tests/test_rope_hip.py	Adds idtype parametrization, MLA coverage, and re-enables cos/sin-cache testing with HIP-appropriate tolerances.
tests/rocm_tests/test_activation_hip.py	Removes ad-hoc __main__ execution block from the test file.

---

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

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Comments suppressed due to low confidence (1)

include/flashinfer/attention/generic/pos_enc.cuh:1236

• Similarly, the Q-noRoPE section unconditionally loads/stores vec_size elements per thread with no tail guard. For no_rope_dim values that aren’t multiples of rope_dim, this can read/write out of bounds. Add a tail predicate (or validate no_rope_dim % rope_dim == 0 up front) so the last chunk doesn’t access beyond q_nope_* tensors.

    } else {
      uint32_t q_nope_start = k_nope_end + (IS_MLA ? 0u : num_kv_heads);
      uint32_t q_head_idx = (by - q_nope_start) / no_rope_chunks;
      uint32_t nope_chunk_idx = (by - q_nope_start) % no_rope_chunks;
      uint32_t elem_offset = nope_chunk_idx * rope_chunk_size;

      DType* q_nope_in_ptr =
          q_nope_in + get_elem_offset_impl(idx, q_head_idx, elem_offset, q_nope_in_stride_n,
                                           q_nope_in_stride_h);
      QuantType* q_nope_out_ptr =
          q_nope_out + get_elem_offset_impl(idx, q_head_idx, elem_offset, q_nope_out_stride_n,
                                            q_nope_out_stride_h);

      vec_t q_nope_vec;
      q_nope_vec.cast_load(q_nope_in_ptr + tx * vec_size);
#pragma unroll

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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