feat(kda): add packed-input CuTe decode kernel

[ameynaik-hub]

Packed-input CuTe-DSL KDA T=1 decode kernel for B200

Adds flashinfer.kda_kernels.packed_kda_decode_cute.run_packed_kda_decode_cute, a CuTe-DSL decode kernel for serving-native packed KDA (Kimi Delta Attention) T=1 inputs: packed bf16 QKV rows [B, 3*12*128], raw gate/beta logits, fp32 internal math, bf16 state pool [N, 12, 128, 128] updated in place, state_indices rows outside the pool (negative or past the end) produce zero output and leave the pool untouched. CUDA-graph capture/replay safe; runs on the caller's current stream via TVM-FFI.

Kernel design

A pipelined implementation tuned per batch size:

• cp.async shared-memory ring (B≥24): 128-thread CTAs stream the bf16 state through a 4–5-slot ring of 4KB chunks (cp.async.cg, L1 bypass) with a barrier-free per-thread pipeline and LDS double-buffering, so in-flight read volume is not bounded by the register file.
• Fused output projection: o = (h·d)·q + vn·(k·q) with k·q computed once at staging — removes one butterfly reduction tree per row.
• bf16 register economy: state stays packed bf16 in registers; unpack is a shift/mask bit trick and repack a single cvt.rn.bf16x2.f32 (full-rate ALU instead of the conversion pipe), feeding packed f32x2 FMAs.
• Per-batch policy (_select_config): register-prefetch kernel with 128-thread CTAs (B≤11) or 32-thread CTAs (B∈[12,23]); cp.async kernel with half-head tiles (B∈[24,37]) or whole-head tiles (B≥38). tile_v= forcing maps onto tuned schedules for sweeps.

Benchmark results

B200, driver 595.58.03, CUDA 13.0, benchmarks/bench_packed_kda_decode.py (CUPTI kernel timing, warmup 100, iterations 100, default clocks). "Previous" is the earlier single-warp variant of this PR at the same protocol; medians in µs:

B	direct (prev → this)	cuda_graph (prev → this)	speedup (direct)
1	2.75 → 2.37	2.50 → 2.11	1.16×
8	3.46 → 3.20	3.23 → 2.98	1.08×
16	4.53 → 4.38	4.42 → 4.29	1.03×
31	5.89 → 5.89	5.70 → 5.76	1.00×
32	6.00 → 6.06	5.82 → 5.89	0.99× (tie)
64	9.47 → 9.09	9.25 → 8.90	1.04×
128	15.55 → 14.32	15.26 → 14.05	1.09×
256	35.14 → 31.87	35.26 → 32.10	1.10×
512	68.16 → 63.82	68.13 → 63.89	1.07×

No regression at any batch size (B=31/32 within run-to-run noise). Correctness at every batch is checked against an fp32 torch reference before timing (--refcheck equivalent is built into the script); max output error is at the bf16 quantization level.

How to run the benchmark

# full sweep, both launch modes
python benchmarks/bench_packed_kda_decode.py --warmup 100 --iterations 100

# specific batches / direct launches only / JSON report
python benchmarks/bench_packed_kda_decode.py \
    --batch-size 1 32 512 --mode direct \
    --warmup 100 --iterations 100 --json results.json

# force a tile schedule (benchmark override; default lets the policy pick)
python benchmarks/bench_packed_kda_decode.py --tile-v 64

# cold-L2 timing
python benchmarks/bench_packed_kda_decode.py --cold-l2

Requires a B200 (exact CC 10.0) and CUDA ≥12.8; pip install -U cupti-python for CUPTI timing (falls back to CUDA events).

Tests:

pytest tests/kda/test_packed_kda_decode_cute.py -m "" -v

36 tests: reference match B=1–512, all forced tiles, sanitizer (odd-stride) schedules, shifted/misaligned tensors for every argument, out-of-range slots, all-inactive bitwise no-op, CUDA-graph replay with changed inputs and indices, current-stream semantics, and a 512-step fp64 drift diagnostic.

T=1 fast path inside recurrent_kda (no new public API)

The kernel takes q/k/v/g/beta as five independent strided tensors — the packed layout only ever existed in the launcher. A new launch_unpacked_kda_decode_cute entry point exposes this, and run_recurrent_kda now routes eligible T=1 decode calls to it: T=1, H=HV=12, K=V=128, bf16, SM100a, use_gate_in_kernel=True with lower_bound=-5, beta_is_logit=True, in-kernel QK L2 norm, 1-D ssm_state_indices, no spec/varlen/GQA/final-state. Ineligible calls use the existing kernels unchanged. Toggle: FLASHINFER_KDA_T1_FAST_PATH=0 disables (default enabled).

