From af660901c85086691755f1e716c3e5b9af0382a3 Mon Sep 17 00:00:00 2001
From: santosh-agebold <12378904+shotsan@users.noreply.github.com>
Date: Thu, 9 Jul 2026 10:36:03 -0700
Subject: [PATCH 01/19] feat(data): extract verifier KV caches for Gemma4 MTP

Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
---
 src/speculators/train/data.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index 4e612373b..71da53470 100644
--- a/src/speculators/train/data.py
+++ b/src/speculators/train/data.py
@@ -85,6 +85,8 @@ def create_empty_sample(
         "hidden_states": torch.empty(0, num_target_layers * hidden_size, dtype=dtype),
         "input_ids": torch.empty(0, dtype=torch.long),
         "verifier_last_hidden_states": torch.empty(0, hidden_size, dtype=dtype),
+        "verifier_kv_last_local": torch.empty(0, dtype=dtype),
+        "verifier_kv_last_global": torch.empty(0, dtype=dtype),
         "loss_mask": torch.empty(0, dtype=torch.bool),
         "lengths": torch.tensor([0], dtype=torch.long),
         "position_ids": torch.arange(0, dtype=torch.long),
@@ -104,12 +106,17 @@ def standardize_data_v1(data: dict[str, Any]) -> dict[str, Any]:
     #  ],
     # }
 
-    return {
+    res = {
         "hidden_states": torch.cat(data["hidden_states"][:-1], dim=-1),
         "input_ids": data["input_ids"],
         "verifier_last_hidden_states": data["hidden_states"][-1],
         "loss_mask": data["loss_mask"],
     }
+    if "verifier_kv_last_local" in data:
+        res["verifier_kv_last_local"] = data["verifier_kv_last_local"]
+    if "verifier_kv_last_global" in data:
+        res["verifier_kv_last_global"] = data["verifier_kv_last_global"]
+    return res
 
 
 def _has_multimodal_content(messages: list[dict]) -> bool:
@@ -183,9 +190,11 @@ def __getitem__(self, index) -> BatchType | None:
         #  "loss_mask": [seq_len],
         # }
 
-        # Convert hidden states to the correct dtype
+        # Convert hidden states and KV caches to the correct dtype
         data = {
-            k: v.to(self.hidden_states_dtype) if "hidden_states" in k else v
+            k: v.to(self.hidden_states_dtype)
+            if ("hidden_states" in k or "verifier_kv" in k)
+            else v
             for k, v in data.items()
         }
 
@@ -357,7 +366,7 @@ def _get_raw_data(self, index):
             )
             return None
 
-        return {
+        res = {
             "hidden_states": loaded_hs["hidden_states"][:, :-1].flatten(
                 1
             ),  # [seq_len, 3 * hidden_size]
@@ -367,6 +376,11 @@ def _get_raw_data(self, index):
             ],  # [seq_len, hidden_size]
             "loss_mask": self.data[index]["loss_mask"],  # [seq_len]
         }
+        if "verifier_kv_last_local" in loaded_hs:
+            res["verifier_kv_last_local"] = loaded_hs["verifier_kv_last_local"]
+        if "verifier_kv_last_global" in loaded_hs:
+            res["verifier_kv_last_global"] = loaded_hs["verifier_kv_last_global"]
+        return res
 
 
 class SampleFileDataset(BaseDataset):

From 3d362e0095b04195c616dbea09f25fc1dfb6aeb2 Mon Sep 17 00:00:00 2001
From: santosh-agebold <12378904+shotsan@users.noreply.github.com>
Date: Fri, 10 Jul 2026 10:47:33 -0700
Subject: [PATCH 02/19] test(data): add unit tests for KV caches and fix schema
 docstrings

Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
---
 src/speculators/models/mtp/config.py | 13 +++++++++
 src/speculators/train/data.py        |  6 ++++
 tests/unit/train/test_data.py        | 43 ++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/src/speculators/models/mtp/config.py b/src/speculators/models/mtp/config.py
index 1529827a3..7b5ac5375 100644
--- a/src/speculators/models/mtp/config.py
+++ b/src/speculators/models/mtp/config.py
@@ -52,6 +52,19 @@ class MTPSpeculatorConfig(SpeculatorModelConfig):
         ),
     )
 
+    num_centroids: int | None = Field(
+        default=None,
+        description=(
+            "Number of centroids for a centroid-masked multi-level LM head. "
+            "If None, standard flat LM head is used."
+        ),
+    )
+
+    top_k_centroids: int = Field(
+        default=32,
+        description="Number of top centroids to select during multi-level LM head decoding.",
+    )
+
     @property
     def hidden_size(self) -> int:
         return self.transformer_layer_config.hidden_size  # type: ignore[return-value]
diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index 71da53470..23d2c5155 100644
--- a/src/speculators/train/data.py
+++ b/src/speculators/train/data.py
@@ -72,6 +72,8 @@ def create_empty_sample(
     #     "hidden_states": [seq_len, num_target_layers * hidden_size],
     #     "input_ids": [seq_len],
     #     "verifier_last_hidden_states": [seq_len, hidden_size],
+    #     "verifier_kv_last_local": [seq_len, ...],
+    #     "verifier_kv_last_global": [seq_len, ...],
     #     "loss_mask": [seq_len],
     #     "lengths": [1],
     #     "position_ids": [seq_len],
@@ -187,6 +189,8 @@ def __getitem__(self, index) -> BatchType | None:
         #  "hidden_states": [seq_len, 3 * hidden_size],
         #  "input_ids": [seq_len],
         #  "verifier_last_hidden_states": [seq_len, hidden_size],
+        #  "verifier_kv_last_local": [seq_len, ...],
+        #  "verifier_kv_last_global": [seq_len, ...],
         #  "loss_mask": [seq_len],
         # }
 
@@ -210,6 +214,8 @@ def __getitem__(self, index) -> BatchType | None:
         #     "hidden_states": [seq_len, 3 * hidden_size],
         #     "input_ids": [seq_len],
         #     "verifier_last_hidden_states": [seq_len, hidden_size],
+        #     "verifier_kv_last_local": [seq_len, ...],
+        #     "verifier_kv_last_global": [seq_len, ...],
         #     "loss_mask": [seq_len],
         #     "lengths": [1],
         #     "position_ids": [seq_len],
diff --git a/tests/unit/train/test_data.py b/tests/unit/train/test_data.py
index 54a40e9f3..f3d10c2eb 100644
--- a/tests/unit/train/test_data.py
+++ b/tests/unit/train/test_data.py
@@ -472,6 +472,49 @@ def test_arrow_dataset_default_split_ratio_does_not_crash(tmp_path: Path):
     assert arrow_ds._map_to_file_idx(5) == 5
 
 
+def test_verifier_kvs_survive_pipeline():
+    """Test that optional verifier KV caches survive standardize_data_v1 and collate_fn."""
+    # 1. Create a dummy v1 offline data dictionary with KV caches
+    v1_data = {
+        "input_ids": torch.tensor([0, 1, 2], dtype=torch.long),
+        "loss_mask": torch.tensor([0, 1, 1], dtype=torch.long),
+        "hidden_states": [
+            torch.tensor([[0.0, 0.1], [1.0, 1.1], [2.0, 2.1]]),  # Layer 0
+            torch.tensor([[10.0, 10.1], [11.0, 11.1], [12.0, 12.1]]),  # Verifier last hs
+        ],
+        "verifier_kv_last_local": torch.tensor([[100.0], [101.0], [102.0]]),
+        "verifier_kv_last_global": torch.tensor([[200.0], [201.0], [202.0]]),
+    }
+
+    # 2. Pass through standardize_data_v1
+    standardized = standardize_data_v1(v1_data)
+
+    assert "verifier_kv_last_local" in standardized
+    assert "verifier_kv_last_global" in standardized
+    assert torch.equal(standardized["verifier_kv_last_local"], v1_data["verifier_kv_last_local"])
+
+    # 3. Add lengths and position_ids (simulate BaseDataset.__getitem__)
+    standardized["lengths"] = torch.tensor([3], dtype=torch.long)
+    standardized["position_ids"] = torch.tensor([0, 1, 2], dtype=torch.long)
+
+    # 4. Pass through collate_fn
+    collate_fn = create_collate_fn(max_len=5, hidden_size=2, num_target_layers=1)
+    batch = [standardized]
+
+    collated = collate_fn(batch)
+
+    # 5. Verify the KV caches survived collation and were properly padded
+    assert "verifier_kv_last_local" in collated
+    assert "verifier_kv_last_global" in collated
+
+    local_kv = collated["verifier_kv_last_local"]
+    assert local_kv.shape == (1, 5, 1)  # [batch=1, max_len=5, ...]
+
+    # First 3 positions should match original, last 2 should be padded with 0
+    expected_local = torch.tensor([[[100.0], [101.0], [102.0], [0.0], [0.0]]])
+    assert torch.equal(local_kv, expected_local)
+
+
 def test_arrow_dataset_on_generate_cache_creates_hidden_states_dir(tmp_path: Path):
     """on_generate="cache" must create the cache dir when cache() is called —
     otherwise shutil.move into it raises FileNotFoundError, which _maybe_generate_hs

From e6aafa4fbdab80a2db2d830747b9cc688a5e51b7 Mon Sep 17 00:00:00 2001
From: santosh-agebold <12378904+shotsan@users.noreply.github.com>
Date: Fri, 10 Jul 2026 11:56:38 -0700
Subject: [PATCH 03/19] feat(mtp): add multi-level LM head for Gemma4 MTP

Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
---
 src/speculators/models/mtp/core.py |  21 +++++
 src/speculators/models/mtp/head.py |  39 ++++++++
 tests/unit/models/test_mtp_head.py | 138 +++++++++++++++++++++++++++++
 3 files changed, 198 insertions(+)
 create mode 100644 src/speculators/models/mtp/head.py
 create mode 100644 tests/unit/models/test_mtp_head.py

diff --git a/src/speculators/models/mtp/core.py b/src/speculators/models/mtp/core.py
index a3d3c67e6..d6d217ee1 100644
--- a/src/speculators/models/mtp/core.py
+++ b/src/speculators/models/mtp/core.py
@@ -63,6 +63,8 @@ class MTPDraftModel(DraftVocabMixin, SpeculatorModel):
         "verifier_lm_head.weight",
         "t2d",
         "d2t",
+        "masked_embedding.centroids.weight",
+        "masked_embedding.token_ordering",
     ]
 
     t2d: torch.Tensor | None
@@ -87,6 +89,17 @@ def __init__(self, config: MTPSpeculatorConfig) -> None:
 
         self.post_init()
 
+        if getattr(self.config, "num_centroids", None) is not None:
+            from speculators.models.mtp.head import MultiLevelLMHead
+
+            self.masked_embedding = MultiLevelLMHead(
+                hidden_size=tc.hidden_size,
+                vocab_size=tc.vocab_size,
+                num_centroids=self.config.num_centroids,
+            )
+        else:
+            self.masked_embedding = None
+
     @property
     def layers(self) -> nn.ModuleList:
         """Expose mtp_layers for FSDP wrapping compatibility."""
@@ -208,12 +221,20 @@ def forward(
                 step_targets = step_targets.clone()
                 step_targets[step_mask == 0] = _IGNORE_INDEX
             weight = step_weights[step] if step_weights is not None else 1.0
+
             unreduced = nn.functional.cross_entropy(
                 logits.permute(0, 2, 1),
                 step_targets,
                 ignore_index=_IGNORE_INDEX,
                 reduction="none",
             )
+
+            if getattr(self, "masked_embedding", None) is not None:
+                unreduced += self.masked_embedding.compute_centroid_loss(
+                    hidden_states=mtp_output,
+                    targets=step_targets,
+                    ignore_index=_IGNORE_INDEX,
+                )
             valid_count = (step_targets != _IGNORE_INDEX).sum()
             step_loss = weight * unreduced.sum() / valid_count.clamp(min=1)
             total_loss = total_loss + step_loss
diff --git a/src/speculators/models/mtp/head.py b/src/speculators/models/mtp/head.py
new file mode 100644
index 000000000..d94faf3ca
--- /dev/null
+++ b/src/speculators/models/mtp/head.py
@@ -0,0 +1,39 @@
+"""Multi-level LM head for Gemma4 MTP."""
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class MultiLevelLMHead(nn.Module):
+    """Auxiliary centroid projection for Gemma4 MTP."""
+
+    def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int):
+        super().__init__()
+        self.tokens_per_centroid = vocab_size // num_centroids
+        self.centroids = nn.Linear(hidden_size, num_centroids, bias=False)
+        # Token ordering mapping (centroid subsets -> token IDs)
+        self.register_buffer("token_ordering", torch.arange(vocab_size, dtype=torch.long))
+
+    def compute_centroid_loss(
+        self, hidden_states: torch.Tensor, targets: torch.Tensor, ignore_index: int = -100
+    ) -> torch.Tensor:
+        """Compute the unreduced centroid cross-entropy loss."""
+        valid_mask = targets != ignore_index
+        valid_hidden = hidden_states[valid_mask]
+        valid_targets = targets[valid_mask]
+
+        loss = torch.zeros_like(targets, dtype=hidden_states.dtype)
+        if valid_hidden.shape[0] == 0:
+            return loss
+
+        # Inverse mapping: token_id -> centroid_id
+        inverse = torch.empty_like(self.token_ordering)
+        positions = torch.arange(self.token_ordering.shape[0], device=self.token_ordering.device)
+        inverse[self.token_ordering] = positions // self.tokens_per_centroid
+        
+        centroid_logits = self.centroids(valid_hidden)
+        target_centroids = inverse[valid_targets]
+        loss[valid_mask] = F.cross_entropy(centroid_logits, target_centroids, reduction="none")
+        
+        return loss
diff --git a/tests/unit/models/test_mtp_head.py b/tests/unit/models/test_mtp_head.py
new file mode 100644
index 000000000..cdabac45f
--- /dev/null
+++ b/tests/unit/models/test_mtp_head.py
@@ -0,0 +1,138 @@
+"""Unit tests for MultiLevelLMHead in Gemma4 MTP."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from speculators.models.mtp.head import MultiLevelLMHead
+
+
+@pytest.fixture
+def head_setup():
+    hidden_size = 16
+    vocab_size = 256
+    num_centroids = 4
+    head = MultiLevelLMHead(hidden_size, vocab_size, num_centroids)
+    return head
+
+
+def test_initialization(head_setup):
+    """Test that the head initializes and tokens_per_centroid is computed correctly."""
+    head = head_setup
+    assert head.tokens_per_centroid == 256 // 4
+    assert head.centroids.in_features == 16
+    assert head.centroids.out_features == 4
+    assert head.token_ordering.shape == (256,)
+    
+    # Check that sequential token mapping works
+    assert torch.equal(head.token_ordering, torch.arange(256))
+
+
+def test_compute_centroid_loss_basic(head_setup):
+    """Test basic loss computation without ignore_index."""
+    head = head_setup
+    batch_size, seq_len = 2, 3
+    hidden_states = torch.randn(batch_size, seq_len, 16)
+    
+    # Targets for the first 3 centroids
+    targets = torch.tensor([
+        [0, 64, 128], # belongs to centroids 0, 1, 2
+        [192, 10, 70] # belongs to centroids 3, 0, 1
+    ])
+    expected_centroids = torch.tensor([
+        [0, 1, 2],
+        [3, 0, 1]
+    ])
+    
+    loss = head.compute_centroid_loss(hidden_states, targets)
+    
+    # Loss should have the same shape as targets
+    assert loss.shape == (batch_size, seq_len)
+    
+    # Manually compute the expected loss
+    flat_hidden = hidden_states.view(-1, 16)
+    expected_logits = head.centroids(flat_hidden)
+    expected_loss = F.cross_entropy(
+        expected_logits, 
+        expected_centroids.view(-1), 
+        reduction="none"
+    ).view(batch_size, seq_len)
+    
+    assert torch.allclose(loss, expected_loss)
+
+
+def test_compute_centroid_loss_with_ignore_index(head_setup):
+    """Test that ignore_index is properly masked out."""
+    head = head_setup
+    batch_size, seq_len = 2, 3
+    hidden_states = torch.randn(batch_size, seq_len, 16)
+    
+    targets = torch.tensor([
+        [0, -100, 128], 
+        [-100, -100, 70]
+    ])
+    expected_centroids = torch.tensor([
+        [0, -100, 2],
+        [-100, -100, 1]
+    ])
+    
+    loss = head.compute_centroid_loss(hidden_states, targets, ignore_index=-100)
+    
+    # Loss should be 0.0 where ignore_index was used
+    assert loss[0, 1].item() == 0.0
+    assert loss[1, 0].item() == 0.0
+    assert loss[1, 1].item() == 0.0
+    
+    # Check valid positions
+    valid_mask = targets != -100
+    valid_hidden = hidden_states[valid_mask]
+    valid_centroids = expected_centroids[valid_mask]
+    
+    expected_logits = head.centroids(valid_hidden)
+    expected_loss = F.cross_entropy(
+        expected_logits, 
+        valid_centroids, 
+        reduction="none"
+    )
+    
+    assert torch.allclose(loss[valid_mask], expected_loss)
+
+
+def test_compute_centroid_loss_all_ignored(head_setup):
+    """Test behavior when all targets are ignore_index."""
+    head = head_setup
+    batch_size, seq_len = 2, 3
+    hidden_states = torch.randn(batch_size, seq_len, 16)
+    
+    targets = torch.full((batch_size, seq_len), -100, dtype=torch.long)
+    
+    loss = head.compute_centroid_loss(hidden_states, targets, ignore_index=-100)
+    
+    assert loss.shape == (batch_size, seq_len)
+    assert torch.all(loss == 0.0)
+
+
+def test_compute_centroid_loss_with_custom_ordering(head_setup):
+    """Test that the token_ordering buffer correctly scrambles the target mapping."""
+    head = head_setup
+    
+    # Reverse the token ordering
+    head.token_ordering.copy_(torch.arange(255, -1, -1))
+    
+    hidden_states = torch.randn(1, 1, 16)
+    
+    # The token '0' is now at the very end of the array (position 255)
+    # Position 255 belongs to centroid 255 // 64 = 3
+    targets = torch.tensor([[0]])
+    
+    loss = head.compute_centroid_loss(hidden_states, targets)
+    
+    # Manually check that target centroid is 3
+    expected_logits = head.centroids(hidden_states.view(-1, 16))
+    expected_loss = F.cross_entropy(
+        expected_logits, 
+        torch.tensor([3]), 
+        reduction="none"
+    ).view(1, 1)
+    
+    assert torch.allclose(loss, expected_loss)

From 4e898bd98c407f34f9477ecd052ce2f5df79a350 Mon Sep 17 00:00:00 2001
From: santosh-agebold <12378904+shotsan@users.noreply.github.com>
Date: Mon, 13 Jul 2026 13:09:14 -0700
Subject: [PATCH 04/19] fix(mtp): address PR 767 CodeRabbit and collaborator
 comments

Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
---
 src/speculators/models/mtp/config.py | 14 +++++++-
 src/speculators/models/mtp/core.py   |  3 +-
 src/speculators/models/mtp/head.py   | 22 ++++++++----
 src/speculators/train/data.py        |  2 +-
 tests/unit/models/test_mtp_head.py   | 54 ++++++++++++++++++++++++++++
 tests/unit/train/test_data.py        | 52 +++++++++++++++++++++++++++
 6 files changed, 137 insertions(+), 10 deletions(-)

diff --git a/src/speculators/models/mtp/config.py b/src/speculators/models/mtp/config.py
index 7b5ac5375..845681ce0 100644
--- a/src/speculators/models/mtp/config.py
+++ b/src/speculators/models/mtp/config.py
@@ -2,6 +2,7 @@
 
 from typing import Any, Literal
 
+import pydantic
 from pydantic import Field, field_serializer, field_validator
 from transformers import AutoConfig, PretrainedConfig
 from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
@@ -60,11 +61,22 @@ class MTPSpeculatorConfig(SpeculatorModelConfig):
         ),
     )
 
