From 28cd462db209c0eb48855b71d8be30a2a4a3aa0a Mon Sep 17 00:00:00 2001 From: Shiki Wu Date: Sat, 3 Oct 2026 07:57:32 -0700 Subject: [PATCH 1/5] support FP32 for lm head to reduce TIM Signed-off-by: Shiki Wu --- tensorrt_llm/_torch/model_config.py | 4 + tensorrt_llm/_torch/models/modeling_utils.py | 2 + tensorrt_llm/_torch/modules/embedding.py | 65 +++++++- .../_torch/pyexecutor/model_loader.py | 2 + tensorrt_llm/llmapi/llm_args.py | 12 ++ .../usage/llm_args_golden_manifest.json | 9 ++ .../test_lists/test-db/l0_h100.yml | 2 + .../modules/test_lm_head_output_dtype.py | 143 ++++++++++++++++++ .../api_stability/references/llm.yaml | 4 + tests/unittest/llmapi/test_lm_head_dtype.py | 81 ++++++++++ 10 files changed, 322 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/_torch/modules/test_lm_head_output_dtype.py create mode 100644 tests/unittest/llmapi/test_lm_head_dtype.py diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index d0e3d697383f..50ef60c56be1 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -233,6 +233,10 @@ class ModelConfig(Generic[TConfig]): # draft models so greedy draft sampling can do a lighter TP gather. Defaults # to True to preserve behavior for every non-draft model. lm_head_gather_output: bool = True + # Dtype of the logits written by the LM head GEMM (LlmArgs.lm_head_dtype). + # None keeps the model dtype; torch.float32 writes the float32 accumulator + # directly instead of rounding the logits to bf16/fp16 first. + lm_head_dtype: Optional[torch.dtype] = None lora_config: Optional["LoraConfig"] = None sparse_attention_config: Optional["SparseAttentionConfig"] = None kv_cache_compression_config: Optional["KvCacheCompressionConfig"] = None diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 21905050ebf0..5e7925dfc5d5 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -578,6 +578,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], hidden_size, dtype=config.pretrained_config.torch_dtype, quant_config=lm_head_quant_config, + output_dtype=config.lm_head_dtype, ) else: if (hasattr(config, 'lora_config') @@ -607,6 +608,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], use_custom_cublas_mm=getattr(model, 'use_custom_cublas_mm', False), quant_config=lm_head_quant_config, + output_dtype=config.lm_head_dtype, ) if self.has_custom_lm_head: diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index a707372590a2..93d035c0639d 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/tensorrt_llm/_torch/modules/embedding.py @@ -14,6 +14,28 @@ from .linear import Linear, TensorParallelMode +def _matmul_with_output_dtype(input: torch.Tensor, weight: torch.Tensor, + output_dtype: torch.dtype) -> torch.Tensor: + """Compute ``input @ weight.T`` and write the result in ``output_dtype``. + + ``input`` and ``weight`` stay in the model dtype; cuBLAS accumulates in + float32 and stores the result directly in ``output_dtype``, so the output + is not rounded to the model dtype first. + """ + out_features = weight.shape[0] + output_shape = (*input.shape[:-1], out_features) + input_2d = input.reshape(-1, input.shape[-1]) + if input_2d.shape[0] == 0: + return input.new_empty(output_shape, dtype=output_dtype) + # cublas_mm reads mat_b as column-major with leading dimension K. + assert weight.is_contiguous(), "LM head weight must be contiguous" + output = torch.ops.trtllm.cublas_mm(input_2d.contiguous(), + weight.t(), + None, + out_dtype=output_dtype) + return output.view(output_shape) + + class LMHead(Linear): """LM head layer. @@ -23,6 +45,11 @@ class LMHead(Linear): dtype (Optional[torch.dtype]): type of the parameters. mapping (Optional[Mapping]): parallelism configuration. If not provided, the embedding is not parallelized. + output_dtype (Optional[torch.dtype]): dtype of the logits written by + the GEMM. None keeps the parameter dtype. torch.float32 keeps the + weights and inputs in the parameter dtype but writes the float32 + accumulator directly, so logits are not rounded to bf16/fp16. + Only supported for unquantized, non-ROW-parallel heads. """ def __init__( @@ -36,6 +63,7 @@ def __init__( reduce_output: bool = True, use_custom_cublas_mm: bool = False, quant_config=None, + output_dtype: Optional[torch.dtype] = None, ): local_in_features = embedding_dim local_out_features = num_embeddings @@ -119,6 +147,21 @@ def __init__( self.weight = Parameter(torch.empty(weight_shape, dtype=dtype)) self.register_parameter("bias", None) + if output_dtype is not None: + if self.has_any_quant: + raise NotImplementedError( + "LMHead output_dtype is only supported for unquantized " + f"LM heads (quant_algo={quant_config.quant_algo})") + if self.tp_mode == TensorParallelMode.ROW: + raise NotImplementedError( + "LMHead output_dtype does not support ROW tensor-parallel " + "mode") + self.output_dtype = output_dtype + + def _writes_output_dtype(self) -> bool: + return (self.output_dtype is not None + and self.output_dtype != self.weight.dtype) + @property def vocab_size_padded(self) -> int: if self.tp_mode == TensorParallelMode.COLUMN and self.gather_output: @@ -141,8 +184,12 @@ def forward( slice_width = ceil_div(self.out_features, tp_size) slice_start = tp_rank * slice_width slice_end = min((tp_rank + 1) * slice_width, self.out_features) - output = F.linear(input, self.weight[slice_start:slice_end, :], - None) + weight = self.weight[slice_start:slice_end, :] + if self._writes_output_dtype(): + output = _matmul_with_output_dtype(input, weight, + self.output_dtype) + else: + output = F.linear(input, weight, None) else: output = super().forward(input, all_reduce_params=all_reduce_params) if (self.tp_mode == TensorParallelMode.COLUMN and self.gather_output @@ -151,6 +198,20 @@ def forward( return output + def apply_linear(self, + input, + bias, + lora_params: Optional[dict] = None, + layer_idx: Optional[int] = None): + # Linear.forward runs the GEMM through apply_linear for the COLUMN and + # non-parallel modes, so the TP all-gather and the vocab-padding slice + # in forward() apply unchanged to the output_dtype logits. + if self._writes_output_dtype(): + assert bias is None, "LMHead has no bias" + return _matmul_with_output_dtype(input, self.weight, + self.output_dtype) + return super().apply_linear(input, bias, lora_params, layer_idx) + def skip_forward( self, input: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 13e5a6c1ec87..cd685393a4e6 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -1922,6 +1922,8 @@ def _load_and_validate_config( multimodal_config=self.llm_args.multimodal_config, use_cute_dsl_bf16_bmm=self.llm_args.use_cute_dsl_bf16_bmm, use_cute_dsl_bf16_gemm=self.llm_args.use_cute_dsl_bf16_gemm, + lm_head_dtype=(torch.float32 if self.llm_args.lm_head_dtype + == "float32" else None), ) # Only pass model_kwargs if it's explicitly set (not None) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 825dfc6ae102..c7d84390c03f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5922,6 +5922,18 @@ def init_multimodal_config(cls, v): "Separated-routing MoE backends only.", status="beta") + lm_head_dtype: Literal["auto", "float32"] = Field( + default="auto", + description= + "Dtype of the logits written by the LM head GEMM. 'auto' keeps the " + "model dtype, so logits are rounded to bf16/fp16 before being upcast " + "to float32. 'float32' keeps the LM head weights and inputs in the " + "model dtype but writes the float32 GEMM accumulator directly, so " + "logits and logprobs carry no bf16/fp16 rounding (e.g. to reduce " + "train/inference logprob mismatch in RL). Unquantized LM heads of " + "decoder-only models only.", + status="prototype") + moe_config: MoeConfig = Field(default_factory=MoeConfig, description="MoE config.", status="beta") diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 3a6cc7982e90..73893bf07e36 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -825,6 +825,15 @@ "kind": "categorical", "path": "layer_wise_benchmarks_config.calibration_mode" }, + { + "allowed_values": [ + "auto", + "float32" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "lm_head_dtype" + }, { "allowed_values": [ "auto", diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 31659d64e83c..7c2552d62978 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -39,6 +39,7 @@ l0_h100: - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules - unittest/_torch/modules/test_w4a16_nvfp4_linear.py + - unittest/_torch/modules/test_lm_head_output_dtype.py # ------------- MoE components tests --------------- - unittest/_torch/moe/test_moe_load_balancer.py - unittest/_torch/moe/test_moe_routing.py @@ -144,6 +145,7 @@ l0_h100: - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[PYTHON-mla-ctx_fp16_gen_fp16] - unittest/llmapi/test_llm_telemetry.py - unittest/llmapi/test_return_routed_experts.py + - unittest/llmapi/test_lm_head_dtype.py - unittest/scaffolding - unittest/usage/test_transport.py - unittest/usage/test_e2e_capture.py diff --git a/tests/unittest/_torch/modules/test_lm_head_output_dtype.py b/tests/unittest/_torch/modules/test_lm_head_output_dtype.py new file mode 100644 index 000000000000..eef9e0d8d833 --- /dev/null +++ b/tests/unittest/_torch/modules/test_lm_head_output_dtype.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""LMHead(output_dtype=torch.float32), which backs LlmArgs.lm_head_dtype="float32". + +The weights and inputs stay bf16; the GEMM writes its float32 accumulator instead of +rounding the logits to bf16. The float64 reference uses the same bf16 values, so the +float32 path should match it up to accumulation order, while the default path is off +by the bf16 rounding of each logit. +""" + +import pytest +import torch + +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +VOCAB_SIZE = 32000 +HIDDEN_SIZE = 1024 +DTYPE = torch.bfloat16 + + +def _make_inputs(num_tokens: int): + torch.manual_seed(0) + hidden = torch.randn(num_tokens, HIDDEN_SIZE, dtype=DTYPE) + # Logits with a std of about 4, where bf16 rounding is on the order of 1e-2. + weight = (torch.randn(VOCAB_SIZE, HIDDEN_SIZE) * 4 / HIDDEN_SIZE**0.5).to(DTYPE) + return hidden.cuda(), weight + + +def _make_lm_head(weight: torch.Tensor, output_dtype, **kwargs) -> LMHead: + lm_head = LMHead(VOCAB_SIZE, HIDDEN_SIZE, dtype=DTYPE, output_dtype=output_dtype, **kwargs) + lm_head.load_weights([{"weight": weight}]) + return lm_head.cuda() + + +def _reference_logits(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return hidden.double() @ weight.cuda().double().t() + + +@pytest.mark.parametrize( + "tp_mode", [None, TensorParallelMode.COLUMN], ids=["no_tp", "column_gather"] +) +@pytest.mark.parametrize("num_tokens", [0, 1, 7, 128]) +def test_float32_output_matches_reference(num_tokens: int, tp_mode): + hidden, weight = _make_inputs(num_tokens) + kwargs = dict( + mapping=Mapping(), tensor_parallel_mode=tp_mode, gather_output=True, reduce_output=False + ) + logits_bf16 = _make_lm_head(weight, None, **kwargs)(hidden) + logits_fp32 = _make_lm_head(weight, torch.float32, **kwargs)(hidden) + reference = _reference_logits(hidden, weight) + + assert logits_bf16.dtype == DTYPE + assert logits_fp32.dtype == torch.float32 + assert logits_fp32.shape == logits_bf16.shape == reference.shape + if num_tokens == 0: + return + + torch.testing.assert_close(logits_fp32.double(), reference, rtol=1e-4, atol=1e-3) + # The logits are no longer bf16 values ... + assert not torch.equal(logits_fp32, logits_fp32.to(DTYPE).float()) + # ... and the error left is far below the bf16 rounding error, for logits and logprobs. + err_bf16 = (logits_bf16.double() - reference).abs().max() + err_fp32 = (logits_fp32.double() - reference).abs().max() + assert err_fp32 * 50 < err_bf16 + logprobs_ref = torch.log_softmax(reference, dim=-1) + lp_err_bf16 = (torch.log_softmax(logits_bf16.double(), dim=-1) - logprobs_ref).abs().max() + lp_err_fp32 = (torch.log_softmax(logits_fp32.double(), dim=-1) - logprobs_ref).abs().max() + assert lp_err_fp32 * 10 < lp_err_bf16 + + +def test_float32_output_one_dim_input(): + # LogitsProcessor passes hidden_states[-1] (1-D) when there is no attention metadata. + hidden, weight = _make_inputs(1) + logits = _make_lm_head(weight, torch.float32)(hidden[0]) + assert logits.shape == (VOCAB_SIZE,) + assert logits.dtype == torch.float32 + torch.testing.assert_close( + logits.double(), _reference_logits(hidden, weight)[0], rtol=1e-4, atol=1e-3 + ) + + +def test_float32_output_spec_decoding_head_slice(): + # With lm_head TP in attention DP, the spec-decoding head slices the vocab at forward time. + hidden, weight = _make_inputs(16) + mapping = Mapping(enable_attention_dp=True, enable_lm_head_tp_in_adp=True) + lm_head = _make_lm_head( + weight, + torch.float32, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + reduce_output=False, + ) + lm_head_tp = Mapping(world_size=2, tp_size=2, rank=1) + logits = lm_head(hidden, mapping_lm_head_tp=lm_head_tp, is_spec_decoding_head=True) + reference = _reference_logits(hidden, weight)[:, VOCAB_SIZE // 2 :] + assert logits.dtype == torch.float32 + torch.testing.assert_close(logits.double(), reference, rtol=1e-4, atol=1e-3) + + +def test_output_dtype_rejects_quantized_head(): + with pytest.raises(NotImplementedError, match="unquantized"): + LMHead( + VOCAB_SIZE, + HIDDEN_SIZE, + dtype=DTYPE, + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8), + output_dtype=torch.float32, + ) + + +def test_output_dtype_rejects_row_parallel(): + with pytest.raises(NotImplementedError, match="ROW"): + LMHead( + VOCAB_SIZE, + HIDDEN_SIZE, + dtype=DTYPE, + tensor_parallel_mode=TensorParallelMode.ROW, + output_dtype=torch.float32, + ) + + +def test_llm_args_lm_head_dtype(tmp_path): + assert TorchLlmArgs(model=str(tmp_path)).lm_head_dtype == "auto" + assert TorchLlmArgs(model=str(tmp_path), lm_head_dtype="float32").lm_head_dtype == "float32" + with pytest.raises(ValueError): + TorchLlmArgs(model=str(tmp_path), lm_head_dtype="float16") diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 447c0a1e329f..6f34ef605bc3 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -171,6 +171,10 @@ methods: annotation: bool default: False status: beta + lm_head_dtype: + annotation: Literal['auto', 'float32'] + default: auto + status: prototype moe_config: annotation: tensorrt_llm.llmapi.llm_args.MoeConfig status: beta diff --git a/tests/unittest/llmapi/test_lm_head_dtype.py b/tests/unittest/llmapi/test_lm_head_dtype.py new file mode 100644 index 000000000000..03657527449e --- /dev/null +++ b/tests/unittest/llmapi/test_lm_head_dtype.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end test for LlmArgs.lm_head_dtype. + +With the default LM head, logits are rounded to the model dtype before the float32 +upcast, so every returned logit is exactly representable in the model dtype. With +lm_head_dtype="float32" the LM head GEMM writes its float32 accumulator, so the +returned logits are not rounded, and greedy decoding still follows them. +""" + +import json +import os + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM +from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.sampling_params import SamplingParams + +_MODEL = os.path.join("llama-models-v2", "TinyLlama-1.1B-Chat-v1.0") + + +def _model_dir() -> str: + root = llm_models_root() + path = os.path.join(root, _MODEL) if root else None + if not path or not os.path.isdir(path): + pytest.skip(f"{_MODEL} not available under llm_models_root()") + return path + + +def _rounded_to(logits: torch.Tensor, dtype: torch.dtype) -> bool: + return torch.equal(logits, logits.to(dtype).float()) + + +@pytest.mark.parametrize("lm_head_dtype", ["auto", "float32"]) +def test_lm_head_dtype_returned_logits(lm_head_dtype: str): + model_dir = _model_dir() + with open(os.path.join(model_dir, "config.json")) as f: + model_dtype = getattr(torch, json.load(f)["torch_dtype"]) + + sampling_params = SamplingParams( + max_tokens=8, + temperature=0.0, + return_context_logits=True, + return_generation_logits=True, + ) + with LLM( + model=model_dir, + lm_head_dtype=lm_head_dtype, + max_batch_size=8, + # Context logits are not recomputed for reused KV cache blocks. + kv_cache_config=KvCacheConfig(max_tokens=10000, enable_block_reuse=False), + ) as llm: + output = llm.generate(["The capital of France is"], sampling_params)[0] + + completion = output.outputs[0] + context_logits = output.context_logits + generation_logits = completion.generation_logits + assert context_logits.dtype == torch.float32 + assert generation_logits.dtype == torch.float32 + assert context_logits.shape[0] == len(output.prompt_token_ids) + assert generation_logits.shape[0] == len(completion.token_ids) + + rounded = lm_head_dtype == "auto" + assert _rounded_to(context_logits, model_dtype) == rounded + assert _rounded_to(generation_logits, model_dtype) == rounded + assert torch.argmax(generation_logits, dim=-1).tolist() == completion.token_ids From 55aa431fc6acd461a7bf420c5d394bbfa05842f0 Mon Sep 17 00:00:00 2001 From: Shiki Wu Date: Mon, 12 Oct 2026 23:35:58 -0700 Subject: [PATCH 2/5] [None][fix] Match loader dtype resolution in LM head test Signed-off-by: Shiki Wu --- tests/unittest/llmapi/test_lm_head_dtype.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unittest/llmapi/test_lm_head_dtype.py b/tests/unittest/llmapi/test_lm_head_dtype.py index d84694bc8da9..e63b647d6ab0 100644 --- a/tests/unittest/llmapi/test_lm_head_dtype.py +++ b/tests/unittest/llmapi/test_lm_head_dtype.py @@ -22,12 +22,14 @@ import json import os +from types import SimpleNamespace import pytest import torch from utils.llm_data import llm_models_root from tensorrt_llm import LLM +from tensorrt_llm._torch.pyexecutor.config_utils import resolve_hf_torch_dtype from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.sampling_params import SamplingParams @@ -53,7 +55,8 @@ def test_lm_head_output_dtype_returned_logits(lm_head_output_dtype: str) -> None """Check context and decode precision through the public LLM API.""" model_dir = _model_dir() with open(os.path.join(model_dir, "config.json")) as f: - model_dtype = getattr(torch, json.load(f)["torch_dtype"]) + model_config = SimpleNamespace(**json.load(f)) + model_dtype = resolve_hf_torch_dtype(model_config) or torch.bfloat16 sampling_params = SamplingParams( max_tokens=8, From 2d15d15945a2aab9fe28e29694cccb59aa5ce93b Mon Sep 17 00:00:00 2001 From: Shiki Wu Date: Wed, 14 Oct 2026 22:47:44 -0700 Subject: [PATCH 3/5] [None][refactor] Unify speculative LM head matmul paths Signed-off-by: Shiki Wu --- tensorrt_llm/_torch/modules/embedding.py | 20 +++++++++---------- .../modules/test_lm_head_output_dtype.py | 20 +++++++++++++------ 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index cbdac020a461..e527f8aab629 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/tensorrt_llm/_torch/modules/embedding.py @@ -18,17 +18,19 @@ def _matmul_with_output_dtype(input: torch.Tensor, weight: torch.Tensor, - output_dtype: torch.dtype) -> torch.Tensor: - """Compute ``input @ weight.T`` and write the result in ``output_dtype``. + output_dtype: torch.dtype | None) -> torch.Tensor: + """Compute ``input @ weight.T`` with an optional output dtype override. - ``input`` and ``weight`` stay in the model dtype; the GEMM accumulates in - float32 and stores the result directly in ``output_dtype``, so the output - is not rounded to the model dtype first. + None keeps the model dtype. An FP32 override stores the result directly + in FP32 without first rounding it to the BF16/FP16 model dtype. """ out_features = weight.shape[0] output_shape = (*input.shape[:-1], out_features) input_2d = input.reshape(-1, input.shape[-1]) - output = torch.mm(input_2d, weight.t(), out_dtype=output_dtype) + mm_kwargs = {} + if output_dtype is not None and output_dtype != weight.dtype: + mm_kwargs["out_dtype"] = output_dtype + output = torch.mm(input_2d, weight.t(), **mm_kwargs) return output.view(output_shape) @@ -182,11 +184,7 @@ def forward( slice_start = tp_rank * slice_width slice_end = min((tp_rank + 1) * slice_width, self.out_features) weight = self.weight[slice_start:slice_end, :] - if self._writes_output_dtype(): - output = _matmul_with_output_dtype(input, weight, - self.output_dtype) - else: - output = F.linear(input, weight, None) + output = _matmul_with_output_dtype(input, weight, self.output_dtype) else: output = super().forward(input, all_reduce_params=all_reduce_params) if (self.tp_mode == TensorParallelMode.COLUMN and self.gather_output diff --git a/tests/unittest/_torch/modules/test_lm_head_output_dtype.py b/tests/unittest/_torch/modules/test_lm_head_output_dtype.py index c53b78957c85..16fefac631ff 100644 --- a/tests/unittest/_torch/modules/test_lm_head_output_dtype.py +++ b/tests/unittest/_torch/modules/test_lm_head_output_dtype.py @@ -110,14 +110,18 @@ def test_float32_output_one_dim_input() -> None: ) -def test_float32_output_spec_decoding_head_slice() -> None: - """Keep FP32 output precision when the speculative head slices the vocabulary.""" +@pytest.mark.parametrize("output_dtype", [None, DTYPE, torch.float32], ids=["auto", "bf16", "fp32"]) +@pytest.mark.parametrize("num_tokens", [0, 1, 16]) +def test_output_dtype_spec_decoding_head_slice( + num_tokens: int, output_dtype: torch.dtype | None +) -> None: + """Preserve default and FP32 output precision for a sliced speculative head.""" # With lm_head TP in attention DP, the spec-decoding head slices the vocab at forward time. - hidden, weight = _make_inputs(16) + hidden, weight = _make_inputs(num_tokens) mapping = Mapping(enable_attention_dp=True, enable_lm_head_tp_in_adp=True) lm_head = _make_lm_head( weight, - torch.float32, + output_dtype, mapping=mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, gather_output=True, @@ -126,8 +130,12 @@ def test_float32_output_spec_decoding_head_slice() -> None: lm_head_tp = Mapping(world_size=2, tp_size=2, rank=1) logits = lm_head(hidden, mapping_lm_head_tp=lm_head_tp, is_spec_decoding_head=True) reference = _reference_logits(hidden, weight)[:, VOCAB_SIZE // 2 :] - assert logits.dtype == torch.float32 - torch.testing.assert_close(logits.double(), reference, rtol=1e-4, atol=1e-3) + assert logits.dtype == (output_dtype or DTYPE) + assert logits.shape == reference.shape + if output_dtype == torch.float32: + torch.testing.assert_close(logits.double(), reference, rtol=1e-4, atol=1e-3) + else: + torch.testing.assert_close(logits, reference.to(DTYPE)) def test_output_dtype_rejects_quantized_head() -> None: From 5c1f3dc50bd474634c8cbc0ce98a79406a6201ce Mon Sep 17 00:00:00 2001 From: Shiki Wu Date: Mon, 19 Oct 2026 02:17:54 -0700 Subject: [PATCH 4/5] [None][fix] Preserve LM head dtype annotation in API docs Signed-off-by: Shiki Wu --- tensorrt_llm/llmapi/llm_args.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index fef0c6c71b50..5515ae8e0a3f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6289,7 +6289,8 @@ def init_multimodal_config(cls, v): "rounding of logits before logprob computation. Weights and inputs " "remain in the model dtype. Only supported for unquantized LM heads " "of decoder-only models.", - status="prototype") + status="prototype", + json_schema_extra={"type": "Literal['auto', 'float32']"}) moe_config: MoeConfig = Field(default_factory=MoeConfig, description="MoE config.", From b17ce09faadf73fc3b7abeac016a79ed4842d583 Mon Sep 17 00:00:00 2001 From: Shiki Wu Date: Mon, 19 Oct 2026 08:35:14 -0700 Subject: [PATCH 5/5] [None][test] Sync upstream BOLT subprocess timeout fix Signed-off-by: Shiki Wu --- tests/unittest/scripts/test_bolt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/scripts/test_bolt.py b/tests/unittest/scripts/test_bolt.py index e53d8f7ab2d8..e16d16a24480 100644 --- a/tests/unittest/scripts/test_bolt.py +++ b/tests/unittest/scripts/test_bolt.py @@ -120,7 +120,7 @@ def _compile_so(path: Path, emit_relocs: bool) -> bool: cmd = ["gcc", "-shared", "-fPIC", "-o", str(path), str(src)] if emit_relocs: cmd.insert(3, "-Wl,--emit-relocs") - return subprocess.run(cmd, capture_output=True).returncode == 0 + return subprocess.run(cmd, capture_output=True, timeout=60).returncode == 0 needs_toolchain = pytest.mark.skipif(