The fast path indexes the state pool in-kernel (replacing the host-side slot gather/scatter), and additionally accepts padded state pools and strided q/k/v/g views (zero-copy from a fused-QKV projection GEMM). Measured through the public recurrent_kda API with a state cache (same protocol):

B	fast path off	fast path on	speedup
8	66.0µs	4.2µs	15.9×
64	88.5µs	11.1µs	8.0×
512	638.0µs	63.7µs	10.0×

Unpacked dispatch is bitwise-identical to the packed entry point on equal inputs (same cubin, different base pointers).

Both recurrent_kda calling conventions are covered: the raw form (use_gate_in_kernel=True, lower_bound=-5, beta_is_logit=True — the kernel computes the decay and sigmoid in fp32) and the API-default pre-computed form (log-space g, pre-sigmoided beta), which compiles a kernel variant with decay = exp(g). Pre-computed mode measures 3.9µs (B=8) / 63.9µs (B=512) on the fast path.

For attribution (per-kernel profile at B=512): the generic path's decode kernel itself takes 81.1µs (so kernel-vs-kernel the gain is 1.27×, matching its no-indices benchmark of 83.4µs); the remainder of its 638µs is the wrapper's state round-trip — a 69µs vectorized gather plus a 463µs elementwise index_copy_ scatter running at 0.87 TB/s. The scatter is independently fixable in the wrapper, but even with a perfect scatter the gather/scatter design pays two extra passes over the state (~126µs at B=512); the fast path's in-kernel indexing removes those passes entirely.

DRAM utilization

The kernel is DRAM-bound at large batch: per step it must read and rewrite the full B×12×128×128 bf16 state (786KB/row round trip; ~402MB at B=512) plus ~15KB/row of activations. From the benchmark's logical_TB/s column (compulsory bytes ÷ median time), against the 8 TB/s HBM3e peak:

B	effective DRAM (prev)	effective DRAM (this PR)
64	5.42 TB/s (68%)	5.65 TB/s (71%)
128	6.60 TB/s (82%)*	7.17 TB/s (90%)*
256	5.84 TB/s (73%)	6.44 TB/s (81%)
512	6.02 TB/s (75%)	6.43 TB/s (80%)

* B=128 exceeds the cold-memory ceiling because this benchmark's state pool (~50MB) partially survives in the 126MB L2 across iterations; B≥256 is the true DRAM-resident regime.

Nsight Compute confirms the gain is pure memory-pipeline efficiency, not traffic: at B=512 both kernels move identical bytes (dram__bytes_read.sum 207.7MB, in-kernel writes ~140MB with the remainder draining from L2 after kernel end), while this variant sustains ~8% higher dram__throughput.avg.pct_of_peak_sustained_elapsed (47.7% vs 44.1% under locked profiling clocks). At ~80% of peak on interleaved cold-L2 read/write, the kernel is at the observed machine ceiling for this access pattern — remaining headroom would require reducing bytes (e.g., fp8 state), not scheduling.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds a Blackwell-only packed KDA T=1 CuTe decoder with multiple kernel schedules, public validation and exports, extensive CUDA correctness tests, and a configurable direct or CUDA Graph benchmark with JSON reporting.

Changes

Packed KDA Decode

Layer / File(s)	Summary
Packed decode kernels flashinfer/kda_kernels/packed_kda_decode_cute.py	Adds register-prefetch, shared-memory, and persistent kernels with packed input handling, recurrent state updates, output generation, and asynchronous data movement.
Compilation, dispatch, and public validation flashinfer/kda_kernels/packed_kda_decode_cute.py, flashinfer/kda_kernels/__init__.py	Adds cached CuTe compilation, tile and stride selection, B200 validation, public runners, and guarded package exports.
Reference-based correctness coverage tests/kda/test_packed_kda_decode_cute.py	Adds deterministic fixtures, reference computation, tile checks, and output/state comparisons.
Layout, lifecycle, and numerical validation tests/kda/test_packed_kda_decode_cute.py	Tests shifted addresses, inactive rows, CUDA Graph replay, non-default streams, changed inputs, and long-run FP64 accuracy.
Benchmark generation and reporting benchmarks/bench_packed_kda_decode.py	Adds randomized cases, correctness checks, direct and graph timing, CLI validation, throughput reporting, and optional JSON output.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Runner as run_packed_kda_decode_cute
  participant Dispatch as launch_packed_kda_decode_cute
  participant Kernel as CuTe decode kernel
  participant State as Indexed recurrent state
  Runner->>Dispatch: Validate inputs and select schedule
  Dispatch->>Kernel: Launch on current CUDA stream
  Kernel->>State: Update valid state rows
  Kernel-->>Runner: Return output tensor

Loading

Possibly related PRs

