feat(cake_kda): add serving-native packed decode backend

[yyihuang]

Description

This PR adds a serving-native CAKE KDA decode backend behind a new independent
flashinfer.packed_kda_decode API. It consumes post-convolution packed QKV,
raw gate and beta logits, and an indexed BF16 recurrent-state pool directly,
so serving integrations do not need to split/copy Q/K/V, preprocess gate/beta,
or gather/scatter state around the recurrent update.

This API is intentionally separate from recurrent_kda (already-split inputs)
and fused_kda_decode (which also owns convolution and gated RMSNorm).

Exact contract

• T=1, H=HV=12, K=V=128, BF16 data/state;
• packed QKV [B, 4608], raw gate [B, 1536], raw beta [B, 12];
• FP32 A_log[12] and dt_bias[1536];
• indexed state [N, 12, 128, 128], with compact inner dimensions and an
  arbitrary positive/disjoint outer slot stride;
• contiguous int32 state indices; unique in-bounds active slots are a caller
  contract, and -1 is an inactive graph-padding row;
• fixed scale=1/sqrt(128), L2 epsilon 1e-6, and lower_bound=-5 semantics;
• canonical caller-owned output [B, 1, 12, 128] and the caller's current
  PyTorch CUDA stream.

The public call is strict: unsupported devices or tensor contracts raise and
never silently route to another implementation.

Frozen schedules and architecture boundary

• B < 32: tile8 schedule, grid (192, B, 1);
• B >= 32: tile16 schedule, grid (96, B, 1);
• block (32, 1, 1), zero dynamic shared memory;
• exact sm_100a and exact sm_103a JIT/AOT modules; no sm_100f claim.

The frozen generated bodies are protected by raw-body SHA-256 tests:

• tile8: d0de8869242d09bf0c1c4840a7fd73dcd32835050cdc08db58b19a2c7506d0da;
• tile16: d8a446e42da47e2d8cd05139c77efe9c970f2d36394b68b49649beb6bc2bbfbe.

Integration and tests

• independent TVM-FFI binding and target-aware JIT/AOT registration;
• public API, top-level export, trace definition including the in-place state
  side effect, and API documentation;
• selector, frozen-source, binding, and AOT tests;
• GPU correctness coverage for B=1/8/16/31/32/64/128/256/512, production
  strides, non-identity and -1 indices, untouched state/padding, current
  stream, CUDA Graph changed-input/index replay, and a 512-step FP64 diagnostic;
• strict CUPTI direct/graph benchmark with clean-source and loaded-binary
  digests.

Validation

PR head: 826ed5f2fdc301440aa1ec799066f2fa3dfb63d5.

The performance-evidence commit is its direct parent
89a5d3fcee8f8ddcff9da5547d095ffd73c1ef79. The follow-up only pre-includes
math_constants.h outside the generated-type rename boundary and sorts a
Python __all__; it does not change either frozen body, the binding checks,
launch parameters, or public API. The exact PR head was revalidated on both
architectures after that review follow-up.

• focused CPU/JIT/AOT/facade tests: 27 passed;
• pre-commit on all changed files: passed;
• frozen tile8/tile16 body hashes: exact match.

Exported public-API correctness

Target	GPU	JIT tests	GPU tests	Worst output max abs	Worst state max abs	Sanitizer
sm_100a	NVIDIA B200	19 passed	26 passed	1.5259e-5	2.4414e-4	synccheck + memcheck: 0 errors
sm_103a	NVIDIA GB300	19 passed	26 passed	7.6294e-6	4.8828e-4	synccheck + memcheck: 0 errors

Both targets cover the complete batch sweep, production strides, non-identity
and inactive indices, current-stream launches, CUDA Graph changed-input replay,
and the 512-step FP64 diagnostic.

Exported public API vs exact SGLang packed Triton

