From 50dc9758b68e96ccefc208bde325d22fc66c8cb0 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 15 Sep 2026 03:56:14 +0000 Subject: [PATCH 1/4] [None][chore] Core: MoE backend enum, deepgemm JIT warmup, postproc metrics, cleanup Ports the non-Rubin-specific core changes from the internal Rubin branch. Public API / config: - Add "CUTEDSL_FC12" to MoeConfig.backend. The backend implementation (CuteDslFc12FusedMoE + moe_resolution registration) lands separately; this only opens the config gate. llm_args_golden_manifest.json is updated to match. Autotune / JIT warmup: - Add deep_gemm_jit_warmup_buckets() and use it for the two DeepGemm runners whose only purpose is to drive JIT warmup. A step-16 M grid over the whole range is required: the SM100 layout heuristic selects on ceil_div(m, block_m) with every candidate block_m a multiple of 16, and the last-wave-utilization tie-break keeps oscillating at high M, so a coarse high-M band silently skips layouts that then compile mid-inference (nvcc fork under the GIL, stalling every attention-DP rank). Fp8BlockScalingGemmRunner also gets exclude_from_cache so a warm disk cache cannot short-circuit the warmup. Executor: - Propagate decoding_iter, avg_decoded_tokens_per_iter and cached_tokens through PostprocWorker.Output, so a result served by a postproc worker reports the same metrics as the in-process path. - When TRTLLM_WORKER_DISABLE_GC=1, also disable dynamo's post-compile gc.collect(1): with automatic GC off it walks every object allocated since the previous compile, costing seconds per recompile. Models: - Kimi K2.5: replicate the vision tower whenever cp_size > 1. Helix carries its parallelism in cp with tp_size=1, so the existing tp-only check never fired and the folded world size collapsed below the rank range. - dwdp: resolve the MoE wrapper model-agnostically, so Kimi K3's block_sparse_moe/routed_experts spelling is handled alongside DeepSeek's mlp/experts. Misc: - CuTe DSL compatibility shim for legacy cute.core.ThrCopy/ThrMma and cute.make_fragment, needed by QuACK and Transformer Engine against the pinned CUTLASS DSL. - Drop the unused include from six thop GEMM translation units. - Do not fail collection of tests/integration/defs when torch._inductor is unavailable. - Unit test for submit.py's replace_env_in_file. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- cpp/tensorrt_llm/thop/fp4Gemm.cpp | 1 - cpp/tensorrt_llm/thop/fp4GemmTrtllmGen.cpp | 3 +- .../thop/fp4xFp8GemmTrtllmGen.cpp | 3 +- .../thop/fp8BatchedGemmTrtllmGen.cpp | 3 +- .../thop/fp8PerTensorScalingTrtllmGenGemm.cpp | 3 +- cpp/tensorrt_llm/thop/fp8RowwiseGemm.cpp | 3 +- tensorrt_llm/__init__.py | 23 +++++++++++++ .../_torch/custom_ops/torch_custom_ops.py | 13 +++++--- .../_torch/models/modeling_kimi_k25.py | 15 +++++++-- tensorrt_llm/_torch/modules/dwdp/setup.py | 33 +++++++++++++++++-- tensorrt_llm/_torch/utils.py | 32 ++++++++++++++++++ tensorrt_llm/executor/postproc_worker.py | 15 +++++++-- tensorrt_llm/executor/result.py | 4 +++ tensorrt_llm/executor/worker.py | 5 +++ tensorrt_llm/llmapi/llm_args.py | 4 +-- .../usage/llm_args_golden_manifest.json | 3 +- tests/integration/defs/__init__.py | 7 ++-- tests/unittest/scripts/test_perf_submit.py | 24 ++++++++++++++ 18 files changed, 167 insertions(+), 27 deletions(-) diff --git a/cpp/tensorrt_llm/thop/fp4Gemm.cpp b/cpp/tensorrt_llm/thop/fp4Gemm.cpp index f0066c714624..6df24e501fd2 100644 --- a/cpp/tensorrt_llm/thop/fp4Gemm.cpp +++ b/cpp/tensorrt_llm/thop/fp4Gemm.cpp @@ -27,7 +27,6 @@ #endif #include -#include #include #include diff --git a/cpp/tensorrt_llm/thop/fp4GemmTrtllmGen.cpp b/cpp/tensorrt_llm/thop/fp4GemmTrtllmGen.cpp index 1c9ac017fb1a..4993f8826ac1 100644 --- a/cpp/tensorrt_llm/thop/fp4GemmTrtllmGen.cpp +++ b/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 -#include #include diff --git a/cpp/tensorrt_llm/thop/fp4xFp8GemmTrtllmGen.cpp b/cpp/tensorrt_llm/thop/fp4xFp8GemmTrtllmGen.cpp index b657b92eb34b..f22b262cca91 100644 --- a/cpp/tensorrt_llm/thop/fp4xFp8GemmTrtllmGen.cpp +++ b/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 -#include #include diff --git a/cpp/tensorrt_llm/thop/fp8BatchedGemmTrtllmGen.cpp b/cpp/tensorrt_llm/thop/fp8BatchedGemmTrtllmGen.cpp index f3da650a9408..c7b46e37a061 100644 --- a/cpp/tensorrt_llm/thop/fp8BatchedGemmTrtllmGen.cpp +++ b/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 -#include #include diff --git a/cpp/tensorrt_llm/thop/fp8PerTensorScalingTrtllmGenGemm.cpp b/cpp/tensorrt_llm/thop/fp8PerTensorScalingTrtllmGenGemm.cpp index 7f044a198edb..940652a52fc9 100644 --- a/cpp/tensorrt_llm/thop/fp8PerTensorScalingTrtllmGenGemm.cpp +++ b/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 -#include #include diff --git a/cpp/tensorrt_llm/thop/fp8RowwiseGemm.cpp b/cpp/tensorrt_llm/thop/fp8RowwiseGemm.cpp index b87e7973935e..a63a0da311be 100644 --- a/cpp/tensorrt_llm/thop/fp8RowwiseGemm.cpp +++ b/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 -#include #include #include diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 516e31ad2c52..814aa9874244 100644 --- a/tensorrt_llm/__init__.py +++ b/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__ diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index febfe31293e5..507fb4c4614c 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/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, ) @@ -2114,12 +2114,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( diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index f630f276fc3d..1a8db4251764 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -372,6 +372,11 @@ 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 @@ -379,12 +384,18 @@ 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, ) diff --git a/tensorrt_llm/_torch/modules/dwdp/setup.py b/tensorrt_llm/_torch/modules/dwdp/setup.py index e8a1cb740a7a..b15f4330e69f 100644 --- a/tensorrt_llm/_torch/modules/dwdp/setup.py +++ b/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,13 +1136,22 @@ 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 @@ -1132,8 +1159,8 @@ def _get_moe_and_experts( 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) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index da6797268ffe..84a393c1b8dd 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/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]) diff --git a/tensorrt_llm/executor/postproc_worker.py b/tensorrt_llm/executor/postproc_worker.py index 184b9f924c8e..b8943aac24ab 100644 --- a/tensorrt_llm/executor/postproc_worker.py +++ b/tensorrt_llm/executor/postproc_worker.py @@ -88,6 +88,9 @@ class Output(NamedTuple): should_abort: bool = False finish_reason: Optional[str] = None num_generated_tokens: Optional[int] = None + decoding_iter: int = 0 + avg_decoded_tokens_per_iter: Optional[float] = None + cached_tokens: int = 0 def __init__( self, @@ -237,8 +240,9 @@ async def handle_single_input(inp: PostprocWorker.Input, self._records.pop(client_id, None) return try: - is_final = inp.rsp.result.is_final if is_llm_response( - inp.rsp) else True + response_result = inp.rsp.result if is_llm_response( + inp.rsp) else None + is_final = response_result.is_final if response_result else True res, metrics, perf_metrics, disaggregated_params = await self._handle_input( inp) record = self._records.get(client_id) @@ -264,6 +268,13 @@ async def handle_single_input(inp: PostprocWorker.Input, should_abort=should_abort, finish_reason=finish_reason, num_generated_tokens=num_generated_tokens, + decoding_iter=getattr(response_result, "decoding_iter", + 0), + avg_decoded_tokens_per_iter=getattr( + response_result, "avg_decoded_tokens_per_iter", + None), + cached_tokens=getattr(response_result, "cached_tokens", + 0), )) if is_final: self._records.pop(client_id, None) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 854bc4452872..be3370b64570 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -541,6 +541,10 @@ def _handle_response(self, if isinstance(response, PostprocWorker.Output): self._done = response.is_final + self.decoding_iter = response.decoding_iter + self.avg_decoded_tokens_per_iter = ( + response.avg_decoded_tokens_per_iter) + self.cached_tokens = response.cached_tokens if isinstance(response.res, CompletionOutput): # in streaming mode self._outputs[0] = response.res diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 40cee5b65b52..a9fa5d7a6c88 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -422,6 +422,11 @@ def notify_proxy_threads_to_quit(): # Optionally disable GC (default: not disabled) if os.getenv("TRTLLM_WORKER_DISABLE_GC", "0") == "1": gc.disable() + # With automatic GC off, dynamo's post-compile gc.collect(1) walks every + # object allocated since the previous compile (seconds per recompile). + if "TORCH_DYNAMO_RUN_GC_AFTER_COMPILE" not in os.environ: + import torch._dynamo.config + torch._dynamo.config.run_gc_after_compile = False with worker: try: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 59c5030de54c..5bc2560353cd 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1605,8 +1605,8 @@ class MoeConfig(StrictBaseModel): """Configuration for MoE.""" backend: Literal[ "AUTO", "CUTLASS", "CUTEDSL", "TRTLLM", "DEEPGEMM", "DENSEGEMM", - "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL"] = Field( + "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", + "CUTEDSL_FC12"] = Field( default='AUTO', description="MoE backend to use. " "AUTO selects default backend based on model. It currently doesn\'t always give the best choice for all scenarios. The capabilities of auto selection will be improved in future releases." diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e4f0bcc10f7e..73f6d9913f57 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -897,7 +897,8 @@ "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL" + "MEGAMOE_CUTEDSL", + "CUTEDSL_FC12" ], "capture_policy": "literal", "kind": "categorical", diff --git a/tests/integration/defs/__init__.py b/tests/integration/defs/__init__.py index 23c84ca923fc..3b38d9b04bcf 100644 --- a/tests/integration/defs/__init__.py +++ b/tests/integration/defs/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -19,4 +19,7 @@ # But if the import happens lazily after the test starts, pytest will think you leaked # the thread. We thus do the import here to prevent thread leak issues cropping up when messing # with the import statements in tests. -from torch._inductor import lowering # NOQA +try: + from torch._inductor import lowering # NOQA +except ModuleNotFoundError: + pass diff --git a/tests/unittest/scripts/test_perf_submit.py b/tests/unittest/scripts/test_perf_submit.py index 48e0922fb947..cc3034d05340 100644 --- a/tests/unittest/scripts/test_perf_submit.py +++ b/tests/unittest/scripts/test_perf_submit.py @@ -148,6 +148,30 @@ def test_example_worker_environment_exports_positive_concurrency(example_submit_ assert worker_environment["TLLM_BENCHMARK_REQ_QUEUES_SIZE"] == "4301" +def test_example_replace_env_in_file_replaces_all_variables( + example_submit_module: ModuleType, + tmp_path: Path, +) -> None: + config_path = tmp_path / "task.yaml" + config_path.write_text( + "model_root: LLM_MODELS_ROOT\nhf_home: HF_HOME\n", + encoding="utf-8", + ) + + output_dir = example_submit_module.replace_env_in_file( + tmp_path, + config_path, + { + "LLM_MODELS_ROOT": "/models", + "HF_HOME": "/cache", + }, + ) + + assert (Path(output_dir) / config_path.name).read_text(encoding="utf-8") == ( + "model_root: /models\nhf_home: /cache\n" + ) + + def test_ci_submit_selects_same_least_duration_shard_as_pytest_split( ci_submit_module: ModuleType, tmp_path: Path, From ec8d25994b1f92beb9f502a8796be38958b4ccd8 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 15 Sep 2026 08:54:05 +0000 Subject: [PATCH 2/4] [None][fix] dispatch CuteDSL MoE backends by subclass in create_moe_backend PR5 adds `CUTEDSL_FC12` to the `MoeConfig.backend` literal, but `create_moe_backend` matched the CuteDSL branch by class identity: elif moe_cls in (CuteDslFusedMoE, CuteDslB12xFusedMoE): `CuteDslFc12FusedMoE` subclasses `CuteDslFusedMoE` and takes the same argument set, so selecting the newly-legal backend string would fall through to the `Unsupported moe backend` raise. Match by `issubclass`, as the DeepGEMM and MegaMoE branches below already do for the same reason. Every class dispatched after this branch derives from `MoEImplBase` / `MoE`, so the wider match steals nothing; with no CuteDSL subclass on this branch the change is inert until the FC12 backend lands. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/moe/fused_moe/create_moe.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py index 3464268b5fbd..b1cd3025eeb5 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py @@ -178,7 +178,12 @@ def create_moe_backend( layer_idx=layer_idx, activation=activation, ) - elif moe_cls in (CuteDslFusedMoE, CuteDslB12xFusedMoE): + # ``issubclass`` for the same reason as the DeepGEMM and MegaMoE branches + # below: the Rubin FC12 backend subclasses ``CuteDslFusedMoE`` and takes the + # same argument set, so it must not fall through to the ``Unsupported moe + # backend`` raise. Every class dispatched after this branch derives from + # ``MoEImplBase`` / ``MoE``, so widening the match steals nothing. + elif issubclass(moe_cls, (CuteDslFusedMoE, CuteDslB12xFusedMoE)): # The narrower CuteDsl argument set: these kernels take no expert bias. return moe_cls( routing_method=routing_method, From 88e1d8980df40a13083ed12805228db16e17258b Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 00:58:05 +0000 Subject: [PATCH 3/4] [None][fix] add the fp8_prequantized_swap_ab_gemm op The pre-quantized FP8 path in Linear.apply selects torch.ops.trtllm.fp8_prequantized_swap_ab_gemm whenever the activation scale is int32 on SM100f with deep_gemm enabled, but no registration for that op exists. Without it every DeepSeek V4 run on GB200/GB300 aborts during executor init: AttributeError: '_OpNamespace' 'trtllm' object has no attribute 'fp8_prequantized_swap_ab_gemm' Unlike fp8_swap_ab_gemm, the activation and its packed scale both carry the M dimension, so the tuning config ties input 1 dim 0 to input 0 dim 0, and the runner restores the MN-major packed-scale stride that DeepGemm requires after the autotuner recreates the constrained integer tensor contiguously. Verified on GB300: the DeepSeek V4 DSpark disaggregated accuracy guard (1p1d, DEP4) reaches GSM8K 96.21 with this op present and fails at executor init without it. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/custom_ops/torch_custom_ops.py | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 507fb4c4614c..4e67ed02fd66 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -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, From 104123b08532fa09a673f126d5b959018863e11d Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 02:47:54 +0000 Subject: [PATCH 4/4] [None][chore] drop CUTEDSL_FC12 now that the backend is not in this series The CuteDslFc12FusedMoE backend has been removed from the MoE PR: the custom op it drives was lost in the rebase, so it could never be selected. Nothing in this series consumes the backend string, so the enum value would be a `MoeConfig.backend` literal that resolves to an empty candidate list. Removes "CUTEDSL_FC12" from the literal and from the golden manifest, and reverts ec8d25994b -- the `issubclass` dispatch in create_moe_backend existed only so the FC12 subclass would not fall through to the "Unsupported moe backend" raise. With no CuteDSL subclass in the tree the identity check is exact again. Both come back with the backend and its op. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/moe/fused_moe/create_moe.py | 7 +------ tensorrt_llm/llmapi/llm_args.py | 4 ++-- tensorrt_llm/usage/llm_args_golden_manifest.json | 3 +-- 3 files changed, 4 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py index b1cd3025eeb5..3464268b5fbd 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py @@ -178,12 +178,7 @@ def create_moe_backend( layer_idx=layer_idx, activation=activation, ) - # ``issubclass`` for the same reason as the DeepGEMM and MegaMoE branches - # below: the Rubin FC12 backend subclasses ``CuteDslFusedMoE`` and takes the - # same argument set, so it must not fall through to the ``Unsupported moe - # backend`` raise. Every class dispatched after this branch derives from - # ``MoEImplBase`` / ``MoE``, so widening the match steals nothing. - elif issubclass(moe_cls, (CuteDslFusedMoE, CuteDslB12xFusedMoE)): + elif moe_cls in (CuteDslFusedMoE, CuteDslB12xFusedMoE): # The narrower CuteDsl argument set: these kernels take no expert bias. return moe_cls( routing_method=routing_method, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5bc2560353cd..59c5030de54c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1605,8 +1605,8 @@ class MoeConfig(StrictBaseModel): """Configuration for MoE.""" backend: Literal[ "AUTO", "CUTLASS", "CUTEDSL", "TRTLLM", "DEEPGEMM", "DENSEGEMM", - "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", - "CUTEDSL_FC12"] = Field( + "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL"] = Field( default='AUTO', description="MoE backend to use. " "AUTO selects default backend based on model. It currently doesn\'t always give the best choice for all scenarios. The capabilities of auto selection will be improved in future releases." diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 73f6d9913f57..e4f0bcc10f7e 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -897,8 +897,7 @@ "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL", - "CUTEDSL_FC12" + "MEGAMOE_CUTEDSL" ], "capture_policy": "literal", "kind": "categorical",