• flashinfer-ai/flashinfer#4262: Adds a related B200/SM100a BF16 recurrent KDA implementation for optimized multi-token prefill.
• flashinfer-ai/flashinfer#4378: Extends the packed KDA decode backend and benchmark area with a serving-native implementation.
• flashinfer-ai/flashinfer#4445: Adds an alternative CuTe implementation and tests for the same packed KDA T=1 decode path.

Suggested reviewers: sricketts, aleozlx, yzh119

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 32.35% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely identifies the main change: adding a packed-input CuTe KDA decode kernel.
Description check	✅ Passed	The description is detailed and covers the kernel design, behavior, tests, benchmarks, usage, and performance objectives.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (7)

flashinfer/kda_kernels/packed_kda_decode_cute.py (6)

846-857: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider supported_compute_capability instead of the hand-rolled _check_b200.

_check_b200 reimplements architecture gating that the repository already provides as a decorator. This API targets a single backend and a single compute capability, so supported_compute_capability([100]) is the idiomatic mechanism. The CUDA 12.8 minimum-version test at Lines 853-857 stays as an explicit check.

_check_b200 also runs at Line 939, after all shape validation and after output is allocated at Line 931. On unsupported hardware the wrapper allocates a buffer and then raises. Move the hardware gate before the allocation, or let the decorator handle it.

Based on learnings: for single-backend, architecture-gated APIs that exclusively target a specific compute capability, prefer supported_compute_capability([...]) over backend_requirement.

Also applies to: 939-942

🤖 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/packed_kda_decode_cute.py` around lines 846 - 857,
Replace the hand-rolled _check_b200 architecture validation with the
supported_compute_capability([100]) decorator on the targeted API, while keeping
the explicit CUDA 12.8 version check. Remove the redundant _check_b200 calls,
ensuring capability gating occurs before output allocation in the function
around the existing API definition.

Source: Learnings

---

427-428: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Add a static assertion for the CTA-kernel tile range.

iterations_per_group evaluates to 0 when tile_v is 8, 16, or 32. The kernel then writes no output and silently leaves output uninitialized. Today only _get_compiled_kernel at Line 820 prevents that, through the tile_v in (8, 16, 32) dispatch tuple. Add a compile-time assertion here so a future change to the dispatch tuple fails at compile time instead of producing garbage output.

🛡️ Proposed guard

     rows_per_group: cutlass.Constexpr = tile_v // _NUM_GROUPS
     iterations_per_group: cutlass.Constexpr = rows_per_group // _ILP_ROWS
+    # The CTA kernel only covers tiles that give every group at least one
+    # full _ILP_ROWS iteration; smaller tiles route to the warp kernel.
+    assert iterations_per_group * _ILP_ROWS * _NUM_GROUPS == tile_v

🤖 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/packed_kda_decode_cute.py` around lines 427 - 428, Add
a compile-time assertion immediately after iterations_per_group in the CTA
kernel setup, requiring it to be greater than zero (equivalently, rejecting
tile_v values that produce zero iterations). Keep the existing
_get_compiled_kernel dispatch unchanged while ensuring unsupported tile ranges
fail during compilation before output is written.

---

700-713: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Apply cute.assume to the mixed and gate row strides here too.

_packed_kda_decode_warp_launch narrows mixed_qkv.stride[0] and raw_gate.stride[0] with cute.assume(..., divby=_ELEMS_PER_LANE) at Lines 615-619. This launch omits that step and only narrows the state stride. The Python wrapper already verifies both strides at Lines 946-948, so the same fact is available. Without the assumption, cute.autovec_copy at Lines 460-465 may fall back to scalar loads because the row base offset has unknown alignment.

⚡ Proposed change

     batch = state_indices.shape[0]
     num_v_tiles: cutlass.Constexpr = _HEAD_DIM // tile_v
 
+    mixed_stride = mixed_qkv.stride[0]
+    gate_stride = raw_gate.stride[0]
+    if cutlass.const_expr(use_aligned_io):
+        # The wrapper verifies both row strides; expose the fact so the
+        # per-lane slices lower to vector transactions.
+        mixed_stride = cute.assume(mixed_stride, divby=_ELEMS_PER_LANE)
+        gate_stride = cute.assume(gate_stride, divby=_ELEMS_PER_LANE)
     mixed_layout = cute.make_tensor(
         mixed_qkv.iterator,
         cute.make_layout(
             (_ELEMS_PER_LANE, _LANES_PER_ROW, 3, _HEADS, batch),
-            stride=(1, _ELEMS_PER_LANE, _GATE_WIDTH, _HEAD_DIM, mixed_qkv.stride[0]),
+            stride=(1, _ELEMS_PER_LANE, _GATE_WIDTH, _HEAD_DIM, mixed_stride),
         ),
     )
     gate_layout = cute.make_tensor(
         raw_gate.iterator,
         cute.make_layout(
             (_ELEMS_PER_LANE, _LANES_PER_ROW, _HEADS, batch),
-            stride=(1, _ELEMS_PER_LANE, _HEAD_DIM, raw_gate.stride[0]),
+            stride=(1, _ELEMS_PER_LANE, _HEAD_DIM, gate_stride),
         ),
     )

