Repository navigation
[Perf][GEMM] Read less DRAM at the power cap: one raster group for q-b, fewest CTAs on uniform coop2 grids - #2794
Merged
Conversation
…partition and B does not At the power limit DRAM reads lower the clock. On ds-v3 prefill q-b (4096x24576x1536 NT) B overflows one L2 partition while A fits, so a 16-row group re-reads each B tile from DRAM once per group. One group over all rows reads it once: DRAM reads drop from 168 MB to 94 MB and back-to-back time by 1.0-1.5%, level with cuBLASLt-best; busy time is unchanged or better. _narrow_tiles bounds its split as a 16-row group does, so its pick does not change. Part of tile-ai#2791.
…rm coop2 grids A persistent coop2 or coop2_swap grid with uniform tile widths now launches ceil(tiles / ceil(tiles / sm_count)) CTAs instead of sm_count. The busiest CTA runs as many tiles either way, while each wave reads fewer distinct panels. - Manifest busy time: every row level or faster; 4096x4096x4096 TT and 4096x7168x4096 TN 0.4-0.6% faster. - Back to back at the power limit: the same two rows 0.3-1.0% faster. - 16 off-manifest coop2/coop2_swap shapes: 0.4-2.8% faster in busy time. Ping-pong grids, which the selector already sizes to fill their waves, keep sm_count: they measured no busy gain and up to 1.3% slower back to back. Narrow and mixed-width grids keep sm_count, which their widths were chosen for. Part of tile-ai#2791
zhen8838
approved these changes
Oct 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part of #2791.
Summary
_group_size_m: when A fits in one L2 partition and B does not, one raster group spans every row, so each B tile is read from DRAM once instead of once per 16-row group. On ds-v3 prefill q-b (4096×24576×1536 NT), DRAM reads drop from 168 MB to 94 MB._narrow_tilesbounds its split as a 16-row group does, so its pick does not change._persistent_ctas: a coop2 or coop2_swap grid with uniform tile widths launchesceil(tiles / ceil(tiles / sm_count))CTAs instead ofsm_count. The busiest CTA runs as many tiles either way, and each wave reads fewer distinct panels. Ping-pong, narrow and mixed-width grids keepsm_count.Test plan
pre-commit runon the changed filespytest tests/ops/test_gemm.pyon H200: 212 passed, 24 skippedBenchmark
Environment: H200, CUDA 13.2, PyTorch 2.13.0, TileLang 0.1.12
GemmFwdOp, busy time (
bench_gemm.py)main and this PR alternating, 4 rounds each, on two idle H200s (GPU A / GPU B); medians in µs. Ratio = best baseline / TileOPs.
Every other manifest row keeps its config; on both GPUs all 24 rows read within noise of main (time new / main 0.989–1.008).
GemmFwdOp, back to back at the 700 W cap
Back-to-back calls of one op, clocks unlocked, 1 s warm-up then 1.5 s timed, alternating rounds; ratio = best baseline / TileOPs.
Off the manifest, the fewest-CTA grid is 0.4–2.8% faster in busy time on 15 of 16 coop2 / coop2_swap shapes across NT, TN and TT, and up to 2.5% faster back to back. Applied to ping-pong it measured no busy gain and 0.8–1.3% slower back to back on 3072×5120×4096, so ping-pong keeps
sm_count.Takeaways: q-b reaches cuBLASLt-best back to back, and the two TN/TT rows go from level to ahead. o-proj, down, mlp-up and wgrad-q-b still trail back to back; no selector change closes them without a busy-time regression, so they stay open under #2791 for kernel work.