[https://nvbugs/6550749][fix] Size the DeepGEMM warmup bucket stride to the 16-token config quantum - #17242
Conversation
…token config quantum
deep_gemm_gen_tuning_buckets exists so DeepGEMM JIT-compiles every kernel
config a workload needs during autotuning at startup, keeping nvcc out of
the measured window. It walked M with range(128, x, 128).
DeepGEMM's SM100 config selection changes every 16 tokens, so a stride-128
grid sampled 1 band in 8 and left the other 7 cold. Whichever live iteration
first landed in a cold band paid ~2.2s of nvcc mid-inference.
Why 16 is correct rather than merely sufficient: sm100.hpp has exactly three
M-dependent sites, and every breakpoint sits at M % 16 == 1.
* :62-63 the non-swap_ab BLOCK_M pick (m<=32 -> 32; m<=64 -> 64; else 128),
stepping at M = 33 and M = 65;
* :230 the only use of M in config scoring, ceil_div(expected_m, block_m),
stepping at M = k*block_m + 1 -- and every candidate block_m is a
multiple of 16 (swap_ab uses lcm(16, block_m_multiple_of),
non-swap_ab uses {32, 64, 128});
* expected_m equals m on this path.
num_stages is a closed form over block_m/block_n/smem and num_sms is constant,
so neither adds a breakpoint. Every band is therefore an aligned
[16k+1, 16k+16] cell containing exactly one multiple of 16, so a stride-16 grid
provably touches every band for any (N, K). A stride-1 residual sweep confirmed
it: stride 128 left 7 bands cold, stride 32 left 1 (M=385, a band exactly 16
wide), stride 16 left 0.
Second fix, same expression: round the top up to a whole quantum and include
it. The only multiple of 16 inside a band is its top, so half-open
range(128, x, 16) never warmed the band containing max_num_tokens itself -- the
M every full batch runs at. That hole predates this change (it was 128 wide) and
is why Fp8BlockScalingGemmRunner, which pins tune_max_num_tokens=4096, had
[3969, 4096] cold.
Measured on minimax_m2.5_fp8-bench-pytorch-float8-maxbs:512-maxnt:2048-
input_output_len:128,128-gpus:4, GB300, one node, byte-identical dataset, sole
variable the JIT cache state:
Total Latency iter 140 host_step_time cubins built
before, cold cache 6061.27 ms 2316.62 ms -
before, warm cache 3888.29 ms 114.08 ms -
after, cold cache NVIDIA#1 3978.59 ms 114.16 ms 37
after, cold cache NVIDIA#2 3929.05 ms 118.88 ms 37
after, warm cache 3989.01 ms 113.68 ms 0
Iteration 140 alone was 101.4% of the end-to-end gap: iters 132-139 pack the
2048-token cap and iter 140 gets the ragged 540-token leftover (ctx 416 + gen
124), whose config (BLOCK_M=144, num_stages=8) was never warmed. This is why
the bug was filed as instability -- whether a run pays it depends on request
packing, not on the commit.
After the change cold and warm are indistinguishable; both cold reps came in
marginally faster than warm, so no cache-state signal remains. The warm rep
compiling 0 cubins while holding 37 settles completeness independently of any
timing. The compiles moving to startup is visible in the cache dir alone: the
last kernel.cubin was written 20s after the first benchmark iteration before,
and 18s before it after.
Cost is +4 compiles (33 -> 37 cubins) for 46 -> 264 buckets: the number of
distinct configs is bounded by the BLOCK_M ladder, not by how densely M is
sampled. The extra buckets are warm GEMM calls at ~4.3ms each during startup
autotuning, and they cannot change which kernel steady-state inference picks --
selection is a function of M either way.
The max(x, 4096) floor is left alone and now carries a comment recording why it
is load-bearing: fp8SwapABGemmRunner leaves tune_max_num_tokens unset, so
autotuner.py passes the current input size rather than a maximum.
Adds tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py (25 tests, no
GPU). The function had no coverage before, which is how a stride this coarse
survived; the tests assert the property that matters -- for every reachable M
some bucket selects the same config -- exhaustively at stride 1 over every
BLOCK_M DeepGEMM can choose. They found the top-band hole, which no benchmark
here could: this bug's case is maxnt:2048 and never reaches that band.
Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
WalkthroughDeepGEMM tuning buckets now support declared maximums and use a 16-token quantum above 128. FC1, FC2, and FP8 tuning callbacks use declared maximums. CPU-only tests cover bounds, coverage, spacing, legacy behavior, and measured cold bands. ChangesDeepGEMM tuning bucket generation
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant TuningConfig
participant BucketGenerator
participant DeepGEMMAutotuner
TuningConfig->>BucketGenerator: configure x_is_declared_max=True
BucketGenerator->>BucketGenerator: generate bounded 16-token tuning buckets
BucketGenerator-->>DeepGEMMAutotuner: return tuning bucket tuple
DeepGEMMAutotuner->>DeepGEMMAutotuner: profile configured buckets
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/utils.py (1)
365-385: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the public return type and API docstring.
Add
-> tuple[int, ...]todeep_gemm_gen_tuning_buckets. Add a Google-style docstring that documentsxand the returned bucket sequence. As per coding guidelines, “Annotate every function” and “Prefer docstrings for external interfaces.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/utils.py` around lines 365 - 385, Add the public return annotation tuple[int, ...] to deep_gemm_gen_tuning_buckets and add a concise Google-style docstring documenting the x parameter and returned tuning-bucket sequence. Keep the existing bucket-generation logic unchanged.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py`:
- Line 1: Add the standard required NVIDIA copyright header at the beginning of
test_deep_gemm_tuning_buckets.py, placing it before the module docstring.
Preserve the existing deep_gemm_gen_tuning_buckets coverage test content
unchanged.
---
Nitpick comments:
In `@tensorrt_llm/_torch/utils.py`:
- Around line 365-385: Add the public return annotation tuple[int, ...] to
deep_gemm_gen_tuning_buckets and add a concise Google-style docstring
documenting the x parameter and returned tuning-bucket sequence. Keep the
existing bucket-generation logic unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 537aa206-1877-4541-bb41-dcb3535e4662
📒 Files selected for processing (2)
tensorrt_llm/_torch/utils.pytests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py
| # max_num_tokens cold, which is the single most likely M of all (every | ||
| # full batch hits it). Pre-fix the same hole was 128 wide. | ||
| top = -(-x // DEEP_GEMM_BLOCK_M_QUANTUM) * DEEP_GEMM_BLOCK_M_QUANTUM | ||
| buckets += tuple(range(128, top + 1, DEEP_GEMM_BLOCK_M_QUANTUM)) |
There was a problem hiding this comment.
This will increase the tuning numbers a lot. Is it expected?
There was a problem hiding this comment.
Yes, expected — but you're right that it's more than it needs to be, and I'd like to fix half of it.
Measured cost. 46 -> 264 buckets, and the autotune window (FMHA JIT warmup marker -> "Run warmup") goes 107.6s -> 110.3s, so +2.7s (+2.5%) at ~12.2ms per extra bucket. Compiles only go 33 -> 37 cubins: the distinct-config count is bounded by the BLOCK_M ladder (~17/shape), not by bucket density, so the extra buckets are warm GEMM calls, not nvcc. They also cannot change steady-state kernel choice — selection is a function of M either way.
Where the waste actually is. You're right that the grid is mostly redundant: 264 buckets collapse to 37 distinct configs (86% duplicates). But the stride isn't the culprit — a greedy minimum cover over the conservative BLOCK_M superset needs 263 buckets, versus my 264. Stride 16 is essentially minimal for a guaranteed cover, because a band can be exactly 16 wide (measured: stride 32 leaves M=385 cold).
The real waste is the pre-existing max(x, 4096) floor. This workload is maxnt:2048, so:
| buckets | reachable | unreachable | |
|---|---|---|---|
maxnt:1024 |
264 | 72 | 192 |
maxnt:2048 |
264 | 136 | 128 (~1.6s of the +2.7s) |
maxnt:4096 |
264 | 264 | 0 |
So ~48% of the buckets on this case tune M values the workload can never reach. My standalone coverage probe measured exactly that reachable-clamped stride-16 set — 136 buckets — and it came back with residual 0 on a stride-1 sweep, i.e. provably complete for both live shapes. Same coverage, half the buckets, ~1.6s cheaper.
Why I didn't just drop the floor here. fp8SwapABGemmRunner leaves tune_max_num_tokens=None, so autotuner.py:1579-1582 hands this function the current input size, not a maximum. Remove the floor and a small first call (say M=64) warms nothing above 120 — reintroducing the same class of hole. The floor is a workaround for the runner not declaring its max.
The clean fix, if you're happy with it: plumb a real max through, exactly as the neighbouring runners already do — Fp8BlockScalingGemmRunner pins tune_max_num_tokens=4096 and MoERunner 8192 (torch_custom_ops.py:86, 233, 1957, 2552). fp8_swap_ab_gemm gets a tune_max_num_tokens arg, linear.py:1152 passes the model's real max_num_tokens, and the floor goes away. Then the count is 136 at maxnt:2048 and 264 only when the workload genuinely reaches 4096 — strictly fewer buckets than today in every case, and it removes the guessing.
That's a slightly wider change than a warmup-stride bugfix, so I didn't fold it in unasked. Happy to do it in this PR, or land the stride fix (which is what removes the 2.2s mid-inference nvcc stall and the +55.9% end-to-end) and follow up with the plumbing — your call.
One note on the diff you're looking at: the second half of the change is also a bugfix, not just a stride change. range(128, x, 16) is half-open, so the band containing max_num_tokens itself was never warmed — the M every full batch runs at. That hole was 128 wide before this change and it's why Fp8BlockScalingGemmRunner had [3969, 4096] cold.
…ocstring Addresses review feedback: NVIDIA copyright header on the new test file, and a Tuple[int, ...] return annotation plus a Google-style docstring on deep_gemm_gen_tuning_buckets (Tuple was not previously imported). Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
|
Seconding @hyukn's question with numbers — the tuning cost is very uneven across callers, and the two that pay most are not the ones this PR is about.
The two this PR targets are cheap. The other two are not. The CuTe DSL NVFP4 runners ( And most of that work is unreachable. Both set It also changes runtime behavior on that path. Suggestion: keep the dense stride for the JIT-warmup callers without imposing it on the CuTe DSL ones — a separate generator, or take the bound as a parameter so Two things that check out, so nobody re-derives them: the new test file is collected ( |
…n the DeepGEMM warmup buckets The stride-16 warmup grid applied `max(x, 4096)` unconditionally. That floor is only needed by callers that leave `tune_max_num_tokens` unset (fp8SwapABGemmRunner receives the current input size, not a maximum), and it silently overrode callers that do declare one. The two cute_dsl MoE ops declare `tune_max_num_tokens=512` and, unlike the DeepGEMM warmup runners, are multi-tactic — so every bucket is profiled per tactic. They were being handed 264 buckets, 224 of them above their own declared max. Those are unreachable: `AutoTuner._find_nearest_profile` looks up with `min(M, tune_max_num_tokens)`, so a bucket above the declared max can never be selected. Honouring the declaration takes them to 40 buckets with identical coverage of every reachable M. Callers without a declared max are byte-identical to before, so the nvbug 6550749 fix is unaffected. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
|
Thanks again for pushing on this — the concern found a real bug, not just a cost question. Your objection surfaced a regression I had introduced. The generator has four
This is perf-neutral by construction, not by measurement.
Its cost is the +2.7s I measured: +2.5% on startup, 12.2ms per bucket, and only Happy to split the cute_dsl thinning into its own PR if you'd rather review the |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py`:
- Around line 252-255: Replace the self-comparison in the test loop with
independent assertions for the expected bucket boundaries or counts returned by
deep_gemm_gen_tuning_buckets. Cover representative inputs, including at least
one value at or above 128, while retaining coverage of the legacy default
behavior where appropriate.
- Around line 203-229: Define and enforce the input contract for
tune_max_num_tokens in deep_gemm_gen_tuning_buckets: either support arbitrary
declared maxima by adding boundary tests such as 64, 127, and 513 with expected
bucket limits, or validate and reject values below 128 and not aligned to 16.
Keep ModelConfig.max_num_tokens and TuningConfig.tune_max_num_tokens consistent
with the chosen contract.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 616da2e0-fef4-4543-a50f-ea68e9bfa992
📒 Files selected for processing (4)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/utils.pytests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/utils.py
| for x in (64, 128, 512, 2048, 4096, 8192, 9000): | ||
| assert deep_gemm_gen_tuning_buckets(x) == deep_gemm_gen_tuning_buckets( | ||
| x, x_is_declared_max=False | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Assert legacy behavior independently.
deep_gemm_gen_tuning_buckets(x) and deep_gemm_gen_tuning_buckets(x, x_is_declared_max=False) use the same implementation. This comparison only detects a change to the default argument. It does not detect a regression in the default bucket sequence.
Assert stable expected boundary values or counts for representative inputs, including an input at or above 128.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py` around lines 252
- 255, Replace the self-comparison in the test loop with independent assertions
for the expected bucket boundaries or counts returned by
deep_gemm_gen_tuning_buckets. Cover representative inputs, including at least
one value at or above 128, while retaining coverage of the legacy default
behavior where appropriate.
Source: Path instructions
|
Two things on the new commit. 1. The floor now also fires below 128, which it never did on if not x_is_declared_max:
x = max(x, 4096)
x = min(x, 8192)
if x >= 128:On I suspect you want this — it is what the docstring already claims the floor does — and the in-process profiling cache ( 2. assert deep_gemm_gen_tuning_buckets(x) == deep_gemm_gen_tuning_buckets(x, x_is_declared_max=False)compares the function to itself. The one assertion with content, The rest checks out: |
…8 guard The previous commit hoisted `max(x, 4096)` above the `if x >= 128` guard while adding the declared-max flag. That silently changed a path this fix has no reason to touch: on main a current input size below 128 returns 15 buckets (max 120) because the guard short-circuits before the floor is reached, but at that commit f(64) returned 264 buckets (max 4096). The path is live. fp8SwapABGemmRunner declares no tune_max_num_tokens, so the autotuner passes it the current input M rather than a maximum, and it carries `exclude_from_cache=True`, so the sweep re-runs on every process start. The cost was never measured — the measured +2.7s figure was taken before the hoist — so restore main's ordering rather than keep an unmeasured widening. Also replaces test_default_is_unchanged_for_callers_without_a_declared_max, which could not fail: it asserted f(x) == f(x, x_is_declared_max=False), comparing the function to its own default. Its one assertion with content pinned the widened behavior, so it would have reported a restore as the regression. The replacement asserts main's literal output below 128, and a mutation test confirms it fails when the hoist is reintroduced. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com>
|
Both findings are correct, and (1) is worse than you framed it. Fixed in 1. The floor firing below 128 — confirmed, and not intended. Your numbers are exact: if x >= 128:
if not x_is_declared_max:
x = max(x, 4096)
x = min(x, 8192)You offered "make it deliberate and say so, or restore the ordering." Restoring is 2. The tautological test — confirmed, and it was worse than dead weight.
Replaced with On your first comment, the tactic-count asymmetry you quantified is what
Not splitting the thinning out, per your call. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py`:
- Around line 246-272: Register test_below_128_is_unchanged_from_main and
test_the_floor_still_fires_from_128_up in the corresponding test-list files
under both the test-db and qa directories. Preserve the existing list format and
placement conventions so both CI and manual QA pipelines discover and execute
the new tests.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 112b4524-7ad8-459f-981e-4d6170d715f8
📒 Files selected for processing (2)
tensorrt_llm/_torch/utils.pytests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/utils.py
| @pytest.mark.parametrize("x", [1, 8, 32, 64, 120, 127]) | ||
| def test_below_128_is_unchanged_from_main(x): | ||
| """A sub-128 current input size must return ONLY the low buckets. | ||
|
|
||
| Both clamps live inside the ``x >= 128`` guard, so this path is | ||
| byte-identical to main. That ordering is load-bearing and easy to lose: | ||
| hoisting ``max(x, 4096)`` above the guard turns f(64) from 15 buckets into | ||
| 264. Nothing in this fix needs that, and it is an unmeasured startup cost on | ||
| a live path -- ``fp8SwapABGemmRunner`` declares no ``tune_max_num_tokens``, | ||
| so the autotuner hands it the *current* M, and ``exclude_from_cache=True`` | ||
| means the sweep re-runs on every process start. | ||
|
|
||
| Asserted against main's literal output, not against the function's own | ||
| default: comparing ``f(x)`` to ``f(x, x_is_declared_max=False)`` compares the | ||
| function to itself and cannot fail. | ||
| """ | ||
| assert deep_gemm_gen_tuning_buckets(x) == tuple(range(8, 128, 8)) | ||
|
|
||
|
|
||
| def test_the_floor_still_fires_from_128_up(): | ||
| """Above the guard the floor must still apply -- it is not dead code. | ||
|
|
||
| Complements the test above, so the sub-128 short-circuit cannot be mistaken | ||
| for the floor having been removed. A first call at M=128 warms out to 4096. | ||
| """ | ||
| assert max(deep_gemm_gen_tuning_buckets(128)) >= 4096 | ||
| assert max(deep_gemm_gen_tuning_buckets(129)) >= 4096 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
test_file='tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py'
test_stem='test_deep_gemm_tuning_buckets'
test_functions=(
'test_below_128_is_unchanged_from_main'
'test_the_floor_still_fires_from_128_up'
)
for list_root in \
tests/integration/test_lists/test-db \
tests/integration/test_lists/qa
do
echo "== ${list_root} =="
if [[ ! -d "$list_root" ]]; then
echo "MISSING: ${list_root}"
continue
fi
rg -n -F "$test_file" "$list_root" || true
rg -n -F "$test_stem" "$list_root" || true
for test_function in "${test_functions[@]}"; do
rg -n -F "$test_function" "$list_root" || true
done
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 238
Register the new tests in both test-list directories before merge.
The new tests test_below_128_is_unchanged_from_main and test_the_floor_still_fires_from_128_up are not registered in tests/integration/test_lists/test-db/ (CI) or tests/integration/test_lists/qa/ (manual QA). Without registration, these tests will not run in the CI or QA pipelines. Add both test functions to the appropriate list files.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py` around lines 246
- 272, Register test_below_128_is_unchanged_from_main and
test_the_floor_still_fires_from_128_up in the corresponding test-list files
under both the test-db and qa directories. Preserve the existing list format and
placement conventions so both CI and manual QA pipelines discover and execute
the new tests.
Source: Path instructions
|
One small thing left, and it is latent rather than live: A declared max that is not a multiple of 16 still gets a bucket above it. Not blocking on that. From my side the change is good; I'm leaving approval until hyukn's thread is closed out, since the numbers I posted were an answer to his question and he should get to sign it off. |
|
Minimax2.5 is low priority model. Close this bug. |
3 participants
[https://nvbugs/6550749][fix] Size the DeepGEMM warmup bucket stride to the 16-token config quantum
Summary
deep_gemm_gen_tuning_bucketsexists to make DeepGEMM JIT-compile every kernelconfig a workload will need during autotuning at startup, so no live
iteration pays an nvcc stall. It walked M with
range(128, x, 128).DeepGEMM's SM100 config selection changes every 16 tokens, so a stride-128
grid sampled 1 band in 8 and left the other 7 unwarmed. Whichever live iteration
first landed in a cold band paid ~2.2 s of
nvccinside the measured window.On
minimax_m2.5_fp8-bench-pytorch-float8-maxbs:512-maxnt:2048-input_output_len:128,128-gpus:4(GB300) that was +55.9 % end-to-end, and it was filed as a perf instability
bug because whether a run pays it depends on request packing: iterations 132–139
pack the full 2048-token cap, iteration 140 gets the ragged 540-token leftover
(ctx 416 + gen 124), whose config was never warmed.
Two changes, both to the same one-line expression:
half-open
range(128, x, 16)never warmed the band containingmax_num_tokensitself — the M every full batch runs at. This is asecond, pre-existing hole (128 wide before this change); it is why
Fp8BlockScalingGemmRunner, which pinstune_max_num_tokens=4096, had theband
[3969, 4096]cold.Why stride 16 is correct, not just empirically sufficient
csrc/jit_kernels/heuristics/sm100.hpphas exactly three M-dependent sites, andevery breakpoint sits at
M % 16 == 1::62-63— the non-swap_abBLOCK_Mpick (m<=32 → 32; m<=64 → 64; else 128), stepping at M=33 and M=65.:230— the only use of M in config scoring,ceil_div(get_expected_m(), layout.block_m), stepping atM = k*block_m + 1. Every candidateblock_mis a multiple of 16 (swap_ab useslcm(16, get_block_m_multiple_of()); non-swap_ab uses{32, 64, 128}).expected_mequalsmon this path (sm100_fp8_fp4_gemm_1d1d.hpp:188,config.hpp:27).num_stagesis a closed form overblock_m/block_n/smem (:176-217), thetmem check reads
block_m/block_nonly (:123-128), andnum_smsisconstant — none adds a breakpoint.
Every band is therefore an aligned
[16k+1, 16k+16]cell, and each such cellcontains exactly one multiple of 16. A stride-16 grid provably touches every
band for any (N, K), not merely the shapes measured here. Correspondingly, a
stride of 32 must miss any band that is exactly 16 wide.
Measurements (GB300, aws-cmh)
One node, byte-identical dataset, sole variable the JIT cache state:
host_step_timeIteration 140 alone accounted for 101.4 % of the end-to-end gap (all 136 other
iterations combined were −19.4 ms). After the change a cold cache costs the same
as a warm one on that iteration, to within 0.08 ms.
End to end, cold and warm are now indistinguishable — both cold reps came in
marginally faster than the warm one, so the residual ±1.5 % is run-to-run noise
with no cache-state signal left in it. Iteration 140 does not appear in any rep's
worst-iteration list; the slowest non-startup iteration is 131 in all three, and
it is highest in the warm rep (220.5 ms vs 206.5 / 201.3), which is what noise
looks like. The warm rep compiling zero cubins while holding 37 also settles
completeness independently: the stride-16 warmup set left nothing for a later
process to discover.
The compiles moving to startup is also visible in the cache directory alone,
independent of any timing:
kernel.cubinwrittenBoth cold reps compiled an identical 37 cubins, so the warmup set is
deterministic.
A stride-1 residual sweep, one fresh cache and one process per candidate:
Cost
Tripling the bucket count costs +4 compiles, not +200: the number of distinct
configs is bounded by the
BLOCK_Mladder (~17 per shape), not by how densely Mis sampled. In this workload the DeepGEMM cache goes from 33 cubins to 37. The
extra buckets are warm GEMM calls at ~4.3 ms each, they run during autotuning at
startup, and they cannot change which kernel steady-state inference selects —
selection is a function of M either way.
Stride 8 was measured and rejected: identical compiles and identical coverage for
+120 buckets and +30 s of nvcc. Stride 16 is chosen because 16 is the band
quantum, not because a sweep happened to find no holes there.
Notes
max(x, 4096)floor is deliberately left alone, and a comment now records why it is load-bearing:fp8SwapABGemmRunnerleavestune_max_num_tokensunset, soautotuner.pypasses the current input size rather than a maximum. Dropping the floor would make a small first call (say M=64) warm nothing above 120.map_to_tuning_bucketsround-up (only rewrites the autotuner cache key, leaving raw M and the JIT untouched);set_block_size_multiple_of(128)(process-global, changes steady-state kernel choice for every model).Test coverage
deep_gemm_gen_tuning_bucketshad no test coverage, which is how a stridethis coarse survived. Adds
tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py(25 tests, no GPU required), asserting the property that matters — for every
reachable M, some bucket selects the same config — exhaustively at stride 1 over
every
BLOCK_MDeepGEMM can choose. It is auto-collected by CI, sinceunittest/_torch/miscis listed as a directory inl0_h100.yml(pre_merge) andl0_b200.yml(post_merge).The top-band hole was found by these tests, not by a benchmark: the bug's own
case is
maxnt:2048and never reaches that band.One assertion form is deliberately avoided as vacuous:
ceil_div(512, 144) == ceil_div(540, 144), so "some bucket shares M=540's band index" passes on thepre-fix buckets even though M=540 demonstrably compiled a fresh kernel — the
heuristic picks a different
BLOCK_Mat 512 than at 540. The tests assert abucket lands physically inside the 16-wide band instead.
Dev Engineer Review
deep_gemm_gen_tuning_bucketsto use 16-token strides above 128.max(x, 4096)floor for callers without a declared maximum.x_is_declared_maxsupport for callers with a maximum token count.QA Engineer Review
test_every_reachable_m_shares_a_band_with_some_buckettest_stride_matches_the_block_m_quantumtest_lower_clamp_survives_a_small_first_calltest_the_band_containing_max_num_tokens_is_warmedtest_low_buckets_are_unchangedtest_buckets_are_sorted_and_uniquetest_measured_cold_band_now_has_a_bucket_inside_ittest_declared_max_is_not_raised_to_the_floortest_declared_max_still_covers_every_reachable_mtest_declared_max_thins_the_multi_tactic_moe_casetest_below_128_is_unchanged_from_maintest_the_floor_still_fires_from_128_uptests/integration/test_lists/undertest-db/orqa/.