diff --git a/tests/models/language/generation/test_gemma.py b/tests/models/language/generation/test_gemma.py index 2375fc2ab7..c2c8e88e42 100644 --- a/tests/models/language/generation/test_gemma.py +++ b/tests/models/language/generation/test_gemma.py @@ -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: diff --git a/vllm/model_executor/models/gemma3n.py b/vllm/model_executor/models/gemma3n.py index f981b0707b..56d6473578 100644 --- a/vllm/model_executor/models/gemma3n.py +++ b/vllm/model_executor/models/gemma3n.py @@ -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( [ diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 7f850f4960..961709a7e7 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -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",