PR #37870: fix: add qk_layernorm support for Phi models

• Repository: vllm-project/vllm
• Author: gambletan
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/37870

Description (problem / solution / changelog)

Summary

• Fixes #37852 — Phi qk_layernorm was unsupported in vLLM, causing silent correctness issues for Phi checkpoints with config.qk_layernorm=True.
• Adds conditional q_layernorm / k_layernorm (nn.LayerNorm(head_dim)) modules to PhiAttention, applied after QKV projection and split, before rotary embedding — matching the Transformers reference implementation exactly.
• No-op for existing models where qk_layernorm is False (the default).

Transformers reference

In transformers/models/phi/modeling_phi.py, PhiAttention does:

# __init__
self.qk_layernorm = config.qk_layernorm
if self.qk_layernorm:
    self.q_layernorm = nn.LayerNorm(self.head_dim, ...)
    self.k_layernorm = nn.LayerNorm(self.head_dim, ...)

# forward
if self.qk_layernorm:
    query_states = self.q_layernorm(query_states)
    key_states = self.k_layernorm(key_states)

This PR mirrors the same behavior in vLLM, following the existing pattern used by PersimmonAttention in vllm/model_executor/models/persimmon.py.

Changes

In vllm/model_executor/models/phi.py (PhiAttention):

__init__: Read config.qk_layernorm (defaults to False). If True, create self.q_layernorm and self.k_layernorm as nn.LayerNorm(head_size).

forward: After qkv.chunk(3), if qk_layernorm is enabled:

1. Reshape q/k from [seq_len, hidden_size] to [seq_len, num_heads, head_dim]
2. Apply per-head LayerNorm
3. Merge back to [seq_len, hidden_size]
4. Then proceed to rotary embedding as before

Test plan

• Verify with a Phi model checkpoint that has qk_layernorm=True (e.g., compare logits against HF Transformers output)
• Verify existing Phi models (qk_layernorm=False) are unaffected (no new modules created, identical code path)

🤖 Generated with Claude Code

Changed files

• vllm/model_executor/models/phi.py (modified, +13/-0)