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

flashinfer/__init__.py

@@ -259,6 +259,22 @@
     sys.modules["flashinfer.prefill"] = sys.modules["flashinfer.prefill_rocm"]
     sys.modules["flashinfer.decode"] = sys.modules["flashinfer.decode_rocm"]
 
+    # Cascade imports must come after the sys.modules injection above so that
+    # cascade.py's relative imports of flashinfer.decode / flashinfer.prefill
+    # resolve to the ROCm implementations.
+    from .cascade import (
+        BatchDecodeWithSharedPrefixPagedKVCacheWrapper as BatchDecodeWithSharedPrefixPagedKVCacheWrapper,
+    )
+    from .cascade import (
+        BatchPrefillWithSharedPrefixPagedKVCacheWrapper as BatchPrefillWithSharedPrefixPagedKVCacheWrapper,
+    )
+    from .cascade import (
+        MultiLevelCascadeAttentionWrapper as MultiLevelCascadeAttentionWrapper,
+    )
+    from .cascade import merge_state as merge_state
+    from .cascade import merge_state_in_place as merge_state_in_place
+    from .cascade import merge_states as merge_states
+
     from .utils import next_positive_power_of_2 as next_positive_power_of_2
     from .utils import use_torch_custom_ops_enabled as use_torch_custom_ops_enabled
 else:

flashinfer/cascade.py

