@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f2d45f26bd

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

Signed-off-by: Josephasafg <ajgard7@gmail.com> Co-authored-by: minyichen <96753147+cyc00518@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>

Asaf Gardin committed Sep 7, 2026 at 17:04 UTC f2d45f26bd6a2c841ffbd0030ebaefe87c274e58

3 files changed +222 -99

tests/models/language/generation/test_gemma.py

+66

@@ -10,6 +10,11 @@ import torch

10 from vllm.config import VllmConfig

11 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

12 from vllm.model_executor.models import gemma

13 +from vllm.model_executor.models.gemma3n import (

14 + Gemma3nTextModel,

15 + _kv_sharing_weights_mapper,

16 +)

17 +from vllm.model_executor.models.gemma4 import Gemma4Model

19 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:

56 assert torch.equal(model.lm_head.weight[:4], lm_head_weight)

59 +@pytest.mark.cpu_test

60 +def test_gemma4_kv_shared_layer_loads_plain_q_proj() -> None:

61 + """KV-shared layers have q_proj instead of a packed qkv_proj; their

62 + redundant K/V tensors in original checkpoints have no parameter and are

63 + skipped rather than failing the load."""

64 + model = torch.nn.Module()

65 + model.config = SimpleNamespace(num_experts=0)

66 + model.start_layer, model.end_layer = 0, 2

67 + model.layers = torch.nn.ModuleList([torch.nn.Module(), torch.nn.Module()])

68 + for layer in model.layers:

69 + layer.self_attn = torch.nn.Module()

70 + model.layers[0].self_attn.qkv_proj = torch.nn.Linear(2, 6, bias=False)

71 + model.layers[1].self_attn.q_proj = torch.nn.Linear(2, 2, bias=False)

72 + shards: list[str] = []

73 + model.layers[0].self_attn.qkv_proj.weight.weight_loader = (

74 + lambda param, weight, shard_id: shards.append(shard_id)

75 + )

76 +

77 + q_weight = torch.full((2, 2), 2.0)

78 + weights = [

79 + (f"layers.{i}.self_attn.{tensor}.weight", q_weight)

80 + for i in (0, 1)

81 + for tensor in ("q_proj", "k_proj", "v_proj")

82 + ] + [("layers.1.self_attn.k_norm.weight", torch.ones(2))]

83 + loaded = Gemma4Model.load_weights(cast(Gemma4Model, model), weights)

84 +

85 + assert shards == ["q", "k", "v"]

86 + assert "layers.1.self_attn.q_proj.weight" in loaded

87 + assert not any(name.startswith("layers.1.self_attn.k") for name in loaded)

88 + assert not any(name.startswith("layers.1.self_attn.v") for name in loaded)

89 + assert torch.equal(model.layers[1].self_attn.q_proj.weight, q_weight)

90 +

91 +

92 +@pytest.mark.cpu_test

93 +def test_gemma3n_kv_shared_layer_mapper() -> None:

94 + """Only non-shared layers pack q/k/v into qkv_proj; KV-shared layers keep

95 + q_proj and drop the redundant K/V tensors original checkpoints ship."""

96 + config = SimpleNamespace(num_hidden_layers=4, num_kv_shared_layers=2)

97 + mapper = Gemma3nTextModel.hf_to_vllm_mapper | _kv_sharing_weights_mapper(config)

98 + weights = [

99 + (f"layers.{i}.self_attn.{tensor}.weight", torch.empty(0))

100 + for i in (1, 3)

101 + for tensor in ("q_proj", "k_proj", "v_proj", "k_norm", "o_proj")

102 + ]

103 +

104 + mapped = [

105 + (name, getattr(weight, "shard_id", None))

106 + for name, weight in mapper.apply(weights)

107 + ]

108 +

109 + assert mapped == [

110 + ("layers.1.self_attn.qkv_proj.weight", "q"),

111 + ("layers.1.self_attn.qkv_proj.weight", "k"),

112 + ("layers.1.self_attn.qkv_proj.weight", "v"),

113 + ("layers.1.self_attn.k_norm.weight", None),

114 + ("layers.1.self_attn.o_proj.weight", None),

115 + ("layers.3.self_attn.q_proj.weight", None),

116 + ("layers.3.self_attn.o_proj.weight", None),

117 + ]

