[None][feat] Rubin K3 merge-back cleanup: CFT release op, MegaMoE multi-node bootstrap, DeepGEMM JIT warmup

cpp/tensorrt_llm/thop/moe/moeAlltoAllOp.cpp

@@ -299,7 +299,15 @@ torch::Tensor moeA2AInitializeOp(torch::Tensor const& workspace, int64_t epRank,
 // CFT Handle-Based Counted Writes Initialization
 // ============================================================================
 
-// Static CftLeManager — lives for the process lifetime (like workspace).
+// Static CftLeManager — one per process, bound to a single workspace.
+//
+// The manager owns a logical endpoint bound to the workspace's MNNVL
+// allocation, so it must be destroyed before that allocation is freed and its
+// virtual address is recycled; otherwise a later binding could resolve to an
+// endpoint left over from a dead allocation. The Python NVLinkOneSided
+// teardown calls moe_a2a_cft_release while the workspace is still alive.
+// MoeAlltoAll otherwise retains its shared workspaces for the process
+// lifetime, so the manager normally survives until static destruction.
 static std::unique_ptr<tensorrt_llm::kernels::moe_comm::CftLeManager> g_cft_manager;
 
 // Initialize CFT Logical Endpoints by binding the LE to the MNNVL workspace.
@@ -396,6 +404,45 @@ void moeA2ACftInitializeOp(torch::Tensor const& workspace, int64_t workspaceMemH
     }
 }
 
+// Release the CFT logical endpoint before its backing workspace is freed.
+//
+// Idempotent, and a no-op unless the manager is actually bound to this
+// workspace's rank region: the caller passes the workspace it is tearing down,
+// and a manager bound to some other allocation must outlive that teardown.
+// Destroying the manager here — rather than at static destruction — keeps the
+// endpoint from outliving the virtual address it is bound to.
+void moeA2ACftReleaseOp(torch::Tensor const& workspace, int64_t epRank)
+{
+    CHECK_TH_CUDA(workspace);
+    CHECK_TYPE(workspace, torch::kUInt8);
+    TORCH_CHECK(workspace.dim() == 2, "workspace must be a 2D tensor of shape [epSize, sizePerRank]");
+    TORCH_CHECK(epRank >= 0 && epRank < workspace.size(0), "epRank must be in the range [0, epSize)");
+
+    if (!g_cft_manager)
+    {
+        return;
+    }
+
+    // An uninitialized manager holds no endpoint (initialization threw part
+    // way through); drop it unconditionally so a retry starts clean.
+    if (!g_cft_manager->isInitialized())
+    {
+        g_cft_manager.reset();
+        return;
+    }
+
+    CUdeviceptr workspaceRankPtr
+        = reinterpret_cast<CUdeviceptr>(workspace.data_ptr<uint8_t>() + epRank * workspace.stride(0));
+    if (g_cft_manager->getLocalBackingPtr() != workspaceRankPtr)
+    {
+        return;
+    }
+
+    // ~CftLeManager runs destroy(): unbind, destroy the local and imported
+    // endpoints, and release the reserved LE id block.
+    g_cft_manager.reset();
+}
+
 // MoE All-to-All Dispatch Operation
 // This operation dispatches tokens and their associated payloads to different expert ranks.
 //
@@ -1062,6 +1109,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, module)
     module.def(
         "moe_a2a_cft_initialize(Tensor(a!) workspace, int workspace_mem_handle, "
         "int workspace_size_per_rank, int ep_rank, int ep_size) -> ()");
+    module.def("moe_a2a_cft_release(Tensor(a!) workspace, int ep_rank) -> ()");
     module.def(
         "moe_a2a_initialize(Tensor(a!) workspace, int ep_rank, int ep_size, int max_num_tokens_per_rank, "
         "int? eplb_stats_num_experts=None, bool can_use_cft_counted_writes=False) -> Tensor");
@@ -1088,4 +1136,5 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, module)
     module.impl(
         "moe_a2a_get_combine_payload_tensor", &tensorrt_llm::torch_ext::moe_comm::moeA2AGetCombinePayloadTensorOp);
     module.impl("moe_a2a_cft_initialize", &tensorrt_llm::torch_ext::moe_comm::moeA2ACftInitializeOp);
