From c91f8d93a49b56db701f1cd0cfb1a41535ce26dc Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Tue, 4 Aug 2026 16:08:32 -0700 Subject: [PATCH] fix(cake_kda): support non-aligned recurrent prefill heads Pad beta TMA rows to the next eight-head boundary and teach the binding pack kernel to use the dynamic padded stride. Add H=12 eager, packed, full-plus-tail, and CUDA graph correctness coverage.\n\nAI-assisted-by: OpenAI Codex --- csrc/kda/flashkda_binding_common.cuh | 40 ++++++----- docs/api/kda_prefill.rst | 6 +- flashinfer/kda_prefill.py | 8 ++- tests/jit/test_flash_kda_jit.py | 4 ++ tests/kda/test_recurrent_kda_prefill.py | 95 +++++++++++++++++++++++-- 5 files changed, 125 insertions(+), 28 deletions(-) diff --git a/csrc/kda/flashkda_binding_common.cuh b/csrc/kda/flashkda_binding_common.cuh index 94254d22629..bb9b82b5ccb 100644 --- a/csrc/kda/flashkda_binding_common.cuh +++ b/csrc/kda/flashkda_binding_common.cuh @@ -40,18 +40,23 @@ constexpr size_t kTensorMapCount = 6; constexpr size_t kTensorMapAlignment = 64; static_assert(sizeof(CUtensorMap) == 128); constexpr size_t kDescriptorStorageBytes = kTensorMapCount * sizeof(CUtensorMap); -constexpr int64_t kBetaTmaMinHeads = 8; +constexpr int64_t kBetaTmaHeadsPerBox = 8; + +inline int64_t RoundUpBetaTmaHeads(int64_t num_heads) { + return (num_heads / kBetaTmaHeadsPerBox + + static_cast(num_heads % kBetaTmaHeadsPerBox != 0)) * + kBetaTmaHeadsPerBox; +} static __global__ void PackBetaForTmaKernel(const __nv_bfloat16* beta, __nv_bfloat16* beta_tma, - int64_t token_count, int64_t padded_token_count, - int32_t num_heads) { + int64_t token_count, int64_t padded_elements, + int64_t num_heads, int64_t padded_num_heads) { const int64_t linear_index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t padded_elements = padded_token_count * kBetaTmaMinHeads; if (linear_index >= padded_elements) { return; } - const int64_t token_index = linear_index / kBetaTmaMinHeads; - const int32_t head_index = static_cast(linear_index % kBetaTmaMinHeads); + const int64_t token_index = linear_index / padded_num_heads; + const int64_t head_index = linear_index % padded_num_heads; __nv_bfloat16 value = __float2bfloat16(0.0f); if (token_index < token_count && head_index < num_heads) { value = beta[token_index * num_heads + head_index]; @@ -228,14 +233,14 @@ inline int64_t CheckCommonInputs(const TensorView& q, const TensorView& k, const TVM_FFI_ICHECK(beta.ndim() >= 2 && beta.size(beta.ndim() - 1) == num_heads && beta.numel() == token_count * num_heads) << "beta must match flattened [tokens, H]"; - const int64_t beta_tma_heads = std::max(num_heads, 8); + const int64_t beta_tma_heads = RoundUpBetaTmaHeads(num_heads); TVM_FFI_ICHECK(beta_tma.ndim() >= 2 && beta_tma.size(beta_tma.ndim() - 1) == beta_tma_heads && beta_tma.numel() % beta_tma_heads == 0 && beta_tma.numel() / beta_tma_heads >= std::max(token_count, 32)) - << "beta_tma must have at least [max(tokens, 32), max(H, 8)] " + << "beta_tma must have at least [max(tokens, 32), round_up(H, 8)] " "storage"; CheckNoPartialOverlapOrExactAlias(beta, "beta", beta_tma, "beta_tma"); - if (num_heads < kBetaTmaMinHeads) { + if (beta_tma_heads != num_heads) { CheckNoOverlap(beta_tma, "beta_tma", q, "q"); CheckNoOverlap(beta_tma, "beta_tma", k, "k"); CheckNoOverlap(beta_tma, "beta_tma", v, "v"); @@ -314,23 +319,22 @@ inline int64_t CheckCommonInputs(const TensorView& q, const TensorView& k, const inline void PackBetaForTmaIfNeeded(const TensorView& beta, const TensorView& beta_tma, int64_t num_heads, cudaStream_t stream) { - // Full chunks TMA-load an eight-head beta box. Only H<8 requires a - // materialized row-padded source; H>=8 aliases beta whenever a full chunk - // exists, while shorter inputs stay entirely on the direct-load tail path. - if (num_heads >= kBetaTmaMinHeads) { + // Full chunks TMA-load an eight-head beta box, so any partial final group + // needs a materialized row padded to the next eight-head boundary. + const int64_t padded_num_heads = beta_tma.size(beta_tma.ndim() - 1); + if (padded_num_heads == num_heads) { return; } const int64_t token_count = beta.numel() / num_heads; - const int64_t padded_token_count = beta_tma.numel() / kBetaTmaMinHeads; - const int64_t padded_elements = padded_token_count * kBetaTmaMinHeads; + const int64_t padded_elements = beta_tma.numel(); constexpr int32_t kThreads = 256; - const int64_t blocks_i64 = (padded_elements + kThreads - 1) / kThreads; + const int64_t blocks_i64 = (padded_elements - 1) / kThreads + 1; TVM_FFI_ICHECK(blocks_i64 > 0 && blocks_i64 <= std::numeric_limits::max()) << "beta TMA pack grid.x is out of range: " << blocks_i64; PackBetaForTmaKernel<<(blocks_i64), kThreads, 0, stream>>>( reinterpret_cast(beta.data_ptr()), - reinterpret_cast<__nv_bfloat16*>(beta_tma.data_ptr()), token_count, padded_token_count, - static_cast(num_heads)); + reinterpret_cast<__nv_bfloat16*>(beta_tma.data_ptr()), token_count, padded_elements, + num_heads, padded_num_heads); CheckCuda(cudaGetLastError(), "PackBetaForTmaKernel launch"); } diff --git a/docs/api/kda_prefill.rst b/docs/api/kda_prefill.rst index a36194a0858..78234ac601e 100644 --- a/docs/api/kda_prefill.rst +++ b/docs/api/kda_prefill.rst @@ -108,5 +108,7 @@ stream. When an explicit workspace is used with ``initial_state=None`` and ``output_final_state=True``, the returned final state is workspace-owned stable scratch. Otherwise an explicitly supplied ``initial_state`` is updated -directly in place by the frozen kernel. The small-head ``H < 8`` path captures -the beta copy into workspace-owned padded storage before the frozen launch. +directly in place by the frozen kernel. Head counts that are not divisible by +eight capture the beta copy into workspace-owned storage padded to the next +eight-head boundary before the frozen launch. The public beta and state shapes +keep the caller's original head count. diff --git a/flashinfer/kda_prefill.py b/flashinfer/kda_prefill.py index 719560dfda5..1e5b69f6d99 100644 --- a/flashinfer/kda_prefill.py +++ b/flashinfer/kda_prefill.py @@ -35,7 +35,7 @@ from .jit.flash_kda import FlashKDATarget, FlashKDAVariant _FLASH_KDA_HEAD_DIM = 128 -_FLASH_KDA_BETA_TMA_MIN_HEADS = 8 +_FLASH_KDA_BETA_TMA_HEADS_PER_BOX = 8 _FLASH_KDA_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} _FLASH_KDA_DESCRIPTOR_STORAGE_BYTES = 6 * 128 _flash_kda_tensor_cache: dict[tuple, torch.Tensor] = {} @@ -381,7 +381,11 @@ def _beta_tma_source( total_tokens = batch_size * seq_len beta_flat = beta.reshape(total_tokens, num_heads) padded_tokens = max(total_tokens, 32) - padded_heads = max(num_heads, _FLASH_KDA_BETA_TMA_MIN_HEADS) + padded_heads = ( + (num_heads + _FLASH_KDA_BETA_TMA_HEADS_PER_BOX - 1) + // _FLASH_KDA_BETA_TMA_HEADS_PER_BOX + * _FLASH_KDA_BETA_TMA_HEADS_PER_BOX + ) if padded_tokens == total_tokens and padded_heads == num_heads: return beta_flat shape = (padded_tokens, padded_heads) diff --git a/tests/jit/test_flash_kda_jit.py b/tests/jit/test_flash_kda_jit.py index 01b70770698..0b0f07b7be5 100644 --- a/tests/jit/test_flash_kda_jit.py +++ b/tests/jit/test_flash_kda_jit.py @@ -190,6 +190,10 @@ def test_flash_kda_descriptor_workspace_contract(): assert "major == 10 && (minor == 0 || minor == 3)" in common_text assert "CheckFlashKDATarget" in common_text assert "PackBetaForTmaKernel" in common_text + assert "RoundUpBetaTmaHeads(num_heads)" in common_text + assert "padded_num_heads == num_heads" in common_text + assert "linear_index / padded_num_heads" in common_text + assert "linear_index % padded_num_heads" in common_text assert ( 'CheckNoPartialOverlapOrExactAlias(beta, "beta", beta_tma, "beta_tma")' in common_text diff --git a/tests/kda/test_recurrent_kda_prefill.py b/tests/kda/test_recurrent_kda_prefill.py index 82f123a7460..76ff26b1268 100644 --- a/tests/kda/test_recurrent_kda_prefill.py +++ b/tests/kda/test_recurrent_kda_prefill.py @@ -280,7 +280,12 @@ def test_multi_token_gqa_stays_on_existing_backend(cuda_device, monkeypatch): @pytest.mark.parametrize( ("packed", "num_heads", "expected_variant"), - [(False, 64, "m64"), (True, 64, "m128"), (True, 2, "m128")], + [ + (False, 64, "m64"), + (True, 64, "m128"), + (True, 2, "m128"), + (False, 12, "m128"), + ], ) @pytest.mark.parametrize( ("compute_capability", "expected_target"), @@ -339,7 +344,7 @@ def get_module(variant, target): assert args[4].data_ptr() == inputs["beta"].data_ptr() assert args[5].shape == ( max(inputs["q"].numel() // (num_heads * 128), 32), - max(num_heads, 8), + (num_heads + 7) // 8 * 8, ) assert args[8].dtype == torch.int64 assert args[9].dtype == torch.int32 @@ -353,7 +358,7 @@ def get_module(variant, target): assert math.isclose(args[18], 128**-0.5) assert args[19] == -5.0 assert args[20] == int(torch.cuda.current_stream(cuda_device).cuda_stream) - if num_heads < 8: + if num_heads % 8 != 0: assert args[5].data_ptr() != inputs["beta"].data_ptr() @@ -776,6 +781,81 @@ def test_frozen_prefill_h6_full_tma_chunk_matches_reference(flash_kda_device): ) +@pytest.mark.parametrize("seq_len", [32, 33]) +def test_frozen_prefill_h12_tma_chunks_match_reference(flash_kda_device, seq_len): + inputs = _make_inputs( + seq_lens=[seq_len], + num_heads=12, + packed=False, + initial_state=True, + seed=2012 + seq_len, + ) + reference_inputs = { + **inputs, + "initial_state": inputs["initial_state"].clone(), + } + expected_output, expected_state = _reference(reference_inputs) + output = torch.empty_like(inputs["q"]) + + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=True, + ) + + assert actual_output.data_ptr() == output.data_ptr() + assert actual_state is inputs["initial_state"] + torch.testing.assert_close( + actual_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + actual_state.float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) + + +def test_frozen_prefill_h12_packed_matches_reference(flash_kda_device): + inputs = _make_inputs( + seq_lens=[32, 3], + num_heads=12, + packed=True, + initial_state=True, + seed=2047, + ) + reference_inputs = { + **inputs, + "initial_state": inputs["initial_state"].clone(), + } + expected_output, expected_state = _reference(reference_inputs) + output = torch.empty_like(inputs["q"]) + + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=True, + ) + + assert actual_output.data_ptr() == output.data_ptr() + assert actual_state is inputs["initial_state"] + torch.testing.assert_close( + actual_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + actual_state.float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) + + def test_frozen_prefill_m64_matches_reference(flash_kda_device): inputs = _make_inputs( seq_lens=[2], @@ -898,13 +978,16 @@ def test_frozen_prefill_cuda_graph_capture_and_replay( ) -def test_frozen_prefill_h6_full_chunk_graph_refreshes_beta(flash_kda_device): +@pytest.mark.parametrize("num_heads", [6, 12]) +def test_frozen_prefill_non_aligned_heads_graph_refreshes_beta( + flash_kda_device, num_heads +): inputs = _make_inputs( seq_lens=[32], - num_heads=6, + num_heads=num_heads, packed=False, initial_state=True, - seed=2033, + seed=2033 + num_heads, ) initial_state_seed = inputs["initial_state"].clone() output = torch.empty_like(inputs["q"])