118 +

119 +

120 @pytest.mark.parametrize("model", MODELS)

121 def test_dummy_loader(vllm_runner, monkeypatch, model: str) -> None:

122 with monkeypatch.context() as m:

vllm/model_executor/models/gemma3n.py

+79 -37

@@ -307,15 +307,36 @@ class Gemma3nAttention(nn.Module):

307 self.q_size = self.num_heads * self.head_dim

308 self.kv_size = self.num_kv_heads * self.head_dim

309

310 - self.qkv_proj = QKVParallelLinear(

311 - hidden_size,

312 - self.head_dim,

313 - self.total_num_heads,

314 - self.total_num_kv_heads,

315 - bias=config.attention_bias,

316 - quant_config=quant_config,

317 - prefix=f"{prefix}.qkv_proj",

310 + layer_idx = extract_layer_index(prefix)

311 + first_kv_shared_layer_idx = (

312 + config.num_hidden_layers - config.num_kv_shared_layers

313 )

314 + self.is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx

315 +

316 + if self.is_kv_shared_layer:

317 + # K/V come from the target layer's cache, so like HF this layer

318 + # only has q_proj (and no k_norm/v_norm).

319 + self.q_proj = ColumnParallelLinear(

320 + hidden_size,

321 + self.total_num_heads * self.head_dim,

322 + bias=config.attention_bias,

323 + quant_config=quant_config,

324 + prefix=f"{prefix}.q_proj",

325 + )

326 + else:

327 + self.qkv_proj = QKVParallelLinear(

328 + hidden_size,

329 + self.head_dim,

330 + self.total_num_heads,

331 + self.total_num_kv_heads,

332 + bias=config.attention_bias,

333 + quant_config=quant_config,

334 + prefix=f"{prefix}.qkv_proj",

335 + )

336 + self.k_norm = RMSNorm(hidden_size=self.head_dim, eps=config.rms_norm_eps)

337 + self.v_norm = RMSNorm(

338 + hidden_size=self.head_dim, eps=config.rms_norm_eps, has_weight=False

339 + )

340 self.o_proj = RowParallelLinear(

341 self.total_num_heads * self.head_dim,

342 hidden_size,

@@ -324,12 +345,7 @@ class Gemma3nAttention(nn.Module):

345 prefix=f"{prefix}.o_proj",

346 )

347 self.q_norm = RMSNorm(hidden_size=self.head_dim, eps=config.rms_norm_eps)

327 - self.k_norm = RMSNorm(hidden_size=self.head_dim, eps=config.rms_norm_eps)

328 - self.v_norm = RMSNorm(

329 - hidden_size=self.head_dim, eps=config.rms_norm_eps, has_weight=False

330 - )

348

332 - layer_idx = extract_layer_index(prefix)

349 layer_type = config.layer_types[layer_idx]

350 is_sliding = layer_type == "sliding_attention"

351 self.sliding_window = config.sliding_window if is_sliding else None

@@ -346,13 +362,8 @@ class Gemma3nAttention(nn.Module):

362 if is_sliding:

363 rope_parameters["rope_theta"] = config.rope_local_base_freq

364

349 - first_kv_shared_layer_idx = (

350 - config.num_hidden_layers - config.num_kv_shared_layers

351 - )

352 - self.is_kv_shared = layer_idx >= first_kv_shared_layer_idx

353 -

365 kv_sharing_target_layer_name = None

355 - if self.is_kv_shared:

366 + if self.is_kv_shared_layer:

367 # Last full attention layer is 1 before sharing

368 # Last sliding attention layer is 2 before sharing

369 offset = 2 if self.sliding_window is not None else 1

@@ -405,21 +416,30 @@ class Gemma3nAttention(nn.Module):

416 hidden_states: torch.Tensor,

417 **kwargs,

418 ) -> torch.Tensor:

408 - qkv, _ = self.qkv_proj(hidden_states)

409 - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

410 -

411 - q = q.unflatten(-1, (self.num_heads, self.head_dim))

412 - q = self.q_norm(q)

413 - q = q.flatten(-2, -1)

