[Tests][Spec Decode] Add gemma4 MTP acceptance rates test (#47920)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Jul 27, 2026 at 16:15 UTC
272abd5f486967f1fb9db7ca7504f8c34235ef50
4 files changed
+120
-48
.buildkite/test_areas/spec_decode.yaml
+12
@@ -188,3 +188,15 @@ steps:
188
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
189
commands:
190
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
191
+
192
+- label: Spec Decode Acceptance Rates Nightly
193
+ key: spec-decode-acceptance-rates-nightly
194
+ timeout_in_minutes: 60
195
+ device: h200_35gb
196
+ optional: true
197
+ source_file_dependencies:
198
+ - vllm/v1/spec_decode/
199
+ - vllm/v1/worker/gpu/spec_decode/
200
+ - tests/v1/e2e/spec_decode/
201
+ commands:
202
+ - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "acceptance_rates"
tests/v1/e2e/spec_decode/test_spec_decode.py
+73
-42
@@ -1393,51 +1393,81 @@ def load_and_process_dataset(data_name: str):
1393
return dataset
1394
1395
1396
-@pytest.fixture
1397
-def dflash_config():
1398
- target_model = "Qwen/Qwen3-8B"
1399
- draft_model = "z-lab/Qwen3-8B-DFlash-b16"
1400
-
1401
- return dict(
1402
- model=target_model,
1403
- trust_remote_code=True,
1404
- speculative_config={
1405
- "method": "dflash",
1406
- "model": draft_model,
1407
- "num_speculative_tokens": 16,
1408
- "max_model_len": 32768,
1409
- },
1410
- max_model_len=32768,
1411
- max_num_seqs=128,
1412
- gpu_memory_utilization=0.85,
1413
- enforce_eager=False,
1414
- disable_log_stats=False,
1415
- )
1416
-
1417
-
1396
+@pytest.mark.parametrize(
1397
+ ["spec_config", "expected_acceptance_lengths", "chat_template_kwargs"],
1398
+ [
1399
+ pytest.param(
1400
+ dict(
1401
+ model="Qwen/Qwen3-8B",
1402
+ trust_remote_code=True,
1403
+ speculative_config={
1404
+ "method": "dflash",
1405
+ "model": "z-lab/Qwen3-8B-DFlash-b16",
1406
+ "num_speculative_tokens": 16,
1407
+ "max_model_len": 32768,
1408
+ },
1409
+ max_model_len=32768,
1410
+ max_num_seqs=128,
1411
+ gpu_memory_utilization=0.85,
1412
+ enforce_eager=False,
1413
+ disable_log_stats=False,
1414
+ ),
1415
+ # All scores from Table 1 in https://arxiv.org/pdf/2602.06036
1416
+ {
1417
+ "mt-bench": 4.24,
1418
+ "humaneval": 6.50,
1419
+ # runs with a subset of prompts so extra wide tol here
1420
+ "gsm8k": 6.54 * 0.975,
1421
+ },
1422
+ {"enable_thinking": False},
1423
+ id="dflash",
1424
+ ),
1425
+ pytest.param(
1426
+ dict(
1427
+ model="google/gemma-4-E4B-it",
1428
+ trust_remote_code=True,
1429
+ speculative_config={
1430
+ "method": "mtp",
1431
+ "model": "google/gemma-4-E4B-it-assistant",
1432
+ "num_speculative_tokens": 2,
1433
+ "max_model_len": 32768,
1434
+ },
1435
+ max_model_len=32768,
1436
+ # Skip multimodal profiling; this is a text-only eval.
1437
+ limit_mm_per_prompt={"image": 0, "audio": 0},
1438
+ disable_log_stats=False,
1439
+ ),
1440
+ {
1441
+ "mt-bench": 2.28,
1442
+ "humaneval": 2.68,
1443
+ # runs with a subset of prompts so extra wide tol here
1444
+ "gsm8k": 2.67 * 0.975,
1445
+ },
1446
+ {},
1447
+ id="gemma4",
1448
+ ),
1449
+ ],
1450
+)
1451
@pytest.mark.parametrize("use_mrv2", [False, True])
1419
-def test_dflash_acceptance_rates(
1420
- monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config
1452
+def test_acceptance_rates(
1453
+ monkeypatch: pytest.MonkeyPatch,
1454
+ spec_config: dict[str, Any],
1455
+ expected_acceptance_lengths: dict[str, float],
1456
+ chat_template_kwargs: dict[str, Any],
1457
+ use_mrv2: bool,
1458
):
1459
"""
1423
- E2E test for DFlash (block diffusion) speculative decoding.
1424
- Runs acceptance rate validation on GSM8k, MT-Bench, and HumanEval
1425
- comparing against baseline results from the paper (Table 1).
1426
- See https://github.com/z-lab/dflash/blob/main/benchmark_sglang.py for methodology.
1460
+ E2E acceptance-rate validation for speculative decoding.
1461
+
1462
+ Drives one or more datasets (keyed in ``expected_acceptance_lengths``)
1463
+ through the spec decode engine and asserts the mean acceptance length
1464
+ stays within tolerance of the reference figure for each dataset.
1465
"""
1466
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
1429
-
1430
- spec_llm = LLM(**dflash_config)
1467
+ spec_llm = LLM(**spec_config)
1468
1469
max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k
1470
1434
- # All scores from Table 1 in https://arxiv.org/pdf/2602.06036
1435
- expected_acceptance_lengths = {
1436
- "mt-bench": 4.24,
1437
- "humaneval": 6.50,
1438
- "gsm8k": 6.54 * 0.975, # runs with a subset of prompts so extra wide tol here
1439
- }
1440
-
1471
tokenizer = spec_llm.get_tokenizer()
1472
for dataset_name, expected_len in expected_acceptance_lengths.items():
1473
dataset = load_and_process_dataset(dataset_name)
@@ -1452,10 +1482,11 @@ def test_dflash_acceptance_rates(
1482
[{"role": "user", "content": user_content}],
1483
tokenize=False,
1484
add_generation_prompt=True,
1455
- enable_thinking=False,
1485
+ **chat_template_kwargs,
1486
)
1487
1458
- # Temp=0, MaxTokens=2048 from the paper
1488
+ # Greedy (temp=0) so acceptance length is deterministic and comparable
1489
+ # across runs.
1490
spec_llm.generate(
1491
[prompt_text],
1492
SamplingParams(temperature=0, max_tokens=2048),
@@ -1467,17 +1498,17 @@ def test_dflash_acceptance_rates(
1498
acceptance_lengths.append(acceptance_len)
1499
1500
mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths)
1470
- # Fairly tight tolerance of 95% against the paper's figures,
1501
+ # Fairly tight tolerance of 95% against the reference figures,
1502
# watching for regressions. Can be relaxed if test is flaky but be sure to
1503
# check for genuine issues such as #40727.
1504
expected_len = expected_len * 0.95
1505
print(
1475
- f"DFlash acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}"
1506
+ f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}"
1507
f" (expected at least {expected_len:.2f})"
1508
)
1509
1510
assert mean_acceptance_length >= expected_len, (
1480
- f"DFlash acceptance_len for {dataset_name} is below expected threshold:"
1511
+ f"acceptance_len for {dataset_name} is below expected threshold: "
1512
f"{mean_acceptance_length:.2f} < {expected_len:.2f}"
1513
)
1514
vllm/v1/spec_decode/gemma4.py
+24
@@ -14,6 +14,7 @@ import torch
14
import torch.nn as nn
15
16
from vllm.config import VllmConfig, get_layers_from_vllm_config, replace
17
+from vllm.distributed.parallel_state import get_pp_group
18
from vllm.logger import init_logger
19
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
20
from vllm.v1.attention.backend import CommonAttentionMetadata
@@ -165,6 +166,29 @@ class Gemma4Proposer(SpecDecodeBaseProposer):
166
)
167
return base
168
169
+ def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None:
170
+ """Gemma4 MTP requires dim-mismatched embedding sharing.
171
+
172
+ The draft checkpoint's embed_tokens is a draft-dim placeholder
173
+ (tied to lm_head so load_weights populates both); the model
174
+ expects it to be replaced by the target's backbone-dim embedding,
175
+ so the base class's embedding-dim equality guard must not apply.
176
+ """
177
+ if get_pp_group().world_size != 1:
178
+ return
179
+ inner_model = getattr(target_language_model, "model", None)
180
+ target_embed_tokens = getattr(inner_model, "embed_tokens", None)
181
+ if target_embed_tokens is None:
182
+ raise AttributeError(
183
+ "Target model does not have an 'embed_tokens' attribute"
184
+ )
185
+ del self.model.model.embed_tokens
186
+ self.model.model.embed_tokens = target_embed_tokens
187
+ logger.info(
188
+ "Gemma4 MTP: sharing target model's backbone-dim embed_tokens "
189
+ "with the draft model."
190
+ )
191
+
192
def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:
193
"""Gemma4 MTP always keeps its own draft-dim lm_head.
194
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+11
-6
@@ -446,10 +446,16 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
446
)
447
last_hidden_states = last_hidden_states[:num_reqs]
448
449
+ sample_positions = positions
450
+ if not self.advance_draft_positions:
451
+ # The forward pass holds positions fixed (Q-only, shared target KV),
452
+ # but Gumbel sampling still needs the absolute draft position.
453
+ sample_positions = positions + self.current_draft_step
454
+
455
# Sample the draft tokens.
456
draft_tokens = self.sample_draft(
457
last_hidden_states,
452
- positions,
458
+ sample_positions,
459
idx_mapping,
460
self.temperature,
461
self.seeds,
@@ -630,6 +636,9 @@ def _prepare_decode_inputs_kernel(
636
draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride)
637
tl.store(input_ids_ptr + req_idx, draft_token)
638
639
+ target_seq_len = tl.load(target_seq_lens_ptr + req_idx)
640
+ num_rejected = tl.load(num_rejected_ptr + req_idx)
641
+ seq_len = target_seq_len - num_rejected
642
if ADVANCE_DRAFT_POSITIONS:
643
# Compute position and seq_lens.
644
# NOTE(woosuk): To prevent out-of-range access, we clamp these values
@@ -637,12 +646,8 @@ def _prepare_decode_inputs_kernel(
646
position = tl.load(positions_ptr + req_idx)
647
position = tl.minimum(position + 1, max_model_len - 1)
648
tl.store(positions_ptr + req_idx, position)
640
-
641
- target_seq_len = tl.load(target_seq_lens_ptr + req_idx)
642
- num_rejected = tl.load(num_rejected_ptr + req_idx)
643
- seq_len = target_seq_len - num_rejected
649
seq_len = tl.minimum(seq_len + 1, max_model_len)
645
- tl.store(seq_lens_ptr + req_idx, seq_len)
650
+ tl.store(seq_lens_ptr + req_idx, seq_len)
651
652
653
def prepare_decode_inputs(