From 570c57e3695ba70f34b680ed3d3a5f6e20e4ef66 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:57:34 -0700 Subject: [PATCH 1/4] [https://nvbugs/6550749][fix] Size DeepGEMM warmup buckets to the 16-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 #1 3978.59 ms 114.16 ms 37 after, cold cache #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> --- tensorrt_llm/_torch/utils.py | 55 +++++- .../misc/test_deep_gemm_tuning_buckets.py | 174 ++++++++++++++++++ 2 files changed, 228 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 56051035fd34..2c6603ca3ddb 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -323,13 +323,66 @@ def get_last_power_of_2_num_tokens_buckets(max_num_tokens) -> List[int]: return tuple(num_token_buckets[::-1]) +# DeepGEMM selects (BLOCK_M, BLOCK_N, num_stages, swap_ab) as a pure function +# of M, and every breakpoint of that function sits at M % 16 == 1. Two sites in +# csrc/jit_kernels/heuristics/sm100.hpp are the *entire* M dependence: +# * :62-63 the non-swap_ab BLOCK_M pick, which steps at M = 33 and M = 65; +# * :230 ceil_div(expected_m, block_m) in the config score, which steps 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}). +# num_stages is a closed form over block_m/block_n only, and num_sms is +# constant, so neither adds a breakpoint. Each band is therefore an aligned +# union of [16k+1, 16k+16] cells, and every such cell contains exactly one +# multiple of 16 -- so a stride-16 grid provably touches every band, for any +# (N, K), not merely the shapes measured below. +# +# A stride of 128 leaves ~1 band in 8 unwarmed. Those bands are then JIT'd by +# whichever live iteration first lands in one, putting ~2.2s of nvcc *inside* +# the measured window. Sizing the stride to the 16-token quantum is what makes +# the warmup complete; the numbers here are measured, not assumed: +# +# stride buckets nvcc compiles nvcc time bands left unwarmed +# 128 46 28 64s 7 <-- nvbug 6550749 +# 32 76 33 88s 1 (misses M=385, a band +# exactly 16 wide) +# 16 136 34 81s 0 COMPLETE +# +# (Those counts come from a standalone sweep bounded at M <= 2048; the clamp +# below makes the production list larger -- 264 buckets -- without adding +# compiles, since it adds no new BLOCK_M rungs.) +# +# Tripling the bucket count costs +4 compiles, not +200: the number of distinct +# configs is bounded by the BLOCK_M ladder (~17 per shape), not by how densely +# M is sampled. End to end on the nvbug 6550749 case the DeepGEMM cache goes +# from 33 cubins to 37. The extra buckets are warm GEMM calls at ~4.3ms each, +# they run during autotuning at startup, and they cannot change which kernel +# steady-state inference picks -- selection is a function of M either way. +# See tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py. +DEEP_GEMM_BLOCK_M_QUANTUM = 16 + + def deep_gemm_gen_tuning_buckets(x: int): buckets = tuple(range(8, 128, 8)) # Clamp x to be between 4096 and 8192. + # + # The lower clamp is load-bearing and must not be "tightened" to the real + # max_num_tokens: fp8SwapABGemmRunner leaves tune_max_num_tokens unset, so + # autotuner.py passes the *current input size* here rather than a maximum + # (see autotuner.py, `Use the current input size as the opt value`). Drop + # the floor and a small first call (say M=64) would warm nothing above 120 + # and every larger M would JIT mid-iteration. if x >= 128: x = min(x, 8192) x = max(x, 4096) - buckets += tuple(range(128, x, 128)) + # Round the top up to a whole quantum, and include it. The only multiple + # of 16 inside a band [16k+1, 16k+16] is its *top*, so a bucket >= x is + # required to warm the band that x itself lives in -- and a half-open + # range(128, x, 16) provides none. That left the band containing + # 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)) return buckets diff --git a/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py new file mode 100644 index 000000000000..5581947893ce --- /dev/null +++ b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py @@ -0,0 +1,174 @@ +"""Coverage tests for deep_gemm_gen_tuning_buckets (nvbug 6550749). + +The bucket list exists for exactly one purpose: to make DeepGEMM JIT-compile +every kernel config it will need *during autotuning at startup*, so that no +live iteration ever pays an ~2.2s nvcc stall. A bucket list with holes silently +defeats that -- the workload still runs, just with a multi-second spike on +whichever iteration first lands in an unwarmed band. On +minimax_m2.5_fp8 128,128 gpus:4 (GB300) that spike was 2316.6ms on a single +iteration, 101.4% of a 55.9% end-to-end regression, and it presented as +*instability* because whether a run pays it depends on request packing. + +So the property under test is coverage, not a literal bucket list: for every M +the runtime can produce, some bucket must select the same DeepGEMM config. +Config selection steps only at M % 16 == 1 (see the derivation next to +DEEP_GEMM_BLOCK_M_QUANTUM in tensorrt_llm/_torch/utils.py), which makes that +property checkable in pure Python with no GPU. +""" + +import pytest + +from tensorrt_llm._torch.utils import DEEP_GEMM_BLOCK_M_QUANTUM, deep_gemm_gen_tuning_buckets + + +# Mirrors csrc/jit_kernels/heuristics/sm100.hpp:62-63 and :230: the band index +# of M under a given BLOCK_M. Two M values sharing a band select byte-identical +# kernels, so warming either one warms both. +def _band(m: int, block_m: int) -> int: + return -(-m // block_m) # ceil_div + + +# Every BLOCK_M DeepGEMM can choose on SM100. swap_ab walks +# lcm(16, block_m_multiple_of=1)..256 step 16 (256 is dropped by the tmem check +# at :123-128, but including it here only makes the test stricter); non-swap_ab +# picks from {32, 64, 128}. +_BLOCK_M_CANDIDATES = tuple(range(16, 257, 16)) + (32, 64, 128) + + +@pytest.mark.parametrize("max_num_tokens", [128, 512, 2048, 4096, 8192]) +def test_every_reachable_m_shares_a_band_with_some_bucket(max_num_tokens): + """The coverage property, checked exhaustively at stride 1. + + For each candidate BLOCK_M, every reachable M must land in a band that at + least one bucket also lands in. This is what "the warmup is complete" + means; it is not a statement about bucket count or spacing. + """ + buckets = deep_gemm_gen_tuning_buckets(max_num_tokens) + + for block_m in _BLOCK_M_CANDIDATES: + warmed = {_band(b, block_m) for b in buckets} + holes = [m for m in range(1, max_num_tokens + 1) if _band(m, block_m) not in warmed] + assert not holes, ( + f"max_num_tokens={max_num_tokens} BLOCK_M={block_m}: " + f"{len(holes)} M values fall in bands no bucket warms, e.g. " + f"{holes[:8]}. Each is an ~2.2s in-window nvcc stall (nvbug " + f"6550749). If the stride grew past {DEEP_GEMM_BLOCK_M_QUANTUM}, " + f"that is the cause." + ) + + +def test_stride_matches_the_block_m_quantum(): + """Guard the specific constant, so a future widening is a deliberate act. + + A stride of 128 (the pre-fix value) leaves ~1 band in 8 cold. Because the + bands are 16 wide, only a stride that divides 16 can be complete. + """ + buckets = deep_gemm_gen_tuning_buckets(4096) + high = [b for b in buckets if b >= 128] + strides = {b - a for a, b in zip(high, high[1:])} + + assert strides == {DEEP_GEMM_BLOCK_M_QUANTUM}, ( + f"expected a uniform stride of {DEEP_GEMM_BLOCK_M_QUANTUM} above 128, got {sorted(strides)}" + ) + assert 16 % DEEP_GEMM_BLOCK_M_QUANTUM == 0, ( + f"DeepGEMM config bands are 16 wide; a stride of " + f"{DEEP_GEMM_BLOCK_M_QUANTUM} cannot sample every band" + ) + + +def test_lower_clamp_survives_a_small_first_call(): + """The max(x, 4096) floor must stay -- it is not dead code. + + fp8SwapABGemmRunner leaves tune_max_num_tokens unset, so the autotuner + passes the *current input size* here, not a maximum. Without the floor a + small first call would warm nothing above 120 and every larger M would JIT + mid-iteration -- the same bug in a different disguise. + """ + for first_call_m in (1, 8, 64, 129, 540, 2048): + buckets = deep_gemm_gen_tuning_buckets(first_call_m) + if first_call_m >= 128: + assert max(buckets) >= 4096, ( + f"first call M={first_call_m} warmed only up to " + f"{max(buckets)}; the lower clamp is gone" + ) + + +@pytest.mark.parametrize("max_num_tokens", [4096, 5000, 5001, 6144, 8000, 8191, 8192, 9000]) +def test_the_band_containing_max_num_tokens_is_warmed(max_num_tokens): + """The top band must be covered -- it is the most-visited M, not an edge. + + Every full batch runs at M == max_num_tokens, so of all the bands this one + is the likeliest to be hit. A half-open range(128, x, 16) covers none of it: + the only multiple of 16 in a band [16k+1, 16k+16] is its *top*, so the list + must reach a bucket >= x. Pre-fix this same hole was 128 wide. + """ + buckets = deep_gemm_gen_tuning_buckets(max_num_tokens) + effective = min(max(max_num_tokens, 4096), 8192) + band_start = ((effective - 1) // DEEP_GEMM_BLOCK_M_QUANTUM) * DEEP_GEMM_BLOCK_M_QUANTUM + 1 + + assert any(band_start <= b < band_start + DEEP_GEMM_BLOCK_M_QUANTUM for b in buckets), ( + f"max_num_tokens={max_num_tokens} (effective {effective}) " + f"lives in band [{band_start}, " + f"{band_start + DEEP_GEMM_BLOCK_M_QUANTUM - 1}] but the " + f"highest bucket is {max(buckets)}: a full batch JITs " + f"mid-iteration" + ) + + +def test_low_buckets_are_unchanged(): + """M < 128 was already covered; keep it byte-identical.""" + buckets = deep_gemm_gen_tuning_buckets(4096) + assert [b for b in buckets if b < 128] == list(range(8, 128, 8)) + + +def test_buckets_are_sorted_and_unique(): + """The autotuner de-dupes into a set, but duplicates would waste profiling + iterations and mask a generator bug.""" + for max_num_tokens in (128, 512, 2048, 4096, 8192): + buckets = deep_gemm_gen_tuning_buckets(max_num_tokens) + assert list(buckets) == sorted(buckets) + assert len(set(buckets)) == len(buckets) + + +# The M values measured compiling on a *warmed* GB300 cache before the fix, +# with the BLOCK_M the heuristic actually chose for each (job 2815335, shapes +# N=8192,K=3072 and N=3072,K=6144). M=540 is the one that landed inside the +# measured window and caused the 2316.6ms spike on iter 140; the other seven +# are the full residual set the stride-1 sweep found. +# +# These are pinned as *band starts*, not as bucket membership: what makes a +# band cold is that no bucket falls anywhere inside it, and each band here +# begins at its listed M and runs 16 wide. +_MEASURED_COLD_BANDS = ( + (540, 144), + (129, 80), + (161, 96), + (193, 112), + (257, 144), + (385, 128), + (401, 80), + (1729, 128), +) + + +@pytest.mark.parametrize("m,block_m", _MEASURED_COLD_BANDS) +def test_measured_cold_band_now_has_a_bucket_inside_it(m, block_m): + """Each band measured cold on GB300 must now contain a bucket. + + Note the assertion is "a bucket lands in [16k+1, 16k+16] around M", not + "some bucket shares ceil_div(M, block_m)". The weaker form is vacuous: at + BLOCK_M=144, ceil_div(512, 144) == ceil_div(540, 144) == 4, so the old + stride-128 buckets appear to cover M=540 -- yet 540 demonstrably compiled a + fresh kernel, because the heuristic picks a *different* BLOCK_M at 512 than + at 540. Only a bucket physically inside the band warms it. + """ + del block_m # documents the measured selection; not needed by the check + band_start = ((m - 1) // 16) * 16 + 1 + band = range(band_start, band_start + 16) + buckets = deep_gemm_gen_tuning_buckets(2048) + + assert any(b in band for b in buckets), ( + f"no bucket inside [{band.start}, {band.stop - 1}], the band " + f"containing M={m}: that band JIT-compiles on a live iteration " + f"(nvbug 6550749)" + ) From 247dcaf3c78e40502c636b31ebdfc279100a9cdc Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:46:14 -0700 Subject: [PATCH 2/4] [https://nvbugs/6550749][fix] Add copyright header, return type and docstring 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> --- tensorrt_llm/_torch/utils.py | 16 ++++++++++++++-- .../_torch/misc/test_deep_gemm_tuning_buckets.py | 14 ++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 2c6603ca3ddb..ef6729d48519 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -4,7 +4,7 @@ import threading from dataclasses import dataclass from enum import Enum, IntEnum -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple import torch from torch.nn import functional as F @@ -362,7 +362,19 @@ def get_last_power_of_2_num_tokens_buckets(max_num_tokens) -> List[int]: DEEP_GEMM_BLOCK_M_QUANTUM = 16 -def deep_gemm_gen_tuning_buckets(x: int): +def deep_gemm_gen_tuning_buckets(x: int) -> Tuple[int, ...]: + """Generate the M values to autotune, so DeepGEMM JIT-compiles every config. + + Args: + x: Upper bound on M. Note this is the *current input size* rather than a + maximum for callers that leave ``tune_max_num_tokens`` unset; see the + lower-clamp comment below. + + Returns: + Ascending M values to profile. Stride is + ``DEEP_GEMM_BLOCK_M_QUANTUM`` above 128, which touches every distinct + DeepGEMM config for any (N, K), and the top of the range is included. + """ buckets = tuple(range(8, 128, 8)) # Clamp x to be between 4096 and 8192. # diff --git a/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py index 5581947893ce..58d1e7ef96ee 100644 --- a/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py +++ b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py @@ -1,3 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Coverage tests for deep_gemm_gen_tuning_buckets (nvbug 6550749). The bucket list exists for exactly one purpose: to make DeepGEMM JIT-compile From 4be5ddd5fc6fedc9e13292cd62117301060251b7 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 4 Aug 2026 00:31:33 -0700 Subject: [PATCH 3/4] [https://nvbugs/6550749][fix] Honour a declared tune_max_num_tokens in the DeepGEMM warmup buckets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 8 ++- .../_torch/custom_ops/torch_custom_ops.py | 5 +- tensorrt_llm/_torch/utils.py | 42 +++++++---- .../misc/test_deep_gemm_tuning_buckets.py | 70 +++++++++++++++++++ 4 files changed, 108 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index e53d58693a4b..23a3aaef07c3 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -4707,7 +4707,9 @@ def get_tuning_config(self) -> TuningConfig: if key not in self.tuning_config_cache: self.tuning_config_cache[key] = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( - 0, 0, deep_gemm_gen_tuning_buckets), ), + 0, 0, + functools.partial(deep_gemm_gen_tuning_buckets, + x_is_declared_max=True)), ), constraint_specs=(ConstraintSpec(2, 0, fp4_scale_infer_shape), ), use_cold_l2_cache=True, @@ -5105,7 +5107,9 @@ def get_tuning_config(self) -> TuningConfig: if key not in self.tuning_config_cache: self.tuning_config_cache[key] = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( - 0, 0, deep_gemm_gen_tuning_buckets), ), + 0, 0, + functools.partial(deep_gemm_gen_tuning_buckets, + x_is_declared_max=True)), ), constraint_specs=( ConstraintSpec(2, 0, fp4_scale_infer_shape), ConstraintSpec(4, 0, lambda shapes: shapes[0][0]), diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index f3b0dc0476d3..903b3111cfda 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -17,7 +17,7 @@ import os import threading from dataclasses import replace -from functools import lru_cache +from functools import lru_cache, partial from typing import ClassVar, List, Mapping, Optional, Tuple, Union import torch @@ -1957,7 +1957,8 @@ def _( class Fp8BlockScalingGemmRunner(TunableRunner): tuning_config = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( - 0, 0, deep_gemm_gen_tuning_buckets), ), + 0, 0, partial(deep_gemm_gen_tuning_buckets, + x_is_declared_max=True)), ), tune_max_num_tokens=4096, ) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index ef6729d48519..68fa6b6dd710 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -362,13 +362,21 @@ def get_last_power_of_2_num_tokens_buckets(max_num_tokens) -> List[int]: DEEP_GEMM_BLOCK_M_QUANTUM = 16 -def deep_gemm_gen_tuning_buckets(x: int) -> Tuple[int, ...]: +def deep_gemm_gen_tuning_buckets(x: int, + x_is_declared_max: bool = False + ) -> Tuple[int, ...]: """Generate the M values to autotune, so DeepGEMM JIT-compiles every config. Args: - x: Upper bound on M. Note this is the *current input size* rather than a - maximum for callers that leave ``tune_max_num_tokens`` unset; see the - lower-clamp comment below. + x: Upper bound on M. + x_is_declared_max: True when the caller sets ``tune_max_num_tokens``, so + ``x`` is that declared maximum rather than the current input size. + Buckets above it are then dropped instead of profiled: M is clamped + with ``min(M, tune_max_num_tokens)`` for the runtime lookup (see + ``AutoTuner._find_nearest_profile``), so a higher bucket can never be + a reachable cache key. Pass it via ``functools.partial`` on the spec + rather than repeating the numeric limit, which would desync from a + ``tune_max_num_tokens`` that is reassigned per call. Returns: Ascending M values to profile. Stride is @@ -376,17 +384,25 @@ def deep_gemm_gen_tuning_buckets(x: int) -> Tuple[int, ...]: DeepGEMM config for any (N, K), and the top of the range is included. """ buckets = tuple(range(8, 128, 8)) - # Clamp x to be between 4096 and 8192. + # Clamp x to be at most 8192, and -- only when we are guessing -- at least + # 4096. # - # The lower clamp is load-bearing and must not be "tightened" to the real - # max_num_tokens: fp8SwapABGemmRunner leaves tune_max_num_tokens unset, so - # autotuner.py passes the *current input size* here rather than a maximum - # (see autotuner.py, `Use the current input size as the opt value`). Drop - # the floor and a small first call (say M=64) would warm nothing above 120 - # and every larger M would JIT mid-iteration. - if x >= 128: - x = min(x, 8192) + # The lower clamp is load-bearing for callers that leave + # tune_max_num_tokens unset (fp8SwapABGemmRunner): autotuner.py then passes + # the *current input size* here rather than a maximum (see autotuner.py, + # `Use the current input size as the opt value`). Drop the floor there and a + # small first call (say M=64) would warm nothing above 120, so every larger + # M would JIT mid-iteration. + # + # A caller that *does* declare a maximum needs no such guess, and profiling + # past it is pure waste -- the runtime lookup is clamped to the declared max, + # so those buckets are unreachable. Honouring it keeps the two cute_dsl MoE + # ops (which declare 512, and unlike the DeepGEMM warmup runners are + # multi-tactic) at 40 buckets instead of 264. + if not x_is_declared_max: x = max(x, 4096) + x = min(x, 8192) + if x >= 128: # Round the top up to a whole quantum, and include it. The only multiple # of 16 inside a band [16k+1, 16k+16] is its *top*, so a bucket >= x is # required to warm the band that x itself lives in -- and a half-open diff --git a/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py index 58d1e7ef96ee..38f53d24a005 100644 --- a/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py +++ b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py @@ -186,3 +186,73 @@ def test_measured_cold_band_now_has_a_bucket_inside_it(m, block_m): f"containing M={m}: that band JIT-compiles on a live iteration " f"(nvbug 6550749)" ) + + +# --------------------------------------------------------------------------- +# x_is_declared_max: honour a caller's tune_max_num_tokens instead of the floor. +# +# The 4096 floor exists only because fp8SwapABGemmRunner declares no +# tune_max_num_tokens, so it receives the *current input size*. A caller that +# does declare one gets its buckets clamped there instead: AutoTuner +# ._find_nearest_profile looks up with min(M, tune_max_num_tokens), so a bucket +# above the declared max can never be selected. Profiling it is pure cost, and +# for the two cute_dsl MoE ops (declared max 512, and multi-tactic, so every +# bucket is profiled per tactic) that was 224 of 264 buckets. +# --------------------------------------------------------------------------- + +_DECLARED_MAXES = (128, 256, 512, 1024, 2048, 4096, 8192) + + +@pytest.mark.parametrize("declared_max", _DECLARED_MAXES) +def test_declared_max_is_not_raised_to_the_floor(declared_max): + """No bucket may exceed a declared max (it would be unreachable).""" + buckets = deep_gemm_gen_tuning_buckets(declared_max, x_is_declared_max=True) + + over = [b for b in buckets if b > declared_max] + assert not over, ( + f"buckets {over} exceed the declared max {declared_max}; the runtime " + f"lookup clamps M to it, so these are profiled but never selected" + ) + + +@pytest.mark.parametrize("declared_max", _DECLARED_MAXES) +def test_declared_max_still_covers_every_reachable_m(declared_max): + """Thinning must not reintroduce a hole below the declared max.""" + buckets = deep_gemm_gen_tuning_buckets(declared_max, x_is_declared_max=True) + + for m in range(128, declared_max + 1): + band_start = ((m - 1) // 16) * 16 + 1 + band = range(band_start, band_start + 16) + assert any(b in band for b in buckets), ( + f"declared_max={declared_max}: M={m} is reachable but its band " + f"[{band.start}, {band.stop - 1}] holds no bucket" + ) + + +def test_declared_max_thins_the_multi_tactic_moe_case(): + """The cute_dsl MoE ops declare 512; that must cost far less than the floor. + + Pins the reduction hyukn asked for on PR #17242 so a later edit cannot + silently restore the floor for callers that declare a maximum. + """ + honoured = deep_gemm_gen_tuning_buckets(512, x_is_declared_max=True) + floored = deep_gemm_gen_tuning_buckets(512) + + assert len(honoured) == 40, len(honoured) + assert len(floored) == 264, len(floored) + assert max(honoured) == 512 + + +def test_default_is_unchanged_for_callers_without_a_declared_max(): + """fp8SwapABGemmRunner's path must be byte-identical to the floor form. + + It receives the current input size, not a maximum, so dropping the floor + there would warm nothing above the first call's M. + """ + 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 + ) + + # A small first call still warms the full floor range. + assert max(deep_gemm_gen_tuning_buckets(64)) == 4096 From 7f90d2ef7dda64f39a406146ce14b3c5673c10a3 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 4 Aug 2026 03:05:55 -0700 Subject: [PATCH 4/4] [https://nvbugs/6550749][fix] Keep the bucket clamps inside the x>=128 guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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> --- tensorrt_llm/_torch/utils.py | 12 +++++-- .../misc/test_deep_gemm_tuning_buckets.py | 36 +++++++++++++------ 2 files changed, 34 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 68fa6b6dd710..de45ae641cad 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -399,10 +399,16 @@ def deep_gemm_gen_tuning_buckets(x: int, # so those buckets are unreachable. Honouring it keeps the two cute_dsl MoE # ops (which declare 512, and unlike the DeepGEMM warmup runners are # multi-tactic) at 40 buckets instead of 264. - if not x_is_declared_max: - x = max(x, 4096) - x = min(x, 8192) + # + # Both clamps stay INSIDE the `x >= 128` guard, as on main: a caller whose + # current input size is below 128 keeps returning just the low buckets. See + # test_below_128_is_unchanged_from_main -- hoisting the floor above the guard + # silently turned f(64) from 15 buckets into 264, an unmeasured startup cost + # on a path this fix has no reason to touch. if x >= 128: + if not x_is_declared_max: + x = max(x, 4096) + x = min(x, 8192) # Round the top up to a whole quantum, and include it. The only multiple # of 16 inside a band [16k+1, 16k+16] is its *top*, so a bucket >= x is # required to warm the band that x itself lives in -- and a half-open diff --git a/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py index 38f53d24a005..df8186d54e6f 100644 --- a/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py +++ b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py @@ -243,16 +243,30 @@ def test_declared_max_thins_the_multi_tactic_moe_case(): assert max(honoured) == 512 -def test_default_is_unchanged_for_callers_without_a_declared_max(): - """fp8SwapABGemmRunner's path must be byte-identical to the floor form. - - It receives the current input size, not a maximum, so dropping the floor - there would warm nothing above the first call's M. +@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. """ - 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 - ) + assert deep_gemm_gen_tuning_buckets(x) == tuple(range(8, 128, 8)) + - # A small first call still warms the full floor range. - assert max(deep_gemm_gen_tuning_buckets(64)) == 4096 +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