feat(mtp): implement QueryOnlyGemma2Attention for Gemma4 MTP draft models#768
feat(mtp): implement QueryOnlyGemma2Attention for Gemma4 MTP draft models#768shotsan wants to merge 14 commits into
Conversation
shotsan
changed the title
|
Important Review skippedIgnore keyword(s) in the title. ⛔ Ignored keywords (3)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds optional Gemma2 support, verifier KV-cache propagation through training data and MTP attention, and centroid-masked loss computation with configuration and unit tests. ChangesGemma2 MTP and centroid-loss support
Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/speculators/train/data.py (1)
70-97: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftEmpty verifier KV placeholders need the real feature rank.
create_empty_sample()buildsverifier_kv_last_localandverifier_kv_last_globalas 1D empty tensors, but the collate fallback treats them like the per-sample KV tensors and can pass a rank-mismatched batch intoQueryOnlyGemma2Attention.forward, which will fail in the all-samples-failed fallback path. Thread the KV feature dimension through this fallback so the empty sample matches the real[seq_len, feature_dim]shape.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/speculators/train/data.py` around lines 70 - 97, create_empty_sample currently creates 1D empty verifier KV tensors, but downstream collation expects [seq_len, feature_dim]. Add a KV feature-dimension parameter to create_empty_sample, pass the real dimension from the collate fallback, and construct verifier_kv_last_local and verifier_kv_last_global as empty tensors with shape [0, feature_dim]. Update all callers accordingly.src/speculators/models/base_components.py (1)
72-120: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winKeep the qwen3_next/qwen3_5_text/qwen3_5_moe_text registrations at
src/speculators/models/base_components.py:72-120—src/speculators/models/mtp/model_definitions.pyonly adds MTP variants when these base entries exist, andMTPDraftModelresolves throughresolve_model_type(...); removing them makes those model types unsupported.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/speculators/models/base_components.py` around lines 72 - 120, Keep the qwen3_next, qwen3_5_text, and qwen3_5_moe_text entries in the model_classes registry, including their existing ModelComponents mappings and ImportError handling. Verify that resolve_model_type(...) and MTPDraftModel continue to resolve MTP variants through these base registrations.
🧹 Nitpick comments (3)
src/speculators/models/mtp/head.py (1)
30-37: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRecompute the inverse mapping only when
token_orderingchanges.
inverseis rebuilt with a full O(vocab_size) scatter on everycompute_centroid_losscall (once per MTP step, per training step). Sincetoken_orderingis a buffer that changes rarely (if ever, during training), caching the inverse mapping (recomputed only whentoken_orderingis updated) would avoid this repeated allocation/scatter on the hot training path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/speculators/models/mtp/head.py` around lines 30 - 37, Cache the token-to-centroid inverse mapping instead of rebuilding it in every compute_centroid_loss call. Add a cached buffer or equivalent tied to token_ordering, initialize it from token_ordering, and refresh it whenever token_ordering is updated; then have compute_centroid_loss reuse the cache for valid_targets.tests/unit/train/test_data.py (1)
474-516: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winGood coverage for the normal path; consider also covering the empty-sample fallback.
This test verifies KV caches survive collation for real samples, but doesn't cover the
create_empty_sample/all-batch-failed fallback path, where the KV placeholder shape diverges from real samples (seedata.pycomment). Adding a test usingcreate_empty_sample+create_collate_fnwould catch that regression. As per path instructions, tests should cover edge cases specific to speculative decoding and new code paths introduced in the PR.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/train/test_data.py` around lines 474 - 516, Add a test covering the all-batch-failed empty-sample fallback for verifier KV caches. Use create_empty_sample to construct an empty item, pass it through create_collate_fn, and assert verifier_kv_last_local/global exist with shapes and padding compatible with real samples, preventing placeholder shape divergence.Source: Path instructions
src/speculators/models/base_components.py (1)
15-22: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFragile
except NameErrorguard for optional Gemma2 support.Catching
NameErrorto detect a prior failed import is a fragile pattern: it will silently swallow any unrelatedNameErrorraised inside thetryblock (e.g., a future typo when extending this block), masking real bugs instead of surfacing them at import time.♻️ Suggested clearer pattern
+_HAS_GEMMA2 = False try: from transformers.models.gemma2.modeling_gemma2 import ( Gemma2DecoderLayer, Gemma2RMSNorm, Gemma2RotaryEmbedding, ) + _HAS_GEMMA2 = True except ImportError: pass ... -try: +if _HAS_GEMMA2: model_classes["gemma2"] = ModelComponents( Gemma2DecoderLayer, Gemma2DecoderLayer, Gemma2RMSNorm, Gemma2RotaryEmbedding, ) -except NameError: - passAlso applies to: 62-71
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/speculators/models/base_components.py` around lines 15 - 22, Replace the fragile `NameError`-based optional Gemma2 detection in `base_components.py` with an explicit availability flag initialized before the import and set only after `Gemma2DecoderLayer`, `Gemma2RMSNorm`, and `Gemma2RotaryEmbedding` import successfully; catch only `ImportError`, and use that flag wherever these classes are referenced, including the corresponding block around lines 62–71.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/speculators/models/mtp/config.py`:
- Around line 55-67: In the configuration model containing num_centroids and
top_k_centroids, add an after-model validator that requires num_centroids to be
positive and, when set, evenly divide vocab_size; raise a clear validation error
for invalid values before constructing MultiLevelLMHead. Preserve the existing
None behavior for the standard flat LM head.
In `@src/speculators/models/mtp/head.py`:
- Around line 11-16: Fix centroid assignment for vocabularies not evenly
divisible by num_centroids: update the initialization in __init__ and the
token-to-centroid mapping used by compute_centroid_loss so every token ID maps
within 0..num_centroids-1, including remainder tokens. Ensure the mapping
distributes or explicitly assigns all vocab_size tokens without producing
centroid index num_centroids, and keep token_ordering consistent with that
mapping.
In `@src/speculators/models/mtp/model_definitions.py`:
- Around line 263-278: Update the MTP attention mask construction to call
create_sliding_window_causal_mask when self.sliding_window is set, while
retaining create_causal_mask otherwise; do not rely on forwarding sliding_window
through QueryOnlyGemma2Attention. Add a regression test verifying that tokens
outside the configured window are masked for local Gemma2 attention.
In `@tests/unit/models/test_mtp_gemma2.py`:
- Around line 37-60: Strengthen test_query_only_gemma2_attention_local_kv and
the related tests so they verify cache routing rather than only output shape:
run paired forwards while varying only the selected verifier cache, assert the
output changes, then vary the inactive cache and assert the output remains
unchanged. Cover both local and global cache paths introduced by the PR, using
distinct cache values and the appropriate configuration or inputs to select each
path.
- Around line 18-34: Update the test fixture to instantiate Gemma2Config instead
of PretrainedConfig, supplying the required Gemma2-specific fields such as
layer_types and attn_logit_softcapping. Configure the global-attention case with
layer_types set to full_attention rather than relying only on sliding_window =
None, while preserving the settings needed by QueryOnlyGemma2Attention.
- Around line 48-50: The RoPE test fixtures use an incorrect extra singleton
dimension. In both tests around the rotary embedding setup, update the `cos` and
`sin` tensors to shape `[batch_sz, seq_len, 16]` instead of `[1, 1, seq_len,
16]`, matching the input expected by `apply_rotary_pos_emb`.
In `@tests/unit/models/test_mtp_head.py`:
- Around line 10-16: The test fixture only covers evenly divisible vocabulary
and centroid sizes. Add a PyTest case alongside `head_setup` using values where
`vocab_size % num_centroids != 0`, then exercise the relevant `MultiLevelLMHead`
forward or vocab-mapping behavior and assert outputs remain valid without
out-of-range centroid indices, covering the mapping boundary for speculative
decoding.
---
Outside diff comments:
In `@src/speculators/models/base_components.py`:
- Around line 72-120: Keep the qwen3_next, qwen3_5_text, and qwen3_5_moe_text
entries in the model_classes registry, including their existing ModelComponents
mappings and ImportError handling. Verify that resolve_model_type(...) and
MTPDraftModel continue to resolve MTP variants through these base registrations.
In `@src/speculators/train/data.py`:
- Around line 70-97: create_empty_sample currently creates 1D empty verifier KV
tensors, but downstream collation expects [seq_len, feature_dim]. Add a KV
feature-dimension parameter to create_empty_sample, pass the real dimension from
the collate fallback, and construct verifier_kv_last_local and
verifier_kv_last_global as empty tensors with shape [0, feature_dim]. Update all
callers accordingly.
---
Nitpick comments:
In `@src/speculators/models/base_components.py`:
- Around line 15-22: Replace the fragile `NameError`-based optional Gemma2
detection in `base_components.py` with an explicit availability flag initialized
before the import and set only after `Gemma2DecoderLayer`, `Gemma2RMSNorm`, and
`Gemma2RotaryEmbedding` import successfully; catch only `ImportError`, and use
that flag wherever these classes are referenced, including the corresponding
block around lines 62–71.
In `@src/speculators/models/mtp/head.py`:
- Around line 30-37: Cache the token-to-centroid inverse mapping instead of
rebuilding it in every compute_centroid_loss call. Add a cached buffer or
equivalent tied to token_ordering, initialize it from token_ordering, and
refresh it whenever token_ordering is updated; then have compute_centroid_loss
reuse the cache for valid_targets.
In `@tests/unit/train/test_data.py`:
- Around line 474-516: Add a test covering the all-batch-failed empty-sample
fallback for verifier KV caches. Use create_empty_sample to construct an empty
item, pass it through create_collate_fn, and assert
verifier_kv_last_local/global exist with shapes and padding compatible with real
samples, preventing placeholder shape divergence.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9b5a90a1-e15e-4ef1-9fa5-b442f53e3ff8
📒 Files selected for processing (9)
src/speculators/models/base_components.pysrc/speculators/models/mtp/config.pysrc/speculators/models/mtp/core.pysrc/speculators/models/mtp/head.pysrc/speculators/models/mtp/model_definitions.pysrc/speculators/train/data.pytests/unit/models/test_mtp_gemma2.pytests/unit/models/test_mtp_head.pytests/unit/train/test_data.py
| 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.", | ||
| ) | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Consider validating num_centroids against vocab_size here.
num_centroids has no validation (e.g., must be > 0 and ideally divide vocab_size evenly). As shown in head.py, a non-divisible combination causes out-of-range centroid targets. A model_validator(mode="after") here (with access to vocab_size) would be a natural place to enforce this contract before it reaches MultiLevelLMHead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/speculators/models/mtp/config.py` around lines 55 - 67, In the
configuration model containing num_centroids and top_k_centroids, add an
after-model validator that requires num_centroids to be positive and, when set,
evenly divide vocab_size; raise a clear validation error for invalid values
before constructing MultiLevelLMHead. Preserve the existing None behavior for
the standard flat LM head.
| 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)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Out-of-range centroid index when vocab_size is not evenly divisible by num_centroids.
tokens_per_centroid = vocab_size // num_centroids truncates; remainder tokens then map to a centroid id equal to num_centroids (out of the valid 0..num_centroids-1 range used by self.centroids), causing F.cross_entropy in compute_centroid_loss to index out of bounds for any target token in that trailing remainder. E.g., vocab_size=10, num_centroids=3 maps token position 9 to centroid 3, which doesn't exist.
🐛 Suggested fix
def __init__(self, hidden_size: int, vocab_size: int, num_centroids: int):
super().__init__()
+ if num_centroids <= 0 or vocab_size % num_centroids != 0:
+ raise ValueError(
+ f"vocab_size ({vocab_size}) must be evenly divisible by "
+ f"num_centroids ({num_centroids})."
+ )
self.tokens_per_centroid = vocab_size // num_centroids📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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 __init__(self, hidden_size: int, vocab_size: int, num_centroids: int): | |
| super().__init__() | |
| if num_centroids <= 0 or vocab_size % num_centroids != 0: | |
| raise ValueError( | |
| f"vocab_size ({vocab_size}) must be evenly divisible by " | |
| f"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)) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/speculators/models/mtp/head.py` around lines 11 - 16, Fix centroid
assignment for vocabularies not evenly divisible by num_centroids: update the
initialization in __init__ and the token-to-centroid mapping used by
compute_centroid_loss so every token ID maps within 0..num_centroids-1,
including remainder tokens. Ensure the mapping distributes or explicitly assigns
all vocab_size tokens without producing centroid index num_centroids, and keep
token_ordering consistent with that mapping.
| 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, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant files and inspect the surrounding implementation.
git ls-files src/speculators/models/mtp/model_definitions.py src/speculators/models/mtp/core.py
printf '\n--- model_definitions outline ---\n'
ast-grep outline src/speculators/models/mtp/model_definitions.py --view expanded || true
printf '\n--- core outline ---\n'
ast-grep outline src/speculators/models/mtp/core.py --view expanded || true
printf '\n--- model_definitions excerpt ---\n'
sed -n '220,320p' src/speculators/models/mtp/model_definitions.py
printf '\n--- core excerpt ---\n'
sed -n '1,260p' src/speculators/models/mtp/core.pyRepository: vllm-project/speculators
Length of output: 16216
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search for mask creation helpers and Gemma2-specific handling.
rg -n "create_(causal|sliding_window_causal)_mask|sliding_window|Gemma2|attention_mask" src/speculators/models/mtp src/speculators/models -g '*.py'Repository: vllm-project/speculators
Length of output: 10863
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the upstream Gemma2 attention implementation reference in installed docs/source if available in repo deps.
python3 - <<'PY'
import os, sys, subprocess, textwrap, json, pathlib
candidates = [
"transformers/src/transformers/models/gemma2/modeling_gemma2.py",
"site-packages/transformers/models/gemma2/modeling_gemma2.py",
]
for c in candidates:
if os.path.exists(c):
print("FOUND", c)
with open(c, "r", encoding="utf-8") as f:
lines = f.readlines()
for i, line in enumerate(lines, 1):
if "def eager_attention_forward" in line or "sliding_window" in line or "attention_mask" in line:
start = max(1, i-20)
end = min(len(lines), i+120)
print(f"\n--- {c}:{start}-{end} ---")
for j in range(start, end+1):
print(f"{j}:{lines[j-1].rstrip()}")
break
PYRepository: vllm-project/speculators
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Try to inspect the relevant raw upstream source via web if the repository doesn't vendor it.
# This is a read-only query to validate whether eager Gemma2 attention relies on the passed attention_mask.
python3 - <<'PY'
import urllib.request, ssl
url = "https://raw.githubusercontent.com/huggingface/transformers/v4.56.1/src/transformers/models/gemma2/modeling_gemma2.py"
ctx = ssl.create_default_context()
with urllib.request.urlopen(url, context=ctx, timeout=20) as resp:
text = resp.read().decode("utf-8")
for needle in ["def eager_attention_forward", "sliding_window", "attention_mask"]:
idx = text.find(needle)
if idx != -1:
start = max(0, idx - 600)
end = min(len(text), idx + 2200)
print(text[start:end])
print("\n--- END SNIP ---\n")
PYRepository: vllm-project/speculators
Length of output: 2756
🌐 Web query:
Transformers 4.56.1 Gemma2 eager_attention_forward sliding_window attention_mask source
💡 Result:
In the Hugging Face Transformers library (v4.56.1 and surrounding versions), the eager_attention_forward function in the Gemma2 model implementation serves as a fallback or default attention mechanism when the _attn_implementation is not set to optimized alternatives like SDPA or Flash Attention [1][2][3]. Key technical details regarding your query: 1. sliding_window: In Gemma2, the sliding_window parameter is defined in the configuration (defaulting to 4096) [4] and is passed to the attention forward pass based on whether the specific layer is configured as a sliding_attention layer [5]. However, in the standard eager_attention_forward implementation for Gemma2, the sliding_window argument is accepted but is typically not explicitly used to mask or limit attention within that specific function, which instead performs full attention (with potential softcapping and causal masking) [1][2][3]. The sliding window behavior is often handled by external mechanisms or specific model architecture configurations that define layer types [5]. 2. attention_mask: The attention_mask in eager_attention_forward is used to augment the attention scores (attn_weights) before the softmax operation [1][2][3]. In the Gemma2 implementation, this is typically done via simple addition (attn_weights = attn_weights + attention_mask) after applying the attention logit softcapping [1][2][3]. 3. Source Location: The implementation can be found in the transformers library under src/transformers/models/gemma2/modeling_gemma2.py [6][2][3]. Gemma2 is unique in that it interleaves local (sliding window) and global attention across its layers, and it utilizes Grouped Query Attention (GQA) and logit softcapping (via the softcap parameter) to maintain stability during training and inference [1][4].
Citations:
- 1: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma2/modular_gemma2.py
- 2: https://github.com/huggingface/transformers/blob/v4.52.3/src/transformers/models/gemma2/modeling_gemma2.py
- 3: https://github.com/huggingface/transformers/blob/c472755e/src/transformers/models/gemma2/modeling_gemma2.py
- 4: https://huggingface.co/docs/transformers/v4.56.1/en/model_doc/gemma2
- 5: https://github.com/huggingface/transformers/blob/v4.57.2/src/transformers/models/gemma2/modeling_gemma2.py
- 6: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma2/modeling_gemma2.py
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import inspect
import transformers.masking_utils as mu
from transformers import PretrainedConfig
print("transformers", __import__("transformers").__version__)
print("\ncreate_causal_mask signature:")
print(inspect.signature(mu.create_causal_mask))
print("\ncreate_causal_mask source excerpt:")
src = inspect.getsource(mu.create_causal_mask)
print(src[:5000])
PYRepository: vllm-project/speculators
Length of output: 5708
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import inspect
import transformers.masking_utils as mu
for name in ["create_sliding_window_causal_mask", "create_causal_mask"]:
fn = getattr(mu, name, None)
print(f"{name}:", fn)
if fn is not None:
print("signature:", inspect.signature(fn))
try:
src = inspect.getsource(fn)
print(src[:2500])
except Exception as e:
print("source unavailable:", e)
print("-" * 80)
PYRepository: vllm-project/speculators
Length of output: 6430
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import inspect
import transformers.masking_utils as mu
src = inspect.getsource(mu.create_causal_mask)
for needle in ["sliding_window", "create_sliding_window_causal_mask", "config.sliding_window", "mask_kwargs"]:
print(f"\n=== {needle} ===")
idx = src.find(needle)
print(idx)
if idx != -1:
start = max(0, idx - 800)
end = min(len(src), idx + 1800)
print(src[start:end])
PYRepository: vllm-project/speculators
Length of output: 293
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import inspect
import transformers.masking_utils as mu
for name in ["create_sliding_window_causal_mask", "create_causal_mask"]:
fn = getattr(mu, name, None)
print(f"\n### {name}")
if fn is None:
print("MISSING")
continue
try:
print(inspect.signature(fn))
except Exception as e:
print("signature error:", e)
try:
src = inspect.getsource(fn)
print(src[:6000])
except Exception as e:
print("source error:", e)
PYRepository: vllm-project/speculators
Length of output: 13143
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect how verifier KV tensors are produced/passed, and whether local/global tensors differ in sequence length.
rg -n "verifier_kv_last_local|verifier_kv_last_global|sliding_window" src/speculators -g '*.py'Repository: vllm-project/speculators
Length of output: 5648
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
paths = [
Path("src/speculators/models/mtp/core.py"),
Path("src/speculators/models/mtp/model_definitions.py"),
]
for path in paths:
txt = path.read_text(encoding="utf-8").splitlines()
for needle in ["verifier_kv_last_local", "verifier_kv_last_global", "sliding_window"]:
hits = [i+1 for i,l in enumerate(txt) if needle in l]
if hits:
print(f"\n{path} :: {needle} -> {hits}")
for i in hits:
start = max(1, i-8)
end = min(len(txt), i+12)
print(f"--- lines {start}-{end} ---")
for j in range(start, end+1):
print(f"{j}: {txt[j-1]}")
PYRepository: vllm-project/speculators
Length of output: 9459
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether there are already tests around MTP/Gemma2 masking.
rg -n "MTP|Gemma2|sliding_window|create_sliding_window_causal_mask|create_causal_mask" tests src -g '*.py'Repository: vllm-project/speculators
Length of output: 29084
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import inspect
import transformers.models.gemma2.modeling_gemma2 as mg
print("transformers", __import__("transformers").__version__)
for name in ["eager_attention_forward", "Gemma2Attention"]:
obj = getattr(mg, name, None)
print(f"\n### {name}: {obj}")
if obj is None:
continue
try:
src = inspect.getsource(obj)
for needle in ["sliding_window", "attention_mask", "attn_weights"]:
print(f"\n-- {needle} --")
idx = src.find(needle)
print(idx)
if idx != -1:
start = max(0, idx - 800)
end = min(len(src), idx + 2400)
print(src[start:end])
except Exception as e:
print("source error:", e)
PYRepository: vllm-project/speculators
Length of output: 9338
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,220p' tests/unit/models/test_mtp_gemma2.pyRepository: vllm-project/speculators
Length of output: 2843
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("tests/unit/models/test_mtp_gemma2.py")
print(path.read_text(encoding="utf-8"))
PYRepository: vllm-project/speculators
Length of output: 2844
🏁 Script executed:
#!/bin/bash
set -euo pipefail
cat -n tests/unit/models/test_mtp_gemma2.pyRepository: vllm-project/speculators
Length of output: 3438
Use a sliding-window causal mask for Gemma2 local attention. QueryOnlyGemma2Attention only consumes attention_mask; the forwarded sliding_window kwarg does not narrow keys in eager Gemma2, so create_causal_mask still lets local layers attend over the full causal prefix. Switch the MTP mask build to create_sliding_window_causal_mask when self.sliding_window is set, and add a regression test that tokens outside the window are masked.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/speculators/models/mtp/model_definitions.py` around lines 263 - 278,
Update the MTP attention mask construction to call
create_sliding_window_causal_mask when self.sliding_window is set, while
retaining create_causal_mask otherwise; do not rely on forwarding sliding_window
through QueryOnlyGemma2Attention. Add a regression test verifying that tokens
outside the configured window are masked for local Gemma2 attention.
Source: Path instructions
| 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 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the test file around the referenced lines
wc -l tests/unit/models/test_mtp_gemma2.py
cat -n tests/unit/models/test_mtp_gemma2.py | sed -n '1,140p'
# Locate any Gemma2 config/test usage in the repo
rg -n "Gemma2Config|Gemma2Attention|attn_logit_softcapping|layer_types|sliding_window|query_pre_attn_scalar" tests src -g '!**/node_modules/**'
# If the transformers source is vendored or referenced in-repo, inspect the relevant model code
fd -a "modeling_gemma2.py|configuration_gemma2.py" .Repository: vllm-project/speculators
Length of output: 11035
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the wrapper that instantiates Gemma2Attention
wc -l src/speculators/models/mtp/model_definitions.py
cat -n src/speculators/models/mtp/model_definitions.py | sed -n '180,320p'
# Check where Gemma2 config-specific fields are read in the repo
rg -n "attn_logit_softcapping|query_pre_attn_scalar|layer_types|sliding_window|Gemma2Attention" src/speculators/models -g '!**/node_modules/**'
# Inspect any tests or fixtures that already use Gemma2Config or similar patterns
rg -n "Gemma2Config|PretrainedConfig\(" tests src -g '!**/node_modules/**'Repository: vllm-project/speculators
Length of output: 10522
🌐 Web query:
Transformers 4.56.1 Gemma2Attention __init__ attn_logit_softcapping layer_types Gemma2Config source
💡 Result:
In the Hugging Face Transformers library, the Gemma2Attention class implements specialized attention mechanisms for the Gemma 2 model, including the use of attn_logit_softcapping. The attn_logit_softcapping parameter serves as a scaling factor when applying tanh-based softcapping to attention scores [1][2]. Within the init method of the Gemma2Attention class, this value is initialized directly from the Gemma2Config object [3][4][5]: self.attn_logit_softcapping = self.config.attn_logit_softcapping The Gemma2Config class defines this as an optional float, defaulting to 50.0 [6][1]. Additionally, the Gemma2Attention class uses the layer_types configuration list to determine the attention behavior for each specific layer [7][3][4]. The layer type (e.g., "sliding_attention" or "full_attention") is retrieved during initialization [7][5]: self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None This layer_type is then used to set the sliding_window attribute for that attention instance [7][8][4]. These mechanisms are central to the Gemma 2 architecture as implemented in the transformers library [3][9][2].
Citations:
- 1: https://huggingface.co/docs/transformers/v5.3.0/en/model_doc/gemma2
- 2: https://github.com/huggingface/transformers/blob/96d3ecce/src/transformers/models/gemma2/configuration_gemma2.py
- 3: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma2/modeling_gemma2.py
- 4: https://github.com/huggingface/transformers/blob/v5.4.0/src/transformers/models/gemma2/modeling_gemma2.py
- 5: https://github.com/huggingface/transformers/blob/v5.1.0/src/transformers/models/gemma2/modeling_gemma2.py
- 6: https://github.com/huggingface/transformers/blob/v5.6.0/src/transformers/models/gemma2/configuration_gemma2.py
- 7: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma2/modular_gemma2.py
- 8: https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/gemma2/modeling_gemma2.py
- 9: https://huggingface.co/docs/transformers/en/model_doc/gemma2
🌐 Web query:
Transformers v4.56.1 src/transformers/models/gemma2/modeling_gemma2.py Gemma2Attention __init__ attn_logit_softcapping layer_types
💡 Result:
In the Hugging Face Transformers library, the Gemma2Attention class implements attention mechanisms designed for the Gemma 2 architecture, which includes specific handling for soft-capping and varying attention types [1][2]. attn_logit_softcapping This parameter represents a scaling factor used for logit soft-capping within the attention layers [3][4]. Soft-capping is a technique used to stabilize training by preventing attention scores (logits) from growing excessively large [5][6]. It is implemented using the formula: logits = soft_cap * tanh(logits / soft_cap) The default value for attn_logit_softcapping is 50.0 [3][7]. In the Gemma2Attention.init method, this value is assigned to self.attn_logit_softcapping from the configuration, where it is later passed to the attention forward pass to constrain the logits [1][2]. layer_types The layer_types configuration attribute is used to determine the specific attention implementation for a given layer index [1][2]. Gemma 2 models interleave different types of attention—specifically sliding window attention and standard (often full-quadratic or global) attention—to balance performance and long-context capabilities [5][6][4]. In the init method of Gemma2Attention, the code typically checks the layer_types configuration at the current layer_idx to set the sliding_window behavior [1][2]: self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None This ensures that only layers designated as "sliding_attention" apply the restricted window size, while others operate without that constraint [1][2]. Older versions of the implementation sometimes used a modulo operation (e.g., if not bool(layer_idx % 2)) to achieve this interleaving, but modern implementations favor the explicit layer_types list for clarity and flexibility [8][9][2].
Citations:
- 1: https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/gemma2/modeling_gemma2.py
- 2: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma2/modular_gemma2.py
- 3: https://huggingface.co/docs/transformers/en/model_doc/gemma2
- 4: https://huggingface.co/docs/transformers/model_doc/gemma2
- 5: https://github.com/huggingface/blog/blob/main/gemma2.md
- 6: https://arxiv.org/html/2408.00118
- 7: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma2/configuration_gemma2.py
- 8: https://github.com/huggingface/transformers/blob/v4.52.1/src/transformers/models/gemma2/modeling_gemma2.py
- 9: https://github.com/huggingface/transformers/blob/v4.52.3/src/transformers/models/gemma2/modeling_gemma2.py
Use Gemma2Config here
QueryOnlyGemma2Attention inherits Gemma2Attention, so this fixture needs Gemma2-specific config such as layer_types and attn_logit_softcapping; PretrainedConfig omits them and can fail during construction before either path runs. For the global-path case, switch the layer type to full_attention instead of only setting sliding_window = None.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/models/test_mtp_gemma2.py` around lines 18 - 34, Update the test
fixture to instantiate Gemma2Config instead of PretrainedConfig, supplying the
required Gemma2-specific fields such as layer_types and attn_logit_softcapping.
Configure the global-attention case with layer_types set to full_attention
rather than relying only on sliding_window = None, while preserving the settings
needed by QueryOnlyGemma2Attention.
| 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) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Assert which verifier cache is selected.
Both caches have identical shapes and only the output shape is checked, so these tests pass if routing is inverted, ignored, or replaced with zeros. Run paired forwards that vary only the active cache and confirm changes to the inactive cache leave output unchanged.
As per path instructions, ensure new code paths introduced in the PR are covered.
Also applies to: 63-85
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/models/test_mtp_gemma2.py` around lines 37 - 60, Strengthen
test_query_only_gemma2_attention_local_kv and the related tests so they verify
cache routing rather than only output shape: run paired forwards while varying
only the selected verifier cache, assert the output changes, then vary the
inactive cache and assert the output remains unchanged. Cover both local and
global cache paths introduced by the PR, using distinct cache values and the
appropriate configuration or inputs to select each path.
Source: Path instructions
| # Mock position embeddings (cos, sin) | ||
| cos = torch.randn(1, 1, seq_len, 16) | ||
| sin = torch.randn(1, 1, seq_len, 16) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== test file =="
sed -n '1,140p' tests/unit/models/test_mtp_gemma2.py
echo
echo "== search for Gemma2 / rotary usage =="
rg -n "apply_rotary_pos_emb|cos = torch.randn|sin = torch.randn|rope|Gemma2" tests src . -g '!**/.git/**'Repository: vllm-project/speculators
Length of output: 20400
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== mtp Gemma2 implementation =="
sed -n '200,280p' src/speculators/models/mtp/model_definitions.py
echo
echo "== dflash rotary helper (same pattern) =="
sed -n '1,220p' src/speculators/models/dflash/model_definitions.pyRepository: vllm-project/speculators
Length of output: 12186
🏁 Script executed:
python3 - <<'PY'
import inspect
from transformers.models.gemma2.modeling_gemma2 import apply_rotary_pos_emb
print("signature:", inspect.signature(apply_rotary_pos_emb))
print("\nsource:")
src = inspect.getsource(apply_rotary_pos_emb)
print(src)
PYRepository: vllm-project/speculators
Length of output: 1692
🏁 Script executed:
python3 - <<'PY'
import torch
from transformers.models.gemma2.modeling_gemma2 import apply_rotary_pos_emb
B, Hq, Hk, S, D = 2, 4, 2, 5, 16
q = torch.randn(B, Hq, S, D)
k = torch.randn(B, Hk, S, D)
cases = {
"expected_3d": (torch.randn(B, S, D), torch.randn(B, S, D)),
"current_4d": (torch.randn(1, 1, S, D), torch.randn(1, 1, S, D)),
}
for name, (cos, sin) in cases.items():
print(f"\n{name}: input cos shape={tuple(cos.shape)}")
try:
q_embed, k_embed = apply_rotary_pos_emb(q, k, cos, sin)
print(" q_embed:", tuple(q_embed.shape))
print(" k_embed:", tuple(k_embed.shape))
except Exception as e:
print(" error:", type(e).__name__, e)
PYRepository: vllm-project/speculators
Length of output: 352
Use the Gemma2 RoPE input shape here. apply_rotary_pos_emb expects cos/sin as [batch, seq_len, head_dim]; the extra singleton axis makes the rotated tensors 5D and relies on unintended broadcasting. Update both tests to build cos and sin as [batch_sz, seq_len, 16].
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/models/test_mtp_gemma2.py` around lines 48 - 50, The RoPE test
fixtures use an incorrect extra singleton dimension. In both tests around the
rotary embedding setup, update the `cos` and `sin` tensors to shape `[batch_sz,
seq_len, 16]` instead of `[1, 1, seq_len, 16]`, matching the input expected by
`apply_rotary_pos_emb`.
| @pytest.fixture | ||
| def head_setup(): | ||
| hidden_size = 16 | ||
| vocab_size = 256 | ||
| num_centroids = 4 | ||
| head = MultiLevelLMHead(hidden_size, vocab_size, num_centroids) | ||
| return head |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Add a test for non-divisible vocab_size/num_centroids.
All tests use vocab_size=256, num_centroids=4 (evenly divisible). Add a case where vocab_size % num_centroids != 0 to cover the vocab-mapping boundary — this would have caught the out-of-range centroid index issue flagged in head.py. As per path instructions, PyTest tests should cover edge cases specific to speculative decoding such as vocab mapping boundaries.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/models/test_mtp_head.py` around lines 10 - 16, The test fixture
only covers evenly divisible vocabulary and centroid sizes. Add a PyTest case
alongside `head_setup` using values where `vocab_size % num_centroids != 0`,
then exercise the relevant `MultiLevelLMHead` forward or vocab-mapping behavior
and assert outputs remain valid without out-of-range centroid indices, covering
the mapping boundary for speculative decoding.
Source: Path instructions
orestis-z
left a comment
orestis-z
left a comment
There was a problem hiding this comment.
Design looks solid. The query-only cross-attention into verifier KV caches is the right approach for Gemma4's shared-cache architecture, and the kwargs propagation path through Gemma2DecoderLayer.forward → self.self_attn is confirmed working (the installed transformers 5.12.1 passes **kwargs through). Agree with @coderabbitai on the sliding-window mask gap (the eager attention path silently drops the sliding_window kwarg, so local layers see the full causal prefix during training) and on strengthening the attention tests to assert cache routing, not just output shape. Recommend approving once those are addressed.
🤖 Generated with Claude Code using the /pr-review skill
20c69b3 to
fb6b2ae
Compare
|
Thanks for the detailed review @coderabbitai and @orestis-z! I've pushed a new commit addressing all the remaining edge cases and test gaps:
|
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -e
echo "=== base_components.py HAS_GEMMA2 ==="
rg -n "HAS_GEMMA2|gemma2" src/speculators/models/base_components.py
echo "=== QueryOnlyGemma2Attention sliding window ==="
rg -n "sliding_window|triu|tril|eager_attention_forward" src/speculators/models/mtp/model_definitions.py
echo "=== test_mtp_gemma2.py ==="
cat -n tests/unit/models/test_mtp_gemma2.py
echo "=== create_empty_sample kv_feature_dim ==="
rg -n "kv_feature_dim|create_empty_sample" src/speculators/train/data.py -A5Length of output: 8088 🏁 Script executed: #!/bin/bash
sed -n '190,300p' src/speculators/models/mtp/model_definitions.pyLength of output: 5229
One small, non-blocking observation: the sliding-window masking fix in Nice work overall — this looks ready from my side pending any outstanding feedback from (=^･ω･^=) ﾍ(=^･ω･^)ﾉ |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
gh pr diff 768 --repo vllm-project/speculators > /tmp/pr768.diff || true
wc -l /tmp/pr768.diff
sed -n '1,400p' /tmp/pr768.diffLength of output: 17682 🏁 Script executed: #!/bin/bash
sed -n '400,851p' /tmp/pr768.diffLength of output: 16990 🏁 Script executed: #!/bin/bash
python3 -c "
import transformers, inspect
print(transformers.__version__)
from transformers.models.gemma2 import modeling_gemma2 as m
src = inspect.getsource(m.Gemma2Attention.__init__)
print(src)
"
</dev/nullLength of output: 1134 🏁 Script executed: #!/bin/bash
python3 -c "
import transformers
print(transformers.__version__)
"
pip show transformers 2>/dev/null | head -5
python3 - <<'EOF'
import inspect
from transformers.models.gemma2 import modeling_gemma2 as m
import re
src = inspect.getsource(m)
# find sliding_window occurrences
for i, line in enumerate(src.splitlines()):
if 'sliding_window' in line:
print(i, line)
EOFLength of output: 1294 🏁 Script executed: #!/bin/bash
python3 - <<'EOF'
import inspect
from transformers.models.gemma2 import modeling_gemma2 as m
print(inspect.getsource(m.eager_attention_forward))
EOFLength of output: 1525 🏁 Script executed: #!/bin/bash
rg -n 'top_k_centroids|centroid_intermediate_top_k' --type=pyLength of output: 381 🏁 Script executed: #!/bin/bash
sed -n '1,120p' src/speculators/models/mtp/head.py 2>/dev/null
echo "---"
rg -n 'token_ordering_inv|_keys_to_ignore_on_load_missing' src/speculators/models/mtp/head.py src/speculators/models/mtp/core.pyLength of output: 2958
Two smaller items I noticed while reviewing, non-blocking:
Overall this looks good to merge from my side; the sliding-window/KV-routing fixes are sound. Nice work. ✏️ Learnings added
|
|
Thanks for the incredibly thorough review @coderabbitai! I have addressed all three of your non-blocking items:
I just pushed these updates to this branch. We should be totally good to go! |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
gh pr diff 768 --repo vllm-project/speculators > /tmp/pr768_new.diff || true
echo "=== token_ordering_inv persistent=False ==="
rg -n "token_ordering_inv|persistent" -A2 -B2 /tmp/pr768_new.diff | head -60
echo "=== new sliding window mask test ==="
rg -n "test_query_only_gemma2_attention_sliding_window_mask" -A60 /tmp/pr768_new.diff | head -120Length of output: 4112
All three items check out. This looks good to merge from my side. (=^･ω･^=) ﾍ(=^･ω･^)ﾉ |
orestis-z
left a comment
orestis-z
left a comment
There was a problem hiding this comment.
Both outstanding items from my previous review are addressed:
- Sliding-window mask: The eager attention path now explicitly enforces the window constraint on the causal mask, so local layers no longer silently see the full causal prefix.
- Cache-routing tests: Tests now verify that changing the active cache changes output while changing the inactive cache does not, and a dedicated sliding-window mask test confirms tokens outside the window are correctly masked.
The coderabbitai feedback (NameError guard, centroid validation, inverse caching, empty-sample KV shape) has also been addressed. LGTM. Recommend approving.
🤖 Generated with Claude Code using the /pr-review skill
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
The query-only cross-attention design for Gemma4 MTP is clean, and the kwargs propagation path through Gemma2DecoderLayer.forward(**kwargs) into QueryOnlyGemma2Attention is correct. One shape-mismatch bug in core.py will crash Gemma2 MTP training at runtime -- details in the inline comment.
🤖 Generated with Claude Code using the /pr-review skill
| verifier_kv_last_local=kwargs.get("verifier_kv_last_local"), | ||
| verifier_kv_last_global=kwargs.get("verifier_kv_last_global"), |
There was a problem hiding this comment.
Bug: verifier KV caches are not sliced to valid_len, causing a shape mismatch during attention.
All other sequence-dimension tensors in the loop are sliced to valid_len (e.g., step_hidden = current_hidden[:, :valid_len] at L202, and the causal mask is created for valid_len at L191–198). The KV caches are passed through at full seq_len.
Inside eager_attention_forward the attention scores have shape [batch, heads, valid_len, seq_len] but the causal mask has shape [batch, 1, valid_len, valid_len]. The element-wise add crashes:
RuntimeError: The size of tensor a (seq_len) must match the size of
tensor b (valid_len) at non-singleton dimension 3
(I reproduced this locally with seq_len=10, valid_len=6.)
Fix: slice once before the loop and pass the sliced tensors:
# before the for-loop
step_kv_local = kwargs.get("verifier_kv_last_local")
step_kv_global = kwargs.get("verifier_kv_last_global")
if step_kv_local is not None:
step_kv_local = step_kv_local[:, :valid_len]
if step_kv_global is not None:
step_kv_global = step_kv_global[:, :valid_len]then pass step_kv_local / step_kv_global at L214–215 instead of re-reading from kwargs.
The unit tests in test_mtp_gemma2.py do not catch this because they feed KV tensors with the same seq_len as the query and never go through MTPDraftModel.forward() (which is where the slicing to valid_len happens).
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
…dels 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 vllm-project#586, vllm-project#758 Signed-off-by: santosh-agebold <12378904+shotsan@users.noreply.github.com> Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
79f837f to
b9101e9
Compare
Signed-off-by: shotsan <shotsan@users.noreply.github.com>
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
3 participants
Context
This PR introduces the final architecture modifications required for Gemma4 MTP draft models, completing the 3-part rollout strategy outlined in Issue #586.
Unlike Qwen-style MTP which maintains its own full self-attention layers and distinct KV caches, Gemma4 MTP borrows the verifier's local and global KV caches directly, stripping K/V projections out of its own layers.
Dependencies:
(Note: Commits from feat(mtp): add multi-level LM head for Gemma4 MTP (2/3) #767 currently appear in this PR as this was branched from it for stacking. This will resolve once feat(mtp): add multi-level LM head for Gemma4 MTP (2/3) #767 is merged).
Implementation Details
Gemma2DecoderLayer,Gemma2RMSNorm, andGemma2RotaryEmbeddingto thebase_components.pyregistry.QueryOnlyGemma2Attention: Built a specialized customized self-attention module that overrides the baseGemma2Attention. It zeroes out K/V linear projections and instead cross-attends intoverifier_kv_last_localandverifier_kv_last_global.sliding_windowsize.Testing
tests/unit/models/test_mtp_gemma2.pyexplicitly mock out the KV cache tensor routing, verifying that local vs. global tensor routing succeeds based on window configuration and shape propagation works as expected.