🤖 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/packed_kda_decode_cute.py` around lines 700 - 713,
Apply cute.assume with divby=_ELEMS_PER_LANE to mixed_qkv.stride[0] and
raw_gate.stride[0] before constructing mixed_layout and gate_layout, matching
_packed_kda_decode_warp_launch. Use the narrowed stride values in both layout
stride tuples so cute.autovec_copy can rely on the verified row alignment.

---

154-163: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Extract the duplicated state-row load into a helper.

The same load pattern appears three times in this kernel: the prologue at Lines 154-163, the loop body at Lines 240-247, and the CTA kernel at Lines 531-538. Each copy repeats the local_tile construction and the aligned/unaligned branch. A small Python helper that the DSL inlines keeps the vectorization behavior and removes the duplication.

Also applies to: 232-247

🤖 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/packed_kda_decode_cute.py` around lines 154 - 163,
Extract the repeated state-row loading logic into a shared helper, including
cute.local_tile construction and the use_aligned_io-dependent autovec_copy or
element-wise copy. Replace the implementations in the prologue, loop body, and
CTA kernel while preserving their existing indices, buffers, and vectorization
behavior.

---

61-63: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Document the batch threshold rationale.

The value 26 is a tuned crossover point on the dispatch hot path. Add a short comment that records the measurement basis and the alternative tile choice, so future tuning does not need re-derivation.

As per coding guidelines: "For performance-critical hot paths, document the rationale for special algorithmic choices and potential alternatives in a code comment."

🤖 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/packed_kda_decode_cute.py` around lines 61 - 63, Add a
concise comment above _select_tile_v documenting that the batch threshold of 26
is a measured performance crossover on the dispatch hot path, and state the
alternative tile choices (8 below the threshold, 16 at or above it). Preserve
the existing selection logic.

Source: Coding guidelines

---

860-872: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add @flashinfer_api(trace=...) to the entrypoint.

run_packed_kda_decode_cute is now a public KDA-kernel entrypoint exported from flashinfer.kda_kernels, but it only has @torch.no_grad(). Use @flashinfer_api with an optional trace= argument so tensor inputs are logged before the kernel and fi_trace() can generate benchmark JSON.

🤖 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/packed_kda_decode_cute.py` around lines 860 - 872,
Decorate the public run_packed_kda_decode_cute entrypoint with flashinfer_api,
preserving torch.no_grad and exposing the optional trace parameter so tensor
inputs are logged and fi_trace() can generate benchmark JSON.

Sources: Coding guidelines, Learnings

tests/kda/test_packed_kda_decode_cute.py (1)

353-357: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Seed the dt_bias tensor.

Every other tensor in this file comes from a seeded torch.Generator. This torch.randn call uses the global RNG, so the test values change between runs. Reproduce failures by passing a seeded generator.

♻️ Proposed change

+    dt_bias_generator = torch.Generator(device=packed_kda_cute_device).manual_seed(
+        20261191
+    )
     dt_bias_storage = torch.randn(
         _GATE_WIDTH + 1,
         dtype=torch.float32,
         device=packed_kda_cute_device,
+        generator=dt_bias_generator,
     )

🤖 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 `@tests/kda/test_packed_kda_decode_cute.py` around lines 353 - 357, Update the
dt_bias_storage initialization in the packed KDA test to use the same seeded
torch.Generator as the other tensors in the file, passing it to torch.randn so
test values remain deterministic and reproducible.

🤖 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 `@benchmarks/bench_packed_kda_decode.py`:
- Around line 269-281: Update the benchmark flow around _run, _capture, and
bench_gpu_time so each warmup and measured decode starts from restored case
state. Build separate reset and run CUDA graphs, replay the reset graph
immediately before every timed run replay, and apply the equivalent
reset-before-timing behavior in direct mode rather than passing one stateful
callable unchanged.

In `@flashinfer/kda_kernels/packed_kda_decode_cute.py`:
- Around line 917-918: Guard state-pool indexing in both decode kernels,
including the kernel around the current wrapper validation and
_packed_kda_decode_kernel: pass the state pool row count into each kernel and
require each state_indices value to be within [0, state.shape[0]) before
indexing, while preserving the existing requested_slot >= 0 liveness behavior.
Do not add a host-side range check or synchronization; alternatively, explicitly
document the caller’s in-range slot contract if that is the chosen design.
- Around line 767-790: The fake tensors in the decode setup must only declare
16-byte alignment when the runtime gate verified the corresponding
pointer/layout. Update mixed_qkv, raw_gate, state, raw_beta, and constant
tensors to use element-width/default alignment in unaligned paths, and do not
promote raw_beta unless its own pointer is covered by the alignment check;
preserve 16-byte assumptions only for explicitly validated inputs.