414 - k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))

415 - k = self.k_norm(k)

416 - k = k.flatten(-2, -1)

417 - v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))

418 - v = self.v_norm(v)

419 - v = v.flatten(-2, -1)

420 -

421 - q, k = self.rotary_emb(positions, q, k)

422 - attn_output = self.attn(q, k, v)

419 + if self.is_kv_shared_layer:

420 + # Shared KV: only Q is projected; K/V come from the target layer.

421 + q, _ = self.q_proj(hidden_states)

422 + q = q.unflatten(-1, (self.num_heads, self.head_dim))

423 + q = self.q_norm(q)

424 + q = q.flatten(-2, -1)

425 + q, _ = self.rotary_emb(positions, q, None)

426 + attn_output = self.attn(q, None, None)

427 + else:

428 + qkv, _ = self.qkv_proj(hidden_states)

429 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

430 +

431 + q = q.unflatten(-1, (self.num_heads, self.head_dim))

432 + q = self.q_norm(q)

433 + q = q.flatten(-2, -1)

434 + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))

435 + k = self.k_norm(k)

436 + k = k.flatten(-2, -1)

437 + v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))

438 + v = self.v_norm(v)

439 + v = v.flatten(-2, -1)

440 +

441 + q, k = self.rotary_emb(positions, q, k)

442 + attn_output = self.attn(q, k, v)

443

444 output, _ = self.o_proj(attn_output)

445 return output

@@ -795,6 +815,28 @@ class Gemma3nCrossDecoder(nn.Module):

815 return hidden_states

816

817

818 +def _kv_sharing_weights_mapper(config: Gemma3nTextConfig) -> WeightsMapper:

819 + """KV-shared layers only have q_proj, so qkv_proj packing applies to the

820 + other layers. Original checkpoints still ship K/V tensors for the shared

821 + layers (fine-tuned ones omit them); those are dropped."""

822 + first_kv_shared_layer_idx = config.num_hidden_layers - config.num_kv_shared_layers

823 + return WeightsMapper(

824 + orig_to_new_substr={

825 + f"layers.{i}.self_attn.{name}.": None

826 + for i in range(first_kv_shared_layer_idx, config.num_hidden_layers)

827 + for name in ("k_proj", "v_proj", "k_norm")

828 + },

829 + orig_to_new_stacked={

830 + f"layers.{i}.self_attn.{shard}_proj.": (

831 + f"layers.{i}.self_attn.qkv_proj.",

832 + shard,

833 + )

834 + for i in range(first_kv_shared_layer_idx)

835 + for shard in ("q", "k", "v")

836 + },

837 + )

838 +

839 +

840 # This disables torch.compile if --kv-sharing-fast-prefill passed

841 @support_torch_compile(

842 enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill

@@ -811,9 +853,6 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

853 "altup_projections.": "self_decoder.altup_projections.",

854 },

855 orig_to_new_stacked={

814 - ".q_proj": (".qkv_proj", "q"),

815 - ".k_proj": (".qkv_proj", "k"),

816 - ".v_proj": (".qkv_proj", "v"),

856 ".gate_proj": (".gate_up_proj", 0),

857 ".up_proj": (".gate_up_proj", 1),

858 },

@@ -826,6 +865,9 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

865 quant_config = vllm_config.quant_config

866 self.config = config

867 self.quant_config = quant_config

868 + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | _kv_sharing_weights_mapper(

869 + config

870 + )

871

