[yyihuang]

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 changes program selection on sm_107a (R200) only. On R200, four rows are 1.020x to 1.047x faster with block-scaled inputs than the previous Cake version. Three M * K dispatch ranges move to the band-path forms from #7492, and two of them also use the packed f32x2 sigmoid from #7396. Forward rowwise in [2^24, 2^26) now runs the two-row form with the called band path at 4 CTAs per SM. Backward rowwise in [2^18, 2^26) runs the packed-sigmoid form with the called band path at 8 CTAs per SM. Backward colwise in [2^23, 2^24) runs the packed-sigmoid form with the ffs band path at 8 CTAs per SM. The package goes from 57 to 60 programs. The backward rowwise packed-sigmoid program no longer ships for sm_107a but still serves sm_103a. Outputs are byte-identical to the previous version, and the public API is unchanged.

The baseline is the previous Cake version of this kernel, which is the package from #7770. The method is the same as in #7770, with 8 paired rounds. 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.

Four R200 rows change program. With block-scaled inputs, all four are faster than the run-to-run (A/A) spread: 1.047x (4096 x 2048 backward colwise), 1.042x (4096 x 7168 forward rowwise), 1.031x (4096 x 2048 backward rowwise) and 1.020x (4096 x 7168 backward rowwise). The other 26 block-scaled rows stay within the spread, and block-scaled ratios range from 0.999x to 1.047x. With unit-normal inputs, all 30 rows are ties within the spread, from 0.999x to 1.001x. No row is slower than the spread. The table covers R200 (sm_107a, 208 SMs) on CUDA 13.4, with five shapes and six passes. The measured package has the same machine code as the package in this PR. There are no B200 or GB300 tables because the sm_100a and sm_103a programs have the same machine code as in the previous version.

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	2.11	2.11	1.000x	2.25	2.25	1.002x
128 x 128	fwd colwise	2.26	2.26	1.000x	2.41	2.41	0.999x
128 x 128	fwd both	2.34	2.33	1.001x	2.64	2.64	0.999x
128 x 128	bwd rowwise	2.04	2.04	1.000x	2.26	2.26	1.000x
128 x 128	bwd colwise	2.15	2.15	1.001x	2.24	2.23	1.002x
128 x 128	bwd both	2.27	2.27	0.999x	2.40	2.40	1.001x
4096 x 2048	fwd rowwise	7.15	7.15	1.000x	8.20	8.17	1.004x
4096 x 2048	fwd colwise	7.30	7.29	1.000x	8.75	8.74	1.000x
4096 x 2048	fwd both	8.71	8.71	1.000x	10.32	10.30	1.002x
4096 x 2048	bwd rowwise	10.53	10.52	1.001x	11.01	10.68	1.031x
4096 x 2048	bwd colwise	11.05	11.04	1.001x	12.28	11.73	1.047x
4096 x 2048	bwd both	13.24	13.24	1.000x	13.88	13.85	1.002x
4096 x 7168	fwd rowwise	18.39	18.37	1.001x	20.37	19.55	1.042x
4096 x 7168	fwd colwise	18.83	18.82	1.000x	21.46	21.47	1.000x
4096 x 7168	fwd both	21.91	21.91	1.000x	25.80	25.79	1.000x
4096 x 7168	bwd rowwise	28.89	28.88	1.000x	29.56	28.98	1.020x
4096 x 7168	bwd colwise	28.82	28.82	1.000x	30.23	30.22	1.000x
4096 x 7168	bwd both	35.91	35.91	1.000x	36.90	36.84	1.002x
16384 x 7168	fwd rowwise	59.06	59.07	1.000x	65.89	65.91	1.000x
16384 x 7168	fwd colwise	60.41	60.41	1.000x	69.03	69.01	1.000x
16384 x 7168	fwd both	75.10	75.09	1.000x	79.41	79.41	1.000x
16384 x 7168	bwd rowwise	97.26	97.26	1.000x	101.09	101.11	1.000x
16384 x 7168	bwd colwise	99.36	99.36	1.000x	107.40	107.43	1.000x
16384 x 7168	bwd both	132.00	132.02	1.000x	134.08	133.99	1.001x
131072 x 8192	fwd rowwise	481.61	481.61	1.000x	532.50	532.46	1.000x
131072 x 8192	fwd colwise	514.00	513.99	1.000x	561.53	561.77	1.000x
131072 x 8192	fwd both	657.84	657.91	1.000x	665.77	665.79	1.000x
131072 x 8192	bwd rowwise	831.94	831.92	1.000x	864.37	864.39	1.000x
131072 x 8192	bwd colwise	864.92	864.93	1.000x	917.28	917.31	1.000x
131072 x 8192	bwd both	1183.11	1183.16	1.000x	1196.51	1196.42	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).