@@ -15,15 +15,21 @@
 """
 
 import functools
+import os
 from typing import List, Optional, Tuple, Union
 
 import torch
 
 from .decode import BatchDecodeWithPagedKVCacheWrapper
+from .device_utils import IS_HIP
 from .jit.cascade import gen_cascade_module
 from .prefill import BatchPrefillWithPagedKVCacheWrapper, single_prefill_with_kv_cache
 from .utils import register_custom_op, register_fake_op
 
+_HIP_FUSED_CASCADE = (
+    IS_HIP and os.environ.get("FLASHINFER_HIP_FUSED_CASCADE", "0") == "1"
+)
+
 
 @functools.cache
 def get_cascade_module():
@@ -537,8 +543,13 @@ def run(
             return_lse=True,
         )
         for wrapper in self._batch_prefill_wrappers[:-1]:
-            out_i, lse_i = wrapper.run(q, paged_kv_cache, return_lse=True)
-            merge_state_in_place(out, lse, out_i, lse_i)
+            if _HIP_FUSED_CASCADE:
+                out, lse = wrapper.run(
+                    q, paged_kv_cache, return_lse=True, partial_state=(out, lse)
+                )
+            else:
+                out_i, lse_i = wrapper.run(q, paged_kv_cache, return_lse=True)
+                merge_state_in_place(out, lse, out_i, lse_i)
 
         return out
 

flashinfer/csrc_rocm/batch_prefill.cu

@@ -204,7 +204,9 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
                                      at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices,
                                      at::Tensor paged_kv_last_page_len, at::Tensor o,
                                      std::optional<at::Tensor> maybe_lse, int64_t mask_mode_code,
-                                     int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS) {
+                                     int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS,
+                                     std::optional<at::Tensor> maybe_partial_o = std::nullopt,
+                                     std::optional<at::Tensor> maybe_partial_lse = std::nullopt) {
   PrefillPlanInfo plan_info;
   plan_info.FromVector(tensor_to_vec(plan_info_vec));
   QKVLayout kv_layout = static_cast<QKVLayout>(layout);
@@ -267,6 +269,10 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
         params.o = static_cast<DTypeO*>(o.data_ptr());
 
         params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.partial_o =
+            maybe_partial_o ? static_cast<DTypeO*>(maybe_partial_o->data_ptr()) : nullptr;
+        params.partial_lse =
+            maybe_partial_lse ? static_cast<float*>(maybe_partial_lse->data_ptr()) : nullptr;
         params.num_qo_heads = num_qo_heads;
         params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads);
         params.q_stride_n = q_stride_n;

flashinfer/csrc_rocm/batch_prefill_customize_config.jinja

@@ -90,6 +90,8 @@ struct PagedParams {
   IdType* q_indptr;
   DTypeO* o;
   float* lse;
+  DTypeO* partial_o = nullptr;
+  float* partial_lse = nullptr;
   uint_fastdiv group_size;
 
   {{ additional_params_decl }}

flashinfer/csrc_rocm/batch_prefill_jit_pybind.cu

@@ -37,7 +37,9 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
                                      at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices,
                                      at::Tensor paged_kv_last_page_len, at::Tensor o,
                                      std::optional<at::Tensor> maybe_lse, int64_t mask_mode_code,
-                                     int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS);
+                                     int64_t layout, int64_t window_left ADDITIONAL_FUNC_PARAMS,
+                                     std::optional<at::Tensor> maybe_partial_o = std::nullopt,
+                                     std::optional<at::Tensor> maybe_partial_lse = std::nullopt);
 
 TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
   // Batch-request prefill attention with KV-Cache plan

flashinfer/csrc_rocm/cascade.cu

@@ -131,3 +131,86 @@ void merge_states(at::Tensor v, at::Tensor s, at::Tensor v_merged, at::Tensor s_
 
   TORCH_CHECK(success, "MergeStates kernel launch failed: unsupported data type");
 }
+
+void variable_length_merge_states(at::Tensor v, at::Tensor s, at::Tensor indptr,
+                                  at::Tensor v_merged, at::Tensor s_merged) {
+  CHECK_INPUT(v);
+  CHECK_INPUT(s);
+  CHECK_INPUT(indptr);
+  auto device = v.device();
+  CHECK_EQ(s.device(), device);
+  CHECK_EQ(indptr.device(), device);
+  CHECK_DIM(3, v);
+  CHECK_DIM(2, s);
+  CHECK_DIM(1, indptr);
+  CHECK_EQ(v.size(0), s.size(0));
+  CHECK_EQ(v.size(1), s.size(1));
+  unsigned int num_heads = v.size(1);
+  unsigned int head_dim = v.size(2);
+  unsigned int seq_len = indptr.size(0) - 1;
+
+  const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(v.device());
+  const hipStream_t stream = at::hip::getCurrentHIPStream();
+  bool success = DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(v.scalar_type(), c_type, [&] {
+    hipError_t status = VariableLengthMergeStates(
+        static_cast<c_type*>(v.data_ptr()), static_cast<float*>(s.data_ptr()),
+        static_cast<int32_t*>(indptr.data_ptr()), static_cast<c_type*>(v_merged.data_ptr()),
+        static_cast<float*>(s_merged.data_ptr()), seq_len, /*seq_len_ptr=*/nullptr, num_heads,
+        head_dim, stream);
+    TORCH_CHECK(status == hipSuccess,
+                "VariableLengthMergeStates kernel launch failed: ", hipGetErrorString(status));
+    return true;
+  });
+  TORCH_CHECK(success, "VariableLengthMergeStates kernel launch failed: unsupported data type");
+}
+
+void attention_sum(at::Tensor v, at::Tensor v_sum, int64_t num_index_sets) {
+  CHECK_INPUT(v);
+  CHECK_INPUT(v_sum);
+  CHECK_EQ(v.device(), v_sum.device());
+  CHECK_DIM(4, v);
+  CHECK_DIM(3, v_sum);
+  unsigned int seq_len = v.size(0);
+  unsigned int num_heads = v.size(2);
+  unsigned int head_dim = v.size(3);
+
+  const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(v.device());
+  const hipStream_t stream = at::hip::getCurrentHIPStream();
+  bool success = DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(v.scalar_type(), c_type, [&] {
+    hipError_t status =
+        AttentionSum(static_cast<c_type*>(v.data_ptr()), static_cast<c_type*>(v_sum.data_ptr()),
+                     static_cast<uint32_t>(num_index_sets), seq_len, num_heads, head_dim, stream);
+    TORCH_CHECK(status == hipSuccess,
+                "AttentionSum kernel launch failed: ", hipGetErrorString(status));
+    return true;
+  });
+  TORCH_CHECK(success, "AttentionSum kernel launch failed: unsupported data type");
+}
+
+void variable_length_attention_sum(at::Tensor v, at::Tensor indptr, at::Tensor v_sum) {
+  CHECK_INPUT(v);
+  CHECK_INPUT(indptr);
+  CHECK_INPUT(v_sum);
+  auto device = v.device();
+  CHECK_EQ(indptr.device(), device);
+  CHECK_EQ(v_sum.device(), device);
+  CHECK_DIM(3, v);
+  CHECK_DIM(1, indptr);
+  CHECK_DIM(3, v_sum);
+  unsigned int num_heads = v.size(1);
+  unsigned int head_dim = v.size(2);
+  unsigned int seq_len = indptr.size(0) - 1;
+
+  const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(v.device());
+  const hipStream_t stream = at::hip::getCurrentHIPStream();
+  bool success = DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(v.scalar_type(), c_type, [&] {
+    hipError_t status = VariableLengthAttentionSum(
+        static_cast<c_type*>(v.data_ptr()), static_cast<int32_t*>(indptr.data_ptr()),
+        static_cast<c_type*>(v_sum.data_ptr()), seq_len, /*seq_len_ptr=*/nullptr, num_heads,
+        head_dim, stream);
+    TORCH_CHECK(status == hipSuccess,
+                "VariableLengthAttentionSum kernel launch failed: ", hipGetErrorString(status));
+    return true;
+  });
+  TORCH_CHECK(success, "VariableLengthAttentionSum kernel launch failed: unsupported data type");
+}

flashinfer/csrc_rocm/flashinfer_cascade_binding.cu

@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) 2026 Advanced Micro Devices, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <optional>
+
+#include "pytorch_extension_utils.h"
+
+void merge_state(at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b,
+                 at::Tensor v_merged, at::Tensor s_merged);
+
+void merge_state_in_place(at::Tensor v, at::Tensor s, at::Tensor v_other, at::Tensor s_other,
+                          std::optional<at::Tensor> mask);
+
+void merge_states(at::Tensor v, at::Tensor s, at::Tensor v_merged, at::Tensor s_merged);
+
+void variable_length_merge_states(at::Tensor v, at::Tensor s, at::Tensor indptr,
+                                  at::Tensor v_merged, at::Tensor s_merged);
+
+void attention_sum(at::Tensor v, at::Tensor v_sum, int64_t num_index_sets);
+
+void variable_length_attention_sum(at::Tensor v, at::Tensor indptr, at::Tensor v_sum);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  m.def("merge_state", merge_state);
+  m.def("merge_state_in_place", merge_state_in_place);
+  m.def("merge_states", merge_states);
+  m.def("variable_length_merge_states", variable_length_merge_states);
+  m.def("attention_sum", attention_sum);
+  m.def("variable_length_attention_sum", variable_length_attention_sum);
+}

flashinfer/prefill_rocm.py

@@ -670,6 +670,8 @@ def paged_run(
         cum_seq_lens_q: Optional[torch.Tensor] = None,
         cum_seq_lens_kv: Optional[torch.Tensor] = None,
         sinks: Optional[torch.Tensor] = None,
+        maybe_partial_o: Optional[torch.Tensor] = None,
+        maybe_partial_lse: Optional[torch.Tensor] = None,
     ) -> None:
         if backend != "fa2":
             logger.warning(
@@ -704,6 +706,8 @@ def paged_run(
             1.0 / rope_scale,  # rope_rcp_scale
             1.0 / rope_theta,  # rope_rcp_theta
             # token_pos_in_items_len,  # Not supported by HIP FA2 kernels
+            maybe_partial_o,
+            maybe_partial_lse,
         )
 
     @register_fake_op(f"flashinfer::{uri}_paged_run")
@@ -746,6 +750,8 @@ def _fake_paged_run(
         batch_size: Optional[int] = None,
         cum_seq_lens_q: Optional[torch.Tensor] = None,
         cum_seq_lens_kv: Optional[torch.Tensor] = None,
+        maybe_partial_o: Optional[torch.Tensor] = None,
+        maybe_partial_lse: Optional[torch.Tensor] = None,
     ) -> None:
         pass
 
@@ -2061,6 +2067,7 @@ def run(
         enable_pdl: Optional[bool] = None,
         window_left: Optional[int] = None,
         sinks: Optional[torch.Tensor] = None,
+        partial_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         r"""Compute batch prefill/append attention between query and paged kv-cache.
 