In `@tests/kda/test_packed_kda_decode_cute.py`:
- Around line 42-49: Update the packed_kda_cute_device fixture to remove the
torch.cuda.is_available() guard and replace the raw get_device_capability check
with the flashinfer.utils architecture helper or matching `@backend_requirement`.
Preserve the exact compute-capability 10.0 requirement while including the
existing is_sm100a_supported requirement used by the archive convention.

---

Nitpick comments:
In `@flashinfer/kda_kernels/packed_kda_decode_cute.py`:
- Around line 846-857: Replace the hand-rolled _check_b200 architecture
validation with the supported_compute_capability([100]) decorator on the
targeted API, while keeping the explicit CUDA 12.8 version check. Remove the
redundant _check_b200 calls, ensuring capability gating occurs before output
allocation in the function around the existing API definition.
- Around line 427-428: Add a compile-time assertion immediately after
iterations_per_group in the CTA kernel setup, requiring it to be greater than
zero (equivalently, rejecting tile_v values that produce zero iterations). Keep
the existing _get_compiled_kernel dispatch unchanged while ensuring unsupported
tile ranges fail during compilation before output is written.
- Around line 700-713: Apply cute.assume with divby=_ELEMS_PER_LANE to
mixed_qkv.stride[0] and raw_gate.stride[0] before constructing mixed_layout and
gate_layout, matching _packed_kda_decode_warp_launch. Use the narrowed stride
values in both layout stride tuples so cute.autovec_copy can rely on the
verified row alignment.
- Around line 154-163: Extract the repeated state-row loading logic into a
shared helper, including cute.local_tile construction and the
use_aligned_io-dependent autovec_copy or element-wise copy. Replace the
implementations in the prologue, loop body, and CTA kernel while preserving
their existing indices, buffers, and vectorization behavior.
- Around line 61-63: Add a concise comment above _select_tile_v documenting that
the batch threshold of 26 is a measured performance crossover on the dispatch
hot path, and state the alternative tile choices (8 below the threshold, 16 at
or above it). Preserve the existing selection logic.
- Around line 860-872: Decorate the public run_packed_kda_decode_cute entrypoint
with flashinfer_api, preserving torch.no_grad and exposing the optional trace
parameter so tensor inputs are logged and fi_trace() can generate benchmark
JSON.

In `@tests/kda/test_packed_kda_decode_cute.py`:
- Around line 353-357: Update the dt_bias_storage initialization in the packed
KDA test to use the same seeded torch.Generator as the other tensors in the
file, passing it to torch.randn so test values remain deterministic and
reproducible.

🪄 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: 7e279b9f-1fe9-4e4c-bf0d-240854dbfc0a

📥 Commits

Reviewing files that changed from the base of the PR and between b1d9585 and f701c92.

📒 Files selected for processing (4)

• benchmarks/bench_packed_kda_decode.py
• flashinfer/kda_kernels/__init__.py
• flashinfer/kda_kernels/packed_kda_decode_cute.py
• tests/kda/test_packed_kda_decode_cute.py

[coderabbitai]

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 `@tests/kda/test_packed_kda_decode_cute.py`:
- Around line 187-188: Handle an empty state pool in _reference_step: either
return zero output without mutation when state.shape[0] is zero, preserving
inactive semantics, or add explicit public API validation that rejects empty
pools before index clamping. Add a test covering the chosen contract, including
empty-state behavior or validation failure.

🪄 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: 22b3c554-87a3-4bcf-aa6a-ecbf1b419f2c

📥 Commits

Reviewing files that changed from the base of the PR and between f701c92 and d15c132.

📒 Files selected for processing (2)

• flashinfer/kda_kernels/packed_kda_decode_cute.py
• tests/kda/test_packed_kda_decode_cute.py

🚧 Files skipped from review as they are similar to previous changes (1)

• flashinfer/kda_kernels/packed_kda_decode_cute.py

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Support or reject an empty state pool.

The dispatch validation accepts state.shape[0] == 0. Line 188 then clamps indices to [-1], and Line 189 cannot select a row from the empty tensor.

The documented contract treats every out-of-range index as inactive. An empty pool therefore must produce zero output without mutation, or the public API must reject an empty pool explicitly. Update _reference_step and add an empty-pool test, or add and test explicit API validation.