872 self.altup_unembed_projections = nn.ModuleList(

873 [

vllm/model_executor/models/gemma4.py

+77 -62

@@ -413,18 +413,41 @@ class Gemma4Attention(nn.Module):

413 # Q/K norms with learnable weights handle scaling implicitly.

414 self.scaling = 1.0

415

416 - # QKVParallelLinear handles GQA correctly for all layer types.

417 - # k_eq_v layers load K weights into both K and V slots via

418 - # _weight_iterator remapping — no structural difference needed.

419 - self.qkv_proj = QKVParallelLinear(

420 - hidden_size,

421 - self.head_dim,

422 - self.total_num_heads,

423 - self.total_num_kv_heads,

424 - bias=config.attention_bias,

425 - quant_config=quant_config,

426 - prefix=f"{prefix}.qkv_proj",

416 + layer_idx = extract_layer_index(prefix)

417 + num_kv_shared_layers = getattr(config, "num_kv_shared_layers", 0)

418 + first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared_layers

419 + self.is_kv_shared_layer = (

420 + num_kv_shared_layers > 0 and layer_idx >= first_kv_shared_layer_idx

421 )

422 +

423 + if self.is_kv_shared_layer:

424 + # K/V come from the target layer's cache, so like HF this layer

425 + # only has q_proj (and no k_norm/v_norm).

426 + self.q_proj = ColumnParallelLinear(

427 + hidden_size,

428 + self.total_num_heads * self.head_dim,

429 + bias=config.attention_bias,

430 + quant_config=quant_config,

431 + prefix=f"{prefix}.q_proj",

432 + )

433 + else:

434 + # QKVParallelLinear handles GQA correctly for all layer types.

435 + # k_eq_v layers load K weights into both K and V slots via

436 + # _weight_iterator remapping — no structural difference needed.

437 + self.qkv_proj = QKVParallelLinear(

438 + hidden_size,

439 + self.head_dim,

440 + self.total_num_heads,

441 + self.total_num_kv_heads,

442 + bias=config.attention_bias,

443 + quant_config=quant_config,

444 + prefix=f"{prefix}.qkv_proj",

445 + )

446 + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

447 + # V norm: no learnable scale (pure normalization only)

448 + self.v_norm = RMSNorm(

449 + self.head_dim, eps=config.rms_norm_eps, has_weight=False

450 + )

451 self.o_proj = RowParallelLinear(

452 self.total_num_heads * self.head_dim,

453 hidden_size,

@@ -433,14 +456,10 @@ class Gemma4Attention(nn.Module):

456 prefix=f"{prefix}.o_proj",

457 )

458

436 - # Q/K norms: output = norm(x) * weight (learnable per-head scale)

459 + # Q norm: output = norm(x) * weight (learnable per-head scale)

460 self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

438 - self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

439 - # V norm: no learnable scale (pure normalization only)

440 - self.v_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False)

461

462 # Determine layer type and sliding window

443 - layer_idx = extract_layer_index(prefix)

463 layer_type = config.layer_types[layer_idx]

464 self.is_sliding = layer_type == "sliding_attention"

465 sliding_window = config.sliding_window if self.is_sliding else None

@@ -466,30 +485,25 @@ class Gemma4Attention(nn.Module):

485 # KV sharing: layers in the last `num_kv_shared_layers` share KV

486 # cache with earlier layers of the same type.

487 kv_sharing_target_layer_name = None

469 - self.is_kv_shared_layer = False

470 - num_kv_shared_layers = getattr(config, "num_kv_shared_layers", 0)

471 - if num_kv_shared_layers > 0:

472 - first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared_layers

473 - if layer_idx >= first_kv_shared_layer_idx:

474 - self.is_kv_shared_layer = True

475 - # Find the last non-shared layer of the same attention type

476 - prev_layers = config.layer_types[:first_kv_shared_layer_idx]

477 - current_layer_type = config.layer_types[layer_idx]

478 - kv_shared_layer_index = (

479 - len(prev_layers) - 1 - prev_layers[::-1].index(current_layer_type)

480 - )

481 - if kv_shared_layer_index >= 0:

482 - if ".layers." in prefix:

483 - param_name_before_layers = prefix.split(".layers.")[0]

484 - else:

485 - raise ValueError(

486 - "Unexpected prefix format for Gemma4Attention: "

487 - f"'{prefix}'. Expected to contain '.layers.'."

488 - )

489 - kv_sharing_target_layer_name = (

490 - f"{param_name_before_layers}.layers."

491 - f"{kv_shared_layer_index}.self_attn.attn"

488 + if self.is_kv_shared_layer:

489 + # Find the last non-shared layer of the same attention type

490 + prev_layers = config.layer_types[:first_kv_shared_layer_idx]

491 + current_layer_type = config.layer_types[layer_idx]

492 + kv_shared_layer_index = (

493 + len(prev_layers) - 1 - prev_layers[::-1].index(current_layer_type)

494 + )

495 + if kv_shared_layer_index >= 0:

496 + if ".layers." in prefix:

497 + param_name_before_layers = prefix.split(".layers.")[0]

498 + else:

499 + raise ValueError(

500 + "Unexpected prefix format for Gemma4Attention: "

501 + f"'{prefix}'. Expected to contain '.layers.'."

502 )

503 + kv_sharing_target_layer_name = (

504 + f"{param_name_before_layers}.layers."

505 + f"{kv_shared_layer_index}.self_attn.attn"

506 + )

507

508 self.rotary_emb = get_rope(

509 self.head_dim,

@@ -524,19 +538,24 @@ class Gemma4Attention(nn.Module):

538 hidden_states: torch.Tensor,

539 **kwargs,

540 ) -> torch.Tensor:

527 - # Unified QKV path (works for both k_eq_v and standard layers).

528 - # For k_eq_v, K weights are loaded into both K and V slots of

529 - # qkv_proj, so V == K automatically.

530 - qkv, _ = self.qkv_proj(hidden_states)

531 - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

532 -

533 - # Q norm (always applied)

534 - q = q.unflatten(-1, (self.num_heads, self.head_dim))

535 - q = self.q_norm(q)

536 - q = q.flatten(-2, -1)

537 -

538 - if not self.is_kv_shared_layer:

539 - # Non-shared: apply K norm + RoPE, V norm

541 + if self.is_kv_shared_layer:

542 + # Shared KV: only Q is projected; K/V come from the target layer.

543 + q, _ = self.q_proj(hidden_states)

544 + q = q.unflatten(-1, (self.num_heads, self.head_dim))

545 + q = self.q_norm(q)

546 + q = q.flatten(-2, -1)

547 + q, _ = self.rotary_emb(positions, q, None)

548 + attn_output = self.attn(q, None, None)

549 + else:

550 + # For k_eq_v, K weights are loaded into both K and V slots of

551 + # qkv_proj, so V == K automatically.

552 + qkv, _ = self.qkv_proj(hidden_states)

553 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

554 +

555 + q = q.unflatten(-1, (self.num_heads, self.head_dim))

556 + q = self.q_norm(q)

557 + q = q.flatten(-2, -1)

558 +

559 k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))

560 k = self.k_norm(k)

561 k = k.flatten(-2, -1)

@@ -545,11 +564,8 @@ class Gemma4Attention(nn.Module):

564 v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))

