Repository navigation
perf(cake_nvfp4_svdquant): K=96 MMA builds on sm_103a, aliased-inbox split-K and M 129 tuning - #7516
Merged
Merged
Conversation
…oute rows Registry records of the two-CTA programs built on the SM103 K=96 block-scaled MMA (shipped for sm_103a only), of a register-bias-path build of the whole-tile program, of SM100 64-row builds with the 256-byte L2 promotion on the 15-stage ring (shipped for sm_100a only) and of split-K builds whose reduce-scatter inbox shares the ring with the L2 promotion, and the per-architecture route and raster-group rows of the SM103 and SM100 problems that measured faster with them; the records those rows no longer reach are dropped. raster_group() takes the device architecture like select_record(). The backend tests follow the new rows. Co-Authored-By: Claude <[email protected]>
Device and binding units of the SM103 K=96 two-CTA programs, the whole-tile register-bias-path program, the SM100 L2-promoted 64-row programs and the aliased-inbox split-K programs, the regenerated units of the retained programs and the build manifest without the dropped programs, rendered twice with identical output. Co-Authored-By: Claude <[email protected]>
…ts previous build The SM103 route of the (129, 3072, 3072, 64, bias) problem returns to the L2-promoted 64-row build it used before; the 16-stage aliased-ring build keeps the rank-96 and rank-128 rows. No program changes. The backend tests follow the route. Co-Authored-By: Claude <[email protected]>
yyihuang
requested review from
Anerudhan,
Aneureka,
Jingfan-Sun,
StudyingShao,
YangXu1990uiuc,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
dierksen,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
righthandabacus,
saltyminty,
samuellees,
sricketts,
yanqinz2,
yongwww and
yzh119
as code owners
October 27, 2026 00:44
Contributor
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
CUDA kernel variants and bindings flashinfer/cake_nvfp4_svdquant_gemm_backend/csrc/*_binding.cu, flashinfer/cake_nvfp4_svdquant_gemm_backend/csrc/*_kernel.cu |
Generated bindings add checked host-side integer arithmetic and updated tensor-map setup. Kernel variants add K=96 block16 MMA paths, direct BF16 output stores in selected variants, and revised TMA/cluster barrier phase tracking and stage handling. |
Build catalog and architecture dispatch flashinfer/cake_nvfp4_svdquant_gemm_backend/build.json, flashinfer/cake_nvfp4_svdquant_gemm_backend/cake_jit.py, tests/gemm/test_cake_nvfp4_svdquant_gemm_backend.py |
The manifest and JIT records are refreshed with new artifacts and SM103a K=96 variants. Exact routes and raster-group overrides are architecture-specific. Tests check the added mappings and build options. |
Priority:
Estimated code review effort: 5 (Critical) | ~120 minutes
Sequence Diagram(s)
sequenceDiagram
participant Caller
participant CakeJIT
participant TVMFFIBinding
participant CUDAKernel
Caller->>CakeJIT: select build using shape and architecture
CakeJIT->>TVMFFIBinding: provide selected kernel and launch plan
TVMFFIBinding->>CUDAKernel: validate inputs and launch
CUDAKernel-->>Caller: write BF16 output
Suggested reviewers: yzh119
Merge Risk: 
2cbb5
-
This PR adds and re-routes tuned GEMM builds for sm_100a and sm_103a, with outputs reported as bitwise identical to the previous package. The review found no concrete defect. The one thing static review cannot check is the K=96 MMA encodings, which depend on the author's reported runs and the requested test job.
Pre-merge checks: 
-
Check name Status Explanation Resolution Description check The description states the purpose, the kernel variant, the affected shapes, and the per-shape benchmark results for GB300 and GB200. It does not follow the repository template. It omits the Related I… Add a Related Issues entry (or state none). Complete the Pre-commit Checks and Tests checklist items, and state whether tests/gemm/test_cake_nvfp4_svdquant_gemm_backend.pypassed. Include the test output or a link to the bot run requested…Docstring Coverage Docstring coverage is 2.96% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Documentation-only edits and deletions within survi… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Title check The title names the main change: K=96 MMA builds for sm_103a, aliased-inbox split-K, and M=129 tuning. It fits the changeset, which adds K=96 sm_103a builds, split-K builds for M=129 K=12288 shapes, a… 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: Title check
-
Explanation
The title names the main change: K=96 MMA builds for sm_103a, aliased-inbox split-K, and M=129 tuning. It fits the changeset, which adds K=96 sm_103a builds, split-K builds for M=129 K=12288 shapes, and retuned M=129 K=3072 routes. The title uses the conventional type(scope) prefix and reads as one sentence.
-
Full details: Description check
-
Explanation
The description states the purpose, the kernel variant, the affected shapes, and the per-shape benchmark results for GB300 and GB200. It does not follow the repository template. It omits the Related Issues section, the pre-commit and test checklist items, and the Reviewer Notes section. The experimental section is correctly omitted because the change is not experimental. The description also does not say whether the tests were run; the PR objectives note a requested test run with no result attached.
Resolution
Add a Related Issues entry (or state none). Complete the Pre-commit Checks and Tests checklist items, and state whether
tests/gemm/test_cake_nvfp4_svdquant_gemm_backend.pypassed. Include the test output or a link to the bot run requested in the PR discussion. Add a short Reviewer Notes section if reviewers should focus on specific areas, such as theraster_groupsignature change toarch: str | Nonethat can raiseTypeError, or the sm_100a dispatch-only changes.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 2.96% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 270 functions across 46 files. (41 skipped: 1 unsupported, 4 too large, 36 over the file limit.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches: 💡 1
-
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
🧪 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.
Comment @coderabbitai help to get the list of available commands.
Collaborator
Author
|
@flashinfer-bot run tests/gemm/test_cake_nvfp4_svdquant_gemm_backend.py |
1 task done
yzh119
approved these changes
Oct 27, 2026
yyihuang
enabled auto-merge (squash)
October 27, 2026 05:29
This was referenced Oct 27, 2026
Merged
yyihuang
added a commit
that referenced
this pull request
Oct 28, 2026
…3072, 12288, rank 32) (#7575) `flashinfer.mm_nvfp4_svdquant(..., backend="cake")` computes the NVFP4 SVDQuant GEMM `out = BF16(alpha * (deq(A) @ deq(B)^T + D @ L1^T) + bias)`. In this kernel, A and B are NVFP4 (E2M1) with E4M3 block scales. The low-rank D / L1, the bias and the output are BF16, and accumulation is FP32. Its Cake backend has one data-type / layout variant with two fused forms (with and without the fused bias), across LoRA ranks 32–128. Large-M shapes run on a 2-CTA persistent program or a 2-CTA tail-split program. M = 129 shapes run on a 64-row program or a swap-AB split-K program. This PR changes one shape on sm_100a, (M, N, K, rank, bias) = (4096, 3072, 12288, 32, 1), which gets a new build of the 2-CTA tail-split program. On 152 SMs (76 CTA pairs), this shape has 192 output tiles of 256 x 256: two full waves and a 40-tile tail. The previous build ran the tail as a third wave of 40 whole tiles on 40 of the 76 pairs. The new build runs the tail as one wave of 76 units. Eight of them are 192-column pieces, made by splitting each of two groups of three adjacent tiles into four. The other 68 are half tiles of 128 columns. The dispatch table routes only this shape to the new build, and every other shape keeps its previous kernel (‡ rows). Outputs are bitwise identical to the previous package on every shape (same K order, FP32 accumulation, single BF16 rounding). The baseline is the previous Cake version of this kernel on FlashInfer main, shipped in #7516. Against it, the changed shape is 1.064× (session s1, table below) / 1.068× (session s2) faster on GB200. **NVIDIA GB200 (sm_100a), driver 580.126.20, CUDA 13.0 runtime, nvcc 13.4, torch 2.14.1+cu130; median kernel time over 5 ABBA groups x 400 ms, CUPTI kernel-active, cold L2; previous Cake = the previous Cake version on FlashInfer main, this PR = this package; speedup = previous / this PR; session s1.** | M | N | K | rank | bias | previous Cake (µs) | this PR (µs) | speedup vs previous Cake | |---:|---:|---:|---:|---:|---:|---:|---:| | 4096 | 3072 | 12288 | 32 | 1 | 62.98 | 59.17 | 1.064 | | 129 | 3072 | 3072 | 32 | 1 | 5.89 | 5.89 | 1.000 ‡ | | 129 | 3072 | 3072 | 32 | 0 | 5.82 | 5.82 | 1.000 ‡ | | 129 | 3072 | 3072 | 64 | 1 | 5.95 | 5.95 | 1.000 ‡ | | 129 | 3072 | 3072 | 96 | 1 | 6.11 | 6.11 | 1.000 ‡ | | 129 | 3072 | 3072 | 128 | 1 | 6.21 | 6.21 | 1.000 ‡ | | 129 | 3072 | 3072 | 128 | 0 | 6.08 | 6.08 | 1.000 ‡ | | 129 | 3072 | 12288 | 32 | 0 | 10.70 | 10.72 | 0.999 ‡ | | 129 | 3072 | 12288 | 32 | 1 | 10.72 | 10.72 | 1.000 ‡ | | 129 | 3072 | 12288 | 128 | 0 | 10.82 | 10.82 | 1.000 ‡ | | 129 | 3072 | 12288 | 128 | 1 | 10.82 | 10.82 | 1.000 ‡ | | 6912 | 3072 | 3072 | 32 | 1 | 28.42 | 28.61 | 0.994 ‡ | | 6912 | 3072 | 3072 | 64 | 1 | 28.83 | 28.94 | 0.996 ‡ | | 6912 | 3072 | 3072 | 96 | 1 | 30.18 | 30.14 | 1.001 ‡ | | 6912 | 3072 | 3072 | 128 | 1 | 30.72 | 30.70 | 1.001 ‡ | | 6912 | 3072 | 12288 | 32 | 1 | 98.67 | 98.70 | 1.000 ‡ | | 6912 | 3072 | 12288 | 128 | 1 | 101.02 | 101.01 | 1.000 ‡ | | 4096 | 3072 | 3072 | 32 | 1 | 19.49 | 19.49 | 1.000 ‡ | | 4096 | 12288 | 3072 | 32 | 1 | 64.19 | 64.24 | 0.999 ‡ | | 6889 | 3072 | 3072 | 32 | 1 | 28.16 | 28.10 | 1.002 ‡ | | 6889 | 3072 | 12288 | 32 | 1 | 98.62 | 98.62 | 1.000 ‡ | | 6889 | 12288 | 3072 | 32 | 1 | 104.54 | 104.53 | 1.000 ‡ | | 9216 | 3072 | 3072 | 32 | 1 | 34.91 | 34.85 | 1.002 ‡ | | 9216 | 3072 | 12288 | 32 | 1 | 124.32 | 124.22 | 1.001 ‡ | | 9216 | 12288 | 3072 | 32 | 1 | 136.48 | 136.56 | 0.999 ‡ | | 16384 | 3072 | 3072 | 32 | 1 | 64.10 | 63.87 | 1.004 ‡ | | 16384 | 3072 | 12288 | 32 | 1 | 207.89 | 207.86 | 1.000 ‡ | | 16384 | 12288 | 3072 | 32 | 1 | 230.96 | 230.99 | 1.000 ‡ | ‡ This shape keeps the previous package's kernel on this GPU; its ratio is the same-kernel reading of the session. This kernel has no other Cake variant, and the family ships on sm_100a (GB200) and sm_103a (GB300), with sm_103a unchanged by this PR. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Added a specialized execution schedule for a rank-32 matrix multiplication workload on supported SM100a hardware. * Schedule selection now considers the device’s SM count alongside workload shape and launch requirements. Eligible devices use the specialized schedule; other cases continue to use an existing schedule. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Avery Huang <[email protected]> Co-authored-by: Claude <[email protected]>
yyihuang
added a commit
that referenced
this pull request
Oct 28, 2026
…test expectations (#7632) ## Summary This PR fixes three test failures from the scheduled nightly runs of the test suite. It changes no kernel sources: - On toolkits older than CUDA 13.4, NVFP4 context attention programs now fail closed instead of failing in ptxas. - An SVDQuant tail-split test now checks the actual schedule rule. - The warp-decode PDL test now reads graph edges without a cuda-bindings use-after-free. ## Changes **NVFP4 context programs require CUDA 13.4** (`flashinfer/cake_fmha.py`, `tests/attention/test_cake_fmha_context_programs.py`) - The NVFP4 programs under `csrc/cake_fmha_context_nvfp4/` (Q511 and Q511_SCALE_PTR, SM100A/SM103A) dequantize E2M1 pages with the PTX ISA 9.4 `mul.rn.satfinite.e4m3x4.e2m1x4.e4m3x4`. Only CUDA 13.4+ supports this instruction. With CUDA 12.9 and 13.0, ptxas rejects it (`Illegal modifier '.satfinite' for instruction 'mul'`). This failed `test_cake_context_quantized_programs_match_reference[case2-Q511]`, `[case3-Q511_SCALE_PTR]` and the six `test_cake_context_programs_match_the_jit_route_bitwise[nvfp4_*]` rows on B200, GB200 and GB300. CUDA 13.4 passed. - Add `CONTEXT_NVFP4_PROGRAMS_MIN_CUDA = "13.4"` and `context_nvfp4_programs_unavailable()`, which returns a reason string or `None` via `is_cuda_version_at_least`. - On an older toolkit, `get_cake_fmha_context_module` routes an NVFP4 program route to compat_v1 with a `RuntimeWarning` ("failed closed to compat_v1"), as the NVFP4 JIT route already does. `_load_cake_fmha_context_program` raises a `RuntimeError` naming the program and the reason instead of launching a build that ptxas rejects. `flashinfer/jit/cake_dsv4.py` already requires CUDA 13.4 for the same instruction in the DSv4 NVFP4 decode kernels. - Tests: on such toolkits, `_skip_without_nvfp4_programs()` skips the NVFP4 rows and states the reason. The new `test_cake_context_nvfp4_programs_fail_closed_without_their_toolkit` monkeypatches the check. It asserts that the loader raises and that the route falls back to compat_v1 with the warning. **SVDQuant tail-split routing test** (`tests/gemm/test_cake_nvfp4_svdquant_gemm_backend.py`) - `test_tailsplit_rows_split_on_this_device` required every `pair_tailsplit` row to split its tail (`split == 2 and full < units`) on 148/152-SM devices. Since #7516, the row M 6889, N 12288, K 3072, rank 32, bias dispatches there on sm_100a. - On 148 SMs (74 CTA pairs) the row has 1296 pair tiles (27 × 48): 17 full waves plus 38 tail tiles. - 2 × 38 = 76 half units exceed 74 pairs, so `tailsplit_units` keeps whole tiles (`tail_split = 1`). - B200 therefore failed with `AssertionError: ('M6889-N12288-K3072-r32-bias', 148)` under CUDA 13.0 and 13.4. - The test now asserts the schedule rule. A partial last wave must exist (`tail > 0`). If `2 * tail <= pairs`, the tail becomes `2 * tail` half units; otherwise all tiles stay whole. - `test_tailsplit_units_schedule` adds `units(m=6889, n=12288, sm_count=148) == (1296, 1296, 1296, 1)` and `units(m=6889, n=12288, sm_count=152) == (1296, 1300, 1292, 2)`. - There is no loader or kernel change, and the #7516 dispatch stays as is. **Warp-decode PDL graph edges** (`tests/moe/test_cake_warp_decode_pdl_gpu.py`) - `test_first_kernel_launch_dependency[direct|route_pack|fused_route_pack]` read captured-graph edges with `cuda.bindings.driver.cuGraphGetEdges` (`cuGraphGetEdges_v2` in the CUDA 12 bindings). - cuda-bindings before 12.9.7 and 13.3.0 return `CUgraphEdgeData` wrappers backed by a buffer that is freed when the call returns (NVIDIA/cuda-python#1804, fixed in NVIDIA/cuda-python#2110). - In the CUDA 13.0 environment (cuda-bindings 13.0.3), the test saw edge types 255, 250, 76, 178, 40, 57, ... instead of 1 (programmatic) or 0 (default). - The 12.9.7 and 13.4.2 bindings passed. - The new `_graph_edges(graph_handle)` calls `cuGraphGetEdges_v2` from `libcuda.so.1` via ctypes. It uses a `CUgraphEdgeData` Structure (`from_port`, `to_port`, `type`, `reserved[5]`). - `_kernel_graph` still uses cuda-bindings for node enumeration, node types and kernel names, which the bug does not affect. ## Validation Validation is in flight and will be summarized here once the GPU matrix (B200, GB200, GB300 with the CUDA 12.9, 13.0 and 13.4 environments) on the three test files and the public CI run have completed. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Compatibility** * NVFP4 context programs require CUDA 13.4 or newer. On older toolkits, loading these programs reports the version requirement, and module selection warns before falling back to a compatible implementation. This makes the toolkit limitation clearer and allows supported workloads to use the fallback where available. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Avery Huang <[email protected]> Co-authored-by: Claude Fable 5.1 <[email protected]>
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
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.
flashinfer.mm_nvfp4_svdquant(..., backend="cake")computes the NVFP4 SVDQuant GEMMout = BF16(alpha * (deq(A) @ deq(B)^T + D @ L1^T) + bias). In this kernel, A and B are NVFP4 (E2M1) with E4M3 block scales. The low-rank D / L1, the bias and the output are BF16, and accumulation is FP32. Its Cake backend has one data-type / layout variant with two fused forms (with and without the fused bias), across LoRA ranks 32 / 64 / 96 / 128. That variant ships as the vendorable packageflashinfer/cake_nvfp4_svdquant_gemm_backend.This PR is a further tuning round of the same variant. Its outputs are bitwise identical to the previous package on every shape. Shapes below are written as (M, N, K, rank, bias).
kind::mxf4nvf4.block16) with K = 96 instead of K = 64, as K 96 + 96 + 64 per 256-wide K stage. On (16384, 3072, 12288, 32, 1), the tail-split program uses whole K = 768 units.Every other shape keeps its previous kernel build, marked ‡ below. Compared with the previous Cake release, the changed shapes run 1.023–1.095× faster on GB300 (22 shapes) and 1.005–1.046× faster on GB200 (14 shapes, this session). "Previous Cake" is the previous Cake release of this kernel, measured in the same session as this PR's package.
Speedup of this PR's Cake kernel over the previous Cake kernel on NVIDIA GB300 (sm_103a), CUDA 13.4 nvcc (CUDA 13.0 runtime), median kernel time (CUPTI kernel-active, cold L2), session s1
‡ This shape keeps the previous package's kernel on this GPU; its ratio is the same-kernel reading of the session.
Speedup of this PR's Cake kernel over the previous Cake kernel on NVIDIA GB200 (sm_100a), CUDA 13.4 nvcc (CUDA 13.0 runtime), median kernel time (CUPTI kernel-active, cold L2), session s1
‡ This shape keeps the previous package's kernel on this GPU; its ratio is the same-kernel reading of the session.
This kernel has no other Cake variant, and the family ships on sm_100a and sm_103a.
Summary by CodeRabbit
New Features
Bug Fixes