🤖 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 `@tests/kda/test_packed_kda_decode_cute.py` around lines 187 - 188, Handle an
empty state pool in _reference_step: either return zero output without mutation
when state.shape[0] is zero, preserving inactive semantics, or add explicit
public API validation that rejects empty pools before index clamping. Add a test
covering the chosen contract, including empty-state behavior or validation
failure.

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

flashinfer/kda_kernels/packed_kda_decode_cute_claude.py (1)

1338-1338: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Honor or remove use_packed_fma in the persistent kernel.

_kda_packed_t1_persist_kernel declares use_packed_fma at Line 1338 but never reads it. Pass 1 at Line 1629 and pass 2 at Line 1669 always call cute.arch.fma_packed_f32x2. The other two kernels branch on cutlass.const_expr(use_packed_fma), so the persistent kernel silently ignores the flag while still consuming a _get_compiled cache-key dimension.

_check_b200 restricts execution to compute capability 10.0, so no invalid instruction is generated today. Either add the const_expr branch for consistency, or drop the parameter from the persistent kernel and its launch wrapper so the contract matches the behavior.

Also applies to: 1626-1679

🤖 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/packed_kda_decode_cute_claude.py` at line 1338, Update
_kda_packed_t1_persist_kernel and its launch wrapper so use_packed_fma matches
actual behavior: either branch both pass-1 and pass-2 FMA calls on
cutlass.const_expr(use_packed_fma), or remove the unused parameter and
corresponding _get_compiled cache-key dimension throughout the persistent-kernel
launch path.

🤖 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/packed_kda_decode_cute_claude.py`:
- Around line 2133-2145: Validate all tuning environment overrides before
returning from the configuration helper and before tracing in
launch_packed_kda_decode_cute. Reject invalid TILE_V, ILP, THREADS, STAGES, and
CHUNKR values when they violate the kernel’s divisibility and minimum-value
invariants, including supported tile sizes, full V coverage, positive integral
row/group iterations, stages at least two, and positive THREADS_PER_ROW. Raise a
clear configuration error instead of allowing invalid schedules to reach
run_packed_kda_decode_cute_claude.
- Around line 1049-1094: Update launch_packed_kda_decode_cute to enforce that
private_ring, tma_read, and bulk_store are mutually exclusive, rejecting or
normalizing invalid combinations before kernel launch. Compute the final mode
flags there and pass those same bulk_store, tma_read, and private_ring values to
_get_compiled instead of recomputing them inline, ensuring staging setup and
wait/refill dispatch select one consistent mode.

---

Nitpick comments:
In `@flashinfer/kda_kernels/packed_kda_decode_cute_claude.py`:
- Line 1338: Update _kda_packed_t1_persist_kernel and its launch wrapper so
use_packed_fma matches actual behavior: either branch both pass-1 and pass-2 FMA
calls on cutlass.const_expr(use_packed_fma), or remove the unused parameter and
corresponding _get_compiled cache-key dimension throughout the persistent-kernel
launch path.

🪄 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: dfaa5ff9-b0e8-4c19-abf3-367df4d7f3f5

📥 Commits

Reviewing files that changed from the base of the PR and between d15c132 and 8a75f829a43fc655384a426570603e365c580ee1.

📒 Files selected for processing (4)

• benchmarks/bench_packed_kda_decode.py
• flashinfer/kda_kernels/__init__.py
• flashinfer/kda_kernels/packed_kda_decode_cute_claude.py
• tests/kda/test_packed_kda_decode_cute.py

🚧 Files skipped from review as they are similar to previous changes (1)

• flashinfer/kda_kernels/init.py

[coderabbitai]

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Enforce exclusivity between private_ring, tma_read, and bulk_store.

The wait/refill dispatch treats the three staging modes as alternatives, but bulk_store and tma_read are selected independently of private_ring in launch_packed_kda_decode_cute. Two combinations corrupt the ring:

• private_ring=True with bulk_store=True: the bulk-store write-back at Line 1269-1292 still runs, but _cp_async_bulk_wait_read only exists in the else branch at Line 1086-1093. A refill can overwrite a ring slot while its shared-to-global bulk store is still reading it. The state written to the pool and the state read into registers both become undefined. FLASHINFER_PACKED_KDA_PRIVRING defaults to "1", so setting FLASHINFER_PACKED_KDA_BULK=1 alone reaches this combination.
• private_ring=True with tma_read=True: the mbarriers are initialized and all slots are prefilled at Line 920-939, but the private-ring branch never waits on those mbarriers and refills the same slots with cp.async. Two producers then write the same ring slots.

Reject or normalize the invalid combinations at the launch site.

🐛 Proposed fix in `launch_packed_kda_decode_cute`

