Skip to content

perf(cake_dense_projection_gemm): 212-SM sm_107a router weight-gradient windows run the head/tail program with the A parts in tensor memory - #8517

Merged
yyihuang merged 3 commits into
flashinfer-ai:mainfrom
yyihuang:cake-dense-projection-gemm-router-ts-sm107a-212
Nov 3, 2026
Merged

yyihuang merged 3 commits into
flashinfer-ai:mainfrom
yyihuang:cake-dense-projection-gemm-router-ts-sm107a-212

Conversation

@yyihuang

@yyihuang yyihuang commented Nov 3, 2026 •

Copy link
Copy Markdown
Collaborator

The flashinfer.experimental.dense_projection_gemm family has three kinds of GEMM. They are BF16-input dense projection GEMMs, BF16-input batched projection GEMMs, and an FP32 router GEMM (256 experts, BF16x3 emulation). Each kind has forward, input-gradient and weight-gradient programs. deterministic=True (the default) calls dense_projection_det_gemm. deterministic=False runs this package's split-K / stream-K programs. This PR changes one variant: the router weight gradient (FP32 out, 256 experts) on sm_107a parts with exactly 212 SMs, at 16172 to 16384 tokens.

At hidden 5120, this window already exists on FlashInfer main with plan "H0.5x4" (106 head and 424 tail units). In that window, router_fp32_gemm_ts_r6p6_atomic_nn_t_ht_c2_nreg replaces router_fp32_gemm_atomic_nn_t_ht_c3_nreg. The new program is the head/tail program that already ships for sm_103a. It keeps the three BF16 parts of A in tensor memory. The partition, the products and their accumulation order are unchanged.

At hidden 6144, this PR adds a new window with plan "H0.5x4". That window runs only the tensor-memory program. On FlashInfer main, this row runs the deterministic program at 212 SMs.

ROUTER_HEAD_TAIL_TEMPLATES_BY_SM selects the program by SM count. 208-SM and 216-SM sm_107a parts keep their programs and plans. sm_100a and sm_103a also keep their programs and plans. Hidden 4096 keeps the deterministic plan. The forward and input gradient are unchanged.

These results are on R200, against FlashInfer main. Each range below covers all runs. At hidden 5120:

  • The router weight gradient gives det/new 1.299-1.360, previous/new 1.048-1.091 and SOTA/new 1.809-1.875.
  • bwd gives det/new 1.130-1.191, previous/new 1.025-1.064 and SOTA/new 1.446-1.524.
  • 2fwd+bwd gives det/new 1.057-1.084, previous/new 1.010-1.031 and SOTA/new 1.732-1.807.

At hidden 6144:

  • The router weight gradient gives det/new 1.147-1.201 and SOTA/new 1.827-1.915. Previous Cake runs the deterministic program here, so its weight-gradient column equals det Cake.
  • bwd gives det/new 1.042-1.100, previous/new 1.041-1.098 and SOTA/new 1.454-1.521.
  • 2fwd+bwd gives det/new 1.015-1.043, previous/new 1.019-1.045 and SOTA/new 1.706-1.840.

The forward runs the same program in every Cake arm. Its previous/new ratio is 0.979-1.008 at hidden 5120 and 0.994-1.017 at hidden 6144.

An A/A control runs the new program against a second prepared copy of itself. It ran in 6 of the runs at each hidden size. On the weight gradient, it spans 0.994-1.009. On bwd, it spans 1.001-1.030 at hidden 5120 and 1.011-1.037 at hidden 6144. On 2fwd+bwd, it spans 0.994-1.010 at hidden 5120 and 0.941-1.009 at hidden 6144.