565 v = self.v_norm(v)

566 v = v.flatten(-2, -1)

548 - else:

549 - # Shared: only apply RoPE to Q

550 - q = self.rotary_emb(positions, q, k)[0]

567 + attn_output = self.attn(q, k, v)

568

552 - attn_output = self.attn(q, k, v)

569 output, _ = self.o_proj(attn_output)

570

571 return output

@@ -1438,9 +1454,9 @@ class Gemma4Model(nn.Module, EagleModelMixin):

1454 if shard_name not in name:

1455 continue

1456 stacked_name = name.replace(shard_name, param_name)

1441 - # k_eq_v layers use separate q_proj/k_proj instead of

1442 - # packed qkv_proj. If the stacked param doesn't exist,

1443 - # skip this mapping and fall through to direct load.

1457 + # KV-shared layers have a plain q_proj instead of a packed

1458 + # qkv_proj: fall through to the direct load, which also

1459 + # skips the redundant K/V tensors original checkpoints ship.

1460 if stacked_name not in params_dict:

1461 continue

1462 if is_pp_missing_parameter(stacked_name, self):

@@ -1534,9 +1550,8 @@ class Gemma4ForCausalLM(

1550 ".moe.experts.down_proj": ".moe.down_proj",

1551 },

1552 )

1537 - # Note: qkv_proj packing applies to non-k_eq_v layers (sliding

1538 - # attention and full attention without k_eq_v). k_eq_v layers use

1539 - # separate q_proj + k_proj without packing.

1553 + # KV-shared layers only have q_proj, so qkv_proj packing applies to

1554 + # the other layers.

1555 packed_modules_mapping = {

1556 "qkv_proj": [

1557 "q_proj",