Conversation
Use one-warp for D128 single-token decode below 1,920 sequence-heads and grouped-CTA at/above that cutoff, based on Kimi K3 B300 sweeps. Retain the existing D64 minimum-grid policy and add backend-selection regression tests.
Use the same T=1 selection rule for both head dims: one-warp below a sequence-head cutoff, grouped-CTA at/above it. Set D128 to 1,920 and D64 to 7,680 from B300 crossover sweeps, replacing the old inverted D64 minimum-grid policy. Extend backend-selection tests for both cutoffs.
Replace per-dim constants and helper lookup with GROUPED_MIN_SEQUENCE_HEADS and use NUM_TOKENS consistently at both dispatch call sites. Table-drive the backend-selection regression tests.
Cover B∈{1,4,32} at H=96/D=128 with in-kernel gated decode and lower_bound=-5.0
using the existing vLLM-style cu_seqlens correctness test.
Exercise the vLLM logits path in test_vllm_decode for K3-shaped decode cases so beta sigmoid stays in-kernel rather than only in spec-decode tests.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughRecurrent KDA now resolves explicit or identity state indices for dense and packed decoding. Indexed execution updates the caller-owned state pool and returns compact standard-decode state. Dispatch adds architecture- and workload-specific kernel schedules. Tests and trace callers cover the new contracts. ChangesRecurrent KDA state indexing and dispatch
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant recurrent_kda
participant run_recurrent_kda
participant recurrent_kda_decode_kernel
participant StatePool
Caller->>recurrent_kda: provide decode inputs and optional ssm_state_indices
recurrent_kda->>run_recurrent_kda: pass state indices and decode inputs
run_recurrent_kda->>run_recurrent_kda: select one-warp or grouped schedule
run_recurrent_kda->>recurrent_kda_decode_kernel: launch with resolved state indices
recurrent_kda_decode_kernel->>StatePool: update indexed state slots
run_recurrent_kda->>StatePool: return full pool or state[ssi]
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
|
/bot run tests/kda |
|
[SUCCESS] Pipeline #59980611: 18/18 executed test jobs passed |
djmmoss
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/kda_kernels/recurrent_kda.py (1)
1409-1449: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winNo validation that standard-decode
ssm_state_indicesare non-negative before indexing an uninitializedout_buf.Docs reserve
-1for packed/spec rows only (not the plain[B]standard-decode format), but nothing here enforces that. If a negative index slips through:out_bufistorch.empty(...)-allocated (line 1449, never zeroed in this branch), and the one-warp kernel'sis_activegate (lines 433-435) skips writing that row's output entirely — leaving uninitialized memory in the result. The grouped-CTA kernel's equivalent path (_grouped_kda_kernel, lines 689-692, same file) explicitly zero-fills the analogous case, so which failure mode occurs is dependent on the_use_one_warpcutoff this PR changes.🛡️ Suggested guard
grid_seqs = B sequence_heads = grid_seqs * HV use_one_warp = _use_one_warp(K, NUM_TOKENS, sequence_heads) cu_seqlens_i32 = None + if (ssm_state_indices < 0).any(): + raise ValueError( + "ssm_state_indices must be non-negative for standard decode " + "(no cu_seqlens); -1 is reserved for padded packed rows" + )🤖 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 `@flashinfer/kda_kernels/recurrent_kda.py` around lines 1409 - 1449, Validate standard-decode ssm_state_indices before allocating or using state/output buffers: reject any negative entry because this [B]-shaped path does not permit the reserved -1 sentinel. Add the guard in the non-cu_seqlens branch alongside the existing shape validation, using the existing ssm_state_indices and B symbols, and raise a clear ValueError before kernel execution.
🧹 Nitpick comments (1)
flashinfer/kda_kernels/recurrent_kda.py (1)
1240-1244: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring calls the returned indexed state a "view"; it's actually a gather copy.
state[ssi](fancy/advanced indexing, see line 1643) allocates a new tensor rather than aliasingstate's storage. Calling it a "view" here could mislead callers into assuming in-place edits to the returned tensor propagate back into the caller-owned pool, which they do not.📝 Suggested wording fix
- - state: Updated state if ``output_final_state=True``, else - ``None``. Indexed standard decode returns the compact - ``[B, HV, V, K]`` view while updating the full caller-owned pool - in place. For batched spec decode without ``cu_seqlens``, this is - the packed checkpoint state pool used by the shim. + - state: Updated state if ``output_final_state=True``, else + ``None``. Indexed standard decode returns a compact + ``[B, HV, V, K]`` gather (a new tensor, not a view) while the + full caller-owned pool is updated in place. For batched spec + decode without ``cu_seqlens``, this is the packed checkpoint + state pool used by the shim.🤖 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 `@flashinfer/kda_kernels/recurrent_kda.py` around lines 1240 - 1244, Update the return-state documentation for the indexed standard decode in the relevant docstring to describe the compact [B, HV, V, K] result as a gathered copy rather than a view, and clarify that edits to it do not modify the caller-owned state pool. Leave the behavior and other state descriptions unchanged.
🤖 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.
Outside diff comments:
In `@flashinfer/kda_kernels/recurrent_kda.py`:
- Around line 1409-1449: Validate standard-decode ssm_state_indices before
allocating or using state/output buffers: reject any negative entry because this
[B]-shaped path does not permit the reserved -1 sentinel. Add the guard in the
non-cu_seqlens branch alongside the existing shape validation, using the
existing ssm_state_indices and B symbols, and raise a clear ValueError before
kernel execution.
---
Nitpick comments:
In `@flashinfer/kda_kernels/recurrent_kda.py`:
- Around line 1240-1244: Update the return-state documentation for the indexed
standard decode in the relevant docstring to describe the compact [B, HV, V, K]
result as a gathered copy rather than a view, and clarify that edits to it do
not modify the caller-owned state pool. Leave the behavior and other state
descriptions unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 17137426-641e-40d1-9dc5-8a75987642f9
📒 Files selected for processing (6)
flashinfer/kda_decode.pyflashinfer/kda_kernels/recurrent_kda.pyflashinfer/trace/templates/kda.pytests/kda/test_recurrent_kda.pytests/trace/example.pytests/trace/test_fi_trace.py
|
/bot run tests/kda |
|
[SUCCESS] Pipeline #60413256: 18/18 executed test jobs passed |
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 `@flashinfer/kda_kernels/recurrent_kda.py`:
- Around line 31-33: Update the dispatch documentation near _use_one_warp to
state that multi-token workloads generally use grouped-CTA, except SM100
workloads where num_tokens == 3 with gating or num_tokens == 4 without gating
and without GQA use the one-warp kernel.
🪄 Autofix (Beta)
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: 66165540-5493-4921-b175-4fd02839b7ad
📒 Files selected for processing (5)
flashinfer/kda_decode.pyflashinfer/kda_kernels/recurrent_kda.pyflashinfer/trace/templates/kda.pytests/kda/test_recurrent_kda.pytests/trace/example.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/trace/example.py
- flashinfer/trace/templates/kda.py
|
/bot run tests/kda |
|
[SUCCESS] Pipeline #60944357: 18/18 executed test jobs passed |
| elif ssm_state_indices is not None: | ||
| state = initial_state[ssm_state_indices].contiguous() | ||
| copy_back_indices = ssm_state_indices | ||
| max_idx = int(ssi.max().item()) + 1 |
There was a problem hiding this comment.
This will leave a D2H sync everytime .item() is called + can break cuda graph capture. You can add
if ssm_state_indices is None:
state = torch.zeros(B, HV, V, K, device=device, dtype=torch.bfloat16)
else:
max_idx = ...
| ) -> bool: | ||
| """Select the measured kernel architecture for the active GPU.""" | ||
| if compute_capability not in TUNED_DISPATCH_COMPUTE_CAPABILITIES: | ||
| return num_tokens == 1 and sequence_heads < GROUPED_MIN_SEQUENCE_HEADS[head_dim] |
There was a problem hiding this comment.
I believe with this change it is now:
| sequence_heads (D=128, T=1) | pre-PR | PR fallback |
|---|---|---|
| < 128 | grouped | one-warp |
| 128 … 1919 | one-warp | one-warp |
>= 1920 |
one-warp | grouped |
Is this expected?
Signed-off-by: Duncan Moss <djm.moss@gmail.com>
Signed-off-by: Duncan Moss <djm.moss@gmail.com>
|
[SUCCESS] Pipeline #60944357: 18/18 executed test jobs passed |
Pull request was closed
4 participants
Summary
removing the full-state gather and scatter around standard decode. Dense calls
reuse the same path through a cached identity mapping.
architecture and its row/reduction or key/value split schedule from the workload.
boundaries.
B200 performance
BF16 cold-L2 CUPTI timing with CUDA Graphs disabled; each row is the median of five
independent rounds. This compares upstream pre-PR
recurrent_kdaat76c583655against this PR at
6039adef. Both columns use the CuTe DSL backend.The current implementation passed output and full-state correctness across all 95
shapes.