The weight-gradient and bwd gains over det Cake are larger than the A/A spread. At hidden 5120 the previous Cake arm is also a separately prepared copy, and the A/A control shows that such a copy can be up to 3 % slower on bwd. Against the second copy of the new program, previous/new bwd is 1.007-1.038 (in the 6 A/A runs). So the hidden-5120 bwd previous/new ratios in the table (1.045-1.051) may be up to about 3 % too high. The hidden-5120 2fwd+bwd gain over det Cake is also larger than the A/A spread. The hidden-5120 2fwd+bwd previous/new ratio (1.010-1.031) is about as large as its A/A spread (0.994-1.010). At hidden 6144, the 2fwd+bwd gain over both Cake baselines (1.015-1.045) is smaller than its A/A spread (0.941-1.009). That row is parity with no regression, not a measured gain.

  • Rows: FP32 router GEMM, hidden 5120 and 6144, 256 experts, T 16172, 16231 and 16384 tokens (the windows this PR changes) on a 212-SM sm_107a part. Ops are wgrad (router weight gradient, the changed op), bwd (input gradient + weight gradient) and 2fwd+bwd (checkpointed step: 2 x forward + bwd). Modes are eager (calls of the prepared launches) and graph (CUDA graph replays).
  • Latency: The table reports median us of the op's GPU span from CUPTI activity records.
    • The span runs from the first kernel begin to the last kernel end of every launch of the op. It includes the BF16x3 split and memset launches.
    • L2 is flushed before every iteration. Programmatic dependent launch is at the package default.
    • Hidden 5120 has 12 runs on four GPUs. Hidden 6144 has 10 runs on four GPUs.
    • Half of the runs use one long-lived process for all rows. The other half use a fresh process per row in reverse row order.
    • In each run, each arm runs 15 paired groups x 20 timed steps in rotated arm order. The per-run cell is the median of the per-group medians. A table cell is the median of the per-run cells.
  • Baselines: det Cake = deterministic=True (dense_projection_det_gemm).
    • Previous Cake = deterministic=False on FlashInfer main. At hidden 5120, it runs router_fp32_gemm_atomic_nn_t_ht_c3_nreg at the same plan. At hidden 6144, it runs the deterministic program, which is main's plan for that row.
    • SOTA (every GEMM of the op) = cuBLAS 13.8.0 cublasSgemm with CUBLAS_FP32_EMULATED_BF16X9_MATH.
    • Plain FP32 cuBLAS takes 1.80-2.71x the SOTA time on the hidden-5120 rows. TF32 and BF16 miss the router's FP32 accuracy budget. This makes SOTA the strongest applicable library baseline.
  • Setup: NVIDIA R200, sm_107a, 212 SMs, driver 620.43, CUDA 13.4 (nvcc 13.4.92), PyTorch 2.14.1+cu130, cuBLAS 13.8.0.
    • The GPUs ran at their software power cap part of the time. nvidia-smi was sampled every 200 ms only in 6 of the 12 hidden-5120 runs and 6 of the 10 hidden-6144 runs (the runs on two of the four GPUs). It reported clock event reason 0x4 in 17-47 % of the non-idle samples.
    • In those samples, the median SM clock was 2388-2400 MHz (2424 MHz maximum).
    • A ratio is the quotient of the two us columns. A ratio above 1.00 means this PR is faster.
  • Accuracy: The new program's relative Frobenius error against FP64 is within the 2e-6 budget. This holds in every run at both hidden sizes.
