[https://nvbugs/6550749][fix] Size the DeepGEMM warmup bucket stride to the 16-token config quantum

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]),

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,
     )
 

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