[None][Perf] Resolve the runtime JIT issue

tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py

@@ -1005,11 +1005,7 @@ def _chunk_kda_fwd(
         if chunk_indices is None:
             chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
         NT = len(chunk_indices)
-        # Any chunk count is launchable: the varlen scheduler rounds up
-        # (`total_cgs_per_head = ceil(num_chunks / CHUNKS_PER_BLOCK)`), launches
-        # the constant `NUM_SMS` grid, and guards every chunk loop with
-        # `chunk_idx < num_chunks`. Only eqlen needs divisibility, and it gets
-        # that from its own `CPB_BT` pad above.
+        # The varlen scheduler supports any chunk count.
         N_seqs = len(cu_seqlens) - 1
     else:
         NT = T // BT

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

@@ -50,7 +50,7 @@
 from ..modules.multi_stream_utils import do_multi_stream
 from ..modules.swiglu import silu_and_mul_kernel
 from ..utils import (ActivationType, deep_gemm_gen_tuning_buckets,
-                     fp4_scale_infer_shape,
+                     deep_gemm_jit_warmup_buckets, fp4_scale_infer_shape,
                      get_last_power_of_2_num_tokens_buckets,
                      get_power_of_2_num_tokens_buckets,
                      is_nvfp4_marlin_supported_sm, last_positive_power_of_2,
@@ -2077,8 +2077,6 @@ def _(
     return input.new_empty((M, N), dtype=output_dtype)
 
 
-# deep_gemm_gen_tuning_buckets is imported from ..utils
-
 _USE_FUSED_FP8_QUANT_PACK = os.environ.get("TRTLLM_FUSED_FP8_QUANT_PACK",
                                            "1") == "1"
 
@@ -2150,7 +2148,7 @@ class fp8SwapABGemmRunner(TunableRunner):
     # every process startup.
     tuning_config = TuningConfig(
         dynamic_tensor_specs=(DynamicTensorSpec(
-            0, 0, deep_gemm_gen_tuning_buckets), ),
+            0, 0, deep_gemm_jit_warmup_buckets), ),
         exclude_from_cache=True,
     )
 
@@ -2343,12 +2341,17 @@ def _(
     return input.new_empty((input.size(0), weight.size(0)), dtype=output_dtype)
 
 
-# The runner is used to trigger deepgemm jit during autotune.
+# The runner is used to trigger deepgemm jit during autotune. Only Hopper has
+# work to do: on SM100 this GEMM dispatches to TrtllmGenGemmRunner's prebuilt
+# cubins and compiles nothing.
 class Fp8BlockScalingGemmRunner(TunableRunner):
+    # Without exclude_from_cache, a warm disk cache short-circuits tuning and
+    # the JIT warmup never runs.
     tuning_config = TuningConfig(
         dynamic_tensor_specs=(DynamicTensorSpec(
-            0, 0, deep_gemm_gen_tuning_buckets), ),
+            0, 0, deep_gemm_jit_warmup_buckets), ),
         tune_max_num_tokens=4096,
+        exclude_from_cache=True,
     )
 
     def get_valid_tactics(

tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py

@@ -30,7 +30,6 @@
 import triton
 import triton.language as tl
 
-from ..fla.index import prepare_chunk_indices
 from . import _kda_decode
 
 try:
@@ -430,19 +429,11 @@ def can_use_indexed_prefill(
         cu_seqlens: Optional[torch.Tensor],
         num_sequences: int,
         num_tokens: int,
-        chunk_indices: Optional[torch.Tensor] = None,
-        chunk_size: int = 64,
     ) -> bool:
         """Return whether prefill can update this V-first pool directly."""
         if self.prefill_kernel_path != "optimized" or num_tokens == 0:
             return False
         if cu_seqlens is not None:
-            if chunk_indices is None:
-                from fla.ops.utils.index import prepare_chunk_indices
-
-                chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
-            if chunk_indices.shape[0] < 4:
-                return False
             if cu_seqlens.shape[0] - 1 != num_sequences:
                 return False
         if state_pool.dtype != torch.float32 or state_pool.ndim != 4:
@@ -523,20 +514,6 @@ def prefill_chunk_kda(
                 "Indexed KDA prefill requires state_indices and does not accept initial_state."
             )
         use_optimized = self.prefill_kernel_path == "optimized" and use_indexed_state
-        if use_optimized and cu_seqlens is not None:
-            if chunk_indices is None:
-                chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
-            # The persistent K123 scheduler needs at least 4 total chunks
-            # (cgs_per_head = NT // 4 cooperative groups per head). The
-            # eqlen path guarantees this by padding to a 256-token multiple
-            # inside the op; varlen has no such pad, so small varlen
-            # batches (short-prompt contexts, NT < 4) launch with a
-            # zero-size grid -> DSLCudaRuntimeError. Route them to the FLA
-            # reference path (negligible perf impact at these sizes). The
-            # check must happen before dispatch: the FLA fallback applies
-            # Q/K normalization and beta sigmoid in its own kernels.
-            if chunk_indices.shape[0] < 4:
-                use_optimized = False
 
         if use_indexed_state and not use_optimized:
             raise RuntimeError(

tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py

@@ -40,6 +40,9 @@ def _prepare_kda_chunk_indices(
 class KimiK3MambaMetadata(Mamba2Metadata):
     """Mamba metadata extended with Kimi K3 KDA preparation."""
 
+    # The KDA prefill kernel keys its compile cache on chunk alignment.
+    prefill_chunk_alignment = KDA_PREFILL_CHUNK_SIZE
+
     def __init__(self, max_batch_size: int, chunk_size: int, max_num_tokens: int) -> None:
         super().__init__(max_batch_size, chunk_size, max_num_tokens)
         # Every chunk contains at least one token, so this is a hard upper bound.

tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py

@@ -654,7 +654,6 @@ def forward_prefill(
             cu_seqlens=cu_seqlens,
             num_sequences=num_prefills,
             num_tokens=x2d.shape[0],
-            chunk_indices=chunk_indices,
         )
         recurrent_in = None
         if use_indexed_state:

tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py

@@ -15,7 +15,7 @@
 
 import contextlib
 import math
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 import triton
@@ -324,6 +324,11 @@ class Mamba2Metadata:
     # inference from a stray shell export or a forked worker.
     _warmup_force_initial_states: bool = False
 
+    # Set by a subclass whose prefill kernel compiles a separate variant
+    # depending on whether every context length is a multiple of this value.
+    # Warmup reads it to cover both variants.
+    prefill_chunk_alignment: Optional[int] = None
+
     @classmethod
     @contextlib.contextmanager
     def force_initial_states_for_warmup(cls):

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -2133,16 +2133,9 @@ def trtllm_gen_fmha_jit_warmup():
         else:
             logger.debug("Skipped TRTLLM-Gen FMHA JIT warmup for Ctx kernels")
 
-        model_type = getattr(self.model.model_config.pretrained_config,
-                             "model_type", None)
-        if can_run_general_warmup and model_type in ("kimi_k3", "kimi_linear"):
-            # Kimi's one-token context takes the NT < 4 FLA fallback and does
-            # not compile the optimized single-sequence K123 variant. A
-            # non-aligned five-chunk context enters the pure K123 path.
-            _KIMI_KDA_PREFILL_WARMUP_TOKENS = 257
-            logger.info("Adding Kimi KDA pure-prefill warmup with "
-                        f"{_KIMI_KDA_PREFILL_WARMUP_TOKENS} context tokens")
-            warmup_requests_configs.append((_KIMI_KDA_PREFILL_WARMUP_TOKENS, 0))
+        # Kimi K3 / Kimi Linear are warmed by ``_run_mamba_hybrid_warmup``, not
+        # here: they always run on a ``MambaHybridCacheManager``, so
+        # ``can_run_general_warmup`` is False and the configs above never run.
 
         if self.guided_decoder is None and can_run_general_warmup:
             # The cute_dsl_mla FMHA lib now only support the generation-only batch, we need to warmup the TRTLLM-Gen FMHA lib for the mixed context+generation batch.
@@ -2404,8 +2397,9 @@ def _run_mamba_hybrid_warmup(self,
         This method runs two extra forward passes to compile those variants
         during warmup:
 
-        1. ``least_requests=False`` — splits ``curr_max_num_tokens`` into many
-           short sequences, forcing the multi-seq path of
+        1. ``least_requests=False``, chunk-ragged — splits
+           ``curr_max_num_tokens`` into many short sequences, forcing the
+           multi-seq path of
            ``cu_seqlens_to_chunk_indices_offsets_triton`` and its
            ``_cu_seqlens_triton_kernel``.
         2. ``least_requests=False`` inside
@@ -2459,10 +2453,19 @@ def _run_mamba_hybrid_warmup(self,
         logger.info(
             "Running Mamba hybrid warmup (multi-seq + HAS_INITSTATES=True)...")
 
-        # (num_tokens, num_gen_requests, least_requests, force_initstates)
+        # A model whose prefill kernel specializes on chunk alignment declares
+        # it on its Mamba metadata class; the two passes then cover both
+        # variants, for free since alignment is independent of HAS_INITSTATES.
+        # Resolved the same way the runtime does, so warmup can't prime the
+        # alignment variant of a class the runtime never instantiates.
+        metadata_cls = resolve_mamba_metadata_cls(self.model)
+        chunk_alignment = metadata_cls.prefill_chunk_alignment
+
+        # (num_tokens, num_gen_requests, least_requests, force_initstates,
+        #  chunk_aligned)
         mamba_warmup_shapes = [
-            (capped_num_tokens, 0, False, False),
-            (capped_num_tokens, 0, False, True),
+            (capped_num_tokens, 0, False, False, False),
+            (capped_num_tokens, 0, False, True, True),
         ]
 
         autotuner_enabled = self.llm_args.enable_autotuner
@@ -2471,8 +2474,8 @@ def _run_mamba_hybrid_warmup(self,
                         if autotuner_enabled else contextlib.nullcontext())
 
         with self.no_cuda_graph(), autotune_ctx:
-            for (num_tokens_i, num_gen_requests_i, least_req_i,
-                 force_init_i) in mamba_warmup_shapes:
+            for (num_tokens_i, num_gen_requests_i, least_req_i, force_init_i,
+                 chunk_aligned_i) in mamba_warmup_shapes:
                 init_ctx = (Mamba2Metadata.force_initial_states_for_warmup()
                             if force_init_i else contextlib.nullcontext())
                 shape = (f"Mamba hybrid, num_tokens={num_tokens_i}, "
@@ -2484,7 +2487,9 @@ def _run_mamba_hybrid_warmup(self,
                             resource_manager,
                             num_tokens_i,
                             num_gen_requests_i,
-                            least_requests=least_req_i)
+                            least_requests=least_req_i,
+                            chunk_alignment=chunk_alignment,
+                            chunk_aligned=chunk_aligned_i)
                     except torch.OutOfMemoryError as e:
                         if self._is_distributed_forward():
                             raise
@@ -3171,13 +3176,42 @@ def _release_batch_context(self, batch: Optional[ScheduledRequests],
                     if spec_resource_manager is not None:
                         spec_resource_manager.free_resources(req)
 
+    @staticmethod
+    def _apply_chunk_alignment(ctx_token_nums: List[int], alignment: int,
+                               aligned: bool) -> Optional[List[int]]:
+        """Rewrite warmup context lengths onto one side of ``alignment``.
+
+        Returns lengths that are all multiples of ``alignment`` (``aligned``) or
+        that include at least one non-multiple, or None if that would leave the
+        batch empty. Only shrinks or drops sequences.
+        """
+        if aligned:
+            floored = [n - n % alignment for n in ctx_token_nums]
+            kept = [n for n in floored if n > 0]
+            return kept or None
+        if any(n % alignment != 0 for n in ctx_token_nums):
+            return list(ctx_token_nums)
+        # Every length is a multiple; one token off the last breaks it.
+        if ctx_token_nums[-1] <= 1:
+            return None
+        return ctx_token_nums[:-1] + [ctx_token_nums[-1] - 1]
+
     def _create_warmup_request(
             self,
             resource_manager: ResourceManager,
             num_tokens: int,
             num_gen_requests: int,
-            least_requests: bool = True) -> Optional[ScheduledRequests]:
-        """Creates a generic dummy ScheduledRequests object for warmup."""
+            least_requests: bool = True,
+            chunk_alignment: Optional[int] = None,
+            chunk_aligned: bool = False) -> Optional[ScheduledRequests]:
+        """Creates a generic dummy ScheduledRequests object for warmup.
+
+        ``chunk_alignment``, when set, rewrites the context lengths onto one
+        side of it, so a kernel specializing on chunk alignment can be primed
+        for both variants instead of whichever one the split happens to hit.
+        Only shrinks or drops sequences, so the block estimate below stays a
+        safe over-estimate.
+        """
         kv_cache_manager = resource_manager.get_resource_manager(
             self.kv_cache_manager_key)
         draft_kv_cache_manager = self._get_draft_kv_cache_manager(
@@ -3268,6 +3302,20 @@ def blocks_for_seq(num_tokens: int) -> int:
             if num_left_over_tokens > 0:
                 ctx_token_nums.append(num_left_over_tokens)
 
+            if chunk_alignment:
+                adjusted = self._apply_chunk_alignment(ctx_token_nums,
+                                                       chunk_alignment,
+                                                       chunk_aligned)
+                if adjusted is None:
+                    logger.debug(
+                        f"Warmup batch of {ctx_token_nums} cannot be made "
+                        f"{'aligned' if chunk_aligned else 'ragged'} modulo "
+                        f"{chunk_alignment}; that variant will compile on the "
+                        f"first request needing it.")
+                else:
+                    ctx_token_nums = adjusted
+                    num_ctx_requests = len(ctx_token_nums)
+
             ctx_requests = kv_cache_manager.add_dummy_requests(
                 list(range(num_ctx_requests)),
                 token_nums=ctx_token_nums,

tensorrt_llm/_torch/utils.py

@@ -428,6 +428,38 @@ def deep_gemm_gen_tuning_buckets(x: int):
     return buckets
 
 
+def deep_gemm_jit_warmup_buckets(max_m: int):
+    """M grid for the DeepGemm runners that exist only to drive JIT warmup.
+
+    DeepGemm picks its tile layout from a heuristic over M and compiles one
+    kernel per selected layout. A layout no bucket selects gets compiled
+    mid-inference instead, and DeepGemm forks nvcc while holding the GIL, so
+    that compile stalls every rank of the attention-DP group.
+
+    Step 16 is exactly the right spacing, and it is needed over the whole
+    range. In ``deepgemm/csrc/jit_kernels/heuristics/sm100.hpp`` the selection
+    depends on M only through ``ceil_div(m, block_m)``, and every candidate
+    ``block_m`` is a multiple of 16, so the choice is constant on each window
+    ``[16j + 1, 16j + 16]``: one sample per window misses nothing, and anything
+    coarser skips whole windows.
+
+    A coarse high-M band is *not* safe -- ``compare`` tie-breaks on
+    ``last_wave_util = num_blocks % num_sms``, which keeps oscillating. At
+    148 SMs and ``n=128, k=512``, ``M in [2305, 2368]`` selects a layout of its
+    own (``block_m=16``: one wave, best last-wave utilization) that a step-128
+    grid steps over, sampling 2304 and 2432.
+    """
+    # A worker whose M never leaves the low band -- a disagg GEN worker runs at
+    # batch x MTP tokens -- must not be pulled up to the 4096 floor. Measured
+    # cost of doing so: +283 s of GEN autotune, +1096 tuning-cache entries.
+    if max_m < 128:
+        return tuple(range(8, 128, 8))
+    max_m = max(min(max_m, 8192), 4096)
+    low = range(8, 128, 8)
+    dense = range(128, max_m, 16)
+    return tuple(low) + tuple(dense) + (max_m, )
+
+
 def fp4_scale_infer_shape(input_shapes: List[List[int]]) -> int:
     """Calculate the swizzled scale size for a packed FP4 input tensor."""
     unpacked_shape = list(input_shapes[0])

tests/unittest/_torch/executor/test_distributed_warmup_oom.py

@@ -434,6 +434,10 @@ def _mamba_engine(*, world_size: int = 1, dwdp_size: int = 0) -> tuple[PyTorchMo
     engine.original_max_draft_len = 0
     engine.llm_args = SimpleNamespace(enable_autotuner=False)
     engine.no_cuda_graph = contextlib.nullcontext
+    # The warmup resolves its chunk-alignment variant off the model's Mamba
+    # metadata class; a model declaring none takes the ``Mamba2Metadata``
+    # default, i.e. no alignment rewrite.
+    engine.model = SimpleNamespace()
     batch = object()
     engine._release_batch_context = lambda *_a, **_kw: _released_batch(batch)
 

tests/unittest/_torch/misc/test_autotuner.py

@@ -37,7 +37,8 @@
                                            StaticDim, TunableRunner,
                                            TuningConfig, autotune)
 from tensorrt_llm._torch.distributed import Distributed
-from tensorrt_llm._torch.utils import (get_power_of_2_num_tokens_buckets,
+from tensorrt_llm._torch.utils import (deep_gemm_jit_warmup_buckets,
+                                       get_power_of_2_num_tokens_buckets,
                                        next_positive_power_of_2)
 from tensorrt_llm.bindings.internal.runtime import delay_kernel
 from tensorrt_llm.logger import logger
@@ -1867,3 +1868,29 @@ def test_visual_gen_autotune_dist_world_allgather(mpi_pool_executor):
                               *zip(*[(world_size, )] * world_size)))
     for got in results:
         assert got == ["rank0", "rank1"]
+
+
+def test_deep_gemm_jit_warmup_buckets():
+    # A bucket left out leaves a DeepGemm layout uncompiled, so it compiles
+    # under traffic while holding the GIL. The layout is constant on each
+    # 16-aligned window of M, so every window must hold a bucket.
+    for max_m in (4096, 6144, 8192):
+        b = deep_gemm_jit_warmup_buckets(max_m)
+        assert list(b) == sorted(set(b)), "must be sorted and duplicate-free"
+        assert b[-1] == max_m, "the bound itself must be a bucket"
+
+        assert all(hi - lo <= 16 for lo, hi in zip(b, b[1:])), \
+            "no gap may exceed 16: it would skip a 16-aligned window of M"
+        windows = {(m - 1) // 16 for m in b}
+        missing = [j for j in range((max_m - 1) // 16 + 1) if j not in windows]
+        assert not missing, \
+            f"windows with no bucket (M from {[16 * j + 1 for j in missing]})"
+
+    # Clamping keeps a small bound from skipping warmup for larger M,
+    # and a huge one from exploding the bucket count.
+    assert deep_gemm_jit_warmup_buckets(256) == deep_gemm_jit_warmup_buckets(
+        4096)
+    # But a decode-only worker (M = batch x MTP) stays in the low band.
+    assert max(deep_gemm_jit_warmup_buckets(32)) < 128
+    assert deep_gemm_jit_warmup_buckets(16384) == deep_gemm_jit_warmup_buckets(
+        8192)