Median GPU span from strict CUPTI, cold L2, 30 rounds per backend, same inputs,
and alternating AB/BA order. The baseline is the exact raw SGLang packed Triton
serving kernel (source SHA-256 cf5e980d86174a631bcd0872f8ebf7a6ab7c21c3435c097f66f8ba0e686debc0),
not a copied CAKE-internal timing table.

B200 / sm_100a

B	Direct CAKE / Triton (us)	Direct speedup	CUDA Graph CAKE / Triton (us)	Graph speedup
1	3.6800 / 4.7360	1.2870x	3.8395 / 4.8160	1.2543x
8	4.5120 / 5.5355	1.2268x	4.7035 / 5.4720	1.1634x
16	5.6960 / 6.5275	1.1460x	5.8875 / 6.5440	1.1115x
31	8.4480 / 8.7360	1.0341x	8.6400 / 8.6720	1.0037x
32	8.6885 / 8.7840	1.0110x	8.8635 / 8.7675	0.9892x
64	13.7280 / 15.1525	1.1038x	13.8080 / 15.2000	1.1008x
128	22.2880 / 24.2240	1.0869x	22.2080 / 24.5120	1.1037x
256	39.8555 / 41.9840	1.0534x	40.0000 / 41.7440	1.0436x
512	75.0550 / 77.5995	1.0339x	75.0240 / 77.3760	1.0313x
Geomean		1.1058x		1.0863x

GB300 / sm_103a

B	Direct CAKE / Triton (us)	Direct speedup	CUDA Graph CAKE / Triton (us)	Graph speedup
1	4.0160 / 4.6080	1.1474x	4.0645 / 4.4640	1.0983x
8	4.5760 / 5.3600	1.1713x	4.6720 / 5.3760	1.1507x
16	5.6485 / 6.4000	1.1330x	5.8080 / 6.4000	1.1019x
31	8.7040 / 8.1600	0.9375x	8.8800 / 8.1920	0.9225x
32	8.3525 / 8.2560	0.9884x	8.5280 / 8.2560	0.9681x
64	13.2480 / 14.3520	1.0833x	13.4400 / 14.1920	1.0560x
128	21.0240 / 22.8320	1.0860x	21.1200 / 23.0565	1.0917x
256	37.7125 / 40.2400	1.0670x	37.9675 / 40.1760	1.0582x
512	70.4320 / 73.5685	1.0445x	70.6720 / 73.8725	1.0453x
Geomean		1.0708x		1.0526x

Every timed sample on both GPUs contained exactly one launch and one kernel
activity with zero inter-kernel gap. Across all batches and both modes, worst
output max abs was 1.5259e-5, worst state max abs was 4.8828e-4, and
unselected state slots and padding were bitwise preserved. Regressions at
GB300 B31/B32 and B200 graph B32 are retained rather than hidden by geomeans.

A same-configuration cross-run Kimi-K3 SGLang comparison reports CAKE output
throughput 1985.439 tok/s versus the recorded Triton baseline
1979.075 tok/s (1.0032x), while median TPOT is effectively flat
(30.814 ms versus 30.806 ms, 0.9998x). GSM8K medians are 0.975 for
CAKE and 0.980 for Triton, passing the configured quality gate. Because the
measurements use different source snapshots and allocations and the delta is
noise-scale, they do not establish a material E2E speedup. Full runs and
methodology are in the SGLang PR.

SGLang integration: sgl-project/sglang#33647

Related to #4254.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a packed Kimi K3 T=1 recurrent decode path. The change includes exact-target CUDA kernels, JIT/AOT integration, public API and trace support, Blackwell correctness tests, and a strict CUPTI benchmark with JSON provenance reporting.

Changes

Packed KDA decode

