@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70be6aae19

[Feature][Spec Decode] Support sampling mask replay for MRV2 MTP (#59359)

Signed-off-by: Cuiping Cheng <chengcuipingswu@163.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Cuiping Cheng <chengcuipingswu@163.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

vx120 committed Oct 1, 2026 at 13:51 UTC 70be6aae19587a251987c405b5bdf2d540da5991

13 files changed +583 -68

docs/training/sampling_mask.md

+4 -1

@@ -63,9 +63,12 @@ The mask is also available via the `/inference/v1/generate` HTTP endpoint:

63 | `top_k > 0` | Bounds mask size; pure top-p can produce vocab-sized masks |

64 | Model Runner V2 | Required by the async D2H copy pipeline |

66 +Speculative decoding is supported; each emitted token gets the support of the

67 +target distribution at its position.

68 +

69 The engine rejects unsupported combinations at startup or request time:

68 -- Speculative decoding

71 +- Speculative decoding with adaptive verification

72 - Diffusion models

73 - Custom logits processors (engine-level `--logits-processors`)

tests/test_config.py

+91

@@ -147,6 +147,97 @@ def test_rocm_mm_prefix_lm_disables_chunked_mm_input(

147 assert config.scheduler_config.disable_chunked_mm_input is expected

148

149

150 +def _sampling_replay_config(

151 + *,

152 + return_sampling_mask: bool = True,

153 + use_v2_model_runner: bool = True,

154 + speculative_method: str | None = None,

155 + rejection_sample_method: str = "standard",

156 + adaptive: bool = False,

157 + is_diffusion: bool = False,

158 + logits_processors: list[str] | None = None,

159 + logprobs_mode: str = "processed_logprobs",

160 +):

161 + speculative_config = None

162 + if speculative_method is not None:

163 + speculative_config = SimpleNamespace(

164 + method=speculative_method,

165 + enable_adaptive_verification=adaptive,

166 + rejection_sample_method=rejection_sample_method,

167 + )

168 + return SimpleNamespace(

169 + model_config=SimpleNamespace(

170 + return_sampling_mask=return_sampling_mask,

171 + is_diffusion=is_diffusion,

172 + logits_processors=logits_processors or [],

173 + logprobs_mode=logprobs_mode,

174 + ),

175 + use_v2_model_runner=use_v2_model_runner,

176 + speculative_config=speculative_config,

177 + )

178 +

179 +

180 +@pytest.mark.parametrize(

181 + ("config", "message"),

182 + [

183 + (_sampling_replay_config(), None),

184 + *(

185 + (_sampling_replay_config(speculative_method=method), None)

186 + for method in (

187 + "mtp",

188 + "eagle",

189 + "eagle3",

190 + "dflash",

191 + "dspark",

192 + "draft_model",

193 + )

194 + ),

195 + *(

196 + (

197 + _sampling_replay_config(

198 + speculative_method="mtp",

199 + rejection_sample_method=rejection_sample_method,

200 + ),

201 + None,

202 + )

203 + for rejection_sample_method in ("standard", "block", "synthetic")

204 + ),

205 + (

206 + _sampling_replay_config(

207 + return_sampling_mask=False, speculative_method="dspark", adaptive=True

208 + ),

209 + None,

210 + ),

211 + (

212 + _sampling_replay_config(speculative_method="dspark", adaptive=True),

213 + "requires fixed verification boundaries",

214 + ),

215 + (

216 + _sampling_replay_config(is_diffusion=True),

217 + "does not support diffusion models",

218 + ),

219 + (

220 + _sampling_replay_config(logits_processors=["custom"]),

221 + "does not support custom logits processors",

222 + ),

223 + (

224 + _sampling_replay_config(logprobs_mode="raw_logprobs"),

225 + "requires logprobs_mode='processed_logprobs'",

226 + ),

227 + (

228 + _sampling_replay_config(use_v2_model_runner=False),

229 + "requires Model Runner V2",

230 + ),

231 + ],

232 +)

233 +def test_sampling_replay_config(config, message):

234 + if message is None:

235 + VllmConfig._verify_sampling_replay_config(config)

236 + else:

237 + with pytest.raises(ValueError, match=message):

238 + VllmConfig._verify_sampling_replay_config(config)

239 +

240 +

241 def test_kda_recoverssm_derivation_is_revalidated():

242 config = SimpleNamespace(

243 cache_config=SimpleNamespace(

tests/v1/core/test_scheduler.py

+43

@@ -908,6 +908,49 @@ def test_update_from_output_routes_sampling_masks_by_request():

908 assert all(out.new_sampling_mask.offsets is None for out in outputs)

909

910

911 +def test_update_from_output_routes_multi_position_sampling_masks():

912 + scheduler = create_scheduler()

913 + scheduler.return_sampling_mask = True

914 + requests = create_requests(num_requests=2, max_tokens=10)

915 + for req in requests:

916 + req.num_computed_tokens = req.num_tokens

917 + scheduler.requests[req.request_id] = req

918 + scheduler.running.append(req)

919 + req.status = RequestStatus.RUNNING

920 +

921 + scheduler_output = SchedulerOutput(

922 + scheduled_new_reqs=[],

923 + scheduled_cached_reqs=CachedRequestData.make_empty(),

924 + num_scheduled_tokens={req.request_id: 1 for req in requests},

925 + total_num_scheduled_tokens=2,

926 + scheduled_encoder_inputs={},

927 + scheduled_spec_decode_tokens={},

928 + num_common_prefix_blocks=[],

929 + finished_req_ids=set(),

930 + free_encoder_mm_hashes=[],

931 + )

932 + model_output = ModelRunnerOutput(

933 + req_ids=[req.request_id for req in requests],

934 + req_id_to_index={req.request_id: i for i, req in enumerate(requests)},

935 + sampled_token_ids=[[1, 2], [3, 4, 5]],

936 + logprobs=None,

937 + prompt_logprobs_dict={},

938 + pooler_output=[],

939 + sampling_masks=SamplingMaskLists(

940 + token_ids=np.array([1, 6, 2, 3, 7, 8, 4, 5, 9], dtype=np.int32),

941 + offsets=np.array([0, 2, 3, 6, 7, 9]),

942 + cu_num_generated_tokens=[0, 2, 5],

943 + ),

944 + )

945 +

946 + outputs = scheduler.update_from_output(scheduler_output, model_output)[0].outputs

947 +

948 + assert [out.new_sampling_mask.to_nested_list() for out in outputs] == [

949 + [[1, 6], [2]],

950 + [[3, 7, 8], [4], [5, 9]],

951 + ]

952 +

953 +

954 def test_stop_via_update_from_output():

955 """Test stopping behavior through update_from_output."""

956 scheduler = create_scheduler(num_speculative_tokens=1)

tests/v1/engine/test_output_processor.py

+26

@@ -35,6 +35,7 @@ from vllm.v1.engine.output_processor import (

35 RequestState,

36 )

37 from vllm.v1.metrics.stats import IterationStats, PrefillStats, SchedulerStats

38 +from vllm.v1.outputs import SamplingMaskLists

41 @pytest.mark.parametrize("flat_logprobs", [False, True])

@@ -62,6 +63,31 @@ def test_delta_output_without_new_tokens_returns_empty_logprobs(

63 assert len(output.logprobs) == 0

66 +def test_completion_output_preserves_each_sampling_mask_position() -> None:

67 + state = RequestState.__new__(RequestState)

68 + state.detokenizer = MagicMock()

69 + state.detokenizer.get_next_output_text.return_value = ""

70 + state.logprobs_processor = MagicMock()

71 + state.logprobs_processor.logprobs = None

72 + state.logprobs_processor.cumulative_logprob = None

73 + state.output_kind = RequestOutputKind.DELTA

74 + state.request_index = 0

75 + state.sampling_mask_chunks = [

76 + SamplingMaskLists(

77 + token_ids=np.array([10, 11, 20]),

78 + offsets=np.array([0, 2, 3]),

79 + ),

80 + SamplingMaskLists(token_ids=np.array([30, 31, 32])),

81 + ]

82 + state.routed_experts_chunks = []

83 + state.spec_decode_metrics = None

84 +

85 + output = state._new_completion_output([1, 2, 3], FinishReason.LENGTH, None)

86 +

87 + assert output.sampling_mask is not None

88 + assert output.sampling_mask.token_ids == [[10, 11], [20], [30, 31, 32]]

89 +

90 +

91 def _ref_convert_id_to_token(

92 tokenizer: TokenizerLike,

93 token_id: int,

tests/v1/test_outputs.py

+67 -3

@@ -13,6 +13,7 @@ from vllm.v1.outputs import (

13 LogprobsLists,

14 LogprobsTensors,

15 ModelRunnerOutput,

16 + SamplingMaskLists,

17 )

18 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p

19 from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors

@@ -63,8 +64,6 @@ def test_logprobs_tensors_tolists_with_tensor_boundaries():

66 def test_sampling_mask_lists_to_nested_list():

66 - from vllm.v1.outputs import SamplingMaskLists

67 -

67 mask = SamplingMaskLists(

68 token_ids=np.array([10, 11, 12, 20, 21]),

69 offsets=np.array([0, 3, 5]),

@@ -76,6 +75,22 @@ def test_sampling_mask_lists_to_nested_list():

75 assert SamplingMaskLists(np.array([7, 9])).to_nested_list() == [[7, 9]]

78 +def test_sampling_mask_lists_slices_multiple_positions_by_request():

79 + masks = SamplingMaskLists(

80 + token_ids=np.array([10, 11, 20, 30, 31, 32, 40]),

81 + offsets=np.array([0, 2, 3, 6, 7]),

82 + cu_num_generated_tokens=[0, 1, 3, 4],

83 + )

84 +

85 + single = masks.slice_request(0, 1)

86 + multi = masks.slice_request(1, 2)

87 +

88 + assert single.to_nested_list() == [[10, 11]]

89 + assert single.offsets is None

90 + assert multi.to_nested_list() == [[20], [30, 31, 32]]

91 + assert multi.cu_num_generated_tokens is None

92 +

93 +

94 @pytest.mark.parametrize("max_num_kept", [512, 20_001])

95 @pytest.mark.skipif(

96 current_platform.is_xpu(),

@@ -102,7 +117,10 @@ def test_sampling_mask_tensors_match_finite_support(max_num_kept):

117 expected[1] = expected[7] = []

118

119 tensors = SamplingMaskTensors.from_logits(

105 - logits.to(DEVICE_TYPE), num_sampled_tokens.to(DEVICE_TYPE), max_num_kept

120 + logits.to(DEVICE_TYPE),

121 + torch.arange(len(sizes) + 1, device=DEVICE_TYPE, dtype=torch.int32),

122 + num_sampled_tokens.to(DEVICE_TYPE),

123 + max_num_kept,

124 )

125

126 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():

145

146 tensors = SamplingMaskTensors.from_logits(

147 processed_logits,

148 + cu_num_logits=torch.tensor([0, 1], dtype=torch.int32, device=DEVICE_TYPE),

149 num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE),

150 max_num_kept=3,

151 )

@@ -135,6 +154,50 @@ def test_sampling_mask_matches_processed_top_k_top_p_support():

154 assert result.to_nested_list() == [expected_token_ids]

155

156

157 +def test_sampling_mask_tensors_uses_request_boundaries():

158 + logits = torch.full((6, 8), -float("inf"), device=DEVICE_TYPE)

159 + logits[0, [1, 3]] = 0

160 + logits[2, [2, 4, 6]] = 0

161 + logits[5, [0, 7]] = 0

162 + counts = torch.tensor([1, 0, 1], device=DEVICE_TYPE, dtype=torch.int32)

163 +

164 + tensors = SamplingMaskTensors.from_logits(

165 + logits,

166 + torch.tensor([0, 2, 5, 6], device=DEVICE_TYPE, dtype=torch.int32),

167 + counts,

168 + max_num_kept=3,

169 + )

170 +

171 + assert tensors.tolists().to_nested_list() == [[1, 3], [], [0, 7]]

172 +

173 +

174 +def test_sampling_mask_tensors_multirow_request_layout():

175 + rows_per_request = 4

176 + counts = np.array([1, 2, 4, 0], dtype=np.int32)

177 + logits = torch.full((16, 16), -float("inf"), device=DEVICE_TYPE)

178 + expected = []

179 + for row in [0, 4, 5, 8, 9, 10, 11]:

180 + kept = torch.tensor(

181 + [row % 16, (row + 3) % 16, (row + 7) % 16], device=DEVICE_TYPE

182 + ).unique(sorted=True)

183 + logits[row, kept] = 0

184 + expected.append(kept.tolist())

185 +

186 + tensors = SamplingMaskTensors.from_logits(

187 + logits,

188 + torch.arange(0, 17, 4, device=DEVICE_TYPE, dtype=torch.int32),

189 + torch.from_numpy(counts).to(DEVICE_TYPE),

190 + max_num_kept=2,

191 + rows_per_request=rows_per_request,

192 + )

193 + result = tensors.to_cpu_nonblocking().tolists(counts)

194 +

195 + assert tensors.token_ids.shape == (16, 2)

196 + assert tensors.rows_per_request == rows_per_request

197 + assert result.to_nested_list() == expected

198 + assert result.cu_num_generated_tokens == [0, 1, 3, 7, 7]

199 +

200 +

201 def test_sampling_mask_preserves_top_k_boundary_ties():

202 """When the kept support is wider than `max_num_kept` (e.g. a top-k

203 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():

210

211 tensors = SamplingMaskTensors.from_logits(

212 processed_logits,

213 + cu_num_logits=torch.tensor([0, 1], dtype=torch.int32, device=DEVICE_TYPE),

214 num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE),

215 max_num_kept=3,

216 )

tests/v1/worker/test_gpu_rejection_sampler_chunking.py

+207 -29

@@ -10,12 +10,53 @@ import torch

11 from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode

12 from vllm.platforms import current_platform

13 +from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors

14 from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (

15 RejectionSampler,

16 _iter_request_chunks,

17 )

20 +def _make_input_batch(

21 + cu_num_logits_np: np.ndarray,

22 + idx_mapping_np: np.ndarray,

23 + device: torch.device,

24 +) -> SimpleNamespace:

25 + """Minimal InputBatch fields read by RejectionSampler._verify_in_chunks."""

26 + num_reqs = len(idx_mapping_np)

27 + num_logits_per_req = np.diff(cu_num_logits_np)

28 + return SimpleNamespace(

29 + num_reqs=num_reqs,

30 + cu_num_logits_np=cu_num_logits_np,

31 + cu_num_logits=torch.from_numpy(cu_num_logits_np).to(device),

32 + idx_mapping_np=idx_mapping_np,

33 + idx_mapping=torch.from_numpy(idx_mapping_np).to(device),

34 + expanded_idx_mapping=torch.from_numpy(

35 + np.repeat(idx_mapping_np, num_logits_per_req)

36 + ).to(device),

37 + expanded_local_pos=torch.from_numpy(

38 + np.concatenate(

39 + [np.arange(count, dtype=np.int32) for count in num_logits_per_req]

40 + )

41 + ).to(device),

42 + seq_lens_cpu_upper_bound=torch.full((num_reqs,), 64, dtype=torch.int32),

43 + )

44 +

45 +

46 +def _make_rejection_sampler(

47 + sampler: SimpleNamespace, num_speculative_steps: int

48 +) -> RejectionSampler:

49 + """RejectionSampler for standard, fixed-boundary, unwatermarked verification."""

50 + rejection_sampler = object.__new__(RejectionSampler)

51 + rejection_sampler.sampler = sampler

52 + rejection_sampler.num_speculative_steps = num_speculative_steps

53 + rejection_sampler.enable_adaptive_verification = False

54 + rejection_sampler.synthetic_conditional_rates = None

55 + rejection_sampler.use_block_verification = False

56 + rejection_sampler.watermark_key = None

57 + return rejection_sampler

58 +

59 +

60 def test_iter_request_chunks_preserves_request_boundaries():

61 cu_num_logits = np.array([0, 3, 4, 11, 13], dtype=np.int32)

@@ -33,28 +74,11 @@ def test_chunked_scores_match_full_batch(logprobs_mode: str):

74 cu_num_logits_np = np.array([0, 3, 4, 8, 10], dtype=np.int32)

75 num_logits_per_req = np.diff(cu_num_logits_np)

76 idx_mapping_np = np.array([7, 2, 9, 1], dtype=np.int32)

36 - input_batch = SimpleNamespace(

37 - num_reqs=4,

38 - cu_num_logits_np=cu_num_logits_np,

39 - cu_num_logits=torch.from_numpy(cu_num_logits_np).to(device),

40 - idx_mapping_np=idx_mapping_np,

41 - idx_mapping=torch.from_numpy(idx_mapping_np).to(device),

42 - expanded_idx_mapping=torch.from_numpy(

43 - np.repeat(idx_mapping_np, num_logits_per_req)

44 - ).to(device),

45 - expanded_local_pos=torch.from_numpy(

46 - np.concatenate(

47 - [np.arange(count, dtype=np.int32) for count in num_logits_per_req]

48 - )

49 - ).to(device),

50 - seq_lens_cpu_upper_bound=torch.from_numpy(

51 - np.array([10, 20, 30, 40], dtype=np.int32)

52 - ),

77 + input_batch = _make_input_batch(cu_num_logits_np, idx_mapping_np, device)

78 + rejection_sampler = _make_rejection_sampler(

79 + SimpleNamespace(logprobs_mode=logprobs_mode, return_sampling_mask=False),

80 + num_speculative_steps=3,

81 )

54 - rejection_sampler = object.__new__(RejectionSampler)

55 - rejection_sampler.sampler = SimpleNamespace(logprobs_mode=logprobs_mode)

56 - rejection_sampler.num_speculative_steps = 3

57 - rejection_sampler.enable_adaptive_verification = False

83 def fake_verify(

84 self,

@@ -75,14 +99,16 @@ def test_chunked_scores_match_full_batch(logprobs_mode: str):

99 rejection_sampler._verify = MethodType(fake_verify, rejection_sampler)

100 logits = torch.arange(170, dtype=torch.float32, device=device).view(10, 17)

101

78 - sampled, num_sampled, chunked_logprobs = rejection_sampler._verify_in_chunks(

79 - logits,

80 - input_batch,

81 - draft_logits=None,

82 - draft_sampled=torch.arange(10, device=device),

83 - pos=torch.arange(10, device=device),

84 - max_chunk_logits=5,

85 - max_num_logprobs=2,

102 + sampled, num_sampled, chunked_logprobs, sampling_mask_tensors = (

103 + rejection_sampler._verify_in_chunks(

104 + logits,

105 + input_batch,

106 + draft_logits=None,

107 + draft_sampled=torch.arange(10, device=device),

108 + pos=torch.arange(10, device=device),

109 + max_chunk_logits=5,

110 + max_num_logprobs=2,

111 + )

112 )

113 score_logits = logits + 1 if logprobs_mode in PROCESSED_LOGPROBS_MODES else logits

114 full_logprobs = rejection_sampler._get_logprobs_tensors(

@@ -111,3 +137,155 @@ def test_chunked_scores_match_full_batch(logprobs_mode: str):

137 chunked_logprobs.cu_num_generated_tokens

138 == full_logprobs.cu_num_generated_tokens

139 )

140 + assert sampling_mask_tensors is None

141 +

142 +

143 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")

144 +def test_replay_on_off_preserves_rejection_sampling_and_rng(monkeypatch):

145 + device = torch.device("cuda")

146 + num_reqs = 3

147 + num_speculative_steps = 3

148 + rows_per_request = num_speculative_steps + 1

149 + vocab_size = 8

150 + cu_num_logits_np = np.arange(0, 13, rows_per_request, dtype=np.int32)

151 + idx_mapping_np = np.arange(num_reqs, dtype=np.int32)

152 + input_batch = _make_input_batch(cu_num_logits_np, idx_mapping_np, device)

153 + processed_logits = torch.zeros((12, vocab_size), device=device)

154 + processed_logits[0, 0] = -float("inf")

155 + processed_logits[rows_per_request + 1, 1] = -float("inf")

156 + draft_logits = torch.zeros(

157 + (num_reqs, num_speculative_steps, vocab_size), device=device

158 + )

159 + draft_sampled = torch.zeros(

160 + (num_reqs, rows_per_request), dtype=torch.int64, device=device

161 + )

162 + draft_sampled[:, 1:] = torch.tensor([0, 1, 2], device=device)

163 + draft_sampled = draft_sampled.flatten()

164 + sampler = SimpleNamespace(

165 + logprobs_mode="processed_logprobs",

166 + return_sampling_mask=False,

167 + sampling_states=SimpleNamespace(

168 + top_k=SimpleNamespace(np=np.full(num_reqs, vocab_size)),

169 + temperature=SimpleNamespace(gpu=torch.ones(num_reqs, device=device)),

170 + seeds=SimpleNamespace(

171 + gpu=torch.arange(num_reqs, dtype=torch.int64, device=device)

172 + ),

173 + ),

174 + use_fp64_gumbel=False,

175 + apply_sampling_params=lambda logits, *_: logits,

176 + )

177 + rejection_sampler = _make_rejection_sampler(sampler, num_speculative_steps)

178 + pos = torch.arange(12, dtype=torch.int32, device=device)

179 +

180 + pack_calls: list[tuple[list[int], int, int]] = []

181 + pack = SamplingMaskTensors.from_logits.__func__

182 +

183 + def track_pack(

184 + cls,

185 + logits,

186 + cu_num_logits,

187 + num_sampled_tokens,

188 + max_num_kept,

189 + rows_per_request=1,

190 + ):

191 + pack_calls.append(

192 + (cu_num_logits.cpu().tolist(), max_num_kept, rows_per_request)

193 + )

194 + return pack(

195 + cls,

196 + logits,

197 + cu_num_logits,

198 + num_sampled_tokens,

199 + max_num_kept,

200 + rows_per_request,

201 + )

202 +

203 + monkeypatch.setattr(

204 + SamplingMaskTensors,

205 + "from_logits",

206 + classmethod(track_pack),

207 + )

208 +

209 + rng_before = torch.cuda.get_rng_state()

210 + off_sampled, off_counts, off_logprobs, off_masks = (

211 + rejection_sampler._verify_in_chunks(

212 + processed_logits,

213 + input_batch,

214 + draft_logits,

215 + draft_sampled,

216 + pos,

217 + max_chunk_logits=8,

218 + max_num_logprobs=0,

219 + )

220 + )

221 + rng_after_off = torch.cuda.get_rng_state()

222 + assert pack_calls == []

223 +

224 + sampler.return_sampling_mask = True

225 + on_sampled, on_counts, on_logprobs, on_masks = rejection_sampler._verify_in_chunks(

226 + processed_logits,

227 + input_batch,

228 + draft_logits,

229 + draft_sampled,

230 + pos,

231 + max_chunk_logits=8,

232 + max_num_logprobs=0,

233 + )

234 + rng_after_on = torch.cuda.get_rng_state()

235 +

236 + assert pack_calls == [

237 + ([0, 4, 8], vocab_size, rows_per_request),

238 + ([0, 4], vocab_size, rows_per_request),

239 + ]

240 + assert off_masks is None and on_masks is not None

241 + assert on_masks.rows_per_request == rows_per_request

242 + assert on_masks.token_ids.shape == (num_reqs * rows_per_request, vocab_size)

243 + assert on_masks.packed_mask.shape[0] == num_reqs * rows_per_request

244 + assert off_counts.tolist() == [1, 2, 4]

245 + assert torch.equal(off_counts, on_counts)

246 + counts = off_counts.cpu().tolist()

247 + for req_idx, count in enumerate(counts):

248 + assert torch.equal(off_sampled[req_idx, :count], on_sampled[req_idx, :count])

249 +

250 + mapped_rows = [

251 + int(input_batch.cu_num_logits_np[req_idx]) + slot_idx

252 + for req_idx, count in enumerate(counts)

253 + for slot_idx in range(count)

254 + ]

255 + assert off_logprobs is not None and on_logprobs is not None

256 + assert torch.equal(

257 + off_logprobs.logprob_token_ids[mapped_rows],

258 + on_logprobs.logprob_token_ids[mapped_rows],

259 + )

260 + assert torch.equal(

261 + off_logprobs.logprobs[mapped_rows], on_logprobs.logprobs[mapped_rows]

262 + )

263 + assert torch.equal(

264 + off_logprobs.selected_token_ranks[mapped_rows],

265 + on_logprobs.selected_token_ranks[mapped_rows],

266 + )

267 + assert torch.equal(rng_before, rng_after_off)

268 + assert torch.equal(rng_after_off, rng_after_on)

269 +

270 + masks = (

271 + on_masks.to_cpu_nonblocking().tolists(on_counts.cpu().numpy()).to_nested_list()

272 + )

273 + expected_masks = [

274 + torch.isfinite(processed_logits[row]).nonzero().flatten().tolist()

275 + for row in mapped_rows

276 + ]

277 + emitted = [

278 + int(on_sampled[req_idx, slot_idx])

279 + for req_idx, count in enumerate(on_counts.cpu().tolist())

280 + for slot_idx in range(count)

281 + ]

282 + expected_logprobs = torch.stack(

283 + [

284 + torch.log_softmax(processed_logits[row], dim=-1)[token_id]

285 + for row, token_id in zip(mapped_rows, emitted)

286 + ]

287 + )

288 +

289 + assert masks == expected_masks

290 + assert all(token_id in mask for token_id, mask in zip(emitted, masks))

291 + assert torch.equal(on_logprobs.logprobs[mapped_rows, 0], expected_logprobs)

vllm/config/vllm.py

+8 -2

@@ -1265,9 +1265,15 @@ class VllmConfig:

1265 return

1266 if not self.use_v2_model_runner:

1267 raise ValueError("sampling distribution replay requires Model Runner V2")

1268 - if self.speculative_config is not None:

1268 + speculative_config = self.speculative_config

1269 + if (

1270 + speculative_config is not None

1271 + and speculative_config.enable_adaptive_verification

1272 + ):

1273 raise ValueError(

1270 - "sampling distribution replay does not support speculative decoding"

1274 + "sampling distribution replay with speculative decoding "

1275 + "requires fixed verification boundaries; disable adaptive "

1276 + "verification"

1277 )

1278 if model_config.is_diffusion:

1279 raise ValueError(

vllm/v1/engine/output_processor.py

+5 -1

@@ -441,7 +441,11 @@ class RequestState:

441 sampling_mask = None

442 if finished and self.sampling_mask_chunks:

443 sampling_mask = SamplingMask(

444 - [chunk.token_ids.tolist() for chunk in self.sampling_mask_chunks]

444 + [

445 + position

446 + for chunk in self.sampling_mask_chunks

447 + for position in chunk.to_nested_list()

448 + ]

449 )

450

451 # Concatenate routed experts on finish

vllm/v1/outputs.py

+10 -4

@@ -57,19 +57,25 @@ class LogprobsLists(NamedTuple):

59 class SamplingMaskLists(NamedTuple):

60 - """CSR sampling masks; a step slice holds one position (``offsets=None``)."""

60 + """CSR sampling masks; a request slice may contain multiple positions."""

62 # [num_kept_tokens]

63 token_ids: np.ndarray

64 # [num_positions + 1], or None for a single position

65 offsets: np.ndarray | None = None

66 - # Unused with one position per request; kept for the wire layout.

66 + # [num_requests + 1] for multi-position request batches.

67 cu_num_generated_tokens: list[int] | None = None

69 def slice_request(self, req_idx: int, num_positions: int) -> "SamplingMaskLists":

70 - assert num_positions == 1 and self.offsets is not None

70 + assert self.offsets is not None

71 + cu = self.cu_num_generated_tokens

72 + start = req_idx if cu is None else cu[req_idx]

73 + end = start + num_positions

74 + lo, hi = self.offsets[start], self.offsets[end]

75 + if num_positions == 1:

76 + return SamplingMaskLists(self.token_ids[lo:hi])

77 return SamplingMaskLists(

72 - self.token_ids[self.offsets[req_idx] : self.offsets[req_idx + 1]]

78 + self.token_ids[lo:hi], self.offsets[start : end + 1] - lo

79 )

81 def to_nested_list(self) -> list[list[int]]:

vllm/v1/worker/gpu/async_utils.py

+1 -1

@@ -185,7 +185,7 @@ class AsyncOutput(AsyncModelRunnerOutput):

185

186 if self.sampling_mask_tensors is not None:

187 self.model_runner_output.sampling_masks = (

188 - self.sampling_mask_tensors.tolists()

188 + self.sampling_mask_tensors.tolists(self.num_sampled_tokens_np)

189 )

190

191 if self.num_nans is not None:

vllm/v1/worker/gpu/sample/output.py

+70 -21

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from __future__ import annotations

5 +from collections.abc import Sequence

6 from dataclasses import dataclass

7 from typing import NamedTuple

@@ -27,6 +28,7 @@ def _compact_sampling_mask_kernel(

28 logits_ptr,

29 logits_row_stride,

30 logits_col_stride,

31 + cu_num_logits_ptr,

32 num_sampled_tokens_ptr,

33 token_ids_ptr,

34 token_ids_row_stride,

@@ -35,25 +37,32 @@ def _compact_sampling_mask_kernel(

37 counts_ptr,

38 vocab_size,

39 max_num_kept,

40 + ROWS_PER_REQUEST: tl.constexpr,

41 BLOCK_SIZE: tl.constexpr,

42 ):

43 """Per row: first ``max_num_kept`` finite-logit ids, the full count, the bitmask."""

41 - req_idx = tl.program_id(0)

42 - is_active = tl.load(num_sampled_tokens_ptr + req_idx) > 0

44 + output_row = tl.program_id(0)

45 + req_idx = output_row // ROWS_PER_REQUEST

46 + slot_idx = output_row % ROWS_PER_REQUEST

47 + if slot_idx >= tl.load(num_sampled_tokens_ptr + req_idx):

48 + tl.store(counts_ptr + output_row, 0)

49 + return

50 + source_row = tl.load(cu_num_logits_ptr + req_idx) + slot_idx

51 +

52 count = tl.zeros((), dtype=tl.int32)

54 for start_idx in range(0, vocab_size, BLOCK_SIZE):

55 offsets = start_idx + tl.arange(0, BLOCK_SIZE)

56 logits = tl.load(

48 - logits_ptr + req_idx * logits_row_stride + offsets * logits_col_stride,

57 + logits_ptr + source_row * logits_row_stride + offsets * logits_col_stride,

58 mask=offsets < vocab_size,

59 other=-float("inf"),

60 )

52 - keep = (logits > -float("inf")) & (logits < float("inf")) & is_active

61 + keep = (logits > -float("inf")) & (logits < float("inf"))

62 keep_i32 = keep.to(tl.int32)

63 pos = count + tl.cumsum(keep_i32, axis=0) - keep_i32

64 tl.store(

56 - token_ids_ptr + req_idx * token_ids_row_stride + pos,

65 + token_ids_ptr + output_row * token_ids_row_stride + pos,

66 offsets.to(tl.int32),

67 mask=keep & (pos < max_num_kept),

68 )

@@ -62,12 +71,12 @@ def _compact_sampling_mask_kernel(

71 bits = tl.reshape(keep_i32, (BLOCK_SIZE // 8, 8)) << tl.arange(0, 8)[None, :]

72 byte_offsets = start_idx // 8 + tl.arange(0, BLOCK_SIZE // 8)

73 tl.store(

65 - packed_mask_ptr + req_idx * packed_mask_row_stride + byte_offsets,

74 + packed_mask_ptr + output_row * packed_mask_row_stride + byte_offsets,

75 tl.sum(bits, axis=1).to(tl.uint8),

76 mask=byte_offsets < tl.cdiv(vocab_size, 8),

77 )

70 - tl.store(counts_ptr + req_idx, count)

79 + tl.store(counts_ptr + output_row, count)

82 # Bounds the [num_reqs, width] int32 buffer; wider rows use the bitmask instead.

@@ -78,37 +87,46 @@ class SamplingMaskTensors(NamedTuple):

87 """Device-side masks pending async D2H: compact ids, plus the bitmask as

88 the exact fallback for rows wider than ``max_num_kept``."""

81 - # [num_requests, max_num_kept]

90 + # [num_requests * rows_per_request, max_num_kept]

91 token_ids: torch.Tensor

83 - # [num_requests, ceil(vocab_size / 8)]

92 + # [num_requests * rows_per_request, ceil(vocab_size / 8)]

93 packed_mask: torch.Tensor

85 - # [num_requests]

94 + # [num_requests * rows_per_request]

95 counts: torch.Tensor

96 vocab_size: int

97 + # Fixed request-major output slots. Ordinary sampling has one slot.

98 + rows_per_request: int = 1

100 @classmethod

101 def from_logits(

102 cls,

103 logits: torch.Tensor,

104 + cu_num_logits: torch.Tensor,

105 num_sampled_tokens: torch.Tensor,

106 max_num_kept: int,

107 + rows_per_request: int = 1,

108 ) -> SamplingMaskTensors:

96 - """Capture the finite-logit support of every row with a sampled token."""

97 - num_reqs, vocab_size = logits.shape

109 + """Capture committed target supports in fixed request-major slots."""

110 + num_reqs = num_sampled_tokens.shape[0]

111 + vocab_size = logits.shape[1]

112 max_num_kept = min(max_num_kept, vocab_size, MAX_COMPACT_SUPPORT)

113 + num_output_rows = num_reqs * rows_per_request

114 device = logits.device

115

116 token_ids = torch.empty(

102 - (num_reqs, max_num_kept), dtype=torch.int32, device=device

117 + (num_output_rows, max_num_kept), dtype=torch.int32, device=device

118 )

119 packed_mask = torch.empty(

105 - (num_reqs, (vocab_size + 7) // 8), dtype=torch.uint8, device=device

120 + (num_output_rows, (vocab_size + 7) // 8),

121 + dtype=torch.uint8,

122 + device=device,

123 )

107 - counts = torch.empty(num_reqs, dtype=torch.int32, device=device)

108 - _compact_sampling_mask_kernel[(num_reqs,)](

124 + counts = torch.empty(num_output_rows, dtype=torch.int32, device=device)

125 + _compact_sampling_mask_kernel[(num_output_rows,)](

126 logits,

127 logits.stride(0),

128 logits.stride(1),

129 + cu_num_logits,

130 num_sampled_tokens,

131 token_ids,

132 token_ids.stride(0),

@@ -117,9 +135,23 @@ class SamplingMaskTensors(NamedTuple):

135 counts,

136 vocab_size,

137 max_num_kept,

138 + ROWS_PER_REQUEST=rows_per_request,

139 BLOCK_SIZE=8192,

140 )

122 - return cls(token_ids, packed_mask, counts, vocab_size)

141 + return cls(token_ids, packed_mask, counts, vocab_size, rows_per_request)

142 +

143 + @classmethod

144 + def cat(cls, chunks: Sequence[SamplingMaskTensors]) -> SamplingMaskTensors:

145 + """Join request-contiguous chunks that share one output layout."""

146 + first = chunks[0]

147 + assert len({(c.vocab_size, c.rows_per_request) for c in chunks}) == 1

148 + return cls(

149 + torch.cat([chunk.token_ids for chunk in chunks]),

150 + torch.cat([chunk.packed_mask for chunk in chunks]),

151 + torch.cat([chunk.counts for chunk in chunks]),

152 + first.vocab_size,

153 + first.rows_per_request,

154 + )

155

156 def to_cpu_nonblocking(self) -> SamplingMaskTensors:

157 if self.token_ids.device.type == "cpu":

@@ -129,15 +161,31 @@ class SamplingMaskTensors(NamedTuple):

161 self.packed_mask.to("cpu", non_blocking=True),

162 self.counts.to("cpu", non_blocking=True),

163 self.vocab_size,

164 + self.rows_per_request,

165 )

166

134 - def tolists(self) -> SamplingMaskLists:

135 - """CSR over all requests; rows without a sampled token are empty."""

167 + def tolists(

168 + self, num_sampled_tokens: np.ndarray | None = None

169 + ) -> SamplingMaskLists:

170 + """Convert fixed output rows to the scheduler's CSR representation."""

171 counts = self.counts.cpu().numpy()

172 token_ids = self.token_ids.cpu().numpy()

173 packed_mask = self.packed_mask.cpu().numpy()

174 width = token_ids.shape[1]

175

176 + cu_num_generated_tokens = None

177 + if self.rows_per_request == 1:

178 + sampled_rows = np.arange(len(counts))

179 + else:

180 + assert num_sampled_tokens is not None

181 + active_slots = (

182 + np.arange(self.rows_per_request)[None, :] < num_sampled_tokens[:, None]

183 + )

184 + sampled_rows = np.flatnonzero(active_slots)

185 + cu_num_generated_tokens = np.cumsum(

186 + np.concatenate(([0], num_sampled_tokens))

187 + ).tolist()

188 +

189 def support(row: int) -> np.ndarray:

190 if counts[row] <= width:

191 return token_ids[row, : counts[row]]

@@ -147,7 +195,8 @@ class SamplingMaskTensors(NamedTuple):

195 )

196 return np.flatnonzero(bits).astype(np.int32, copy=False)

197

150 - supports = [support(row) for row in range(len(counts))]

198 + supports = [support(row) for row in sampled_rows]

199 offsets = np.zeros(len(supports) + 1, dtype=np.int64)

200 np.cumsum([len(s) for s in supports], out=offsets[1:])

153 - return SamplingMaskLists(np.concatenate(supports), offsets)

201 + flat = np.concatenate([np.empty(0, dtype=np.int32), *supports])

202 + return SamplingMaskLists(flat, offsets, cu_num_generated_tokens)

vllm/v1/worker/gpu/sample/sampler.py

+4 -1

@@ -203,7 +203,10 @@ class Sampler:

203 # Size by the validated top_k batch max; wider supports use the bitmask.

204 max_num_kept = int(np.max(self.sampling_states.top_k.np[idx_mapping_np]))

205 sampling_mask_tensors = SamplingMaskTensors.from_logits(

206 - processed_logits, num_sampled, max_num_kept

206 + processed_logits,

207 + input_batch.cu_num_logits,

208 + num_sampled,

209 + max_num_kept,

210 )

211

212 # These are GPU tensors.

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

+47 -5

@@ -18,7 +18,7 @@ from vllm.v1.worker.gpu.input_batch import (

18 )

19 from vllm.v1.worker.gpu.metrics.logits import get_num_nans

20 from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores

21 -from vllm.v1.worker.gpu.sample.output import SamplerOutput

21 +from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors

22 from vllm.v1.worker.gpu.sample.sampler import Sampler

23 from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS

24 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (

@@ -296,7 +296,12 @@ class RejectionSampler:

296 max_chunk_logits: int,

297 max_num_logprobs: int,

298 verify_draft_sampled: torch.Tensor | None = None,

299 - ) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]:

299 + ) -> tuple[

300 + torch.Tensor,

301 + torch.Tensor,

302 + LogprobsTensors | None,

303 + SamplingMaskTensors | None,

304 + ]:

305 if verify_draft_sampled is None:

306 verify_draft_sampled = draft_sampled

307 cu_num_logits_np = input_batch.cu_num_logits_np

@@ -316,6 +321,14 @@ class RejectionSampler:

321 sampled_chunks: list[torch.Tensor] = []

322 num_sampled_chunks: list[torch.Tensor] = []

323 logprobs_chunks: list[LogprobsTensors] = []

324 + sampling_mask_chunks: list[SamplingMaskTensors] = []

325 + max_num_kept = None

326 + if self.sampler.return_sampling_mask:

327 + max_num_kept = int(

328 + np.max(

329 + self.sampler.sampling_states.top_k.np[input_batch.idx_mapping_np]

330 + )

331 + )

332

333 for start, end in request_chunks:

334 lo = int(cu_num_logits_np[start])

@@ -346,13 +359,31 @@ class RejectionSampler:

359 )

360 if chunk_logprobs is not None:

361 logprobs_chunks.append(chunk_logprobs)

362 + if max_num_kept is not None:

363 + sampling_mask_chunks.append(

364 + SamplingMaskTensors.from_logits(

365 + processed_logits,

366 + chunk_cu_num_logits,

367 + num_sampled,

368 + max_num_kept,

369 + self.num_speculative_steps + 1,

370 + )

371 + )

372 del processed_logits

373 sampled_chunks.append(sampled)

374 num_sampled_chunks.append(num_sampled)

375

376 if len(sampled_chunks) == 1:

377 logprobs_tensors = logprobs_chunks[0] if logprobs_chunks else None

355 - return sampled_chunks[0], num_sampled_chunks[0], logprobs_tensors

378 + sampling_mask_tensors = (

379 + sampling_mask_chunks[0] if sampling_mask_chunks else None

380 + )

381 + return (

382 + sampled_chunks[0],

383 + num_sampled_chunks[0],

384 + logprobs_tensors,

385 + sampling_mask_tensors,

386 + )

387

388 logprobs_tensors = None

389 if logprobs_chunks:

@@ -366,7 +397,12 @@ class RejectionSampler:

397

398 sampled = torch.cat(sampled_chunks)

399 num_sampled = torch.cat(num_sampled_chunks)

369 - return sampled, num_sampled, logprobs_tensors

400 + sampling_mask_tensors = (

401 + SamplingMaskTensors.cat(sampling_mask_chunks)

402 + if sampling_mask_chunks

403 + else None

404 + )

405 + return sampled, num_sampled, logprobs_tensors, sampling_mask_tensors

406

407 def __call__(

408 self,

@@ -400,7 +436,12 @@ class RejectionSampler:

436 input_batch.idx_mapping_np

437 )

438 chunk_logit_limit = get_max_chunk_logits(logits.shape[1])

403 - sampled, num_sampled, logprobs_tensors = self._verify_in_chunks(

439 + (

440 + sampled,

441 + num_sampled,

442 + logprobs_tensors,

443 + sampling_mask_tensors,

444 + ) = self._verify_in_chunks(

445 logits,

446 input_batch,

447 draft_logits,

@@ -425,4 +466,5 @@ class RejectionSampler:

466 num_nans=num_nans,

467 num_sampled=num_sampled,

468 num_rejected=num_rejected,

469 + sampling_mask_tensors=sampling_mask_tensors,

470 )