fix(cake_kda): support non-aligned recurrent prefill head counts - #4351
Conversation
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
📝 WalkthroughWalkthroughFlashKDA beta TMA padding now rounds head counts to multiples of eight. C++ validation and packing use the rounded width and full padded storage. Documentation and recurrent prefill tests cover non-aligned head counts. ChangesBeta TMA padding and packing
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant BetaTmaSource
participant BetaTmaValidation
participant BetaTmaLauncher
participant BetaPackingKernel
BetaTmaSource->>BetaTmaValidation: provide rounded beta TMA width
BetaTmaValidation->>BetaTmaLauncher: validate padded storage
BetaTmaLauncher->>BetaPackingKernel: pass padded dimensions and storage extent
BetaPackingKernel->>BetaTmaLauncher: pack beta TMA storage
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/kda/flashkda_binding_common.cuh`:
- Line 243: Update the beta TMA packing guards around beta_tma_heads and the
PackBetaForTmaKernel call to skip packing only when beta and beta_tma are
exactly aliased, applying the existing overlap exemptions only in that case. For
every separate beta_tma destination, pack its full numel rather than relying on
padded_num_heads. Add a regression case covering H=8 with fewer than 32 tokens.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: d41a495b-e56a-473b-832a-eadb1e7c31c0
📒 Files selected for processing (5)
csrc/kda/flashkda_binding_common.cuhdocs/api/kda_prefill.rstflashinfer/kda_prefill.pytests/jit/test_flash_kda_jit.pytests/kda/test_recurrent_kda_prefill.py
| "storage"; | ||
| CheckNoPartialOverlapOrExactAlias(beta, "beta", beta_tma, "beta_tma"); | ||
| if (num_heads < kBetaTmaMinHeads) { | ||
| if (beta_tma_heads != num_heads) { |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Use exact storage aliasing to decide whether to skip packing.
When H is divisible by eight and token_count < 32, flashinfer/kda_prefill.py allocates a separate [32, H] beta_tma workspace. padded_num_heads == num_heads is then true, but beta_tma does not contain beta values. Line 325 skips PackBetaForTmaKernel, and the frozen kernel reads uninitialized workspace values.
Use exact beta/beta_tma storage aliasing for this early return. Apply the overlap exemptions at Line 243 only for that exact alias. Pack the full beta_tma.numel() for every separate destination. Add a regression case with H=8 and fewer than 32 tokens.
Proposed fix
- if (beta_tma_heads != num_heads) {
+ const bool beta_tma_exact_alias =
+ beta.data_ptr() == beta_tma.data_ptr() && beta.numel() == beta_tma.numel();
+ if (!beta_tma_exact_alias) {
CheckNoOverlap(beta_tma, "beta_tma", q, "q");
// ...
}
- if (padded_num_heads == num_heads) {
+ const bool beta_tma_exact_alias =
+ beta.data_ptr() == beta_tma.data_ptr() && beta.numel() == beta_tma.numel();
+ if (beta_tma_exact_alias) {
return;
}Also applies to: 324-326
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@csrc/kda/flashkda_binding_common.cuh` at line 243, Update the beta TMA
packing guards around beta_tma_heads and the PackBetaForTmaKernel call to skip
packing only when beta and beta_tma are exactly aliased, applying the existing
overlap exemptions only in that case. For every separate beta_tma destination,
pack its full numel rather than relying on padded_num_heads. Add a regression
case covering H=8 with fewer than 32 tokens.
yzh119
left a comment
yzh119
left a comment
There was a problem hiding this comment.
Any changes in terms of performance？
yzh119
left a comment
yzh119
left a comment
There was a problem hiding this comment.
Any changes on performance？
|
/bot run tests/kda |
4 participants
Follow-up to merged #4262 and #4313.
What changed
round_up(H, 8)instead of only paddingH < 8;H, state shape, launch grid, and frozen M64/M128 CUDA bodies unchanged;The frozen kernels load beta in 8-head TMA boxes. For H=12, the original BF16 row stride is 24 bytes and
cuTensorMapEncodeTiledrejects it. Padding the descriptor source to 16 heads gives a 32-byte row stride; heads 12–15 are padding only and are never assigned CTAs.This generalizes the fix to every positive head count that is not divisible by eight, while aligned head counts retain the existing zero-copy beta path.
Correctness
Both runs used the public
flashinfer.recurrent_kdafacade and compared BF16 output plus the complete final state against the PyTorch reference withatol=rtol=1e-2.62 passed, 0 skipped, 0 failed62 passed, 0 skipped, 0 failedThe test gate runs:
H=12 coverage includes fixed T=32, fixed T=33 (one full TMA chunk plus the direct-load tail), packed sequence lengths
[32, 3], in-place initial/final state, and CUDA graph replay after beta is changed. The JIT contract test also verifies that the frozen generated M64/M128 bodies remain unchanged.pre-commit run --files <changed files>passes.Performance
The H=12 path was benchmarked through the public
flashinfer.recurrent_kdafacade with fallback forbidden. Speedup is the official FlashKDA raw GPU span divided by this PR's public-API GPU span. The baseline is the same official FlashKDA source used for #4262:MoonshotAI/FlashKDA@d2ff19a, with CUTLASS5c149f5.Measurements use strict CUPTI first-to-last correlated compute-kernel span, cold L2, no CUDA Graph, and two independent 128-sample blocks in symmetric ABCCBA order. The PR span includes both the beta pack and frozen M128 recurrence kernels. All six benchmark shapes passed output and complete-final-state correctness against the official peer with BF16
atol=rtol=1e-2.[512] x 32[128] x 8[512][8192][1300, 547, 2048, 963, 271, 3063][1024] x 8unsupported / N/AUpstream main fails before kernel launch with
cuTensorMapEncodeTiled failed for beta_tma with CUresult=1, so there is no valid upstream H=12 timing or speedup claim. This PR leaves the frozen compute kernel unchanged; it adds the required beta packing only for non-8-aligned head counts.Related to #4254.
Summary by CodeRabbit
Bug Fixes
Documentation
Tests