perf(cake_gated_act_mxfp8): narrow tiles for small problems and four rows per thread for large ones by yyihuang · Pull Request #7346 · flashinfer-ai/flashinfer

yyihuang · 2026-10-25T16:49:07Z

This PR updates the Cake programs behind FlashInfer's fused SwiGLU + MXFP8 quantization. The kernel computes silu(gate) * up from a BF16 input [M, 2K] and quantizes the result to MXFP8. The output uses E4M3 values with one E8M0 scale (RCEIL rounding) per 32 values, and the scales are stored in the swizzled 128 x 4 layout. There is one data type (BF16 in, MXFP8 E4M3 + E8M0 out) 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, which produces gate and up gradients [M, 2K] from grad_output [M, K]. Each one has rowwise, colwise or both outputs. The package flashinfer.cake_gated_act_mxfp8_backend runs one program per call, chosen by M * K per architecture. This PR updates the programs of all six forms on sm_100a, sm_103a and sm_107a:

Below the first size threshold, a 32 x 32 CTA tile runs the forward on sm_100a and sm_103a and the backward on all three architectures.
On sm_103a, the colwise forward runs a 32 x 64 tile from 2^22 to 2^24 elements.
A four-rows-per-thread program runs at the large sm_107a shapes and at the largest sm_103a forward shapes with both outputs.