Layer / File(s)	Summary
CUDA kernels and binding csrc/kda/flashkda_packed_t1_binding.cuh, csrc/kda/flashkda_packed_t1_tile*.cu	Adds tile8 and tile16 kernels with BF16 state updates, inactive-row handling, aligned access paths, and strict tensor and alias validation.
JIT and AOT module integration flashinfer/jit/flash_kda_packed_t1.py, flashinfer/aot.py, tests/jit/test_flash_kda_packed_t1_jit.py	Adds target-specific module generation, batch-based variant selection, exact SM100a/SM103a checks, AOT registration, and JIT integrity tests.
Public API and trace wiring flashinfer/kda_kernels/*, flashinfer/kda_decode.py, flashinfer/trace/templates/kda.py, flashinfer/__init__.py, docs/api/kda_decode.rst	Adds packed_kda_decode, its backend and trace template, top-level export, fixed-shape contract, and API documentation.
Correctness and contract validation tests/kda/test_packed_kda_decode.py, tests/trace/example.py	Tests numerical output, state mutation, inactive rows, output handling, streams, CUDA graphs, target selection, and trace usage.
CUPTI benchmark and reporting benchmarks/bench_packed_kda_decode.py	Adds deterministic reference validation, direct and graph timing, source and binary provenance checks, and JSON result reporting.

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

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant packed_kda_decode
  participant run_packed_kda_decode
  participant JITModule
  participant CUDAKernel
  participant StatePool
  Caller->>packed_kda_decode: pass packed QKV, gate, beta, and state
  packed_kda_decode->>run_packed_kda_decode: forward decode request
  run_packed_kda_decode->>JITModule: select target and batch variant
  JITModule->>CUDAKernel: launch on current CUDA stream
  CUDAKernel->>StatePool: read and update recurrent state
  CUDAKernel-->>run_packed_kda_decode: produce output tensor
  run_packed_kda_decode-->>Caller: return output

Loading

Possibly related PRs

• flashinfer-ai/flashinfer#4279: Adds related frozen CUDA recurrent-KDA decode backends and JIT/AOT integration.

Suggested labels: run-ci, op: misc, model: dsv4

Suggested reviewers: yzh119, aleozlx, djmmoss

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 30.43% 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 new serving-native packed KDA decode backend, which is the main change.
Description check	✅ Passed	The description thoroughly explains the API contract, implementation, testing, validation results, and related work.

✨ 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: 3

🧹 Nitpick comments (4)

benchmarks/bench_packed_kda_decode.py (2)

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

Replace the assigned lambdas with local functions.

Ruff reports E731 for Lines 341 and 342.

♻️ Proposed fix

     if mode == "direct":
-        warmup_run = lambda: _run(warmup_case)
-        measured_run = lambda: _run(case)
+
+        def warmup_run():
+            return _run(warmup_case)
+
+        def measured_run():
+            return _run(case)
+
     elif mode == "cuda_graph":

🤖 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 `@benchmarks/bench_packed_kda_decode.py` around lines 339 - 350, Replace the
lambda assignments in _make_timing_runners for direct mode with local def
functions that invoke _run on warmup_case and case, preserving the returned
warmup_run and measured_run behavior.

Source: Linters/SAST tools

---

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

Document the 6 warmup-call budget used by CUPTI timing.

warmup_calls = 6 + warmup_iters encodes the extra staged_run invocation made while bench_gpu_time estimates execution time: one initial overhead-exclusion call plus five estimation iterations. cold_l2_cache=True does not add extra runner calls here, but use_cuda_graph=True would change this path. Add an inline comment stating which internal calls this budget covers so helper changes catch the issue early and the benchmark does not stop on a stale constant.

🤖 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 `@benchmarks/bench_packed_kda_decode.py` around lines 353 - 380, Document the
warmup_calls calculation in _timed_sample with an inline comment explaining that
the fixed 6 covers one overhead-exclusion invocation plus five execution-time
estimation iterations, while warmup_iters accounts for dry-run calls. Note that
cold_l2_cache does not add runner calls here and this budget assumes
use_cuda_graph is disabled.

tests/kda/test_packed_kda_decode.py (1)

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

Add a CUDA-version guard to the skip condition.

The fixture only checks the compute capability. flashinfer/kda_kernels/packed_kda_decode.py also requires CUDA 12.8 for sm100a and CUDA 12.9 for sm103a, as the parametrized cases in Lines 654-662 show. On a CC 10.0 or 10.3 device with an older CUDA toolkit, _target_for_device raises RuntimeError and every GPU test fails instead of skipping.

Add the version check to the fixture so the tests skip on unsupported toolkits.

♻️ Proposed guard

     device = torch.device("cuda")
-    if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3)):
+    capability = torch.cuda.get_device_capability(device)
+    if capability not in ((10, 0), (10, 3)):
         pytest.skip(
             "packed KDA T=1 requires exact CC 10.0 (SM100a) or CC 10.3 (SM103a)"
         )
+    required_cuda = "12.8" if capability == (10, 0) else "12.9"
+    if not is_cuda_version_at_least(required_cuda):
+        pytest.skip(f"packed KDA T=1 on CC {capability} requires CUDA {required_cuda}")
     return device

Import the helper from flashinfer.utils.

As per coding guidelines: "architecture-specific tests must skip unsupported GPUs using the appropriate flashinfer.utils checks or backend capability APIs".

🤖 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.py` around lines 45 - 54, Update the
packed_kda_device fixture to also validate the installed CUDA version using the
appropriate helper from flashinfer.utils, alongside the existing torch.cuda
availability and compute-capability checks. Skip CC 10.0 devices unless CUDA
12.8 is supported and CC 10.3 devices unless CUDA 12.9 is supported, preserving
the current behavior for supported configurations.

Source: Coding guidelines

flashinfer/jit/flash_kda_packed_t1.py (1)

75-80: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Derive the selector threshold from the variant metadata.

FLASH_KDA_PACKED_T1_VARIANT_METADATA declares batch_min and batch_max, but _variant_for_batch hard-codes 32. The threshold now lives in two places. A future retune must change both, and nothing in the module enforces that they agree.

♻️ Proposed refactor

 def _variant_for_batch(batch: int) -> FlashKDAPackedT1Variant:
     """Select the frozen schedule using only host-visible shape metadata."""
 
     if batch <= 0:
         raise ValueError(f"packed KDA T=1 batch must be positive, got {batch}")
-    return "tile16" if batch >= 32 else "tile8"
+    for variant in FLASH_KDA_PACKED_T1_VARIANTS:
+        metadata = FLASH_KDA_PACKED_T1_VARIANT_METADATA[variant]
+        if batch >= metadata.batch_min and (
+            metadata.batch_max is None or batch <= metadata.batch_max
+        ):
+            return variant
+    raise ValueError(f"no packed KDA T=1 schedule covers batch {batch}")

🤖 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/jit/flash_kda_packed_t1.py` around lines 75 - 80, Update
_variant_for_batch to derive the batch-selection threshold from
FLASH_KDA_PACKED_T1_VARIANT_METADATA rather than hard-coding 32, using the
declared variant batch bounds to preserve tile8 below the boundary and tile16 at
or above it.

🤖 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`:
- Line 56: Validate TYPED_SOURCE_SHA256 against the SHA-256 digest of
csrc/kda/flashkda_packed_t1_binding.cuh from disk before publishing the
benchmark report. Extend the existing verification flow around
_verify_frozen_body_hashes() to raise on mismatch, ensuring report generation
cannot proceed with an incorrect typed-source hash.

In `@csrc/kda/flashkda_packed_t1_binding.cuh`:
- Around line 32-55: Pre-include <math_constants.h> alongside the existing CUDA
headers, before the fixed-width type and tensor-map rename macros in the
binding. Keep the rename region and generated-body inclusion unchanged so the
frozen body’s guarded nested include becomes a no-op.

In `@flashinfer/jit/flash_kda_packed_t1.py`:
- Around line 236-246: Reorder the entries in __all__ so
FLASH_KDA_PACKED_T1_VARIANTS appears before
FLASH_KDA_PACKED_T1_VARIANT_METADATA, satisfying RUF022 while leaving all
exported symbols unchanged.

---

Nitpick comments:
In `@benchmarks/bench_packed_kda_decode.py`:
- Around line 339-350: Replace the lambda assignments in _make_timing_runners
for direct mode with local def functions that invoke _run on warmup_case and
case, preserving the returned warmup_run and measured_run behavior.
- Around line 353-380: Document the warmup_calls calculation in _timed_sample
with an inline comment explaining that the fixed 6 covers one overhead-exclusion
invocation plus five execution-time estimation iterations, while warmup_iters
accounts for dry-run calls. Note that cold_l2_cache does not add runner calls
here and this budget assumes use_cuda_graph is disabled.

In `@flashinfer/jit/flash_kda_packed_t1.py`:
- Around line 75-80: Update _variant_for_batch to derive the batch-selection
threshold from FLASH_KDA_PACKED_T1_VARIANT_METADATA rather than hard-coding 32,
using the declared variant batch bounds to preserve tile8 below the boundary and
tile16 at or above it.

In `@tests/kda/test_packed_kda_decode.py`:
- Around line 45-54: Update the packed_kda_device fixture to also validate the
installed CUDA version using the appropriate helper from flashinfer.utils,
alongside the existing torch.cuda availability and compute-capability checks.
Skip CC 10.0 devices unless CUDA 12.8 is supported and CC 10.3 devices unless
CUDA 12.9 is supported, preserving the current behavior for supported
configurations.

🪄 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: ae6f3a64-588f-45e9-b0d1-34106e163c53

📥 Commits

Reviewing files that changed from the base of the PR and between 307e2fe and 89a5d3f.

📒 Files selected for processing (15)

• benchmarks/bench_packed_kda_decode.py
• csrc/kda/flashkda_packed_t1_binding.cuh
• csrc/kda/flashkda_packed_t1_tile16.cu
• csrc/kda/flashkda_packed_t1_tile8.cu
• docs/api/kda_decode.rst
• flashinfer/__init__.py
• flashinfer/aot.py
• flashinfer/jit/flash_kda_packed_t1.py
• flashinfer/kda_decode.py
• flashinfer/kda_kernels/__init__.py
• flashinfer/kda_kernels/packed_kda_decode.py
• flashinfer/trace/templates/kda.py
• tests/jit/test_flash_kda_packed_t1_jit.py
• tests/kda/test_packed_kda_decode.py
• tests/trace/example.py

[ameynaik-hub]

#4417 I have this PR which kda decode kernel based on my kernel gdn decode kernel which is ~10.7% better than this version.

B	Direct CAKE (µs)	Direct CuTe (µs)	Direct speedup	Graph CAKE (µs)	Graph CuTe (µs)	Graph speedup
1	4.1600	3.9355	1.0570×	4.0965	3.9040	1.0493×
8	4.8640	4.6720	1.0411×	4.9600	4.7685	1.0402×
16	6.1280	5.8240	1.0522×	6.1440	5.8725	1.0462×
31	8.9600	7.6480	1.1715×	9.0240	7.7760	1.1605×
32	9.0240	7.7440	1.1653×	9.2160	7.9360	1.1613×
64	14.1760	12.3845	1.1447×	14.3040	12.5760	1.1374×
128	22.8160	20.3200	1.1228×	22.9120	20.4480	1.1205×
256	40.7675	36.5440	1.1156×	40.9280	36.7040	1.1151×
512	76.4165	68.9445	1.1084×	76.6725	69.1525	1.1087×
Geomean			1.1078×			1.1034×

Methodology:

• Strict CUPTI GPU-span timing
• Cold L2
• 100 measured rounds per backend and configuration
• 100 configured warmup iterations per round
• Identical inputs

[yyihuang]

Superseded by #4445. That PR contains this draft's serving-native packed T=1 decode API and frozen SM100a/SM103a kernels, plus the strided-prefill, indexed-state, native-checkpoint, and later build-boundary hardening work. Closing this draft to keep review on the consolidated PR; this branch and its review history remain available.