@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 157bcb7c48

[DSpark] Support pipeline-parallel targets in aggregated serving (#56956)

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Kimi Code <noreply@moonshot.cn> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

Gabriel Wu committed Sep 23, 2026 at 19:52 UTC 157bcb7c489689dd34cf28d9c9970a465d326a03
11 files changed +629 -28
tests/kernels/moe/test_topk_softplus_sqrt.py
+130
@@ -339,6 +339,136 @@ def test_dsv4_fast_topk(
339 )
340
341
342 +@pytest.mark.skipif(
343 + not current_platform.is_cuda(),
344 + reason="The DeepSeek V4 fast path is CUDA-only.",
345 +)
346 +def test_dsv4_fast_topk_padding_uint32_falls_back(monkeypatch: pytest.MonkeyPatch):
347 + """Padded rows need the -1 sentinel, which uint32 cannot represent: the
348 + router must skip the dsv4 fast path and still route the real rows."""
349 + torch.manual_seed(0)
350 + num_tokens = 17
351 + num_experts = 256
352 + hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda")
353 + gating_output = torch.randn(
354 + (num_tokens, num_experts), dtype=torch.float32, device="cuda"
355 + )
356 + correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
357 + is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
358 + is_padding[1::2] = True
359 + gating_output[is_padding] = float("nan")
360 +
361 + monkeypatch.setattr(
362 + "vllm.model_executor.layers.fused_moe.router."
363 + "fused_topk_bias_router._get_padding_mask",
364 + lambda _: is_padding,
365 + )
366 + # uint32 + padding would trip dsv4_topk's signed-indices assertion; the
367 + # generic path must take over instead.
368 + topk_weights, topk_ids = fused_topk_bias(
369 + hidden_states=hidden_states,
370 + gating_output=gating_output,
371 + scoring_func="sqrtsoftplus",
372 + e_score_correction_bias=correction_bias,
373 + topk=6,
374 + renormalize=True,
375 + indices_type=torch.uint32,
376 + routed_scaling_factor=1.5,
377 + )
378 +
379 + assert topk_ids.dtype == torch.uint32
380 + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
381 + gating_output=gating_output[~is_padding],
382 + topk=6,
383 + renormalize=True,
384 + routed_scaling_factor=1.5,
385 + e_score_correction_bias=correction_bias,
386 + )
387 + # uint32 CUDA tensors do not support boolean-mask indexing; widen first.
388 + torch.testing.assert_close(
389 + topk_ids.to(torch.int64)[~is_padding],
390 + topk_ids_ref.to(torch.int64),
391 + atol=0,
392 + rtol=0,
393 + )
394 + torch.testing.assert_close(
395 + topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5
396 + )
397 +
398 +
399 +@pytest.mark.skipif(
400 + not current_platform.is_cuda(),
401 + reason="The DeepSeek V4 fast path is CUDA-only.",
402 +)
403 +def test_dsv4_fast_topk_padding(monkeypatch: pytest.MonkeyPatch):
404 + """Verify the DSV4 fast path removes graph-padding rows from routing."""
405 + torch.manual_seed(0)
406 + num_tokens = 17
407 + num_experts = 256
408 + hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda")
409 + gating_output = torch.randn(
410 + (num_tokens, num_experts), dtype=torch.float32, device="cuda"
411 + )
412 + correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
413 + is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
414 + is_padding[1::2] = True
415 + gating_output[is_padding] = float("nan")
416 +
417 + monkeypatch.setattr(
418 + "vllm.model_executor.layers.fused_moe.router."
419 + "fused_topk_bias_router._get_padding_mask",
420 + lambda _: is_padding,
421 + )
422 + topk_weights, topk_ids = fused_topk_bias(
423 + hidden_states=hidden_states,
424 + gating_output=gating_output,
425 + scoring_func="sqrtsoftplus",
426 + e_score_correction_bias=correction_bias,
427 + topk=6,
428 + renormalize=True,
429 + routed_scaling_factor=1.5,
430 + )
431 +
432 + assert torch.equal(topk_ids[is_padding], torch.full_like(topk_ids[is_padding], -1))
433 + assert torch.equal(
434 + topk_weights[is_padding], torch.zeros_like(topk_weights[is_padding])
435 + )
436 +
437 + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
438 + gating_output=gating_output[~is_padding],
439 + topk=6,
440 + renormalize=True,
441 + routed_scaling_factor=1.5,
442 + e_score_correction_bias=correction_bias,
443 + )
444 + torch.testing.assert_close(topk_ids[~is_padding], topk_ids_ref, atol=0, rtol=0)
445 + torch.testing.assert_close(
446 + topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5
447 + )
448 +
449 + # The mask buffer is persistent under CUDA graph replay, but its contents
450 + # change with every batch. Verify that the kernel reads those contents at
451 + # runtime rather than specializing on the mask captured above.
452 + graph = torch.cuda.CUDAGraph()
453 + with torch.cuda.graph(graph):
454 + graph_weights, graph_ids = dsv4_topk(
455 + gating_output,
456 + correction_bias,
457 + torch.int32,
458 + 1.5,
459 + is_padding=is_padding,
460 + )
461 +
462 + is_padding.logical_not_()
463 + graph.replay()
464 + assert torch.equal(
465 + graph_ids[is_padding], torch.full_like(graph_ids[is_padding], -1)
466 + )
467 + assert torch.equal(
468 + graph_weights[is_padding], torch.zeros_like(graph_weights[is_padding])
469 + )
470 +
471 +
472 @pytest.mark.skipif(
473 not current_platform.is_cuda_alike(),
474 reason="This test is skipped on non-CUDA platform.",
tests/v1/spec_decode/test_dflash_prepare_inputs.py
+149
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from contextlib import contextmanager
5 from types import SimpleNamespace
6
7 import numpy as np
@@ -8,7 +9,9 @@ import pytest
9 import torch
10
11 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
12 +from vllm.v1.worker.gpu.spec_decode.dflash import speculator as dflash_speculator
13 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
14 + DFlashSpeculator,
15 prepare_dflash_inputs,
16 )
17
@@ -33,6 +36,7 @@ def _run_prepare(
36 input_buffers = SimpleNamespace(
37 input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
38 positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device),
39 + is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device),
40 query_start_loc=torch.full(
41 (max_num_reqs + 1,), -1, dtype=torch.int32, device=device
42 ),
@@ -140,6 +144,114 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix():
144 assert out.temperature[2].item() == 1.0
145 assert out.seeds[2].item() == 17
146
147 + assert not out.input_buffers.is_padding[:3].any()
148 + assert out.input_buffers.is_padding[3:].all()
149 + assert out.input_buffers.input_ids[3:].cpu().tolist() == [0] * 13
150 + assert out.input_buffers.positions[3:].cpu().tolist() == [0] * 13
151 +
152 +
153 +def test_prepare_dflash_inputs_compacts_noncontiguous_request_slots():
154 + device = torch.device("cuda")
155 + max_num_reqs = 4
156 + max_num_tokens = 16
157 + num_speculative_steps = 3
158 + input_buffers = SimpleNamespace(
159 + input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
160 + positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device),
161 + is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device),
162 + query_start_loc=torch.full(
163 + (max_num_reqs + 1,), -1, dtype=torch.int32, device=device
164 + ),
165 + seq_lens=torch.full((max_num_reqs,), -1, dtype=torch.int32, device=device),
166 + )
167 + input_batch = SimpleNamespace(
168 + num_reqs=2,
169 + num_scheduled_tokens=np.array([4, 4], dtype=np.int32),
170 + positions=torch.tensor(
171 + [10, 11, 12, 13, 20, 21, 22, 23],
172 + dtype=torch.int64,
173 + device=device,
174 + ),
175 + query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32, device=device),
176 + # Active batch rows are compact, while request state remains in slots 3 and 1.
177 + idx_mapping=torch.tensor([3, 1], dtype=torch.int32, device=device),
178 + )
179 + query_slot_mapping = torch.full(
180 + (max_num_tokens,), -2, dtype=torch.int64, device=device
181 + )
182 + context_positions = torch.full(
183 + (max_num_tokens,), -1, dtype=torch.int64, device=device
184 + )
185 + context_slot_mapping = torch.full(
186 + (max_num_tokens,), -2, dtype=torch.int64, device=device
187 + )
188 + sample_indices = torch.full(
189 + (max_num_reqs * num_speculative_steps,),
190 + -1,
191 + dtype=torch.int64,
192 + device=device,
193 + )
194 + sample_pos = torch.full_like(sample_indices, -1)
195 + sample_idx_mapping = torch.full(
196 + sample_indices.shape, -1, dtype=torch.int32, device=device
197 + )
198 + temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device)
199 + seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device)
200 + input_temperature = torch.tensor(
201 + [0.0, 0.5, 0.0, 1.0], dtype=torch.float32, device=device
202 + )
203 + input_seeds = torch.tensor([0, 11, 0, 33], dtype=torch.int64, device=device)
204 + last_sampled = torch.tensor([0, 77, 0, 99], dtype=torch.int64, device=device)
205 + next_prefill_tokens = torch.zeros_like(last_sampled)
206 + block_table = torch.tensor(
207 + [[0, 0, 7, 8, 9, 10, 11, 12], [0, 0, 13, 14, 15, 16, 17, 18]],
208 + dtype=torch.int32,
209 + device=device,
210 + )
211 +
212 + prepare_dflash_inputs(
213 + input_buffers,
214 + query_slot_mapping,
215 + context_positions,
216 + context_slot_mapping,
217 + sample_indices,
218 + sample_pos,
219 + sample_idx_mapping,
220 + temperature,
221 + seeds,
222 + input_batch,
223 + torch.tensor([1, 1], dtype=torch.int32, device=device),
224 + torch.tensor([2, 1], dtype=torch.int32, device=device),
225 + last_sampled,
226 + next_prefill_tokens,
227 + input_temperature,
228 + input_seeds,
229 + block_table,
230 + 4,
231 + 0,
232 + 1,
233 + 1,
234 + 123,
235 + num_speculative_steps,
236 + num_speculative_steps,
237 + max_num_reqs,
238 + max_num_tokens,
239 + 128,
240 + sample_from_anchor=True,
241 + )
242 + torch.accelerator.synchronize()
243 +
244 + # Query rows follow compact batch order, but every persistent state lookup
245 + # follows idx_mapping instead of accidentally using the compact row index.
246 + assert input_buffers.input_ids[:6].cpu().tolist() == [99, 123, 123, 77, 123, 123]
247 + assert input_buffers.positions[:6].cpu().tolist() == [12, 13, 14, 23, 24, 25]
248 + assert sample_indices[:6].cpu().tolist() == [0, 1, 2, 3, 4, 5]
249 + assert sample_idx_mapping[:6].cpu().tolist() == [3, 3, 3, 1, 1, 1]
250 + assert temperature.cpu().tolist() == [0.0, 0.5, 0.0, 1.0]
251 + assert seeds.cpu().tolist() == [0, 11, 0, 33]
252 + assert not input_buffers.is_padding[:6].any()
253 + assert input_buffers.is_padding[6:].all()
254 +
255
256 def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp():
257 out = _run_prepare(
@@ -174,3 +286,40 @@ def test_prepare_dflash_inputs_never_writes_the_null_block():
286 PAD_SLOT_ID,
287 PAD_SLOT_ID,
288 ]
289 +
290 +
291 +def test_dflash_forward_context_receives_draft_padding_mask(monkeypatch):
292 + device = torch.device("cuda")
293 + input_buffers = SimpleNamespace(
294 + input_ids=torch.tensor([11, 12, 0, 0], dtype=torch.int32, device=device),
295 + positions=torch.tensor([7, 8, 0, 0], dtype=torch.int64, device=device),
296 + is_padding=torch.tensor([False, False, True, True], device=device),
297 + )
298 + observed = None
299 +
300 + @contextmanager
301 + def fake_set_forward_context(*args, **kwargs):
302 + nonlocal observed
303 + observed = kwargs["is_padding"].clone()
304 + yield
305 +
306 + monkeypatch.setattr(
307 + dflash_speculator, "set_forward_context", fake_set_forward_context
308 + )
309 + speculator = SimpleNamespace(
310 + input_buffers=input_buffers,
311 + vllm_config=SimpleNamespace(),
312 + model=lambda **kwargs: kwargs["input_ids"],
313 + )
314 +
315 + result = DFlashSpeculator._run_model(
316 + speculator,
317 + num_tokens=4,
318 + attn_metadata=None,
319 + slot_mappings=None,
320 + num_tokens_across_dp=None,
321 + )
322 +
323 + assert result.tolist() == [11, 12, 0, 0]
324 + assert observed is not None
325 + assert observed.tolist() == [False, False, True, True]
tests/v1/worker/test_pp_utils.py
+149 -3
@@ -5,16 +5,35 @@
5 from unittest.mock import Mock
6
7 import numpy as np
8 +import pytest
9 +import torch
10
9 -from vllm.v1.worker.gpu import pp_utils
11 +from vllm.v1.worker.gpu import model_runner, pp_utils
12
13
12 -def _batch(num_computed, prefill_len, num_scheduled):
14 +def _cuda_handler(max_sample_len=6):
15 + handler = object.__new__(pp_utils.PPHandler)
16 + handler.is_last_rank = True
17 + handler.disabled = False
18 + handler.max_sample_len = max_sample_len
19 + handler.last_rank = 1
20 + handler.broadcast_group = Mock()
21 + handler.device = torch.device("cuda")
22 + handler.main_stream = torch.cuda.current_stream()
23 + handler.broadcast_stream = torch.cuda.Stream()
24 + return handler
25 +
26 +
27 +def _batch(num_computed, prefill_len, num_scheduled, idx_mapping=None):
28 + num_reqs = len(num_computed)
29 + if idx_mapping is None:
30 + idx_mapping = list(range(num_reqs))
31 return Mock(
14 - num_reqs=len(num_computed),
32 + num_reqs=num_reqs,
33 num_computed_tokens_np=np.array(num_computed, dtype=np.int32),
34 prefill_len_np=np.array(prefill_len, dtype=np.int32),
35 num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32),
36 + idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64),
37 )
38
39
@@ -81,3 +100,130 @@ def test_decode_row_ahead_of_a_prefill_chunk():
100
101 assert mask is not None
102 assert mask.tolist() == [True, False]
103 +
104 +
105 +def test_disabled_handler_skips_broadcast_and_receive(monkeypatch):
106 + """While disabled (warmup), neither side enqueues a broadcast op."""
107 + sent = []
108 + monkeypatch.setattr(
109 + pp_utils.torch.distributed,
110 + "broadcast",
111 + lambda *args, **kwargs: sent.append((args, kwargs)),
112 + )
113 +
114 + handler = object.__new__(pp_utils.PPHandler)
115 + handler.set_disabled(True)
116 +
117 + handler.is_last_rank = False
118 + assert handler.receive(Mock()) is False
119 +
120 + handler.is_last_rank = True
121 + assert handler.broadcast(Mock(), Mock(), Mock(), Mock()) is None
122 +
123 + assert sent == []
124 +
125 + handler.set_disabled(False)
126 + assert handler.disabled is False
127 +
128 +
129 +def test_alloc_combined_keeps_unbind_views_16_byte_aligned():
130 + """Triton specializes on pointer alignment: an unaligned `num_rejected`
131 + would compile a second `_post_update_kernel` variant at serving time,
132 + where the in-flight broadcast NCCL kernel can block the module load."""
133 + for num_reqs in range(1, 9):
134 + combined = pp_utils._alloc_combined(num_reqs, torch.device("cpu"))
135 + num_sampled, num_rejected = combined.unbind(dim=0)
136 + assert num_sampled.data_ptr() % 16 == 0
137 + assert num_rejected.data_ptr() % 16 == 0
138 + assert combined.shape[1] >= num_reqs
139 +
140 +
141 +def test_warmup_pp_decode_update_matches_serving_specialization(monkeypatch):
142 + """The warmup launch must hit the same triton specialization as serving.
143 +
144 + A mismatch means the first real ``update_pp_decode_requests`` recompiles
145 + mid-serving, where the in-flight broadcast NCCL kernel blocks the CUDA
146 + module load and deadlocks the pipeline.
147 + """
148 + calls = []
149 + monkeypatch.setattr(model_runner, "post_update", lambda *args: calls.append(args))
150 +
151 + runner = object.__new__(model_runner.GPUModelRunner)
152 + runner.device = torch.device("cpu")
153 + runner.pp_handler = Mock(max_sample_len=3)
154 + runner.req_states = Mock()
155 +
156 + runner.warmup_pp_decode_update()
157 +
158 + assert len(calls) == 1
159 + args = calls[0]
160 + idx_mapping, _, _, output_bin_counts = args[:4]
161 + sampled_tokens, num_sampled, num_rejected, query_start_loc = args[4:8]
162 + broadcast_drafts, draft_tokens_out = args[10:12]
163 + assert idx_mapping.tolist() == [-1] and idx_mapping.dtype == torch.int64
164 + assert output_bin_counts is None
165 + assert query_start_loc is None
166 + assert sampled_tokens.shape == (1, 3) and sampled_tokens.dtype == torch.int64
167 + assert num_sampled.dtype == torch.int32
168 + assert num_rejected.dtype == torch.int32
169 + # Spec-enabled PP handlers receive drafts over the broadcast; the warmup
170 + # must compile that specialization (non-None draft pointers) too.
171 + assert broadcast_drafts.shape == (1, 2) and broadcast_drafts.dtype == torch.int64
172 + assert draft_tokens_out is runner.req_states.draft_tokens
173 +
174 +
175 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA stream")
176 +def test_broadcast_pads_plain_sampler_rows_to_max_sample_len(monkeypatch):
177 + """The wire shape must not depend on whether the batch carried drafts:
178 + the receiver always allocates [num_reqs, max_sample_len], and a NCCL
179 + broadcast with mismatched counts hangs the receiver."""
180 + sent = []
181 + monkeypatch.setattr(
182 + pp_utils.torch.distributed,
183 + "broadcast",
184 + lambda tensor, **kw: sent.append(tensor),
185 + )
186 + handler = _cuda_handler()
187 + batch = _batch(num_computed=[10], prefill_len=[8], num_scheduled=[1])
188 +
189 + handler.broadcast(
190 + torch.zeros(1, 1, dtype=torch.int64, device="cuda"), # plain sampler
191 + torch.ones(1, dtype=torch.int32, device="cuda"),
192 + torch.zeros(1, dtype=torch.int32, device="cuda"),
193 + batch,
194 + )
195 +
196 + assert sent[0].shape == (1, 6)
197 + assert sent[1].shape == (2, 4)
198 + torch.accelerator.synchronize()
199 +
200 +
201 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA stream")
202 +def test_broadcast_drafts_gathers_fresh_rows_from_the_table(monkeypatch):
203 + sent = []
204 + monkeypatch.setattr(
205 + pp_utils.torch.distributed,
206 + "broadcast",
207 + lambda tensor, **kw: sent.append(tensor),
208 + )
209 + handler = _cuda_handler()
210 + # The batch's single row maps to request-state row 2.
211 + batch = _batch(num_computed=[10], prefill_len=[8], num_scheduled=[1])
212 + batch.idx_mapping = torch.tensor([2], dtype=torch.int64, device="cuda")
213 + table = torch.arange(20, dtype=torch.int64, device="cuda").view(4, 5)
214 +
215 + handler.broadcast_drafts(table, batch)
216 +
217 + assert sent[0].shape == (1, 5)
218 + # The payload is a gather into a fresh tensor: propose() overwrites its
219 + # persistent buffer on the next step, possibly before this send completes.
220 + assert sent[0].data_ptr() != table.data_ptr()
221 + assert sent[0].cpu().tolist() == [table[2].cpu().tolist()]
222 +
223 + # An all-prefill batch sends nothing (receive() enqueues nothing either).
224 + sent.clear()
225 + prefill_batch = _batch(num_computed=[0], prefill_len=[4096], num_scheduled=[448])
226 + prefill_batch.idx_mapping = prefill_batch.idx_mapping.cuda()
227 + handler.broadcast_drafts(table, prefill_batch)
228 + assert sent == []
229 + torch.accelerator.synchronize()
vllm/model_executor/layers/fused_moe/router/dsv4_topk.py
+18
@@ -41,6 +41,7 @@ if current_platform.is_cuda():
41 def _dsv4_topk_kernel(
42 gating_output_ptr,
43 correction_bias_ptr,
44 + is_padding_ptr,
45 topk_weights_ptr,
46 topk_ids_ptr,
47 routed_scaling_factor,
@@ -49,6 +50,7 @@ if current_platform.is_cuda():
50 image_sentinel_lo,
51 NUM_EXPERTS: tl.constexpr,
52 BLOCK_N: tl.constexpr,
53 + HAS_PADDING: tl.constexpr,
54 HAS_VL: tl.constexpr,
55 launch_pdl: tl.constexpr,
56 ):
@@ -106,6 +108,11 @@ if current_platform.is_cuda():
108 output_mask = topk_offsets < 6
109 output_offsets = row * 6 + topk_offsets
110
111 + if HAS_PADDING:
112 + is_padding = tl.load(is_padding_ptr + row)
113 + selected_weights = tl.where(is_padding, 0.0, selected_weights)
114 + selected_ids = tl.where(is_padding, -1, selected_ids)
115 +
116 if launch_pdl:
117 tl.extra.cuda.gdc_launch_dependents()
118
@@ -118,11 +125,20 @@ def dsv4_topk(
125 correction_bias: torch.Tensor,
126 indices_dtype: torch.dtype,
127 routed_scaling_factor: float,
128 + is_padding: torch.Tensor | None = None,
129 input_ids: torch.Tensor | None = None,
130 bias_vl: torch.Tensor | None = None,
131 image_sentinel_lo: int = 0,
132 ) -> tuple[torch.Tensor, torch.Tensor]:
133 num_tokens, num_experts = gating_output.shape
134 + if is_padding is not None:
135 + assert is_padding.dtype == torch.bool
136 + assert is_padding.shape == (num_tokens,)
137 + assert is_padding.device == gating_output.device
138 + assert is_padding.is_contiguous()
139 + assert indices_dtype in (torch.int32, torch.int64), (
140 + "Padding requires a signed indices dtype for the -1 sentinel."
141 + )
142 has_vl = bias_vl is not None and image_sentinel_lo > 0
143 if bias_vl is not None:
144 assert input_ids is not None, "bias_vl routing requires input_ids"
@@ -136,6 +152,7 @@ def dsv4_topk(
152 _dsv4_topk_kernel[(num_tokens,)](
153 gating_output,
154 correction_bias,
155 + is_padding,
156 topk_weights,
157 topk_ids,
158 routed_scaling_factor,
@@ -144,6 +161,7 @@ def dsv4_topk(
161 image_sentinel_lo,
162 NUM_EXPERTS=num_experts,
163 BLOCK_N=triton.next_power_of_2(num_experts),
164 + HAS_PADDING=is_padding is not None,
165 HAS_VL=has_vl,
166 num_warps=1,
167 launch_pdl=current_platform.is_arch_support_pdl(),
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+5
@@ -148,9 +148,13 @@ def fused_topk_bias(
148 )
149
150 output_indices_dtype = torch.int32 if indices_type is None else indices_type
151 + padding_mask = _get_padding_mask(gating_output.shape[0])
152 if (
153 scoring_func == "sqrtsoftplus"
154 and hash_indices_table is None
155 + # dsv4_topk marks padded rows with a -1 sentinel, which unsigned
156 + # indices cannot represent; keep the generic path for that pair.
157 + and (padding_mask is None or output_indices_dtype != torch.uint32)
158 and can_use_dsv4_topk(
159 gating_output,
160 e_score_correction_bias,
@@ -165,6 +169,7 @@ def fused_topk_bias(
169 e_score_correction_bias,
170 output_indices_dtype,
171 routed_scaling_factor,
172 + is_padding=padding_mask,
173 input_ids=input_tokens,
174 bias_vl=bias_vl,
175 image_sentinel_lo=image_sentinel_lo,
vllm/v1/worker/gpu/input_batch.py
+26
@@ -555,6 +555,11 @@ def _post_update_kernel(
555 all_token_ids_ptr,
556 all_token_ids_stride,
557 total_len_ptr,
558 + broadcast_drafts_ptr,
559 + broadcast_drafts_stride,
560 + draft_tokens_ptr,
561 + draft_tokens_stride,
562 + num_spec,
563 ):
564 req_id = tl.program_id(0)
565 req_state_idx = tl.load(idx_mapping_ptr + req_id)
@@ -562,6 +567,17 @@ def _post_update_kernel(
567 # Filter rows with negative index entries.
568 return
569
570 + if broadcast_drafts_ptr is not None:
571 + # PP path: adopt the draft tokens proposed by the last rank's
572 + # speculator so the next verification step embeds the real drafts.
573 + for i in range(num_spec):
574 + token_id = tl.load(
575 + broadcast_drafts_ptr + req_id * broadcast_drafts_stride + i
576 + )
577 + tl.store(
578 + draft_tokens_ptr + req_state_idx * draft_tokens_stride + i, token_id
579 + )
580 +
581 total_len = tl.load(total_len_ptr + req_state_idx)
582 num_sampled = tl.load(num_sampled_ptr + req_id)
583 if num_sampled > 0:
@@ -622,6 +638,11 @@ def post_update(
638 all_token_ids: torch.Tensor,
639 # [max_num_reqs]
640 total_len: torch.Tensor,
641 + # [num_reqs, num_spec]; drafts broadcast from the last PP rank. Only
642 + # passed on non-last PP ranks, which never run the speculator.
643 + broadcast_drafts: torch.Tensor | None = None,
644 + # [max_num_reqs, num_spec]
645 + draft_tokens_out: torch.Tensor | None = None,
646 ) -> None:
647 num_reqs = idx_mapping.shape[0]
648 _post_update_kernel[(num_reqs,)](
@@ -638,6 +659,11 @@ def post_update(
659 all_token_ids,
660 all_token_ids.stride(0),
661 total_len,
662 + broadcast_drafts,
663 + broadcast_drafts.stride(0) if broadcast_drafts is not None else 0,
664 + draft_tokens_out,
665 + draft_tokens_out.stride(0) if draft_tokens_out is not None else 0,
666 + broadcast_drafts.shape[1] if broadcast_drafts is not None else 0,
667 num_warps=1,
668 )
669
vllm/v1/worker/gpu/model_runner.py
+63 -5
@@ -860,6 +860,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
860 # dummy run the eagle speculator's propose to ensure DP/EP sync.
861 if self.speculator is not None:
862 assert self.sampler is not None
863 + assert hidden_states is not None
864 self.step_timing.drafter_start()
865 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
866 if self.speculator.supports_mm_inputs:
@@ -875,7 +876,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
876 spec_hidden_states = hidden_states
877 if hasattr(self.model, "get_mtp_target_hidden_states"):
878 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
878 - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]
879 + if pre_hc_hidden_states is not None:
880 + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]
881 if isinstance(self.sampler, GPUWatermarkSampler):
882 self.speculator.prepare_watermarking(
883 self.sampler._get_contexts(input_batch.idx_mapping),
@@ -1107,6 +1109,44 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1109 if outputs is not None:
1110 self.postprocess_sampled(**outputs)
1111
1112 + def warmup_pp_decode_update(self) -> None:
1113 + """JIT-compile the kernel behind ``update_pp_decode_requests``.
1114 +
1115 + That path only runs on real steps, so the warmup steps never reach it
1116 + on non-last PP ranks. Its first triton compile must not happen
1117 + mid-serving: the in-flight sampled-token broadcast keeps a NCCL kernel
1118 + spinning on this device, which blocks the CUDA module load and
1119 + deadlocks the pipeline. An all -1 idx_mapping makes this a no-op.
1120 + The freshly allocated int32 tensors are 16-byte aligned, matching the
1121 + padded views `PPHandler` produces at serving time (triton specializes
1122 + on pointer alignment).
1123 + """
1124 + assert self.pp_handler is not None
1125 + num_spec = self.pp_handler.max_sample_len - 1
1126 + broadcast_drafts = (
1127 + torch.zeros((1, num_spec), dtype=torch.int64, device=self.device)
1128 + if num_spec > 0
1129 + else None
1130 + )
1131 + post_update(
1132 + torch.full((1,), -1, dtype=torch.int64, device=self.device),
1133 + self.req_states.num_computed_tokens.gpu,
1134 + self.req_states.last_sampled_tokens,
1135 + None,
1136 + torch.zeros(
1137 + (1, self.pp_handler.max_sample_len),
1138 + dtype=torch.int64,
1139 + device=self.device,
1140 + ),
1141 + torch.zeros(1, dtype=torch.int32, device=self.device),
1142 + torch.zeros(1, dtype=torch.int32, device=self.device),
1143 + None,
1144 + self.req_states.all_token_ids.gpu,
1145 + self.req_states.total_len.gpu,
1146 + broadcast_drafts,
1147 + self.req_states.draft_tokens if broadcast_drafts is not None else None,
1148 + )
1149 +
1150 def add_requests(self, scheduler_output: SchedulerOutput) -> None:
1151 for new_req_data in scheduler_output.scheduled_new_reqs:
1152 assert new_req_data.prefill_token_ids is not None
@@ -1588,6 +1628,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1628 num_sampled: torch.Tensor,
1629 num_rejected: torch.Tensor,
1630 query_start_loc: torch.Tensor | None = None,
1631 + broadcast_drafts: torch.Tensor | None = None,
1632 ) -> None:
1633 # Update the number of computed tokens.
1634 output_bin_counts = None
@@ -1605,6 +1646,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1646 query_start_loc,
1647 self.req_states.all_token_ids.gpu,
1648 self.req_states.total_len.gpu,
1649 + broadcast_drafts,
1650 + self.req_states.draft_tokens if broadcast_drafts is not None else None,
1651 )
1652
1653 self.model_state.postprocess_state(
@@ -2125,7 +2168,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2168 spec_hidden_states = draft_hidden_states
2169 if hasattr(self.model, "get_mtp_target_hidden_states"):
2170 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
2128 - spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
2171 + if pre_hc_hidden_states is not None:
2172 + spec_hidden_states = pre_hc_hidden_states[
2173 + : draft_hidden_states.size(0)
2174 + ]
2175 if isinstance(self.sampler, GPUWatermarkSampler):
2176 self.speculator.prepare_watermarking(
2177 self.sampler._get_contexts(input_batch.idx_mapping),
@@ -2147,8 +2193,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2193 dp_sync=dp_sync,
2194 mm_inputs=mm_inputs,
2195 )
2150 - self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
2151 - if self.adaptive_verification is not None:
2196 + if draft_tokens is not None:
2197 + self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
2198 + if self.pp_handler is not None:
2199 + # Earlier stages never run the speculator; ship the
2200 + # drafts so their next verification step embeds the real
2201 + # draft tokens instead of stale buffer contents.
2202 + self.pp_handler.broadcast_drafts(
2203 + self.req_states.draft_tokens, input_batch
2204 + )
2205 + if draft_tokens is not None and self.adaptive_verification is not None:
2206 self.adaptive_verification.record_confidences(
2207 self.speculator.draft_token_confidence_probs, input_batch
2208 )
@@ -2160,7 +2214,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2214 input_batch,
2215 self.req_states.draft_tokens[input_batch.idx_mapping],
2216 )
2163 - if self.pp_handler is not None:
2217 + if self.pp_handler is not None and self.speculator is None:
2218 + # When a speculator ran, the propose() path above already
2219 + # broadcast the fresh drafts. Broadcasting here as well would
2220 + # double-post on the pp_broadcast group and misalign the
2221 + # recv FIFO on earlier stages, hanging the pipeline.
2222 self.pp_handler.broadcast_drafts(
2223 self.req_states.draft_tokens, input_batch
2224 )
vllm/v1/worker/gpu/pp_utils.py
+65 -20
@@ -46,6 +46,18 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
46 return produces_sample if produces_sample.any() else None
47
48
49 +def _alloc_combined(num_reqs: int, device: torch.device) -> torch.Tensor:
50 + """Allocate the (2, N) int32 buffer broadcast alongside sampled tokens.
51 +
52 + The inner dim is padded to a multiple of 4 so that both `unbind` views
53 + stay 16-byte aligned for any `num_reqs`: triton specializes on pointer
54 + alignment, and a misaligned `num_rejected` would JIT-compile a second
55 + `_post_update_kernel` variant at serving time. The padding is broadcast
56 + but never read. Sender and receiver must both use this allocation.
57 + """
58 + return torch.empty(2, -(-num_reqs // 4) * 4, dtype=torch.int32, device=device)
59 +
60 +
61 class PPHandler:
62 """Runs the PP sampled-token broadcast/recv on a side stream so the
63 default stream isn't gated by the matching peer call. Step T's recv is
@@ -86,6 +98,14 @@ class PPHandler:
98 )
99 self.aux_hidden_state_relay_keys: tuple[str, ...] = ()
100
101 + # Warmup steps run the pipeline with synthetic batches whose outputs are
102 + # discarded; the sampled-token broadcast is disabled there so its
103 + # side-stream NCCL ops cannot overlap the next step's activation p2p.
104 + self.disabled = False
105 +
106 + def set_disabled(self, disabled: bool) -> None:
107 + self.disabled = disabled
108 +
109 def on_req_idx_freed(self, req_idx: int) -> None:
110 self.req_idx_gen_np[req_idx] += 1
111
@@ -157,29 +177,15 @@ class PPHandler:
177 num_sampled=slot.num_sampled,
178 num_rejected=slot.num_rejected,
179 idx_mapping=idx_mapping,
180 + broadcast_drafts=slot.draft_tokens,
181 )
182
162 - def broadcast_drafts(
163 - self, draft_tokens: torch.Tensor, input_batch: InputBatch
164 - ) -> None:
165 - """Broadcast draft proposals so non-last ranks can embed real token ids."""
166 - assert self.is_last_rank
167 - if compute_need_sampled_mask(input_batch) is None:
168 - return
169 - with torch.cuda.stream(self.broadcast_stream):
170 - self.broadcast_stream.wait_stream(self.main_stream)
171 - send = draft_tokens[input_batch.idx_mapping].contiguous()
172 - # Must record the idx_mapping tensor since it was allocated
173 - # on the main stream.
174 - input_batch.idx_mapping.record_stream(self.broadcast_stream)
175 - torch.distributed.broadcast(
176 - send, src=self.last_rank, group=self.broadcast_group
177 - )
178 -
183 def receive(self, input_batch: InputBatch) -> bool:
184 """Returns True iff sampled tokens need to be gathered from *all*
185 requests in the batch."""
186 assert not self.is_last_rank
187 + if self.disabled:
188 + return False
189 need_sampled_mask = compute_need_sampled_mask(input_batch)
190 if need_sampled_mask is None:
191 # Leave this step's reserved slot as None.
@@ -195,7 +201,7 @@ class PPHandler:
201 sampled_tokens = torch.empty(
202 num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device
203 )
198 - combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device)
204 + combined = _alloc_combined(num_reqs, self.device)
205 torch.distributed.broadcast(
206 sampled_tokens, src=self.last_rank, group=self.broadcast_group
207 )
@@ -213,6 +219,7 @@ class PPHandler:
219 torch.distributed.broadcast(
220 draft_tokens, src=self.last_rank, group=self.broadcast_group
221 )
222 +
223 event = self.broadcast_stream.record_event()
224 num_sampled, num_rejected = combined.unbind(dim=0)
225 # Must record_stream since these were allocated on broadcast stream but
@@ -242,11 +249,16 @@ class PPHandler:
249 input_batch: InputBatch,
250 ) -> None:
251 assert self.is_last_rank
245 - if compute_need_sampled_mask(input_batch) is None:
252 + if self.disabled:
253 + return
254 + mask = compute_need_sampled_mask(input_batch)
255 + if mask is None:
256 # No request needs sampled outputs for a subsequent decode step.
257 return
258
259 assert sampled_token_ids.dtype == torch.int64
260 + assert num_sampled.dtype == torch.int32
261 + assert num_rejected.dtype == torch.int32
262
263 if current_platform.is_xpu():
264 self.main_stream.synchronize()
@@ -262,9 +274,42 @@ class PPHandler:
274 src=self.last_rank,
275 group=self.broadcast_group,
276 )
265 - combined = torch.stack((num_sampled, num_rejected), dim=0)
277 + combined = _alloc_combined(num_sampled.shape[0], self.device)
278 + combined[0, : num_sampled.shape[0]] = num_sampled
279 + combined[1, : num_sampled.shape[0]] = num_rejected
280 torch.distributed.broadcast(
281 combined, src=self.last_rank, group=self.broadcast_group
282 )
283 for tensor in (sampled_token_ids, num_sampled, num_rejected):
284 tensor.record_stream(self.broadcast_stream)
285 +
286 + def broadcast_drafts(
287 + self, draft_token_table: torch.Tensor, input_batch: InputBatch
288 + ) -> None:
289 + """Broadcast the speculator's freshly proposed draft tokens.
290 +
291 + Runs after propose() on the last rank; the send is stream-ordered
292 + after broadcast()'s sends, matching receive()'s enqueue order. The
293 + payload is gathered from the runner's draft table (just updated from
294 + propose()'s output) into a fresh compact tensor: the speculator
295 + overwrites its own persistent buffer on the next step, possibly
296 + before this async send completes.
297 + """
298 + assert self.is_last_rank
299 + if self.disabled or self.max_sample_len == 1:
300 + return
301 + if compute_need_sampled_mask(input_batch) is None:
302 + return
303 + assert draft_token_table.dtype == torch.int64
304 + assert draft_token_table.shape[1] == self.max_sample_len - 1
305 +
306 + # Gather on the main stream so the payload is ordered after this
307 + # step's table update and before any later one.
308 + drafts = draft_token_table[input_batch.idx_mapping]
309 + assert drafts.shape[0] == input_batch.num_reqs
310 + with torch.cuda.stream(self.broadcast_stream):
311 + self.broadcast_stream.wait_stream(self.main_stream)
312 + torch.distributed.broadcast(
313 + drafts, src=self.last_rank, group=self.broadcast_group
314 + )
315 + drafts.record_stream(self.broadcast_stream)
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+7
@@ -256,6 +256,7 @@ class DFlashSpeculator(DraftModelSpeculator):
256 num_tokens_across_dp=num_tokens_across_dp,
257 slot_mapping=slot_mappings,
258 batch_descriptor=batch_descriptor,
259 + is_padding=self.input_buffers.is_padding[:num_tokens],
260 ):
261 last_hidden_states = self.model(
262 input_ids=self.input_buffers.input_ids[:num_tokens],
@@ -487,6 +488,7 @@ def _prepare_dflash_inputs_kernel(
488 # Outputs
489 out_input_ids_ptr,
490 out_query_positions_ptr,
491 + out_is_padding_ptr,
492 out_query_start_loc_ptr,
493 out_seq_lens_ptr,
494 out_query_slot_mapping_ptr,
@@ -617,6 +619,7 @@ def _prepare_dflash_inputs_kernel(
619 )
620
621 tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
622 + tl.store(out_is_padding_ptr + query_idx, False, mask=is_query)
623 clamped_query_pos = tl.minimum(query_pos, max_model_len - 1)
624 tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query)
625 tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
@@ -679,6 +682,9 @@ def _prepare_dflash_inputs_kernel(
682 for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE):
683 block = i + tl.arange(0, BLOCK_SIZE)
684 mask = block < max_num_tokens
685 + tl.store(out_input_ids_ptr + block, 0, mask=mask)
686 + tl.store(out_query_positions_ptr + block, 0, mask=mask)
687 + tl.store(out_is_padding_ptr + block, True, mask=mask)
688 tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask)
689
690
@@ -730,6 +736,7 @@ def prepare_dflash_inputs(
736 _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
737 input_buffers.input_ids,
738 input_buffers.positions,
739 + input_buffers.is_padding,
740 input_buffers.query_start_loc,
741 input_buffers.seq_lens,
742 query_slot_mapping,
vllm/v1/worker/gpu/warmup.py
+5
@@ -445,6 +445,11 @@ def _warmup_kernels(
445 for step_indices, step_spec_flags in decode_steps:
446 _run_decode_step(step_indices, step_spec_flags)
447
448 + # The deferred PP post-update path only runs on real steps, so the steps
449 + # above never JIT-compile its kernel on non-last ranks.
450 + if not model_runner.is_last_pp_rank and model_runner.pp_handler is not None:
451 + model_runner.warmup_pp_decode_update()
452 +
453 # Clean up - process finish_req_ids.
454 cleanup_output = SchedulerOutput.make_empty()
455 cleanup_output.finished_req_ids = set(req_ids)
vllm/v1/worker/gpu_worker.py
+12
@@ -805,6 +805,15 @@ class Worker(WorkerBase):
805
806 @instrument(span_name="Warmup (GPU)")
807 def compile_or_warm_up_model(self) -> CompilationTimes:
808 + # All warmup phases below run synthetic steps whose sampled outputs are
809 + # discarded. The PP sampled-token broadcast would carry no payload, and
810 + # its side-stream NCCL ops can overlap the next step's activation p2p
811 + # and deadlock the pipeline, so keep it disabled for the whole warmup
812 + # window and restore it before serving.
813 + pp_handler = getattr(self.model_runner, "pp_handler", None)
814 + if pp_handler is not None:
815 + pp_handler.set_disabled(True)
816 +
817 warmup_sizes: list[int] = []
818
819 if (
@@ -983,6 +992,9 @@ class Worker(WorkerBase):
992 # gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
993 enable_gpu_sync_check()
994
995 + if pp_handler is not None:
996 + pp_handler.set_disabled(False)
997 +
998 return CompilationTimes(
999 language_model=self.compilation_config.compilation_time,
1000 encoder=self.compilation_config.encoder_compilation_time,