From 99c9ad2e4f0fe63f3f29c690f0627db47a613ba8 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 23 Sep 2026 17:30:16 -0700 Subject: [PATCH] Revert "[DSpark] Support pipeline-parallel targets in aggregated serving (#56956)" This reverts commit 157bcb7c489689dd34cf28d9c9970a465d326a03. --- tests/kernels/moe/test_topk_softplus_sqrt.py | 130 --------------- .../spec_decode/test_dflash_prepare_inputs.py | 149 ----------------- tests/v1/worker/test_pp_utils.py | 152 +----------------- .../layers/fused_moe/router/dsv4_topk.py | 18 --- .../router/fused_topk_bias_router.py | 5 - vllm/v1/worker/gpu/input_batch.py | 26 --- vllm/v1/worker/gpu/model_runner.py | 68 +------- vllm/v1/worker/gpu/pp_utils.py | 85 +++------- .../gpu/spec_decode/dflash/speculator.py | 7 - vllm/v1/worker/gpu/warmup.py | 5 - vllm/v1/worker/gpu_worker.py | 12 -- 11 files changed, 28 insertions(+), 629 deletions(-) diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index bb8856f83a4d..bc438c66931b 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -339,136 +339,6 @@ def test_dsv4_fast_topk( ) -@pytest.mark.skipif( - not current_platform.is_cuda(), - reason="The DeepSeek V4 fast path is CUDA-only.", -) -def test_dsv4_fast_topk_padding_uint32_falls_back(monkeypatch: pytest.MonkeyPatch): - """Padded rows need the -1 sentinel, which uint32 cannot represent: the - router must skip the dsv4 fast path and still route the real rows.""" - torch.manual_seed(0) - num_tokens = 17 - num_experts = 256 - hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda") - gating_output = torch.randn( - (num_tokens, num_experts), dtype=torch.float32, device="cuda" - ) - correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") - is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda") - is_padding[1::2] = True - gating_output[is_padding] = float("nan") - - monkeypatch.setattr( - "vllm.model_executor.layers.fused_moe.router." - "fused_topk_bias_router._get_padding_mask", - lambda _: is_padding, - ) - # uint32 + padding would trip dsv4_topk's signed-indices assertion; the - # generic path must take over instead. - topk_weights, topk_ids = fused_topk_bias( - hidden_states=hidden_states, - gating_output=gating_output, - scoring_func="sqrtsoftplus", - e_score_correction_bias=correction_bias, - topk=6, - renormalize=True, - indices_type=torch.uint32, - routed_scaling_factor=1.5, - ) - - assert topk_ids.dtype == torch.uint32 - topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( - gating_output=gating_output[~is_padding], - topk=6, - renormalize=True, - routed_scaling_factor=1.5, - e_score_correction_bias=correction_bias, - ) - # uint32 CUDA tensors do not support boolean-mask indexing; widen first. - torch.testing.assert_close( - topk_ids.to(torch.int64)[~is_padding], - topk_ids_ref.to(torch.int64), - atol=0, - rtol=0, - ) - torch.testing.assert_close( - topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5 - ) - - -@pytest.mark.skipif( - not current_platform.is_cuda(), - reason="The DeepSeek V4 fast path is CUDA-only.", -) -def test_dsv4_fast_topk_padding(monkeypatch: pytest.MonkeyPatch): - """Verify the DSV4 fast path removes graph-padding rows from routing.""" - torch.manual_seed(0) - num_tokens = 17 - num_experts = 256 - hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda") - gating_output = torch.randn( - (num_tokens, num_experts), dtype=torch.float32, device="cuda" - ) - correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") - is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda") - is_padding[1::2] = True - gating_output[is_padding] = float("nan") - - monkeypatch.setattr( - "vllm.model_executor.layers.fused_moe.router." - "fused_topk_bias_router._get_padding_mask", - lambda _: is_padding, - ) - topk_weights, topk_ids = fused_topk_bias( - hidden_states=hidden_states, - gating_output=gating_output, - scoring_func="sqrtsoftplus", - e_score_correction_bias=correction_bias, - topk=6, - renormalize=True, - routed_scaling_factor=1.5, - ) - - assert torch.equal(topk_ids[is_padding], torch.full_like(topk_ids[is_padding], -1)) - assert torch.equal( - topk_weights[is_padding], torch.zeros_like(topk_weights[is_padding]) - ) - - topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( - gating_output=gating_output[~is_padding], - topk=6, - renormalize=True, - routed_scaling_factor=1.5, - e_score_correction_bias=correction_bias, - ) - torch.testing.assert_close(topk_ids[~is_padding], topk_ids_ref, atol=0, rtol=0) - torch.testing.assert_close( - topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5 - ) - - # The mask buffer is persistent under CUDA graph replay, but its contents - # change with every batch. Verify that the kernel reads those contents at - # runtime rather than specializing on the mask captured above. - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_weights, graph_ids = dsv4_topk( - gating_output, - correction_bias, - torch.int32, - 1.5, - is_padding=is_padding, - ) - - is_padding.logical_not_() - graph.replay() - assert torch.equal( - graph_ids[is_padding], torch.full_like(graph_ids[is_padding], -1) - ) - assert torch.equal( - graph_weights[is_padding], torch.zeros_like(graph_weights[is_padding]) - ) - - @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="This test is skipped on non-CUDA platform.", diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index f3f94e1a398f..16d6d8e516df 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import contextmanager from types import SimpleNamespace import numpy as np @@ -9,9 +8,7 @@ import torch from vllm.v1.attention.backends.utils import PAD_SLOT_ID -from vllm.v1.worker.gpu.spec_decode.dflash import speculator as dflash_speculator from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( - DFlashSpeculator, prepare_dflash_inputs, ) @@ -36,7 +33,6 @@ def _run_prepare( input_buffers = SimpleNamespace( input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device), positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device), - is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device), query_start_loc=torch.full( (max_num_reqs + 1,), -1, dtype=torch.int32, device=device ), @@ -144,114 +140,6 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix(): assert out.temperature[2].item() == 1.0 assert out.seeds[2].item() == 17 - assert not out.input_buffers.is_padding[:3].any() - assert out.input_buffers.is_padding[3:].all() - assert out.input_buffers.input_ids[3:].cpu().tolist() == [0] * 13 - assert out.input_buffers.positions[3:].cpu().tolist() == [0] * 13 - - -def test_prepare_dflash_inputs_compacts_noncontiguous_request_slots(): - device = torch.device("cuda") - max_num_reqs = 4 - max_num_tokens = 16 - num_speculative_steps = 3 - input_buffers = SimpleNamespace( - input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device), - positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device), - is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device), - query_start_loc=torch.full( - (max_num_reqs + 1,), -1, dtype=torch.int32, device=device - ), - seq_lens=torch.full((max_num_reqs,), -1, dtype=torch.int32, device=device), - ) - input_batch = SimpleNamespace( - num_reqs=2, - num_scheduled_tokens=np.array([4, 4], dtype=np.int32), - positions=torch.tensor( - [10, 11, 12, 13, 20, 21, 22, 23], - dtype=torch.int64, - device=device, - ), - query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32, device=device), - # Active batch rows are compact, while request state remains in slots 3 and 1. - idx_mapping=torch.tensor([3, 1], dtype=torch.int32, device=device), - ) - query_slot_mapping = torch.full( - (max_num_tokens,), -2, dtype=torch.int64, device=device - ) - context_positions = torch.full( - (max_num_tokens,), -1, dtype=torch.int64, device=device - ) - context_slot_mapping = torch.full( - (max_num_tokens,), -2, dtype=torch.int64, device=device - ) - sample_indices = torch.full( - (max_num_reqs * num_speculative_steps,), - -1, - dtype=torch.int64, - device=device, - ) - sample_pos = torch.full_like(sample_indices, -1) - sample_idx_mapping = torch.full( - sample_indices.shape, -1, dtype=torch.int32, device=device - ) - temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device) - seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device) - input_temperature = torch.tensor( - [0.0, 0.5, 0.0, 1.0], dtype=torch.float32, device=device - ) - input_seeds = torch.tensor([0, 11, 0, 33], dtype=torch.int64, device=device) - last_sampled = torch.tensor([0, 77, 0, 99], dtype=torch.int64, device=device) - next_prefill_tokens = torch.zeros_like(last_sampled) - block_table = torch.tensor( - [[0, 0, 7, 8, 9, 10, 11, 12], [0, 0, 13, 14, 15, 16, 17, 18]], - dtype=torch.int32, - device=device, - ) - - prepare_dflash_inputs( - input_buffers, - query_slot_mapping, - context_positions, - context_slot_mapping, - sample_indices, - sample_pos, - sample_idx_mapping, - temperature, - seeds, - input_batch, - torch.tensor([1, 1], dtype=torch.int32, device=device), - torch.tensor([2, 1], dtype=torch.int32, device=device), - last_sampled, - next_prefill_tokens, - input_temperature, - input_seeds, - block_table, - 4, - 0, - 1, - 1, - 123, - num_speculative_steps, - num_speculative_steps, - max_num_reqs, - max_num_tokens, - 128, - sample_from_anchor=True, - ) - torch.accelerator.synchronize() - - # Query rows follow compact batch order, but every persistent state lookup - # follows idx_mapping instead of accidentally using the compact row index. - assert input_buffers.input_ids[:6].cpu().tolist() == [99, 123, 123, 77, 123, 123] - assert input_buffers.positions[:6].cpu().tolist() == [12, 13, 14, 23, 24, 25] - assert sample_indices[:6].cpu().tolist() == [0, 1, 2, 3, 4, 5] - assert sample_idx_mapping[:6].cpu().tolist() == [3, 3, 3, 1, 1, 1] - assert temperature.cpu().tolist() == [0.0, 0.5, 0.0, 1.0] - assert seeds.cpu().tolist() == [0, 11, 0, 33] - assert not input_buffers.is_padding[:6].any() - assert input_buffers.is_padding[6:].all() - def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp(): out = _run_prepare( @@ -286,40 +174,3 @@ def test_prepare_dflash_inputs_never_writes_the_null_block(): PAD_SLOT_ID, PAD_SLOT_ID, ] - - -def test_dflash_forward_context_receives_draft_padding_mask(monkeypatch): - device = torch.device("cuda") - input_buffers = SimpleNamespace( - input_ids=torch.tensor([11, 12, 0, 0], dtype=torch.int32, device=device), - positions=torch.tensor([7, 8, 0, 0], dtype=torch.int64, device=device), - is_padding=torch.tensor([False, False, True, True], device=device), - ) - observed = None - - @contextmanager - def fake_set_forward_context(*args, **kwargs): - nonlocal observed - observed = kwargs["is_padding"].clone() - yield - - monkeypatch.setattr( - dflash_speculator, "set_forward_context", fake_set_forward_context - ) - speculator = SimpleNamespace( - input_buffers=input_buffers, - vllm_config=SimpleNamespace(), - model=lambda **kwargs: kwargs["input_ids"], - ) - - result = DFlashSpeculator._run_model( - speculator, - num_tokens=4, - attn_metadata=None, - slot_mappings=None, - num_tokens_across_dp=None, - ) - - assert result.tolist() == [11, 12, 0, 0] - assert observed is not None - assert observed.tolist() == [False, False, True, True] diff --git a/tests/v1/worker/test_pp_utils.py b/tests/v1/worker/test_pp_utils.py index 8dfdf2e6d480..cfc92479f57e 100644 --- a/tests/v1/worker/test_pp_utils.py +++ b/tests/v1/worker/test_pp_utils.py @@ -5,35 +5,16 @@ from unittest.mock import Mock import numpy as np -import pytest -import torch -from vllm.v1.worker.gpu import model_runner, pp_utils +from vllm.v1.worker.gpu import pp_utils -def _cuda_handler(max_sample_len=6): - handler = object.__new__(pp_utils.PPHandler) - handler.is_last_rank = True - handler.disabled = False - handler.max_sample_len = max_sample_len - handler.last_rank = 1 - handler.broadcast_group = Mock() - handler.device = torch.device("cuda") - handler.main_stream = torch.cuda.current_stream() - handler.broadcast_stream = torch.cuda.Stream() - return handler - - -def _batch(num_computed, prefill_len, num_scheduled, idx_mapping=None): - num_reqs = len(num_computed) - if idx_mapping is None: - idx_mapping = list(range(num_reqs)) +def _batch(num_computed, prefill_len, num_scheduled): return Mock( - num_reqs=num_reqs, + num_reqs=len(num_computed), num_computed_tokens_np=np.array(num_computed, dtype=np.int32), prefill_len_np=np.array(prefill_len, dtype=np.int32), num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32), - idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64), ) @@ -100,130 +81,3 @@ def test_decode_row_ahead_of_a_prefill_chunk(): assert mask is not None assert mask.tolist() == [True, False] - - -def test_disabled_handler_skips_broadcast_and_receive(monkeypatch): - """While disabled (warmup), neither side enqueues a broadcast op.""" - sent = [] - monkeypatch.setattr( - pp_utils.torch.distributed, - "broadcast", - lambda *args, **kwargs: sent.append((args, kwargs)), - ) - - handler = object.__new__(pp_utils.PPHandler) - handler.set_disabled(True) - - handler.is_last_rank = False - assert handler.receive(Mock()) is False - - handler.is_last_rank = True - assert handler.broadcast(Mock(), Mock(), Mock(), Mock()) is None - - assert sent == [] - - handler.set_disabled(False) - assert handler.disabled is False - - -def test_alloc_combined_keeps_unbind_views_16_byte_aligned(): - """Triton specializes on pointer alignment: an unaligned `num_rejected` - would compile a second `_post_update_kernel` variant at serving time, - where the in-flight broadcast NCCL kernel can block the module load.""" - for num_reqs in range(1, 9): - combined = pp_utils._alloc_combined(num_reqs, torch.device("cpu")) - num_sampled, num_rejected = combined.unbind(dim=0) - assert num_sampled.data_ptr() % 16 == 0 - assert num_rejected.data_ptr() % 16 == 0 - assert combined.shape[1] >= num_reqs - - -def test_warmup_pp_decode_update_matches_serving_specialization(monkeypatch): - """The warmup launch must hit the same triton specialization as serving. - - A mismatch means the first real ``update_pp_decode_requests`` recompiles - mid-serving, where the in-flight broadcast NCCL kernel blocks the CUDA - module load and deadlocks the pipeline. - """ - calls = [] - monkeypatch.setattr(model_runner, "post_update", lambda *args: calls.append(args)) - - runner = object.__new__(model_runner.GPUModelRunner) - runner.device = torch.device("cpu") - runner.pp_handler = Mock(max_sample_len=3) - runner.req_states = Mock() - - runner.warmup_pp_decode_update() - - assert len(calls) == 1 - args = calls[0] - idx_mapping, _, _, output_bin_counts = args[:4] - sampled_tokens, num_sampled, num_rejected, query_start_loc = args[4:8] - broadcast_drafts, draft_tokens_out = args[10:12] - assert idx_mapping.tolist() == [-1] and idx_mapping.dtype == torch.int64 - assert output_bin_counts is None - assert query_start_loc is None - assert sampled_tokens.shape == (1, 3) and sampled_tokens.dtype == torch.int64 - assert num_sampled.dtype == torch.int32 - assert num_rejected.dtype == torch.int32 - # Spec-enabled PP handlers receive drafts over the broadcast; the warmup - # must compile that specialization (non-None draft pointers) too. - assert broadcast_drafts.shape == (1, 2) and broadcast_drafts.dtype == torch.int64 - assert draft_tokens_out is runner.req_states.draft_tokens - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA stream") -def test_broadcast_pads_plain_sampler_rows_to_max_sample_len(monkeypatch): - """The wire shape must not depend on whether the batch carried drafts: - the receiver always allocates [num_reqs, max_sample_len], and a NCCL - broadcast with mismatched counts hangs the receiver.""" - sent = [] - monkeypatch.setattr( - pp_utils.torch.distributed, - "broadcast", - lambda tensor, **kw: sent.append(tensor), - ) - handler = _cuda_handler() - batch = _batch(num_computed=[10], prefill_len=[8], num_scheduled=[1]) - - handler.broadcast( - torch.zeros(1, 1, dtype=torch.int64, device="cuda"), # plain sampler - torch.ones(1, dtype=torch.int32, device="cuda"), - torch.zeros(1, dtype=torch.int32, device="cuda"), - batch, - ) - - assert sent[0].shape == (1, 6) - assert sent[1].shape == (2, 4) - torch.accelerator.synchronize() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA stream") -def test_broadcast_drafts_gathers_fresh_rows_from_the_table(monkeypatch): - sent = [] - monkeypatch.setattr( - pp_utils.torch.distributed, - "broadcast", - lambda tensor, **kw: sent.append(tensor), - ) - handler = _cuda_handler() - # The batch's single row maps to request-state row 2. - batch = _batch(num_computed=[10], prefill_len=[8], num_scheduled=[1]) - batch.idx_mapping = torch.tensor([2], dtype=torch.int64, device="cuda") - table = torch.arange(20, dtype=torch.int64, device="cuda").view(4, 5) - - handler.broadcast_drafts(table, batch) - - assert sent[0].shape == (1, 5) - # The payload is a gather into a fresh tensor: propose() overwrites its - # persistent buffer on the next step, possibly before this send completes. - assert sent[0].data_ptr() != table.data_ptr() - assert sent[0].cpu().tolist() == [table[2].cpu().tolist()] - - # An all-prefill batch sends nothing (receive() enqueues nothing either). - sent.clear() - prefill_batch = _batch(num_computed=[0], prefill_len=[4096], num_scheduled=[448]) - prefill_batch.idx_mapping = prefill_batch.idx_mapping.cuda() - handler.broadcast_drafts(table, prefill_batch) - assert sent == [] - torch.accelerator.synchronize() diff --git a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py index 1df95bc6119f..dbd29f3d08ab 100644 --- a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py +++ b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py @@ -41,7 +41,6 @@ def can_use_dsv4_topk( def _dsv4_topk_kernel( gating_output_ptr, correction_bias_ptr, - is_padding_ptr, topk_weights_ptr, topk_ids_ptr, routed_scaling_factor, @@ -50,7 +49,6 @@ def _dsv4_topk_kernel( image_sentinel_lo, NUM_EXPERTS: tl.constexpr, BLOCK_N: tl.constexpr, - HAS_PADDING: tl.constexpr, HAS_VL: tl.constexpr, launch_pdl: tl.constexpr, ): @@ -108,11 +106,6 @@ def _dsv4_topk_kernel( output_mask = topk_offsets < 6 output_offsets = row * 6 + topk_offsets - if HAS_PADDING: - is_padding = tl.load(is_padding_ptr + row) - selected_weights = tl.where(is_padding, 0.0, selected_weights) - selected_ids = tl.where(is_padding, -1, selected_ids) - if launch_pdl: tl.extra.cuda.gdc_launch_dependents() @@ -125,20 +118,11 @@ def dsv4_topk( correction_bias: torch.Tensor, indices_dtype: torch.dtype, routed_scaling_factor: float, - is_padding: torch.Tensor | None = None, input_ids: torch.Tensor | None = None, bias_vl: torch.Tensor | None = None, image_sentinel_lo: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens, num_experts = gating_output.shape - if is_padding is not None: - assert is_padding.dtype == torch.bool - assert is_padding.shape == (num_tokens,) - assert is_padding.device == gating_output.device - assert is_padding.is_contiguous() - assert indices_dtype in (torch.int32, torch.int64), ( - "Padding requires a signed indices dtype for the -1 sentinel." - ) has_vl = bias_vl is not None and image_sentinel_lo > 0 if bias_vl is not None: assert input_ids is not None, "bias_vl routing requires input_ids" @@ -152,7 +136,6 @@ def dsv4_topk( _dsv4_topk_kernel[(num_tokens,)]( gating_output, correction_bias, - is_padding, topk_weights, topk_ids, routed_scaling_factor, @@ -161,7 +144,6 @@ def dsv4_topk( image_sentinel_lo, NUM_EXPERTS=num_experts, BLOCK_N=triton.next_power_of_2(num_experts), - HAS_PADDING=is_padding is not None, HAS_VL=has_vl, num_warps=1, launch_pdl=current_platform.is_arch_support_pdl(), diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index d464d73ba17d..257ede27d5fa 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -148,13 +148,9 @@ def fused_topk_bias( ) output_indices_dtype = torch.int32 if indices_type is None else indices_type - padding_mask = _get_padding_mask(gating_output.shape[0]) if ( scoring_func == "sqrtsoftplus" and hash_indices_table is None - # dsv4_topk marks padded rows with a -1 sentinel, which unsigned - # indices cannot represent; keep the generic path for that pair. - and (padding_mask is None or output_indices_dtype != torch.uint32) and can_use_dsv4_topk( gating_output, e_score_correction_bias, @@ -169,7 +165,6 @@ def fused_topk_bias( e_score_correction_bias, output_indices_dtype, routed_scaling_factor, - is_padding=padding_mask, input_ids=input_tokens, bias_vl=bias_vl, image_sentinel_lo=image_sentinel_lo, diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 5437405d7fdc..75bbef5d8d04 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -555,11 +555,6 @@ def _post_update_kernel( all_token_ids_ptr, all_token_ids_stride, total_len_ptr, - broadcast_drafts_ptr, - broadcast_drafts_stride, - draft_tokens_ptr, - draft_tokens_stride, - num_spec, ): req_id = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + req_id) @@ -567,17 +562,6 @@ def _post_update_kernel( # Filter rows with negative index entries. return - if broadcast_drafts_ptr is not None: - # PP path: adopt the draft tokens proposed by the last rank's - # speculator so the next verification step embeds the real drafts. - for i in range(num_spec): - token_id = tl.load( - broadcast_drafts_ptr + req_id * broadcast_drafts_stride + i - ) - tl.store( - draft_tokens_ptr + req_state_idx * draft_tokens_stride + i, token_id - ) - total_len = tl.load(total_len_ptr + req_state_idx) num_sampled = tl.load(num_sampled_ptr + req_id) if num_sampled > 0: @@ -638,11 +622,6 @@ def post_update( all_token_ids: torch.Tensor, # [max_num_reqs] total_len: torch.Tensor, - # [num_reqs, num_spec]; drafts broadcast from the last PP rank. Only - # passed on non-last PP ranks, which never run the speculator. - broadcast_drafts: torch.Tensor | None = None, - # [max_num_reqs, num_spec] - draft_tokens_out: torch.Tensor | None = None, ) -> None: num_reqs = idx_mapping.shape[0] _post_update_kernel[(num_reqs,)]( @@ -659,11 +638,6 @@ def post_update( all_token_ids, all_token_ids.stride(0), total_len, - broadcast_drafts, - broadcast_drafts.stride(0) if broadcast_drafts is not None else 0, - draft_tokens_out, - draft_tokens_out.stride(0) if draft_tokens_out is not None else 0, - broadcast_drafts.shape[1] if broadcast_drafts is not None else 0, num_warps=1, ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 75c9c412cf67..5c617533f182 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -862,7 +862,6 @@ def _dummy_run( # dummy run the eagle speculator's propose to ensure DP/EP sync. if self.speculator is not None: assert self.sampler is not None - assert hidden_states is not None self.step_timing.drafter_start() mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None if self.speculator.supports_mm_inputs: @@ -878,8 +877,7 @@ def _dummy_run( spec_hidden_states = hidden_states if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - if pre_hc_hidden_states is not None: - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler._get_contexts(input_batch.idx_mapping), @@ -1111,44 +1109,6 @@ def update_pp_decode_requests(self): if outputs is not None: self.postprocess_sampled(**outputs) - def warmup_pp_decode_update(self) -> None: - """JIT-compile the kernel behind ``update_pp_decode_requests``. - - That path only runs on real steps, so the warmup steps never reach it - on non-last PP ranks. Its first triton compile must not happen - mid-serving: the in-flight sampled-token broadcast keeps a NCCL kernel - spinning on this device, which blocks the CUDA module load and - deadlocks the pipeline. An all -1 idx_mapping makes this a no-op. - The freshly allocated int32 tensors are 16-byte aligned, matching the - padded views `PPHandler` produces at serving time (triton specializes - on pointer alignment). - """ - assert self.pp_handler is not None - num_spec = self.pp_handler.max_sample_len - 1 - broadcast_drafts = ( - torch.zeros((1, num_spec), dtype=torch.int64, device=self.device) - if num_spec > 0 - else None - ) - post_update( - torch.full((1,), -1, dtype=torch.int64, device=self.device), - self.req_states.num_computed_tokens.gpu, - self.req_states.last_sampled_tokens, - None, - torch.zeros( - (1, self.pp_handler.max_sample_len), - dtype=torch.int64, - device=self.device, - ), - torch.zeros(1, dtype=torch.int32, device=self.device), - torch.zeros(1, dtype=torch.int32, device=self.device), - None, - self.req_states.all_token_ids.gpu, - self.req_states.total_len.gpu, - broadcast_drafts, - self.req_states.draft_tokens if broadcast_drafts is not None else None, - ) - def add_requests(self, scheduler_output: SchedulerOutput) -> None: for new_req_data in scheduler_output.scheduled_new_reqs: assert new_req_data.prefill_token_ids is not None @@ -1630,7 +1590,6 @@ def postprocess_sampled( num_sampled: torch.Tensor, num_rejected: torch.Tensor, query_start_loc: torch.Tensor | None = None, - broadcast_drafts: torch.Tensor | None = None, ) -> None: # Update the number of computed tokens. output_bin_counts = None @@ -1648,8 +1607,6 @@ def postprocess_sampled( query_start_loc, self.req_states.all_token_ids.gpu, self.req_states.total_len.gpu, - broadcast_drafts, - self.req_states.draft_tokens if broadcast_drafts is not None else None, ) self.model_state.postprocess_state( @@ -2170,10 +2127,7 @@ def sample_tokens( spec_hidden_states = draft_hidden_states if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - if pre_hc_hidden_states is not None: - spec_hidden_states = pre_hc_hidden_states[ - : draft_hidden_states.size(0) - ] + spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)] if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler._get_contexts(input_batch.idx_mapping), @@ -2195,16 +2149,8 @@ def sample_tokens( dp_sync=dp_sync, mm_inputs=mm_inputs, ) - if draft_tokens is not None: - self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - if self.pp_handler is not None: - # Earlier stages never run the speculator; ship the - # drafts so their next verification step embeds the real - # draft tokens instead of stale buffer contents. - self.pp_handler.broadcast_drafts( - self.req_states.draft_tokens, input_batch - ) - if draft_tokens is not None and self.adaptive_verification is not None: + self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens + if self.adaptive_verification is not None: self.adaptive_verification.record_confidences( self.speculator.draft_token_confidence_probs, input_batch ) @@ -2216,11 +2162,7 @@ def sample_tokens( input_batch, self.req_states.draft_tokens[input_batch.idx_mapping], ) - if self.pp_handler is not None and self.speculator is None: - # When a speculator ran, the propose() path above already - # broadcast the fresh drafts. Broadcasting here as well would - # double-post on the pp_broadcast group and misalign the - # recv FIFO on earlier stages, hanging the pipeline. + if self.pp_handler is not None: self.pp_handler.broadcast_drafts( self.req_states.draft_tokens, input_batch ) diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 862b86e5a107..c2aaa1fba25a 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -46,18 +46,6 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: return produces_sample if produces_sample.any() else None -def _alloc_combined(num_reqs: int, device: torch.device) -> torch.Tensor: - """Allocate the (2, N) int32 buffer broadcast alongside sampled tokens. - - The inner dim is padded to a multiple of 4 so that both `unbind` views - stay 16-byte aligned for any `num_reqs`: triton specializes on pointer - alignment, and a misaligned `num_rejected` would JIT-compile a second - `_post_update_kernel` variant at serving time. The padding is broadcast - but never read. Sender and receiver must both use this allocation. - """ - return torch.empty(2, -(-num_reqs // 4) * 4, dtype=torch.int32, device=device) - - class PPHandler: """Runs the PP sampled-token broadcast/recv on a side stream so the default stream isn't gated by the matching peer call. Step T's recv is @@ -98,14 +86,6 @@ def __init__( ) self.aux_hidden_state_relay_keys: tuple[str, ...] = () - # Warmup steps run the pipeline with synthetic batches whose outputs are - # discarded; the sampled-token broadcast is disabled there so its - # side-stream NCCL ops cannot overlap the next step's activation p2p. - self.disabled = False - - def set_disabled(self, disabled: bool) -> None: - self.disabled = disabled - def on_req_idx_freed(self, req_idx: int) -> None: self.req_idx_gen_np[req_idx] += 1 @@ -177,15 +157,29 @@ def get_prev_sampled_outputs( num_sampled=slot.num_sampled, num_rejected=slot.num_rejected, idx_mapping=idx_mapping, - broadcast_drafts=slot.draft_tokens, ) + def broadcast_drafts( + self, draft_tokens: torch.Tensor, input_batch: InputBatch + ) -> None: + """Broadcast draft proposals so non-last ranks can embed real token ids.""" + assert self.is_last_rank + if compute_need_sampled_mask(input_batch) is None: + return + with torch.cuda.stream(self.broadcast_stream): + self.broadcast_stream.wait_stream(self.main_stream) + send = draft_tokens[input_batch.idx_mapping].contiguous() + # Must record the idx_mapping tensor since it was allocated + # on the main stream. + input_batch.idx_mapping.record_stream(self.broadcast_stream) + torch.distributed.broadcast( + send, src=self.last_rank, group=self.broadcast_group + ) + def receive(self, input_batch: InputBatch) -> bool: """Returns True iff sampled tokens need to be gathered from *all* requests in the batch.""" assert not self.is_last_rank - if self.disabled: - return False need_sampled_mask = compute_need_sampled_mask(input_batch) if need_sampled_mask is None: # Leave this step's reserved slot as None. @@ -201,7 +195,7 @@ def receive(self, input_batch: InputBatch) -> bool: sampled_tokens = torch.empty( num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device ) - combined = _alloc_combined(num_reqs, self.device) + combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device) torch.distributed.broadcast( sampled_tokens, src=self.last_rank, group=self.broadcast_group ) @@ -219,7 +213,6 @@ def receive(self, input_batch: InputBatch) -> bool: torch.distributed.broadcast( draft_tokens, src=self.last_rank, group=self.broadcast_group ) - event = self.broadcast_stream.record_event() num_sampled, num_rejected = combined.unbind(dim=0) # Must record_stream since these were allocated on broadcast stream but @@ -249,16 +242,11 @@ def broadcast( input_batch: InputBatch, ) -> None: assert self.is_last_rank - if self.disabled: - return - mask = compute_need_sampled_mask(input_batch) - if mask is None: + if compute_need_sampled_mask(input_batch) is None: # No request needs sampled outputs for a subsequent decode step. return assert sampled_token_ids.dtype == torch.int64 - assert num_sampled.dtype == torch.int32 - assert num_rejected.dtype == torch.int32 if current_platform.is_xpu(): self.main_stream.synchronize() @@ -274,42 +262,9 @@ def broadcast( src=self.last_rank, group=self.broadcast_group, ) - combined = _alloc_combined(num_sampled.shape[0], self.device) - combined[0, : num_sampled.shape[0]] = num_sampled - combined[1, : num_sampled.shape[0]] = num_rejected + combined = torch.stack((num_sampled, num_rejected), dim=0) torch.distributed.broadcast( combined, src=self.last_rank, group=self.broadcast_group ) for tensor in (sampled_token_ids, num_sampled, num_rejected): tensor.record_stream(self.broadcast_stream) - - def broadcast_drafts( - self, draft_token_table: torch.Tensor, input_batch: InputBatch - ) -> None: - """Broadcast the speculator's freshly proposed draft tokens. - - Runs after propose() on the last rank; the send is stream-ordered - after broadcast()'s sends, matching receive()'s enqueue order. The - payload is gathered from the runner's draft table (just updated from - propose()'s output) into a fresh compact tensor: the speculator - overwrites its own persistent buffer on the next step, possibly - before this async send completes. - """ - assert self.is_last_rank - if self.disabled or self.max_sample_len == 1: - return - if compute_need_sampled_mask(input_batch) is None: - return - assert draft_token_table.dtype == torch.int64 - assert draft_token_table.shape[1] == self.max_sample_len - 1 - - # Gather on the main stream so the payload is ordered after this - # step's table update and before any later one. - drafts = draft_token_table[input_batch.idx_mapping] - assert drafts.shape[0] == input_batch.num_reqs - with torch.cuda.stream(self.broadcast_stream): - self.broadcast_stream.wait_stream(self.main_stream) - torch.distributed.broadcast( - drafts, src=self.last_rank, group=self.broadcast_group - ) - drafts.record_stream(self.broadcast_stream) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 3f661bbda575..a5710829e0ff 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -256,7 +256,6 @@ def _run_model( num_tokens_across_dp=num_tokens_across_dp, slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, - is_padding=self.input_buffers.is_padding[:num_tokens], ): last_hidden_states = self.model( input_ids=self.input_buffers.input_ids[:num_tokens], @@ -488,7 +487,6 @@ def _prepare_dflash_inputs_kernel( # Outputs out_input_ids_ptr, out_query_positions_ptr, - out_is_padding_ptr, out_query_start_loc_ptr, out_seq_lens_ptr, out_query_slot_mapping_ptr, @@ -619,7 +617,6 @@ def _prepare_dflash_inputs_kernel( ) tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query) - tl.store(out_is_padding_ptr + query_idx, False, mask=is_query) clamped_query_pos = tl.minimum(query_pos, max_model_len - 1) tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query) tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query) @@ -682,9 +679,6 @@ def _prepare_dflash_inputs_kernel( for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) mask = block < max_num_tokens - tl.store(out_input_ids_ptr + block, 0, mask=mask) - tl.store(out_query_positions_ptr + block, 0, mask=mask) - tl.store(out_is_padding_ptr + block, True, mask=mask) tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask) @@ -736,7 +730,6 @@ def prepare_dflash_inputs( _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)]( input_buffers.input_ids, input_buffers.positions, - input_buffers.is_padding, input_buffers.query_start_loc, input_buffers.seq_lens, query_slot_mapping, diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 2b7f6b6d494e..56f7d438fc3f 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -445,11 +445,6 @@ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: for step_indices, step_spec_flags in decode_steps: _run_decode_step(step_indices, step_spec_flags) - # The deferred PP post-update path only runs on real steps, so the steps - # above never JIT-compile its kernel on non-last ranks. - if not model_runner.is_last_pp_rank and model_runner.pp_handler is not None: - model_runner.warmup_pp_decode_update() - # Clean up - process finish_req_ids. cleanup_output = SchedulerOutput.make_empty() cleanup_output.finished_req_ids = set(req_ids) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 58af8211278e..3d4c5f3d5cc5 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -805,15 +805,6 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: @instrument(span_name="Warmup (GPU)") def compile_or_warm_up_model(self) -> CompilationTimes: - # All warmup phases below run synthetic steps whose sampled outputs are - # discarded. The PP sampled-token broadcast would carry no payload, and - # its side-stream NCCL ops can overlap the next step's activation p2p - # and deadlock the pipeline, so keep it disabled for the whole warmup - # window and restore it before serving. - pp_handler = getattr(self.model_runner, "pp_handler", None) - if pp_handler is not None: - pp_handler.set_disabled(True) - warmup_sizes: list[int] = [] if ( @@ -992,9 +983,6 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # gate so subsequent `execute_model` / `sample_tokens` calls enforce it. enable_gpu_sync_check() - if pp_handler is not None: - pp_handler.set_disabled(False) - return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time,