Repository navigation
perf(cake_gated_act_mxfp8): register-fitted band-path and packed-sigmoid programs in eight M * K ranges (SM100 / SM103) - #7568
Merged
yyihuang merged 2 commits intoOct 28, 2026
Conversation
…oid programs in eight M * K ranges (SM100 / SM103) Eight per-architecture M * K ranges (sm_100a 5, sm_103a 3) now select a program whose register budget is fitted to a CTA count per SM. Below M * K = 2^20, backward rowwise and backward both take a 32 x 32 tile band-path program on sm_100a and sm_103a. On sm_100a, forward colwise, backward colwise and backward both split at 2^27; below the split they take a band-path program from 2^26 (forward) or 2^24 (backward). On sm_103a, backward rowwise L1::no_allocate from 2^26 moves from the 12-CTA build, which spills, to the packed-sigmoid build at 10 CTAs. Outputs are byte-identical; the public API is unchanged. The package goes from 45 to 50 programs (6 new, 1 removed); sm_107a programs are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
yyihuang
requested review from
Anerudhan,
Aneureka,
Jingfan-Sun,
YangXu1990uiuc,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
dierksen,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
righthandabacus,
saltyminty,
samuellees,
sricketts,
yongwww and
yzh119
as code owners
October 27, 2026 09:30
flashinfer-bot
added
the
op: misc
1 task done
Contributor
|
Important Review skippedReview was skipped as selected files did not have any reviewable changes.
 |
