Repository navigation
perf(cake_gated_act_mxfp8): band-path program for the SM107 forward both entry below 2^20 - #8023
Conversation
…oth entry below 2^20 - SM107 forward both, M * K below 2^20: the plain program with the call band path at 4 CTAs per SM (new record). - From 2^20 up to 2^23 the previous plain program still runs; its record now starts at 2^20 on SM107. Larger sizes and SM100 / SM103 are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
flashinfer-bot
added
the
op: misc
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Gated activation kernel flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu |
Adds BF16 gated-activation computation, including a fast path and an alternate path that recomputes pairs with zero gate factors. |
Row and column quantization flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu |
Computes row and column maxima and scales, converts values to saturated E4M3, and writes both output layouts. |
Kernel binding and launch flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu |
Validates tensor and grid properties, obtains the device stream, launches the kernel, and reports validation or launch errors. |
Library registration and dispatch flashinfer/cake_gated_act_mxfp8_backend/build.json, flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py |
Adds the sm_107a build manifest entry and updates fwd_both routing. The new fallback handles calls below the threshold-selected record. |
Priority:
Estimated code review effort: 4 (Complex) | ~60 minutes
Change: Refactor
Sequence Diagram(s)
sequenceDiagram
participant fwd_both
participant run
participant kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950
participant row_output
participant col_output
fwd_both->>run: selects and invokes the library entry
run->>kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950: launches on the device stream
kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950->>row_output: writes row-oriented E4M3 data and scales
kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950->>col_output: writes column-oriented E4M3 data and scales
Suggested labels: op: misc
Suggested reviewers: aleozlx, sricketts
Merge Risk: 
c11c8
-
This change adds a new program selection for sm_107a forward-both calls below 2^20 elements and moves one existing selection bound. No concrete merge-blocking risk was found in the supplied changes. Hardware validation on R200 remains with the author's reported measurements.
Pre-merge checks: 
-
Check name Status Explanation Resolution Description check The description gives a detailed change summary, scope, benchmark setup, and performance results. However, it omits the required template sections and checklist, including related issues, pre-commit s… Add the required template headings and checklist. State related issues or indicate none, report pre-commit and test results, declare whether full-scope CI is required with a reason, complete the experimental-track section when applicable, a… Docstring Coverage Docstring coverage is 6.25% 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. -
Check name Status Explanation Title check The title clearly identifies the performance change: adding a band-path program for SM107 forward-both calls below 2^20. 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 a detailed change summary, scope, benchmark setup, and performance results. However, it omits the required template sections and checklist, including related issues, pre-commit status, explicit test status, CI scope, and reviewer notes.
Resolution
Add the required template headings and checklist. State related issues or indicate none, report pre-commit and test results, declare whether full-scope CI is required with a reason, complete the experimental-track section when applicable, and add reviewer notes or indicate none.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 6.25% 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 16 functions across 3 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.
|
@flashinfer-bot run tests/utils/test_cake_gated_act_mxfp8.py tests/utils/test_cake_gated_act_mxfp8_api.py |
The Cake fused SwiGLU + MXFP8 quantization kernel computes silu(gate) * up from a BF16 input [M, 2K] and quantizes the result to MXFP8. It stores one E8M0 scale (RCEIL rounding) per 32 E4M3 values, 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_quantizeand the backward form isflashinfer.silu_and_mul_mxfp8_quantize_backward. Each 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 program selection for one forward entry on sm_107a (R200). When M * K is below 2^20, forward both now runs the called band path at 4 CTAs per SM. From 2^20 up to 2^23, the previous program still runs, with no band-path or occupancy setting. Sizes above that are unchanged. The sm_100a and sm_103a programs are unchanged. The package goes from 74 to 75 per-arch programs across 62 records. One record is new and applies to sm_107a only, and one existing record now starts at 2^20 on sm_107a. The other libraries keep the same machine code.The baseline is the previous Cake version of this kernel, which is the package on FlashInfer main from #7880. With block-scaled inputs the changed owner row is faster: 128 x 128 forward both is 1.076x (2.36 us -> 2.20 us). The same entry also covers 256 x 256 and 512 x 512 forward both. Measured outside the owner list, these are 1.085x and 1.081x with block-scaled inputs. With unit-normal inputs the changed rows tie within the run-to-run spread (128 x 128 forward both is 1.004x). Every other row runs the same machine code as before and stays within run-to-run noise. Over all 30 rows, unit normals range from 0.998x to 1.004x and block-scaled inputs from 0.999x to 1.076x.
The table covers the owner rows on NVIDIA R200 (sm_107a, 216 SMs) with CUDA 13.4: 5 shapes x 6 passes, 30 rows. Both versions run in the same process on the same GPU, in 8 alternating paired rounds, with a cold L2 before every launch. Latency is the mean over pairs of each sample's median CUPTI kernel time, in microseconds. Speedup is the previous latency divided by this PR's latency, so a value above 1 means this PR is faster. Block-scaled inputs are normals with per-32-block scales spread over a wide exponent range. Outputs are byte-identical to the previous version, and the public API is unchanged.
NVIDIA R200 (sm_107a, 216 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