+    private_ring = os.environ.get("FLASHINFER_PACKED_KDA_PRIVRING", "1") == "1"
+    bulk_store = (
+        os.environ.get("FLASHINFER_PACKED_KDA_BULK", "0") == "1"
+        and n_stages > 0
+        and chunk_rows == _NUM_GROUPS * ilp_rows
+    )
+    tma_read = os.environ.get("FLASHINFER_PACKED_KDA_TMA", "0") == "1" and n_stages > 0
+    if bulk_store or tma_read:
+        # The barriered ring owns the bulk-store completion waits and the
+        # mbarrier waits; the private ring implements neither.
+        private_ring = False
+    if bulk_store and tma_read:
+        raise ValueError(
+            "FLASHINFER_PACKED_KDA_BULK and FLASHINFER_PACKED_KDA_TMA are "
+            "mutually exclusive"
+        )

Then pass bulk_store, tma_read, and private_ring to _get_compiled instead of recomputing them inline.

🤖 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/packed_kda_decode_cute_claude.py` around lines 1049 -
1094, Update launch_packed_kda_decode_cute to enforce that private_ring,
tma_read, and bulk_store are mutually exclusive, rejecting or normalizing
invalid combinations before kernel launch. Compute the final mode flags there
and pass those same bulk_store, tma_read, and private_ring values to
_get_compiled instead of recomputing them inline, ensuring staging setup and
wait/refill dispatch select one consistent mode.

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Validate the tuning environment knobs against the kernel divisibility invariants.

The knobs are applied without checking the invariants that the kernels assume. Every affected expression is a floor division evaluated at trace time, so an invalid value silently reduces the iteration count instead of failing. The kernel then updates only part of the state slot and writes only part of the output. Because run_packed_kda_decode_cute_claude allocates the output with new_empty, the unwritten rows keep uninitialized data.

Concrete failures:

• FLASHINFER_PACKED_KDA_THREADS=48 sets num_groups=3, so ROWS_PER_GROUP = tile_v // num_groups = 16 // 3 = 5 at Line 444 drops one V row per group.
• FLASHINFER_PACKED_KDA_ILP=3 makes ITERS = ROWS_PER_GROUP // ilp_rows = 2 // 3 = 0 at Line 445, so the register-prefetch kernel performs no work and writes no output.
• FLASHINFER_PACKED_KDA_TILE_V is not checked against _SUPPORTED_TILE_V, so a value that does not divide V makes num_v_tiles = V // tile_v at Line 1879 cover fewer than V rows.
• FLASHINFER_PACKED_KDA_STAGES=1 makes D = n_stages - 1 = 0 at Line 919, so the prologue issues no cp.async and the main loop consumes ring slots that were never filled.
• FLASHINFER_PACKED_KDA_CHUNKR=256 makes THREADS_PER_ROW = _NUM_THREADS // CHUNK_ROWS = 0 at Line 806, so COL_SPAN = K // THREADS_PER_ROW raises ZeroDivisionError during tracing.

Reject invalid knob values instead of tracing an incorrect schedule.

🐛 Proposed validation

     tile_v, ilp_rows, num_groups, stages, evict = cfg
     if _TILE_V_OVERRIDE:
         tile_v = int(_TILE_V_OVERRIDE)
+        if tile_v not in _SUPPORTED_TILE_V:
+            raise ValueError(
+                f"FLASHINFER_PACKED_KDA_TILE_V must be one of {_SUPPORTED_TILE_V}, "
+                f"got {tile_v}"
+            )
         ilp_rows = max(1, min(4, tile_v // _NUM_GROUPS))
         num_groups = _NUM_GROUPS
     ilp_env = os.environ.get("FLASHINFER_PACKED_KDA_ILP")
     if ilp_env:
         ilp_rows = int(ilp_env)
     threads_env = os.environ.get("FLASHINFER_PACKED_KDA_THREADS")
     if threads_env:
+        threads = int(threads_env)
+        if threads % _LANES_PER_ROW != 0:
+            raise ValueError(
+                "FLASHINFER_PACKED_KDA_THREADS must be a multiple of "
+                f"{_LANES_PER_ROW}, got {threads}"
+            )
-        num_groups = int(threads_env) // _LANES_PER_ROW
+        num_groups = threads // _LANES_PER_ROW
+    if num_groups < 1 or tile_v % num_groups != 0:
+        raise ValueError(f"num_groups={num_groups} must divide tile_v={tile_v}")
+    rows_per_group = tile_v // num_groups
+    if ilp_rows < 1 or rows_per_group % ilp_rows != 0:
+        raise ValueError(
+            f"ilp_rows={ilp_rows} must divide tile_v // num_groups={rows_per_group}"
+        )
     if os.environ.get("FLASHINFER_PACKED_KDA_EVICT_FIRST"):

And in launch_packed_kda_decode_cute:

     chunk_rows = max(chunk_rows, _NUM_GROUPS * ilp_rows)
+    if chunk_rows > _NUM_THREADS or tile_v % chunk_rows != 0:
+        raise ValueError(
+            f"chunk_rows={chunk_rows} must divide tile_v={tile_v} and be at most "
+            f"{_NUM_THREADS}"
+        )
     iters = tile_v // chunk_rows
     if pool_div != 8 or iters < 2:
         n_stages = 0
     if n_stages:
         n_stages = min(
             int(os.environ.get("FLASHINFER_PACKED_KDA_STAGES", str(n_stages))),
             iters + 1,
         )
+        if n_stages < 2:
+            raise ValueError(
+                f"FLASHINFER_PACKED_KDA_STAGES must be at least 2, got {n_stages}"
+            )
         num_groups = _NUM_GROUPS  # smem kernel is fixed at 128 threads

Also applies to: 2200-2228

🤖 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/packed_kda_decode_cute_claude.py` around lines 2133 -
2145, Validate all tuning environment overrides before returning from the
configuration helper and before tracing in launch_packed_kda_decode_cute. Reject
invalid TILE_V, ILP, THREADS, STAGES, and CHUNKR values when they violate the
kernel’s divisibility and minimum-value invariants, including supported tile
sizes, full V coverage, positive integral row/group iterations, stages at least
two, and positive THREADS_PER_ROW. Raise a clear configuration error instead of
allowing invalid schedules to reach run_packed_kda_decode_cute_claude.

