[Feature][Spec Decode] Support sampling mask replay for MRV2 MTP

docs/training/sampling_mask.md

@@ -63,9 +63,12 @@ The mask is also available via the `/inference/v1/generate` HTTP endpoint:
 | `top_k > 0` | Bounds mask size; pure top-p can produce vocab-sized masks |
 | Model Runner V2 | Required by the async D2H copy pipeline |
 
-The engine rejects unsupported combinations at startup or request time:
+Sampling-mask replay with speculative decoding is supported only for
+fixed-boundary MTP with `rejection_sample_method="standard"`. The engine
+rejects unsupported combinations at startup or request time, including:
 
-- Speculative decoding
+- Adaptive verification, non-MTP draft methods, and synthetic or block
+  verification
 - Diffusion models
 - Custom logits processors (engine-level `--logits-processors`)
 
@@ -76,8 +79,10 @@ The engine rejects unsupported combinations at startup or request time:
    logits to `-inf`.
 2. After sampling, `torch.isfinite(processed_logits)` identifies the surviving
    token IDs — this is the sampling mask.
-3. The mask is transferred GPU → CPU asynchronously alongside sampled tokens.
-4. On request completion, per-step masks are merged and converted to
+3. For MTP, the engine returns aligned processed-target sampling support for
+   every actually emitted accepted, recovered, or bonus token.
+4. The mask is transferred GPU → CPU asynchronously alongside sampled tokens.
+5. On request completion, per-step masks are merged and converted to
    `list[list[int]]` for the response.
 
 ## RL training usage

tests/test_config.py

@@ -56,6 +56,103 @@ def _write_json(path: Path, value: object) -> None:
     path.write_text(json.dumps(value), encoding="utf-8")
 
 
