From 251af6f4eab2af24fad6fcf3ea5b0c0b7c86bcd2 Mon Sep 17 00:00:00 2001 From: Khang Le Duy Date: Fri, 30 Oct 2026 16:17:36 +0700 Subject: [PATCH] [Bugfix][Spec Decode] Build MRV2 draft attention groups with the allocated kernel block sizes The Model Runner V2 speculator re-picked each KV cache group's kernel block size from the drafter's backends alone. When the drafter shares a group with target layers whose backend supports only smaller kernel pages, the KV cache is allocated and viewed with the target's split pages while draft decode steps 2+ built unsplit block tables, so they read the wrong KV and acceptance after the first draft position collapsed (#64652). Pass the runner's kernel block sizes to the speculator, as Model Runner V1 already does for its drafter. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Khang Le Duy --- .../eagle/test_eagle_correctness.py | 43 +++++++++++++++++++ vllm/v1/worker/gpu/attn_utils.py | 4 +- vllm/v1/worker/gpu/model_runner.py | 1 + .../gpu/spec_decode/dflash/speculator.py | 2 + vllm/v1/worker/gpu/spec_decode/speculator.py | 2 + .../spec_decode/standalone_ar/speculator.py | 2 + 6 files changed, 53 insertions(+), 1 deletion(-) diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py b/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py index bb8ff8e57f27..6f2029df2141 100644 --- a/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py +++ b/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py @@ -10,6 +10,7 @@ from vllm import SamplingParams from vllm.platforms import current_platform +from ..utils import get_spec_decode_metric_value, get_test_prompts from .utils import _run_eagle_correctness @@ -170,3 +171,45 @@ def test_eagle_correctness_medium( attn_backend, vllm_runner, ) + + +@single_gpu_only +@pytest.mark.skipif( + not current_platform.is_cuda() or current_platform.is_device_capability_family(100), + reason="FLASHINFER splits 128-token blocks into kernel pages only below SM100", +) +def test_qwen3_eagle3_split_kernel_pages_keep_draft_acceptance( + monkeypatch: pytest.MonkeyPatch, vllm_runner +): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + prompts = get_test_prompts(mm_enabled=False) + sampling_params = SamplingParams(temperature=0, max_tokens=128) + acceptance_per_pos = {} + for block_size in (64, 128): + with vllm_runner( + "Qwen/Qwen3-8B", + block_size=block_size, + max_model_len=2048, + attention_config={"backend": "FLASHINFER"}, + speculative_config={ + "method": "eagle3", + "model": "AngelSlim/Qwen3-8B_eagle3", + "num_speculative_tokens": 3, + "attention_backend": "TRITON_ATTN", + }, + disable_log_stats=False, + ) as runner: + runner.llm.chat(prompts, sampling_params) + metrics = runner.llm.get_metrics() + num_drafts = get_spec_decode_metric_value( + metrics, "vllm:spec_decode_num_drafts" + ) + accepted_per_pos = next( + m + for m in metrics + if m.name == "vllm:spec_decode_num_accepted_tokens_per_pos" + ) + acceptance_per_pos[block_size] = [ + n / num_drafts for n in accepted_per_pos.values + ] + assert acceptance_per_pos[128] == pytest.approx(acceptance_per_pos[64], abs=0.05) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 03a81881b755..bcde0e405eb8 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -200,6 +200,7 @@ def init_attn_backend( device: torch.device, active_layer_names: set[str] | None = None, draft_layer_names: set[str] | None = None, + kernel_block_sizes: list[int] | None = None, ) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]: # Phase 1: discover attention groups for each kv cache group. attn_groups: list[list[AttentionGroup]] = [] @@ -266,7 +267,8 @@ def init_attn_backend( # Phase 2: pick a kernel block size per kv cache group that is supported # by all backends within that group. - kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups) + if kernel_block_sizes is None: + kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups) if current_platform.is_xpu(): for kv_cache_group_id, selected_kernel_block_size in enumerate( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 268e137882ff..d3a2d8192649 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -800,6 +800,7 @@ def initialize_kv_cache( self.block_tables, self.input_buffers, self.attn_groups, + self.kernel_block_sizes, ) if self.speculator is not None: # After set_attn, so the speculator can size its cudagraph mode diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index d12a8aa1e9d1..e47482f706dd 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -250,6 +250,7 @@ def set_attn( block_tables: BlockTables, target_input_buffers: InputBuffers, target_attn_groups: list[list[AttentionGroup]], + kernel_block_sizes: list[int], ) -> None: super().set_attn( model_state, @@ -257,6 +258,7 @@ def set_attn( block_tables, target_input_buffers, target_attn_groups, + kernel_block_sizes, ) # FlashAttention's AOT split schedule is wrong for a windowed drafter, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index bcd6bfd608e4..bfc4e7540574 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -337,6 +337,7 @@ def set_attn( block_tables: BlockTables, target_input_buffers: InputBuffers, target_attn_groups: list[list[AttentionGroup]], + kernel_block_sizes: list[int], ) -> None: self.model_state = model_state self.kv_cache_config = kv_cache_config @@ -345,6 +346,7 @@ def set_attn( self.attn_vllm_config, self.device, active_layer_names=self.draft_attn_layer_names, + kernel_block_sizes=kernel_block_sizes, ) self.block_tables = block_tables # The target model runner's buffers and attention groups. Draft diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index dc8ff567d0c4..6f463a7deb32 100644 --- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -85,6 +85,7 @@ def set_attn( block_tables: BlockTables, target_input_buffers: InputBuffers, target_attn_groups: list[list[AttentionGroup]], + kernel_block_sizes: list[int], ) -> None: super().set_attn( model_state, @@ -92,6 +93,7 @@ def set_attn( block_tables, target_input_buffers, target_attn_groups, + kernel_block_sizes, ) self.expanded_slot_mappings = torch.empty( block_tables.num_kv_cache_groups,