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

[KKothuri]

Purpose

Inkling folds the muP width divisor into a bespoke lm_head torch.addmm that emits logits in the input (bf16) dtype, silently dropping an fp32 head_dtype (--hf-overrides '{"head_dtype": "float32"}') whenever muP is active (the normal served-checkpoint path).

fp32 logits are required for RL training-inference consistency. In speculative decoding the draft's logits drive the rejection-sampling acceptance distribution, so the MTP draft must match the target — otherwise the target runs fp32 while the draft runs bf16.

This fixes both bespoke muP paths:

• InklingLogitsProcessor._base_forward (target)
• InklingMTP.compute_logits (MTP draft)

When an fp32 head_dtype is requested, both now fold the muP divisor directly into the fp32-accumulate lm_head GEMM via torch.addmm(..., alpha=1/mup, out_dtype=float32) — the addmm.dtype overload accumulates bf16 inputs into fp32 without materializing an fp32 weight copy, and scales in the epilogue. This is bit-for-bit identical to the dtype-aware projection (_get_logits → _apply_head) followed by an elementwise multiply, with no extra kernel. That elementwise-multiply projection is kept as the fallback for cases the fused op does not cover (non-CUDA cast path, quantized head, embedding bias). The default (model-dtype) serving path is unchanged, so there is no perf regression there.

Note: the general spec-decode draft path is not affected — a draft using the default LogitsProcessor already inherits the target's head_dtype via the current-config context (verified at runtime; _create_draft_vllm_config keeps the target model_config). Inkling was the sole breakage because it bypasses the default projection.

muP is folded into the fp32 GEMM (not a bypass)

The fp32 branch is not a muP bypass. muP is applied as the alpha scalar of the same fused addmm that computes the logits — out = alpha * (hidden @ w.t()) with alpha = 1/mup — so it costs nothing beyond the projection and mutates no shared state.

An earlier revision applied muP as a separate logits * (1/mup) after the projection, on the assumption that the fp32 accumulate path (torch.mm(out_dtype=fp32)) had no scalar to fold into. That is not the case: the CUDA addmm.dtype overload takes both alpha and out_dtype, and folding through it is bit-identical to the unfused projection + multiply (verified on an RTX 5080; max abs diff 0.0). The separate multiply now only runs in the fallback cases above. Shared can_fold_fp32_head / mup_addmm helpers keep the target LP and the MTP draft in sync.

Not a duplicate

Searched open issues/PRs for Inkling fp32 / head_dtype / mup logits and head_dtype draft speculative — none address this. #48768 (Inkling multi-depth MTP) is unrelated follow-up work.

Testing

tests/models/inkling/test_logits_fp32_head.py (5 tests) — asserts fp32 output for both the target InklingLogitsProcessor and InklingMTP.compute_logits under head_dtype=float32, that the default path stays bf16 (fast addmm preserved), and (CUDA-guarded) that the fp32 fold is bit-for-bit identical to the unfused dtype-aware projection + multiply.

pytest tests/models/inkling/test_logits_fp32_head.py -v
# 5 passed (fold-equivalence test runs on GPU, skipped on CPU-only CI)

Confirmed the fp32 tests fail on unpatched source (bf16 output) and pass after the fix; the default-path tests confirm bf16 is unchanged. pre-commit run (ruff, ruff-format, mypy) is clean.

Model evaluation

The change is a no-op for the default (model-dtype) serving path, so standard serving output is unaffected. The fp32 path is opt-in for RL training-inference consistency. A full Inkling-checkpoint eval was not run (large NVFP4 checkpoint not available locally); the unit tests assert the exact fp32 dtype/value contract at the production method boundary.

AI assistance

This change was prepared with AI assistance. I have reviewed every changed line.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.