GPU shape op mode det Cake us previous Cake us SOTA us new us det/new previous/new SOTA/new
R200 router 5120 x 256, T 16172 wgrad eager 132.97 107.88 183.18 99.97 1.330 1.079 1.832
R200 router 5120 x 256, T 16172 wgrad graph 132.16 106.44 183.19 98.28 1.345 1.083 1.864
R200 router 5120 x 256, T 16172 bwd eager 271.1 243.8 341.9 233.3 1.162 1.045 1.465
R200 router 5120 x 256, T 16172 bwd graph 267.8 241.8 341.8 230.9 1.160 1.047 1.480
R200 router 5120 x 256, T 16172 2fwd+bwd eager 524.2 502.2 866.0 489.5 1.071 1.026 1.769
R200 router 5120 x 256, T 16172 2fwd+bwd graph 518.3 496.2 858.4 485.3 1.068 1.022 1.769
R200 router 5120 x 256, T 16231 wgrad eager 133.02 108.20 183.69 100.23 1.327 1.080 1.833
R200 router 5120 x 256, T 16231 wgrad graph 131.53 106.78 184.08 98.64 1.333 1.083 1.866
R200 router 5120 x 256, T 16231 bwd eager 272.9 245.2 343.4 233.8 1.167 1.049 1.469
R200 router 5120 x 256, T 16231 bwd graph 268.2 242.6 343.5 230.9 1.162 1.051 1.488
R200 router 5120 x 256, T 16231 2fwd+bwd eager 525.1 500.7 873.2 489.9 1.072 1.022 1.782
R200 router 5120 x 256, T 16231 2fwd+bwd graph 518.6 495.0 856.9 483.1 1.073 1.025 1.774
R200 router 5120 x 256, T 16384 wgrad eager 134.09 108.73 185.41 101.41 1.322 1.072 1.828
R200 router 5120 x 256, T 16384 wgrad graph 133.08 107.36 185.28 99.84 1.333 1.075 1.856
R200 router 5120 x 256, T 16384 bwd eager 272.6 246.5 346.1 235.0 1.160 1.049 1.473
R200 router 5120 x 256, T 16384 bwd graph 268.3 244.1 346.1 232.5 1.154 1.050 1.489
R200 router 5120 x 256, T 16384 2fwd+bwd eager 530.2 505.1 879.5 495.2 1.071 1.020 1.776
R200 router 5120 x 256, T 16384 2fwd+bwd graph 522.6 498.0 865.5 487.3 1.072 1.022 1.776
R200 router 6144 x 256, T 16172 wgrad eager 135.30 135.30 212.05 114.24 1.184 1.184 1.856
R200 router 6144 x 256, T 16172 wgrad graph 133.56 133.56 211.78 112.43 1.188 1.188 1.884
R200 router 6144 x 256, T 16172 bwd eager 293.2 293.1 402.8 271.8 1.079 1.078 1.482
R200 router 6144 x 256, T 16172 bwd graph 289.3 290.0 403.1 269.8 1.072 1.075 1.494
R200 router 6144 x 256, T 16172 2fwd+bwd eager 592.7 595.8 1019.4 574.9 1.031 1.036 1.773
R200 router 6144 x 256, T 16172 2fwd+bwd graph 584.2 587.6 1004.7 568.5 1.028 1.034 1.767
R200 router 6144 x 256, T 16231 wgrad eager 135.41 135.41 213.31 114.26 1.185 1.185 1.867
R200 router 6144 x 256, T 16231 wgrad graph 133.59 133.59 213.44 112.24 1.190 1.190 1.902
R200 router 6144 x 256, T 16231 bwd eager 295.4 293.8 403.8 273.2 1.081 1.075 1.478
R200 router 6144 x 256, T 16231 bwd graph 291.0 290.2 403.9 271.7 1.071 1.068 1.487
R200 router 6144 x 256, T 16231 2fwd+bwd eager 592.0 594.8 1033.0 573.5 1.032 1.037 1.801
R200 router 6144 x 256, T 16231 2fwd+bwd graph 582.9 586.1 1012.2 565.9 1.030 1.036 1.789
R200 router 6144 x 256, T 16384 wgrad eager 136.66 136.66 215.06 114.63 1.192 1.192 1.876
R200 router 6144 x 256, T 16384 wgrad graph 134.69 134.69 215.13 113.10 1.191 1.191 1.902
R200 router 6144 x 256, T 16384 bwd eager 296.3 295.7 405.9 273.9 1.082 1.080 1.482
R200 router 6144 x 256, T 16384 bwd graph 292.6 292.3 407.1 271.4 1.078 1.077 1.500
R200 router 6144 x 256, T 16384 2fwd+bwd eager 597.4 599.5 1039.8 578.2 1.033 1.037 1.798
R200 router 6144 x 256, T 16384 2fwd+bwd graph 589.3 590.6 1023.3 571.4 1.031 1.034 1.791

The other variants of the family keep their programs, and the family ships on sm_100a, sm_103a and sm_107a. These other variants are the projection GEMMs, the router forward and input gradient, and the router weight gradient on sm_100a, sm_103a, 208-SM and 216-SM sm_107a parts, at hidden 4096 and outside these windows.