The public API is unchanged, and every output byte is identical. The baseline is the previous Cake version of this kernel, which is the Cake package FlashInfer main ships (#7305).

Speedup is previous Cake latency / this PR's latency. Latency is the mean CUPTI kernel time in microseconds, with a cold L2 before every launch, and both versions run in the same process on the same GPU in alternating paired runs. On B200, only the six 128 x 128 rows run a new program, at 1.144-1.260x. The other 24 rows run the same program on both sides, and their ratios (0.978-1.014x) are run-to-run variation. On GB300, nine rows run a new program at 1.005-1.121x: the six 128 x 128 rows, 4096 x 2048 fwd colwise, and 16384 x 7168 and 131072 x 8192 fwd both. The other 21 rows run the same program on both sides, at 0.998-1.001x. On R200, twelve rows run a new program at 1.017-1.116x: the three 128 x 128 backward rows, 4096 x 7168 fwd both, 16384 x 7168 fwd rowwise / colwise / both and bwd colwise, and 131072 x 8192 fwd rowwise / colwise / both and bwd colwise. The other 18 rows run the same program on both sides, at 0.999-1.001x.

NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.4.92

M x K	Pass	Previous Cake (µs)	This PR (µs)	Speedup
128 x 128	fwd rowwise	2.53	2.06	1.230x
128 x 128	fwd colwise	2.67	2.12	1.260x
128 x 128	fwd both	2.78	2.25	1.234x
128 x 128	bwd rowwise	2.56	2.24	1.144x
128 x 128	bwd colwise	2.85	2.36	1.206x
128 x 128	bwd both	2.92	2.53	1.152x
4096 x 2048	fwd rowwise	9.60	9.60	1.001x
4096 x 2048	fwd colwise	10.37	10.37	1.000x
4096 x 2048	fwd both	11.53	11.53	1.000x
4096 x 2048	bwd rowwise	13.83	13.83	1.000x
4096 x 2048	bwd colwise	14.91	14.90	1.000x
4096 x 2048	bwd both	18.04	18.04	1.000x
4096 x 7168	fwd rowwise	25.30	25.29	1.000x
4096 x 7168	fwd colwise	26.47	26.47	1.000x
4096 x 7168	fwd both	31.31	31.31	1.000x
4096 x 7168	bwd rowwise	39.34	39.34	1.000x
4096 x 7168	bwd colwise	41.06	41.07	1.000x
4096 x 7168	bwd both	50.41	50.40	1.000x
16384 x 7168	fwd rowwise	89.91	89.92	1.000x
16384 x 7168	fwd colwise	91.71	91.72	1.000x
16384 x 7168	fwd both	108.54	108.53	1.000x
16384 x 7168	bwd rowwise	141.22	141.29	1.000x
16384 x 7168	bwd colwise	145.29	145.30	1.000x
16384 x 7168	bwd both	187.17	187.21	1.000x
131072 x 8192	fwd rowwise	765.89	764.82	1.001x
131072 x 8192	fwd colwise	821.69	815.25	1.008x
131072 x 8192	fwd both	1003.00	1025.67	0.978x
131072 x 8192	bwd rowwise	1247.49	1241.54	1.005x
131072 x 8192	bwd colwise	1391.55	1371.77	1.014x
131072 x 8192	bwd both	1728.92	1746.36	0.990x

NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4.92

M x K	Pass	Previous Cake (µs)	This PR (µs)	Speedup
128 x 128	fwd rowwise	2.13	1.94	1.096x
128 x 128	fwd colwise	2.26	2.02	1.121x
128 x 128	fwd both	2.35	2.12	1.111x
128 x 128	bwd rowwise	2.16	2.14	1.010x
128 x 128	bwd colwise	2.42	2.23	1.087x
128 x 128	bwd both	2.48	2.41	1.031x
4096 x 2048	fwd rowwise	9.02	9.03	1.000x
4096 x 2048	fwd colwise	9.88	9.83	1.005x
4096 x 2048	fwd both	10.97	10.97	1.000x
4096 x 2048	bwd rowwise	13.27	13.26	1.001x
4096 x 2048	bwd colwise	13.99	13.99	1.000x
4096 x 2048	bwd both	16.81	16.81	1.000x
4096 x 7168	fwd rowwise	24.42	24.44	0.999x
4096 x 7168	fwd colwise	25.52	25.52	1.000x
4096 x 7168	fwd both	30.30	30.30	1.000x
4096 x 7168	bwd rowwise	37.75	37.81	0.998x
4096 x 7168	bwd colwise	39.95	39.92	1.001x
4096 x 7168	bwd both	47.77	47.77	1.000x
16384 x 7168	fwd rowwise	87.35	87.28	1.001x
16384 x 7168	fwd colwise	90.02	90.02	1.000x
16384 x 7168	fwd both	106.60	105.34	1.012x
16384 x 7168	bwd rowwise	135.79	135.81	1.000x
16384 x 7168	bwd colwise	139.71	139.69	1.000x
16384 x 7168	bwd both	175.33	175.25	1.000x
131072 x 8192	fwd rowwise	742.29	741.71	1.001x
131072 x 8192	fwd colwise	756.04	756.05	1.000x
131072 x 8192	fwd both	923.27	904.89	1.020x
131072 x 8192	bwd rowwise	1174.31	1174.17	1.000x
131072 x 8192	bwd colwise	1222.94	1223.02	1.000x
131072 x 8192	bwd both	1552.67	1552.57	1.000x

NVIDIA R200 (sm_107a, 208 SMs), CUDA 13.4.92

M x K	Pass	Previous Cake (µs)	This PR (µs)	Speedup
128 x 128	fwd rowwise	2.11	2.11	1.000x
128 x 128	fwd colwise	2.28	2.28	0.999x
128 x 128	fwd both	2.34	2.34	1.001x
128 x 128	bwd rowwise	2.24	2.05	1.090x
128 x 128	bwd colwise	2.42	2.17	1.116x
128 x 128	bwd both	2.56	2.30	1.112x
4096 x 2048	fwd rowwise	7.22	7.22	1.000x
4096 x 2048	fwd colwise	7.34	7.34	1.000x
4096 x 2048	fwd both	8.76	8.76	1.001x
4096 x 2048	bwd rowwise	10.50	10.49	1.000x
4096 x 2048	bwd colwise	11.07	11.06	1.000x
4096 x 2048	bwd both	13.34	13.33	1.000x
4096 x 7168	fwd rowwise	18.48	18.48	1.000x
4096 x 7168	fwd colwise	18.99	19.00	0.999x
4096 x 7168	fwd both	22.84	21.96	1.040x
4096 x 7168	bwd rowwise	28.84	28.82	1.001x
4096 x 7168	bwd colwise	29.04	29.05	1.000x
4096 x 7168	bwd both	36.09	36.11	1.000x
16384 x 7168	fwd rowwise	63.67	59.12	1.077x
16384 x 7168	fwd colwise	63.19	61.00	1.036x
16384 x 7168	fwd both	78.43	75.13	1.044x
16384 x 7168	bwd rowwise	97.26	97.26	1.000x
16384 x 7168	bwd colwise	100.35	98.63	1.017x
16384 x 7168	bwd both	131.91	131.93	1.000x
131072 x 8192	fwd rowwise	528.30	481.99	1.096x
131072 x 8192	fwd colwise	533.65	515.98	1.034x
131072 x 8192	fwd both	682.08	658.88	1.035x
131072 x 8192	bwd rowwise	831.87	831.91	1.000x
131072 x 8192	bwd colwise	882.75	864.51	1.021x
131072 x 8192	bwd both	1183.12	1183.07	1.000x

Cake ships this kernel only in this BF16-to-MXFP8 form (the six forms above), and the package ships on sm_100a, sm_103a and sm_107a.

Summary by CodeRabbit

New Features
- Added broader GPU kernel coverage for gated-activation MXFP8 operations, including forward and gradient calculations across additional tile sizes and architectures.
- Expanded support for large input sizes and added more flexible kernel selection based on input dimensions.
Bug Fixes
- Improved fallback selection when an architecture has no specific rule, and added checks for ambiguous kernel-selection rules.

…rows per thread for large ones Below its first size threshold the package now runs a 32 x 32 CTA tile: the forward on SM100 and SM103, the backward on SM100, SM103 and SM107. The SM103 colwise forward runs a 32 x 64 tile from 2^22 to 2^24 elements. A four-rows-per-thread build runs at the large SM107 shapes and at the largest SM103 forward shapes with both outputs. A record without a threshold on an architecture is that architecture's fallback for its route. Every output byte is unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

coderabbitai · 2026-10-25T16:49:31Z

Important

Review skipped

Review was skipped as selected files did not have any reviewable changes.

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: 615c4e38-9461-4965-a69b-ab785dfe3c0f

Commits

Reviewing files that changed from the base of the PR and between c740bad and 4a4b1a6.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

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: ee6ba9d4-1d8a-40ad-8835-80f4d45069f5

Commits

Reviewing files that changed from the base of the PR and between b8ed186 and c740bad.

Files selected for processing: 26

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_2513ecc12f03a6081c2e_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_2513ecc12f03a6081c2e_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_34e92b6e8ad60c86843d_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_34e92b6e8ad60c86843d_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_49303287790d1f318420_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_49303287790d1f318420_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_4c31b87b790c95ed0b77_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_70d1cafb2bb8a175a516_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_70d1cafb2bb8a175a516_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_7e191474abd37ba82e73_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_7e191474abd37ba82e73_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_8c43730c1e3d76ce89f1_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_8c43730c1e3d76ce89f1_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_8da8bc14e105ce409f74_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_8da8bc14e105ce409f74_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_bbe04581406928a16cbc_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_bbe04581406928a16cbc_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_bf857c6b036c4025d3fd_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_bf857c6b036c4025d3fd_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_d482e0e2dc5d60ddeb62_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_d482e0e2dc5d60ddeb62_kernel.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_f1232b8d8f9b3a7f66cf_binding.cu
flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_f1232b8d8f9b3a7f66cf_kernel.cu
tests/utils/test_cake_gated_act_mxfp8.py

Files with no reviewable changes: 1

flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_4c31b87b790c95ed0b77_kernel.cu

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

Walkthrough

This change adds generated gated-activation MXFP8 CUDA kernel variants and bindings. It updates architecture-specific tile and threshold selection, fallback handling, and build records. Tests add larger entry shapes and accept the updated tile configurations.

Changes: Gated-activation forward kernels

Layer / File(s)	Summary
Forward activation kernels `flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_{34e92...,70d1...,7e191...,f123...}_{binding,kernel}.cu`, `flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_{bbe045...,bf857...}_binding.cu`, `flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_d482..._binding.cu`, `flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_4c31..._kernel.cu`	Adds forward BF16 gated-activation kernels with E4M3 outputs and scales, plus TVM-FFI bindings. The `d482...` binding now targets its matching kernel name, and the prior `4c31...` kernel source is removed.

Gated-activation gradient kernels

Layer / File(s)	Summary
Gradient kernels `flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_{2513...,8c437...,8da8...}_{binding,kernel}.cu`, `flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_493..._binding.cu`	Adds gradient kernels that compute BF16 activation gradients and quantize outputs to E4M3. The kernels include fast and fallback computation paths, scale calculation, and output writes. Their bindings validate tensor and launch parameters before invoking the kernels.

Architecture-specific registry and selection

Layer / File(s)	Summary
Architecture-specific registry and selection `flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py`, `flashinfer/cake_gated_act_mxfp8_backend/build.json`, `tests/utils/test_cake_gated_act_mxfp8.py`	Adds and updates architecture-specific build records and route programs. Selection treats a missing architecture threshold as a fallback, selects the eligible program with the highest threshold, and checks for a unique fallback. Tests cover additional input shapes and the revised tile and block configurations.

Priority: Normal

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

Change: Refactor

Suggested labels: performance

Suggested reviewers: sricketts, aleozlx, yongwww

Merge Risk: Minimal · up to c740b

This PR adds architecture-specific tile and threshold selection for the fused SwiGLU + MXFP8 quantization kernels. It leaves the public API unchanged. No actionable merge-blocking risk was identified, and the reported performance gains rest on the author's own measurements.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 7.86% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 140 functions across 20 files. (1 skipped:…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title is a conventional-commit style summary that names the two main changes: narrower tiles for small problems and four rows per thread for large ones. Both appear in the raw summary and the PR o…
Description check		The description covers the purpose (updating the Cake programs for fused SwiGLU + MXFP8 quantization), the six forms, per-architecture changes, benchmark methodology, and tables of results. It also st…
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 is a conventional-commit style summary that names the two main changes: narrower tiles for small problems and four rows per thread for large ones. Both appear in the raw summary and the PR objectives (32x32 CTA tiles below the first threshold, 32x64 colwise forward tile on sm_103a, and a four-rows-per-thread program for large shapes). The title is specific and matches the changeset, so it passes. It is slightly long at 99 characters.

Full details: Description check

Explanation

The description covers the purpose (updating the Cake programs for fused SwiGLU + MXFP8 quantization), the six forms, per-architecture changes, benchmark methodology, and tables of results. It also states that the public API is unchanged and output bytes are identical. It does not follow the repository template. The 'Related Issues' section is missing (it does reference PR #7305 as baseline, but no issue link), and the pre-commit and test checklist items are not addressed. The Experimental Track section is not relevant here and can be left out. The description is substantive, so it is mostly complete, but the template sections for related issues, checklist status, and test status are absent.

Resolution

Consider adding the template's Related Issues section (or stating 'None'), and confirming the pre-commit and test checklist items. Stating which tests were run (for example, the updated tests/utils/test_cake_gated_act_mxfp8.py) would also help reviewers.

Full details: Docstring Coverage

Explanation

Docstring coverage is 7.86% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 140 functions across 20 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

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

yyihuang · 2026-10-25T18:55:39Z

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

yyihuang · 2026-10-25T19:56:30Z

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

… / SM103 / SM107) (#7396) 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`, and the backward form is `flashinfer.silu_and_mul_mxfp8_quantize_backward`. 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 switches 18 of these per-architecture programs to packed versions: 4 on sm_100a, 9 on sm_103a and 5 on sm_107a. Together they cover all six forms. In the packed versions, the sigmoid's FP32 multiply, subtract and both Newton FMA steps run on pairs of elements with f32x2 instructions. The exponentials and reciprocals still run one element at a time. Every output byte is identical, and the public API is unchanged. The baseline is the previous Cake version of this kernel, which is the package FlashInfer main ships (#7346). 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, and 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, five rows run a new program: 4096 x 2048 fwd colwise, fwd both, bwd rowwise and bwd colwise, and 4096 x 7168 bwd rowwise. On GB300, 18 rows run a new program: all six 4096 x 2048 rows, every 4096 x 7168 row except fwd colwise, 16384 x 7168 fwd rowwise, fwd both, bwd colwise and bwd both, and 131072 x 8192 fwd both, bwd colwise and bwd both. On R200, ten rows run a new program: 4096 x 2048 and 4096 x 7168 bwd rowwise, colwise and both, 16384 x 7168 fwd colwise and bwd both, and 131072 x 8192 fwd 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.92 | 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.03 | 2.03 | 1.000x | 2.38 | 2.38 | 1.001x | | 128 x 128 | fwd colwise | 2.09 | 2.09 | 1.000x | 2.41 | 2.41 | 1.001x | | 128 x 128 | fwd both | 2.23 | 2.23 | 1.000x | 2.56 | 2.56 | 1.000x | | 128 x 128 | bwd rowwise | 2.22 | 2.22 | 1.000x | 2.61 | 2.60 | 1.001x | | 128 x 128 | bwd colwise | 2.33 | 2.33 | 1.000x | 2.66 | 2.66 | 1.000x | | 128 x 128 | bwd both | 2.51 | 2.51 | 1.000x | 2.90 | 2.89 | 1.001x | | 4096 x 2048 | fwd rowwise | 9.58 | 9.58 | 1.000x | 10.99 | 10.99 | 1.000x | | 4096 x 2048 | fwd colwise | 10.47 | 10.30 | 1.016x | 13.12 | 12.68 | 1.035x | | 4096 x 2048 | fwd both | 11.57 | 11.50 | 1.006x | 14.51 | 14.35 | 1.011x | | 4096 x 2048 | bwd rowwise | 13.95 | 13.90 | 1.004x | 15.21 | 14.94 | 1.018x | | 4096 x 2048 | bwd colwise | 14.84 | 14.82 | 1.002x | 18.65 | 18.33 | 1.017x | | 4096 x 2048 | bwd both | 17.83 | 17.83 | 1.000x | 20.59 | 20.58 | 1.000x | | 4096 x 7168 | fwd rowwise | 25.22 | 25.23 | 1.000x | 28.58 | 28.58 | 1.000x | | 4096 x 7168 | fwd colwise | 26.42 | 26.43 | 1.000x | 34.41 | 34.42 | 1.000x | | 4096 x 7168 | fwd both | 31.36 | 31.36 | 1.000x | 38.88 | 38.88 | 1.000x | | 4096 x 7168 | bwd rowwise | 39.34 | 39.24 | 1.003x | 42.58 | 41.72 | 1.021x | | 4096 x 7168 | bwd colwise | 41.29 | 41.29 | 1.000x | 49.01 | 49.01 | 1.000x | | 4096 x 7168 | bwd both | 50.41 | 50.42 | 1.000x | 54.80 | 54.87 | 0.999x | | 16384 x 7168 | fwd rowwise | 90.07 | 90.05 | 1.000x | 101.92 | 101.93 | 1.000x | | 16384 x 7168 | fwd colwise | 91.93 | 91.94 | 1.000x | 112.75 | 112.74 | 1.000x | | 16384 x 7168 | fwd both | 109.38 | 109.36 | 1.000x | 134.21 | 134.18 | 1.000x | | 16384 x 7168 | bwd rowwise | 140.77 | 140.72 | 1.000x | 147.28 | 147.24 | 1.000x | | 16384 x 7168 | bwd colwise | 144.49 | 144.51 | 1.000x | 171.42 | 171.37 | 1.000x | | 16384 x 7168 | bwd both | 186.16 | 186.12 | 1.000x | 195.26 | 195.37 | 0.999x | | 131072 x 8192 | fwd rowwise | 761.03 | 757.92 | 1.004x | 907.19 | 905.57 | 1.002x | | 131072 x 8192 | fwd colwise | 794.40 | 806.78 | 0.985x | 994.34 | 991.28 | 1.003x | | 131072 x 8192 | fwd both | 1000.01 | 1001.98 | 0.998x | 1214.29 | 1211.58 | 1.002x | | 131072 x 8192 | bwd rowwise | 1227.09 | 1228.15 | 0.999x | 1486.79 | 1482.47 | 1.003x | | 131072 x 8192 | bwd colwise | 1372.76 | 1385.67 | 0.991x | 1619.29 | 1622.29 | 0.998x | | 131072 x 8192 | bwd both | 1728.21 | 1728.68 | 1.000x | 1978.19 | 1947.35 | 1.016x | ### NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4.92 | 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.94 | 1.94 | 1.003x | 2.24 | 2.23 | 1.001x | | 128 x 128 | fwd colwise | 2.01 | 2.01 | 1.002x | 2.31 | 2.32 | 0.999x | | 128 x 128 | fwd both | 2.11 | 2.11 | 1.000x | 2.42 | 2.42 | 1.001x | | 128 x 128 | bwd rowwise | 2.12 | 2.12 | 1.000x | 2.45 | 2.46 | 0.999x | | 128 x 128 | bwd colwise | 2.22 | 2.22 | 1.000x | 2.52 | 2.52 | 1.001x | | 128 x 128 | bwd both | 2.40 | 2.40 | 1.000x | 2.76 | 2.76 | 1.000x | | 4096 x 2048 | fwd rowwise | 9.02 | 9.01 | 1.001x | 10.40 | 10.12 | 1.028x | | 4096 x 2048 | fwd colwise | 9.82 | 9.61 | 1.022x | 12.43 | 12.32 | 1.009x | | 4096 x 2048 | fwd both | 10.98 | 10.89 | 1.008x | 13.60 | 13.28 | 1.024x | | 4096 x 2048 | bwd rowwise | 13.25 | 13.23 | 1.002x | 14.25 | 14.09 | 1.011x | | 4096 x 2048 | bwd colwise | 13.98 | 13.92 | 1.004x | 17.39 | 17.05 | 1.020x | | 4096 x 2048 | bwd both | 16.80 | 16.71 | 1.006x | 19.35 | 19.16 | 1.010x | | 4096 x 7168 | fwd rowwise | 24.43 | 24.26 | 1.007x | 26.96 | 26.19 | 1.029x | | 4096 x 7168 | fwd colwise | 25.52 | 25.52 | 1.000x | 31.65 | 31.66 | 1.000x | | 4096 x 7168 | fwd both | 30.30 | 30.00 | 1.010x | 35.61 | 34.61 | 1.029x | | 4096 x 7168 | bwd rowwise | 37.74 | 37.61 | 1.004x | 39.83 | 39.43 | 1.010x | | 4096 x 7168 | bwd colwise | 39.91 | 39.74 | 1.004x | 45.76 | 45.17 | 1.013x | | 4096 x 7168 | bwd both | 47.76 | 47.70 | 1.001x | 50.77 | 49.83 | 1.019x | | 16384 x 7168 | fwd rowwise | 87.28 | 86.46 | 1.010x | 95.77 | 95.01 | 1.008x | | 16384 x 7168 | fwd colwise | 89.98 | 89.98 | 1.000x | 102.74 | 102.75 | 1.000x | | 16384 x 7168 | fwd both | 105.32 | 105.03 | 1.003x | 123.03 | 119.22 | 1.032x | | 16384 x 7168 | bwd rowwise | 135.78 | 135.80 | 1.000x | 139.70 | 139.69 | 1.000x | | 16384 x 7168 | bwd colwise | 139.69 | 139.23 | 1.003x | 158.04 | 155.30 | 1.018x | | 16384 x 7168 | bwd both | 175.17 | 175.22 | 1.000x | 180.17 | 176.88 | 1.019x | | 131072 x 8192 | fwd rowwise | 741.29 | 741.25 | 1.000x | 819.82 | 819.83 | 1.000x | | 131072 x 8192 | fwd colwise | 756.02 | 756.01 | 1.000x | 869.90 | 869.93 | 1.000x | | 131072 x 8192 | fwd both | 904.81 | 905.38 | 0.999x | 1043.09 | 1006.06 | 1.037x | | 131072 x 8192 | bwd rowwise | 1173.91 | 1174.54 | 0.999x | 1209.91 | 1209.79 | 1.000x | | 131072 x 8192 | bwd colwise | 1222.90 | 1218.40 | 1.004x | 1376.10 | 1351.01 | 1.019x | | 131072 x 8192 | bwd both | 1552.01 | 1553.08 | 0.999x | 1586.77 | 1559.45 | 1.018x | ### NVIDIA R200 (sm_107a, 216 SMs), CUDA 13.4.92 | 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.86 | 1.86 | 1.001x | 2.12 | 2.12 | 1.000x | | 128 x 128 | fwd colwise | 2.03 | 2.04 | 1.000x | 2.28 | 2.28 | 0.999x | | 128 x 128 | fwd both | 2.09 | 2.10 | 0.997x | 2.36 | 2.36 | 1.000x | | 128 x 128 | bwd rowwise | 2.10 | 2.10 | 1.001x | 2.31 | 2.31 | 1.001x | | 128 x 128 | bwd colwise | 2.22 | 2.21 | 1.001x | 2.40 | 2.40 | 1.000x | | 128 x 128 | bwd both | 2.35 | 2.35 | 1.002x | 2.54 | 2.54 | 0.999x | | 4096 x 2048 | fwd rowwise | 7.12 | 7.12 | 1.000x | 8.09 | 8.10 | 0.999x | | 4096 x 2048 | fwd colwise | 7.12 | 7.13 | 0.999x | 9.35 | 9.35 | 1.000x | | 4096 x 2048 | fwd both | 8.58 | 8.58 | 1.000x | 10.05 | 10.04 | 1.001x | | 4096 x 2048 | bwd rowwise | 10.11 | 10.12 | 0.999x | 10.75 | 10.70 | 1.005x | | 4096 x 2048 | bwd colwise | 10.71 | 10.67 | 1.004x | 12.08 | 11.89 | 1.016x | | 4096 x 2048 | bwd both | 12.70 | 12.64 | 1.005x | 14.18 | 14.04 | 1.010x | | 4096 x 7168 | fwd rowwise | 17.73 | 17.73 | 1.000x | 19.95 | 19.95 | 1.000x | | 4096 x 7168 | fwd colwise | 18.04 | 18.04 | 1.000x | 22.14 | 22.14 | 1.000x | | 4096 x 7168 | fwd both | 20.94 | 20.95 | 1.000x | 25.29 | 25.33 | 0.999x | | 4096 x 7168 | bwd rowwise | 27.09 | 27.04 | 1.002x | 28.74 | 28.50 | 1.009x | | 4096 x 7168 | bwd colwise | 27.73 | 27.63 | 1.004x | 30.47 | 30.20 | 1.009x | | 4096 x 7168 | bwd both | 33.97 | 33.90 | 1.002x | 36.06 | 35.77 | 1.008x | | 16384 x 7168 | fwd rowwise | 56.74 | 56.73 | 1.000x | 64.44 | 64.45 | 1.000x | | 16384 x 7168 | fwd colwise | 58.88 | 57.73 | 1.020x | 68.45 | 67.13 | 1.020x | | 16384 x 7168 | fwd both | 70.43 | 70.44 | 1.000x | 75.70 | 75.75 | 0.999x | | 16384 x 7168 | bwd rowwise | 92.89 | 92.88 | 1.000x | 98.60 | 98.59 | 1.000x | | 16384 x 7168 | bwd colwise | 94.04 | 94.08 | 1.000x | 103.08 | 103.03 | 1.000x | | 16384 x 7168 | bwd both | 123.48 | 123.33 | 1.001x | 128.22 | 127.33 | 1.007x | | 131072 x 8192 | fwd rowwise | 461.96 | 461.97 | 1.000x | 514.87 | 513.92 | 1.002x | | 131072 x 8192 | fwd colwise | 488.52 | 483.29 | 1.011x | 550.19 | 539.52 | 1.020x | | 131072 x 8192 | fwd both | 609.64 | 609.62 | 1.000x | 620.94 | 621.04 | 1.000x | | 131072 x 8192 | bwd rowwise | 796.23 | 796.19 | 1.000x | 842.92 | 842.89 | 1.000x | | 131072 x 8192 | bwd colwise | 809.44 | 809.45 | 1.000x | 867.29 | 867.65 | 1.000x | | 131072 x 8192 | bwd both | 1095.31 | 1094.10 | 1.001x | 1126.13 | 1119.32 | 1.006x | 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 processing on additional GPU architectures. * **Improvements** * Updated forward and backward processing paths with packed-sigmoid variants and refined activation arithmetic. * Expanded the available kernel implementations for gated-activation MXFP8 workloads.  Co-authored-by: Yingyi Huang <averyh@nvidia.com>

flashinfer-bot added the op: misc norm, activation, sampling, RoPE, quantization, etc. label Oct 25, 2026

yzh119 approved these changes Oct 25, 2026

View reviewed changes

yyihuang mentioned this pull request Oct 25, 2026

[CAKE] Activation & Quantization Kernel Progress Tracker #6598

Open

1 task done

flashinfer-bot added run-ci and removed run-ci labels Oct 25, 2026

flashinfer-bot added the run-ci label Oct 25, 2026

Merge branch 'main' into cake-gated-act-mxfp8-narrow-tiles

4a4b1a6

flashinfer-bot added run-ci and removed run-ci labels Oct 25, 2026

yyihuang merged commit 0b352a6 into flashinfer-ai:main Oct 25, 2026
19 checks passed

yyihuang mentioned this pull request Oct 26, 2026

perf(cake_gated_act_mxfp8): packed f32x2 sigmoid where it pays (SM100 / SM103 / SM107) #7396

Merged

yyihuang mentioned this pull request Nov 5, 2026

[Cake 12/13] gated_act_mxfp8: fused SwiGLU + MXFP8 quantization, forward and backward radixark/miles-ext-kernels#13

Merged

Navigation Menu

Repository navigation

perf(cake_gated_act_mxfp8): narrow tiles for small problems and four rows per thread for large ones - #7346

perf(cake_gated_act_mxfp8): narrow tiles for small problems and four rows per thread for large ones#7346
yyihuang merged 2 commits into
flashinfer-ai:mainfrom
yyihuang:cake-gated-act-mxfp8-narrow-tiles

yyihuang commented Oct 25, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Oct 25, 2026 •

edited

Loading

Review skipped

Uh oh!

yyihuang commented Oct 25, 2026

Uh oh!

yyihuang commented Oct 25, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

yyihuang commented Oct 25, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.4.92

NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4.92

NVIDIA R200 (sm_107a, 208 SMs), CUDA 13.4.92

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Oct 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

yyihuang commented Oct 25, 2026

Uh oh!

yyihuang commented Oct 25, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yyihuang commented Oct 25, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Oct 25, 2026 •

edited

Loading