| Layer / File(s) | Summary |
|---|---|
CUDA kernels and FFI bindings flashinfer/cake_gated_act_mxfp8_backend/csrc/* |
New CUDA kernels compute gated-activation outputs or gradients and quantize results to E4M3 with scales. Their TVM-FFI bindings validate tensors and launch kernels. An existing kernel and binding are renamed; the kernel’s arithmetic, reductions, and launch setting are also updated. |
Architecture-specific library entries and dispatch flashinfer/cake_gated_act_mxfp8_backend/build.json, flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py |
The manifest adds and removes sm_100a and sm_103a library entries. The host registry updates architecture-specific forward and backward records, route membership, and selection thresholds. |
Priority:
Estimated code review effort: 4 (Complex) | ~60 minutes
Change: Refactor
Sequence Diagram(s)
sequenceDiagram
participant Registry as cake_jit registry
participant Binding as TVM-FFI binding
participant Kernel as CUDA kernel
Registry->>Binding: Select and invoke architecture-specific kernel
Binding->>Kernel: Validate tensors and launch on the current stream
Kernel-->>Binding: Write MXFP8 outputs and scales
Suggested labels: op: misc, run-ci
Suggested reviewers: sricketts, aleozlx, yongwww
Merge Risk: 
db076
-
This change retunes which fused SwiGLU plus MXFP8 quantization programs are selected for specific input sizes on Blackwell GPUs, and adds the supporting kernels and bindings. No concrete defect was found in the supplied evidence. Registry-to-manifest consistency was checked, and the public API and output bytes are reported unchanged. Normal test runs should confirm it before merge.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 4.55% 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 clearly identifies the performance change: registering band-path and packed-sigmoid programs for SM100 and SM103 across selected M × K ranges. Description check The description provides a detailed summary of the kernel-selection changes, affected architectures, performance rationale, benchmark results, package changes, and API/output compatibility. It does no… 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 provides a detailed summary of the kernel-selection changes, affected architectures, performance rationale, benchmark results, package changes, and API/output compatibility. It does not include the formal template headings, related issue links, or explicit test and pre-commit checklist results, but the main description is sufficiently complete.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 4.55% 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 66 functions across 11 files. (1 skipped: 1 unsupported.)
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/utils/test_cake_gated_act_mxfp8.py tests/utils/test_cake_gated_act_mxfp8_api.py |
yzh119
approved these changes
Oct 27, 2026
yyihuang
enabled auto-merge (squash)
October 27, 2026 14:33
Collaborator
Author
|
@flashinfer-bot run tests/utils/test_cake_gated_act_mxfp8.py tests/utils/test_cake_gated_act_mxfp8_api.py |
yyihuang
added a commit
that referenced
this pull request
Oct 28, 2026
…ge SM100 problems and band-path fallbacks (SM100 / SM107) (#7660) The Cake fused SwiGLU + MXFP8 quantization kernel computes silu(gate) * up from a BF16 input [M, 2K] and quantizes the result to MXFP8. Each group of 32 E4M3 values gets one E8M0 scale (RCEIL rounding), and the scales are stored in the swizzled 128 x 4 layout. The kernel has one data type and six forms. The forward form is `flashinfer.silu_and_mul_mxfp8_quantize`, with output [M, K]. The backward form is `flashinfer.silu_and_mul_mxfp8_quantize_backward`. It computes the gate and up gradients [M, 2K] from the input [M, 2K] and grad_output [M, K], and quantizes them to MXFP8 in the same way. Each one can write rowwise, colwise or both outputs. `flashinfer.cake_gated_act_mxfp8_backend` runs one program per call and picks it by M * K for each architecture. This PR makes 10 changes to program selection: 6 on sm_100a, 4 on sm_107a and none on sm_103a. Every changed range now runs a program that uses one of the guarded-sigmoid band-path forms from #7492. Each of these programs also has a register budget fitted to a number of CTAs per SM. The three forms are a loop over only the flagged pairs, a non-inlined call for each flagged pair, and an inline branch for each pair. On sm_100a, four large-problem ranges get a new program that also uses the packed f32x2 sigmoid form from #7396. Forward rowwise from M * K = 2^28 keeps `no_allocate` loads and 8 CTAs per SM, and adds the inline form. Forward colwise from 2^27 keeps two rows per thread and `no_allocate` loads, and runs at 5 CTAs per SM with the call form. Forward both from 2^24 keeps two rows per thread, `no_allocate` loads and 4 CTAs per SM, and adds the flagged-pair loop. Backward rowwise from 2^28 keeps two rows per thread and `no_allocate` loads, and runs at 7 CTAs per SM with the call form. Backward colwise no longer splits at 2^27. Its program for [2^24, 2^27) already used the packed sigmoid and the call form at 6 CTAs per SM. That program now serves every size from 2^24 up. The removed program for 2^27 and up had neither the packed sigmoid nor a band path. Below 2^21, the backward colwise fallback keeps its 32 x 32 tile and runs at 12 CTAs per SM with the inline form. On sm_107a, all four changes are in the fallback programs, which serve the smallest sizes of each route. Forward rowwise below 2^23 takes a new program at 4 CTAs per SM with the inline form. Forward colwise gets a new split at 2^20. Below the split, a new program runs at 4 CTAs per SM with the inline form. The previous fallback keeps [2^20, 2^23). Below 2^19, backward colwise and backward both take new 32 x 32 tile programs. Backward colwise runs at 9 CTAs per SM with the inline form, and backward both runs at 8 CTAs per SM with the flagged-pair loop. The package goes from 50 to 54 programs: 9 are new and 5 are removed. The same arithmetic runs on the same operands, so every output byte is identical. The public API is unchanged. Block-scaled inputs gain most, because more warps take the guarded path there. With block-scaled inputs, all eight B200 rows that change program run faster, by up to 1.135x. All four such R200 rows also run faster, by up to 1.068x. With unit normals, three B200 rows run faster and all other rows stay within run-to-run spread. No row on any GPU runs slower by more than its run-to-run spread. The baseline is the previous Cake version of this kernel, which is the package FlashInfer main ships (#7568). Both versions run in the same process on the same GPU, in alternating paired runs, with a cold L2 before every launch. Each table row is one shape and pass. There are two column groups: unit-normal inputs, and block-scaled inputs (normals with per-32-block scales spread over a wide exponent range). Each speedup sits next to both latencies. Speedup is previous Cake latency divided by this PR's latency. Latency is the mean, over the paired samples, of each sample's median CUPTI kernel time per call, in microseconds. A speedup above 1 means this PR is faster. On B200, eight rows change program: 128 x 128 bwd colwise, 4096 x 7168 and 16384 x 7168 fwd both, and 131072 x 8192 fwd rowwise, fwd colwise, fwd both, bwd rowwise and bwd colwise. On GB300, no row changes program, because the sm_103a programs are byte-identical to the previous version. On R200, four rows change program: 128 x 128 fwd rowwise, fwd colwise, bwd colwise and bwd both. All other rows run the same program on both sides, so their ratios are run-to-run variation. ### NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.4 | M x K | Pass | Unit normals: previous Cake (µs) | This PR (µs) | Speedup | Block-scaled: previous Cake (µs) | This PR (µs) | Speedup | |---|---|---:|---:|---:|---:|---:|---:| | 128 x 128 | fwd rowwise | 2.16 | 2.16 | 1.000x | 2.39 | 2.39 | 1.000x | | 128 x 128 | fwd colwise | 2.21 | 2.21 | 1.000x | 2.43 | 2.43 | 1.001x | | 128 x 128 | fwd both | 2.35 | 2.35 | 1.001x | 2.58 | 2.58 | 1.001x | | 128 x 128 | bwd rowwise | 2.29 | 2.29 | 1.000x | 2.48 | 2.48 | 1.000x | | 128 x 128 | bwd colwise | 2.43 | 2.41 | 1.008x | 2.69 | 2.54 | 1.057x | | 128 x 128 | bwd both | 2.58 | 2.58 | 0.999x | 2.74 | 2.73 | 1.000x | | 4096 x 2048 | fwd rowwise | 9.75 | 9.75 | 1.000x | 11.11 | 11.11 | 1.000x | | 4096 x 2048 | fwd colwise | 10.41 | 10.42 | 1.000x | 12.81 | 12.81 | 1.000x | | 4096 x 2048 | fwd both | 11.60 | 11.60 | 1.000x | 13.64 | 13.64 | 1.000x | | 4096 x 2048 | bwd rowwise | 14.09 | 14.09 | 1.000x | 14.55 | 14.55 | 1.000x | | 4096 x 2048 | bwd colwise | 15.02 | 15.02 | 1.000x | 18.52 | 18.52 | 1.000x | | 4096 x 2048 | bwd both | 18.07 | 18.07 | 1.000x | 20.79 | 20.79 | 1.000x | | 4096 x 7168 | fwd rowwise | 25.41 | 25.41 | 1.000x | 28.77 | 28.77 | 1.000x | | 4096 x 7168 | fwd colwise | 26.57 | 26.57 | 1.000x | 32.15 | 32.15 | 1.000x | | 4096 x 7168 | fwd both | 31.61 | 31.30 | 1.010x | 39.21 | 35.70 | 1.098x | | 4096 x 7168 | bwd rowwise | 39.42 | 39.43 | 1.000x | 40.57 | 40.57 | 1.000x | | 4096 x 7168 | bwd colwise | 41.45 | 41.46 | 1.000x | 46.10 | 46.09 | 1.000x | | 4096 x 7168 | bwd both | 50.55 | 50.57 | 1.000x | 51.95 | 51.89 | 1.001x | | 16384 x 7168 | fwd rowwise | 90.21 | 90.21 | 1.000x | 102.27 | 102.26 | 1.000x | | 16384 x 7168 | fwd colwise | 92.23 | 92.24 | 1.000x | 101.68 | 101.68 | 1.000x | | 16384 x 7168 | fwd both | 109.92 | 108.90 | 1.009x | 134.95 | 122.95 | 1.098x | | 16384 x 7168 | bwd rowwise | 140.52 | 140.56 | 1.000x | 144.24 | 144.24 | 1.000x | | 16384 x 7168 | bwd colwise | 144.77 | 144.76 | 1.000x | 159.77 | 159.73 | 1.000x | | 16384 x 7168 | bwd both | 186.34 | 186.33 | 1.000x | 185.76 | 185.64 | 1.001x | | 131072 x 8192 | fwd rowwise | 760.65 | 760.72 | 1.000x | 918.82 | 809.53 | 1.135x | | 131072 x 8192 | fwd colwise | 792.55 | 792.26 | 1.000x | 995.49 | 921.81 | 1.080x | | 131072 x 8192 | fwd both | 1008.78 | 976.90 | 1.033x | 1218.22 | 1138.17 | 1.070x | | 131072 x 8192 | bwd rowwise | 1220.48 | 1218.62 | 1.002x | 1476.10 | 1392.03 | 1.060x | | 131072 x 8192 | bwd colwise | 1377.48 | 1364.09 | 1.010x | 1620.12 | 1521.75 | 1.065x | | 131072 x 8192 | bwd both | 1717.42 | 1712.86 | 1.003x | 1956.95 | 1960.92 | 0.998x | ### NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4 | M x K | Pass | Unit normals: previous Cake (µs) | This PR (µs) | Speedup | Block-scaled: previous Cake (µs) | This PR (µs) | Speedup | |---|---|---:|---:|---:|---:|---:|---:| | 128 x 128 | fwd rowwise | 2.06 | 2.05 | 1.001x | 2.33 | 2.34 | 0.999x | | 128 x 128 | fwd colwise | 2.13 | 2.12 | 1.001x | 2.41 | 2.41 | 1.001x | | 128 x 128 | fwd both | 2.23 | 2.23 | 0.998x | 2.51 | 2.51 | 0.999x | | 128 x 128 | bwd rowwise | 2.19 | 2.19 | 1.001x | 2.29 | 2.29 | 0.999x | | 128 x 128 | bwd colwise | 2.34 | 2.33 | 1.001x | 2.62 | 2.62 | 1.000x | | 128 x 128 | bwd both | 2.47 | 2.47 | 0.999x | 2.62 | 2.62 | 1.001x | | 4096 x 2048 | fwd rowwise | 9.14 | 9.14 | 1.000x | 10.18 | 10.18 | 1.000x | | 4096 x 2048 | fwd colwise | 9.73 | 9.73 | 1.001x | 12.35 | 12.35 | 1.000x | | 4096 x 2048 | fwd both | 11.00 | 11.00 | 1.000x | 13.30 | 13.30 | 1.000x | | 4096 x 2048 | bwd rowwise | 13.55 | 13.55 | 1.000x | 14.33 | 14.34 | 1.000x | | 4096 x 2048 | bwd colwise | 14.21 | 14.21 | 1.000x | 17.14 | 17.15 | 1.000x | | 4096 x 2048 | bwd both | 17.01 | 17.01 | 1.000x | 19.25 | 19.25 | 1.000x | | 4096 x 7168 | fwd rowwise | 24.10 | 24.10 | 1.000x | 26.36 | 26.35 | 1.000x | | 4096 x 7168 | fwd colwise | 25.48 | 25.48 | 1.000x | 31.67 | 31.68 | 1.000x | | 4096 x 7168 | fwd both | 29.98 | 29.97 | 1.000x | 34.63 | 34.63 | 1.000x | | 4096 x 7168 | bwd rowwise | 37.93 | 37.93 | 1.000x | 39.65 | 39.65 | 1.000x | | 4096 x 7168 | bwd colwise | 39.68 | 39.68 | 1.000x | 43.25 | 43.24 | 1.000x | | 4096 x 7168 | bwd both | 47.86 | 47.87 | 1.000x | 50.04 | 50.04 | 1.000x | | 16384 x 7168 | fwd rowwise | 86.95 | 86.95 | 1.000x | 95.09 | 95.09 | 1.000x | | 16384 x 7168 | fwd colwise | 89.96 | 89.96 | 1.000x | 98.66 | 98.67 | 1.000x | | 16384 x 7168 | fwd both | 104.98 | 104.98 | 1.000x | 119.22 | 119.19 | 1.000x | | 16384 x 7168 | bwd rowwise | 135.42 | 135.40 | 1.000x | 139.12 | 139.11 | 1.000x | | 16384 x 7168 | bwd colwise | 139.22 | 139.24 | 1.000x | 148.17 | 148.16 | 1.000x | | 16384 x 7168 | bwd both | 175.45 | 175.43 | 1.000x | 176.82 | 176.80 | 1.000x | | 131072 x 8192 | fwd rowwise | 738.58 | 738.57 | 1.000x | 755.40 | 755.42 | 1.000x | | 131072 x 8192 | fwd colwise | 756.21 | 756.20 | 1.000x | 843.92 | 843.93 | 1.000x | | 131072 x 8192 | fwd both | 906.46 | 906.42 | 1.000x | 1003.39 | 1003.43 | 1.000x | | 131072 x 8192 | bwd rowwise | 1171.09 | 1171.09 | 1.000x | 1203.39 | 1203.37 | 1.000x | | 131072 x 8192 | bwd colwise | 1220.25 | 1220.25 | 1.000x | 1292.73 | 1292.67 | 1.000x | | 131072 x 8192 | bwd both | 1561.61 | 1561.52 | 1.000x | 1559.61 | 1559.62 | 1.000x | ### NVIDIA R200 (sm_107a, 208 SMs), CUDA 13.4 | M x K | Pass | Unit normals: previous Cake (µs) | This PR (µs) | Speedup | Block-scaled: previous Cake (µs) | This PR (µs) | Speedup | |---|---|---:|---:|---:|---:|---:|---:| | 128 x 128 | fwd rowwise | 1.88 | 1.88 | 0.999x | 2.14 | 2.00 | 1.068x | | 128 x 128 | fwd colwise | 2.05 | 2.04 | 1.005x | 2.30 | 2.16 | 1.067x | | 128 x 128 | fwd both | 2.12 | 2.12 | 1.000x | 2.39 | 2.39 | 1.000x | | 128 x 128 | bwd rowwise | 2.12 | 2.12 | 1.000x | 2.33 | 2.32 | 1.000x | | 128 x 128 | bwd colwise | 2.23 | 2.23 | 1.000x | 2.42 | 2.31 | 1.051x | | 128 x 128 | bwd both | 2.37 | 2.35 | 1.007x | 2.56 | 2.48 | 1.032x | | 4096 x 2048 | fwd rowwise | 7.32 | 7.31 | 1.001x | 8.59 | 8.57 | 1.003x | | 4096 x 2048 | fwd colwise | 7.52 | 7.52 | 1.000x | 9.11 | 9.11 | 1.000x | | 4096 x 2048 | fwd both | 8.82 | 8.81 | 1.000x | 10.58 | 10.58 | 1.000x | | 4096 x 2048 | bwd rowwise | 10.49 | 10.49 | 1.000x | 11.05 | 11.05 | 1.000x | | 4096 x 2048 | bwd colwise | 11.04 | 11.04 | 1.000x | 12.51 | 12.51 | 1.000x | | 4096 x 2048 | bwd both | 13.02 | 13.01 | 1.000x | 13.73 | 13.77 | 0.997x | | 4096 x 7168 | fwd rowwise | 18.51 | 18.52 | 1.000x | 20.85 | 20.83 | 1.001x | | 4096 x 7168 | fwd colwise | 18.89 | 18.89 | 1.000x | 21.91 | 21.92 | 1.000x | | 4096 x 7168 | fwd both | 21.83 | 21.83 | 1.000x | 26.08 | 26.07 | 1.000x | | 4096 x 7168 | bwd rowwise | 28.64 | 28.65 | 1.000x | 29.58 | 29.59 | 1.000x | | 4096 x 7168 | bwd colwise | 28.97 | 28.96 | 1.000x | 30.53 | 30.53 | 1.000x | | 4096 x 7168 | bwd both | 35.62 | 35.62 | 1.000x | 36.63 | 36.69 | 0.998x | | 16384 x 7168 | fwd rowwise | 59.46 | 59.47 | 1.000x | 67.11 | 67.16 | 0.999x | | 16384 x 7168 | fwd colwise | 60.32 | 60.30 | 1.000x | 70.06 | 70.04 | 1.000x | | 16384 x 7168 | fwd both | 73.78 | 73.78 | 1.000x | 79.05 | 79.04 | 1.000x | | 16384 x 7168 | bwd rowwise | 97.27 | 97.26 | 1.000x | 101.30 | 101.31 | 1.000x | | 16384 x 7168 | bwd colwise | 98.46 | 98.47 | 1.000x | 107.97 | 107.92 | 1.000x | | 16384 x 7168 | bwd both | 129.25 | 129.26 | 1.000x | 131.18 | 131.28 | 0.999x | | 131072 x 8192 | fwd rowwise | 484.89 | 484.89 | 1.000x | 542.46 | 542.52 | 1.000x | | 131072 x 8192 | fwd colwise | 508.54 | 508.54 | 1.000x | 569.01 | 569.18 | 1.000x | | 131072 x 8192 | fwd both | 646.33 | 646.36 | 1.000x | 654.63 | 654.81 | 1.000x | | 131072 x 8192 | bwd rowwise | 832.90 | 832.90 | 1.000x | 868.23 | 868.25 | 1.000x | | 131072 x 8192 | bwd colwise | 851.75 | 851.76 | 1.000x | 914.99 | 915.04 | 1.000x | | 131072 x 8192 | bwd both | 1150.00 | 1149.98 | 1.000x | 1162.89 | 1162.96 | 1.000x | Cake has no other variants of this kernel beyond these six BF16-to-MXFP8 forms with the swizzled scale layout, and the family ships on sm_100a (B200), sm_103a (GB300) and sm_107a (R200). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Added MXFP8 gated-activation forward and backward kernel variants for additional GPU architectures. * Expanded support for newer GPU hardware, including SM100, SM103, and SM107. * **Improvements** * Updated kernel selection across supported architectures and workload sizes. * Improved handling of numerical edge cases in gated activations and MXFP8 conversion. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude <noreply@anthropic.com>
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.
The Cake fused SwiGLU + MXFP8 quantization kernel computes silu(gate) * up from a BF16 input [M, 2K] and quantizes the result to MXFP8. Each group of 32 E4M3 values gets one E8M0 scale (RCEIL rounding), and the scales are stored in the swizzled 128 x 4 layout. The kernel has one data type and six forms. The forward form is
flashinfer.silu_and_mul_mxfp8_quantize, with output [M, K]. The backward form isflashinfer.silu_and_mul_mxfp8_quantize_backward. It computes the gate and up gradients [M, 2K] from the input [M, 2K] and grad_output [M, K], and quantizes them to MXFP8 in the same way. Each one can write rowwise, colwise or both outputs.flashinfer.cake_gated_act_mxfp8_backendruns one program per call and picks it by M * K for each architecture.This PR changes which program is selected in 8 per-architecture M * K ranges: 5 on sm_100a, 3 on sm_103a and none on sm_107a. In 7 of them the new programs reuse the guarded-sigmoid band-path forms from #7492: a loop over only the flagged pairs, a non-inlined call for each flagged pair, or an inline branch for each pair. Each new program has a register budget fitted to a CTA count per SM. Below M * K = 2^20, which covers the smallest backward shapes, backward rowwise and backward both on sm_100a and sm_103a take a 32 x 32 tile program at 12 CTAs per SM. Rowwise uses the call form on sm_100a and the inline form on sm_103a. Both outputs use the flagged-pair loop on both architectures. On sm_100a, forward colwise, backward colwise and backward both get a new split at M * K = 2^27. Below the split they take a band-path program, starting at 2^26 for forward colwise and at 2^24 for the two backward passes. From 2^27 up they keep the previous program. On sm_103a, backward rowwise from M * K = 2^26 up (the program with
no_allocateloads) used a build at 12 CTAs per SM, which spills 8 bytes to the stack. It now uses the packed f32x2 sigmoid form from #7396 at 10 CTAs per SM, which fits without a spill. The package goes from 45 to 50 programs. 6 programs are new and 1 is removed: the 12-CTA sm_103a build. The same arithmetic runs on the same operands, so every output byte is identical. The public API is unchanged. The band-path programs help block-scaled inputs, where many warps take the guarded path. With unit normals these rows run at about the same speed, and the small 32 x 32 rows run slightly faster. The sm_103a backward rowwise change makes the two large backward rowwise rows on GB300 run slightly faster with both input sets.The baseline is the previous Cake version of this kernel, which is the package FlashInfer main ships (#7492). Both versions run in the same process on the same GPU, in alternating paired runs, with a cold L2 before every launch. Each table row is one shape and pass. There are two column groups: unit-normal inputs, and block-scaled inputs (normals with per-32-block scales spread over a wide exponent range). Each speedup sits next to both latencies. Speedup is previous Cake latency divided by this PR's latency. Latency is the mean, over the paired samples, of each sample's median CUPTI kernel time per call, in microseconds. A speedup above 1 means this PR is faster. On B200, seven rows run a new program: 128 x 128 bwd rowwise and bwd both, 4096 x 7168 bwd colwise and bwd both, and 16384 x 7168 fwd colwise, bwd colwise and bwd both. On GB300, four rows run a new program: 128 x 128 bwd rowwise and bwd both, 16384 x 7168 bwd rowwise and 131072 x 8192 bwd rowwise. On R200, no row runs a new program. All other rows run the same program on both sides, so their ratios are run-to-run variation.
NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.4
NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4
NVIDIA R200 (sm_107a, 208 SMs), CUDA 13.4
Cake has no other variants of this kernel beyond these six BF16-to-MXFP8 forms with the swizzled scale layout, and the family ships on sm_100a (B200), sm_103a (GB300) and sm_107a (R200).
Summary by CodeRabbit