Summary by CodeRabbit

• New Features
  • Added support for gated-activation MXFP8 forward and backward operations on SM107A GPUs, including rowwise and columnwise output paths.
  • Kernel selection now varies by operation and workload size. Some smaller backward rowwise workloads on SM107A use the existing fallback route.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 93b5b442-f5c1-4a2e-9591-256106d4758b

Commits

Reviewing files that changed from the base of the PR and between 0d14819 and 42d7e27.

Files selected for processing: 8

• flashinfer/cake_gated_act_mxfp8_backend/build.json
• flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py
• flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_75c68471d8a4921f432a_binding.cu
• flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_75c68471d8a4921f432a_kernel.cu
• flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_822bead4b44c749f7afe_binding.cu
• flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_822bead4b44c749f7afe_kernel.cu
• flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_f8db01d14808a5afc6da_binding.cu
• flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_f8db01d14808a5afc6da_kernel.cu

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 4 remain after this review.

Walkthrough

The backend adds three SM107A gated-activation MXFP8 kernel variants and their TVM-FFI bindings. JIT routing selects variants by operation and M * K threshold. The build manifest adds library records and removes one existing SM107A record.

Changes: SM107A MXFP8 backend

Layer / File(s)	Summary
Forward-row kernel and binding flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_75c68471d8a4921f432a_kernel.cu, flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_75c68471d8a4921f432a_binding.cu	Adds an SM107A forward-row kernel with fast and fallback activation paths, row-wise E4M3 quantization, and per-region scales. Its binding validates tensors and launch dimensions, then launches the kernel.
Backward kernels and bindings flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_822bead4b44c749f7afe_kernel.cu, flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_822bead4b44c749f7afe_binding.cu, flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_f8db01d14808a5afc6da_kernel.cu, flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_f8db01d14808a5afc6da_binding.cu	Adds two backward kernels that compute gated-activation gradients, use fallback calculations for selected inputs, and write E4M3 outputs with per-group scales. Their bindings validate tensors and launch parameters before launching the kernels.
JIT routing and build records flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py, flashinfer/cake_gated_act_mxfp8_backend/build.json	Adds SM107A JIT records and routes for forward-row, backward-row, and backward-col variants. The routes use thresholds of 2^24, 2^18, and 2^23, respectively. The build manifest adds three SM107A library records and removes one existing record.

Priority: Low

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant cake_jit
  participant TVM_FFI_run
  participant kernel_cake_gated_act_mxfp8_75c68471d8a4921f432a
  cake_jit->>TVM_FFI_run: Select SM107A forward-row variant when M*K >= 2^24
  TVM_FFI_run->>kernel_cake_gated_act_mxfp8_75c68471d8a4921f432a: Validate tensors and launch kernel

Loading

Suggested reviewers: sricketts, aleozlx, yongwww

Merge Risk: Minimal · up to 42d7e

This change routes three SM107A (R200) MXFP8 gated-activation ranges to new kernel variants that the author reports as faster with identical outputs. No concrete merge-blocking risk was found in the supplied context.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 5.88% 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. Documentat…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the performance change, the Cake MXFP8 component, the band-path programs, and the three SM107 entries.
Description check		The description is detailed and directly explains the dispatch changes, affected architectures, benchmarks, output compatibility, and package changes. It does not include the template headings for rel…
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 is detailed and directly explains the dispatch changes, affected architectures, benchmarks, output compatibility, and package changes. It does not include the template headings for related issues, checklist results, or reviewer notes, but the core change description is complete.

Full details: Docstring Coverage

Explanation

Docstring coverage is 5.88% 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 51 functions across 7 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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[yyihuang]

@flashinfer-bot run tests/utils/test_cake_gated_act_mxfp8.py tests/utils/test_cake_gated_act_mxfp8_api.py