+    module.impl("moe_a2a_cft_release", &tensorrt_llm::torch_ext::moe_comm::moeA2ACftReleaseOp);
 }

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

@@ -49,7 +49,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,
                      get_power_of_2_num_tokens_buckets,
@@ -2074,7 +2074,7 @@ def _(
     return input.new_empty((M, N), dtype=output_dtype)
 
 
-# deep_gemm_gen_tuning_buckets is imported from ..utils
+# deep_gemm_jit_warmup_buckets is imported from ..utils
 
 _USE_FUSED_FP8_QUANT_PACK = os.environ.get("TRTLLM_FUSED_FP8_QUANT_PACK",
                                            "1") == "1"
@@ -2147,7 +2147,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,
     )
 
@@ -2195,9 +2195,12 @@ def forward(
 class Fp8PrequantizedSwapABGemmRunner(TunableRunner):
     """Runs DeepGemm with pre-quantized FP8 activations and packed scales."""
 
+    # The same step-16 grid as the other two DeepGemm runners: a layout no
+    # bucket selects is compiled mid-inference instead, and DeepGemm forks
+    # nvcc while holding the GIL.
     tuning_config = TuningConfig(
         dynamic_tensor_specs=(DynamicTensorSpec(
-            0, 0, deep_gemm_gen_tuning_buckets), ),
+            0, 0, deep_gemm_jit_warmup_buckets), ),
         constraint_specs=(ConstraintSpec(
             1, 0, lambda input_shapes: input_shapes[0][0]), ),
         exclude_from_cache=True,
@@ -2340,12 +2343,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/cute_dsl_kernels/rubin/moe/manual_mma_128dp.py