Summary by CodeRabbit

  • New Features
    • Added support for router weight-gradient workloads with hidden size 6144 on 212-SM sm_107a devices.
  • Performance
    • Updated benchmark results for hidden sizes 5120 and 6144, including timing and resident-kernel measurements.
  • Bug Fixes
    • Improved compatibility of tensor-memory operations on sm_107a devices.
  • Tests
    • Expanded coverage for supported device configurations and routing behavior.

…'s TMEM head/tail program for sm_107a

The head/tail program with the BF16 parts of A in tensor memory
(cake_dense_projection_gemm_d384571ea88c28e40e31) serves sm_107a too: its
record's arches gain sm_107a, KERNELS["sm_107a"] names its template and
build.json lists its sm_107a library. The kernel unit is the current render
of the program (with the sm_107a branch of its B shared-memory descriptor
mask); every other generated file is unchanged.
…EM head/tail program on 212-SM sm_107a parts

On sm_107a parts with exactly 212 SMs the weight gradient's head/tail
windows (256 experts, 16172 to 16384 tokens) run the head/tail program
with the BF16 parts of A in tensor memory
(router_fp32_gemm_ts_r6p6_atomic_nn_t_ht_c2_nreg) instead of
router_fp32_gemm_atomic_nn_t_ht_c3_nreg:

- ROUTER_HEAD_TAIL_TEMPLATES_BY_SM maps (arch, sm_count) to the head/tail
  program of the parts of that SM count; router_head_tail_template(arch,
  sm_count) reads it before ROUTER_HEAD_TAIL_TEMPLATES, and
  plan_atomic_router passes the device's SM count.  208-SM and 216-SM
  sm_107a parts keep the three-chunk-buffer program.
- The 212-SM table of ROUTER_RULES_BY_SM gains hidden 6144 "H0.5x4" at
  16172 to 16384 tokens, a window that holds only with the tensor-memory
  program.  Hidden 4096 keeps the deterministic plan.

Tests: the planner, partition and per-SM-count template checks cover the
208, 212 and 216-SM tables; the GPU window test runs the program in the
212-SM windows.  README: the 212-SM rows and their measurements.
@coderabbitai

coderabbitai Bot commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

No actionable comments were generated in the recent review.

Walkthrough

The router planner adds a head/tail weight-gradient route for hidden size 6144 on 212-SM sm_107a devices. Template selection now accounts for SM count. The kernel registry, architecture-specific tensor-memory operations, tests, and measured results are updated.

Changes: Router Weight-Gradient Routing

