[None][feat] Add lm_head_output_dtype for FP32 LM head logits

tensorrt_llm/_torch/model_config.py

@@ -242,6 +242,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_output_dtype).
+    # None keeps the model dtype; torch.float32 writes the float32 accumulator
+    # directly instead of rounding the logits to bf16/fp16 first.
+    lm_head_output_dtype: Optional[torch.dtype] = None
     lora_config: Optional["LoraConfig"] = None
     sparse_attention_config: Optional["SparseAttentionConfig"] = None
     kv_cache_compression_config: Optional["KvCacheCompressionConfig"] = None

tensorrt_llm/_torch/models/modeling_utils.py

@@ -580,6 +580,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_output_dtype,
             )
         else:
             if (hasattr(config, 'lora_config')
@@ -609,6 +610,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_output_dtype,
             )
 
             if self.has_custom_lm_head:

tensorrt_llm/_torch/modules/embedding.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import math
 from typing import Dict, List, Optional, Tuple
 
@@ -14,6 +17,23 @@
 from .linear import Linear, TensorParallelMode
 
 
+def _matmul_with_output_dtype(input: torch.Tensor, weight: torch.Tensor,
+                              output_dtype: torch.dtype | None) -> torch.Tensor:
+    """Compute ``input @ weight.T`` with an optional output dtype override.
+
+    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])
+    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)
+
+
 class LMHead(Linear):
     """LM head layer.
 
@@ -23,6 +43,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 +61,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 +145,22 @@ 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:
+        """Whether the GEMM must store a dtype different from its weights."""
+        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 +183,8 @@ 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, :]
+            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
@@ -151,6 +193,21 @@ def forward(
 
         return output
 
+    def apply_linear(self,
+                     input: torch.Tensor,
+                     bias: Optional[torch.Tensor],
+                     lora_params: Optional[dict] = None,
+                     layer_idx: Optional[int] = None) -> torch.Tensor:
+        """Project hidden states in the configured GEMM output dtype."""
+        # 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,

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -51,6 +51,7 @@
 from ..models.checkpoints.weight_load_plan import WeightLoadPlan
 from ..models.modeling_utils import (DecoderModelForCausalLM, MetaInitMode,
                                      get_registered_model_class, timing_metric)
+from ..modules.embedding import Embedding, LMHead
 from ..modules.low_m_gemm import LOW_M_GEMM_ACTIVE, prepare_low_m_gemm
 from ..moe.fused_moe.moe_load_balancer import (MoeLoadBalancer,
                                                maybe_create_moe_load_balancer)
@@ -296,6 +297,31 @@ def validate_encoder_decoder_tp_scope(model_config: ModelConfig) -> None:
             f"attn_backend='{model_config.attn_backend}'.")
 
 
+def validate_lm_head_output_dtype(model: torch.nn.Module,
+                                  requested_dtype: str) -> None:
+    """Reject an explicit output dtype that the constructed LM heads ignore."""
+    if requested_dtype == "auto":
+        return
+
+    expected_dtype = str_dtype_to_torch(requested_dtype)
+    found_lm_head = False
+    for name, module in model.named_modules():
+        # Embedding inherits LMHead but does not project hidden states to logits.
+        if not isinstance(module, LMHead) or isinstance(module, Embedding):
+            continue
+        found_lm_head = True
+        if module.output_dtype != expected_dtype:
+            raise ValueError(
+                f"lm_head_output_dtype={requested_dtype!r} is not supported "
+                f"by {type(model).__name__}: {name or '<root>'}.output_dtype "
+                f"is {module.output_dtype}, expected {expected_dtype}.")
+
+    if not found_lm_head:
+        raise ValueError(
+            f"lm_head_output_dtype={requested_dtype!r} is not supported by "
+            f"{type(model).__name__}: no LMHead was constructed.")
+
+
 def initialize_dummy_weights(
     model: torch.nn.Module,
     low: float = -1e-3,
@@ -974,6 +1000,9 @@ def load(
                 model = AutoModelForCausalLM.from_config(config)
                 is_meta_init = False
 
+            validate_lm_head_output_dtype(model,
+                                          self.llm_args.lm_head_output_dtype)
+
             requires_standard_hf_loading = any(
                 getattr(module, "_requires_standard_hf_loading", False)
                 for module in model.modules())
@@ -1946,6 +1975,9 @@ 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_output_dtype=(torch.float32
+                                  if self.llm_args.lm_head_output_dtype
+                                  == "float32" else None),
             locality_domain_policy=LocalityDomainPolicy(
                 enabled=self.llm_args.enable_locality_domains),
         )

tensorrt_llm/llmapi/llm_args.py

@@ -6280,6 +6280,18 @@ def init_multimodal_config(cls, v):
         "Separated-routing MoE backends only.",
         status="beta")
 
+    lm_head_output_dtype: Literal["auto", "float32"] = Field(
+        default="auto",
+        description=
+        "Output dtype written by the LM head GEMM. 'auto' uses the model "
+        "dtype before logits are upcast to float32. 'float32' writes the "
+        "float32 accumulator directly, avoiding intermediate bf16/fp16 "
+        "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",
+        json_schema_extra={"type": "Literal['auto', 'float32']"})
+
     moe_config: MoeConfig = Field(default_factory=MoeConfig,
                                   description="MoE config.",
                                   status="beta")

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -866,6 +866,15 @@
       "kind": "categorical",
       "path": "layer_wise_benchmarks_config.calibration_mode"
     },
+    {
+      "allowed_values": [
+        "auto",
+        "float32"
+      ],
+      "capture_policy": "literal",
+      "kind": "categorical",
+      "path": "lm_head_output_dtype"
+    },
     {
       "allowed_values": [
         "auto",
@@ -1379,6 +1388,11 @@
       "kind": "value",
       "path": "sparse_attention_config.fuse_qkv_index_projection"
     },
+    {
+      "capture_policy": "bool",
+      "kind": "value",
+      "path": "sparse_attention_config.gvr_deterministic_topk"
+    },
     {
       "allowed_values": [
         "triton",

tests/integration/test_lists/test-db/l0_h100.yml

@@ -41,6 +41,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
   - unittest/_torch/modules/test_multi_stream_utils.py
   # ------------- MoE components tests ---------------
   - unittest/_torch/moe/test_moe_load_balancer.py
@@ -148,6 +149,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

tests/unittest/_torch/executor/test_model_loader_gms.py

@@ -94,6 +94,7 @@ def _pool_scope(events):
 def _make_loader(monkeypatch, *, events, spec_config=None):
     llm_args = SimpleNamespace(
         load_format=LoadFormat.GMS,
+        lm_head_output_dtype="auto",
         gms_config=SimpleNamespace(socket_path="/tmp/gms.sock", mode="auto", tag="weights"),
     )
     loader = ModelLoader(

tests/unittest/_torch/executor/test_model_loader_mx.py

@@ -547,7 +547,7 @@ def _tiny_profile_registry(*, speculative_mode: str | None = None) -> PostTransf
 
 
 def _make_loader(monkeypatch, *, events, spec_config=None):
-    llm_args = SimpleNamespace(load_format=LoadFormat.AUTO)
+    llm_args = SimpleNamespace(load_format=LoadFormat.AUTO, lm_head_output_dtype="auto")
     loader = ModelLoader(
         llm_args=llm_args,
         mapping=MagicMock(name="mapping"),