[None][chore] Core: MoE backend enum, deepgemm JIT warmup, postproc metrics, cleanup

cpp/tensorrt_llm/thop/fp4Gemm.cpp

@@ -27,7 +27,6 @@
 #endif
 
 #include <ATen/cuda/EmptyTensor.h>
-#include <ATen/native/cuda/Resize.h>
 
 #include <cstddef>
 #include <cuda_fp16.h>

cpp/tensorrt_llm/thop/fp4GemmTrtllmGen.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,7 +19,6 @@
 #include "tensorrt_llm/thop/thUtils.h"
 
 #include <ATen/cuda/EmptyTensor.h>
-#include <ATen/native/cuda/Resize.h>
 
 #include <cuda_fp16.h>
 

cpp/tensorrt_llm/thop/fp4xFp8GemmTrtllmGen.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,7 +19,6 @@
 #include "tensorrt_llm/thop/thUtils.h"
 
 #include <ATen/cuda/EmptyTensor.h>
-#include <ATen/native/cuda/Resize.h>
 
 #include <cuda_fp16.h>
 

cpp/tensorrt_llm/thop/fp8BatchedGemmTrtllmGen.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,7 +21,6 @@
 #include "tensorrt_llm/thop/thUtils.h"
 
 #include <ATen/cuda/EmptyTensor.h>
-#include <ATen/native/cuda/Resize.h>
 
 #include <cuda_fp16.h>
 

cpp/tensorrt_llm/thop/fp8PerTensorScalingTrtllmGenGemm.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,7 +19,6 @@
 #include "tensorrt_llm/thop/thUtils.h"
 
 #include <ATen/cuda/EmptyTensor.h>
-#include <ATen/native/cuda/Resize.h>
 
 #include <cuda_fp16.h>
 

cpp/tensorrt_llm/thop/fp8RowwiseGemm.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,7 +22,6 @@
 #include "tensorrt_llm/thop/userbuffersTensor.h"
 
 #include <ATen/cuda/EmptyTensor.h>
-#include <ATen/native/cuda/Resize.h>
 
 #include <cstddef>
 #include <cuda_fp16.h>

tensorrt_llm/__init__.py

@@ -41,6 +41,29 @@
 # ImportError: libc10.so: cannot open shared object file: No such file or directory
 import torch  # noqa
 
+
+def _setup_cutlass_dsl_compatibility():
+    """Expose legacy CuTe APIs required by TensorRT-LLM and its dependencies."""
+    try:
+        import cutlass.cute as cute
+    except ImportError:
+        return
+
+    # The pinned CUTLASS DSL exposes these types at cute.*, while QuACK and
+    # Transformer Engine still resolve their annotations from cute.core.
+    # Keep this list explicit: copying the full namespace also replaces
+    # cute.core.tuple with the cutlass.cute.tuple module.
+    for name in ("ThrCopy", "ThrMma"):
+        if hasattr(cute, name) and not hasattr(cute.core, name):
+            setattr(cute.core, name, getattr(cute, name))
+
+    # CUTLASS DSL renamed make_fragment to make_rmem_tensor.
+    if hasattr(cute, "make_rmem_tensor") and not hasattr(cute, "make_fragment"):
+        cute.make_fragment = cute.make_rmem_tensor
+
+
+_setup_cutlass_dsl_compatibility()
+
 from .logger import logger
 from .version import __version__
 

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

@@ -45,7 +45,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,
+from ..utils import (ActivationType, deep_gemm_jit_warmup_buckets,
                      fp4_scale_infer_shape,
                      get_last_power_of_2_num_tokens_buckets,
                      is_nvfp4_marlin_supported_sm, last_positive_power_of_2)
@@ -2019,7 +2019,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,
     )
 
