diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 7b13974fcc..0d7cca81ff 100644 --- a/flashinfer/__init__.py +++ b/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: diff --git a/flashinfer/cascade.py b/flashinfer/cascade.py index 267f0d2990..ff562d8d9b 100644 --- a/flashinfer/cascade.py +++ b/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 diff --git a/flashinfer/csrc_rocm/batch_prefill.cu b/flashinfer/csrc_rocm/batch_prefill.cu index e8db3df769..503f212193 100644 --- a/flashinfer/csrc_rocm/batch_prefill.cu +++ b/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 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 maybe_partial_o = std::nullopt, + std::optional maybe_partial_lse = std::nullopt) { PrefillPlanInfo plan_info; plan_info.FromVector(tensor_to_vec(plan_info_vec)); QKVLayout kv_layout = static_cast(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(float_workspace_buffer.data_ptr()); void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); @@ -267,6 +271,10 @@ void BatchPrefillWithPagedKVCacheRun(at::Tensor float_workspace_buffer, params.o = static_cast(o.data_ptr()); params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.partial_o = + maybe_partial_o ? static_cast(maybe_partial_o->data_ptr()) : nullptr; + params.partial_lse = + maybe_partial_lse ? static_cast(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; diff --git a/flashinfer/csrc_rocm/batch_prefill_customize_config.jinja b/flashinfer/csrc_rocm/batch_prefill_customize_config.jinja index d069c4c4f8..bdd832ae99 100644 --- a/flashinfer/csrc_rocm/batch_prefill_customize_config.jinja +++ b/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 }} diff --git a/flashinfer/csrc_rocm/batch_prefill_jit_pybind.cu b/flashinfer/csrc_rocm/batch_prefill_jit_pybind.cu index 024fcde70e..b9fd2f329e 100644 --- a/flashinfer/csrc_rocm/batch_prefill_jit_pybind.cu +++ b/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 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 maybe_partial_o = std::nullopt, + std::optional maybe_partial_lse = std::nullopt); TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { // Batch-request prefill attention with KV-Cache plan diff --git a/flashinfer/csrc_rocm/cascade.cu b/flashinfer/csrc_rocm/cascade.cu index d2cadfb0db..7ad4812238 100644 --- a/flashinfer/csrc_rocm/cascade.cu +++ b/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(v.data_ptr()), static_cast(s.data_ptr()), + static_cast(indptr.data_ptr()), static_cast(v_merged.data_ptr()), + static_cast(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(v.data_ptr()), static_cast(v_sum.data_ptr()), + static_cast(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(v.data_ptr()), static_cast(indptr.data_ptr()), + static_cast(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"); +} diff --git a/flashinfer/csrc_rocm/flashinfer_cascade_binding.cu b/flashinfer/csrc_rocm/flashinfer_cascade_binding.cu new file mode 100644 index 0000000000..9ab007d16b --- /dev/null +++ b/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 + +#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 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); +} diff --git a/flashinfer/prefill_rocm.py b/flashinfer/prefill_rocm.py index c27131053f..2c4988d7b4 100755 --- a/flashinfer/prefill_rocm.py +++ b/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) diff --git a/include/flashinfer/attention/generic/cascade.cuh b/include/flashinfer/attention/generic/cascade.cuh index cc26c3e466..fc08f6ef22 100644 --- a/include/flashinfer/attention/generic/cascade.cuh +++ b/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,30 @@ __device__ __forceinline__ void threadblock_sync_state(state_t& st, DT __syncthreads(); } +// Called by PersistentMergeStatesKernel in persistent.cuh for warp-level state reduction. +template +__device__ __forceinline__ void warp_sync_state(state_t& 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 v; + v.cast_load(v_smem + iter * head_dim + tx * vec_size); + st.merge(v, s, 1); + } +} + template __device__ __forceinline__ void threadblock_sum(vec_t& v, DTypeIn* v_smem) { const uint32_t tx = threadIdx.x, ty = threadIdx.y; @@ -626,11 +650,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; void* args[] = {&v, &s, &v_merged, &s_merged, &num_index_sets, &num_heads}; @@ -681,9 +712,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; diff --git a/include/flashinfer/attention/generic/default_prefill_params.cuh b/include/flashinfer/attention/generic/default_prefill_params.cuh index 9ec70638bb..5f59fb3720 100644 --- a/include/flashinfer/attention/generic/default_prefill_params.cuh +++ b/include/flashinfer/attention/generic/default_prefill_params.cuh @@ -273,6 +273,9 @@ struct BatchPrefillPagedParams { DTypeO* o; float* lse; float* maybe_alibi_slopes; + // Non-null: kernel merges its output with this prior cascade level's state in-register. + DTypeO* partial_o = nullptr; + float* partial_lse = nullptr; uint_fastdiv group_size; uint32_t num_qo_heads; IdType q_stride_n; diff --git a/include/flashinfer/attention/generic/prefill.cuh b/include/flashinfer/attention/generic/prefill.cuh index 9131be0f7f..f48e190a09 100644 --- a/include/flashinfer/attention/generic/prefill.cuh +++ b/include/flashinfer/attention/generic/prefill.cuh @@ -2322,6 +2322,53 @@ __device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( // normalize d normalize_d(o_frag, m, d); +#ifdef PLATFORM_HIP_DEVICE + // Cascade epilogue: merge with a prior cascade level's output in-register. + // Skipped for split-KV chunks (partition_kv=true); those are merged in + // BatchPrefillWithPagedKVCacheDispatched after VariableLengthMergeStates. + if constexpr (AttentionVariant::use_softmax) { + if (params.partial_o != nullptr && !partition_kv) { + if (get_warp_idx_kv(tid.z) == 0) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < NUM_ACCUM_ROWS_PER_THREAD; ++j) { + uint32_t q_idx, r; + group_size.divmod( + qo_packed_idx_base + + (lane_idx / THREADS_PER_BMATRIX_ROW_SET) * NUM_ACCUM_ROWS_PER_THREAD + j + + mma_q * 16, + q_idx, r); + const uint32_t qo_head_idx = kv_head_idx * group_size + r; + const uint32_t qo_idx = q_idx; + if (qo_idx < qo_upper_bound) { + const float s_cur = gpu_iface::math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + const float s_partial = + params.partial_lse[(o_indptr[request_idx] + qo_idx) * num_qo_heads + qo_head_idx]; + const float s_max = fmaxf(s_cur, s_partial); + const float scale_a = exp2f(s_cur - s_max); + const float scale_b = exp2f(s_partial - s_max); + const float inv_denom = 1.0f / (scale_a + scale_b); + const uint32_t po_base = + (o_indptr[request_idx] + qo_idx) * o_stride_n + qo_head_idx * o_stride_h; +#pragma unroll + for (uint32_t mma_d = 0; mma_d < NUM_MMA_D_VO; ++mma_d) { + const float p_o = + (float)params + .partial_o[po_base + mma_d * 16 + lane_idx % THREADS_PER_BMATRIX_ROW_SET]; + o_frag[mma_q][mma_d][j] = + (o_frag[mma_q][mma_d][j] * scale_a + p_o * scale_b) * inv_denom; + } + m[mma_q][j] = static_cast(s_max); + d[mma_q][j] = scale_a + scale_b; + } + } + } + } + } + } +#endif // PLATFORM_HIP_DEVICE + const uint32_t num_kv_chunks = (kv_len_safe + kv_chunk_size - 1) / kv_chunk_size; // write_back @@ -2589,6 +2636,13 @@ gpuError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params FI_GPU_CALL(VariableLengthMergeStates(tmp_v, tmp_s, params.merge_indptr, o, lse, params.max_total_num_rows, params.total_num_rows, num_qo_heads, HEAD_DIM_VO, stream)); +#ifdef PLATFORM_HIP_DEVICE + if (params.partial_o != nullptr) { + FI_GPU_CALL(MergeStateInPlace(o, lse, params.partial_o, params.partial_lse, + params.max_total_num_rows, num_qo_heads, HEAD_DIM_VO, + nullptr, stream)); + } +#endif } else { FI_GPU_CALL(VariableLengthAttentionSum(tmp_v, params.merge_indptr, o, params.max_total_num_rows, params.total_num_rows, diff --git a/tests/rocm_tests/test_activation_hip.py b/tests/rocm_tests/test_activation_hip.py index b4b55893e3..7a42299af8 100644 --- a/tests/rocm_tests/test_activation_hip.py +++ b/tests/rocm_tests/test_activation_hip.py @@ -1,31 +1,6 @@ -""" -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. - -HIP/ROCm tests for fused activation kernels: silu_and_mul, gelu_tanh_and_mul, -gelu_and_mul. - -Each test generates input of shape (num_tokens, 2*d), runs the flashinfer -fused kernel, and compares against a pure-PyTorch reference computed in fp32 -and cast back to the target dtype. - -Shapes cover: - - Small d values (< 64 elements / vec_size, partial last wavefront) - - Typical LLM FFN sizes (LLaMA-2 11008, Mistral 14336, Phi-3 14336 / 2) - - Large d values that hit the 1024-thread blockDim cap - - Non-power-of-two d values (e.g., 5504 = LLaMA-2 11008 // 2) -""" +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 import pytest import torch diff --git a/tests/rocm_tests/test_cascade_hip.py b/tests/rocm_tests/test_cascade_hip.py new file mode 100644 index 0000000000..fb58368115 --- /dev/null +++ b/tests/rocm_tests/test_cascade_hip.py @@ -0,0 +1,349 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +import flashinfer + + +def merge_state_ref(v_a, s_a, v_b, s_b): + """Float32 reference: merge two attention states (s values are log2-based LSE).""" + s_a = s_a.float() + s_b = s_b.float() + v_a = v_a.float() + v_b = v_b.float() + # s values are logsumexp in base 2: s = log2(sum(2^scores)) + m = torch.maximum(s_a, s_b) + scale_a = torch.exp2(s_a - m) # 2^(s_a - m) + scale_b = torch.exp2(s_b - m) # 2^(s_b - m) + denom = scale_a + scale_b + v_merged = ( + v_a * scale_a.unsqueeze(-1) + v_b * scale_b.unsqueeze(-1) + ) / denom.unsqueeze(-1) + s_merged = m + torch.log2(denom) # log2(2^s_a + 2^s_b) + return v_merged, s_merged + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seq_len", [1, 64, 512]) +@pytest.mark.parametrize("num_heads", [4, 16]) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_merge_state(dtype, seq_len, num_heads, head_dim): + torch.manual_seed(42) + atol = 5e-3 if dtype == torch.float16 else 1e-2 + + v_a = torch.randn(seq_len, num_heads, head_dim, dtype=dtype, device="cuda") + s_a = torch.randn(seq_len, num_heads, dtype=torch.float32, device="cuda") + v_b = torch.randn(seq_len, num_heads, head_dim, dtype=dtype, device="cuda") + s_b = torch.randn(seq_len, num_heads, dtype=torch.float32, device="cuda") + + v_ref, s_ref = merge_state_ref(v_a, s_a, v_b, s_b) + v_merged, s_merged = flashinfer.merge_state(v_a, s_a, v_b, s_b) + + torch.testing.assert_close(v_merged.float(), v_ref, rtol=1e-3, atol=atol) + torch.testing.assert_close(s_merged, s_ref, rtol=1e-3, atol=atol) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seq_len", [1, 64, 512]) +@pytest.mark.parametrize("num_heads", [4, 16]) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_merge_state_in_place(dtype, seq_len, num_heads, head_dim): + torch.manual_seed(42) + atol = 5e-3 if dtype == torch.float16 else 1e-2 + + v_a = torch.randn(seq_len, num_heads, head_dim, dtype=dtype, device="cuda") + s_a = torch.randn(seq_len, num_heads, dtype=torch.float32, device="cuda") + v_b = torch.randn(seq_len, num_heads, head_dim, dtype=dtype, device="cuda") + s_b = torch.randn(seq_len, num_heads, dtype=torch.float32, device="cuda") + + v_ref, s_ref = merge_state_ref(v_a, s_a, v_b, s_b) + v_a_copy = v_a.clone() + s_a_copy = s_a.clone() + flashinfer.merge_state_in_place(v_a_copy, s_a_copy, v_b, s_b) + + torch.testing.assert_close(v_a_copy.float(), v_ref, rtol=1e-3, atol=atol) + torch.testing.assert_close(s_a_copy, s_ref, rtol=1e-3, atol=atol) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("num_index_sets", [1, 4, 16, 64, 256]) +@pytest.mark.parametrize("seq_len", [1, 32]) +@pytest.mark.parametrize("num_heads", [4]) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_merge_states(dtype, num_index_sets, seq_len, num_heads, head_dim): + torch.manual_seed(42) + atol = 5e-3 if dtype == torch.float16 else 1e-2 + + v = torch.randn( + seq_len, num_index_sets, num_heads, head_dim, dtype=dtype, device="cuda" + ) + s = torch.randn( + seq_len, num_index_sets, num_heads, dtype=torch.float32, device="cuda" + ) + + # Reference: merge all states iteratively in float32 to avoid accumulation error + v_ref = v[:, 0, :, :].float() + s_ref = s[:, 0, :] + for i in range(1, num_index_sets): + v_ref, s_ref = merge_state_ref(v_ref, s_ref, v[:, i, :, :].float(), s[:, i, :]) + + v_merged, s_merged = flashinfer.merge_states(v, s) + + torch.testing.assert_close(v_merged.float(), v_ref.float(), rtol=1e-3, atol=atol) + torch.testing.assert_close(s_merged, s_ref, rtol=1e-3, atol=atol) + + +@pytest.mark.parametrize("seed", [0]) +@pytest.mark.parametrize("num_tries", [20]) +def test_merge_state_in_place_with_mask(seed, num_tries): + seq_len = 512 + num_heads = 8 + head_dim = 128 + va = torch.randn(seq_len, num_heads, head_dim).half().to("cuda") + sa = torch.randn(seq_len, num_heads, dtype=torch.float32).to("cuda") + vb = torch.randn(seq_len, num_heads, head_dim).half().to("cuda") + sb = torch.randn(seq_len, num_heads, dtype=torch.float32).to("cuda") + va_original = va.clone() + sa_original = sa.clone() + + # No mask + flashinfer.merge_state_in_place(va, sa, vb, sb) + va_merged_ref = va.clone() + sa_merged_ref = sa.clone() + assert not torch.allclose(va_merged_ref, va_original) + assert not torch.allclose(sa_merged_ref, sa_original) + + # Mask all-ones: identical to no mask + mask = torch.ones(seq_len, dtype=torch.bool, device="cuda") + va2 = va_original.clone() + sa2 = sa_original.clone() + flashinfer.merge_state_in_place(va2, sa2, vb, sb, mask=mask) + torch.testing.assert_close(va2, va_merged_ref, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(sa2, sa_merged_ref, rtol=1e-3, atol=1e-3) + + # Mask all-zeros: output unchanged + mask = torch.zeros(seq_len, dtype=torch.bool, device="cuda") + va2 = va_original.clone() + sa2 = sa_original.clone() + flashinfer.merge_state_in_place(va2, sa2, vb, sb, mask=mask) + torch.testing.assert_close(va2, va_original, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(sa2, sa_original, rtol=1e-3, atol=1e-3) + + # Random masks + randgen = torch.Generator(device="cuda") + randgen.manual_seed(seed) + for _ in range(num_tries): + rand_mask = ( + torch.rand(seq_len, generator=randgen, dtype=torch.float32, device="cuda") + > 0.5 + ).to(dtype=torch.bool) + true_indices = rand_mask.nonzero() + false_indices = (rand_mask == 0).nonzero() + va2 = va_original.clone() + sa2 = sa_original.clone() + flashinfer.merge_state_in_place(va2, sa2, vb, sb, mask=rand_mask) + torch.testing.assert_close( + va2[false_indices], va_original[false_indices], rtol=1e-3, atol=1e-3 + ) + torch.testing.assert_close( + sa2[false_indices], sa_original[false_indices], rtol=1e-3, atol=1e-3 + ) + torch.testing.assert_close( + va2[true_indices], va_merged_ref[true_indices], rtol=1e-3, atol=1e-3 + ) + torch.testing.assert_close( + sa2[true_indices], sa_merged_ref[true_indices], rtol=1e-3, atol=1e-3 + ) + + +def _ceil_div(a, b): + return (a + b - 1) // b + + +def _build_cascade_fixture( + batch_size, num_heads, head_dim, shared_kv_len, unique_kv_len, dtype +): + """Build paged-KV fixture for a two-level cascade (shared + unique KV).""" + page_size = 16 + assert shared_kv_len % page_size == 0 + kv_layout = "NHD" + + q = torch.randn(batch_size, num_heads, head_dim, dtype=dtype, device="cuda") + k_shared = torch.randn( + shared_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + v_shared = torch.randn( + shared_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + k_unique = torch.randn( + batch_size * unique_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + v_unique = torch.randn( + batch_size * unique_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + + kv_data = torch.zeros( + _ceil_div(shared_kv_len, page_size) + + batch_size * _ceil_div(unique_kv_len, page_size), + 2, + page_size, + num_heads, + head_dim, + dtype=dtype, + device="cuda", + ) + + shared_kv_indices = torch.arange( + 0, _ceil_div(shared_kv_len, page_size), dtype=torch.int32, device="cuda" + ) + shared_append_indptr = ( + torch.arange(0, 2, dtype=torch.int32, device="cuda") * shared_kv_len + ) + shared_kv_indptr = torch.arange(0, 2, dtype=torch.int32, device="cuda") * _ceil_div( + shared_kv_len, page_size + ) + shared_last_page_len = torch.full( + (1,), (shared_kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda" + ) + flashinfer.append_paged_kv_cache( + k_shared, + v_shared, + *flashinfer.get_batch_indices_positions( + shared_append_indptr, + flashinfer.get_seq_lens(shared_kv_indptr, shared_last_page_len, page_size), + k_shared.shape[0], + ), + kv_data, + shared_kv_indices, + shared_kv_indptr, + shared_last_page_len, + kv_layout, + ) + + unique_kv_indices = torch.arange( + 0, + batch_size * _ceil_div(unique_kv_len, page_size), + dtype=torch.int32, + device="cuda", + ) + _ceil_div(shared_kv_len, page_size) + unique_append_indptr = ( + torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") + * unique_kv_len + ) + unique_kv_indptr = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device="cuda" + ) * _ceil_div(unique_kv_len, page_size) + unique_last_page_len = torch.full( + (batch_size,), + (unique_kv_len - 1) % page_size + 1, + dtype=torch.int32, + device="cuda", + ) + flashinfer.append_paged_kv_cache( + k_unique, + v_unique, + *flashinfer.get_batch_indices_positions( + unique_append_indptr, + flashinfer.get_seq_lens(unique_kv_indptr, unique_last_page_len, page_size), + k_unique.shape[0], + ), + kv_data, + unique_kv_indices, + unique_kv_indptr, + unique_last_page_len, + kv_layout, + ) + + return ( + q, + kv_data, + kv_layout, + shared_kv_indptr, + shared_kv_indices, + shared_last_page_len, + unique_kv_indptr, + unique_kv_indices, + unique_last_page_len, + ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("num_heads", [4, 16]) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_fused_cascade_epilogue(dtype, batch_size, num_heads, head_dim): + """Fused in-kernel cascade merge must agree with standalone merge_state_in_place path.""" + import flashinfer.cascade as cascade_mod + + torch.manual_seed(42) + atol = 5e-3 if dtype == torch.float16 else 1e-2 + shared_kv_len = 128 + unique_kv_len = 17 + page_size = 16 + + ( + q, + kv_data, + kv_layout, + shared_kv_indptr, + shared_kv_indices, + shared_last_page_len, + unique_kv_indptr, + unique_kv_indices, + unique_last_page_len, + ) = _build_cascade_fixture( + batch_size, num_heads, head_dim, shared_kv_len, unique_kv_len, dtype + ) + + def make_wrapper_and_plan(): + w = flashinfer.MultiLevelCascadeAttentionWrapper( + 2, torch.empty(32 * 1024 * 1024, dtype=torch.int8, device="cuda"), kv_layout + ) + qo_indptr_top = torch.tensor([0, q.shape[0]], dtype=torch.int32, device="cuda") + qo_indptr_bottom = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device="cuda" + ) + w.plan( + [qo_indptr_top, qo_indptr_bottom], + [shared_kv_indptr, unique_kv_indptr], + [shared_kv_indices, unique_kv_indices], + [shared_last_page_len, unique_last_page_len], + num_heads, + num_heads, + head_dim, + page_size, + q_data_type=dtype, + ) + return w + + orig_flag = cascade_mod._HIP_FUSED_CASCADE + + cascade_mod._HIP_FUSED_CASCADE = False + try: + w_unfused = make_wrapper_and_plan() + out_unfused = w_unfused.run(q, kv_data) + finally: + cascade_mod._HIP_FUSED_CASCADE = orig_flag + + cascade_mod._HIP_FUSED_CASCADE = True + try: + w_fused = make_wrapper_and_plan() + out_fused = w_fused.run(q, kv_data) + finally: + cascade_mod._HIP_FUSED_CASCADE = orig_flag + + torch.testing.assert_close( + out_fused.float(), out_unfused.float(), rtol=1e-3, atol=atol + ) + + +if __name__ == "__main__": + test_merge_state(torch.float16, 64, 8, 128) + test_merge_state_in_place(torch.float16, 64, 8, 128) + test_merge_states(torch.float16, 16, 32, 8, 128) + test_merge_state_in_place_with_mask(0, 20) + test_fused_cascade_epilogue(torch.float16, 4, 16, 128) + test_fused_cascade_epilogue(torch.bfloat16, 4, 16, 128) diff --git a/tests/rocm_tests/test_shared_prefix_kernels_hip.py b/tests/rocm_tests/test_shared_prefix_kernels_hip.py new file mode 100644 index 0000000000..bb10397269 --- /dev/null +++ b/tests/rocm_tests/test_shared_prefix_kernels_hip.py @@ -0,0 +1,394 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 +# +# Ported from tests/attention/test_shared_prefix_kernels.py + +import pytest +import torch +from jit_utils import gen_decode_attention_modules, gen_prefill_attention_modules + +import flashinfer +from flashinfer.utils import has_flashinfer_jit_cache + + +@pytest.fixture( + autouse=not has_flashinfer_jit_cache(), + scope="module", +) +def warmup_jit(): + flashinfer.jit.build_jit_specs( + gen_decode_attention_modules( + [torch.float16], + [torch.float16], + [128, 256], + [0], + [False], + [False], + ) + + gen_prefill_attention_modules( + [torch.float16], + [torch.float16], + [128, 256], + [0], + [False], + [False], + [False], + ), + verbose=False, + ) + yield + + +def ceil_div(a, b): + return (a + b - 1) // b + + +@pytest.mark.parametrize("stage", ["decode", "append"]) +@pytest.mark.parametrize("batch_size", [12, 17]) +@pytest.mark.parametrize("unique_kv_len", [37, 17]) +@pytest.mark.parametrize("shared_kv_len", [128, 512, 2048]) +@pytest.mark.parametrize("num_heads", [8, 16]) +@pytest.mark.parametrize("causal", [False]) +@pytest.mark.parametrize("head_dim", [128, 256]) +@pytest.mark.parametrize("page_size", [1, 16]) +def test_batch_attention_with_shared_prefix_paged_kv_cache( + stage, + batch_size, + unique_kv_len, + shared_kv_len, + num_heads, + causal, + head_dim, + page_size, +): + if stage == "decode" and causal: + pytest.skip("Causal attention is not required in decode stage") + assert shared_kv_len % page_size == 0 + kv_layout = "NHD" + if stage == "append": + q = torch.randn(batch_size * unique_kv_len, num_heads, head_dim).to(0).half() + q_indptr = torch.arange(0, batch_size + 1).to(0).int() * unique_kv_len + else: + q = torch.randn(batch_size, num_heads, head_dim).to(0).half() + q_indptr = torch.arange(0, batch_size + 1).to(0).int() + k_shared = torch.randn(shared_kv_len, num_heads, head_dim).to(0).half() + v_shared = torch.randn(shared_kv_len, num_heads, head_dim).to(0).half() + k_unique = torch.randn(batch_size * unique_kv_len, num_heads, head_dim).to(0).half() + v_unique = torch.randn(batch_size * unique_kv_len, num_heads, head_dim).to(0).half() + + kv_data = ( + torch.zeros( + ceil_div(shared_kv_len, page_size) + + batch_size * ceil_div(unique_kv_len, page_size), + 2, + page_size, + num_heads, + head_dim, + ) + .to(0) + .half() + ) + shared_kv_indices = torch.arange(0, ceil_div(shared_kv_len, page_size)).to(0).int() + shared_append_indptr = torch.arange(0, 2).to(0).int() * shared_kv_len + shared_kv_indptr = torch.arange(0, 2).to(0).int() * ceil_div( + shared_kv_len, page_size + ) + shared_last_page_len = torch.full( + (1,), (shared_kv_len - 1) % page_size + 1, dtype=torch.int32 + ).to(0) + flashinfer.append_paged_kv_cache( + k_shared, + v_shared, + *flashinfer.get_batch_indices_positions( + shared_append_indptr, + flashinfer.get_seq_lens(shared_kv_indptr, shared_last_page_len, page_size), + k_shared.shape[0], + ), + kv_data, + shared_kv_indices, + shared_kv_indptr, + shared_last_page_len, + kv_layout, + ) + unique_kv_indices = torch.arange( + 0, batch_size * ceil_div(unique_kv_len, page_size) + ).to(0).int() + ceil_div(shared_kv_len, page_size) + unique_append_indptr = torch.arange(0, batch_size + 1).to(0).int() * unique_kv_len + unique_kv_indptr = torch.arange(0, batch_size + 1).to(0).int() * ceil_div( + unique_kv_len, page_size + ) + unique_last_page_len = torch.full( + (batch_size,), (unique_kv_len - 1) % page_size + 1, dtype=torch.int32 + ).to(0) + flashinfer.append_paged_kv_cache( + k_unique, + v_unique, + *flashinfer.get_batch_indices_positions( + unique_append_indptr, + flashinfer.get_seq_lens(unique_kv_indptr, unique_last_page_len, page_size), + k_unique.shape[0], + ), + kv_data, + unique_kv_indices, + unique_kv_indptr, + unique_last_page_len, + kv_layout, + ) + + multi_level_wrapper = flashinfer.MultiLevelCascadeAttentionWrapper( + 2, torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0), kv_layout + ) + if stage == "decode": + shared_prefix_decode_wrapper = ( + flashinfer.BatchDecodeWithSharedPrefixPagedKVCacheWrapper( + torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0), kv_layout + ) + ) + else: + shared_prefix_prefill_wrapper = ( + flashinfer.BatchPrefillWithSharedPrefixPagedKVCacheWrapper( + torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0), kv_layout + ) + ) + + qo_indptr_top = torch.tensor([0, q.shape[0]], dtype=torch.int32).to(0) + if stage == "decode": + qo_indptr_bottom = torch.arange(0, batch_size + 1, dtype=torch.int32).to(0) + multi_level_wrapper.plan( + [qo_indptr_top, qo_indptr_bottom], + [shared_kv_indptr, unique_kv_indptr], + [shared_kv_indices, unique_kv_indices], + [shared_last_page_len, unique_last_page_len], + num_heads, + num_heads, + head_dim, + page_size, + ) + o_multi_level = multi_level_wrapper.run(q, kv_data) + else: + qo_indptr_bottom = ( + torch.arange(0, batch_size + 1, dtype=torch.int32).to(0) * unique_kv_len + ) + multi_level_wrapper.plan( + [qo_indptr_top, qo_indptr_bottom], + [shared_kv_indptr, unique_kv_indptr], + [shared_kv_indices, unique_kv_indices], + [shared_last_page_len, unique_last_page_len], + num_heads, + num_heads, + head_dim, + page_size, + causal=causal, + ) + o_multi_level = multi_level_wrapper.run(q, kv_data) + + if stage == "decode": + shared_prefix_decode_wrapper.begin_forward( + unique_kv_indptr, + unique_kv_indices, + unique_last_page_len, + num_heads, + num_heads, + head_dim, + page_size, + ) + o_two_level = shared_prefix_decode_wrapper.forward( + q, k_shared, v_shared, kv_data + ) + else: + shared_prefix_prefill_wrapper.begin_forward( + q_indptr, + unique_kv_indptr, + unique_kv_indices, + unique_last_page_len, + num_heads, + num_heads, + head_dim, + page_size, + ) + o_two_level = shared_prefix_prefill_wrapper.forward( + q, k_shared, v_shared, kv_data, causal=causal + ) + + torch.testing.assert_close(o_multi_level, o_two_level, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("stage", ["decode", "append"]) +@pytest.mark.parametrize("batch_size", [4, 12]) +@pytest.mark.parametrize("unique_kv_len", [17, 37]) +@pytest.mark.parametrize("shared_kv_len", [128, 512, 2048]) +@pytest.mark.parametrize("num_heads", [8, 16]) +@pytest.mark.parametrize("head_dim", [128, 256]) +@pytest.mark.parametrize("page_size", [16]) +def test_multilevel_cascade_fused_vs_unfused( + dtype, + stage, + batch_size, + unique_kv_len, + shared_kv_len, + num_heads, + head_dim, + page_size, +): + """Fused epilogue path must agree with unfused merge_state_in_place path.""" + import flashinfer.cascade as cascade_mod + + torch.manual_seed(42) + atol = 5e-3 if dtype == torch.float16 else 1e-2 + assert shared_kv_len % page_size == 0 + kv_layout = "NHD" + + if stage == "append": + q = torch.randn( + batch_size * unique_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + else: + q = torch.randn(batch_size, num_heads, head_dim, dtype=dtype, device="cuda") + + k_shared = torch.randn( + shared_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + v_shared = torch.randn( + shared_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + k_unique = torch.randn( + batch_size * unique_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + v_unique = torch.randn( + batch_size * unique_kv_len, num_heads, head_dim, dtype=dtype, device="cuda" + ) + + kv_data = torch.zeros( + ceil_div(shared_kv_len, page_size) + + batch_size * ceil_div(unique_kv_len, page_size), + 2, + page_size, + num_heads, + head_dim, + dtype=dtype, + device="cuda", + ) + shared_kv_indices = torch.arange( + 0, ceil_div(shared_kv_len, page_size), dtype=torch.int32, device="cuda" + ) + shared_append_indptr = ( + torch.arange(0, 2, dtype=torch.int32, device="cuda") * shared_kv_len + ) + shared_kv_indptr = torch.arange(0, 2, dtype=torch.int32, device="cuda") * ceil_div( + shared_kv_len, page_size + ) + shared_last_page_len = torch.full( + (1,), (shared_kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda" + ) + flashinfer.append_paged_kv_cache( + k_shared, + v_shared, + *flashinfer.get_batch_indices_positions( + shared_append_indptr, + flashinfer.get_seq_lens(shared_kv_indptr, shared_last_page_len, page_size), + k_shared.shape[0], + ), + kv_data, + shared_kv_indices, + shared_kv_indptr, + shared_last_page_len, + kv_layout, + ) + + unique_kv_indices = torch.arange( + 0, + batch_size * ceil_div(unique_kv_len, page_size), + dtype=torch.int32, + device="cuda", + ) + ceil_div(shared_kv_len, page_size) + unique_append_indptr = ( + torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") + * unique_kv_len + ) + unique_kv_indptr = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device="cuda" + ) * ceil_div(unique_kv_len, page_size) + unique_last_page_len = torch.full( + (batch_size,), + (unique_kv_len - 1) % page_size + 1, + dtype=torch.int32, + device="cuda", + ) + flashinfer.append_paged_kv_cache( + k_unique, + v_unique, + *flashinfer.get_batch_indices_positions( + unique_append_indptr, + flashinfer.get_seq_lens(unique_kv_indptr, unique_last_page_len, page_size), + k_unique.shape[0], + ), + kv_data, + unique_kv_indices, + unique_kv_indptr, + unique_last_page_len, + kv_layout, + ) + + def make_wrapper_and_plan(): + w = flashinfer.MultiLevelCascadeAttentionWrapper( + 2, + torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda"), + kv_layout, + ) + qo_indptr_top = torch.tensor([0, q.shape[0]], dtype=torch.int32, device="cuda") + if stage == "decode": + qo_indptr_bottom = torch.arange( + 0, batch_size + 1, dtype=torch.int32, device="cuda" + ) + else: + qo_indptr_bottom = ( + torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") + * unique_kv_len + ) + w.plan( + [qo_indptr_top, qo_indptr_bottom], + [shared_kv_indptr, unique_kv_indptr], + [shared_kv_indices, unique_kv_indices], + [shared_last_page_len, unique_last_page_len], + num_heads, + num_heads, + head_dim, + page_size, + q_data_type=dtype, + ) + return w + + # Unfused reference + orig_flag = cascade_mod._HIP_FUSED_CASCADE + cascade_mod._HIP_FUSED_CASCADE = False + try: + w_unfused = make_wrapper_and_plan() + out_unfused = w_unfused.run(q, kv_data) + finally: + cascade_mod._HIP_FUSED_CASCADE = orig_flag + + # Fused path + cascade_mod._HIP_FUSED_CASCADE = True + try: + w_fused = make_wrapper_and_plan() + out_fused = w_fused.run(q, kv_data) + finally: + cascade_mod._HIP_FUSED_CASCADE = orig_flag + + torch.testing.assert_close( + out_fused.float(), out_unfused.float(), rtol=1e-3, atol=atol + ) + + +if __name__ == "__main__": + test_batch_attention_with_shared_prefix_paged_kv_cache( + "decode", 12, 37, 128, 8, False, 128, 16 + ) + test_batch_attention_with_shared_prefix_paged_kv_cache( + "append", 12, 37, 128, 8, False, 128, 16 + ) + test_multilevel_cascade_fused_vs_unfused( + torch.float16, "append", 4, 17, 128, 8, 128, 16 + )