@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 09fe178dba

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

Nick Hill committed Sep 24, 2026 at 08:48 UTC 09fe178dba08e6de4448ade1e37d14ae1dab3acf
11 files changed +28 -629
tests/kernels/moe/test_topk_softplus_sqrt.py
-130
@@ -339,136 +339,6 @@ 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 -
342 @pytest.mark.skipif(
343 not current_platform.is_cuda_alike(),
344 reason="This test is skipped on non-CUDA platform.",
tests/v1/spec_decode/test_dflash_prepare_inputs.py
-149
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from contextlib import contextmanager
4 from types import SimpleNamespace
5
6 import numpy as np
@@ -9,9 +8,7 @@ import pytest
8 import torch
9
10 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
12 -from vllm.v1.worker.gpu.spec_decode.dflash import speculator as dflash_speculator
11 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
14 - DFlashSpeculator,
12 prepare_dflash_inputs,
13 )
14
@@ -36,7 +33,6 @@ def _run_prepare(
33 input_buffers = SimpleNamespace(
34 input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
35 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),
36 query_start_loc=torch.full(
37 (max_num_reqs + 1,), -1, dtype=torch.int32, device=device
38 ),
@@ -144,114 +140,6 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix():
140 assert out.temperature[2].item() == 1.0
141 assert out.seeds[2].item() == 17
142
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 -
143
144 def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp():
145 out = _run_prepare(
@@ -286,40 +174,3 @@ def test_prepare_dflash_inputs_never_writes_the_null_block():
174 PAD_SLOT_ID,
175 PAD_SLOT_ID,
176 ]
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
+3 -149
@@ -5,35 +5,16 @@
5 from unittest.mock import Mock
6
7 import numpy as np
8 -import pytest
9 -import torch
8
11 -from vllm.v1.worker.gpu import model_runner, pp_utils
9 +from vllm.v1.worker.gpu import pp_utils
10
11
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))
12 +def _batch(num_computed, prefill_len, num_scheduled):
13 return Mock(
32 - num_reqs=num_reqs,
14 + num_reqs=len(num_computed),
15 num_computed_tokens_np=np.array(num_computed, dtype=np.int32),
16 prefill_len_np=np.array(prefill_len, dtype=np.int32),
17 num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32),
36 - idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64),
18 )
19
20
@@ -100,130 +81,3 @@ def test_decode_row_ahead_of_a_prefill_chunk():
81
82 assert mask is not None
83 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,7 +41,6 @@ if current_platform.is_cuda():
41 def _dsv4_topk_kernel(
42 gating_output_ptr,
43 correction_bias_ptr,
44 - is_padding_ptr,
44 topk_weights_ptr,
45 topk_ids_ptr,
46 routed_scaling_factor,
@@ -50,7 +49,6 @@ if current_platform.is_cuda():
49 image_sentinel_lo,
50 NUM_EXPERTS: tl.constexpr,
51 BLOCK_N: tl.constexpr,
53 - HAS_PADDING: tl.constexpr,
52 HAS_VL: tl.constexpr,
53 launch_pdl: tl.constexpr,
54 ):
@@ -108,11 +106,6 @@ if current_platform.is_cuda():
106 output_mask = topk_offsets < 6
107 output_offsets = row * 6 + topk_offsets
108
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 -
109 if launch_pdl:
110 tl.extra.cuda.gdc_launch_dependents()
111
@@ -125,20 +118,11 @@ def dsv4_topk(
118 correction_bias: torch.Tensor,
119 indices_dtype: torch.dtype,
120 routed_scaling_factor: float,
128 - is_padding: torch.Tensor | None = None,
121 input_ids: torch.Tensor | None = None,
122 bias_vl: torch.Tensor | None = None,
123 image_sentinel_lo: int = 0,
124 ) -> tuple[torch.Tensor, torch.Tensor]:
125 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 - )
126 has_vl = bias_vl is not None and image_sentinel_lo > 0
127 if bias_vl is not None:
128 assert input_ids is not None, "bias_vl routing requires input_ids"
@@ -152,7 +136,6 @@ def dsv4_topk(
136 _dsv4_topk_kernel[(num_tokens,)](
137 gating_output,
138 correction_bias,
155 - is_padding,
139 topk_weights,
140 topk_ids,
141 routed_scaling_factor,
@@ -161,7 +144,6 @@ def dsv4_topk(
144 image_sentinel_lo,
145 NUM_EXPERTS=num_experts,
146 BLOCK_N=triton.next_power_of_2(num_experts),
164 - HAS_PADDING=is_padding is not None,
147 HAS_VL=has_vl,
148 num_warps=1,
149 launch_pdl=current_platform.is_arch_support_pdl(),
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
-5
@@ -148,13 +148,9 @@ 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])
151 if (
152 scoring_func == "sqrtsoftplus"
153 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)
154 and can_use_dsv4_topk(
155 gating_output,
156 e_score_correction_bias,
@@ -169,7 +165,6 @@ def fused_topk_bias(
165 e_score_correction_bias,
166 output_indices_dtype,
167 routed_scaling_factor,
172 - is_padding=padding_mask,
168 input_ids=input_tokens,
169 bias_vl=bias_vl,
170 image_sentinel_lo=image_sentinel_lo,
vllm/v1/worker/gpu/input_batch.py
-26
@@ -555,11 +555,6 @@ 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,
558 ):
559 req_id = tl.program_id(0)
560 req_state_idx = tl.load(idx_mapping_ptr + req_id)
@@ -567,17 +562,6 @@ def _post_update_kernel(
562 # Filter rows with negative index entries.
563 return
564
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 -
565 total_len = tl.load(total_len_ptr + req_state_idx)
566 num_sampled = tl.load(num_sampled_ptr + req_id)
567 if num_sampled > 0:
@@ -638,11 +622,6 @@ def post_update(
622 all_token_ids: torch.Tensor,
623 # [max_num_reqs]
624 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,
625 ) -> None:
626 num_reqs = idx_mapping.shape[0]
627 _post_update_kernel[(num_reqs,)](
@@ -659,11 +638,6 @@ def post_update(
638 all_token_ids,
639 all_token_ids.stride(0),
640 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,
641 num_warps=1,
642 )
643
vllm/v1/worker/gpu/model_runner.py
+5 -63
@@ -862,7 +862,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
862 # dummy run the eagle speculator's propose to ensure DP/EP sync.
863 if self.speculator is not None:
864 assert self.sampler is not None
865 - assert hidden_states is not None
865 self.step_timing.drafter_start()
866 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
867 if self.speculator.supports_mm_inputs:
@@ -878,8 +877,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
877 spec_hidden_states = hidden_states
878 if hasattr(self.model, "get_mtp_target_hidden_states"):
879 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
881 - if pre_hc_hidden_states is not None:
882 - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]
880 + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]
881 if isinstance(self.sampler, GPUWatermarkSampler):
882 self.speculator.prepare_watermarking(
883 self.sampler._get_contexts(input_batch.idx_mapping),
@@ -1111,44 +1109,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1109 if outputs is not None:
1110 self.postprocess_sampled(**outputs)
1111
1114 - def warmup_pp_decode_update(self) -> None:
1115 - """JIT-compile the kernel behind ``update_pp_decode_requests``.
1116 -
1117 - That path only runs on real steps, so the warmup steps never reach it
1118 - on non-last PP ranks. Its first triton compile must not happen
1119 - mid-serving: the in-flight sampled-token broadcast keeps a NCCL kernel
1120 - spinning on this device, which blocks the CUDA module load and
1121 - deadlocks the pipeline. An all -1 idx_mapping makes this a no-op.
1122 - The freshly allocated int32 tensors are 16-byte aligned, matching the
1123 - padded views `PPHandler` produces at serving time (triton specializes
1124 - on pointer alignment).
1125 - """
1126 - assert self.pp_handler is not None
1127 - num_spec = self.pp_handler.max_sample_len - 1
1128 - broadcast_drafts = (
1129 - torch.zeros((1, num_spec), dtype=torch.int64, device=self.device)
1130 - if num_spec > 0
1131 - else None
1132 - )
1133 - post_update(
1134 - torch.full((1,), -1, dtype=torch.int64, device=self.device),
1135 - self.req_states.num_computed_tokens.gpu,
1136 - self.req_states.last_sampled_tokens,
1137 - None,
1138 - torch.zeros(
1139 - (1, self.pp_handler.max_sample_len),
1140 - dtype=torch.int64,
1141 - device=self.device,
1142 - ),
1143 - torch.zeros(1, dtype=torch.int32, device=self.device),
1144 - torch.zeros(1, dtype=torch.int32, device=self.device),
1145 - None,
1146 - self.req_states.all_token_ids.gpu,
1147 - self.req_states.total_len.gpu,
1148 - broadcast_drafts,
1149 - self.req_states.draft_tokens if broadcast_drafts is not None else None,
1150 - )
1151 -
1112 def add_requests(self, scheduler_output: SchedulerOutput) -> None:
1113 for new_req_data in scheduler_output.scheduled_new_reqs:
1114 assert new_req_data.prefill_token_ids is not None
@@ -1630,7 +1590,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1590 num_sampled: torch.Tensor,
1591 num_rejected: torch.Tensor,
1592 query_start_loc: torch.Tensor | None = None,
1633 - broadcast_drafts: torch.Tensor | None = None,
1593 ) -> None:
1594 # Update the number of computed tokens.
1595 output_bin_counts = None
@@ -1648,8 +1607,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1607 query_start_loc,
1608 self.req_states.all_token_ids.gpu,
1609 self.req_states.total_len.gpu,
1651 - broadcast_drafts,
1652 - self.req_states.draft_tokens if broadcast_drafts is not None else None,
1610 )
1611
1612 self.model_state.postprocess_state(
@@ -2170,10 +2127,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2127 spec_hidden_states = draft_hidden_states
2128 if hasattr(self.model, "get_mtp_target_hidden_states"):
2129 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
2173 - if pre_hc_hidden_states is not None:
2174 - spec_hidden_states = pre_hc_hidden_states[
2175 - : draft_hidden_states.size(0)
2176 - ]
2130 + spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
2131 if isinstance(self.sampler, GPUWatermarkSampler):
2132 self.speculator.prepare_watermarking(
2133 self.sampler._get_contexts(input_batch.idx_mapping),
@@ -2195,16 +2149,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2149 dp_sync=dp_sync,
2150 mm_inputs=mm_inputs,
2151 )
2198 - if draft_tokens is not None:
2199 - self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
2200 - if self.pp_handler is not None:
2201 - # Earlier stages never run the speculator; ship the
2202 - # drafts so their next verification step embeds the real
2203 - # draft tokens instead of stale buffer contents.
2204 - self.pp_handler.broadcast_drafts(
2205 - self.req_states.draft_tokens, input_batch
2206 - )
2207 - if draft_tokens is not None and self.adaptive_verification is not None:
2152 + self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
2153 + if self.adaptive_verification is not None:
2154 self.adaptive_verification.record_confidences(
2155 self.speculator.draft_token_confidence_probs, input_batch
2156 )
@@ -2216,11 +2162,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2162 input_batch,
2163 self.req_states.draft_tokens[input_batch.idx_mapping],
2164 )
2219 - if self.pp_handler is not None and self.speculator is None:
2220 - # When a speculator ran, the propose() path above already
2221 - # broadcast the fresh drafts. Broadcasting here as well would
2222 - # double-post on the pp_broadcast group and misalign the
2223 - # recv FIFO on earlier stages, hanging the pipeline.
2165 + if self.pp_handler is not None:
2166 self.pp_handler.broadcast_drafts(
2167 self.req_states.draft_tokens, input_batch
2168 )
vllm/v1/worker/gpu/pp_utils.py
+20 -65
@@ -46,18 +46,6 @@ 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 -
49 class PPHandler:
50 """Runs the PP sampled-token broadcast/recv on a side stream so the
51 default stream isn't gated by the matching peer call. Step T's recv is
@@ -98,14 +86,6 @@ class PPHandler:
86 )
87 self.aux_hidden_state_relay_keys: tuple[str, ...] = ()
88
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 -
89 def on_req_idx_freed(self, req_idx: int) -> None:
90 self.req_idx_gen_np[req_idx] += 1
91
@@ -177,15 +157,29 @@ class PPHandler:
157 num_sampled=slot.num_sampled,
158 num_rejected=slot.num_rejected,
159 idx_mapping=idx_mapping,
180 - broadcast_drafts=slot.draft_tokens,
160 )
161
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 +
179 def receive(self, input_batch: InputBatch) -> bool:
180 """Returns True iff sampled tokens need to be gathered from *all*
181 requests in the batch."""
182 assert not self.is_last_rank
187 - if self.disabled:
188 - return False
183 need_sampled_mask = compute_need_sampled_mask(input_batch)
184 if need_sampled_mask is None:
185 # Leave this step's reserved slot as None.
@@ -201,7 +195,7 @@ class PPHandler:
195 sampled_tokens = torch.empty(
196 num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device
197 )
204 - combined = _alloc_combined(num_reqs, self.device)
198 + combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device)
199 torch.distributed.broadcast(
200 sampled_tokens, src=self.last_rank, group=self.broadcast_group
201 )
@@ -219,7 +213,6 @@ class PPHandler:
213 torch.distributed.broadcast(
214 draft_tokens, src=self.last_rank, group=self.broadcast_group
215 )
222 -
216 event = self.broadcast_stream.record_event()
217 num_sampled, num_rejected = combined.unbind(dim=0)
218 # Must record_stream since these were allocated on broadcast stream but
@@ -249,16 +242,11 @@ class PPHandler:
242 input_batch: InputBatch,
243 ) -> None:
244 assert self.is_last_rank
252 - if self.disabled:
253 - return
254 - mask = compute_need_sampled_mask(input_batch)
255 - if mask is None:
245 + if compute_need_sampled_mask(input_batch) is None:
246 # No request needs sampled outputs for a subsequent decode step.
247 return
248
249 assert sampled_token_ids.dtype == torch.int64
260 - assert num_sampled.dtype == torch.int32
261 - assert num_rejected.dtype == torch.int32
250
251 if current_platform.is_xpu():
252 self.main_stream.synchronize()
@@ -274,42 +262,9 @@ class PPHandler:
262 src=self.last_rank,
263 group=self.broadcast_group,
264 )
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
265 + combined = torch.stack((num_sampled, num_rejected), dim=0)
266 torch.distributed.broadcast(
267 combined, src=self.last_rank, group=self.broadcast_group
268 )
269 for tensor in (sampled_token_ids, num_sampled, num_rejected):
270 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,7 +256,6 @@ 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],
259 ):
260 last_hidden_states = self.model(
261 input_ids=self.input_buffers.input_ids[:num_tokens],
@@ -488,7 +487,6 @@ def _prepare_dflash_inputs_kernel(
487 # Outputs
488 out_input_ids_ptr,
489 out_query_positions_ptr,
491 - out_is_padding_ptr,
490 out_query_start_loc_ptr,
491 out_seq_lens_ptr,
492 out_query_slot_mapping_ptr,
@@ -619,7 +617,6 @@ def _prepare_dflash_inputs_kernel(
617 )
618
619 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)
620 clamped_query_pos = tl.minimum(query_pos, max_model_len - 1)
621 tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query)
622 tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
@@ -682,9 +679,6 @@ def _prepare_dflash_inputs_kernel(
679 for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE):
680 block = i + tl.arange(0, BLOCK_SIZE)
681 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)
682 tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask)
683
684
@@ -736,7 +730,6 @@ def prepare_dflash_inputs(
730 _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
731 input_buffers.input_ids,
732 input_buffers.positions,
739 - input_buffers.is_padding,
733 input_buffers.query_start_loc,
734 input_buffers.seq_lens,
735 query_slot_mapping,
vllm/v1/worker/gpu/warmup.py
-5
@@ -445,11 +445,6 @@ 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 -
448 # Clean up - process finish_req_ids.
449 cleanup_output = SchedulerOutput.make_empty()
450 cleanup_output.finished_req_ids = set(req_ids)
vllm/v1/worker/gpu_worker.py
-12
@@ -804,15 +804,6 @@ class Worker(WorkerBase):
804
805 @instrument(span_name="Warmup (GPU)")
806 def compile_or_warm_up_model(self) -> CompilationTimes:
807 - # All warmup phases below run synthetic steps whose sampled outputs are
808 - # discarded. The PP sampled-token broadcast would carry no payload, and
809 - # its side-stream NCCL ops can overlap the next step's activation p2p
810 - # and deadlock the pipeline, so keep it disabled for the whole warmup
811 - # window and restore it before serving.
812 - pp_handler = getattr(self.model_runner, "pp_handler", None)
813 - if pp_handler is not None:
814 - pp_handler.set_disabled(True)
815 -
807 warmup_sizes: list[int] = []
808
809 if (
@@ -991,9 +982,6 @@ class Worker(WorkerBase):
982 # gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
983 enable_gpu_sync_check()
984
994 - if pp_handler is not None:
995 - pp_handler.set_disabled(False)
996 -
985 return CompilationTimes(
986 language_model=self.compilation_config.compilation_time,
987 encoder=self.compilation_config.encoder_compilation_time,