fix(cake_kda): support non-aligned recurrent prefill head counts

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<int64_t>(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<int64_t>(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<int32_t>(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<int64_t>(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<int64_t>(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<uint32_t>::max())
       << "beta TMA pack grid.x is out of range: " << blocks_i64;
   PackBetaForTmaKernel<<<static_cast<uint32_t>(blocks_i64), kThreads, 0, stream>>>(
       reinterpret_cast<const __nv_bfloat16*>(beta.data_ptr()),
-      reinterpret_cast<__nv_bfloat16*>(beta_tma.data_ptr()), token_count, padded_token_count,
-      static_cast<int32_t>(num_heads));
+      reinterpret_cast<__nv_bfloat16*>(beta_tma.data_ptr()), token_count, padded_elements,
+      num_heads, padded_num_heads);
   CheckCuda(cudaGetLastError(), "PackBetaForTmaKernel launch");
 }
 

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.

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)

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

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"])