@@ -2064,6 +2064,104 @@ def forward(
         return output
 
 
+class Fp8PrequantizedSwapABGemmRunner(TunableRunner):
+    """Runs DeepGemm with pre-quantized FP8 activations and packed scales."""
+
+    tuning_config = TuningConfig(
+        dynamic_tensor_specs=(DynamicTensorSpec(
+            0, 0, deep_gemm_jit_warmup_buckets), ),
+        constraint_specs=(ConstraintSpec(
+            1, 0, lambda input_shapes: input_shapes[0][0]), ),
+        exclude_from_cache=True,
+    )
+
+    def __init__(self, output_dtype: torch.dtype,
+                 disable_ue8m0_cast: bool) -> None:
+        self.output_dtype = output_dtype
+        self.disable_ue8m0_cast = disable_ue8m0_cast
+
+    def unique_id(self):
+        return (
+            self.output_dtype,
+            self.disable_ue8m0_cast,
+        )
+
+    def get_valid_tactics(
+        self,
+        inputs: List[torch.Tensor],
+        profile: OptimizationProfile,
+    ) -> List[int]:
+        return [0]
+
+    def forward(
+        self,
+        inputs: List[torch.Tensor],
+        tactic: int = -1,
+    ) -> torch.Tensor:
+        del tactic
+        activation, activation_scale, weight, weight_scale = inputs
+        scale_m_aligned = fp4_utils.pad_up(activation_scale.size(0), 4)
+        if activation_scale.stride() != (1, scale_m_aligned):
+            # Dynamic autotuning recreates constrained integer tensors with a
+            # contiguous layout. Restore the MN-major packed-scale stride that
+            # the real quantizers return and DeepGemm requires.
+            normalized_scale = torch.empty_strided(
+                activation_scale.shape, (1, scale_m_aligned),
+                dtype=activation_scale.dtype,
+                device=activation_scale.device)
+            normalized_scale.copy_(activation_scale)
+            activation_scale = normalized_scale
+        output = torch.empty(
+            (activation.size(0), weight.size(0)),
+            device=activation.device,
+            dtype=self.output_dtype,
+        )
+        deep_gemm.fp8_gemm_nt(
+            (activation, activation_scale),
+            (weight, weight_scale),
+            output,
+            disable_ue8m0_cast=self.disable_ue8m0_cast,
+        )
+        return output
+
+
+@torch.library.custom_op("trtllm::fp8_prequantized_swap_ab_gemm",
+                         mutates_args=())
+def fp8_prequantized_swap_ab_gemm(
+    activation: torch.Tensor,
+    activation_scale: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_dtype: torch.dtype = torch.bfloat16,
+    disable_ue8m0_cast: bool = False,
+) -> torch.Tensor:
+    runner = Fp8PrequantizedSwapABGemmRunner(output_dtype, disable_ue8m0_cast)
+    _, best_tactic = AutoTuner.get().choose_one(
+        "trtllm::fp8_prequantized_swap_ab_gemm",
+        [runner],
+        Fp8PrequantizedSwapABGemmRunner.tuning_config,
+        [activation, activation_scale, weight, weight_scale],
+    )
+    return runner(
+        inputs=[activation, activation_scale, weight, weight_scale],
+        tactic=best_tactic,
+    )
+
+
+@fp8_prequantized_swap_ab_gemm.register_fake
+def _(
+    activation: torch.Tensor,
+    activation_scale: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_dtype: torch.dtype = torch.bfloat16,
+    disable_ue8m0_cast: bool = False,
+) -> torch.Tensor:
+    del activation_scale, weight_scale, disable_ue8m0_cast
+    return activation.new_empty((activation.size(0), weight.size(0)),
+                                dtype=output_dtype)
+
+
 @torch.library.custom_op("trtllm::fp8_swap_ab_gemm", mutates_args=())
 def fp8_swap_ab_gemm(
     input: torch.Tensor,
@@ -2114,12 +2212,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/models/modeling_kimi_k25.py

@@ -372,19 +372,30 @@ def _vision_requires_replication(model_config: ModelConfig, num_heads: int) -> b
     mapping = model_config.mapping
     if mapping.enable_attention_dp:
         return True
+    # Helix carries its parallelism in cp with tp_size=1, so a tp-only test
+    # never trips; the tower has no context-parallel form, so any cp > 1
+    # must replicate.
+    if mapping.cp_size > 1:
+        return True
     return (num_heads % mapping.tp_size) != 0
 
 
 def _get_vision_tp_mapping(model_config: ModelConfig, num_heads: int) -> Mapping:
     if not _vision_requires_replication(model_config, num_heads):
         return model_config.mapping
 
+    # Fold every parallel dimension (incl. helix cp) into pp so each rank
+    # runs the tower replicated; without cp the world size collapses below
+    # the rank range under helix.
+    attn_ranks = (
+        model_config.mapping.pp_size * model_config.mapping.tp_size * model_config.mapping.cp_size
+    )
     return Mapping(
-        world_size=model_config.mapping.pp_size * model_config.mapping.tp_size,
+        world_size=attn_ranks,
         rank=model_config.mapping.rank,
         gpus_per_node=model_config.mapping.gpus_per_node,
         tp_size=1,
-        pp_size=model_config.mapping.pp_size * model_config.mapping.tp_size,
+        pp_size=attn_ranks,
     )
 
 

tensorrt_llm/_torch/modules/dwdp/setup.py

@@ -689,7 +689,11 @@ def fixup_moe_backends(
         # ConfigurableMoE has its own ep_size, slot_start, etc. that are used
         # in its forward path.  The backend is the inner module that holds
         # weight parameters.
-        configurable_moe = getattr(layer.mlp, "experts", None)
+        # ``moe_module`` is what _get_moe_and_experts() just resolved, so this
+        # is model-agnostic: on DeepSeek it is layer.mlp and this stays exactly
+        # equivalent to the old getattr(layer.mlp, "experts", None); on K3 it
+        # is layer.block_sparse_moe, which has no ``.mlp`` at all.
+        configurable_moe = _get_configurable_moe(moe_module)
         targets = [experts_module]
         if configurable_moe is not None and configurable_moe is not experts_module:
             targets.insert(0, configurable_moe)
@@ -1110,6 +1114,20 @@ def _get_decoder_model(model: nn.Module) -> nn.Module:
     )
 
 
+def _get_configurable_moe(moe_module: Optional[nn.Module]) -> Optional[nn.Module]:
+    """The ConfigurableMoE wrapper of an MoE module, if the model uses one.
+
+    DeepSeek calls it ``experts``; Kimi K3's ``KimiK3MoERuntime`` calls the
+    same thing ``routed_experts``.  Returns None when the module has neither.
+    """
+    if moe_module is None:
+        return None
+    experts = getattr(moe_module, "experts", None)
+    if experts is None:
+        experts = getattr(moe_module, "routed_experts", None)
+    return experts
+
+
 def _get_moe_and_experts(
     layer: nn.Module,
 ) -> Tuple[Optional[nn.Module], Optional[nn.Module]]:
@@ -1118,22 +1136,31 @@ def _get_moe_and_experts(
     The standard path for DeepSeek is:
         layer.mlp (Deepseekv3MoE) -> .experts (MoE backend)
 
+    Kimi K3 spells the same shape differently:
+        layer.block_sparse_moe (KimiK3MoERuntime) -> .routed_experts (MoE backend)
+
     Returns:
         Tuple of (moe_module, experts_module) where moe_module is the wrapper
         (e.g. Deepseekv3MoE) and experts_module is the backend (e.g.
         CutlassFusedMoE, ConfigurableMoE, etc.). Both may be None if the
         layer is not an MoE layer.
     """
+    # K3's *dense* layers do carry an ``mlp``, but this function is only ever
+    # reached for layer indices that registered themselves from
+    # ConfigurableMoE.__init__, so a dense layer never gets here and the
+    # ``mlp``-first order stays safe.
     mlp = getattr(layer, "mlp", None)
+    if mlp is None:
+        mlp = getattr(layer, "block_sparse_moe", None)
     if mlp is None:
         return None, None
 
     # Check if mlp itself is an MoE backend (has w3_w1_weight)
     if hasattr(mlp, "w3_w1_weight"):
         return mlp, mlp
 
-    # Standard path: mlp.experts
-    experts = getattr(mlp, "experts", None)
+    # Standard path: mlp.experts (K3: block_sparse_moe.routed_experts)
+    experts = _get_configurable_moe(mlp)
     if experts is not None:
         # Prefer the inner backend (ConfigurableMoE wraps it)
         backend = getattr(experts, "backend", None)

tensorrt_llm/_torch/utils.py

@@ -418,6 +418,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])