[Bugfix][Spec Decode] Build MRV2 draft attention groups with the allocated kernel block sizes

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)

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(

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

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -250,13 +250,15 @@ 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,
             kv_cache_config,
             block_tables,
             target_input_buffers,
             target_attn_groups,
+            kernel_block_sizes,
         )
 
         # FlashAttention's AOT split schedule is wrong for a windowed drafter,

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

vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py

@@ -85,13 +85,15 @@ 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,
             kv_cache_config,
             block_tables,
             target_input_buffers,
             target_attn_groups,
+            kernel_block_sizes,
         )
         self.expanded_slot_mappings = torch.empty(
             block_tables.num_kv_cache_groups,