Layer / File(s) Summary
Per-SM router policy and template selection
flashinfer/experimental/dense_projection_gemm/cake_backend.py, tests/experimental/test_cake_dense_projection_gemm.py
The planner adds an H0.5x4 route for hidden size 6144 on 212-SM sm_107a devices at 16,172–16,384 tokens. Template selection uses architecture and SM count. Tests cover the new route, its boundaries, and other SM-count behavior.
sm_107a kernel support and validation
flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/*, flashinfer/experimental/dense_projection_gemm/cake_jit.py, flashinfer/experimental/dense_projection_gemm/build.json, tests/experimental/test_cake_dense_projection_gemm.py, flashinfer/experimental/dense_projection_gemm/README.md
The registered kernel uses architecture-specific tensor-memory operand masks and unsigned addresses. Build metadata and tests add sm_107a support. The README records the 212-SM measurements and compares the tensor-memory and three-chunk-buffer programs.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant plan_atomic_router
  participant router_head_tail_template
  participant cake_jit_registry
  plan_atomic_router->>router_head_tail_template: Select template for architecture and SM count
  router_head_tail_template-->>plan_atomic_router: Return selected head/tail template
  plan_atomic_router->>cake_jit_registry: Resolve template to registered kernel module
Loading

Possibly related issues

Suggested reviewers: yzh119

Merge Risk: Minimal · up to 6e3f8

The new routing matches the documented architecture requirements. No actionable merge-blocking issue remains; merge after normal checks.

Pre-merge checks: Passed 4 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Description check Warning The description follows the template's main sections: a Description covering the change, benchmark results, and setup, plus notes on the A/A control and baselines. It does not fill in the Related Issu… Add a Related Issues entry (the PR cites #8517 in the objectives, so link it or the tracking issue). Complete the Tests checklist and state whether the PR needs full-scope CI. Because the PR changes files under flashinfer/experimental/, com…
Passed Passed checks: 4 passed
Check name Status Explanation
Title check Passed The title names the change accurately. It covers the 212-SM sm_107a router weight-gradient windows and the switch to the head/tail program with A parts in tensor memory, which matches the raw summary …
Docstring Coverage Passed Docstring coverage is 81.82% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati…
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Full details: Title check

Explanation

The title names the change accurately. It covers the 212-SM sm_107a router weight-gradient windows and the switch to the head/tail program with A parts in tensor memory, which matches the raw summary and objectives. It uses the conventional perf(scope) prefix. It is long (139 characters), but it is specific and not misleading.

Full details: Docstring Coverage

Explanation

Docstring coverage is 81.82% which is sufficient. 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 11 functions across 4 files. (2 skipped: 2 unsupported.)

Full details: Description check

Explanation

The description follows the template's main sections: a Description covering the change, benchmark results, and setup, plus notes on the A/A control and baselines. It does not fill in the Related Issues section, the Pre-commit and Tests checklists, the Experimental Track section (which the PR touches via flashinfer/experimental/ and needs an experimental-tests fence), or Reviewer Notes. The experimental track items are required for this PR because it changes code under flashinfer/experimental/. The description is mostly complete on the technical content but leaves required template sections unaddressed.

Resolution

Add a Related Issues entry (the PR cites #8517 in the objectives, so link it or the tracking issue). Complete the Tests checklist and state whether the PR needs full-scope CI. Because the PR changes files under flashinfer/experimental/, complete the Experimental Track section: fill in the tracking issue number and its owner, graduation plan, and the checklist items. Replace the commented placeholder in the experimental-tests fence with the test targets (for example tests/experimental/test_cake_dense_projection_gemm.py), keeping the fence tag unchanged. Update the Pre-commit checklist status if the hooks were run.

Finishing Touches
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

CodeRabbit Verify

  • Run CodeRabbit Verify to test the changes in this pull request.
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: 740fb86a-f373-450f-ae87-a9b9f7a7df20
Commits

Reviewing files that changed from the base of the PR and between afed932 and 6e3f85e.

Files selected for processing: 6
  • flashinfer/experimental/dense_projection_gemm/README.md
  • flashinfer/experimental/dense_projection_gemm/build.json
  • flashinfer/experimental/dense_projection_gemm/cake_backend.py
  • flashinfer/experimental/dense_projection_gemm/cake_jit.py
  • flashinfer/experimental/dense_projection_gemm/csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_d384571ea88c28e40e31_kernel.cu
  • tests/experimental/test_cake_dense_projection_gemm.py

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

@yyihuang

yyihuang commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py

…amples cover

The README's 212-SM clock note said the 200 ms nvidia-smi samples were
taken while the head/tail window rows ran.  Only 6 of the 12 runs at
hidden 5120 and 6 of the 10 at hidden 6144 (those on two of the four
parts) were sampled; the 17-47 % share of clock event reason 0x4 is of
the non-idle samples.  Say so.  No code, test or plan change.
@yyihuang

yyihuang commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py

@yyihuang

yyihuang commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py

@yyihuang
yyihuang enabled auto-merge (squash) November 3, 2026 17:42
@yyihuang
yyihuang merged commit 3b40353 into flashinfer-ai:main Nov 3, 2026
36 of 38 checks passed
yyihuang pushed a commit to yyihuang/flashinfer that referenced this pull request Nov 3, 2026
…ad-windows

Conflicts with the sm_107a 212-SM router weight-gradient windows of flashinfer-ai#8517, resolved keeping both:
- README.md: the router window table footnote keeps the 212-SM clock samples and the GB300 hidden 4096 / 6144 clock
  samples, and the GB300 bf16x9 ratio x1.500-1.640.
- tests/experimental/test_cake_dense_projection_gemm.py: the head/tail program test checks the sm_103a 152-SM windows
  (hidden 4096, 5120, 6144) and the sm_107a 212-SM windows (hidden 5120, 6144) in one window table; the GPU test of the
  tensor-memory head/tail program takes the union of both parameter lists.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants