You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
perf(cake_dense_projection_det_gemm): run the half-tail TMA-store program on the sm_103a 5120 x 17408 weight gradients at T 16231 / 16384 - #8522
This PR adds and updates row rules in the Cake dense_projection_det_gemm package in flashinfer.experimental, the deterministic dense projection GEMM. The package has 127 programs built from 124 templates, and this PR leaves both counts unchanged. Most programs are BF16-operand GEMMs (forward, dgrad and wgrad), and one is an FP32 router GEMM. The PR changes 3 row rules, all for sm_103a, and each applies at a single token count T. It adds one rule for the 5120→17408 weight gradient at T 16231. The other two are the empty sm_103a keys for the 17408→5120 weight gradient at T 16231 and 16384. The previous PR added those keys to keep that row's plan unchanged, and this PR gives them a real rule. All three rules route BF16-output weight-gradient rows to a program that already ships: the half-tail program with a TMA-store epilogue and 12-row raster groups. No program or template is added or removed, the rule count goes from 366 to 367, and sm_100a and sm_107a are unchanged.
The baseline is the previous Cake version of the package on main. Against it, the changed rows have a speedup of 1.013 on GB300 for all three rows and 1.042–1.045 on B300. No changed row is slower on either GPU, and outputs on those rows are bitwise identical to main. All other rows of the 294-row shape list plan the same programs as main.
Speedup of this PR over the previous Cake version on main. GB300 (sm_103a, 152 SMs), B300 (sm_103a, 148 SMs); CUDA 13.4, cold L2, BF16 output. Speedup = previous time / this PR time, and values above 1 mean this PR is faster.
shape (K→N)
op
out
T
GB300 previous µs
GB300 this PR µs
GB300 speedup
B300 previous µs
B300 this PR µs
B300 speedup
17408→5120
wgrad
bf16
16231
1408.1
1390.2
1.013
1612.8
1543.2
1.045
17408→5120
wgrad
bf16
16384
1420.7
1402.2
1.013
1625.6
1555.9
1.045
5120→17408
wgrad
bf16
16231
1427.5
1408.9
1.013
1620.3
1554.2
1.042
The package also ships FP32-output forms, weight-gradient (wgrad) programs and the FP32 router GEMM (run through BF16x3 emulation), on sm_100a, sm_103a and sm_107a.
Summary by CodeRabbit
Performance Improvements
Expanded exact-shape handling for dense projection weight-gradient operations on supported hardware, including additional input sizes and token counts.
Documentation
Updated the documented rule count and coverage for round 2.
…gram on the sm_103a 5120 x 17408 weight gradients at T 16231 / 16384
Re-measured on GB300 with a cooled timing frame, the half-tail program
through the one-slot TMA store with 12-row raster groups (_ht_tma1) is
faster than the row's plan without a key on three sm_103a weight-gradient
cells that round 2 left out:
- the 17408 -> 5120 weight gradient at T 16231 / 16384: its two empty
sm_103a keys now carry the sm_100a key's rule (restated, so an
sm_103a rule keeps winning over the parent lookup);
- the 5120 -> 17408 weight gradient at T 16231: a new T-exact key next to
the existing T 16384 key.
Each key keeps the per-tile K partitions of the row's plan without it, so
every output is bitwise equal to that plan's; every other T, sm_100a and
sm_107a are unchanged. No program text changes. Tests: the sm_103a rule
count, and the T-exact weight-gradient planner and device tests now cover
the three sm_103a rows (program, direct orientation, whole-K plan, FP64
reference, bitwise equality with the plan without the key).
@flashinfer-bot run tests/experimental/test_cake_dense_projection_det_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_gemm.py
No actionable comments were generated in the recent review.
Walkthrough
The planner adds active exact-token weight-gradient rules for additional sm_103a shape and token-count combinations. Tests and documentation update their expected rule coverage.
Changes: Exact-token weight-gradient rules
Layer / File(s)
Summary
Add active sm_103a rules flashinfer/experimental/dense_projection_det_gemm/cake_backend.py
The sm_103a rules add the 17408→5120 orientation at T=16231 and T=16384, and add T=16231 for 5120→17408. The 17408→5120 entries select the 256-row, half-height-tail, one-slot TMA-store program with 12-row raster groups.
Update rule tests and documentation tests/experimental/test_cake_dense_projection_det_gemm.py, flashinfer/experimental/dense_projection_det_gemm/README.md
Tests update the rule counts and require active templates for every listed key. The README updates the round-2 rule count and coverage description.
flashinfer-ai/flashinfer#8491: This PR extends the same sm_103a weight-gradient rule family by activating previously empty keys and adding a T=16231 case.
flashinfer-ai/flashinfer#8441: This PR builds on the exact-token lookup and sm_103a inheritance behavior established by PR #8441.
Suggested labels:run-ci
Suggested reviewers:yzh119
Merge Risk: Minimal · up to 19be3
No concrete merge-blocking issue is established for the added sm_103a rules. Merge readiness remains subject to normal test completion; the supplied CI excerpts do not establish a code failure.
Pre-merge checks: 3 · 2
Failed checks: 2 warnings
Check name
Status
Explanation
Resolution
Description check
The description gives detailed scope, rule changes, performance results, correctness results, and test context. However, it does not follow the repository template and omits the required experimental-…
Update the description to use the repository template. Add the related tracking issue, experimental owner and graduation plan, completed checklist items, and an experimental-tests fence containing the intended targets under tests/experiment…
Docstring Coverage
Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…
Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed checks: 3 passed
Check name
Status
Explanation
Title check
The title clearly identifies the performance change, backend, operation, hardware target, and token counts. It is specific and related to the main changeset, although it is longer than ideal.
Linked Issues check
Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check
Check skipped because no linked issues were found for this pull request.
Full details: Description check
Explanation
The description gives detailed scope, rule changes, performance results, correctness results, and test context. However, it does not follow the repository template and omits the required experimental-track details, including the tracking issue, owner, graduation plan, checklist status, and experimental test targets.
Resolution
Update the description to use the repository template. Add the related tracking issue, experimental owner and graduation plan, completed checklist items, and an experimental-tests fence containing the intended targets under tests/experimental/. State the pre-commit and test results explicitly.
Full details: Docstring Coverage
Explanation
Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 2 functions across 2 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
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.
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 6 remain after this review.
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
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.
This PR adds and updates row rules in the Cake
dense_projection_det_gemmpackage inflashinfer.experimental, the deterministic dense projection GEMM. The package has 127 programs built from 124 templates, and this PR leaves both counts unchanged. Most programs are BF16-operand GEMMs (forward, dgrad and wgrad), and one is an FP32 router GEMM. The PR changes 3 row rules, all forsm_103a, and each applies at a single token count T. It adds one rule for the 5120→17408 weight gradient at T 16231. The other two are the emptysm_103akeys for the 17408→5120 weight gradient at T 16231 and 16384. The previous PR added those keys to keep that row's plan unchanged, and this PR gives them a real rule. All three rules route BF16-output weight-gradient rows to a program that already ships: the half-tail program with a TMA-store epilogue and 12-row raster groups. No program or template is added or removed, the rule count goes from 366 to 367, andsm_100aandsm_107aare unchanged.The baseline is the previous Cake version of the package on main. Against it, the changed rows have a speedup of 1.013 on GB300 for all three rows and 1.042–1.045 on B300. No changed row is slower on either GPU, and outputs on those rows are bitwise identical to main. All other rows of the 294-row shape list plan the same programs as main.
Speedup of this PR over the previous Cake version on main. GB300 (
sm_103a, 152 SMs), B300 (sm_103a, 148 SMs); CUDA 13.4, cold L2, BF16 output. Speedup = previous time / this PR time, and values above 1 mean this PR is faster.The package also ships FP32-output forms, weight-gradient (wgrad) programs and the FP32 router GEMM (run through BF16x3 emulation), on
sm_100a,sm_103aandsm_107a.Summary by CodeRabbit