@@ -2182,6 +2189,20 @@ def run(
         sparse_indptr = self._paged_kv_indptr_buf
 
         assert self._plan_info is not None, "plan info is not initialized"
+        if partial_state is not None:
+            if partial_state[0].dtype != out.dtype:
+                raise ValueError(
+                    f"partial_state dtype {partial_state[0].dtype} must match output dtype {out.dtype}"
+                )
+            if partial_state[0].device != out.device:
+                raise ValueError(
+                    f"partial_state device {partial_state[0].device} must match output device {out.device}"
+                )
+            # Ensure lse is allocated so the kernel can write the merged LSE output.
+            if lse is None:
+                lse = torch.empty(
+                    (q.size(0), q.size(1)), dtype=torch.float32, device=q.device
+                )
         run_args = [
             self._float_workspace_buffer,
             self._int_workspace_buffer,
@@ -2232,6 +2253,9 @@ def run(
                 sinks,
             ]
             if self._backend == "aiter":
+                assert partial_state is None, (
+                    "partial_state (fused cascade epilogue) is only supported with the fa2 backend"
+                )
                 # Pre-computed flat-KV gather info for AITER (None for
                 # natively-supported page sizes).
                 run_args += [
@@ -2240,6 +2264,9 @@ def run(
                     self._aiter_flat_kv_lpl,
                     self._aiter_flat_kv_indices,
                 ]
+            else:
+                po, plse = partial_state if partial_state is not None else (None, None)
+                run_args += [po, plse]
 
         assert self._cached_module is not None, "cached module is not initialized"
         self._cached_module.paged_run(*run_args)

include/flashinfer/attention/generic/cascade.cuh

@@ -1,5 +1,5 @@
 // SPDX-FileCopyrightText: 2023-2025 FlashInfer team.
-// SPDX-FileCopyrightText: 2025 Advanced Micro Devices, Inc.
+// SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
 #pragma once
@@ -133,6 +133,29 @@ __device__ __forceinline__ void threadblock_sync_state(state_t<vec_size>& st, DT
   __syncthreads();
 }
 
+template <uint32_t bdx, uint32_t bdy, uint32_t vec_size, typename DTypeIn>
+__device__ __forceinline__ void warp_sync_state(state_t<vec_size>& st, DTypeIn* v_smem,
+                                                float* s_smem, const uint32_t tx = threadIdx.x,
+                                                const uint32_t ty = threadIdx.y) {
+  constexpr uint32_t head_dim = vec_size * bdx;
+  st.o.cast_store(v_smem + ty * head_dim + tx * vec_size);
+  s_smem[ty] = st.get_lse();
+  st.init();
+#ifdef PLATFORM_HIP_DEVICE
+  __builtin_amdgcn_wave_barrier();
+#else
+  __syncwarp();
+#endif
+
+#pragma unroll
+  for (uint32_t iter = 0; iter < bdy; ++iter) {
+    float s = s_smem[iter];
+    vec_t<float, vec_size> v;
+    v.cast_load(v_smem + iter * head_dim + tx * vec_size);
+    st.merge(v, s, 1);
+  }
+}
+
 template <uint32_t bdx, uint32_t bdy, uint32_t vec_size, typename DTypeIn>
 __device__ __forceinline__ void threadblock_sum(vec_t<float, vec_size>& v, DTypeIn* v_smem) {
   const uint32_t tx = threadIdx.x, ty = threadIdx.y;
@@ -626,11 +649,18 @@ gpuError_t MergeStates(DTypeIn* v, float* s, DTypeO* v_merged, float* s_merged,
     constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U);
     constexpr uint32_t bdx = HEAD_DIM / vec_size;
     if (num_index_sets >= seq_len) {
+#ifdef PLATFORM_HIP_DEVICE
+      // CDNA3 wave-64: fit one wavefront per threadblock for head_dim≤128; stages=1
+      // since pred_load is synchronous (no async pipeline).
+      constexpr uint32_t num_threads = (bdx <= 16) ? 64U : 256U;
+      constexpr uint32_t num_smem_stages = 1;
+#else
       constexpr uint32_t num_threads = 128;
+      constexpr uint32_t num_smem_stages = 4;
+#endif
       constexpr uint32_t bdy = num_threads / bdx;
       dim3 nblks(seq_len, num_heads);
       dim3 nthrs(bdx, bdy);
-      constexpr uint32_t num_smem_stages = 4;
       auto kernel =
           MergeStatesLargeNumIndexSetsKernel<vec_size, bdx, bdy, num_smem_stages, DTypeIn, DTypeO>;
       void* args[] = {&v, &s, &v_merged, &s_merged, &num_index_sets, &num_heads};
@@ -681,9 +711,14 @@ gpuError_t VariableLengthMergeStates(DTypeIn* v, float* s, IdType* indptr, DType
   DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, {
     constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U);
     constexpr uint32_t bdx = HEAD_DIM / vec_size;
+#ifdef PLATFORM_HIP_DEVICE
+    constexpr uint32_t num_threads = (bdx <= 16) ? 64U : 256U;
+    constexpr uint32_t num_smem_stages = 1;
+#else
     constexpr uint32_t num_threads = 128;
-    constexpr uint32_t bdy = num_threads / bdx;
     constexpr uint32_t num_smem_stages = 4;
+#endif
+    constexpr uint32_t bdy = num_threads / bdx;
     uint32_t smem_size =
         num_smem_stages * bdy * head_dim * sizeof(DTypeIn) + num_threads * sizeof(float);
     auto kernel = PersistentVariableLengthMergeStatesKernel<vec_size, bdx, bdy, num_smem_stages,
@@ -715,9 +750,14 @@ gpuError_t VariableLengthAttentionSum(DTypeIn* v, IdType* indptr, DTypeO* v_sum,
   DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, {
     constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U);
     constexpr uint32_t bdx = HEAD_DIM / vec_size;
+#ifdef PLATFORM_HIP_DEVICE
+    constexpr uint32_t num_threads = (bdx <= 16) ? 64U : 256U;
+    constexpr uint32_t num_smem_stages = 1;
+#else
     constexpr uint32_t num_threads = 128;
-    constexpr uint32_t bdy = num_threads / bdx;
     constexpr uint32_t num_smem_stages = 4;
+#endif
+    constexpr uint32_t bdy = num_threads / bdx;
     uint32_t smem_size = num_smem_stages * bdy * head_dim * sizeof(DTypeIn);
     auto kernel = PersistentVariableLengthAttentionSumKernel<vec_size, bdx, bdy, num_smem_stages,
                                                              DTypeIn, DTypeO, IdType>;