+def _sampling_replay_config(
+    *,
+    return_sampling_mask: bool = True,
+    use_v2_model_runner: bool = True,
+    speculative_method: str | None = None,
+    rejection_sample_method: str = "standard",
+    adaptive: bool = False,
+    is_diffusion: bool = False,
+    logits_processors: list[str] | None = None,
+    logprobs_mode: str = "processed_logprobs",
+):
+    speculative_config = None
+    if speculative_method is not None:
+        speculative_config = SimpleNamespace(
+            method=speculative_method,
+            enable_adaptive_verification=adaptive,
+            rejection_sample_method=rejection_sample_method,
+        )
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            return_sampling_mask=return_sampling_mask,
+            is_diffusion=is_diffusion,
+            logits_processors=logits_processors or [],
+            logprobs_mode=logprobs_mode,
+        ),
+        use_v2_model_runner=use_v2_model_runner,
+        speculative_config=speculative_config,
+    )
+
+
+@pytest.mark.parametrize(
+    ("config", "message"),
+    [
+        (_sampling_replay_config(), None),
+        (_sampling_replay_config(speculative_method="mtp"), None),
+        (
+            _sampling_replay_config(
+                return_sampling_mask=False, speculative_method="dflash"
+            ),
+            None,
+        ),
+        (
+            _sampling_replay_config(speculative_method="mtp", adaptive=True),
+            "requires fixed verification boundaries",
+        ),
+        (
+            _sampling_replay_config(speculative_method="eagle"),
+            "currently supports only the MTP speculative method",
+        ),
+        (
+            _sampling_replay_config(speculative_method="dflash"),
+            "currently supports only the MTP speculative method",
+        ),
+        (
+            _sampling_replay_config(speculative_method="dspark"),
+            "currently supports only the MTP speculative method",
+        ),
+        (
+            _sampling_replay_config(
+                speculative_method="mtp",
+                rejection_sample_method="synthetic",
+            ),
+            "only rejection_sample_method='standard'",
+        ),
+        (
+            _sampling_replay_config(
+                speculative_method="mtp",
+                rejection_sample_method="block",
+            ),
+            "only rejection_sample_method='standard'",
+        ),
+        (
+            _sampling_replay_config(is_diffusion=True),
+            "does not support diffusion models",
+        ),
+        (
+            _sampling_replay_config(logits_processors=["custom"]),
+            "does not support custom logits processors",
+        ),
+        (
+            _sampling_replay_config(logprobs_mode="raw_logprobs"),
+            "requires logprobs_mode='processed_logprobs'",
+        ),
+        (
+            _sampling_replay_config(use_v2_model_runner=False),
+            "requires Model Runner V2",
+        ),
+    ],
+)
+def test_sampling_replay_config(config, message):
+    if message is None:
+        VllmConfig._verify_sampling_replay_config(config)
+    else:
+        with pytest.raises(ValueError, match=message):
+            VllmConfig._verify_sampling_replay_config(config)
+
+
 def test_kda_recoverssm_derivation_is_revalidated():
     config = SimpleNamespace(
         cache_config=SimpleNamespace(

tests/v1/core/test_scheduler.py

@@ -715,6 +715,49 @@ def test_update_from_output_routes_sampling_masks_by_request():
     assert all(out.new_sampling_mask.offsets is None for out in outputs)
 
 
+def test_update_from_output_routes_multi_position_sampling_masks():
+    scheduler = create_scheduler()
+    scheduler.return_sampling_mask = True
+    requests = create_requests(num_requests=2, max_tokens=10)
+    for req in requests:
+        req.num_computed_tokens = req.num_tokens
+        scheduler.requests[req.request_id] = req
+        scheduler.running.append(req)
+        req.status = RequestStatus.RUNNING
+
+    scheduler_output = SchedulerOutput(
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        num_scheduled_tokens={req.request_id: 1 for req in requests},
+        total_num_scheduled_tokens=2,
+        scheduled_encoder_inputs={},
+        scheduled_spec_decode_tokens={},
+        num_common_prefix_blocks=[],
+        finished_req_ids=set(),
+        free_encoder_mm_hashes=[],
+    )
+    model_output = ModelRunnerOutput(
+        req_ids=[req.request_id for req in requests],
+        req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
+        sampled_token_ids=[[1, 2], [3, 4, 5]],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+        sampling_masks=SamplingMaskLists(
+            token_ids=np.array([1, 6, 2, 3, 7, 8, 4, 5, 9], dtype=np.int32),
+            offsets=np.array([0, 2, 3, 6, 7, 9]),
+            cu_num_generated_tokens=[0, 2, 5],
+        ),
+    )
+
+    outputs = scheduler.update_from_output(scheduler_output, model_output)[0].outputs
+
+    assert [out.new_sampling_mask.to_nested_list() for out in outputs] == [
+        [[1, 6], [2]],
+        [[3, 7, 8], [4], [5, 9]],
+    ]
+
+
 def test_stop_via_update_from_output():
     """Test stopping behavior through update_from_output"""
     scheduler = create_scheduler(num_speculative_tokens=1)

tests/v1/engine/test_output_processor.py

@@ -5,6 +5,7 @@
 import time
 from unittest.mock import MagicMock
 
+import numpy as np
 import pytest
 
 from tests.v1.engine.utils import (
@@ -33,6 +34,7 @@
     RequestState,
 )
 from vllm.v1.metrics.stats import IterationStats, SchedulerStats
+from vllm.v1.outputs import SamplingMaskLists
 
 
 @pytest.mark.parametrize("flat_logprobs", [False, True])
@@ -60,6 +62,31 @@ def test_delta_output_without_new_tokens_returns_empty_logprobs(
     assert len(output.logprobs) == 0
 
 
+def test_completion_output_preserves_each_sampling_mask_position() -> None:
+    state = RequestState.__new__(RequestState)
+    state.detokenizer = MagicMock()
+    state.detokenizer.get_next_output_text.return_value = ""
+    state.logprobs_processor = MagicMock()
+    state.logprobs_processor.logprobs = None
+    state.logprobs_processor.cumulative_logprob = None
+    state.output_kind = RequestOutputKind.DELTA
+    state.request_index = 0
+    state.sampling_mask_chunks = [
+        SamplingMaskLists(
+            token_ids=np.array([10, 11, 20]),
+            offsets=np.array([0, 2, 3]),
+        ),
+        SamplingMaskLists(token_ids=np.array([30, 31, 32])),
+    ]
+    state.routed_experts_chunks = []
+    state.spec_decode_metrics = None
+
+    output = state._new_completion_output([1, 2, 3], FinishReason.LENGTH, None)
+
+    assert output.sampling_mask is not None
+    assert output.sampling_mask.token_ids == [[10, 11], [20], [30, 31, 32]]
+
+
 def _ref_convert_id_to_token(
     tokenizer: TokenizerLike,
     token_id: int,

tests/v1/test_outputs.py

@@ -13,6 +13,7 @@
     LogprobsLists,
     LogprobsTensors,
     ModelRunnerOutput,
+    SamplingMaskLists,
 )
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
 from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors
@@ -63,8 +64,6 @@ def test_logprobs_tensors_tolists_with_tensor_boundaries():
 
 
 def test_sampling_mask_lists_to_nested_list():
-    from vllm.v1.outputs import SamplingMaskLists
-
     mask = SamplingMaskLists(
         token_ids=np.array([10, 11, 12, 20, 21]),
         offsets=np.array([0, 3, 5]),
@@ -76,6 +75,22 @@ def test_sampling_mask_lists_to_nested_list():
     assert SamplingMaskLists(np.array([7, 9])).to_nested_list() == [[7, 9]]
 
 
+def test_sampling_mask_lists_slices_multiple_positions_by_request():
+    masks = SamplingMaskLists(
+        token_ids=np.array([10, 11, 20, 30, 31, 32, 40]),
+        offsets=np.array([0, 2, 3, 6, 7]),
+        cu_num_generated_tokens=[0, 1, 3, 4],
+    )
+
+    single = masks.slice_request(0, 1)
+    multi = masks.slice_request(1, 2)
+
+    assert single.to_nested_list() == [[10, 11]]
+    assert single.offsets is None
+    assert multi.to_nested_list() == [[20], [30, 31, 32]]
+    assert multi.cu_num_generated_tokens is None
+
+
 @pytest.mark.parametrize("max_num_kept", [512, 20_001])
 @pytest.mark.skipif(
     current_platform.is_xpu(),
@@ -102,7 +117,10 @@ def test_sampling_mask_tensors_match_finite_support(max_num_kept):
     expected[1] = expected[7] = []
 
     tensors = SamplingMaskTensors.from_logits(
-        logits.to(DEVICE_TYPE), num_sampled_tokens.to(DEVICE_TYPE), max_num_kept
+        logits.to(DEVICE_TYPE),
+        torch.arange(len(sizes) + 1, device=DEVICE_TYPE, dtype=torch.int32),
+        num_sampled_tokens.to(DEVICE_TYPE),
+        max_num_kept,
     )
 
     assert tensors.token_ids.shape[1] == min(max_num_kept, MAX_COMPACT_SUPPORT)
@@ -127,6 +145,7 @@ def test_sampling_mask_matches_processed_top_k_top_p_support():
 
     tensors = SamplingMaskTensors.from_logits(
         processed_logits,
+        cu_num_logits=torch.tensor([0, 1], dtype=torch.int32, device=DEVICE_TYPE),
         num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE),
         max_num_kept=3,
     )
@@ -135,6 +154,50 @@ def test_sampling_mask_matches_processed_top_k_top_p_support():
     assert result.to_nested_list() == [expected_token_ids]
 
 
+def test_sampling_mask_tensors_uses_request_boundaries():
+    logits = torch.full((6, 8), -float("inf"), device=DEVICE_TYPE)
+    logits[0, [1, 3]] = 0
+    logits[2, [2, 4, 6]] = 0
+    logits[5, [0, 7]] = 0
+    counts = torch.tensor([1, 0, 1], device=DEVICE_TYPE, dtype=torch.int32)
+
+    tensors = SamplingMaskTensors.from_logits(
+        logits,
+        torch.tensor([0, 2, 5, 6], device=DEVICE_TYPE, dtype=torch.int32),
+        counts,
+        max_num_kept=3,
+    )
+
+    assert tensors.tolists().to_nested_list() == [[1, 3], [], [0, 7]]
+
+
+def test_sampling_mask_tensors_multirow_request_layout():
+    rows_per_request = 4
+    counts = np.array([1, 2, 4, 0], dtype=np.int32)
+    logits = torch.full((16, 16), -float("inf"), device=DEVICE_TYPE)
+    expected = []
+    for row in [0, 4, 5, 8, 9, 10, 11]:
+        kept = torch.tensor(
+            [row % 16, (row + 3) % 16, (row + 7) % 16], device=DEVICE_TYPE
+        ).unique(sorted=True)
+        logits[row, kept] = 0
+        expected.append(kept.tolist())
+
+    tensors = SamplingMaskTensors.from_logits(
+        logits,
+        torch.arange(0, 17, 4, device=DEVICE_TYPE, dtype=torch.int32),
+        torch.from_numpy(counts).to(DEVICE_TYPE),
+        max_num_kept=2,
+        rows_per_request=rows_per_request,
+    )
+    result = tensors.to_cpu_nonblocking().tolists(counts)
+
+    assert tensors.token_ids.shape == (16, 2)
+    assert tensors.rows_per_request == rows_per_request
+    assert result.to_nested_list() == expected
+    assert result.cu_num_generated_tokens == [0, 1, 3, 7, 7]
+
+
 def test_sampling_mask_preserves_top_k_boundary_ties():
     """When the kept support is wider than `max_num_kept` (e.g. a top-k
     boundary tie keeps more than k logits), the mask must fall back to the
@@ -147,6 +210,7 @@ def test_sampling_mask_preserves_top_k_boundary_ties():
 
     tensors = SamplingMaskTensors.from_logits(
         processed_logits,
+        cu_num_logits=torch.tensor([0, 1], dtype=torch.int32, device=DEVICE_TYPE),
         num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE),
         max_num_kept=3,
     )