[Model] Honor fp32 head_dtype in Inkling muP logits (target + MTP draft)

tests/models/inkling/test_logits_fp32_head.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inkling's muP logits path must honor an fp32 ``head_dtype``.
+
+Inkling divides the final logits by a muP width multiplier and, for the served
+checkpoint, folds ``1/mup`` into a bespoke ``addmm`` that emits logits in the
+input (bf16) dtype. That fast path silently dropped an fp32 ``head_dtype``
+(``--hf-overrides '{"head_dtype": "float32"}'``), which RL training-inference
+consistency requires -- for both the target model and the MTP draft, whose
+logits drive the rejection-sampling acceptance distribution.
+"""
+
+import types
+
+import torch
+
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    UnquantizedEmbeddingMethod,
+)
+from vllm.models.inkling.nvidia.logits_processor import InklingLogitsProcessor
+from vllm.models.inkling.nvidia.mtp import InklingMTP
+
+MUP = 8.0
+VOCAB = 64
+HIDDEN = 16
+NUM_TOKENS = 4
+
+
+class _FakeLmHead:
+    def __init__(self, weight: torch.Tensor):
+        self.weight = weight
+        self.quant_method = UnquantizedEmbeddingMethod()
+
+
+def _inputs():
+    torch.manual_seed(0)
+    hidden = torch.randn(NUM_TOKENS, HIDDEN, dtype=torch.bfloat16)
+    weight = torch.randn(VOCAB, HIDDEN, dtype=torch.bfloat16)
+    return hidden, weight
+
+
+def _expected_fp32(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    return torch.nn.functional.linear(hidden.float(), weight.float()) * (1.0 / MUP)
+
+
+def test_target_logits_processor_fp32_head(default_vllm_config):
+    lp = InklingLogitsProcessor(VOCAB, logits_mup_width_multiplier=MUP)
+    lp.head_dtype = torch.float32
+    lp._gather_logits = lambda logits: logits
+
+    hidden, weight = _inputs()
+    logits = lp(_FakeLmHead(weight), hidden)
+
+    assert logits.dtype == torch.float32
+    torch.testing.assert_close(logits, _expected_fp32(hidden, weight))
+
+
+def test_target_logits_processor_default_head_stays_bf16(default_vllm_config):
+    # head_dtype unset -> the muP addmm fast path is preserved (bf16 out).
+    lp = InklingLogitsProcessor(VOCAB, logits_mup_width_multiplier=MUP)
+    assert lp.head_dtype is None
+    lp._gather_logits = lambda logits: logits
+
+    hidden, weight = _inputs()
+    logits = lp(_FakeLmHead(weight), hidden)
+
+    assert logits.dtype == torch.bfloat16
+    torch.testing.assert_close(
+        logits.float(), _expected_fp32(hidden, weight), atol=0.5, rtol=0.05
+    )
+
+
+def _fake_mtp(head_dtype: torch.dtype | None) -> InklingMTP:
+    mtp = InklingMTP.__new__(InklingMTP)
+    torch.nn.Module.__init__(mtp)
+    mtp.config = types.SimpleNamespace(logits_mup_width_multiplier=MUP)
+    lp = InklingLogitsProcessor(VOCAB)
+    lp.head_dtype = head_dtype
+    lp._gather_logits = lambda logits: logits
+    mtp.logits_processor = lp
+    mtp._logits_zero = None
+    return mtp
+
+
+def test_mtp_draft_compute_logits_fp32_head(default_vllm_config):
+    hidden, weight = _inputs()
+    mtp = _fake_mtp(torch.float32)
+    mtp.lm_head = _FakeLmHead(weight)
+
+    logits = mtp.compute_logits(hidden)
+
+    assert logits.dtype == torch.float32
+    torch.testing.assert_close(logits, _expected_fp32(hidden, weight))
+
+
+def test_mtp_draft_compute_logits_default_head_stays_bf16(default_vllm_config):
+    hidden, weight = _inputs()
+    mtp = _fake_mtp(None)
+    mtp.lm_head = _FakeLmHead(weight)
+
+    logits = mtp.compute_logits(hidden)
+
+    assert logits.dtype == torch.bfloat16

vllm/models/inkling/nvidia/logits_processor.py

@@ -107,12 +107,23 @@ def _base_forward(
         mup = self.logits_mup_width_multiplier
         if not mup:
             return super().forward(lm_head, hidden_states, embedding_bias)
+        assert self.soft_cap is None
+        assert self.scale == 1.0
+        # A non-model head dtype (e.g. `--hf-overrides '{"head_dtype":
+        # "float32"}'` for RL training-inference consistency) must be honored.
+        # The fused-alpha ``addmm`` fast path below emits logits in
+        # ``hidden_states``' dtype and would silently drop the promotion, so
+        # route through the dtype-aware head projection and fold in the muP
+        # divisor as an elementwise multiply in that dtype instead.
+        if self.head_dtype is not None and self.head_dtype != hidden_states.dtype:
+            logits = self._get_logits(hidden_states, lm_head, embedding_bias)
+            if logits is not None:
+                logits = logits * (1.0 / mup)
+            return logits
         # Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue):
         # no separate elementwise kernel, no bf16 rounding of scaled logits, and
         # no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no
         # soft cap, unit logits scale.
-        assert self.soft_cap is None
-        assert self.scale == 1.0
         w = lm_head.weight
         if self._logits_zero is None:
             self._logits_zero = w.new_zeros(1)

vllm/models/inkling/nvidia/mtp.py

@@ -273,8 +273,18 @@ def compute_logits(
         mup = self.config.logits_mup_width_multiplier
         if not mup:
             return self.logits_processor(self.lm_head, hidden_states)
-        assert self.logits_processor.soft_cap is None
-        assert self.logits_processor.scale == 1.0
+        lp = self.logits_processor
+        assert lp.soft_cap is None
+        assert lp.scale == 1.0
+        # Honor a non-model head dtype (fp32 lm_head for RL training-inference
+        # consistency) the same way the target's InklingLogitsProcessor does:
+        # the fused-alpha addmm fast path would emit bf16 and drop the
+        # promotion, so the draft's logits must not diverge from the target's.
+        if lp.head_dtype is not None and lp.head_dtype != hidden_states.dtype:
+            logits = lp._get_logits(hidden_states, self.lm_head, None)
+            if logits is not None:
+                logits = logits * (1.0 / mup)
+            return logits
         w = self.lm_head.weight
         if self._logits_zero is None:
             self._logits_zero = w.new_zeros(1)
@@ -285,9 +295,9 @@ def compute_logits(
             beta=0.0,
             alpha=1.0 / mup,
         )
-        logits = self.logits_processor._gather_logits(logits)
+        logits = lp._gather_logits(logits)
         if logits is not None:
-            logits = logits[..., : self.logits_processor.org_vocab_size]
+            logits = logits[..., : lp.org_vocab_size]
         return logits
 
     def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: