iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix][Gemma] Conditionally create KV projections/norms on KV-shared layers (#54917)

Browse Source 
Signed-off-by: Josephasafg <ajgard7@gmail.com>
Co-authored-by: minyichen <96753147+cyc00518@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
This commit is contained in:
Asaf Gardin
2026-09-07 15:04:00 +01:00
committed by GitHub
co-authored by minyichen aoshen02
parent 42801b3a6b
commit f2d45f26bd
3 changed files with 218 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/language/generation/test_gemma.py
+66
View File
@@ -10,6 +10,11 @@ import torch
from vllm.config import VllmConfig
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
from vllm.model_executor.models import gemma
from vllm.model_executor.models.gemma3n import (
Gemma3nTextModel,
_kv_sharing_weights_mapper,
)
from vllm.model_executor.models.gemma4 import Gemma4Model
MODELS = ["google/gemma-2b", "google/gemma-2-2b", "google/gemma-3-4b-it"]
@@ -51,6 +56,67 @@ def test_checkpoint_lm_head_can_override_tied_config(monkeypatch) -> None:
assert torch.equal(model.lm_head.weight[:4], lm_head_weight)
@pytest.mark.cpu_test
def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None:
"""KV-shared layers have q_proj instead of a packed qkv_proj; their
redundant K/V tensors in original checkpoints have no parameter and are
skipped rather than failing the load."""
model = torch.nn.Module()
model.config = SimpleNamespace(num_experts=0)
model.start_layer, model.end_layer = 0, 2
model.layers = torch.nn.ModuleList([torch.nn.Module(), torch.nn.Module()])
for layer in model.layers:
layer.self_attn = torch.nn.Module()
model.layers[0].self_attn.qkv_proj = torch.nn.Linear(2, 6, bias=False)
model.layers[1].self_attn.q_proj = torch.nn.Linear(2, 2, bias=False)
shards: list[str] = []
model.layers[0].self_attn.qkv_proj.weight.weight_loader = (
lambda param, weight, shard_id: shards.append(shard_id)
)
q_weight = torch.full((2, 2), 2.0)
weights = [
(f"layers.{i}.self_attn.{tensor}.weight", q_weight)
for i in (0, 1)
for tensor in ("q_proj", "k_proj", "v_proj")
] + [("layers.1.self_attn.k_norm.weight", torch.ones(2))]
loaded = Gemma4Model.load_weights(cast(Gemma4Model, model), weights)
assert shards == ["q", "k", "v"]
assert "layers.1.self_attn.q_proj.weight" in loaded
assert not any(name.startswith("layers.1.self_attn.k") for name in loaded)
assert not any(name.startswith("layers.1.self_attn.v") for name in loaded)
assert torch.equal(model.layers[1].self_attn.q_proj.weight, q_weight)
@pytest.mark.cpu_test
def test_gemma3n_kv_shared_layer_mapper() -> None:
"""Only non-shared layers pack q/k/v into qkv_proj; KV-shared layers keep
q_proj and drop the redundant K/V tensors original checkpoints ship."""
config = SimpleNamespace(num_hidden_layers=4, num_kv_shared_layers=2)
mapper = Gemma3nTextModel.hf_to_vllm_mapper | _kv_sharing_weights_mapper(config)
weights = [
(f"layers.{i}.self_attn.{tensor}.weight", torch.empty(0))
for i in (1, 3)
for tensor in ("q_proj", "k_proj", "v_proj", "k_norm", "o_proj")
]
mapped = [
(name, getattr(weight, "shard_id", None))
for name, weight in mapper.apply(weights)
]
assert mapped == [
("layers.1.self_attn.qkv_proj.weight", "q"),
("layers.1.self_attn.qkv_proj.weight", "k"),
("layers.1.self_attn.qkv_proj.weight", "v"),
("layers.1.self_attn.k_norm.weight", None),
("layers.1.self_attn.o_proj.weight", None),
("layers.3.self_attn.q_proj.weight", None),
("layers.3.self_attn.o_proj.weight", None),
]
@pytest.mark.parametrize("model", MODELS)
def test_dummy_loader(vllm_runner, monkeypatch, model: str) -> None:
with monkeypatch.context() as m:
vllm/model_executor/models/gemma3n.py
+77 -35
View File
@@ -307,15 +307,36 @@ class Gemma3nAttention(nn.Module):
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.qkv_proj = QKVParallelLinear(
hidden_size,
self.head_dim,
self.total_num_heads,
self.total_num_kv_heads,
bias=config.attention_bias,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
layer_idx = extract_layer_index(prefix)
first_kv_shared_layer_idx = (
config.num_hidden_layers - config.num_kv_shared_layers
)
self.is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx
if self.is_kv_shared_layer:
# K/V come from the target layer's cache, so like HF this layer
# only has q_proj (and no k_norm/v_norm).
self.q_proj = ColumnParallelLinear(
hidden_size,
self.total_num_heads * self.head_dim,
bias=config.attention_bias,
quant_config=quant_config,
prefix=f"{prefix}.q_proj",
)
else:
self.qkv_proj = QKVParallelLinear(
hidden_size,
self.head_dim,
self.total_num_heads,
self.total_num_kv_heads,
bias=config.attention_bias,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
self.k_norm = RMSNorm(hidden_size=self.head_dim, eps=config.rms_norm_eps)
self.v_norm = RMSNorm(
hidden_size=self.head_dim, eps=config.rms_norm_eps, has_weight=False
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
hidden_size,
@@ -324,12 +345,7 @@ class Gemma3nAttention(nn.Module):
prefix=f"{prefix}.o_proj",
)
self.q_norm = RMSNorm(hidden_size=self.head_dim, eps=config.rms_norm_eps)
self.k_norm = RMSNorm(hidden_size=self.head_dim, eps=config.rms_norm_eps)
self.v_norm = RMSNorm(
hidden_size=self.head_dim, eps=config.rms_norm_eps, has_weight=False
)
layer_idx = extract_layer_index(prefix)
layer_type = config.layer_types[layer_idx]
is_sliding = layer_type == "sliding_attention"
self.sliding_window = config.sliding_window if is_sliding else None
@@ -346,13 +362,8 @@ class Gemma3nAttention(nn.Module):
if is_sliding:
rope_parameters["rope_theta"] = config.rope_local_base_freq
first_kv_shared_layer_idx = (
config.num_hidden_layers - config.num_kv_shared_layers
)
self.is_kv_shared = layer_idx >= first_kv_shared_layer_idx
kv_sharing_target_layer_name = None
if self.is_kv_shared:
if self.is_kv_shared_layer:
# Last full attention layer is 1 before sharing
# Last sliding attention layer is 2 before sharing
offset = 2 if self.sliding_window is not None else 1
@@ -405,21 +416,30 @@ class Gemma3nAttention(nn.Module):
hidden_states: torch.Tensor,
**kwargs,
) -> torch.Tensor:
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
if self.is_kv_shared_layer:
# Shared KV: only Q is projected; K/V come from the target layer.
q, _ = self.q_proj(hidden_states)
q = q.unflatten(-1, (self.num_heads, self.head_dim))
q = self.q_norm(q)
q = q.flatten(-2, -1)
q, _ = self.rotary_emb(positions, q, None)
attn_output = self.attn(q, None, None)
else:
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q = q.unflatten(-1, (self.num_heads, self.head_dim))
q = self.q_norm(q)
q = q.flatten(-2, -1)
k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
k = self.k_norm(k)
k = k.flatten(-2, -1)
v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))
v = self.v_norm(v)
v = v.flatten(-2, -1)
q = q.unflatten(-1, (self.num_heads, self.head_dim))
q = self.q_norm(q)
q = q.flatten(-2, -1)
k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
k = self.k_norm(k)
k = k.flatten(-2, -1)
v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))
v = self.v_norm(v)
v = v.flatten(-2, -1)
q, k = self.rotary_emb(positions, q, k)
attn_output = self.attn(q, k, v)
q, k = self.rotary_emb(positions, q, k)
attn_output = self.attn(q, k, v)
output, _ = self.o_proj(attn_output)
return output
@@ -795,6 +815,28 @@ class Gemma3nCrossDecoder(nn.Module):
return hidden_states
def _kv_sharing_weights_mapper(config: Gemma3nTextConfig) -> WeightsMapper:
"""KV-shared layers only have q_proj, so qkv_proj packing applies to the
other layers. Original checkpoints still ship K/V tensors for the shared
layers (fine-tuned ones omit them); those are dropped."""
first_kv_shared_layer_idx = config.num_hidden_layers - config.num_kv_shared_layers
return WeightsMapper(
orig_to_new_substr={
f"layers.{i}.self_attn.{name}.": None
for i in range(first_kv_shared_layer_idx, config.num_hidden_layers)
for name in ("k_proj", "v_proj", "k_norm")
},
orig_to_new_stacked={
f"layers.{i}.self_attn.{shard}_proj.": (
f"layers.{i}.self_attn.qkv_proj.",
shard,
)
for i in range(first_kv_shared_layer_idx)
for shard in ("q", "k", "v")
},
)
# This disables torch.compile if --kv-sharing-fast-prefill passed
@support_torch_compile(
enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill
@@ -811,9 +853,6 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):
"altup_projections.": "self_decoder.altup_projections.",
},
orig_to_new_stacked={
".q_proj": (".qkv_proj", "q"),
".k_proj": (".qkv_proj", "k"),
".v_proj": (".qkv_proj", "v"),
".gate_proj": (".gate_up_proj", 0),
".up_proj": (".gate_up_proj", 1),
},
@@ -826,6 +865,9 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):
quant_config = vllm_config.quant_config
self.config = config
self.quant_config = quant_config
self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | _kv_sharing_weights_mapper(
config
)
self.altup_unembed_projections = nn.ModuleList(
[
vllm/model_executor/models/gemma4.py
+75 -60
View File
@@ -413,18 +413,41 @@ class Gemma4Attention(nn.Module):
# Q/K norms with learnable weights handle scaling implicitly.
self.scaling = 1.0
# QKVParallelLinear handles GQA correctly for all layer types.
# k_eq_v layers load K weights into both K and V slots via
# _weight_iterator remapping — no structural difference needed.
self.qkv_proj = QKVParallelLinear(
hidden_size,
self.head_dim,
self.total_num_heads,
self.total_num_kv_heads,
bias=config.attention_bias,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
layer_idx = extract_layer_index(prefix)
num_kv_shared_layers = getattr(config, "num_kv_shared_layers", 0)
first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared_layers
self.is_kv_shared_layer = (
num_kv_shared_layers > 0 and layer_idx >= first_kv_shared_layer_idx
)
if self.is_kv_shared_layer:
# K/V come from the target layer's cache, so like HF this layer
# only has q_proj (and no k_norm/v_norm).
self.q_proj = ColumnParallelLinear(
hidden_size,
self.total_num_heads * self.head_dim,
bias=config.attention_bias,
quant_config=quant_config,
prefix=f"{prefix}.q_proj",
)
else:
# QKVParallelLinear handles GQA correctly for all layer types.
# k_eq_v layers load K weights into both K and V slots via
# _weight_iterator remapping — no structural difference needed.
self.qkv_proj = QKVParallelLinear(
hidden_size,
self.head_dim,
self.total_num_heads,
self.total_num_kv_heads,
bias=config.attention_bias,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
# V norm: no learnable scale (pure normalization only)
self.v_norm = RMSNorm(
self.head_dim, eps=config.rms_norm_eps, has_weight=False
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
hidden_size,
@@ -433,14 +456,10 @@ class Gemma4Attention(nn.Module):
prefix=f"{prefix}.o_proj",
)
# Q/K norms: output = norm(x) * weight (learnable per-head scale)
# Q norm: output = norm(x) * weight (learnable per-head scale)
self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
# V norm: no learnable scale (pure normalization only)
self.v_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False)
# Determine layer type and sliding window
layer_idx = extract_layer_index(prefix)
layer_type = config.layer_types[layer_idx]
self.is_sliding = layer_type == "sliding_attention"
sliding_window = config.sliding_window if self.is_sliding else None
@@ -466,30 +485,25 @@ class Gemma4Attention(nn.Module):
# KV sharing: layers in the last `num_kv_shared_layers` share KV
# cache with earlier layers of the same type.
kv_sharing_target_layer_name = None
self.is_kv_shared_layer = False
num_kv_shared_layers = getattr(config, "num_kv_shared_layers", 0)
if num_kv_shared_layers > 0:
first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared_layers
if layer_idx >= first_kv_shared_layer_idx:
self.is_kv_shared_layer = True
# Find the last non-shared layer of the same attention type
prev_layers = config.layer_types[:first_kv_shared_layer_idx]
current_layer_type = config.layer_types[layer_idx]
kv_shared_layer_index = (
len(prev_layers) - 1 - prev_layers[::-1].index(current_layer_type)
)
if kv_shared_layer_index >= 0:
if ".layers." in prefix:
param_name_before_layers = prefix.split(".layers.")[0]
else:
raise ValueError(
"Unexpected prefix format for Gemma4Attention: "
f"'{prefix}'. Expected to contain '.layers.'."
)
kv_sharing_target_layer_name = (
f"{param_name_before_layers}.layers."
f"{kv_shared_layer_index}.self_attn.attn"
if self.is_kv_shared_layer:
# Find the last non-shared layer of the same attention type
prev_layers = config.layer_types[:first_kv_shared_layer_idx]
current_layer_type = config.layer_types[layer_idx]
kv_shared_layer_index = (
len(prev_layers) - 1 - prev_layers[::-1].index(current_layer_type)
)
if kv_shared_layer_index >= 0:
if ".layers." in prefix:
param_name_before_layers = prefix.split(".layers.")[0]
else:
raise ValueError(
"Unexpected prefix format for Gemma4Attention: "
f"'{prefix}'. Expected to contain '.layers.'."
)
kv_sharing_target_layer_name = (
f"{param_name_before_layers}.layers."
f"{kv_shared_layer_index}.self_attn.attn"
)
self.rotary_emb = get_rope(
self.head_dim,
@@ -524,19 +538,24 @@ class Gemma4Attention(nn.Module):
hidden_states: torch.Tensor,
**kwargs,
) -> torch.Tensor:
# Unified QKV path (works for both k_eq_v and standard layers).
# For k_eq_v, K weights are loaded into both K and V slots of
# qkv_proj, so V == K automatically.
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
if self.is_kv_shared_layer:
# Shared KV: only Q is projected; K/V come from the target layer.
q, _ = self.q_proj(hidden_states)
q = q.unflatten(-1, (self.num_heads, self.head_dim))
q = self.q_norm(q)
q = q.flatten(-2, -1)
q, _ = self.rotary_emb(positions, q, None)
attn_output = self.attn(q, None, None)
else:
# For k_eq_v, K weights are loaded into both K and V slots of
# qkv_proj, so V == K automatically.
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
# Q norm (always applied)
q = q.unflatten(-1, (self.num_heads, self.head_dim))
q = self.q_norm(q)
q = q.flatten(-2, -1)
q = q.unflatten(-1, (self.num_heads, self.head_dim))
q = self.q_norm(q)
q = q.flatten(-2, -1)
if not self.is_kv_shared_layer:
# Non-shared: apply K norm + RoPE, V norm
k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
k = self.k_norm(k)
k = k.flatten(-2, -1)
@@ -545,11 +564,8 @@ class Gemma4Attention(nn.Module):
v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))
v = self.v_norm(v)
v = v.flatten(-2, -1)
else:
# Shared: only apply RoPE to Q
q = self.rotary_emb(positions, q, k)[0]
attn_output = self.attn(q, k, v)
attn_output = self.attn(q, k, v)
output, _ = self.o_proj(attn_output)
return output
@@ -1438,9 +1454,9 @@ class Gemma4Model(nn.Module, EagleModelMixin):
if shard_name not in name:
continue
stacked_name = name.replace(shard_name, param_name)
# k_eq_v layers use separate q_proj/k_proj instead of
# packed qkv_proj. If the stacked param doesn't exist,
# skip this mapping and fall through to direct load.
# KV-shared layers have a plain q_proj instead of a packed
# qkv_proj: fall through to the direct load, which also
# skips the redundant K/V tensors original checkpoints ship.
if stacked_name not in params_dict:
continue
if is_pp_missing_parameter(stacked_name, self):
@@ -1534,9 +1550,8 @@ class Gemma4ForCausalLM(
".moe.experts.down_proj": ".moe.down_proj",
},
)
# Note: qkv_proj packing applies to non-k_eq_v layers (sliding
# attention and full attention without k_eq_v). k_eq_v layers use
# separate q_proj + k_proj without packing.
# KV-shared layers only have q_proj, so qkv_proj packing applies to
# the other layers.
packed_modules_mapping = {
"qkv_proj": [
"q_proj",