[kahyunnam]

Blocking API-design question: please clarify whether this is intended as a public operation or an internal experimental backend.

The existing KDA surface is already split between the phase-neutral flashinfer.recurrent_kda, the separate flashinfer.kda_decode.recurrent_kda, and flashinfer.fused_kda_decode. This PR currently introduces another pattern: the only callable is the backend-specific flashinfer.kda_kernels.run_packed_kda_decode_cute, exported from kda_kernels.all, with no @flashinfer_api, trace template, top-level facade, or documentation.

Because the packed serving contract differs materially from recurrent_kda, a separate operation may be justified.

• If public, please expose a backend-neutral noun-form API (e.g., packed_recurrent_kda_decode) through the normal facade, decorator, trace, top-level export, and docs, while retaining run_packed_kda_decode_cute as the internal implementation.
• If this is experimental/internal instead, please avoid exporting it as public API and state that explicitly.

(I don’t think this PR needs to solve the pre-existing KDA unification problem, but it should probably not add a third public-entrypoint convention)

[ameynaik-hub]

@kahyunnam

• The public way to use these kernels is the existing flashinfer.recurrent_kda. Its wrapper now routes eligible T=1 decode calls (12 heads, head dim 128, bf16, SM100a, 1-D ssm_state_indices; both raw and pre-computed gate/beta conventions) to the new kernel.
  All other calls (T>1, spec decode, GQA, varlen, other shapes) run the existing kernels unchanged. Kill switch: FLASHINFER_KDA_T1_FAST_PATH=0.
• run_packed_kda_decode_cute is no longer exported from kda_kernels.__all__; its module docstring states it is internal (kept importable by module path for its tests and benchmark).
• So no third public-entrypoint convention is added — eligible calls inherit recurrent_kda's existing @flashinfer_api, trace template, top-level export, and docs.

Benchmarks vs the existing recurrent kernel (B200, CUPTI median, warmup 100 / 100 iters, through the public recurrent_kda API):

Batch	Existing kernel (kernel time only)	Existing full decode step (with state cache)	New kernel	Speedup (kernel / full step)
8	~4.0 µs	63.3 µs	4.1 µs	~1.0× / 15×
64	13.5 µs	87.8 µs	11.0 µs	1.2× / 8×
512	81.1 µs	638.8 µs	63.7 µs	1.27× / 10×

The big full-step gap is not the existing kernel's fault: with ssm_state_indices its wrapper copies each request's state out of the pool and back (at B=512: 69 µs gather + 463 µs elementwise scatter). The new kernel looks up its slot inside the kernel, so
those copies disappear. Kernel-vs-kernel the gain is a more modest 1.2–1.3×.

Correctness: 42 tests pass, including fp32-reference parity for both paths and both gate conventions, toggle behavior, misaligned/strided inputs, out-of-range slots, and CUDA-graph replay.

[kahyunnam]

/bot run tests/kda

[flashinfer-bot]

GitLab MR !1207 has been created, and the CI pipeline #62378647 is currently running. I'll report back once the pipeline job completes.

[flashinfer-bot]

[SUCCESS] Pipeline #62378647: 18/18 executed test jobs passed