[Feature][MRV2] Adapt extract_hidden_states for Model Runner V2

docs/source/user_guide/feature_guide/speculative_decoding.md

@@ -377,6 +377,18 @@ The `extract_hidden_states` method is a special speculative decoding mode that d
 > [!NOTE]
 > This method produces only 1 output token per request. The primary output is the hidden states saved to disk, not the generated text.
 
+Both Model Runner V1 and Model Runner V2 are supported on Ascend. Enable V2 with:
+
+```shell
+export VLLM_USE_V2_MODEL_RUNNER=1
+```
+
+> [!NOTE]
+> Model Runner V2 support reuses upstream vLLM's `ExtractHiddenStatesSpeculator`
+> ([PR #49811](https://github.com/vllm-project/vllm/pull/49811)). Ascend only
+> adds dispatch, aux-hidden enabling, and NPU KV allocate/reshape for
+> `HiddenStateCacheSpec`.
+
 - Offline inference
 
     ```python

tests/e2e/pull_request/one_card/spec_decode/test_extract_hidden_states.py

@@ -23,6 +23,8 @@
 * a hybrid attention model (Qwen3.5-0.8B, GatedDeltaNet + full_attention)
   loaded with dummy weights as a shape/round-trip smoke test. The hybrid case
   mirrors upstream vLLM PR #39949.
+* Model Runner V1 (default) and Model Runner V2 (`VLLM_USE_V2_MODEL_RUNNER=1`),
+  covering the Ascend adaptation of upstream vLLM PR #49811.
 """
 
 from __future__ import annotations
@@ -62,6 +64,8 @@ class ExtractHiddenStatesCase:
     verify_nonzero: bool = True
     # Hybrid smoke test additionally checks the token_ids round-trip.
     verify_token_ids: bool = False
+    # When True, force Model Runner V2 via VLLM_USE_V2_MODEL_RUNNER.
+    use_v2_model_runner: bool = False
 
 
 CASES = [
@@ -110,6 +114,39 @@ class ExtractHiddenStatesCase:
         ),
         id="hybrid_dummy_eager",
     ),
+    pytest.param(
+        ExtractHiddenStatesCase(
+            model_name=DENSE_MODEL,
+            aux_hidden_state_layer_ids=DENSE_AUX_HIDDEN_STATE_LAYER_IDS,
+            prompts=[
+                "Hello, how are you?",
+                "What is machine learning?",
+            ],
+            enforce_eager=True,
+            gpu_memory_utilization=0.8,
+            max_num_seqs=16,
+            use_v2_model_runner=True,
+        ),
+        id="dense_eager_mrv2",
+    ),
+    pytest.param(
+        ExtractHiddenStatesCase(
+            model_name=HYBRID_MODEL,
+            aux_hidden_state_layer_ids=HYBRID_AUX_HIDDEN_STATE_LAYER_IDS,
+            prompts=[
+                "Hello world",
+                "Test prompt with several tokens",
+            ],
+            enforce_eager=True,
+            gpu_memory_utilization=0.4,
+            max_model_len=256,
+            load_format="dummy",
+            verify_nonzero=False,
+            verify_token_ids=True,
+            use_v2_model_runner=True,
+        ),
+        id="hybrid_dummy_eager_mrv2",
+    ),
 ]
 
 
@@ -143,8 +180,13 @@ def _verify_output(output, expected_shape, *, verify_nonzero, verify_token_ids):
 
 
 @pytest.mark.parametrize("case", CASES)
-def test_extract_hidden_states(case: ExtractHiddenStatesCase, sampling_config):
+def test_extract_hidden_states(case: ExtractHiddenStatesCase, sampling_config, monkeypatch):
     """Extract hidden states from the target model and validate the dump."""
+    if case.use_v2_model_runner:
+        monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    else:
+        monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
+
     with tempfile.TemporaryDirectory() as tmpdirname:
         llm_kwargs = dict(
             model=case.model_name,

tests/ut/worker/test_attn_utils_v2.py

@@ -381,3 +381,78 @@ def test_mrv2_builds_shared_dsa_metadata_for_each_execution_mode(
     cache_name = "common_ratio_to_sas_metadata"
     assert calls[0][cache_name] is calls[1][cache_name]
     assert calls[1][cache_name]["first_group"] is True
+
+
+def test_mrv2_allocates_and_reshapes_hidden_state_cache(monkeypatch):
+    """HiddenStateCacheSpec must stay on a single-tensor allocate/reshape path."""
+    from vllm.v1.kv_cache_interface import HiddenStateCacheSpec
+
+    layer_name = "draft.cache_only_layers.36"
+    block_size = 16
+    num_kv_heads = 3
+    head_size = 8
+    num_blocks = 4
+    dtype = torch.bfloat16
+    spec = HiddenStateCacheSpec(
+        block_size=block_size,
+        num_kv_heads=num_kv_heads,
+        head_size=head_size,
+        dtype=dtype,
+    )
+    page_bytes = spec.page_size_bytes
+    tensor_size = num_blocks * page_bytes
+
+    class FakeBackend:
+        @staticmethod
+        def get_kv_cache_shape(num_blocks_, block_size_, num_kv_heads_, head_size_, cache_dtype_str="auto"):
+            return (num_blocks_, block_size_, num_kv_heads_, head_size_)
+
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[KVCacheTensor(size=tensor_size, shared_by=[layer_name])],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                layer_names=[layer_name],
+                kv_cache_spec=spec,
+            )
+        ],
+    )
+
+    monkeypatch.setattr(
+        attn_utils,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(
+            kv_transfer_config=None,
+            model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="qwen3")),
+            quant_config=None,
+            cache_config=SimpleNamespace(cache_dtype="auto"),
+        ),
+    )
+    monkeypatch.setattr(attn_utils, "_is_dsv4_model", lambda _cfg: False)
+    monkeypatch.setattr(attn_utils, "enable_sfa", lambda _cfg: False)
+
+    raw = attn_utils._allocate_kv_cache(kv_cache_config, shared_layers={}, device="cpu")
+    assert isinstance(raw[layer_name], torch.Tensor)
+    assert raw[layer_name].numel() == tensor_size
+
+    attn_groups = [
+        AttentionGroup(
+            backend=FakeBackend,
+            layer_names=[layer_name],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+            metadata_builders=[],
+        )
+    ]
+    reshaped = attn_utils._reshape_kv_cache_v2(
+        attn_groups=attn_groups,
+        kv_cache_raw_tensors=raw,
+        cache_dtype="auto",
+        kernel_block_sizes=[block_size],
+        shared_kv_cache_layers={},
+        kv_cache_config=kv_cache_config,
+    )
+    cache = reshaped[layer_name]
+    assert isinstance(cache, torch.Tensor)
+    assert cache.shape == (num_blocks, block_size, num_kv_heads, head_size)
+    assert cache.dtype == dtype

tests/ut/worker/test_extract_hidden_states_speculator_v2.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+"""Unit tests for Ascend ExtractHiddenStatesSpeculator (Model Runner V2).
+
+Covers Ascend init_speculator dispatch and upstream propose() behavior
+(vLLM PR #49811).
+"""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from typing import Any, cast
+
+import pytest
+import torch
+import vllm.v1.worker.gpu.spec_decode.extract_hidden_states as upstream_spec_module
+
+from vllm_ascend.worker.v2.spec_decode import init_speculator
+from vllm_ascend.worker.v2.spec_decode.extract_hidden_states.speculator import (
+    AscendExtractHiddenStatesSpeculator,
+)
+
+
+class _RecordingModel(torch.nn.Module):
+    def forward(self, *, hidden_states: torch.Tensor) -> None:
+        self.hidden_states = hidden_states.clone()
+
+
+def test_init_requires_greedy_draft_sampling():
+    vllm_config = cast(
+        Any,
+        SimpleNamespace(speculative_config=SimpleNamespace(draft_sample_method="probabilistic")),
+    )
+
+    with pytest.raises(ValueError, match="only supports draft_sample_method='greedy'"):
+        AscendExtractHiddenStatesSpeculator(vllm_config, torch.device("cpu"))
+
+
+def test_init_speculator_dispatches_extract_hidden_states(monkeypatch):
+    vllm_config = cast(
+        Any,
+        SimpleNamespace(speculative_config=SimpleNamespace(method="extract_hidden_states")),
+    )
+    device = torch.device("cpu")
+
+    def fake_speculator(config, target_device):
+        return config, target_device
+
+    monkeypatch.setattr(
+        "vllm_ascend.worker.v2.spec_decode.extract_hidden_states.speculator.AscendExtractHiddenStatesSpeculator",
+        fake_speculator,
+    )
+
+    assert init_speculator(vllm_config, device) == (vllm_config, device)
+
+
+def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch):
+    contexts = []
+
+    def fake_set_forward_context(*args, **kwargs):
+        contexts.append((args, kwargs))
+        return nullcontext()
+
+    monkeypatch.setattr(upstream_spec_module, "set_forward_context", fake_set_forward_context)
+
+    layer_name = "cache_only_layers.2"
+    speculator = object.__new__(AscendExtractHiddenStatesSpeculator)
+    speculator.vllm_config = cast(Any, SimpleNamespace())
+    speculator.num_hidden_states = 2
+    speculator.hidden_states = torch.zeros(4, 2, 3)
+    speculator.draft_attn_layer_names = {layer_name}
+    speculator.model = _RecordingModel()
+
+    input_batch = cast(
+        Any,
+        SimpleNamespace(
+            idx_mapping=torch.tensor([2, 0], dtype=torch.int32),
+            is_padding=torch.zeros(4, dtype=torch.bool),
+        ),
+    )
+    aux_hidden_states = [
+        torch.full((4, 3), 1.0),
+        torch.full((4, 3), 2.0),
+    ]
+    attn_metadata = {layer_name: object(), "target_layer": object()}
+    slot_mappings = {
+        layer_name: torch.arange(4),
+        "target_layer": torch.arange(4),
+    }
+    last_sampled = torch.tensor([[10], [11], [12]], dtype=torch.int64)
+
+    draft_tokens = AscendExtractHiddenStatesSpeculator.propose(
+        speculator,
+        input_batch=input_batch,
+        attn_metadata=attn_metadata,
+        slot_mappings=slot_mappings,
+        last_hidden_states=torch.empty(0),
+        aux_hidden_states=aux_hidden_states,
+        num_sampled=torch.empty(0),
+        num_rejected=torch.empty(0),
+        last_sampled=last_sampled,
+        next_prefill_tokens=torch.empty(0),
+        temperature=torch.empty(0),
+        seeds=torch.empty(0),
+    )
+
+    expected_hidden_states = torch.stack(aux_hidden_states, dim=1)
+    assert torch.equal(speculator.model.hidden_states, expected_hidden_states)
+    assert torch.equal(draft_tokens, torch.tensor([[12], [10]]))
+
+    assert len(contexts) == 1
+    args, kwargs = contexts[0]
+    assert args[0] == {layer_name: attn_metadata[layer_name]}
+    assert kwargs["num_tokens"] == 4
+    assert set(kwargs["slot_mapping"]) == {layer_name}
+    assert torch.equal(kwargs["slot_mapping"][layer_name], slot_mappings[layer_name])
+
+
+def test_propose_requires_aux_hidden_states():
+    speculator = object.__new__(AscendExtractHiddenStatesSpeculator)
+    speculator.num_hidden_states = 2
+    input_batch = cast(Any, SimpleNamespace(idx_mapping=torch.tensor([0], dtype=torch.int32)))
+
+    with pytest.raises(ValueError, match="aux_hidden_states are required"):
+        AscendExtractHiddenStatesSpeculator.propose(
+            speculator,
+            input_batch=input_batch,
+            attn_metadata={},
+            slot_mappings={},
+            last_hidden_states=torch.empty(0),
+            aux_hidden_states=None,
+            num_sampled=torch.empty(0),
+            num_rejected=torch.empty(0),
+            last_sampled=torch.tensor([[10]]),
+            next_prefill_tokens=torch.empty(0),
+            temperature=torch.empty(0),
+            seeds=torch.empty(0),
+        )
+
+
+def test_npu_model_runner_enables_aux_hidden_for_extract_hidden_states():
+    """Document the MRV2 contract: extract_hidden_states needs aux outputs."""
+    speculative_config = SimpleNamespace(method="extract_hidden_states")
+    use_aux_hidden_state_outputs = False
+    if speculative_config.method == "extract_hidden_states":
+        use_aux_hidden_state_outputs = True
+    assert use_aux_hidden_state_outputs is True

vllm_ascend/worker/v2/README.md

@@ -37,3 +37,12 @@ to get specific plans.
     __ of EagleAclGraphManager.
 
     Location: `speculator.AscendEagleSpeculator.init_cudagraph_manager`.
+
+- [x] `extract_hidden_states` (MRV2)
+
+    Why: Upstream vLLM PR #49811 added Model Runner V2 support for
+    `extract_hidden_states`. Ascend thin-wraps upstream
+    `ExtractHiddenStatesSpeculator`, dispatches it via `init_speculator`,
+    forces `use_aux_hidden_state_outputs` when the pinned vLLM omit the
+    method from the GPU allow-list, and keeps HiddenStateCacheSpec on a
+    single-tensor allocate/reshape path.