-    top_k_centroids: int = Field(
+    centroid_intermediate_top_k: int = Field(
         default=32,
         description="Number of top centroids to select during multi-level LM head decoding.",
     )
 
+    @pydantic.model_validator(mode="after")
+    def validate_centroids(self) -> "MTPSpeculatorConfig":
+        if self.num_centroids is not None:
+            if self.num_centroids < 1:
+                raise ValueError(f"num_centroids must be >= 1, got {self.num_centroids}")
+            if self.vocab_size % self.num_centroids != 0:
+                raise ValueError(f"vocab_size ({self.vocab_size}) must be divisible by num_centroids ({self.num_centroids})")
+            if self.centroid_intermediate_top_k > self.num_centroids:
+                raise ValueError(f"centroid_intermediate_top_k ({self.centroid_intermediate_top_k}) cannot exceed num_centroids ({self.num_centroids})")
+        return self
+
     @property
     def hidden_size(self) -> int:
         return self.transformer_layer_config.hidden_size  # type: ignore[return-value]
diff --git a/src/speculators/models/mtp/core.py b/src/speculators/models/mtp/core.py
index d6d217ee1..74f4904a4 100644
--- a/src/speculators/models/mtp/core.py
+++ b/src/speculators/models/mtp/core.py
@@ -69,6 +69,7 @@ class MTPDraftModel(DraftVocabMixin, SpeculatorModel):
 
     t2d: torch.Tensor | None
     d2t: torch.Tensor | None
+    masked_embedding: nn.Module | None
 
     def __init__(self, config: MTPSpeculatorConfig) -> None:
         if config.transformer_layer_config._attn_implementation is None:  # noqa: SLF001
@@ -230,7 +231,7 @@ def forward(
             )
 
             if getattr(self, "masked_embedding", None) is not None:
-                unreduced += self.masked_embedding.compute_centroid_loss(
+                unreduced = unreduced + self.masked_embedding.compute_centroid_loss(
                     hidden_states=mtp_output,
                     targets=step_targets,
                     ignore_index=_IGNORE_INDEX,
diff --git a/src/speculators/models/mtp/head.py b/src/speculators/models/mtp/head.py
index d94faf3ca..58272c866 100644
--- a/src/speculators/models/mtp/head.py
+++ b/src/speculators/models/mtp/head.py
@@ -10,10 +10,22 @@ class MultiLevelLMHead(nn.Module):
 
     def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int):
         super().__init__()
+        if num_centroids <= 0:
+            raise ValueError(f"num_centroids must be positive, got {num_centroids}.")
+        if vocab_size % num_centroids != 0:
+            raise ValueError(f"vocab_size ({vocab_size}) must be divisible by num_centroids ({num_centroids}).")
+        
         self.tokens_per_centroid = vocab_size // num_centroids
         self.centroids = nn.Linear(hidden_size, num_centroids, bias=False)
         # Token ordering mapping (centroid subsets -> token IDs)
         self.register_buffer("token_ordering", torch.arange(vocab_size, dtype=torch.long))
+        self.register_buffer("token_ordering_inv", torch.empty_like(self.token_ordering))
+        self._rebuild_inverse()
+
+    def _rebuild_inverse(self) -> None:
+        """Rebuild the cached inverse mapping from token IDs to centroid IDs."""
+        positions = torch.arange(self.token_ordering.shape[0], device=self.token_ordering.device)
+        self.token_ordering_inv[self.token_ordering] = positions // self.tokens_per_centroid
 
     def compute_centroid_loss(
         self, hidden_states: torch.Tensor, targets: torch.Tensor, ignore_index: int = -100
@@ -27,13 +39,9 @@ def compute_centroid_loss(
         if valid_hidden.shape[0] == 0:
             return loss
 
-        # Inverse mapping: token_id -> centroid_id
-        inverse = torch.empty_like(self.token_ordering)
-        positions = torch.arange(self.token_ordering.shape[0], device=self.token_ordering.device)
-        inverse[self.token_ordering] = positions // self.tokens_per_centroid
-        
         centroid_logits = self.centroids(valid_hidden)
-        target_centroids = inverse[valid_targets]
-        loss[valid_mask] = F.cross_entropy(centroid_logits, target_centroids, reduction="none")
+        target_centroids = self.token_ordering_inv[valid_targets]
+        valid_loss = F.cross_entropy(centroid_logits, target_centroids, reduction="none")
         
+        loss = loss.masked_scatter(valid_mask, valid_loss.to(loss.dtype))
         return loss
diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index 23d2c5155..4feabfddd 100644
--- a/src/speculators/train/data.py
+++ b/src/speculators/train/data.py
@@ -197,7 +197,7 @@ def __getitem__(self, index) -> BatchType | None:
         # Convert hidden states and KV caches to the correct dtype
         data = {
             k: v.to(self.hidden_states_dtype)
-            if ("hidden_states" in k or "verifier_kv" in k)
+            if (k in ("hidden_states", "verifier_last_hidden_states") or k.startswith("verifier_kv_last_"))
             else v
             for k, v in data.items()
         }
diff --git a/tests/unit/models/test_mtp_head.py b/tests/unit/models/test_mtp_head.py
index cdabac45f..ea2bc5bac 100644
--- a/tests/unit/models/test_mtp_head.py
+++ b/tests/unit/models/test_mtp_head.py
@@ -118,6 +118,7 @@ def test_compute_centroid_loss_with_custom_ordering(head_setup):
     
     # Reverse the token ordering
     head.token_ordering.copy_(torch.arange(255, -1, -1))
+    head._rebuild_inverse()
     
     hidden_states = torch.randn(1, 1, 16)
     
@@ -136,3 +137,56 @@ def test_compute_centroid_loss_with_custom_ordering(head_setup):
     ).view(1, 1)
     
     assert torch.allclose(loss, expected_loss)
+
+
+def test_initialization_validation():
+    """Test constructor argument validation."""
+    with pytest.raises(ValueError, match="num_centroids must be positive"):
+        MultiLevelLMHead(16, 256, 0)
+        
+    with pytest.raises(ValueError, match="must be divisible by num_centroids"):
+        MultiLevelLMHead(16, 256, 3)
+
+
+def test_mtp_draft_model_integration_centroid_loss():
+    """Test that MTPDraftModel includes centroid loss when num_centroids is configured."""
+    from speculators.models.mtp.core import MTPDraftModel
+    from speculators.models.mtp.config import MTPSpeculatorConfig
+    from speculators.config import SpeculatorsConfig
+    from speculators.proposals.greedy import GreedyTokenProposalConfig
+    from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+    
+    tc = Qwen2Config(hidden_size=16, vocab_size=256, num_hidden_layers=1, num_attention_heads=2)
+    config = MTPSpeculatorConfig(
+        transformer_layer_config=tc,
+        num_centroids=4,
+        speculators_config=SpeculatorsConfig(
+            algorithm="mtp",
+            proposal_methods=[GreedyTokenProposalConfig(speculative_tokens=1)],
+            default_proposal_method="greedy"
+        )
+    )
+    
+    model = MTPDraftModel(config)
+    assert model.masked_embedding is not None
+    
+    batch_sz, seq_len = 1, 3
+    input_ids = torch.randint(0, 256, (batch_sz, seq_len))
+    hidden_states = torch.randn(batch_sz, seq_len, 16)
+    
+    _, loss, metrics = model(input_ids, hidden_states)
+    
+    # The total loss should include the masked_embedding centroid loss.
+    # We can verify the masked_embedding was called by replacing its compute_centroid_loss with a mock.
+    original_compute = model.masked_embedding.compute_centroid_loss
+    called = False
+    
+    def mock_compute(*args, **kwargs):
+        nonlocal called
+        called = True
+        return original_compute(*args, **kwargs)
+        
+    model.masked_embedding.compute_centroid_loss = mock_compute
+    model(input_ids, hidden_states)
+    
+    assert called, "centroid loss was not computed during forward pass"
diff --git a/tests/unit/train/test_data.py b/tests/unit/train/test_data.py
index f3d10c2eb..92b678700 100644
--- a/tests/unit/train/test_data.py
+++ b/tests/unit/train/test_data.py
@@ -515,6 +515,58 @@ def test_verifier_kvs_survive_pipeline():
     assert torch.equal(local_kv, expected_local)
 
 
+def test_verifier_kvs_mixed_batch():
+    """KV-present and KV-absent samples in one batch should not crash or drop data."""
+    from speculators.train.data import create_collate_fn
+    
+    with_kv = {
+        "input_ids": torch.tensor([0]),
+        "loss_mask": torch.tensor([1]),
+        "hidden_states": torch.tensor([[0.0]]),
+        "verifier_last_hidden_states": torch.tensor([[0.0]]),
+        "verifier_kv_last_local": torch.tensor([[100.0]]),
+        "lengths": torch.tensor([1]),
+        "position_ids": torch.tensor([0]),
+    }
+    without_kv = {k: v for k, v in with_kv.items() if not k.startswith("verifier_kv")}
+    collate_fn = create_collate_fn(max_len=5, hidden_size=1, num_target_layers=1)
+    
+    # Depending on which element is first, it either raises KeyError or drops the field.
+    # We verify the invariant holds.
+    try:
+        collated = collate_fn([with_kv, without_kv])
+    except KeyError:
+        pass
+
+
+def test_verifier_kvs_dtype_cast():
+    """Verify BaseDataset.__getitem__ casts verifier KV tensors when hidden_states_dtype is bfloat16."""
+    from speculators.train.data import BaseDataset
+    
+    class DummyDataset(BaseDataset):
+        def __len__(self):
+            return 1
+            
+        def _get_raw_data(self, idx):
+            return {
+                "input_ids": torch.tensor([0]),
+                "loss_mask": torch.tensor([1]),
+                "hidden_states": torch.tensor([[0.0]], dtype=torch.float32),
+                "verifier_last_hidden_states": torch.tensor([[0.0]], dtype=torch.float32),
+                "verifier_kv_last_local": torch.tensor([[100.0]], dtype=torch.float32),
+            }
+            
+    ds = DummyDataset(
+        max_len=128,
+        hidden_states_dtype=torch.bfloat16,
+    )
+    
+    item = ds[0]
+    assert item["verifier_kv_last_local"].dtype == torch.bfloat16
+    assert item["hidden_states"].dtype == torch.bfloat16
+    assert item["verifier_last_hidden_states"].dtype == torch.bfloat16
+
+
 def test_arrow_dataset_on_generate_cache_creates_hidden_states_dir(tmp_path: Path):
     """on_generate="cache" must create the cache dir when cache() is called —
     otherwise shutil.move into it raises FileNotFoundError, which _maybe_generate_hs

From 8a1ad105e9dbae225d6f74e3c0c26ba5aeb7ee01 Mon Sep 17 00:00:00 2001
From: santosh-agebold <12378904+shotsan@users.noreply.github.com>
Date: Fri, 10 Jul 2026 14:12:00 -0700
Subject: [PATCH 05/19] feat(mtp): implement QueryOnlyGemma2Attention for
 Gemma4 MTP draft models

This commit introduces the architecture modifications required for Gemma4 MTP support.
Unlike Qwen-style MTP which maintains its own full attention layers and KV cache,
Gemma4 MTP borrows the last local and global KV caches from the verifier base model.

- Registered Gemma2 components in base_components.py
- Implemented QueryOnlyGemma2Attention to bypass K/V projections and route external KV caches
- Updated MTPDraftModel.forward to pass verifier_kv_last_local/global to layers
- Added unit tests for KV cache routing in Gemma4 MTP

Relates to #586, #758

Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com>
---
 src/speculators/models/base_components.py     | 18 ++++
 src/speculators/models/mtp/core.py            |  2 +
 .../models/mtp/model_definitions.py           | 88 +++++++++++++++++++
 tests/unit/models/test_mtp_gemma2.py          | 85 ++++++++++++++++++
 4 files changed, 193 insertions(+)
 create mode 100644 tests/unit/models/test_mtp_gemma2.py

diff --git a/src/speculators/models/base_components.py b/src/speculators/models/base_components.py
index 77f6d9d23..f36825b8e 100644
--- a/src/speculators/models/base_components.py
+++ b/src/speculators/models/base_components.py
@@ -12,6 +12,14 @@
     Qwen3RMSNorm,
     Qwen3RotaryEmbedding,
 )
+try:
+    from transformers.models.gemma2.modeling_gemma2 import (
+        Gemma2DecoderLayer,
+        Gemma2RMSNorm,
+        Gemma2RotaryEmbedding,
+    )
+except ImportError:
+    pass
 
 
 class ModelComponents(NamedTuple):
@@ -51,6 +59,16 @@ class ModelComponents(NamedTuple):
     ),
 }
 
+try:
+    model_classes["gemma2"] = ModelComponents(
+        Gemma2DecoderLayer,
+        Gemma2DecoderLayer,
+        Gemma2RMSNorm,
+        Gemma2RotaryEmbedding,
+    )
+except NameError:
+    pass
+
 try:
     from transformers.models.qwen3_next.modeling_qwen3_next import (
         Qwen3NextDecoderLayer,
diff --git a/src/speculators/models/mtp/core.py b/src/speculators/models/mtp/core.py
index 74f4904a4..dcb761f5e 100644
--- a/src/speculators/models/mtp/core.py
+++ b/src/speculators/models/mtp/core.py
@@ -211,6 +211,8 @@ def forward(
                 attention_mask=causal_mask,
                 position_ids=step_pos_ids,
                 position_embeddings=step_pos_emb,
+                verifier_kv_last_local=kwargs.get("verifier_kv_last_local"),
+                verifier_kv_last_global=kwargs.get("verifier_kv_last_global"),
             )
 
             logits = self.lm_head(mtp_output)
diff --git a/src/speculators/models/mtp/model_definitions.py b/src/speculators/models/mtp/model_definitions.py
index cac51b3db..a3cff92ee 100644
--- a/src/speculators/models/mtp/model_definitions.py
+++ b/src/speculators/models/mtp/model_definitions.py
@@ -203,3 +203,91 @@ def __init__(self, config: PretrainedConfig, layer_idx: int = 0) -> None:
     mtp_model_classes["qwen3_5_moe_text"] = base_components.override_components(
         "qwen3_5_moe_text", first_layer_class=Qwen35MoeMTPLayer
     )
+
+if "gemma2" in base_components.model_classes:
+    from typing import Optional, Tuple, Any
+    from transformers.models.gemma2.modeling_gemma2 import (
+        Gemma2DecoderLayer,
+        Gemma2RMSNorm,
+        Gemma2Attention,
+        apply_rotary_pos_emb,
+        eager_attention_forward,
+        ALL_ATTENTION_FUNCTIONS
+    )
+
+    class QueryOnlyGemma2Attention(Gemma2Attention):
+        def __init__(self, config, layer_idx: Optional[int] = None):
+            super().__init__(config, layer_idx)
+            # Strip out K/V projections since we use verifier's KV cache
+            self.k_proj = None
+            self.v_proj = None
+
+        def forward(
+            self,
+            hidden_states: torch.Tensor,
+            position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+            attention_mask: Optional[torch.Tensor],
+            past_key_values: Optional[Any] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            verifier_kv_last_local: Optional[torch.Tensor] = None,
+            verifier_kv_last_global: Optional[torch.Tensor] = None,
+            **kwargs: Any,
+        ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, self.head_dim)
+
+            query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+            is_local = (self.sliding_window is not None and self.sliding_window > 0)
+            kv_tensor = verifier_kv_last_local if is_local else verifier_kv_last_global
+
+            if kv_tensor is None:
+                # Fallback to zeros if not provided (e.g. dummy forward passes)
+                batch_sz, seq_len = input_shape
+                kv_tensor = torch.zeros(
+                    (batch_sz, seq_len, 2, self.num_key_value_heads, self.head_dim),
+                    dtype=query_states.dtype,
+                    device=query_states.device
+                )
+            elif kv_tensor.dim() == 3:
+                batch_sz, seq_len, _ = kv_tensor.shape
+                kv_tensor = kv_tensor.view(batch_sz, seq_len, 2, self.num_key_value_heads, self.head_dim)
+
+            key_states = kv_tensor[..., 0, :, :].transpose(1, 2)
+            value_states = kv_tensor[..., 1, :, :].transpose(1, 2)
+
+            cos, sin = position_embeddings
+            # Only apply RoPE to queries; verifier's KV cache already has RoPE applied
+            query_states, _ = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+            attention_interface = eager_attention_forward
+            if getattr(self.config, "_attn_implementation", "eager") != "eager":
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=self.attention_dropout if self.training else 0.0,
+                scaling=self.scaling,
+                sliding_window=self.sliding_window,
+                softcap=self.attn_logit_softcapping,
+                **kwargs,
+            )
+
+            attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+            attn_output = self.o_proj(attn_output)
+            return attn_output, attn_weights
+
+    class Gemma2MTPLayer(MTPLayerMixin, Gemma2DecoderLayer):
+        def __init__(self, config: PretrainedConfig, layer_idx: int = 0) -> None:
+            super().__init__(config, layer_idx)
+            self._setup_mtp_modules(config, Gemma2RMSNorm)
+            # Replace the standard attention with query-only attention
+            self.self_attn = QueryOnlyGemma2Attention(config, layer_idx)
+
+    mtp_model_classes["gemma2"] = base_components.override_components(
+        "gemma2", first_layer_class=Gemma2MTPLayer
+    )
diff --git a/tests/unit/models/test_mtp_gemma2.py b/tests/unit/models/test_mtp_gemma2.py
new file mode 100644
index 000000000..0b6dd8934
--- /dev/null
+++ b/tests/unit/models/test_mtp_gemma2.py
@@ -0,0 +1,85 @@
+"""Unit tests for Gemma2 MTP Query-Only Attention."""
+
+import pytest
+import torch
+from transformers import PretrainedConfig
+
+from speculators.models.base_components import model_classes
+from speculators.models.mtp.model_definitions import mtp_model_classes
+
+if "gemma2" not in model_classes:
+    pytest.skip("transformers < 4.42 installed, skipping gemma2 tests", allow_module_level=True)
+
+from speculators.models.mtp.model_definitions import QueryOnlyGemma2Attention
+
+
+@pytest.fixture
+def gemma2_config():
+    config = PretrainedConfig(
+        hidden_size=64,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        attention_bias=False,
+        attention_dropout=0.0,
+        _attn_implementation="eager"
+    )
+    # Mock some Gemma2 specific attrs
+    config.sliding_window = 4096
+    config.query_pre_attn_scalar = 224
+    config.hidden_size = 64
+    config.num_attention_heads = 4
+    config.num_key_value_heads = 2
+    config.head_dim = 16
+    return config
+
+
+def test_query_only_gemma2_attention_local_kv(gemma2_config):
+    """Test that local KV cache is used when sliding_window > 0."""
+    attn = QueryOnlyGemma2Attention(gemma2_config, layer_idx=0)
+    
+    batch_sz, seq_len = 2, 5
+    hidden_states = torch.randn(batch_sz, seq_len, 64)
+    
+    # [batch, seq_len, 2, num_kv_heads, head_dim]
+    local_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
+    global_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
+    
+    # Mock position embeddings (cos, sin)
+    cos = torch.randn(1, 1, seq_len, 16)
+    sin = torch.randn(1, 1, seq_len, 16)
+    
+    output, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+        verifier_kv_last_local=local_kv,
+        verifier_kv_last_global=global_kv,
+    )
+    
+    assert output.shape == (batch_sz, seq_len, 64)
+
+
+def test_query_only_gemma2_attention_global_kv(gemma2_config):
+    """Test that global KV cache is used when sliding_window is 0 or None."""
+    gemma2_config.sliding_window = None
+    attn = QueryOnlyGemma2Attention(gemma2_config, layer_idx=0)
+    
+    batch_sz, seq_len = 2, 5
+    hidden_states = torch.randn(batch_sz, seq_len, 64)
+    
+    local_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
+    global_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
+    
+    cos = torch.randn(1, 1, seq_len, 16)
+    sin = torch.randn(1, 1, seq_len, 16)
+    
+    output, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+        verifier_kv_last_local=local_kv,
+        verifier_kv_last_global=global_kv,
+    )
+    
+    assert output.shape == (batch_sz, seq_len, 64)

From 3379b81aabb495fdd15954ab20e1d09ca85fd10d Mon Sep 17 00:00:00 2001
From: shotsan <shotsan@users.noreply.github.com>
Date: Mon, 13 Jul 2026 15:15:08 -0700
Subject: [PATCH 06/19] fix(mtp): address PR 768 CodeRabbit and collaborator
 comments

Signed-off-by: shotsan <shotsan@users.noreply.github.com>
---
 src/speculators/models/base_components.py     |  6 +-
 .../models/mtp/model_definitions.py           |  9 ++-
 src/speculators/train/data.py                 |  9 ++-
 tests/unit/models/test_mtp_gemma2.py          | 77 ++++++++++++++-----
 4 files changed, 75 insertions(+), 26 deletions(-)

diff --git a/src/speculators/models/base_components.py b/src/speculators/models/base_components.py
index f36825b8e..cc5e2120a 100644
--- a/src/speculators/models/base_components.py
+++ b/src/speculators/models/base_components.py
@@ -12,12 +12,14 @@
     Qwen3RMSNorm,
     Qwen3RotaryEmbedding,
 )
+HAS_GEMMA2 = False
 try:
     from transformers.models.gemma2.modeling_gemma2 import (
         Gemma2DecoderLayer,
         Gemma2RMSNorm,
         Gemma2RotaryEmbedding,
     )
+    HAS_GEMMA2 = True
 except ImportError:
     pass
 
@@ -59,15 +61,13 @@ class ModelComponents(NamedTuple):
     ),
 }
 
-try:
+if HAS_GEMMA2:
     model_classes["gemma2"] = ModelComponents(
         Gemma2DecoderLayer,
         Gemma2DecoderLayer,
         Gemma2RMSNorm,
         Gemma2RotaryEmbedding,
     )
-except NameError:
-    pass
 
 try:
     from transformers.models.qwen3_next.modeling_qwen3_next import (
diff --git a/src/speculators/models/mtp/model_definitions.py b/src/speculators/models/mtp/model_definitions.py
index a3cff92ee..cd0ba0eee 100644
--- a/src/speculators/models/mtp/model_definitions.py
+++ b/src/speculators/models/mtp/model_definitions.py
@@ -264,6 +264,14 @@ def forward(
             if getattr(self.config, "_attn_implementation", "eager") != "eager":
                 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
 
+            # In eager mode, sliding_window is often dropped by the interface.
+            # We enforce it directly onto the causal mask here if configured.
+            if self.sliding_window is not None and self.sliding_window > 0 and attention_mask is not None:
+                seq_len = attention_mask.shape[-1]
+                window_mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=attention_mask.device))
+                window_mask = torch.triu(window_mask, diagonal=-self.sliding_window + 1)
+                attention_mask = attention_mask.masked_fill(~window_mask.view(1, 1, seq_len, seq_len), torch.finfo(query_states.dtype).min)
+
             attn_output, attn_weights = attention_interface(
                 self,
                 query_states,
@@ -272,7 +280,6 @@ def forward(
                 attention_mask,
                 dropout=self.attention_dropout if self.training else 0.0,
                 scaling=self.scaling,
-                sliding_window=self.sliding_window,
                 softcap=self.attn_logit_softcapping,
                 **kwargs,
             )
diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index 4feabfddd..9f90b7fdd 100644
--- a/src/speculators/train/data.py
+++ b/src/speculators/train/data.py
@@ -66,7 +66,7 @@ def split_files(datapath: str, ratio: float = 0.9, seed: int = 0):
 
 
 def create_empty_sample(
-    hidden_size: int, num_target_layers: int = 3, dtype: torch.dtype = torch.bfloat16
+    hidden_size: int, num_target_layers: int = 3, dtype: torch.dtype = torch.bfloat16, kv_feature_dim: int = 0
 ):
     # data structure: {
     #     "hidden_states": [seq_len, num_target_layers * hidden_size],
@@ -87,8 +87,8 @@ def create_empty_sample(
         "hidden_states": torch.empty(0, num_target_layers * hidden_size, dtype=dtype),
         "input_ids": torch.empty(0, dtype=torch.long),
         "verifier_last_hidden_states": torch.empty(0, hidden_size, dtype=dtype),
-        "verifier_kv_last_local": torch.empty(0, dtype=dtype),
-        "verifier_kv_last_global": torch.empty(0, dtype=dtype),
+        "verifier_kv_last_local": torch.empty(0, kv_feature_dim, dtype=dtype) if kv_feature_dim > 0 else torch.empty(0, dtype=dtype),
+        "verifier_kv_last_global": torch.empty(0, kv_feature_dim, dtype=dtype) if kv_feature_dim > 0 else torch.empty(0, dtype=dtype),
         "loss_mask": torch.empty(0, dtype=torch.bool),
         "lengths": torch.tensor([0], dtype=torch.long),
         "position_ids": torch.arange(0, dtype=torch.long),
@@ -492,6 +492,7 @@ def create_collate_fn(
     num_target_layers: int = 3,
     dtype: torch.dtype = torch.bfloat16,
     preprocess: Callable[[BatchType], BatchType] | None = None,
+    kv_feature_dim: int = 0,
 ):
     def collate_fn(batch: list[BatchType | None]) -> BatchType:
         # Apply per-sample preprocessing and filter failed samples
@@ -503,7 +504,7 @@ def collate_fn(batch: list[BatchType | None]) -> BatchType:
             # Match the configured `dtype` so the placeholder doesn't crash
             # downstream layers loaded at a different precision (e.g. bf16
             # weights vs fp32 default placeholders).
-            empty = create_empty_sample(hidden_size, num_target_layers, dtype=dtype)
+            empty = create_empty_sample(hidden_size, num_target_layers, dtype=dtype, kv_feature_dim=kv_feature_dim)
             if preprocess:
                 empty = preprocess(empty)
             batch = [empty]
diff --git a/tests/unit/models/test_mtp_gemma2.py b/tests/unit/models/test_mtp_gemma2.py
index 0b6dd8934..f358a557d 100644
--- a/tests/unit/models/test_mtp_gemma2.py
+++ b/tests/unit/models/test_mtp_gemma2.py
@@ -2,7 +2,7 @@
 
 import pytest
 import torch
-from transformers import PretrainedConfig
+from transformers.models.gemma2.configuration_gemma2 import Gemma2Config
 
 from speculators.models.base_components import model_classes
 from speculators.models.mtp.model_definitions import mtp_model_classes
@@ -15,22 +15,20 @@
 
 @pytest.fixture
 def gemma2_config():
-    config = PretrainedConfig(
+    config = Gemma2Config(
         hidden_size=64,
         num_attention_heads=4,
         num_key_value_heads=2,
         head_dim=16,
         attention_bias=False,
         attention_dropout=0.0,
-        _attn_implementation="eager"
     )
-    # Mock some Gemma2 specific attrs
+    # Mock some Gemma2 specific attrs that we override
+    config._attn_implementation = "eager"
     config.sliding_window = 4096
     config.query_pre_attn_scalar = 224
-    config.hidden_size = 64
-    config.num_attention_heads = 4
-    config.num_key_value_heads = 2
-    config.head_dim = 16
+    config.attn_logit_softcapping = 50.0
+    config.layer_types = ["sliding_attention"]
     return config
 
 
@@ -45,23 +43,45 @@ def test_query_only_gemma2_attention_local_kv(gemma2_config):
     local_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
     global_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
     
-    # Mock position embeddings (cos, sin)
-    cos = torch.randn(1, 1, seq_len, 16)
-    sin = torch.randn(1, 1, seq_len, 16)
+    # Mock position embeddings (cos, sin) - shape [batch, seq_len, head_dim]
+    cos = torch.randn(batch_sz, seq_len, 16)
+    sin = torch.randn(batch_sz, seq_len, 16)
     
-    output, _ = attn(
+    output1, _ = attn(
         hidden_states=hidden_states,
         position_embeddings=(cos, sin),
         attention_mask=None,
         verifier_kv_last_local=local_kv,
         verifier_kv_last_global=global_kv,
     )
+    assert output1.shape == (batch_sz, seq_len, 64)
     
-    assert output.shape == (batch_sz, seq_len, 64)
+    # Change global_kv (the inactive cache)
+    global_kv_changed = torch.randn_like(global_kv)
+    output2, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+        verifier_kv_last_local=local_kv,
+        verifier_kv_last_global=global_kv_changed,
+    )
+    assert torch.equal(output1, output2)
+    
+    # Change local_kv (the active cache)
+    local_kv_changed = torch.randn_like(local_kv)
+    output3, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+        verifier_kv_last_local=local_kv_changed,
+        verifier_kv_last_global=global_kv,
+    )
+    assert not torch.equal(output1, output3)
 
 
 def test_query_only_gemma2_attention_global_kv(gemma2_config):
-    """Test that global KV cache is used when sliding_window is 0 or None."""
+    """Test that global KV cache is used when layer_types is full_attention."""
+    gemma2_config.layer_types = ["full_attention"]
     gemma2_config.sliding_window = None
     attn = QueryOnlyGemma2Attention(gemma2_config, layer_idx=0)
     
@@ -71,15 +91,36 @@ def test_query_only_gemma2_attention_global_kv(gemma2_config):
     local_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
     global_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
     
-    cos = torch.randn(1, 1, seq_len, 16)
-    sin = torch.randn(1, 1, seq_len, 16)
+    cos = torch.randn(batch_sz, seq_len, 16)
+    sin = torch.randn(batch_sz, seq_len, 16)
     
-    output, _ = attn(
+    output1, _ = attn(
         hidden_states=hidden_states,
         position_embeddings=(cos, sin),
         attention_mask=None,
         verifier_kv_last_local=local_kv,
         verifier_kv_last_global=global_kv,
     )
+    assert output1.shape == (batch_sz, seq_len, 64)
     
-    assert output.shape == (batch_sz, seq_len, 64)
+    # Change local_kv (the inactive cache)
+    local_kv_changed = torch.randn_like(local_kv)
+    output2, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+        verifier_kv_last_local=local_kv_changed,
+        verifier_kv_last_global=global_kv,
+    )
+    assert torch.equal(output1, output2)
+    
+    # Change global_kv (the active cache)
+    global_kv_changed = torch.randn_like(global_kv)
+    output3, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=None,
+        verifier_kv_last_local=local_kv,
+        verifier_kv_last_global=global_kv_changed,
+    )
+    assert not torch.equal(output1, output3)

From 6fd6179656bfb8b33f9dfabc8a252778bf4685c7 Mon Sep 17 00:00:00 2001
From: shotsan <shotsan@users.noreply.github.com>
Date: Mon, 13 Jul 2026 18:03:43 -0700
Subject: [PATCH 07/19] fix(mtp): address PR 768 CodeRabbit comments

Signed-off-by: shotsan <shotsan@users.noreply.github.com>
---
 src/speculators/models/mtp/head.py   |  2 +-
 tests/unit/models/test_mtp_gemma2.py | 46 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/speculators/models/mtp/head.py b/src/speculators/models/mtp/head.py
index 58272c866..cfb40d178 100644
--- a/src/speculators/models/mtp/head.py
+++ b/src/speculators/models/mtp/head.py
@@ -19,7 +19,7 @@ def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int):
         self.centroids = nn.Linear(hidden_size, num_centroids, bias=False)
         # Token ordering mapping (centroid subsets -> token IDs)
         self.register_buffer("token_ordering", torch.arange(vocab_size, dtype=torch.long))
-        self.register_buffer("token_ordering_inv", torch.empty_like(self.token_ordering))
+        self.register_buffer("token_ordering_inv", torch.empty_like(self.token_ordering), persistent=False)
         self._rebuild_inverse()
 
     def _rebuild_inverse(self) -> None:
diff --git a/tests/unit/models/test_mtp_gemma2.py b/tests/unit/models/test_mtp_gemma2.py
index f358a557d..ff5b34f44 100644
--- a/tests/unit/models/test_mtp_gemma2.py
+++ b/tests/unit/models/test_mtp_gemma2.py
@@ -124,3 +124,49 @@ def test_query_only_gemma2_attention_global_kv(gemma2_config):
         verifier_kv_last_global=global_kv_changed,
     )
     assert not torch.equal(output1, output3)
+
+
+def test_query_only_gemma2_attention_sliding_window_mask(gemma2_config):
+    """Test that sliding window masking actually restricts attention when a mask is provided."""
+    gemma2_config.sliding_window = 2
+    attn = QueryOnlyGemma2Attention(gemma2_config, layer_idx=0)
+    
+    batch_sz, seq_len = 1, 5
+    hidden_states = torch.randn(batch_sz, seq_len, 64)
+    
+    local_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
+    global_kv = torch.randn(batch_sz, seq_len, 2, 2, 16)
+    
+    cos = torch.randn(batch_sz, seq_len, 16)
+    sin = torch.randn(batch_sz, seq_len, 16)
+    
+    # Base causal mask (lower triangular) - 0 for allowed, min_val for masked
+    attention_mask = torch.zeros(batch_sz, 1, seq_len, seq_len)
+    causal_mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))
+    attention_mask = attention_mask.masked_fill(~causal_mask.view(1, 1, seq_len, seq_len), torch.finfo(hidden_states.dtype).min)
+    
+    output1, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=attention_mask,
+        verifier_kv_last_local=local_kv,
+        verifier_kv_last_global=global_kv,
+    )
+    
+    # Perturb local_kv at index 0 (which should be masked out for queries at index >= 2)
+    local_kv_changed = local_kv.clone()
+    local_kv_changed[:, 0, ...] = torch.randn_like(local_kv[:, 0, ...])
+    
+    output2, _ = attn(
+        hidden_states=hidden_states,
+        position_embeddings=(cos, sin),
+        attention_mask=attention_mask,
+        verifier_kv_last_local=local_kv_changed,
+        verifier_kv_last_global=global_kv,
+    )
+    
+    # Output at index 4 should be identical because index 0 is outside the sliding window (window size 2)
+    assert torch.equal(output1[:, 4], output2[:, 4])
+    
+    # Output at index 0 should change since it attends to index 0
+    assert not torch.equal(output1[:, 0], output2[:, 0])

From 727678741b8d8c7d033d3b08c5a4b7cdd2865cf2 Mon Sep 17 00:00:00 2001
From: shotsan <shotsan@users.noreply.github.com>
Date: Mon, 13 Jul 2026 16:13:49 -0700
Subject: [PATCH 08/19] fix(data): harden create_collate_fn for mixed batches

Signed-off-by: shotsan <shotsan@users.noreply.github.com>
---
 src/speculators/train/data.py | 21 +++++++++++++++++++--
 tests/unit/train/test_data.py | 23 +++++++++++++++++------
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index 9f90b7fdd..eaa310034 100644
--- a/src/speculators/train/data.py
+++ b/src/speculators/train/data.py
@@ -509,10 +509,27 @@ def collate_fn(batch: list[BatchType | None]) -> BatchType:
                 empty = preprocess(empty)
             batch = [empty]
 
+        all_keys = set()
+        for b in batch:
+            all_keys.update(b.keys())
+
         collated_data = {}
-        for key in batch[0]:  # type: ignore[union-attr]
+        for key in all_keys:
+            template = next(b[key] for b in batch if key in b)
+            
+            tensors_to_cat = []
+            for b in batch:
+                if key in b:
+                    tensors_to_cat.append(b[key])
+                else:
+                    seq_len = b["input_ids"].shape[0] if "input_ids" in b else 0
+                    dummy_shape = (seq_len,) + template.shape[1:]
+                    tensors_to_cat.append(
+                        torch.zeros(dummy_shape, dtype=template.dtype, device=template.device)
+                    )
+                    
             # Concatenate the tensors along the seq (0th) dimension
-            collated_data[key] = torch.cat([b[key] for b in batch], dim=0)  # type: ignore[index]
+            collated_data[key] = torch.cat(tensors_to_cat, dim=0)
             # shape: [total_seq_len, ...]
 
             if key != "lengths":
diff --git a/tests/unit/train/test_data.py b/tests/unit/train/test_data.py
index 92b678700..2ddaa7b44 100644
--- a/tests/unit/train/test_data.py
+++ b/tests/unit/train/test_data.py
@@ -531,12 +531,20 @@ def test_verifier_kvs_mixed_batch():
     without_kv = {k: v for k, v in with_kv.items() if not k.startswith("verifier_kv")}
     collate_fn = create_collate_fn(max_len=5, hidden_size=1, num_target_layers=1)
     
-    # Depending on which element is first, it either raises KeyError or drops the field.
-    # We verify the invariant holds.
-    try:
-        collated = collate_fn([with_kv, without_kv])
-    except KeyError:
-        pass
+    # Test both orderings
+    for batch in [[with_kv, without_kv], [without_kv, with_kv]]:
+        collated = collate_fn(batch)
+        
+        assert "verifier_kv_last_local" in collated
+        local_kv = collated["verifier_kv_last_local"]
+        assert local_kv.shape == (1, 5, 1)
+        
+        if batch[0] == with_kv:
+            expected = torch.tensor([[[100.0], [0.0], [0.0], [0.0], [0.0]]])
+        else:
+            expected = torch.tensor([[[0.0], [100.0], [0.0], [0.0], [0.0]]])
+            
+        assert torch.equal(local_kv, expected)
 
 
 def test_verifier_kvs_dtype_cast():
@@ -547,6 +555,9 @@ class DummyDataset(BaseDataset):
         def __len__(self):
             return 1
             
+        def _compute_approx_lengths(self):
+            return [1]
+            
         def _get_raw_data(self, idx):
             return {
                 "input_ids": torch.tensor([0]),

From e44624f67037001d12668a0d7ba0bc62fa0c8f51 Mon Sep 17 00:00:00 2001
From: shotsan <shotsan@users.noreply.github.com>
Date: Tue, 14 Jul 2026 13:01:32 -0700
Subject: [PATCH 09/19] fix(mtp): ignore token_ordering_inv in checkpoint
 loading

Signed-off-by: shotsan <shotsan@users.noreply.github.com>
---
 src/speculators/models/mtp/core.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/speculators/models/mtp/core.py b/src/speculators/models/mtp/core.py
index dcb761f5e..c6d597552 100644
--- a/src/speculators/models/mtp/core.py
+++ b/src/speculators/models/mtp/core.py
@@ -65,6 +65,7 @@ class MTPDraftModel(DraftVocabMixin, SpeculatorModel):
         "d2t",
         "masked_embedding.centroids.weight",
         "masked_embedding.token_ordering",
+        "masked_embedding.token_ordering_inv",
     ]
 
     t2d: torch.Tensor | None

From 3df917f2954510d4f4ff18e8bf43ff6cebe36c9c Mon Sep 17 00:00:00 2001
From: shotsan <shotsan@users.noreply.github.com>
Date: Tue, 14 Jul 2026 14:12:27 -0700
Subject: [PATCH 10/19] fix(data): stack separate k and v tensors from PR 674
 dataset format

Signed-off-by: shotsan <shotsan@users.noreply.github.com>
---
 src/speculators/train/data.py | 12 ++++++++----
 tests/unit/train/test_data.py | 26 +++++++++++++++++++++-----
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index eaa310034..30d7855fa 100644
--- a/src/speculators/train/data.py
+++ b/src/speculators/train/data.py
@@ -114,10 +114,14 @@ def standardize_data_v1(data: dict[str, Any]) -> dict[str, Any]:
         "verifier_last_hidden_states": data["hidden_states"][-1],
         "loss_mask": data["loss_mask"],
     }
-    if "verifier_kv_last_local" in data:
-        res["verifier_kv_last_local"] = data["verifier_kv_last_local"]
-    if "verifier_kv_last_global" in data:
-        res["verifier_kv_last_global"] = data["verifier_kv_last_global"]
+    if "kv_last_local_k" in data and "kv_last_local_v" in data:
+        res["verifier_kv_last_local"] = torch.stack(
+            [data["kv_last_local_k"], data["kv_last_local_v"]], dim=1
+        )
+    if "kv_last_global_k" in data and "kv_last_global_v" in data:
+        res["verifier_kv_last_global"] = torch.stack(
+            [data["kv_last_global_k"], data["kv_last_global_v"]], dim=1
+        )
     return res
 
 
diff --git a/tests/unit/train/test_data.py b/tests/unit/train/test_data.py
index 2ddaa7b44..271795606 100644
--- a/tests/unit/train/test_data.py
+++ b/tests/unit/train/test_data.py
@@ -482,8 +482,10 @@ def test_verifier_kvs_survive_pipeline():
             torch.tensor([[0.0, 0.1], [1.0, 1.1], [2.0, 2.1]]),  # Layer 0
             torch.tensor([[10.0, 10.1], [11.0, 11.1], [12.0, 12.1]]),  # Verifier last hs
         ],
-        "verifier_kv_last_local": torch.tensor([[100.0], [101.0], [102.0]]),
-        "verifier_kv_last_global": torch.tensor([[200.0], [201.0], [202.0]]),
+        "kv_last_local_k": torch.tensor([[100.0], [101.0], [102.0]]),
+        "kv_last_local_v": torch.tensor([[110.0], [111.0], [112.0]]),
+        "kv_last_global_k": torch.tensor([[200.0], [201.0], [202.0]]),
+        "kv_last_global_v": torch.tensor([[210.0], [211.0], [212.0]]),
     }
 
     # 2. Pass through standardize_data_v1
@@ -491,7 +493,11 @@ def test_verifier_kvs_survive_pipeline():
 
     assert "verifier_kv_last_local" in standardized
     assert "verifier_kv_last_global" in standardized
-    assert torch.equal(standardized["verifier_kv_last_local"], v1_data["verifier_kv_last_local"])
+
+    expected_local_stack = torch.stack(
+        [v1_data["kv_last_local_k"], v1_data["kv_last_local_v"]], dim=1  # type: ignore[arg-type]
+    )
+    assert torch.equal(standardized["verifier_kv_last_local"], expected_local_stack)
 
     # 3. Add lengths and position_ids (simulate BaseDataset.__getitem__)
     standardized["lengths"] = torch.tensor([3], dtype=torch.long)
@@ -508,10 +514,20 @@ def test_verifier_kvs_survive_pipeline():
     assert "verifier_kv_last_global" in collated
 
     local_kv = collated["verifier_kv_last_local"]
-    assert local_kv.shape == (1, 5, 1)  # [batch=1, max_len=5, ...]
+    assert local_kv.shape == (1, 5, 2, 1)  # [batch=1, max_len=5, ...]
 
     # First 3 positions should match original, last 2 should be padded with 0
-    expected_local = torch.tensor([[[100.0], [101.0], [102.0], [0.0], [0.0]]])
+    expected_local = torch.tensor(
+        [
+            [
+                [[100.0], [110.0]],
+                [[101.0], [111.0]],
+                [[102.0], [112.0]],
+                [[0.0], [0.0]],
+                [[0.0], [0.0]],
+            ]
+        ]
+    )
     assert torch.equal(local_kv, expected_local)
 
 

From c7553393ea7d86f00b8164e4fab5d44c1ae1313f Mon Sep 17 00:00:00 2001
From: shotsan <shotsan@users.noreply.github.com>
Date: Fri, 17 Jul 2026 10:42:41 -0700
Subject: [PATCH 11/19] fix(data): stack separate k and v tensors in
 ArrowDataset as well

Signed-off-by: shotsan <shotsan@users.noreply.github.com>
---
 src/speculators/train/data.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index 30d7855fa..c1ab75f0f 100644
--- a/src/speculators/train/data.py
+++ b/src/speculators/train/data.py
@@ -386,10 +386,14 @@ def _get_raw_data(self, index):
             ],  # [seq_len, hidden_size]
             "loss_mask": self.data[index]["loss_mask"],  # [seq_len]
         }
-        if "verifier_kv_last_local" in loaded_hs:
-            res["verifier_kv_last_local"] = loaded_hs["verifier_kv_last_local"]
-        if "verifier_kv_last_global" in loaded_hs:
-            res["verifier_kv_last_global"] = loaded_hs["verifier_kv_last_global"]
+        if "kv_last_local_k" in loaded_hs and "kv_last_local_v" in loaded_hs:
+            res["verifier_kv_last_local"] = torch.stack(
+                [loaded_hs["kv_last_local_k"], loaded_hs["kv_last_local_v"]], dim=1
+            )
+        if "kv_last_global_k" in loaded_hs and "kv_last_global_v" in loaded_hs:
+            res["verifier_kv_last_global"] = torch.stack(
+                [loaded_hs["kv_last_global_k"], loaded_hs["kv_last_global_v"]], dim=1
+            )
         return res
 
 

From dc35c93373afa29eb3c634be17f3629e14f32b66 Mon Sep 17 00:00:00 2001
From: Beichen-Ma <mabeichen12@gmail.com>
Date: Sat, 27 Jun 2026 23:03:34 +0000
Subject: [PATCH 12/19] feat: add out-of-tree Gemma4 KV connector

---
 .../data_generation/gemma4_kv_connector.py    | 351 ++++++++++++++++++
 tests/e2e/run_gemma4_kv_extraction.py         | 126 +++++++
 tests/e2e/smoke/test_gemma4_kv_connector.py   |  34 ++
 tests/e2e/utils.py                            | 129 +++++++
 4 files changed, 640 insertions(+)
 create mode 100644 src/speculators/data_generation/gemma4_kv_connector.py
 create mode 100644 tests/e2e/run_gemma4_kv_extraction.py
 create mode 100644 tests/e2e/smoke/test_gemma4_kv_connector.py

diff --git a/src/speculators/data_generation/gemma4_kv_connector.py b/src/speculators/data_generation/gemma4_kv_connector.py
new file mode 100644
index 000000000..b61ea520b
--- /dev/null
+++ b/src/speculators/data_generation/gemma4_kv_connector.py
@@ -0,0 +1,351 @@
+"""Out-of-tree KV connector for Gemma4 MTP training-data extraction.
+
+Load it out-of-tree via --kv-transfer-config:
+
+    {
+      "kv_connector": "Gemma4KVConnector",
+      "kv_connector_module_path":
+        "speculators.data_generation.gemma4_kv_connector",
+      "kv_role": "kv_producer",
+      "kv_connector_extra_config": {"shared_storage_path": "<dir>"}
+    }
+"""
+
+import fcntl
+import os
+from collections import defaultdict
+from dataclasses import dataclass
+from functools import partial
+from typing import TYPE_CHECKING, Any
+
+import torch
+from vllm.distributed.communication_op import tensor_model_parallel_gather
+from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import (  # noqa: E501
+    ExampleHiddenStatesConnector,
+    PendingSave,
+    extract_from_kv_cache,
+)
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+_KV_CACHE_NDIM = 5
+_KV_CACHE_KV_AXIS_SIZE = 2
+
+GEMMA4_KV_KEYS = (
+    "kv_last_local_k",
+    "kv_last_local_v",
+    "kv_last_global_k",
+    "kv_last_global_v",
+)
+
+
+def extract_real_kv_from_cache(
+    kv_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    num_tokens: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Extract per-token K and V from a real attention layer's paged cache.
+
+    Args:
+        kv_cache: The layer's paged KV buffer, 5-D as described above.
+        slot_mapping: Per-token physical cache slots (1-D, length >= num_tokens).
+        num_tokens: Number of leading tokens to extract.
+
+    Returns:
+        A (K, V) tuple, each of shape
+        (num_tokens, num_kv_heads, head_size).
+    """
+    _bad_layout_msg = (
+        "Gemma4KVConnector expects the 5-D KV layout "
+        "(num_blocks, 2, block_size, num_kv_heads, head_size) used by the "
+        f"FlashAttention/Triton backends; got shape {tuple(kv_cache.shape)}."
+    )
+    assert kv_cache.dim() == _KV_CACHE_NDIM, _bad_layout_msg
+    assert kv_cache.shape[1] == _KV_CACHE_KV_AXIS_SIZE, _bad_layout_msg
+    block_size = kv_cache.shape[2]
+    block_idx = slot_mapping // block_size
+    block_off = slot_mapping % block_size
+    k = kv_cache[:, 0][block_idx, block_off][:num_tokens]
+    v = kv_cache[:, 1][block_idx, block_off][:num_tokens]
+    return k, v
+
+
+@dataclass
+class Gemma4PendingSave(PendingSave):
+    """PendingSave plus the per-group block_ids of the two verifier KV
+    layers, carried scheduler -> worker so their KV can be extracted."""
+
+    local_block_ids: list[int] | None = None
+    global_block_ids: list[int] | None = None
+
+
+class Gemma4KVConnector(ExampleHiddenStatesConnector):
+    """Hidden-states + verifier-KV extraction connector for Gemma4 MTP training."""
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: "KVConnectorRole",
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config, role, kv_cache_config)
+
+        self._tp_size = vllm_config.parallel_config.tensor_parallel_size
+        pp = vllm_config.parallel_config.pipeline_parallel_size
+        assert pp == 1, (
+            f"Gemma4KVConnector does not support pipeline_parallel_size>1 "
+            f"(got pp={pp})."
+        )
+
+        tcfg = vllm_config.model_config.hf_config.get_text_config()
+        self._local_total_kv_heads = tcfg.num_key_value_heads
+
+        global_kv = getattr(tcfg, "num_global_key_value_heads", None)
+        self._global_total_kv_heads = global_kv or tcfg.num_key_value_heads
+
+        cache_dtype = vllm_config.cache_config.cache_dtype
+        assert cache_dtype == "auto", (
+            "Gemma4KVConnector requires an unquantized KV cache "
+            f"(cache_dtype='auto'); got {cache_dtype!r}. Quantized-cache "
+            "handling (per-token head scales) is not yet supported."
+        )
+
+        self.verifier_local_layer, self._local_group_idx = self._resolve_layer(
+            "sliding_attention"
+        )
+        self.verifier_global_layer, self._global_group_idx = self._resolve_layer(
+            "full_attention"
+        )
+
+        self._local_kv_cache: torch.Tensor | None = None
+        self._global_kv_cache: torch.Tensor | None = None
+
+        logger.info(
+            "Gemma4KVConnector: local=%s (group %d) global=%s (group %d)",
+            self.verifier_local_layer,
+            self._local_group_idx,
+            self.verifier_global_layer,
+            self._global_group_idx,
+        )
+
+    def _resolve_layer(self, attn_type: str) -> tuple[str, int]:
+        """Resolve the last non-KV-shared verifier layer of attn_type.
+
+        Args:
+            attn_type: "sliding_attention" or "full_attention".
+
+        Returns:
+            A (layer_name, kv_cache_group_idx) tuple.
+        """
+        tcfg = self._vllm_config.model_config.hf_config.get_text_config()
+        layer_types = list(getattr(tcfg, "layer_types", []))
+        if not layer_types:
+            raise ValueError(
+                "Gemma4KVConnector requires the verifier config to expose "
+                "'layer_types'; none found."
+            )
+        n_shared = getattr(tcfg, "num_kv_shared_layers", 0)
+        num_non_shared = len(layer_types) - n_shared
+
+        type_to_indices: dict[str, list[int]] = defaultdict(list)
+        for idx, lt in enumerate(layer_types[:num_non_shared]):
+            type_to_indices[lt].append(idx)
+        indices = type_to_indices.get(attn_type)
+        if not indices:
+            counts = {k: len(v) for k, v in type_to_indices.items()}
+            raise ValueError(
+                f"Gemma4KVConnector found no non-KV-shared '{attn_type}' "
+                f"layer; layer-type counts: {counts}"
+            )
+
+        prefix = self._resolve_attn_layer_prefix()
+        layer_name = f"{prefix}.{indices[-1]}.self_attn.attn"
+        return layer_name, self._find_group_idx(layer_name)
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        super().register_kv_caches(kv_caches)
+        for name in (self.verifier_local_layer, self.verifier_global_layer):
+            if name not in kv_caches:
+                raise ValueError(
+                    f"Resolved verifier KV layer {name!r} is not among the "
+                    f"registered kv_caches keys: {sorted(kv_caches)[:8]}..."
+                )
+        self._local_kv_cache = kv_caches[self.verifier_local_layer]
+        self._global_kv_cache = kv_caches[self.verifier_global_layer]
+
+    def _resolve_attn_layer_prefix(self) -> str:
+        assert self._kv_cache_config is not None
+        for group in self._kv_cache_config.kv_cache_groups:
+            for name in group.layer_names:
+                if ".layers." in name and name.endswith(".self_attn.attn"):
+                    return name.split(".layers.")[0] + ".layers"
+        raise ValueError(
+            "Gemma4KVConnector could not resolve the verifier attention-layer "
+            "name prefix (no '*.layers.N.self_attn.attn' layer found)."
+        )
+
+    def _find_group_idx(self, layer_name: str) -> int:
+        assert self._kv_cache_config is not None
+        for i, group in enumerate(self._kv_cache_config.kv_cache_groups):
+            if layer_name in group.layer_names:
+                return i
+        raise ValueError(
+            f"Gemma4KVConnector: layer {layer_name!r} not found in any kv_cache_group."
+        )
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        result = super().request_finished_all_groups(request, block_ids)
+
+        req_id = request.request_id
+        base = self._pending_saves.get(req_id)
+        if base is not None:
+            self._pending_saves[req_id] = Gemma4PendingSave(
+                req_id=base.req_id,
+                filename=base.filename,
+                token_ids=base.token_ids,
+                block_ids=base.block_ids,
+                local_block_ids=list(block_ids[self._local_group_idx]),
+                global_block_ids=list(block_ids[self._global_group_idx]),
+            )
+        return result
+
+    @staticmethod
+    def _slot_mapping_from_blocks(
+        block_ids: list[int], block_size: int
+    ) -> torch.Tensor:
+        block_ids_t = torch.tensor(block_ids, dtype=torch.long)
+        num_blocks = block_ids_t.shape[0]
+        offsets = torch.arange(0, block_size, dtype=torch.long)
+        slot_mapping = (
+            offsets.reshape((1, block_size))
+            + block_ids_t.reshape((num_blocks, 1)) * block_size
+        )
+        return slot_mapping.flatten()
+
+    def _gather_kv_heads(
+        self, x: torch.Tensor, total_kv_heads: int
+    ) -> torch.Tensor | None:
+        """Gather a per-rank KV-head shard onto rank 0 of the TP group..
+
+        Args:
+            x: This rank's shard, shape (num_tokens, num_kv_heads_local, head_size).
+            total_kv_heads: The unsharded KV-head count for this layer, used to
+                dedup under GQA replication.
+
+        Returns:
+            (num_tokens, total_kv_heads, head_size) on rank 0; None on other ranks.
+        """
+        if self._tp_size == 1:
+            return x
+
+        gathered = tensor_model_parallel_gather(x, dst=0, dim=1)
+        if gathered is None:
+            return None
+
+        stride = gathered.shape[1] // total_kv_heads
+        if stride > 1:
+            gathered = gathered[:, ::stride, :]
+
+        return gathered.contiguous()
+
+    def _submit_async_write(self, pending: PendingSave) -> None:
+        num_tokens = pending.token_ids.shape[0]
+
+        gathered_kv: dict[str, torch.Tensor] = {}
+        if isinstance(pending, Gemma4PendingSave) and pending.local_block_ids:
+            for role, cache, blocks, total_heads in (
+                (
+                    "local",
+                    self._local_kv_cache,
+                    pending.local_block_ids,
+                    self._local_total_kv_heads,
+                ),
+                (
+                    "global",
+                    self._global_kv_cache,
+                    pending.global_block_ids,
+                    self._global_total_kv_heads,
+                ),
+            ):
+                assert cache is not None
+                assert blocks is not None
+                slots = self._slot_mapping_from_blocks(blocks, cache.shape[2]).to(
+                    device=cache.device, non_blocking=True
+                )
+                k, v = extract_real_kv_from_cache(cache, slots, num_tokens)
+                # Returns None on non-rank-0; only rank 0 keeps the result.
+                gk = self._gather_kv_heads(k, total_heads)
+                gv = self._gather_kv_heads(v, total_heads)
+                if gk is not None and gv is not None:
+                    gathered_kv[f"kv_last_{role}_k"] = gk
+                    gathered_kv[f"kv_last_{role}_v"] = gv
+
+        # Only rank 0 writes to disk.
+        if not self._is_tp_rank_zero:
+            return
+        assert self._kv_cache is not None
+        if not gathered_kv:
+            logger.warning(
+                "Gemma4KVConnector: req %s has no verifier block_ids; "
+                "writing hidden states only.",
+                pending.req_id,
+            )
+
+        copy_stream = self._get_copy_stream()
+        ready_event = torch.cuda.Event()
+        ready_event.record()
+        copy_stream.wait_event(ready_event)
+
+        tensors: dict[str, torch.Tensor] = {}
+        with torch.cuda.stream(copy_stream):
+            hs_slots = self._slot_mapping_from_blocks(
+                pending.block_ids, self._kv_cache.shape[1]
+            ).to(device=self._kv_cache.device, non_blocking=True)
+            tensors["hidden_states"] = self._pin(
+                extract_from_kv_cache(self._kv_cache, hs_slots, num_tokens)
+            )
+            # Pin the already-gathered (full-head) verifier KV tensors.
+            for key, full in gathered_kv.items():
+                tensors[key] = self._pin(full)
+
+        copy_done = torch.cuda.Event()
+        copy_done.record(copy_stream)
+
+        assert not pending.token_ids.is_cuda
+        tensors["token_ids"] = pending.token_ids.clone()
+
+        prior = self._req_futures.get(pending.req_id)
+        assert prior is None, "Found another KV transfer request with same req_id!"
+
+        os.makedirs(os.path.dirname(pending.filename), exist_ok=True)
+        lock_fd = self._lock_fds.pop(pending.req_id, None)
+        if lock_fd is None and self.use_lock:
+            lock_path = pending.filename + ".lock"
+            lock_fd = os.open(lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644)
+            fcntl.flock(lock_fd, fcntl.LOCK_EX)
+
+        future = self._executor.submit(
+            self._write_tensors, tensors, copy_done, pending.filename, lock_fd
+        )
+        self._req_copy_events[pending.req_id] = copy_done
+        self._req_futures[pending.req_id] = future
+        future.add_done_callback(partial(self._on_write_done, pending.req_id))
+
+    @staticmethod
+    def _pin(gpu_tensor: torch.Tensor) -> torch.Tensor:
+        """Async DtoH copy into pinned host memory. Call inside the copy
+        stream so the copy is ordered with the other per-request copies."""
+        pinned = torch.empty_like(gpu_tensor, device="cpu", pin_memory=True)
+        pinned.copy_(gpu_tensor, non_blocking=True)
+        return pinned
diff --git a/tests/e2e/run_gemma4_kv_extraction.py b/tests/e2e/run_gemma4_kv_extraction.py
new file mode 100644
index 000000000..beac8624a
--- /dev/null
+++ b/tests/e2e/run_gemma4_kv_extraction.py
@@ -0,0 +1,126 @@
+"""Subprocess driver for the Gemma4KVConnector extraction test.
+
+Generic, assertion-free counterpart to run_vllm.py: it is parametrized by
+JSON (--llm-args, --prompts) and dumps raw per-request facts (tensor
+shapes, token alignment, finiteness) to --results-file. All assertions live
+parent-side in tests/e2e/utils.run_gemma4_kv_extraction.
+"""
+
+import argparse
+import json
+import sys
+import traceback
+from pathlib import Path
+from typing import Any
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+# isort: off
+from tests.e2e.run_vllm import LLM, SamplingParams  # noqa: E402
+
+import torch  # noqa: E402
+from vllm.distributed.kv_transfer.kv_connector.v1 import (  # noqa: E402
+    example_hidden_states_connector,
+)
+
+# isort: on
+
+from speculators.data_generation.gemma4_kv_connector import (  # noqa: E402
+    GEMMA4_KV_KEYS as KV_KEYS,
+)
+
+
+def run_extraction(llm_args: dict, prompts: list[str]) -> dict:
+    """Run the extraction forward pass and collect raw per-request facts.
+
+    No assertions: this only reports JSON-serializable facts (shapes, token
+    alignment, finiteness) for the parent process to assert on.
+    """
+    llm = LLM(**llm_args)
+
+    model_config = llm.llm_engine.model_config
+    hidden_size = model_config.get_hidden_size()
+    text_config = model_config.hf_config.get_text_config()
+    local_total_heads = text_config.num_key_value_heads
+    global_total_heads = (
+        getattr(text_config, "num_global_key_value_heads", None) or local_total_heads
+    )
+
+    outputs = llm.generate(prompts, SamplingParams(max_tokens=1, temperature=0.0))
+
+    per_request: list[dict[str, Any]] = []
+    for output in outputs:
+        kv_params = output.kv_transfer_params
+        path = kv_params.get("hidden_states_path") if kv_params else None
+        if path is None:
+            per_request.append({"hidden_states_path": None})
+            continue
+
+        obj = example_hidden_states_connector.load_hidden_states(path)
+        n = len(output.prompt_token_ids)
+        per_request.append(
+            {
+                "hidden_states_path": path,
+                "num_tokens": n,
+                "token_ids_aligned": bool(
+                    torch.equal(obj["token_ids"], torch.tensor(output.prompt_token_ids))
+                ),
+                "hidden_states_shape": list(obj["hidden_states"].shape),
+                "kv": {
+                    key: {
+                        "present": key in obj,
+                        "shape": list(obj[key].shape) if key in obj else None,
+                        "finite": bool(torch.isfinite(obj[key]).all())
+                        if key in obj
+                        else None,
+                    }
+                    for key in KV_KEYS
+                },
+            }
+        )
+
+    return {
+        "num_outputs": len(outputs),
+        "hidden_size": hidden_size,
+        "local_total_heads": local_total_heads,
+        "global_total_heads": global_total_heads,
+        "per_request": per_request,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--llm-args",
+        type=str,
+        required=True,
+        help="JSON-serialized kwargs for LLM instantiation",
+    )
+    parser.add_argument(
+        "--prompts", type=str, required=True, help="JSON-serialized list of prompts"
+    )
+    parser.add_argument(
+        "--results-file",
+        type=Path,
+        required=True,
+        help="File to write the JSON-serialized extraction facts",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    try:
+        results = run_extraction(json.loads(args.llm_args), json.loads(args.prompts))
+        exit_code = 0
+    except Exception as exc:  # noqa: BLE001
+        traceback.print_exc()
+        results = {"ok": False, "error": f"{type(exc).__name__}: {exc}"}
+        exit_code = 1
+
+    with args.results_file.open("w", encoding="utf-8") as f:
+        json.dump(results, f)
+
+    sys.exit(exit_code)
diff --git a/tests/e2e/smoke/test_gemma4_kv_connector.py b/tests/e2e/smoke/test_gemma4_kv_connector.py
new file mode 100644
index 000000000..d54e30b38
--- /dev/null
+++ b/tests/e2e/smoke/test_gemma4_kv_connector.py
@@ -0,0 +1,34 @@
+"""E2E test for the out-of-tree Gemma4KVConnector.
+
+Gemma4KVConnector extends vLLM's ExampleHiddenStatesConnector to also export the
+verifier's last sliding-window (local) and last full-attention (global) layer KV
+caches alongside the hidden states, into one .safetensors per request. This is
+the training-data producer for finetuning Gemma4 MTP draft models, whose
+query-only attention borrows exactly those two verifier KV caches.
+"""
+
+from pathlib import Path
+
+import pytest
+
+from tests.conftest import requires_cuda, requires_multi_gpu
+from tests.e2e.utils import run_gemma4_kv_extraction
+
+MODEL = "google/gemma-4-E4B-it"
+
+
+@pytest.mark.e2e
+@pytest.mark.slow
+@requires_cuda
+def test_gemma4_kv_connector_smoke(tmp_path: Path):
+    """Single-GPU (tp=1) extraction: shapes, alignment, both verifier KV caches."""
+    run_gemma4_kv_extraction(MODEL, tmp_path, tensor_parallel_size=1)
+
+
+@pytest.mark.e2e
+@pytest.mark.slow
+@requires_multi_gpu
+def test_gemma4_kv_connector_tp2(tmp_path: Path):
+    """tp=2: verifier KV heads are sharded per rank; the connector all-gathers
+    (and dedups GQA replicas) so saved tensors carry the full head count."""
+    run_gemma4_kv_extraction(MODEL, tmp_path, tensor_parallel_size=2)
diff --git a/tests/e2e/utils.py b/tests/e2e/utils.py
index 0891818e6..04462e490 100644
--- a/tests/e2e/utils.py
+++ b/tests/e2e/utils.py
@@ -23,6 +23,7 @@
     "launch_vllm_server_context",
     "purge_newfiles",
     "run_data_generation_offline",
+    "run_gemma4_kv_extraction",
     "run_prepare_data",
     "run_stitch_mtp",
     "run_training",
@@ -31,6 +32,13 @@
     "wait_for_server",
 ]
 
+GEMMA4_KV_KEYS = (
+    "kv_last_local_k",
+    "kv_last_local_v",
+    "kv_last_global_k",
+    "kv_last_global_v",
+)
+
 
 def purge_newfiles(fn: Callable[..., Path]):
     """Decorator that turns a Path-returning function into a context manager.
@@ -452,3 +460,124 @@ def run_vllm_engine(
             assert acci >= thresholdi, (
                 f"Acceptance {acci} at token {i} is less than threshold {thresholdi}"
             )
+
+
+def run_gemma4_kv_extraction(
+    model: str,
+    tmp_path: Path,
+    tensor_parallel_size: int = 1,
+    aux_layer_ids: Iterable[int] = (2, 21, 39, 42),
+    max_model_len: int = 2048,
+    gpu_memory_utilization: float = 0.4,
+    timeout: float | None = None,
+):
+    """Drive Gemma4KVConnector extraction via run_gemma4_kv_extraction.py."""
+    logger.info("vLLM Python executable: {}", VLLM_PYTHON)
+
+    driver_file = str(Path(__file__).with_name("run_gemma4_kv_extraction.py"))
+    shared_storage = tmp_path / "kv_out"
+    shared_storage.mkdir(exist_ok=True)
+    results_file = tmp_path / "results.json"
+    aux_layer_ids = list(aux_layer_ids)
+
+    llm_args_dict = {
+        "model": model,
+        "tensor_parallel_size": tensor_parallel_size,
+        "speculative_config": {
+            "method": "extract_hidden_states",
+            "num_speculative_tokens": 1,
+            "draft_model_config": {
+                "hf_config": {"eagle_aux_hidden_state_layer_ids": aux_layer_ids}
+            },
+        },
+        "kv_transfer_config": {
+            "kv_connector": "Gemma4KVConnector",
+            "kv_connector_module_path": (
+                "speculators.data_generation.gemma4_kv_connector"
+            ),
+            "kv_role": "kv_producer",
+            "kv_connector_extra_config": {"shared_storage_path": str(shared_storage)},
+        },
+        "max_model_len": max_model_len,
+        "enforce_eager": True,
+        "enable_chunked_prefill": False,
+        "gpu_memory_utilization": gpu_memory_utilization,
+        "load_format": "dummy",
+    }
+
+    # Short, and a long prompt past the 512 sliding window, in one batch.
+    prompts = [
+        "Hello world",
+        "Test prompt with several tokens",
+        " ".join(["word"] * 600),
+    ]
+
+    command = [
+        VLLM_PYTHON,
+        driver_file,
+        "--llm-args",
+        json.dumps(llm_args_dict),
+        "--prompts",
+        json.dumps(prompts),
+        "--results-file",
+        str(results_file),
+    ]
+    logger.info("run_gemma4_kv_extraction.py command:\n    {}", command)
+
+    env = os.environ.copy()
+    # Avoid a flashinfer JIT-build failure.
+    env["VLLM_USE_FLASHINFER_SAMPLER"] = "0"
+
+    result = subprocess.run(  # noqa: S603
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        check=False,
+        env=env,
+        timeout=timeout,
+    )
+    logger.info(
+        "run_gemma4_kv_extraction.py output:\n{}", indent(result.stdout, "    ")
+    )
+    assert result.returncode == 0, (
+        f"run_gemma4_kv_extraction.py exited with non-zero return code: "
+        f"{result.returncode}"
+    )
+
+    with results_file.open(encoding="utf-8") as f:
+        results = json.load(f)
+
+    assert results["num_outputs"] == len(prompts)
+    assert len(results["per_request"]) == len(prompts)
+    expected_heads = {
+        "kv_last_local_k": results["local_total_heads"],
+        "kv_last_local_v": results["local_total_heads"],
+        "kv_last_global_k": results["global_total_heads"],
+        "kv_last_global_v": results["global_total_heads"],
+    }
+    hidden_size = results["hidden_size"]
+    for req in results["per_request"]:
+        path = req.get("hidden_states_path")
+        assert path is not None, "request produced no hidden_states_path"
+        n = req["num_tokens"]
+        assert req["token_ids_aligned"], "saved token_ids do not match prompt tokens"
+        assert req["hidden_states_shape"] == [n, len(aux_layer_ids), hidden_size]
+        for key in GEMMA4_KV_KEYS:
+            kv = req["kv"][key]
+            assert kv["present"], f"missing {key}"
+            shape = kv["shape"]
+            assert len(shape) == 3, f"{key} expected 3-D, got {shape}"
+            assert shape[0] == n, f"{key} rows {shape[0]} != n_tokens {n}"
+            # full (unsharded) head count even under TP
+            assert shape[1] == expected_heads[key], (
+                f"{key} head count {shape[1]} != unsharded {expected_heads[key]} "
+                f"(tp={tensor_parallel_size})"
+            )
+            assert kv["finite"], f"{key} has non-finite values"
+        # local (sliding) and global (full) caches use different head dims.
+        local_shape = req["kv"]["kv_last_local_k"]["shape"]
+        global_shape = req["kv"]["kv_last_global_k"]["shape"]
+        assert local_shape[1:] != global_shape[1:]
+
+    return results

From 68fecf11ddb8a56d108e3b5d273f0fa88b6c759c Mon Sep 17 00:00:00 2001
From: Beichen-Ma <mabeichen12@gmail.com>
Date: Sun, 28 Jun 2026 04:16:52 +0000
Subject: [PATCH 13/19] fix

---
 .../data_generation/gemma4_kv_connector.py    | 33 ++++++++++---------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/src/speculators/data_generation/gemma4_kv_connector.py b/src/speculators/data_generation/gemma4_kv_connector.py
index b61ea520b..5fe8dc72e 100644
--- a/src/speculators/data_generation/gemma4_kv_connector.py
+++ b/src/speculators/data_generation/gemma4_kv_connector.py
@@ -62,13 +62,12 @@ def extract_real_kv_from_cache(
         A (K, V) tuple, each of shape
         (num_tokens, num_kv_heads, head_size).
     """
-    _bad_layout_msg = (
-        "Gemma4KVConnector expects the 5-D KV layout "
-        "(num_blocks, 2, block_size, num_kv_heads, head_size) used by the "
-        f"FlashAttention/Triton backends; got shape {tuple(kv_cache.shape)}."
-    )
-    assert kv_cache.dim() == _KV_CACHE_NDIM, _bad_layout_msg
-    assert kv_cache.shape[1] == _KV_CACHE_KV_AXIS_SIZE, _bad_layout_msg
+    if kv_cache.dim() != _KV_CACHE_NDIM or kv_cache.shape[1] != _KV_CACHE_KV_AXIS_SIZE:
+        raise ValueError(
+            "Gemma4KVConnector expects the 5-D KV layout "
+            "(num_blocks, 2, block_size, num_kv_heads, head_size) used by the "
+            f"FlashAttention/Triton backends; got shape {tuple(kv_cache.shape)}."
+        )
     block_size = kv_cache.shape[2]
     block_idx = slot_mapping // block_size
     block_off = slot_mapping % block_size
@@ -99,10 +98,11 @@ def __init__(
 
         self._tp_size = vllm_config.parallel_config.tensor_parallel_size
         pp = vllm_config.parallel_config.pipeline_parallel_size
-        assert pp == 1, (
-            f"Gemma4KVConnector does not support pipeline_parallel_size>1 "
-            f"(got pp={pp})."
-        )
+        if pp != 1:
+            raise ValueError(
+                f"Gemma4KVConnector does not support pipeline_parallel_size>1 "
+                f"(got pp={pp})."
+            )
 
         tcfg = vllm_config.model_config.hf_config.get_text_config()
         self._local_total_kv_heads = tcfg.num_key_value_heads
@@ -111,11 +111,12 @@ def __init__(
         self._global_total_kv_heads = global_kv or tcfg.num_key_value_heads
 
         cache_dtype = vllm_config.cache_config.cache_dtype
-        assert cache_dtype == "auto", (
-            "Gemma4KVConnector requires an unquantized KV cache "
-            f"(cache_dtype='auto'); got {cache_dtype!r}. Quantized-cache "
-            "handling (per-token head scales) is not yet supported."
-        )
+        if cache_dtype != "auto":
+            raise ValueError(
+                "Gemma4KVConnector requires an unquantized KV cache "
+                f"(cache_dtype='auto'); got {cache_dtype!r}. Quantized-cache "
+                "handling (per-token head scales) is not yet supported."
+            )
 
         self.verifier_local_layer, self._local_group_idx = self._resolve_layer(
             "sliding_attention"

From cd06a8ce68ea38cd0b8339b8b18697a8a5d5f3cb Mon Sep 17 00:00:00 2001
From: Beichen-Ma <mabeichen12@gmail.com>
Date: Sun, 28 Jun 2026 09:13:58 +0000
Subject: [PATCH 14/19] fix: trim slot_mapping before indexing

---
 src/speculators/data_generation/gemma4_kv_connector.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/speculators/data_generation/gemma4_kv_connector.py b/src/speculators/data_generation/gemma4_kv_connector.py
index 5fe8dc72e..ebaefcfe3 100644
--- a/src/speculators/data_generation/gemma4_kv_connector.py
+++ b/src/speculators/data_generation/gemma4_kv_connector.py
@@ -69,10 +69,12 @@ def extract_real_kv_from_cache(
             f"FlashAttention/Triton backends; got shape {tuple(kv_cache.shape)}."
         )
     block_size = kv_cache.shape[2]
+
+    slot_mapping = slot_mapping[:num_tokens]
     block_idx = slot_mapping // block_size
     block_off = slot_mapping % block_size
-    k = kv_cache[:, 0][block_idx, block_off][:num_tokens]
-    v = kv_cache[:, 1][block_idx, block_off][:num_tokens]
+    k = kv_cache[:, 0][block_idx, block_off]
+    v = kv_cache[:, 1][block_idx, block_off]
     return k, v
 
 

From 4e325d9741deaaec3798d7b9404377040f6d37e0 Mon Sep 17 00:00:00 2001
From: shotsan <shotsan@users.noreply.github.com>
Date: Tue, 14 Jul 2026 14:12:27 -0700
Subject: [PATCH 15/19] fix(data): stack separate k and v tensors from PR 674
 dataset format

Signed-off-by: shotsan <shotsan@users.noreply.github.com>
---
 tests/unit/train/test_data.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit/train/test_data.py b/tests/unit/train/test_data.py
index 271795606..697ace4ae 100644
--- a/tests/unit/train/test_data.py
+++ b/tests/unit/train/test_data.py
@@ -494,6 +494,7 @@ def test_verifier_kvs_survive_pipeline():
     assert "verifier_kv_last_local" in standardized
     assert "verifier_kv_last_global" in standardized
 
+
     expected_local_stack = torch.stack(
         [v1_data["kv_last_local_k"], v1_data["kv_last_local_v"]], dim=1  # type: ignore[arg-type]
     )

From aa6447b91ccbe7ac3ad72196d8ce5a9b14de8fb6 Mon Sep 17 00:00:00 2001
From: Roderick Wu <Roderick-Wu@a100-06.nemg-001.lab.rdu2.dc.redhat.com>
Date: Mon, 20 Jul 2026 14:36:13 -0400
Subject: [PATCH 16/19] add gemma4 components

---
 .gitignore                                    |   5 +-
 src/speculators/models/base_components.py     |  16 +++
 src/speculators/models/mtp/core.py            |  22 +++-
 src/speculators/models/mtp/data.py            |   6 +-
 .../models/mtp/model_definitions.py           | 114 ++++++++++++++++++
 test.sh                                       |  20 +++
 tests/unit/models/test_mtp_data.py            |  35 ++++++
 7 files changed, 210 insertions(+), 8 deletions(-)
 create mode 100644 test.sh

diff --git a/.gitignore b/.gitignore
index cb421e12c..f93acc0dd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -173,4 +173,7 @@ cython_debug/
 .DS_Store
 
 # Visual Studio Code
-.vscode/
\ No newline at end of file
+.vscode/
+
+output/
+output/*
diff --git a/src/speculators/models/base_components.py b/src/speculators/models/base_components.py
index cc5e2120a..6f137a041 100644
--- a/src/speculators/models/base_components.py
+++ b/src/speculators/models/base_components.py
@@ -69,6 +69,22 @@ class ModelComponents(NamedTuple):
         Gemma2RotaryEmbedding,
     )
 
+try:
+    from transformers.models.gemma4.modeling_gemma4 import (
+        Gemma4RMSNorm,
+        Gemma4TextDecoderLayer,
+        Gemma4TextRotaryEmbedding,
+    )
+
+    model_classes["gemma4_text"] = ModelComponents(
+        Gemma4TextDecoderLayer,
+        Gemma4TextDecoderLayer,
+        Gemma4RMSNorm,
+        Gemma4TextRotaryEmbedding,
+    )
+except ImportError:
+    pass
+
 try:
     from transformers.models.qwen3_next.modeling_qwen3_next import (
         Qwen3NextDecoderLayer,
diff --git a/src/speculators/models/mtp/core.py b/src/speculators/models/mtp/core.py
index c6d597552..457cc430a 100644
--- a/src/speculators/models/mtp/core.py
+++ b/src/speculators/models/mtp/core.py
@@ -263,8 +263,8 @@ def from_training_args(  # type: ignore[override]
         if verifier_name_or_path is None:
             raise ValueError(
                 "verifier_name_or_path is required for MTP training. "
-                "The verifier model must contain native MTP weights "
-                "(mtp.* keys) to extract."
+                "It is used to load shared weights (embed_tokens, lm_head) "
+                "and, when available, native MTP layer weights."
             )
 
         config = MTPSpeculatorConfig(
@@ -290,10 +290,20 @@ def from_training_args(  # type: ignore[override]
 
         from speculators.convert.mtp.converter import MTPConverter  # noqa: PLC0415
 
-        state_dict = MTPConverter().convert_to_state_dict(
-            verifier_name_or_path  # type: ignore[arg-type]
-        )
-        model.load_state_dict(state_dict, strict=False)
+        try:
+            state_dict = MTPConverter().convert_to_state_dict(
+                verifier_name_or_path  # type: ignore[arg-type]
+            )
+            model.load_state_dict(state_dict, strict=False)
+        except ValueError as exc:
+            if "No keys with prefix" not in str(exc):
+                raise
+            logger.warning(
+                "Verifier at '%s' has no native MTP weights (mtp.* keys). "
+                "MTP layer will be randomly initialized for training from "
+                "scratch.",
+                verifier_name_or_path,
+            )
 
         model.load_verifier_weights()
         return model
diff --git a/src/speculators/models/mtp/data.py b/src/speculators/models/mtp/data.py
index 2ab74647c..32fd0eec6 100644
--- a/src/speculators/models/mtp/data.py
+++ b/src/speculators/models/mtp/data.py
@@ -9,10 +9,14 @@ def shift_batch_mtp(batch: BatchType) -> BatchType:
     No token-level shifting — the MTP forward pass handles alignment
     internally via per-step offset slicing of input_ids.
     """
-    return {
+    result = {
         "input_ids": batch["input_ids"],
         "hidden_states": batch["verifier_last_hidden_states"],
         "loss_mask": batch["loss_mask"],
         "lengths": batch["lengths"],
         "position_ids": batch["position_ids"],
     }
+    for k in ("verifier_kv_last_local", "verifier_kv_last_global"):
+        if k in batch:
+            result[k] = batch[k]
+    return result
diff --git a/src/speculators/models/mtp/model_definitions.py b/src/speculators/models/mtp/model_definitions.py
index cd0ba0eee..5d7a653a2 100644
--- a/src/speculators/models/mtp/model_definitions.py
+++ b/src/speculators/models/mtp/model_definitions.py
@@ -298,3 +298,117 @@ def __init__(self, config: PretrainedConfig, layer_idx: int = 0) -> None:
     mtp_model_classes["gemma2"] = base_components.override_components(
         "gemma2", first_layer_class=Gemma2MTPLayer
     )
+
+if "gemma4_text" in base_components.model_classes:
+    from typing import Optional, Tuple, Any
+    from transformers.models.gemma4.modeling_gemma4 import (
+        Gemma4TextDecoderLayer,
+        Gemma4RMSNorm,
+        Gemma4TextAttention,
+        Gemma4TextRotaryEmbedding,
+        apply_rotary_pos_emb as gemma4_apply_rotary_pos_emb,
+        eager_attention_forward as gemma4_eager_attention_forward,
+    )
+
+    class Gemma4MTPRotaryEmbedding(Gemma4TextRotaryEmbedding):
+        """Wraps Gemma4TextRotaryEmbedding with a fixed layer_type for MTP.
+
+        Gemma4TextRotaryEmbedding.forward requires a layer_type arg
+        (per-type RoPE params), but MTPDraftModel.forward calls
+        rotary_emb(x, position_ids) without one. This wrapper fixes
+        the layer_type to the MTP layer's type (layer_types[0]).
+        """
+
+        def __init__(self, config: PretrainedConfig, device=None) -> None:
+            super().__init__(config, device)
+            self._mtp_layer_type = config.layer_types[0]
+
+        def forward(
+            self, x: torch.Tensor, position_ids: torch.Tensor, layer_type: str | None = None
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            return super().forward(x, position_ids, layer_type=self._mtp_layer_type)
+
+    class QueryOnlyGemma4TextAttention(Gemma4TextAttention):
+        def __init__(self, config: PretrainedConfig, layer_idx: int = 0) -> None:
+            super().__init__(config, layer_idx)
+            self.k_proj = None
+            self.v_proj = None
+            if hasattr(self, "k_norm"):
+                self.k_norm = None
+            if hasattr(self, "v_norm"):
+                self.v_norm = None
+            self._num_kv_heads = config.num_attention_heads // self.num_key_value_groups
+
+        def forward(
+            self,
+            hidden_states: torch.Tensor,
+            position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+            attention_mask: Optional[torch.Tensor],
+            verifier_kv_last_local: Optional[torch.Tensor] = None,
+            verifier_kv_last_global: Optional[torch.Tensor] = None,
+            **kwargs: Any,
+        ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, self.head_dim)
+
+            query_states = self.q_proj(hidden_states).view(hidden_shape)
+            query_states = self.q_norm(query_states)
+
+            cos, sin = position_embeddings
+            query_states = gemma4_apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2)
+            query_states = query_states.transpose(1, 2)
+
+            kv_tensor = verifier_kv_last_local if self.is_sliding else verifier_kv_last_global
+
+            if kv_tensor is None:
+                batch_sz, seq_len = input_shape
+                kv_tensor = torch.zeros(
+                    (batch_sz, seq_len, 2, self._num_kv_heads, self.head_dim),
+                    dtype=query_states.dtype,
+                    device=query_states.device,
+                )
+            elif kv_tensor.dim() == 3:
+                batch_sz, seq_len, _ = kv_tensor.shape
+                kv_tensor = kv_tensor.view(
+                    batch_sz, seq_len, 2, self._num_kv_heads, self.head_dim
+                )
+
+            key_states = kv_tensor[..., 0, :, :].transpose(1, 2)
+            value_states = kv_tensor[..., 1, :, :].transpose(1, 2)
+
+            if self.is_sliding and attention_mask is not None:
+                seq_len = attention_mask.shape[-1]
+                window_mask = torch.tril(
+                    torch.ones(seq_len, seq_len, dtype=torch.bool, device=attention_mask.device)
+                )
+                window_mask = torch.triu(window_mask, diagonal=-self.sliding_window + 1)
+                attention_mask = attention_mask.masked_fill(
+                    ~window_mask.view(1, 1, seq_len, seq_len),
+                    torch.finfo(query_states.dtype).min,
+                )
+
+            attn_output, attn_weights = gemma4_eager_attention_forward(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=self.attention_dropout if self.training else 0.0,
+                scaling=self.scaling,
+            )
+
+            attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+            attn_output = self.o_proj(attn_output)
+            return attn_output, attn_weights
+
+    class Gemma4TextMTPLayer(MTPLayerMixin, Gemma4TextDecoderLayer):
+        def __init__(self, config: PretrainedConfig, layer_idx: int = 0) -> None:
+            super().__init__(config, layer_idx)
+            self._setup_mtp_modules(config, Gemma4RMSNorm)
+            self.self_attn = QueryOnlyGemma4TextAttention(config, layer_idx)
+
+    mtp_model_classes["gemma4_text"] = base_components.override_components(
+        "gemma4_text",
+        first_layer_class=Gemma4TextMTPLayer,
+        rotary_emb_class=Gemma4MTPRotaryEmbedding,
+    )
diff --git a/test.sh b/test.sh
new file mode 100644
index 000000000..f0230f1c5
--- /dev/null
+++ b/test.sh
@@ -0,0 +1,20 @@
+
+ hf download   inference-optimization/gemma4-31B-responses magpie_gemma-4-31B-it.jsonl   --repo-type dataset   --local-dir ./output/dataset
+
+ python scripts/prepare_data.py   --model google/gemma-4-31B-it   --data ./output/dataset/magpie_gemma-4-31B-it.jsonl   --output ./output/mtp_gemma-4-31B-it   --max-samples 5000   --seq-length 8192
+
+CUDA_VISIBLE_DEVICES=0,1 python scripts/launch_vllm.py   google/gemma-4-31B-it  -- --port 8000 --tensor-parallel-size 2 --max-model-len 8192
+
+CUDA_VISIBLE_DEVICES=2,3 python scripts/train.py \
+  --verifier-name-or-path google/gemma-4-31B-it \
+  --data-path ./output/mtp_gemma-4-31B-it \
+  --vllm-endpoint http://localhost:8000/v1 \
+  --save-path ./output/mtp_qwen3_5_9b/checkpoints \
+  --speculator-type mtp \
+  --num-speculative-steps 3 \
+  --step-weight-beta 0.6 \
+  --epochs 3 \
+  --lr 1e-4 \
+  --total-seq-len 8192 \
+  --on-missing generate \
+  --on-generate delete
\ No newline at end of file
diff --git a/tests/unit/models/test_mtp_data.py b/tests/unit/models/test_mtp_data.py
index 05adec23e..2a6f79738 100644
--- a/tests/unit/models/test_mtp_data.py
+++ b/tests/unit/models/test_mtp_data.py
@@ -23,3 +23,38 @@ def test_shift_batch_mtp():
 
     for key in ("input_ids", "loss_mask", "lengths", "position_ids"):
         assert torch.equal(result[key], batch[key])
+
+
+def test_shift_batch_mtp_forwards_verifier_kv():
+    seq_len, hidden_size = 32, 128
+    kv_dim = 64
+    batch = {
+        "input_ids": torch.randint(0, 1000, (seq_len,)),
+        "verifier_last_hidden_states": torch.randn(seq_len, hidden_size),
+        "loss_mask": torch.ones(seq_len),
+        "lengths": torch.tensor([seq_len]),
+        "position_ids": torch.arange(seq_len),
+        "verifier_kv_last_local": torch.randn(seq_len, 2, kv_dim),
+        "verifier_kv_last_global": torch.randn(seq_len, 2, kv_dim),
+    }
+
+    result = shift_batch_mtp(batch)
+
+    assert torch.equal(result["verifier_kv_last_local"], batch["verifier_kv_last_local"])
+    assert torch.equal(result["verifier_kv_last_global"], batch["verifier_kv_last_global"])
+
+
+def test_shift_batch_mtp_without_kv():
+    seq_len, hidden_size = 32, 128
+    batch = {
+        "input_ids": torch.randint(0, 1000, (seq_len,)),
+        "verifier_last_hidden_states": torch.randn(seq_len, hidden_size),
+        "loss_mask": torch.ones(seq_len),
+        "lengths": torch.tensor([seq_len]),
+        "position_ids": torch.arange(seq_len),
+    }
+
+    result = shift_batch_mtp(batch)
+
+    assert "verifier_kv_last_local" not in result
+    assert "verifier_kv_last_global" not in result

From 7ac04cece9051faf0cad6efb8029e72ca562b4d6 Mon Sep 17 00:00:00 2001
From: Roderick Wu <Roderick-Wu@a100-06.nemg-001.lab.rdu2.dc.redhat.com>
Date: Mon, 20 Jul 2026 15:39:14 -0400
Subject: [PATCH 17/19] fixes

---
 src/speculators/models/mtp/data.py              | 2 +-
 src/speculators/models/mtp/model_definitions.py | 5 ++---
 test.sh                                         | 2 +-
 3 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/speculators/models/mtp/data.py b/src/speculators/models/mtp/data.py
index 32fd0eec6..4da2a65b5 100644
--- a/src/speculators/models/mtp/data.py
+++ b/src/speculators/models/mtp/data.py
@@ -17,6 +17,6 @@ def shift_batch_mtp(batch: BatchType) -> BatchType:
         "position_ids": batch["position_ids"],
     }
     for k in ("verifier_kv_last_local", "verifier_kv_last_global"):
-        if k in batch:
+        if k in batch and batch[k].dim() >= 2:
             result[k] = batch[k]
     return result
diff --git a/src/speculators/models/mtp/model_definitions.py b/src/speculators/models/mtp/model_definitions.py
index 5d7a653a2..94b8068c8 100644
--- a/src/speculators/models/mtp/model_definitions.py
+++ b/src/speculators/models/mtp/model_definitions.py
@@ -241,8 +241,7 @@ def forward(
             is_local = (self.sliding_window is not None and self.sliding_window > 0)
             kv_tensor = verifier_kv_last_local if is_local else verifier_kv_last_global
 
-            if kv_tensor is None:
-                # Fallback to zeros if not provided (e.g. dummy forward passes)
+            if kv_tensor is None or kv_tensor.dim() < 3:
                 batch_sz, seq_len = input_shape
                 kv_tensor = torch.zeros(
                     (batch_sz, seq_len, 2, self.num_key_value_heads, self.head_dim),
@@ -360,7 +359,7 @@ def forward(
 
             kv_tensor = verifier_kv_last_local if self.is_sliding else verifier_kv_last_global
 
-            if kv_tensor is None:
+            if kv_tensor is None or kv_tensor.dim() < 3:
                 batch_sz, seq_len = input_shape
                 kv_tensor = torch.zeros(
                     (batch_sz, seq_len, 2, self._num_kv_heads, self.head_dim),
diff --git a/test.sh b/test.sh
index f0230f1c5..21e5b7c1d 100644
--- a/test.sh
+++ b/test.sh
@@ -3,7 +3,7 @@
 
  python scripts/prepare_data.py   --model google/gemma-4-31B-it   --data ./output/dataset/magpie_gemma-4-31B-it.jsonl   --output ./output/mtp_gemma-4-31B-it   --max-samples 5000   --seq-length 8192
 
-CUDA_VISIBLE_DEVICES=0,1 python scripts/launch_vllm.py   google/gemma-4-31B-it  -- --port 8000 --tensor-parallel-size 2 --max-model-len 8192
+CUDA_VISIBLE_DEVICES=0,1 python scripts/launch_vllm.py   google/gemma-4-31B-it --hidden-states-path /tmp/hidden_states_$(whoami )-- --port 8000  --tensor-parallel-size 2 --max-model-len 8192
 
 CUDA_VISIBLE_DEVICES=2,3 python scripts/train.py \
   --verifier-name-or-path google/gemma-4-31B-it \

From 76325a3ba3f8fc4259f100ed196c754f5aa979fd Mon Sep 17 00:00:00 2001
From: Roderick Wu <Roderick-Wu@a100-06.nemg-001.lab.rdu2.dc.redhat.com>
Date: Tue, 21 Jul 2026 23:21:54 -0400
Subject: [PATCH 18/19] working training

---
 scripts/stitch_mtp.py         | 22 ++++++++++++++++------
 src/speculators/train/data.py | 10 +++++++++-
 test.sh                       | 12 ++++++++++--
 3 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/scripts/stitch_mtp.py b/scripts/stitch_mtp.py
index 0b3fd4a8d..5ccc416aa 100644
--- a/scripts/stitch_mtp.py
+++ b/scripts/stitch_mtp.py
@@ -206,15 +206,25 @@ def _stitch_sharded(
     weight_map: dict[str, str] = index_data["weight_map"]
 
     shard_to_new: dict[str, dict[str, torch.Tensor]] = {}
+    new_keys: dict[str, torch.Tensor] = {}
     for key, tensor in native_weights.items():
         shard = weight_map.get(key)
         if shard is None:
-            raise ValueError(
-                f"Finetuned key '{key}' not found in verifier weight "
-                "map. The finetuned checkpoint may not match the "
-                "verifier."
-            )
-        shard_to_new.setdefault(shard, {})[key] = tensor
+            new_keys[key] = tensor
+        else:
+            shard_to_new.setdefault(shard, {})[key] = tensor
+
+    if new_keys:
+        last_shard = sorted(set(weight_map.values()))[-1]
+        shard_to_new.setdefault(last_shard, {}).update(new_keys)
+        for key in new_keys:
+            weight_map[key] = last_shard
+        with index_path.open("w") as f:
+            json.dump(index_data, f, indent=2)
+        console.print(
+            f"  Added [cyan]{len(new_keys)}[/] new key(s) to shard "
+            f"[dim]{last_shard}[/]"
+        )
 
     with _bar() as progress:
         task = progress.add_task("Stitching shards", total=len(shard_to_new))
diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py
index c1ab75f0f..0e3d10997 100644
--- a/src/speculators/train/data.py
+++ b/src/speculators/train/data.py
@@ -341,7 +341,15 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None:
 
     def _get_raw_data(self, index):
         file_idx = self._map_to_file_idx(index)
-        loaded_hs = self.transfer.get_cached(file_idx)
+        try:
+            loaded_hs = self.transfer.get_cached(file_idx)
+        except Exception as e:
+            warnings.warn(
+                f"Corrupted cached hidden states for sample {index} "
+                f"(file_idx={file_idx}): {e}. Treating as missing.",
+                stacklevel=1,
+            )
+            loaded_hs = None
 
         if loaded_hs is None:
             match self.on_missing:
diff --git a/test.sh b/test.sh
index 21e5b7c1d..d543dd5c7 100644
--- a/test.sh
+++ b/test.sh
@@ -9,7 +9,7 @@ CUDA_VISIBLE_DEVICES=2,3 python scripts/train.py \
   --verifier-name-or-path google/gemma-4-31B-it \
   --data-path ./output/mtp_gemma-4-31B-it \
   --vllm-endpoint http://localhost:8000/v1 \
-  --save-path ./output/mtp_qwen3_5_9b/checkpoints \
+  --save-path ./output/mtp_gemma-4-31B-it/checkpoints \
   --speculator-type mtp \
   --num-speculative-steps 3 \
   --step-weight-beta 0.6 \
@@ -17,4 +17,12 @@ CUDA_VISIBLE_DEVICES=2,3 python scripts/train.py \
   --lr 1e-4 \
   --total-seq-len 8192 \
   --on-missing generate \
-  --on-generate delete
\ No newline at end of file
+  --on-generate delete
+
+
+CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 scripts/train.py  
+ --verifier-name-or-path google/gemma-4-31B-it   --data-path ./output/mtp_gemma-4-31B-it   --vllm-endpoint http://localhost:8000/v1   --save-path ./output/mtp_gemma-4-31B-it/checkpoints   --speculator-type mtp   --num-speculative-steps 3   --step-weight-beta 0.6   --epochs 3   chs 3   --lr 1e-4   --total-seq-len 8192 --on-missing raise
+
+
+
+python scripts/stitch_mtp.py   ./output/mtp_gemma-4-31B-it/checkpoints/checkpoint_best   google/gemma-4-31B-it   --output-path ./output/stitched
\ No newline at end of file

From 0220fede983be1a4947738003538425b3a9bbd9a Mon Sep 17 00:00:00 2001
From: Roderick Wu <Roderick-Wu@a100-06.nemg-001.lab.rdu2.dc.redhat.com>
Date: Thu, 23 Jul 2026 00:28:59 -0400
Subject: [PATCH 19/19] fixes by claude

---
 scripts/export_gemma4_mtp.py                  | 323 ++++++++++++++++++
 src/speculators/models/mtp/core.py            |  11 +-
 .../models/mtp/model_definitions.py           |  57 +++-
 test.sh                                       |   3 +-
 4 files changed, 384 insertions(+), 10 deletions(-)
 create mode 100644 scripts/export_gemma4_mtp.py

diff --git a/scripts/export_gemma4_mtp.py b/scripts/export_gemma4_mtp.py
new file mode 100644
index 000000000..4bc5833e8
--- /dev/null
+++ b/scripts/export_gemma4_mtp.py
@@ -0,0 +1,323 @@
+#!/usr/bin/env python3
+"""
+Export a trained Gemma4 MTP speculator as a standalone vLLM-compatible checkpoint.
+
+Unlike stitch_mtp.py (which merges MTP weights INTO a verifier checkpoint for
+Qwen-style models), this script produces a SEPARATE checkpoint directory that
+vLLM loads as an independent draft model via the Gemma4Speculator path.
+
+Usage:
+    python scripts/export_gemma4_mtp.py ./finetuned-mtp google/gemma-4-31B-it
+
+    # custom output path (defaults to {finetuned}-exported):
+    python scripts/export_gemma4_mtp.py ./finetuned-mtp ./verifier --output-path ./out
+"""
+
+import json
+from pathlib import Path
+from typing import Annotated
+
+import torch
+import typer
+from huggingface_hub import snapshot_download
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import (
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+app = typer.Typer(rich_markup_mode="rich")
+console = Console()
+
+_FROZEN_KEYS = {"embed_tokens.weight", "lm_head.weight"}
+
+_EXACT_REMAP: dict[str, str] = {
+    "mtp_layers.0.pre_projection.": "pre_projection.",
+    "mtp_layers.0.post_projection.": "post_projection.",
+    "mtp_layers.0.final_norm.": "model.norm.",
+}
+
+_PREFIX_REMAP = "mtp_layers.0."
+_PREFIX_TARGET = "model.layers.0."
+
+_TEXT_CONFIG_FIELDS = [
+    "attention_bias",
+    "attention_dropout",
+    "attention_k_eq_v",
+    "bos_token_id",
+    "eos_token_id",
+    "final_logit_softcapping",
+    "global_head_dim",
+    "head_dim",
+    "hidden_activation",
+    "hidden_size",
+    "intermediate_size",
+    "max_position_embeddings",
+    "model_type",
+    "num_attention_heads",
+    "num_key_value_heads",
+    "num_global_key_value_heads",
+    "pad_token_id",
+    "rms_norm_eps",
+    "rope_parameters",
+    "sliding_window",
+    "vocab_size",
+]
+
+
+def _spinner() -> Progress:
+    return Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        TimeElapsedColumn(),
+        console=console,
+    )
+
+
+def _resolve_path(model_path: Path) -> Path:
+    if model_path.exists():
+        return model_path
+    model_id = str(model_path)
+    console.print(
+        f"Path [cyan]{model_id}[/] not found locally, downloading from HuggingFace..."
+    )
+    with _spinner() as progress:
+        progress.add_task(f"Downloading {model_id}", total=None)
+        local_path = snapshot_download(
+            repo_id=model_id,
+            allow_patterns=["*.json", "*.safetensors", "*.bin", "*.index.json"],
+        )
+    console.print(f"  Downloaded to [dim]{local_path}[/]")
+    return Path(local_path)
+
+
+def _load_safetensors(checkpoint_dir: Path) -> dict[str, torch.Tensor]:
+    weights: dict[str, torch.Tensor] = {}
+
+    index_path = checkpoint_dir / "model.safetensors.index.json"
+    if index_path.exists():
+        with index_path.open() as f:
+            weight_map = json.load(f)["weight_map"]
+        for shard in set(weight_map.values()):
+            with safe_open(str(checkpoint_dir / shard), framework="pt") as f:
+                for key in f.keys():  # noqa: SIM118
+                    weights[key] = f.get_tensor(key)
+        return weights
+
+    single = checkpoint_dir / "model.safetensors"
+    if single.exists():
+        with safe_open(str(single), framework="pt") as f:
+            for key in f.keys():  # noqa: SIM118
+                weights[key] = f.get_tensor(key)
+        return weights
+
+    raise FileNotFoundError(f"No safetensors found at {checkpoint_dir}")
+
+
+def _remap_key(key: str) -> str:
+    for src, dst in _EXACT_REMAP.items():
+        if key.startswith(src):
+            return dst + key[len(src):]
+
+    if key.startswith(_PREFIX_REMAP):
+        return _PREFIX_TARGET + key[len(_PREFIX_REMAP):]
+
+    if key == "embed_tokens.weight":
+        return "model.embed_tokens.weight"
+
+    return key
+
+
+def _load_verifier_shared_weights(
+    verifier_dir: Path,
+) -> dict[str, torch.Tensor]:
+    """Load only embed_tokens and lm_head from the verifier."""
+    shared: dict[str, torch.Tensor] = {}
+    target_keys = {
+        "model.embed_tokens.weight",
+        "model.language_model.embed_tokens.weight",
+        "model.language_model.model.embed_tokens.weight",
+        "language_model.model.embed_tokens.weight",
+        "language_model.embed_tokens.weight",
+        "lm_head.weight",
+        "model.language_model.lm_head.weight",
+        "language_model.lm_head.weight",
+    }
+
+    index_path = verifier_dir / "model.safetensors.index.json"
+    if index_path.exists():
+        with index_path.open() as f:
+            weight_map: dict[str, str] = json.load(f)["weight_map"]
+        relevant_shards = {
+            weight_map[k] for k in target_keys if k in weight_map
+        }
+        for shard in relevant_shards:
+            with safe_open(str(verifier_dir / shard), framework="pt") as f:
+                for key in f.keys():  # noqa: SIM118
+                    if key in target_keys:
+                        shared[key] = f.get_tensor(key)
+    else:
+        single = verifier_dir / "model.safetensors"
+        if single.exists():
+            with safe_open(str(single), framework="pt") as f:
+                for key in f.keys():  # noqa: SIM118
+                    if key in target_keys:
+                        shared[key] = f.get_tensor(key)
+
+    result: dict[str, torch.Tensor] = {}
+    for key, tensor in shared.items():
+        if "embed_tokens" in key:
+            result["model.embed_tokens.weight"] = tensor
+        elif "lm_head" in key:
+            result["lm_head.weight"] = tensor
+
+    return result
+
+
+def _build_config(verifier_dir: Path) -> dict:
+    with (verifier_dir / "config.json").open() as f:
+        verifier_config = json.load(f)
+
+    verifier_text = verifier_config.get("text_config", verifier_config)
+
+    text_config: dict = {}
+    for field in _TEXT_CONFIG_FIELDS:
+        if field in verifier_text:
+            text_config[field] = verifier_text[field]
+
+    text_config["num_hidden_layers"] = 1
+    text_config["layer_types"] = ["full_attention"]
+    text_config["enable_moe_block"] = False
+    text_config["hidden_size_per_layer_input"] = 0
+    text_config["vocab_size_per_layer_input"] = 0
+    text_config["tie_word_embeddings"] = True
+
+    hidden_size = verifier_text["hidden_size"]
+    config = {
+        "model_type": "gemma4_assistant",
+        "architectures": ["Gemma4MTPModel"],
+        "backbone_hidden_size": hidden_size,
+        "tie_word_embeddings": True,
+        "text_config": text_config,
+    }
+    return config
+
+
+def export(
+    finetuned_checkpoint: Path,
+    verifier_path: Path,
+    output_path: Path,
+) -> Path:
+    console.print(
+        Panel(
+            f"[bold]Finetuned:[/] {finetuned_checkpoint}\n"
+            f"[bold]Verifier:[/]  {verifier_path}\n"
+            f"[bold]Output:[/]    {output_path}",
+            title="[bold green]Gemma4 MTP Export[/]",
+            border_style="green",
+        )
+    )
+
+    verifier_path = _resolve_path(verifier_path)
+
+    with _spinner() as progress:
+        progress.add_task("Loading finetuned weights", total=None)
+        finetuned = _load_safetensors(finetuned_checkpoint)
+    console.print(f"  Loaded [cyan]{len(finetuned)}[/] finetuned weight tensors")
+
+    trained_keys = {k: v for k, v in finetuned.items() if k not in _FROZEN_KEYS}
+    console.print(
+        f"  Filtered to [cyan]{len(trained_keys)}[/] trained keys "
+        f"(skipped {len(finetuned) - len(trained_keys)} frozen)"
+    )
+
+    exported: dict[str, torch.Tensor] = {}
+    for key, tensor in trained_keys.items():
+        new_key = _remap_key(key)
+        exported[new_key] = tensor
+    console.print(f"  Remapped [cyan]{len(exported)}[/] keys to vLLM format")
+
+    with _spinner() as progress:
+        progress.add_task("Loading verifier shared weights", total=None)
+        shared = _load_verifier_shared_weights(verifier_path)
+    exported.update(shared)
+    console.print(f"  Added [cyan]{len(shared)}[/] shared weight(s) from verifier")
+
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    config = _build_config(verifier_path)
+    config_path = output_path / "config.json"
+    with config_path.open("w") as f:
+        json.dump(config, f, indent=2)
+    console.print(f"  Wrote config to [dim]{config_path}[/]")
+
+    weights_path = output_path / "model.safetensors"
+    with _spinner() as progress:
+        progress.add_task("Saving exported weights", total=None)
+        save_file(exported, str(weights_path))
+    console.print(f"  Saved [cyan]{len(exported)}[/] tensors to [dim]{weights_path}[/]")
+
+    console.print("\n[bold]Exported weight keys:[/]")
+    for key in sorted(exported.keys()):
+        shape = list(exported[key].shape)
+        console.print(f"  [dim]{key}[/]: {shape}")
+
+    console.print(
+        Panel(
+            f"[bold]{output_path}[/]\n\n"
+            "Deploy with:\n"
+            f"  vllm serve <verifier> --speculative-config "
+            f"'{{\"method\":\"mtp\",\"model\":\"{output_path}\","
+            f"\"num_speculative_tokens\":3}}'",
+            title="[bold green]Export complete[/]",
+            border_style="green",
+        )
+    )
+    return output_path
+
+
+@app.command()
+def main(
+    finetuned_checkpoint: Annotated[
+        Path,
+        typer.Argument(
+            help="Path to the finetuned MTP speculator checkpoint.",
+        ),
+    ],
+    verifier_path: Annotated[
+        Path,
+        typer.Argument(
+            help=(
+                "Path to the verifier checkpoint, or a HuggingFace "
+                "model ID (e.g. google/gemma-4-31B-it)."
+            ),
+        ),
+    ],
+    output_path: Annotated[
+        Path | None,
+        typer.Option(
+            help=(
+                "Output directory for the exported checkpoint. "
+                "Defaults to {finetuned}-exported."
+            ),
+        ),
+    ] = None,
+) -> None:
+    """Export a trained Gemma4 MTP speculator for vLLM deployment."""
+    if output_path is None:
+        output_path = Path.cwd() / f"{finetuned_checkpoint.name}-exported"
+
+    export(
+        finetuned_checkpoint=finetuned_checkpoint,
+        verifier_path=verifier_path,
+        output_path=output_path,
+    )
+
+
+if __name__ == "__main__":
+    app()
diff --git a/src/speculators/models/mtp/core.py b/src/speculators/models/mtp/core.py
index 457cc430a..1ee2ecba9 100644
--- a/src/speculators/models/mtp/core.py
+++ b/src/speculators/models/mtp/core.py
@@ -216,7 +216,12 @@ def forward(
                 verifier_kv_last_global=kwargs.get("verifier_kv_last_global"),
             )
 
-            logits = self.lm_head(mtp_output)
+            if isinstance(mtp_output, tuple):
+                draft_hidden, backbone_hidden = mtp_output
+            else:
+                draft_hidden = backbone_hidden = mtp_output
+
+            logits = self.lm_head(draft_hidden)
             all_logits.append(logits)
 
             step_targets = input_ids[:, step + 2 : step + 2 + valid_len]
@@ -235,7 +240,7 @@ def forward(
 
             if getattr(self, "masked_embedding", None) is not None:
                 unreduced = unreduced + self.masked_embedding.compute_centroid_loss(
-                    hidden_states=mtp_output,
+                    hidden_states=draft_hidden,
                     targets=step_targets,
                     ignore_index=_IGNORE_INDEX,
                 )
@@ -244,7 +249,7 @@ def forward(
             total_loss = total_loss + step_loss
             metrics[f"loss_step_{step}"] = step_loss.detach().clone()
 
-            current_hidden = mtp_output
+            current_hidden = backbone_hidden
 
         metrics["loss_sum"] = total_loss.detach().clone()
         metrics["loss_total"] = torch.tensor(1.0, device=device)
diff --git a/src/speculators/models/mtp/model_definitions.py b/src/speculators/models/mtp/model_definitions.py
index 94b8068c8..65a104850 100644
--- a/src/speculators/models/mtp/model_definitions.py
+++ b/src/speculators/models/mtp/model_definitions.py
@@ -315,12 +315,13 @@ class Gemma4MTPRotaryEmbedding(Gemma4TextRotaryEmbedding):
         Gemma4TextRotaryEmbedding.forward requires a layer_type arg
         (per-type RoPE params), but MTPDraftModel.forward calls
         rotary_emb(x, position_ids) without one. This wrapper fixes
-        the layer_type to the MTP layer's type (layer_types[0]).
+        the layer_type to full_attention — the single MTP layer must
+        be full_attention (enforced by Gemma4TextConfig.__post_init__).
         """
 
         def __init__(self, config: PretrainedConfig, device=None) -> None:
             super().__init__(config, device)
-            self._mtp_layer_type = config.layer_types[0]
+            self._mtp_layer_type = "full_attention"
 
         def forward(
             self, x: torch.Tensor, position_ids: torch.Tensor, layer_type: str | None = None
@@ -400,14 +401,60 @@ def forward(
             attn_output = self.o_proj(attn_output)
             return attn_output, attn_weights
 
-    class Gemma4TextMTPLayer(MTPLayerMixin, Gemma4TextDecoderLayer):
+    class Gemma4NativeMTPLayer(Gemma4TextDecoderLayer):
+        """Gemma4 MTP layer matching vLLM's native Gemma4MultiTokenPredictor.
+
+        Uses pre/post projection (no pre-fc norms) and returns a tuple
+        (draft_hidden, backbone_hidden) for split logit/feedback paths.
+        The single MTP layer is always full_attention — Gemma4TextConfig
+        enforces this in __post_init__, and vLLM builds the model
+        accordingly.
+        """
+
         def __init__(self, config: PretrainedConfig, layer_idx: int = 0) -> None:
+            config = copy.copy(config)
+            config.layer_types = ["full_attention"]
             super().__init__(config, layer_idx)
-            self._setup_mtp_modules(config, Gemma4RMSNorm)
             self.self_attn = QueryOnlyGemma4TextAttention(config, layer_idx)
 
+            hidden_size = config.hidden_size
+            self.pre_projection = nn.Linear(2 * hidden_size, hidden_size, bias=False)
+            self.post_projection = nn.Linear(hidden_size, hidden_size, bias=False)
+            self.final_norm = Gemma4RMSNorm(hidden_size, eps=config.rms_norm_eps)
+            self.register_buffer(
+                "normalizer",
+                torch.tensor(hidden_size**0.5),
+                persistent=False,
+            )
+
+        def forward(
+            self,
+            hidden_states: torch.Tensor,
+            token_embeddings: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.Tensor] = None,
+            position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+            **kwargs: Any,
+        ) -> Tuple[torch.Tensor, torch.Tensor]:
+            scaled_embeds = token_embeddings * self.normalizer
+            combined = torch.cat([scaled_embeds, hidden_states], dim=-1)
+            projected = self.pre_projection(combined)
+
+            decoder_output = super().forward(
+                hidden_states=projected,
+                position_embeddings=position_embeddings,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                **kwargs,
+            )
+            hidden = decoder_output[0] if isinstance(decoder_output, tuple) else decoder_output
+
+            draft_hidden = self.final_norm(hidden)
+            backbone_hidden = self.post_projection(draft_hidden)
+            return draft_hidden, backbone_hidden
+
     mtp_model_classes["gemma4_text"] = base_components.override_components(
         "gemma4_text",
-        first_layer_class=Gemma4TextMTPLayer,
+        first_layer_class=Gemma4NativeMTPLayer,
         rotary_emb_class=Gemma4MTPRotaryEmbedding,
     )
diff --git a/test.sh b/test.sh
index d543dd5c7..018b90456 100644
--- a/test.sh
+++ b/test.sh
@@ -20,8 +20,7 @@ CUDA_VISIBLE_DEVICES=2,3 python scripts/train.py \
   --on-generate delete
 
 
-CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 scripts/train.py  
- --verifier-name-or-path google/gemma-4-31B-it   --data-path ./output/mtp_gemma-4-31B-it   --vllm-endpoint http://localhost:8000/v1   --save-path ./output/mtp_gemma-4-31B-it/checkpoints   --speculator-type mtp   --num-speculative-steps 3   --step-weight-beta 0.6   --epochs 3   chs 3   --lr 1e-4   --total-seq-len 8192 --on-missing raise
+CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 scripts/train.py  --verifier-name-or-path google/gemma-4-31B-it   --data-path ./output/mtp_gemma-4-31B-it   --vllm-endpoint http://localhost:8000/v1   --save-path ./output/mtp_gemma-4-31B-it/checkpoints   --speculator-type mtp   --num-speculative-steps 3   --step-weight-beta 0.6   --epochs 3   chs 3   --lr 1e-4   --total-seq-len 8192 --on-missing raise