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..4e67ed02fd66 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, ) @@ -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( 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/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,