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);
@@ -226,6 +228,8 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
     TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0));
     TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1));
   }
+  TORCH_CHECK(maybe_partial_o.has_value() == maybe_partial_lse.has_value(),
+              "partial_o and partial_lse must both be provided or both be absent");
 
   void* float_buffer_ptr = static_cast<void*>(float_workspace_buffer.data_ptr());
   void* int_buffer_ptr = static_cast<void*>(int_workspace_buffer.data_ptr());
@@ -267,6 +271,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,90 @@ 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);
+  TORCH_CHECK(indptr.scalar_type() == at::kInt,
+              "variable_length_merge_states: indptr must be int32, got ", indptr.scalar_type());
+  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);
+  TORCH_CHECK(indptr.scalar_type() == at::kInt,
+              "variable_length_attention_sum: indptr must be int32, got ", indptr.scalar_type());
+  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,28 @@ 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}"
+                )
+            if partial_state[1].dtype != torch.float32:
+                raise ValueError(
+                    f"partial_state lse must be float32, got {partial_state[1].dtype}"
+                )
+            if partial_state[1].device != out.device:
+                raise ValueError(
+                    f"partial_state lse device {partial_state[1].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 +2261,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 +2272,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)