perf(cake_tinygemm2): add a STAGES=16 kernel tier for single-wave large-K shapes - #4423
Merged
bkryu merged 1 commit intoAug 11, 2026
Merged
Conversation
…sor-map ABI Regenerated all variants from the current Loom generator state and merged them into the single TU following the same mechanical transform as the original freeze: - adds the STAGES=16 ring as a third tier: at long K the 8-deep ring's in-flight budget no longer covers the weight-stream latency once the working set is not L2-resident (cold-L2 CUPTI on GB300: +0.7/+1.8/+4.4us vs the 16-deep configuration at K=7168/14336/28672, ~0.2us warm cost); ring selection moves into the binding, mirroring the reference launcher convention — single-wave grids with K>=4608 take stage 16, multi-wave grids keep stage 8 (halved residency loses 2-6us) - all six kernels use the generator's current trailing by-value __grid_constant__ tensor-map pack ABI (one LoomTensorMapPack<2>) - dynamic-SMEM opt-in set once per (kernel, device); the SM100/SM103 verdict cached per device - python dispatch becomes a single hop into the combined op; dispatch gates on exact compute capabilities (10, 0)/(10, 3) so other 10.x devices keep the reference path instead of erroring in the binding - tests: stage16 variants join the per-variant battery; parity shapes add single-wave long-K rows (8, 128, 7168) and (1, 128, 14336) Correctness contract unchanged: bitwise equality with csrc/tinygemm2.cu. On GB300 the six instances are bit-identical to the reference across a 40-shape battery (batch 1-64, K to 7168, M to 4096); at N=8/M=128/K=7168 the new tier measures 4.93us vs the reference's 5.22us (cold-L2 CUPTI).
Contributor
📝 WalkthroughWalkthroughThe PR adds stage16 and stage16 PDL SM100 tinygemm2 kernels. It packs tensor-map arguments, centralizes stage selection in the CUDA binding, updates router GEMM capability checks, and expands parity coverage for large-K inputs. ChangesSM100 tinygemm2 kernel implementation
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant RouterGEMM
participant SM100Binding
participant StageSelector
participant TinyGEMM2Kernel
participant Output
RouterGEMM->>SM100Binding: submit GEMM arguments
SM100Binding->>StageSelector: evaluate reduction size and CTA count
StageSelector->>TinyGEMM2Kernel: launch selected stage and PDL variant
TinyGEMM2Kernel->>Output: store bfloat16 GEMM results
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
xslingcn
marked this pull request as ready for review
August 8, 2026 22:40
xslingcn
requested review from
Aneureka,
StudyingShao,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
August 8, 2026 22:40
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
csrc/tinygemm2_sm100.cu (1)
2199-2212: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winValue-initialize
configandattrs.
cudaLaunchConfig_t config;andcudaLaunchAttribute attrs[1];are declared without initialization. The code assigns every field that the current CUDA header defines, so the launch is correct today. If a later CUDA toolkit adds a field tocudaLaunchConfig_t, the unassigned field carries stack garbage andcudaLaunchKernelExfails in a way that is hard to diagnose. Value-initialization removes that dependency at no runtime cost.♻️ Proposed hardening
if (pdl) { - cudaLaunchConfig_t config; - cudaLaunchAttribute attrs[1]; + cudaLaunchConfig_t config = {}; + cudaLaunchAttribute attrs[1] = {}; config.gridDim = grid;🤖 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/tinygemm2_sm100.cu` around lines 2199 - 2212, Value-initialize the cudaLaunchConfig_t config and cudaLaunchAttribute attrs declarations in the pdl launch path before assigning their fields. Update the declarations used by cudaLaunchKernelEx so any future CUDA struct fields default safely without changing the existing launch configuration.tests/model_optimizations/test_tinygemm2_sm100.py (1)
52-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a shape at the exact
kStage16MinKboundary.The new entries cover K=7168 and K=14336, which land well inside the stage16 tier, and K=4096, which lands in the stage8 tier. No entry sits at K=4608, the exact value of
kStage16MinKincsrc/tinygemm2_sm100.cu. The selection usesin_features >= kStage16MinK, so the boundary itself is the value most likely to break if the constant or the comparison changes. One extra entry pins it.As per coding guidelines: "Write tests for new operations".
♻️ Proposed additional shape
(8, 1024, 1024), + (8, 128, 4608), (8, 128, 7168),🤖 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/model_optimizations/test_tinygemm2_sm100.py` around lines 52 - 63, Add a PARITY_SHAPES entry with K equal to the exact kStage16MinK boundary value, 4608, while keeping the existing shapes unchanged. Ensure the new case exercises the selection path that uses in_features >= kStage16MinK.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@csrc/tinygemm2_sm100.cu`:
- Around line 2199-2212: Value-initialize the cudaLaunchConfig_t config and
cudaLaunchAttribute attrs declarations in the pdl launch path before assigning
their fields. Update the declarations used by cudaLaunchKernelEx so any future
CUDA struct fields default safely without changing the existing launch
configuration.
In `@tests/model_optimizations/test_tinygemm2_sm100.py`:
- Around line 52-63: Add a PARITY_SHAPES entry with K equal to the exact
kStage16MinK boundary value, 4608, while keeping the existing shapes unchanged.
Ensure the new case exercises the selection path that uses in_features >=
kStage16MinK.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 5808b9d0-3027-4580-98bc-1dbe9965ad45
📒 Files selected for processing (3)
csrc/tinygemm2_sm100.cuflashinfer/gemm/routergemm.pytests/model_optimizations/test_tinygemm2_sm100.py
Collaborator
|
/bot run tests/gemm |
Collaborator
bkryu
approved these changes
Aug 11, 2026
bkryu
enabled auto-merge (squash)
August 11, 2026 06:53
Collaborator
|
[SUCCESS] Pipeline #62018422: 18/18 executed test jobs passed |
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
3 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
Adds a third ring-depth tier to the
tinygemm2_sm100family (#4274) andmoves 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.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).
csrc/tinygemm2.culauncher convention: stage 4 for K <= 1024 or gridspast 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.
passes the previous
major == 10predicate but must keep the referencepath instead of erroring in the binding.
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), andend-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.pyextends to the stage-16variants 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):🔍 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
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, stageselection, launch attributes), the dispatcher in
flashinfer/gemm/routergemm.py, and the test additions. The single-TUmerge follows the mechanical transform documented in the file header.
Summary by CodeRabbit
New Features
Bug Fixes
Tests