feat(cake_kda): add optimized B200 recurrent prefill backend - #4262
Merged
yyihuang merged 9 commits intoAug 3, 2026
Conversation
Dispatch the exact BF16 multi-token recurrent KDA contract to exported SM100a M64/M128 kernels while retaining the existing CuTe backend for decode, speculative decode, GQA, and unsupported shapes. Add graph-safe caller workspaces, JIT/AOT integration, frozen-source integrity checks, correctness coverage, trace updates, documentation, and the six-shape CUPTI benchmark. Signed-off-by: Yingyi Huang <averyh@nvidia.com>
Signed-off-by: Yingyi Huang <averyh@nvidia.com>
yyihuang
requested review from
Aneureka,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww and
yzh119
as code owners
July 30, 2026 04:08
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (8)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughAdds SM100a BF16 FlashKDA M64/M128 kernels, TVM bindings, JIT/AOT packaging, recurrent prefill dispatch, workspace and CUDA graph handling, tests, API documentation, and a CUPTI benchmark. ChangesFlashKDA kernel and binding integration
JIT, AOT, and recurrent prefill API
Documentation and benchmarking
Estimated code review effort: 5 (Critical) | ~120 minutes 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
csrc/kda/flashkda_binding_common.cuh (1)
195-200: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider asserting
beta_tma's row length is a multiple of 8 (TMA stride alignment).
EncodeBetaTmaderivesglobal_strides[0] = d1 * sizeof(__nv_bfloat16)from this same dimension.cuTensorMapEncodeTiledrequires global strides to be 16-byte aligned, so anynum_headsthat is not a multiple of 8 (and > 8) makes the encode fail later with only a rawCUresult. TheEncodeQkTmapath already guards this explicitly viad1 % 64 == 0; mirroring that here keeps the error message actionable.♻️ Suggested guard
const int64_t beta_tma_heads = std::max<int64_t>(num_heads, 8); + TVM_FFI_ICHECK(beta_tma_heads % 8 == 0) + << "beta_tma row length must be a multiple of 8 for TMA stride alignment, got " + << beta_tma_heads; TVM_FFI_ICHECK(beta_tma.ndim() >= 2 && beta_tma.size(beta_tma.ndim() - 1) == beta_tma_heads &&🤖 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` around lines 195 - 200, Update the beta_tma validation near beta_tma_heads to require the row dimension used by EncodeBetaTma to be a multiple of 8, while preserving the existing size checks. Include the alignment requirement in the assertion message so invalid num_heads values fail with an actionable error before cuTensorMapEncodeTiled.csrc/kda/flashkda_bf16_fused_m64_binding.cu (1)
17-35: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTypedef-isolation blocks in both bindings silently depend on include ordering. Both generated TUs include
<cuda_bf16.h>and<math_constants.h>after theuint*_t/int*_trenames are in effect; that is only safe becauseflashkda_binding_common.cuhalready resolved those headers (its lines 20-22), so the nested includes hit their guards. Reordering or trimming those includes would textually rewrite system-header declarations with no obvious diagnostic. The root cause is the same in both files; document the dependency at each site.
csrc/kda/flashkda_bf16_fused_m64_binding.cu#L17-L35: extend the existing rationale comment to state thatflashkda_binding_common.cuhmust be included first and must keep providing<cuda_bf16.h>and<math_constants.h>.csrc/kda/flashkda_bf16_fused_m128_binding.cu#L19-L32: line 19 already defers to the M64 binding for the rationale, so no duplicate prose is needed — just confirm the same include-first ordering holds here once the M64 comment is expanded.🤖 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_bf16_fused_m64_binding.cu` around lines 17 - 35, Expand the rationale comment above the typedef isolation in csrc/kda/flashkda_bf16_fused_m64_binding.cu:17-35 to state that flashkda_binding_common.cuh must remain included first and must continue providing cuda_bf16.h and math_constants.h before the generated source is included. In csrc/kda/flashkda_bf16_fused_m128_binding.cu:19-32, make no duplicate comment change; confirm its existing delegation to the M64 rationale preserves the same include-first dependency.flashinfer/kda_decode.py (1)
47-52: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffStream-pointer-keyed caches are never evicted.
_flash_kda_tensor_cacheand_flash_kda_stream_workspacesare keyed on the rawcuda_streaminteger and hold device allocations for process lifetime. Workloads that create/destroy many side streams accumulate one workspace (descriptor storages + beta padding + possible state scratch) and one set of metadata tensors per stream, and a recycled stream pointer silently adopts the previous entry. Contents are stream-independent so results stay correct, but consider keying the metadata cache by device only and bounding/weak-referencing the per-stream workspaces.Also applies to: 328-343
🤖 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_decode.py` around lines 47 - 52, Update _flash_kda_tensor_cache to use device-based keys because its contents are stream-independent, and redesign _flash_kda_stream_workspaces so per-stream allocations are bounded or weakly referenced rather than retained for process lifetime. Ensure recycled stream pointers cannot inherit stale workspace entries and that descriptor, beta-padding, and state-scratch allocations are released when streams disappear.benchmarks/bench_recurrent_kda_prefill.py (1)
362-371: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive
variantfrom the case, not a hardcoded name comparison.
"m64" if case.name == "h64_fixed8192" else "m128"silently misreports the schedule ifCASESgains another M64-eligible shape. Since the variant is part of the published result, make it an explicitCasefield.♻️ Proposed refactor
`@dataclass`(frozen=True) class Case: name: str num_heads: int seq_lens: tuple[int, ...] packed: bool seed: int + variant: str- "variant": "m64" if case.name == "h64_fixed8192" else "m128", + "variant": case.variant,🤖 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_recurrent_kda_prefill.py` around lines 362 - 371, Update the Case definition and its CASES entries to carry an explicit variant field, then set metadata["variant"] from case.variant in the metadata construction block. Remove the case.name comparison so every M64-eligible shape reports its declared variant accurately.
🤖 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_recurrent_kda_prefill.py`:
- Around line 247-253: Update the RuntimeError raised in candidate_run and the
corresponding exhaustion check in the other applicable block to mention the
--state-rotations option, while preserving the existing state index and limit
details so users know which setting to increase.
In `@tests/jit/test_flash_kda_jit.py`:
- Around line 40-48: Ensure the test clears the cached gen_flash_kda_module
result after execution as well as before patching TARGET_CUDA_ARCHS. Prefer
adding a pytest fixture with setup and teardown cache_clear calls, then use it
for this test so specs created under the fake architecture cannot leak into
later tests.
---
Nitpick comments:
In `@benchmarks/bench_recurrent_kda_prefill.py`:
- Around line 362-371: Update the Case definition and its CASES entries to carry
an explicit variant field, then set metadata["variant"] from case.variant in the
metadata construction block. Remove the case.name comparison so every
M64-eligible shape reports its declared variant accurately.
In `@csrc/kda/flashkda_bf16_fused_m64_binding.cu`:
- Around line 17-35: Expand the rationale comment above the typedef isolation in
csrc/kda/flashkda_bf16_fused_m64_binding.cu:17-35 to state that
flashkda_binding_common.cuh must remain included first and must continue
providing cuda_bf16.h and math_constants.h before the generated source is
included. In csrc/kda/flashkda_bf16_fused_m128_binding.cu:19-32, make no
duplicate comment change; confirm its existing delegation to the M64 rationale
preserves the same include-first dependency.
In `@csrc/kda/flashkda_binding_common.cuh`:
- Around line 195-200: Update the beta_tma validation near beta_tma_heads to
require the row dimension used by EncodeBetaTma to be a multiple of 8, while
preserving the existing size checks. Include the alignment requirement in the
assertion message so invalid num_heads values fail with an actionable error
before cuTensorMapEncodeTiled.
In `@flashinfer/kda_decode.py`:
- Around line 47-52: Update _flash_kda_tensor_cache to use device-based keys
because its contents are stream-independent, and redesign
_flash_kda_stream_workspaces so per-stream allocations are bounded or weakly
referenced rather than retained for process lifetime. Ensure recycled stream
pointers cannot inherit stale workspace entries and that descriptor,
beta-padding, and state-scratch allocations are released when streams disappear.
🪄 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: 3b622346-9139-4327-bb75-d8b370dc307b
📒 Files selected for processing (16)
benchmarks/bench_recurrent_kda_prefill.pycsrc/kda/flashkda_bf16_fused_m128.cucsrc/kda/flashkda_bf16_fused_m128_binding.cucsrc/kda/flashkda_bf16_fused_m64.cucsrc/kda/flashkda_bf16_fused_m64_binding.cucsrc/kda/flashkda_binding_common.cuhdocs/api/kda_decode.rstflashinfer/__init__.pyflashinfer/aot.pyflashinfer/jit/__init__.pyflashinfer/jit/flash_kda.pyflashinfer/kda_decode.pyflashinfer/trace/templates/kda.pytests/jit/test_flash_kda_jit.pytests/kda/test_recurrent_kda_prefill.pytests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json
Comment on lines
+247
to
+253
| def candidate_run(): | ||
| state_index = state_cursors["pr"][0] | ||
| if state_index >= state_rotations: | ||
| raise RuntimeError( | ||
| f"PR state rotations exhausted: {state_index} >= {state_rotations}" | ||
| ) | ||
| state_cursors["pr"][0] += 1 |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
State-rotation exhaustion can abort a run mid-benchmark; point the user at the knob.
bench_gpu_time sizes iteration count from --warmup-ms/--bench-ms, so the number of calls per block scales inversely with kernel time. For the fastest shapes a 120 ms block can exceed 512 slots and kill the whole run. Mention --state-rotations in the error so the recovery is obvious.
♻️ Proposed message change
raise RuntimeError(
- f"PR state rotations exhausted: {state_index} >= {state_rotations}"
+ f"PR state rotations exhausted: {state_index} >= {state_rotations}; "
+ "increase --state-rotations or lower --bench-ms/--warmup-ms"
)Also applies to: 326-333
🤖 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_recurrent_kda_prefill.py` around lines 247 - 253, Update the
RuntimeError raised in candidate_run and the corresponding exhaustion check in
the other applicable block to mention the --state-rotations option, while
preserving the existing state index and limit details so users know which
setting to increase.
Comment on lines
+40
to
+48
| monkeypatch.setattr( | ||
| jit_core.current_compilation_context, | ||
| "TARGET_CUDA_ARCHS", | ||
| {(10, "0a")}, | ||
| ) | ||
| flash_kda.gen_flash_kda_module.cache_clear() | ||
|
|
||
| uri = flash_kda.get_flash_kda_uri(variant) | ||
| spec = flash_kda.gen_flash_kda_module(variant) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Clear the JIT spec cache after the test too, not only before.
gen_flash_kda_module is @functools.cached, so the spec generated under the monkeypatched TARGET_CUDA_ARCHS survives after monkeypatch restores the real arch set. Later tests (or any AOT/JIT code) in the same process will reuse a spec built from the fake {(10, "0a")} context.
🧹 Proposed cleanup
monkeypatch.setattr(
jit_core.current_compilation_context,
"TARGET_CUDA_ARCHS",
{(10, "0a")},
)
flash_kda.gen_flash_kda_module.cache_clear()
+ monkeypatch.setattr(
+ flash_kda.gen_flash_kda_module,
+ "__wrapped__",
+ flash_kda.gen_flash_kda_module.__wrapped__,
+ raising=False,
+ )Simpler and preferred — use a fixture/finalizer so the cache is dropped on exit as well:
`@pytest.fixture`
def clean_flash_kda_cache():
flash_kda.gen_flash_kda_module.cache_clear()
yield
flash_kda.gen_flash_kda_module.cache_clear()🤖 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/jit/test_flash_kda_jit.py` around lines 40 - 48, Ensure the test clears
the cached gen_flash_kda_module result after execution as well as before
patching TARGET_CUDA_ARCHS. Prefer adding a pytest fixture with setup and
teardown cache_clear calls, then use it for this test so specs created under the
fake architecture cannot leak into later tests.
5 tasks
This was referenced Jul 30, 2026
3 tasks
This was referenced Aug 4, 2026
yyihuang
pushed a commit
that referenced
this pull request
Aug 5, 2026
…ts with bit-identical outputs (#4274) ## 📌 Description This PR adds `tinygemm2_sm100`: four CAKE-generated SM100/SM103 device kernels specializing `csrc/tinygemm2.cu` (the TensorRT-LLM tinygemm2 port behind `tinygemm_bf16`) for Blackwell, with **bit-identical outputs** and lower latency. The bias path of `tinygemm_bf16` automatically dispatches to them on compute capability 10.0/10.3; every other path and architecture is unchanged, and `FLASHINFER_DISABLE_TINYGEMM2_SM100=1` forces the reference implementation. The kernels are CAKE-generated Loom schedules, following the frozen-generated-source pattern of the CAKE-generated FlashKDA prefill export (#4262): generated kernel sources checked in verbatim (`clang-format off`, provenance headers, no host-library dependency), concatenated into a single translation unit (`csrc/tinygemm2_sm100/tinygemm2_sm100.cu`) with per-variant kernel symbol renames and a small hand-written binding section doing validation, TMA descriptor encode, and launch — one TU, one JIT module, mirroring the incumbent `csrc/tinygemm2.cu` layout. The single-TU restructure is SASS-verified: all four kernels are byte-identical to per-variant builds on both sm_100a and sm_103a (`cuobjdump -sass`, symbol names normalized). The variants are {deep, shallow} pipeline ring x {PDL on, off}: - stage selection at the Python layer follows the measured B200 crossover axes (shallow ring for `K <= 1024` or grids past 2x the SM count); - PDL variants compile the `griddepcontrol` pair in-kernel and launch with programmatic stream serialization, matching the reference kernel's `USE_PDL=true` instantiation; - the activation TMA descriptor allows an out-of-bounds box on the batch axis (TMA zero-fills), covering batch 1-7 decode shapes. **Correctness contract: bitwise equality with `csrc/tinygemm2.cu`, not a tolerance.** On B200: `torch.equal` parity across batch {1,2,4,7,8,13,16,64} x five (M, K) pairs x PDL on/off, per-variant direct launches, PDL back-to-back replay, dispatch/escape-hatch — 23/23; existing `test_tinygemm2.py` unchanged and green (59/59) with the bias path routed to the new backend. Internally, the same generated kernels passed a 239-row bitwise shape sweep and end-to-end gpt-oss-120b serving with zero token flips. **Performance** (torch.profiler kernel-time medians, B200, bias path, 200 launches): (1,128,720) 2.24us vs 2.34us reference; (16,1024,1024) 2.50 vs 2.58; (64,4096,3072) 18.22 vs 32.58 (1.79x). Small decode shapes are launch-bound; a wider internal CUPTI sweep (cold L2, 35 canonical + 239 regression shapes) measured an 18-23% geometric-mean kernel-time reduction. Limitations: the nobias path stays on the reference kernel (no generated nobias variant yet); hardware validation was on B200 — sm_103a is compiled from the textually identical source and gated on, but not yet run on B300. ## 🔍 Related Issues Related to #4254 (long-term CAKE-generated kernel progress tracker). Follows the frozen generated-kernel export pattern introduced in #4262. ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes The generated kernel sections are frozen artifacts (regeneration happens outside this repo, as with FlashKDA); review focus is best spent on the binding section at the tail of `csrc/tinygemm2_sm100/tinygemm2_sm100.cu`, the dispatch in `flashinfer/gemm/routergemm.py`, and the parity tests. The TMA descriptor encode in the binding is field-for-field identical to the reference kernel's own `cuTensorMapEncodeTiled` calls. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added optimized BF16 GEMM execution for eligible SM100/SM103 GPUs. * Added automatic selection among multiple kernel stages, including PDL variants, based on workload and device capabilities. * Added an environment-variable override to disable the optimization and retain the existing execution path. * **Bug Fixes** * Improved routing for bias-enabled BF16 operations while preserving existing behavior on unsupported devices and bias-free operations. * **Tests** * Added coverage for correctness, parity, repeated launches, variant execution, dispatch, and fallback behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Shanli Xing <shanlix@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
yyihuang
added a commit
that referenced
this pull request
Aug 5, 2026
Follow-up to merged #4262 and #4313. ## What changed - pad the beta-only TMA source to `round_up(H, 8)` instead of only padding `H < 8`; - teach the beta pack kernel and binding validation to use the dynamic padded-head stride; - retain the caller-visible `H`, state shape, launch grid, and frozen M64/M128 CUDA bodies unchanged; - add H=12 eager, packed, full-chunk-plus-tail, final-state, and CUDA graph replay coverage. The frozen kernels load beta in 8-head TMA boxes. For H=12, the original BF16 row stride is 24 bytes and `cuTensorMapEncodeTiled` rejects 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_kda` facade and compared BF16 output plus the complete final state against the PyTorch reference with `atol=rtol=1e-2`. | GPU | Compute capability | CUDA | Result | | --- | --- | --- | --- | | NVIDIA B200 | 10.0 | 12.9 | `62 passed, 0 skipped, 0 failed` | | NVIDIA GB300 | 10.3 | 12.9 | `62 passed, 0 skipped, 0 failed` | The test gate runs: ```text tests/jit/test_flash_kda_jit.py tests/kda/test_recurrent_kda_prefill.py ``` 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_kda` facade 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`](MoonshotAI/FlashKDA@d2ff19a), with CUTLASS `5c149f5`. 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`. | H=12 shape | SM100 / B200: PR / baseline | Speedup | SM103 / GB300: PR / baseline | Speedup | | --- | ---: | ---: | ---: | ---: | | packed `[512] x 32` | 136.159 / 240.239 us | **1.7644x** | 128.264 / 233.496 us | **1.8204x** | | packed `[128] x 8` | 23.760 / 46.448 us | **1.9549x** | 25.712 / 52.904 us | **2.0576x** | | fixed `[512]` | 46.184 / 76.383 us | **1.6539x** | 47.712 / 82.240 us | **1.7237x** | | fixed `[8192]` | 514.197 / 814.435 us | **1.5839x** | 487.161 / 779.426 us | **1.5999x** | | mixed `[1300, 547, 2048, 963, 271, 3063]` | 208.647 / 351.550 us | **1.6849x** | 198.216 / 340.913 us | **1.7199x** | | uniform `[1024] x 8` | 82.080 / 162.703 us | **1.9823x** | 78.440 / 162.505 us | **2.0717x** | | **Six-shape geometric mean** | | **1.7645x** | | **1.8238x** | | Comparison | Result | | --- | --- | | FlashInfer upstream main at H=12 | `unsupported / N/A` | Upstream 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. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved beta padding for head counts that are not divisible by eight. - Ensured packed inputs correctly handle larger, non-aligned head counts. - Added validation for padded storage requirements. - **Documentation** - Clarified beta padding behavior and public tensor shapes. - **Tests** - Expanded coverage for 12-head inputs, varied sequence lengths, chunk boundaries, packed inputs, and CUDA graph updates. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Yingyi Huang <averyh@nvidia.com>
5 tasks
bkryu
pushed a commit
that referenced
this pull request
Aug 11, 2026
…ge-K shapes (#4423) ## 📌 Description Adds a third ring-depth tier to the `tinygemm2_sm100` family (#4274) and moves ring selection into the binding. All six kernels are regenerated from the current CAKE generator state, which also updates the tensor-map parameter passing to a single by-value `__grid_constant__` pack. - **New STAGES=16 ring for single-wave large-K shapes.** The 4/8-stage rings were tuned for K up to ~4K; at larger K the 8-deep ring no longer covers the weight-stream latency when the working set is not L2-resident, and trails the reference kernel's own 16-deep configuration. The new tier closes that: at N=8/M=128/K=7168 (bias path) it measures **4.93 µs vs the reference's 5.22 µs on GB300, and 5.31 µs vs 5.56 µs on B200** (cold-L2 CUPTI). - **Ring selection in the binding**, following the reference `csrc/tinygemm2.cu` launcher convention: stage 4 for K <= 1024 or grids past 2x the SM count (unchanged), stage 16 for single-wave grids with K >= 4608 (measured crossover on both GB300 and B200), stage 8 otherwise. The Python dispatcher becomes a single call into the combined op. - **Dispatch gates on exact compute capabilities (10, 0)/(10, 3).** SM107 passes the previous `major == 10` predicate but must keep the reference path instead of erroring in the binding. - The dynamic-SMEM attribute is now set once per (kernel, device) instead of on every launch. **Correctness contract unchanged: bitwise equality with `csrc/tinygemm2.cu`** — verified per-variant and through the dispatcher (`torch.equal`, batch 1-64, K to 7168, M to 4096, on B200 and GB300), and end-to-end in SGLang serving with zero token flips (gpt-oss-120b, and Mistral-Large-3 whose router GEMM sits in the new tier). `tests/model_optimizations/test_tinygemm2_sm100.py` extends to the stage-16 variants and long-K parity shapes. compute-sanitizer synccheck is clean on all six variants. Limitations: unchanged from #4274 — the nobias path stays on the reference kernel. ## 📊 End-to-end serving validation — Mistral-Large-3 (675B FP8), SGLang, TP4, GB300 (CUDA graphs on, production defaults) **Correctness**: 32 fixed greedy prompts — token streams bitwise identical between arms; GSM8K-200: CAKE 0.950 / ref 0.945. **Performance** (CAKE = this PR's kernels via default dispatch, ref = reference `tinygemm2`; throughput in mean output tok/s, ITL is median in ms): | concurrency | throughput (CAKE) | throughput (ref) | ITL (CAKE) | ITL (ref) | |---|---|---|---|---| | 1 | 100.7 | 101.0 | 9.47 | 9.48 | | 8 | 488.6 | 487.9 | 15.47 | 15.48 | | 32 | 1068 | 1073 | 28.90 | 28.92 | | 128 | 3000 | 2997 | 41.31 | 41.44 | ## 🔍 Related Issues Follow-up to #4274. Related to #4254 (CAKE-generated kernel progress tracker). ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes The generated device sections are frozen artifacts (regeneration happens outside this repo, as with #4262/#4274); review focus is best spent on the binding section of `csrc/tinygemm2_sm100.cu` (pack construction, stage selection, launch attributes), the dispatcher in `flashinfer/gemm/routergemm.py`, and the test additions. The single-TU merge follows the mechanical transform documented in the file header. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for larger matrix workloads through new stage16 execution variants. * Automatically selects the appropriate execution stage based on workload size and hardware. * Added direct launch support for all available execution variants. * Expanded compatibility for supported compute capabilities and newer CUDA versions. * **Bug Fixes** * Improved handling of large reduction sizes and deep-ring workloads. * **Tests** * Added coverage for larger K dimensions, including 7168 and 14336, and stage16 variants. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
spectrometerHBH
added a commit
to spectrometerHBH/TIRx-kernels
that referenced
this pull request
Aug 11, 2026
The m128 reference kernel landed in flashinfer-ai/flashinfer#4262 and ships in the installed release, so `from flashinfer.kda import recurrent_kda` resolves on its own. Remove the worktree override, along with the sys.modules purge that existed to displace an already-imported mainline flashinfer, and the README note telling people to set the variable. Also retitle the README to "TIRx kernels"; the distribution name in the pip line is unchanged.
spectrometerHBH
added a commit
to mlc-ai/TIRx-kernels
that referenced
this pull request
Aug 11, 2026
The m128 reference kernel landed in flashinfer-ai/flashinfer#4262 and ships in the installed release, so `from flashinfer.kda import recurrent_kda` resolves on its own. Remove the worktree override, along with the sys.modules purge that existed to displace an already-imported mainline flashinfer, and the README note telling people to set the variable. Also retitle the README to "TIRx kernels"; the distribution name in the pip line is unchanged.
|
Hi @yyihuang after checking the code, I found Cake backend is more like a TVM-FFI jit module warpper to compile the CUDA sources. We are actually did this in flash float : https://github.com/yiakwy-xpu-ml-framework-team/flash-float-jit-kernels
Is that true ? Note my motivation is that after trying to support many subprojects and I found CUDA is the best cross platform languages so that I can support for different devices. |
5 tasks
8 tasks
This was referenced Aug 16, 2026
This was referenced Aug 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
This PR adds an CAKE-generated optimized SM100a BF16 recurrent-KDA prefill backend and
dispatches to it through the existing
recurrent_kdaAPI.The new backend coexists with the current CuTe-DSL implementation. Decode
(
T=1), speculative decode, GQA, state-pool/checkpoint features, unsupportedlayouts and dtypes, and non-B200 devices continue to use the existing backend
unchanged.
Dispatch contract
The optimized path is selected only when all of the following hold:
T > 1, or packed inputwith more tokens than sequences;
[B, T, H, 128]tensors with a sharedhead count, and beta is contiguous BF16
[B, T, H];A_logis contiguous FP32[H], whiledt_biasis contiguous FP32[H, 128]or[H * 128];use_qk_l2norm_in_kernel=True,use_gate_in_kernel=True,beta_is_logit=True, andlower_boundis finite and negative;accepted-token/checkpoint features are disabled.
Packed input uses
B=1and contiguous CUDA int32/int64cu_seqlens. Thebinding consumes int64 offsets; CUDA graph capture therefore requires callers
to supply int64 offsets directly. An optional contiguous CUDA int32
seq_ordercan order packed sequences for better tail utilization.Calls that do not exactly match this contract fall back to the existing
CuTe-DSL path.
Physical schedules and true in-place state update
B=1, H=64.calls.
[N, H, V, K]withV=K=128.initial_stateis passed as both the initial- and final-statekernel pointer. The kernel updates that allocation directly; there is no
state scratch allocation and no post-kernel
copy_/memcpy.(sequence, head). M64 assigns two CTAs whose64-row value partitions are disjoint. Each CTA loads all initial-state rows
it owns before storing those same final-state rows, making exact pointer
aliasing safe.
alias and rejects partial overlap.
initial_state=None, final-state storage is created only ifoutput_final_state=True.The generated source emits
__restrict__for every pointer and has noper-argument alias escape hatch, so the removal of
__restrict__from only thetwo state parameters is an explicitly delimited FlashInfer ABI integration
patch. The frozen-source test restores those two qualifiers, removes the
existing tensor-map acquire integration prologue, and verifies the remaining
generated source against its normalized SHA256.
CUDA graph and tensor-map safety
RecurrentKDAPrefillWorkspaceowns optional final-state scratch for callswithout an initial state, beta padding, and separate 768-byte M64/M128
tensor-map descriptor blocks. It binds to one stream, is warmed eagerly
against the exact tensor signature, and is used by one captured
recurrent_kdainvocation. A workspace that has participated in capture isrejected by later Python calls; graph replay remains valid because it does not
re-enter Python.
Tensor maps are prepared outside capture and published to stable global
storage on the caller stream. Every consumer CTA executes
fence.proxy.tensormap::generic.acquire.gpufor all six maps, followed by aCTA barrier, before any TMA instruction can consume them.
Other integration
seq_orderandprefill_workspace.CUDA graph tests.
including a commit- and binary-verified MoonshotAI/FlashKDA peer.
🧪 Validation
Final validation was run on pushed head
f6c3b0c486787e64f76b2be91ddb239ea3ba66ce, on an NVIDIA B200(compute capability 10.0), PyTorch
2.13.0.dev20260503+cu132, CUDA 13.2,and Python 3.13.13.
compute-sanitizer --tool memcheckERROR SUMMARY: 0 errorscompute-sanitizer --tool synccheckERROR SUMMARY: 0 errorsThe full repository test suite was not run.
Fresh B200 CUPTI benchmark
The candidate is invoked only through the public
flashinfer.kda_decode.recurrent_kdaentry point and performs its state updatein place inside the kernel. Its timed region contains no state memcpy.
The comparison uses six fixed/packed shapes with deterministic seeds, matching
tensor distributions, BF16 state, and
scale=1/sqrt(128). Preinitializedrotating state buffers ensure each timed candidate invocation receives the
same initial state. Allocation, metadata creation, sequence ordering,
state-pool reset, and JIT/cache warmup are outside the measured region.
The peer is
MoonshotAI/FlashKDA._fwd_raw, from source commitd2ff19a6and CUTLASS
5c149f52.The loaded extension SHA256 is
997c3a1d1338f8bf9dba3c1a01386b1b74448214c294d64409454cc11141c04c.The benchmark records the source revision and independently computed extension
digest.
Timing uses CUPTI activity tracing, cold-L2 flushing, no CUDA graph, 20 ms
warmup, and 100 ms measurement per block. Each value is the median of two
independent block medians; the reported scopes occupy symmetric
PR/raw/.../raw/PRpositions. Speedup isFlashKDA / this PR.[8192][1300, 547, 2048, 963, 271, 3063][1024] × 8[8192][1300, 547, 2048, 963, 271, 3063][1024] × 8All six output/state comparisons against FlashKDA passed at
atol=rtol=1e-2; the maximum observed absolute errors were0.0009765625for output and
0.015625for state. The benchmark JSON SHA256 isc34b27f9e8fb4e4dfbc500e6976e63668fcd0bff2bb00086cd516de59491f259.Upstream recurrent-KDA comparison
Pinned upstream main does not implement ordinary multi-token recurrent-KDA
prefill through its public API: fixed
T != 1is rejected, and packed standarddecode processes one token per sequence. Speculative decode has different
checkpoint/state semantics and is not an equivalent workload. The upstream
prefill comparison is therefore unsupported / N/A, rather than timing a
different operation.
🚀 Pull Request Checklist
Reviewer Notes
The main review surfaces are:
integrity normalization.
Related to #4254
Summary by CodeRabbit
New Features
recurrent_kdaAPI supporting decode, speculative decode, and prefill workflows.Documentation
Tests