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 56051035fd34..de45ae641cad 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 @@ -323,13 +323,100 @@ def get_last_power_of_2_num_tokens_buckets(max_num_tokens) -> List[int]: return tuple(num_token_buckets[::-1]) -def deep_gemm_gen_tuning_buckets(x: int): +# 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, + 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. + 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 + ``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. + # Clamp x to be at most 8192, and -- only when we are guessing -- at least + # 4096. + # + # 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. + # + # 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) - 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..df8186d54e6f --- /dev/null +++ b/tests/unittest/_torch/misc/test_deep_gemm_tuning_buckets.py @@ -0,0 +1,272 @@ +# 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 +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)" + ) + + +# --------------------------------------------------------------------------- +# 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 + + +@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