feat: add out-of-tree Gemma4 KV connector#674
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. 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 ChangesGemma4 KV Connector and E2E Tests
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 4
🤖 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/data_generation/gemma4_kv_connector.py`:
- Around line 72-76: The KV extraction logic indexes all entries from
slot_mapping before truncating, so update the slot_mapping handling in this
kv_cache block to validate the length against num_tokens and trim it first. In
the data generation path that computes block_idx and block_off, ensure any extra
mapped slots are ignored and any short mapping is rejected or handled explicitly
before using kv_cache[:, 0] and kv_cache[:, 1]. Keep the fix localized to the
hidden-state/KV save path in gemma4_kv_connector.py so the extracted k and v
rows always align with num_tokens.
- Around line 70-71: The KV layout checks in gemma4_kv_connector should not use
assert because they can be skipped when Python runs with optimizations. Replace
the assertions in the KV validation logic with explicit exception raising in the
same spot, using the existing _bad_layout_msg and the surrounding KV cache
validation path so unsupported pipeline_parallel_size, quantized cache settings,
or malformed kv_cache tensors fail immediately with a clear error.
- Around line 252-260: The current head-stride logic in _gather_kv_heads is
already covered for standard TP cases, so move this from correctness to coverage
by adding a test that exercises total_kv_heads < tensor_parallel_size. The test
should verify the gathered KV head count path in gemma4_kv_connector.py and
assert that a non-divisible gather shape is rejected, while keeping the existing
behavior of deduping contiguous GQA replicas intact.
In `@tests/e2e/utils.py`:
- Around line 508-513: The speculative-decoding helper test is not verifying
that any request actually exceeds the 512-token sliding window, so the overflow
path may never be exercised. Update the test logic around the prompts batch in
the relevant helper(s) to assert via results["per_request"] that at least one
request token count crosses 512, using the existing prompt-generation setup and
the per-request output to confirm the long prompt truly triggers the
sliding-window overflow case.
🪄 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: 3f2cf2b8-34de-4808-9736-673a89209b2c
📒 Files selected for processing (4)
src/speculators/data_generation/gemma4_kv_connector.pytests/e2e/run_gemma4_kv_extraction.pytests/e2e/smoke/test_gemma4_kv_connector.pytests/e2e/utils.py
| 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() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect current gather/dedup coverage and configs mentioning KV head counts.
rg -n "tensor_parallel_size|num_key_value_heads|total_kv_heads|_gather_kv_heads|GQA|dedup" src testsRepository: vllm-project/speculators
Length of output: 4898
Shift this to coverage, not correctness. The code already dedups contiguous GQA replicas in _gather_kv_heads, and the smoke tests cover both TP=1 and TP=2. Add a case with total_kv_heads < tensor_parallel_size to lock down the head-count path and reject any non-divisible gather 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/data_generation/gemma4_kv_connector.py` around lines 252 -
260, The current head-stride logic in _gather_kv_heads is already covered for
standard TP cases, so move this from correctness to coverage by adding a test
that exercises total_kv_heads < tensor_parallel_size. The test should verify the
gathered KV head count path in gemma4_kv_connector.py and assert that a
non-divisible gather shape is rejected, while keeping the existing behavior of
deduping contiguous GQA replicas intact.
Source: Path instructions
| # Short, and a long prompt past the 512 sliding window, in one batch. | ||
| prompts = [ | ||
| "Hello world", | ||
| "Test prompt with several tokens", | ||
| " ".join(["word"] * 600), | ||
| ] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Assert that one request really crosses the 512-token sliding window.
Right now the test assumes the repeated "word" prompt tokenizes past 512, but nothing checks that in results["per_request"]. If tokenization ever stays at or below 512, this helper still passes without exercising the overflow case it is supposed to cover.
Suggested fix
assert results["num_outputs"] == len(prompts)
assert len(results["per_request"]) == len(prompts)
+ assert any(req["num_tokens"] > 512 for req in results["per_request"]), (
+ "expected at least one prompt to exceed the 512-token sliding window"
+ )
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"],As per path instructions, "Ensure PyTest tests are clear, comprehensive, and cover edge cases specific to speculative decoding..."
Also applies to: 551-583
🤖 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/e2e/utils.py` around lines 508 - 513, The speculative-decoding helper
test is not verifying that any request actually exceeds the 512-token sliding
window, so the overflow path may never be exercised. Update the test logic
around the prompts batch in the relevant helper(s) to assert via
results["per_request"] that at least one request token count crosses 512, using
the existing prompt-generation setup and the per-request output to confirm the
long prompt truly triggers the sliding-window overflow case.
Source: Path instructions
orestis-z
left a comment
orestis-z
left a comment
There was a problem hiding this comment.
LGTM. Clean Phase 1 implementation of the vLLM RFC. The TP gather/dedup path, CUDA stream synchronization, and layer resolution logic all look correct. Good test coverage with both tp=1 and tp=2 smoke tests. Recommend approving.
🤖 Generated with Claude Code using the /pr-review skill
| GEMMA4_KV_KEYS = ( | ||
| "kv_last_local_k", | ||
| "kv_last_local_v", | ||
| "kv_last_global_k", | ||
| "kv_last_global_v", | ||
| ) | ||
|
|
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
LGTM. Clean out-of-tree KV connector implementation matching Phase 1 of the vLLM RFC. One non-blocking note: all commits are missing Signed-off-by lines (the DCO CI check is also failing). The author will need to amend or rebase with -s before merge.
🤖 Generated with Claude Code using the /pr-review skill
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.
|
4 participants
Purpose
RFC. Added
Gemma4KVConnector, an out-of-tree vLLM KV connector that produces training data for finetuning Gemma4 MTP draft models.Load it via:
Tests
tests/e2e/run_gemma4_kv_extraction.pyrun_gemma4_kv_extractionhelper intests/e2e/utils.pyChecklist
I have filled in: