Conversation
…n size too vllm-project#55180 switched the tile scheduler to a swizzled CTA order whenever the weight exceeds the L2. Above the L2 that order is not unconditionally better: on a GB10 (24 MiB L2) it loses to the default order across a band of middling activation sizes and wins by up to 3.2x above it, so a weight-only gate leaves a measurable amount on the table. Measured with the swizzle forced both ways at the same shape, six shapes x M = 1024..6144 in steps of 512 (66 cells), four starts, medians, bit-identical on every cell. Summed percentage left on the table against the per-cell best: never swizzle 732.7 weight > L2 (now) 118.8 worst cell 9.9 % + M*K >= 14 MiB 54.0 worst cell 7.4 % + (M <= 2048 || M*K >= 14 MiB) 48.6 worst cell 8.7 % + (M <= 1024 || M*K >= 14 MiB) (this) 35.3 worst cell 5.8 % The worst cells for the weight-only gate are at M = 2560 (5120x5120 -10.4 %, 10240x2560 -9.4 %, 16384x2560 -6.4 %), which an earlier coarser sweep stepping M by 2048 did not sample. The small-M island stops at 1024 rather than 2048 because it is K-dependent: at K = 2560 the swizzle wins at low M, at K = 5120 the default order does, so a wider island gives back 6.9-8.5 % on the 5120-wide weights. The large wins the gate must not lose are all at an activation slab above ~22 MiB (7168x5120 at M = 6144 is +223 %). The kernel launch is unchanged on parts whose L2 holds the weight, so SM120 desktop parts (96-128 MiB L2) are unaffected, as before. Tests: the two existing prefill cases are joined by M = 2560 and M = 1024 on 16384x2560, so all three arms of the gate are exercised. Co-authored-by: Claude <noreply@anthropic.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011SuBgdp87NbfLbiigmzn1z Signed-off-by: Jürgen Schmied <juergenschmied70@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe blockwise FP8 CUTLASS path now uses activation dimensions when selecting CTA swizzle order. The tests add prefill shapes that cover the activation-size gate, small-M gate, default ordering, and non-swap-AB dispatch. ChangesFP8 swizzle gating
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to This updates SM120 blockwise FP8 CTA ordering to account for activation size while preserving default ordering when weights fit in L2. The changed gate branches are covered by added prefill cases, with no remaining merge-readiness risk identified. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
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 |
|
I don't think we will merge this. The new logic looks a bit "overfitting" to my taste. I don't quite understand your table but I think the speedup is not too significant -> likely not noticeable e2e. |
|
Measured it end to end, and the objection was right — closing. Server A/B, six arms, gate as the single variable (both arms route the blockwise GEMM through the same op, differing only in the swizzle the policy picks; one chunk per prompt so M is the prompt length, prefix cache disabled by construction, 3 starts per arm interleaved):
The 7.5k row is the control: both policies emit the same kernel, and it "improves" as much as every row where they differ. The whole signal is an arm-order offset, so the effect is zero within noise. And that null is structural, not a sample-size problem. Over all 154 cells I swept, split by whether this gate changes the decision:
On "overfitting": also correct. On 8 held-out shapes (including two K values absent from the tuning set) the small-M island contributes exactly zero — slab-only and island both score 236.3 against the weight-only rule's 299.1. Its entire benefit was confined to the data it was fitted on. One thing worth separating out, since it is in merged code rather than this PR: Thanks for the quick and correct pushback. |
2 participants
Purpose
#55180 (merged) switches the SM 12.x blockwise FP8 tile scheduler to a swizzled CTA order whenever the weight exceeds the L2. Above the L2 that order is not unconditionally better, so the weight-only gate leaves a measurable amount on the table. This adds the activation term back, with the small-M island that the earlier version of it was missing.
An activation-slab term was in #55180 and was dropped during review, correctly at the time: as written it also kept the default order at small M, where the swizzle wins. The island fixes exactly that case, and the sweep below is a much finer grid than the one that informed the original.
Test Plan
Measured with the swizzle forced both ways at the same shape — a standalone build of this dispatch with
max_swizzle_sizeexposed as an argument — so the two orders are compared directly rather than inferred across a gate. 6 shapes x M = 1024..6144 in steps of 512 = 66 cells, four starts, medians. GB10 (sm_121, 24 MiB L2), TP=1. Then the same 66 cells run through the gate as committed here, checking bit-identity against the stock op and that the gated launch lands on the arm the predicate selects.pytest tests/kernels/quantization/test_cutlass_scaled_mm.py -k blockwiseTest Result
Summed percentage left on the table against the per-cell best, over all 66 cells:
weight > L2(current)+ (m <= 2560 || m*k >= 14 MiB)+ (m*k >= 14 MiB), no island+ (m <= 2048 || m*k >= 14 MiB)+ (m <= 1024 || m*k >= 14 MiB)(this PR)The current gate's worst cells are at M = 2560 — 5120x5120 −10.4 %, 10240x2560 −9.4 %, 16384x2560 −6.4 % — a point the coarser sweep behind #55180 (M stepping by 2048) did not sample. At M = 4096 the 2560-wide weights lose 3–5 %.
The island stops at 1024 rather than 2048 because it is K-dependent: at K = 2560 the swizzle wins at low M, at K = 5120 the default order does, so a wider island gives back 6.9–8.5 % on the 5120-wide weights. M <= 1024 and M <= 1536 are indistinguishable (35.3 vs 35.2 pp); 2048 is not.
What the gate must not lose, and does not: the large wins, all at an activation slab above ~22 MiB — 7168x5120 at M = 6144 +223 %, M = 5632 +199 %, 14336x4096 at M = 6144 +158 %, 5120x5120 at M = 4608 +93 %.
Verification of the code as committed: bit-identical to the stock op on 66/66 cells (twice), and on the 26 cells where the two orders differ by more than 30 % the gated launch tracks the arm the predicate selects, 26/26.
Honest limits: this is one part (GB10, 24 MiB L2). Parts whose L2 holds the weight are unaffected by construction, so SM120 desktop parts (96–128 MiB L2) keep the default order exactly as before. A scalar gate cannot capture all of the structure — 17 of 66 cells still take the worse order, all by 5.8 % or less. And the default order is the noisy one here: its start-to-start spread is median 4.7 % and up to 26 %, against ~1.5 % for the swizzled order, which is why this needed four starts; at two starts the island candidates could not be separated at all.
Essential Elements
gh pr list --searchonscaled_mm_blockwise_sm120,max_swizzle_size,swizzle,blockwise fp8 sm120— no open PR touches this dispatchcc @gau-nernst — this is the follow-up to the gate we simplified in #55180.