[Bugfix][MRV2] Read routed experts from unkeyed full blocks

.buildkite/test_areas/spec_decode.yaml

@@ -40,6 +40,7 @@ steps:
   timeout_in_minutes: 25
   device: h200_35gb
   source_file_dependencies:
+    - vllm/distributed/aux_output_connector/
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
     - tests/v1/e2e/spec_decode/
@@ -55,6 +56,7 @@ steps:
       depends_on:
       - image-build-amd
       source_file_dependencies:
+      - vllm/distributed/aux_output_connector/
       - vllm/v1/spec_decode/
       - vllm/v1/worker/gpu/spec_decode/
       - vllm/model_executor/model_loader/

tests/distributed/aux_output_connector/test_store.py

@@ -547,6 +547,28 @@ def test_logical_buffer_copies_borrowed_block_when_retained():
     buffer.release_block(retained)
 
 
+@pytest.mark.parametrize("token_end", [8, 10])
+@pytest.mark.parametrize("token_start", [0, 2, 5])
+def test_logical_buffer_reads_unkeyed_blocks_and_tail(token_start, token_end):
+    """Full blocks awaiting hashes remain readable after leaving the tail."""
+    buffer = RoutedExpertsBuffer(_DTYPE, _SHAPE, _BLOCK_SIZE, 1, 12, 2)
+    logical = np.arange(token_end * 6, dtype=_DTYPE).reshape(token_end, *_SHAPE)
+    pending = [
+        (start, buffer.retain_block(rows))
+        for start, rows in buffer.capture("request", 0, logical)
+    ]
+
+    output = buffer.read("request", token_start, token_end, pending)
+    np.testing.assert_array_equal(output, logical[token_start:])
+
+    for _, rows in pending:
+        rows[:] = 0
+        buffer.release_block(rows)
+    buffer.reset()
+    buffer.capture("replacement", 0, np.zeros_like(logical[:2]))
+    np.testing.assert_array_equal(output, logical[token_start:])
+
+
 def _make_store(
     *,
     max_bytes: int = 1 << 20,
@@ -1124,9 +1146,10 @@ def test_worker_emits_mid_block_chunked_prefill():
     worker.close()
 
 
-def test_worker_emits_published_block_and_mid_block_tail():
+@pytest.mark.parametrize("token_end", [7, 8, 10])
+def test_worker_emits_published_blocks_with_unkeyed_rows(token_end):
     worker = _make_worker(1)
-    logical = np.arange(7 * 3 * 2, dtype=np.uint8).reshape(7, 3, 2)
+    logical = np.arange(token_end * 6, dtype=np.uint8).reshape(token_end, *_SHAPE)
 
     first = _metadata(
         0,
@@ -1145,7 +1168,7 @@ def test_worker_emits_published_block_and_mid_block_tail():
 
     second = _metadata(
         0,
-        [_request_metadata("request", 6, 1, 0, [])],
+        [_request_metadata("request", 6, token_end - 6, 0, [])],
         {},
     )
     output = _process_output(worker, second, logical[6:], ["request"], np.array([0]))

tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py

@@ -1,18 +1,66 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import numpy as np
 import pytest
 
 from tests.utils import (
     get_attn_backend_list_based_on_platform,
     single_gpu_only,
 )
 from vllm import SamplingParams
+from vllm.inputs import TokensPrompt
 from vllm.platforms import current_platform
 
 from .utils import _run_eagle_correctness
 
 
+@single_gpu_only
+@pytest.mark.skipif(
+    current_platform.is_device_capability_family(100),
+    reason="DeepSeek head_dim=192 not supported on SM100/SM110 (Blackwell)",
+)
+def test_deepseek_eagle_routed_experts_with_cached_aligned_prompt(
+    monkeypatch: pytest.MonkeyPatch, vllm_runner
+):
+    """A cached prefix must include routing from the final, still unkeyed block."""
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    monkeypatch.setenv("VLLM_MLA_DISABLE", "1")
+    with vllm_runner(
+        "eagle618/deepseek-v3-random",
+        max_model_len=256,
+        max_num_batched_tokens=256,
+        max_num_seqs=1,
+        block_size=16,
+        enable_prefix_caching=True,
+        enable_return_routed_experts=True,
+        enforce_eager=True,
+        speculative_config={
+            "method": "eagle",
+            "model": "eagle618/eagle-deepseek-v3-random",
+            "num_speculative_tokens": 1,
+        },
+    ) as runner:
+        prompt = TokensPrompt(prompt_token_ids=[100] * 48)
+        params = SamplingParams(temperature=0, max_tokens=1, ignore_eos=True)
+        cold = runner.llm.generate(prompt, params, use_tqdm=False)[0]
+        warm = runner.llm.generate(prompt, params, use_tqdm=False)[0]
+
+        assert cold.num_cached_tokens == 0
+        assert warm.num_cached_tokens == 32
+        cold_routes = cold.outputs[0].routed_experts
+        warm_routes = warm.outputs[0].routed_experts
+        assert cold_routes is not None and warm_routes is not None
+        model_config = runner.llm.llm_engine.vllm_config.model_config
+        assert warm_routes.shape == (
+            48,
+            model_config.get_total_num_hidden_layers(),
+            model_config.get_num_experts_per_tok(),
+        )
+        assert (warm_routes < model_config.get_num_experts()).all()
+        np.testing.assert_array_equal(warm_routes[:32], cold_routes[:32])
+
+
 @single_gpu_only
 @pytest.mark.skipif(
     current_platform.is_device_capability_family(100),

vllm/distributed/aux_output_connector/routed_experts.py

@@ -117,20 +117,31 @@ def capture(
         return completed
 
     def read(
-        self, request_id: Hashable, token_start: int, token_end: int
+        self,
+        request_id: Hashable,
+        token_start: int,
+        token_end: int,
+        pending_blocks: Sequence[tuple[int, np.ndarray]],
     ) -> np.ndarray:
+        """Copy rows from unkeyed full blocks and the incomplete tail."""
+        blocks = list(pending_blocks)
         tail = self._requests.get(request_id)
-        assert tail is not None, (
-            f"auxiliary output buffer is missing request {request_id}"
-        )
-        local_start = token_start - tail.block_start
-        local_end = token_end - tail.block_start
-        assert 0 <= local_start < local_end <= tail.length, (
+        if tail is not None:
+            blocks.append((tail.block_start, self._rows[tail.slot, : tail.length]))
+        parts = []
+        position = token_start
+        for block_start, rows in blocks:
+            if position < block_start:
+                break
+            part = rows[position - block_start : token_end - block_start]
+            parts.append(part)
+            position += len(part)
+        assert position == token_end, (
             "auxiliary output range is unavailable: "
             f"request={request_id}, range=[{token_start}, {token_end}), "
-            f"available=[{tail.block_start}, {tail.block_start + tail.length})"
+            f"available={[(start, start + len(block)) for start, block in blocks]}"
         )
-        return self._rows[tail.slot, local_start:local_end].copy()
+        return np.concatenate(parts)
 
     def retain_block(self, rows: np.ndarray) -> np.ndarray:
         """Retain one unkeyed block after the current capture call."""

vllm/distributed/aux_output_connector/worker.py

@@ -282,11 +282,14 @@ def _commit_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutpu
                 local_start = emit_start % block_size
                 rows = stored[local_start : local_start + stored_end - emit_start]
                 if stored_end < token_end:
-                    rows = np.concatenate(
-                        (rows, buffer.read(request_id, stored_end, token_end))
+                    suffix = buffer.read(
+                        request_id, stored_end, token_end, state.pending_blocks
                     )
+                    rows = np.concatenate((rows, suffix))
             else:
-                rows = buffer.read(request_id, emit_start, token_end)
+                rows = buffer.read(
+                    request_id, emit_start, token_end, state.pending_blocks
+                )
             outputs[request_id] = AuxRequestOutput(emit_start, rows)
             state.emit_cursor = token_end
         return outputs