@@ -73,7 +73,12 @@ def build_static_idesc_base(
     OR's them in from the SF TMEM addresses.  n_dim is static (non-swapAB) and
     folded in here.
     """
-    assert umma_m in (64, 128, 256), f"Unsupported UMMA_M={umma_m}"
+    # 64 is not encodable here. ``m_dim = umma_m >> 4`` placed at bit 24 only
+    # coincides with the real m_dim field (bits [27,29), holding M >> 7) for
+    # 128 -> bit 27 and 256 -> bit 28. For 64 it yields 4 << 24, i.e. bit 26,
+    # which is _BIT_SFA_LAYOUT: the descriptor would silently select
+    # SFA_128dp_Unique and discard the caller's sfa_layout argument.
+    assert umma_m in (128, 256), f"Unsupported UMMA_M={umma_m}"
     assert umma_k in _K_SIZE_FIELD, f"Unsupported UMMA_K={umma_k}"
     assert 0 <= sfa_layout < 2
 
@@ -108,8 +113,15 @@ def compute_idesc(
     idesc = Int32(static_base)
     sfa_top = Int32(sfa_tmem_addr_i32) & Int32(0xC0000000)
     sfb_top = Int32(sfb_tmem_addr_i32) & Int32(0xC0000000)
-    idesc = idesc | (sfa_top >> Int32(30 - _BIT_A_SF_ID))
-    idesc = idesc | (sfb_top >> Int32(30 - _BIT_B_SF_ID))
+    # Mask after shifting. ``Int32`` is signed, so ``>>`` is an arithmetic
+    # shift: a TMEM address with bit 31 set sign-extends and would leave stray
+    # high bits set. Bit 31 is _BIT_K_SIZE_LO, so for umma_k=128 (k_size 2)
+    # that flips k_size to 3 and the instruction runs the wrong K. Matches
+    # ``compute_idesc`` in
+    # cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py,
+    # which masks both fields; this copy had dropped it.
+    idesc = idesc | ((sfa_top >> Int32(30 - _BIT_A_SF_ID)) & Int32(0x3 << _BIT_A_SF_ID))
+    idesc = idesc | ((sfb_top >> Int32(30 - _BIT_B_SF_ID)) & Int32(0x3 << _BIT_B_SF_ID))
     return idesc
 
 

tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/utils.py

@@ -44,6 +44,7 @@
 # This file is copied and modified from cutlass https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/core.py
 
 import ctypes
+import math
 import os
 from typing import Union
 
@@ -100,6 +101,11 @@ def __init__(
             f"pointer must be {self._assumed_align} bytes aligned"
         )
 
+    def __add__(self, offset: int) -> Pointer:  # type: ignore[override]
+        offset_bytes = offset * self._dtype.width // 8
+        assumed_align = math.gcd(offset_bytes, self._assumed_align)
+        return _Pointer(self._pointer + offset_bytes, self._dtype, self._addr_space, assumed_align)
+
     def size_in_bytes(self) -> int:
         return ctypes.sizeof(ctypes.c_void_p(int(self._pointer)))
 

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/modules/kimi_kda/kimi_kda_mixer.py

@@ -71,7 +71,11 @@ def _cast(tensor: torch.Tensor) -> torch.Tensor:
             continue
         child._apply(_cast)
     for name, param in module.named_parameters(recurse=False):
-        param.data = _cast(param.data)
+        # Pass the parameter, not param.data: reading .data detaches, and
+        # MetaInitMode rejects aten.detach on a meta tensor, so casting the
+        # module's own parameters under meta init raised MetaInitException
+        # before _cast ever got to its is_meta branch.
+        param.data = _cast(param)
 
 
 def _stage_state_rows(ssm_pool: torch.Tensor, slot_indices: torch.Tensor) -> torch.Tensor:

tensorrt_llm/_torch/modules/situ.py

@@ -12,6 +12,10 @@
 import triton.language.extra.libdevice as tldevice  # type: ignore[import]
 from torch import nn
 
+from tensorrt_llm._utils import get_sm_version
+
+from ..flashinfer_utils import get_env_enable_pdl
+
 
 class SituAndMul(nn.Module):
     """SiTU activation with gate/up multiplicative gating.
@@ -59,6 +63,7 @@ def situ_and_mul_kernel(
     linear_beta,
     BLOCK_SIZE: tl.constexpr,
     HAS_LINEAR_BETA: tl.constexpr,
+    LAUNCH_WITH_PDL: tl.constexpr,
 ) -> None:
     """Fused :class:`SituAndMul` on a packed ``[gate | up]`` row layout."""
     i = tl.program_id(axis=0).to(tl.int64)
@@ -67,6 +72,9 @@ def situ_and_mul_kernel(
     o_row_ptr = o_ptr + o_stride * i
     x_row_ptr = x_ptr + x_stride * i
 
+    if LAUNCH_WITH_PDL:
+        tl.extra.cuda.gdc_wait()
+
     offsets = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = offsets < d
 
@@ -80,6 +88,9 @@ def situ_and_mul_kernel(
 
     tl.store(o_row_ptr + offsets, result, mask=mask)
 
+    if LAUNCH_WITH_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
 
 @torch.library.custom_op("trtllm::situ_and_mul", mutates_args=())
 def situ_and_mul(x: torch.Tensor, beta: float, linear_beta: Optional[float] = None) -> torch.Tensor:
@@ -95,6 +106,7 @@ def situ_and_mul(x: torch.Tensor, beta: float, linear_beta: Optional[float] = No
     def grid(meta: Mapping[str, int]) -> tuple[int, int]:
         return (b, triton.cdiv(d, meta["BLOCK_SIZE"]))
 
+    launch_with_pdl = get_env_enable_pdl() and get_sm_version() >= 90
     situ_and_mul_kernel[grid](
         o_ptr=output,
         o_stride=output.stride(0),
@@ -105,6 +117,8 @@ def grid(meta: Mapping[str, int]) -> tuple[int, int]:
         linear_beta=float(linear_beta) if linear_beta is not None else 1.0,
         BLOCK_SIZE=1024,
         HAS_LINEAR_BETA=linear_beta is not None,
+        LAUNCH_WITH_PDL=launch_with_pdl,
+        launch_pdl=launch_with_pdl,
     )
     return output