[hmellor]

Overview

The Transformers modeling backend assumes every causal LM's output head is lm_head, plus logit_scale/final_logit_softcapping from the config. This PR runs the model's own ForCausalLM head instead, generalising #61246 (Ettin / ModernBERT decoder). Depends on #65446.

Claims

• The model's ForCausalLM class (from Transformers or remote code) is built on meta with its base model stubbed to return inputs_embeds and its output embeddings calling vLLM's lm_head, so whatever it applies after the base model (scale, soft cap, dense -> act -> norm transform) runs as written.
• A probe at init sends heads that are only the projection (most models, e.g. Llama, Qwen, Mistral, Gemma) down the existing lm_head path, unchanged from main.
• Fixes the logprobs of models whose scale was ignored on main (Granite, GraniteMoe, HyperCLOVAX logits_scaling; greedy output was unaffected), and lets heads with a transform (Ettin / ModernBERT decoder, Roberta-likes) load.
• Under TP, kept heads run on all-gathered logits, so vocab-wide ops and lm_head LoRA work.
• Heads that cannot run without their base model (e.g. Mamba reads lm_head.weight.dtype) raise at init instead of producing wrong logits.

Validation

Max |logprob difference| vs HF float32 (bf16, 4 prompts x 32 greedy tokens, B200), and throughput with --model-impl transformers:

Model	Native	main	PR TP=1	PR TP=2	Output tok/s main -> PR
ibm-granite/granite-3.3-2b-instruct (scale)	0.051	1.242	0.075	0.109	16,582 -> 17,133
google/gemma-2-2b-it (soft cap)	0.083	0.061	0.061	0.073	15,909 -> 17,226

bf16 noise is ~0.1; throughput is one run each, so the gain is noise rather than a speedup.

vllm bench throughput --model <model> --model-impl transformers --input-len 1024 --output-len 256 --num-prompts 1000 --gpu-memory-utilization 0.6
pytest tests/models/transformers/test_backend.py

Tests: all 58 test_backend.py tests pass on B200 (including 17 output head tests) with the wrapper.py fix from #65446, which test_output_head_transform (Ettin) needs.

TODO before merge: once #66545 (compute_logits_local for batch-sharded sampling) lands, hide compute_logits_local when the kept head needs the full vocab (the runner checks hasattr, so a property raising AttributeError). Elementwise heads (scale, soft cap) can run on the local shard and keep it.

TODO before merge: once #66413 lands (Base.embedding_modules names the input embeddings from get_input_embeddings()), check LoRA adapters that target the model's own output embedding name (e.g. GPT-NeoX embed_out, Roberta-like decoder). This PR maps those checkpoint names to lm_head for weight loading, but CausalMixin.embedding_modules only lists lm_head, so such adapters are probably rejected. The name can come from get_output_embeddings() of the ForCausalLM this PR builds.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hmellor.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hmellor.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork