From 8ce137888fb1e60bfda82501678d9ba77a9817cb Mon Sep 17 00:00:00 2001 From: nilesjarvis Date: Sat, 22 Aug 2026 21:22:34 +0000 Subject: [PATCH 1/5] fix(pp): data-independent PP broadcast schedule with per-step seq validation Replace the data-dependent sampled-token broadcast protocol in PPHandler with a fixed four-broadcast-per-step schedule (header + sampled + counts + draft), a per-step sequence header validated on every receiver, and sentinel payloads for mask-None or skipped-proposal steps. The upstream protocol gated the send on the last rank's own mask while receivers posted the draft recv unconditionally, which FIFO-mismatched collectives and silently hung the pipeline. Receivers now raise RuntimeError with the expected vs actual header the moment a step is skipped cross-rank, instead of deadlocking the engine. Also import faulthandler in WorkerProc so SIGUSR1 dumps all thread stacks to the journal without ptrace (diagnostics parity). --- vllm/v1/executor/multiproc_executor.py | 5 + vllm/v1/worker/gpu/model_runner.py | 11 ++ vllm/v1/worker/gpu/pp_utils.py | 176 +++++++++++++++++-------- 3 files changed, 139 insertions(+), 53 deletions(-) diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index f0d51401cd08..74d1ac34c789 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -906,6 +906,11 @@ def signal_handler(signum, frame): ready_writer.close() ready_writer = None + # Wedge forensics: dump all thread stacks on SIGUSR1 without ptrace. + import faulthandler as _fh + import signal as _sig + _fh.register(_sig.SIGUSR1, all_threads=True) + worker.worker_busy_loop() except Exception: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 45a4e93eb000..a2534439a542 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1663,6 +1663,17 @@ def sample_tokens( self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens if self.pp_handler is not None and broadcast_drafts: self.pp_handler.broadcast_draft_tokens(draft_tokens) + elif self.pp_handler is not None and broadcast_drafts: + # Speculator absent or propose() skipped: send a sentinel draft + # payload so non-last ranks' unconditional draft recv still matches + # this step. Consumers filter via need_sampled_mask at consume time. + sentinel = torch.full( + (input_batch.num_reqs, self.pp_handler.max_draft_len), + -1, + dtype=torch.int64, + device=self.device, + ) + self.pp_handler.broadcast_draft_tokens(sentinel) if self.num_speculative_steps > 0: # Spec-decode and diffusion LLMs both use draft tokens but the latter does diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 570cd1d1e8fb..c3c011b68ffc 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -1,11 +1,43 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Apache-2.0 Section 4(b): modified by Lime Labs EOOD; upstream notices preserved. -# Original overlay modifications only; upstream ownership is not claimed. -"""Pipeline Parallelism utils for V2 Model Runner.""" +"""Pipeline Parallelism utils for V2 Model Runner. + +Fork of the LimeChain pp_utils with a validated broadcast protocol. + +The upstream protocol is data-dependent: each step optionally sends three +broadcasts (sampled tokens, combined num_sampled/num_rejected, and draft +tokens) only when the last rank's `compute_need_sampled_mask()` is non-None, +and the draft broadcast is additionally gated on a successful +`speculator.propose()`. Receivers, however, post all three recvs whenever +their own mask is non-None, unconditionally for the draft recv. NCCL +collectives match FIFO per communicator, so any cross-rank divergence in the +per-step schedule orphans the receivers permanently in spinning recv kernels +while the sender races ahead and blocks on CPU. No timeout or sequencing +exists anywhere in the protocol, so the first mismatch is a silent infinite +hang (today observed as PP8+DSpark wedges under long agentic traffic). + +This fork replaces that with a data-independent protocol: + +- Every step, every rank posts exactly four broadcasts on a dedicated + communicator: a 3-int wire header, sampled tokens, combined counts, and + draft tokens. There is no "skip when mask is None" asymmetry. +- The wire header carries a per-step monotonically increasing sequence number + (per PP group) and a CRC32 of the need-sampled mask, and receivers validate + it before touching any payload. Any desync raises `RuntimeError` naming the + expected vs actual header, failing fast instead of FIFO-matching the wrong + pair of collectives. +- When the step has no real sampled outputs (mask None on the last rank) or + the DSpark proposal failed/skipped, the last rank sends a sentinel payload + (zeroed counts / -1 draft tokens) so receivers still get a matching step. + Consumers filter sentinels out via their own need-sampled-mask state. + +Wire-format is intra-process only (all ranks load the same module), so the +change is compatible with any embedding as long as all ranks use this build. +""" from collections import deque from dataclasses import dataclass +import zlib import numpy as np import torch @@ -19,13 +51,7 @@ def _pad_sampled_tokens_for_broadcast( sampled_token_ids: torch.Tensor, max_sample_len: int ) -> torch.Tensor: - """Return a contiguous, fixed-width sampled-token wire tensor. - - Pipeline ranks must enter NCCL collectives with identical element counts. - Non-speculative steps can produce one token while receivers reserve the - speculative N+1 width, so pad unused positions with the standard -1 - placeholder. ``num_sampled`` remains authoritative for valid positions. - """ + """Return a contiguous, fixed-width sampled-token wire tensor.""" if sampled_token_ids.ndim != 2: raise ValueError( "sampled-token broadcast expects a rank-2 tensor, got " @@ -55,13 +81,7 @@ def _warm_up_broadcast_group( last_rank: int, is_last_rank: bool, ) -> None: - """Initialize the sampled-token NCCL communicator before KV allocation. - - NCCL process groups are lazy. If this communicator first initializes on - the first decode after a large fixed KV pool has been allocated, NCCL may - fail to reserve even its bookkeeping buffers. A synchronized one-element - broadcast makes that memory resident while model setup still has room. - """ + """Initialize the sampled-token NCCL communicator before KV allocation.""" with torch.cuda.stream(stream): marker = torch.zeros(1, dtype=torch.int32, device=device) if is_last_rank: @@ -77,17 +97,13 @@ class PendingRecv: """Per-step slot data for a deferred postprocess on the main stream.""" event: torch.cuda.Event - sampled_tokens: torch.Tensor # [num_reqs, max_sample_len] num_sampled: torch.Tensor # [num_reqs] num_rejected: torch.Tensor # [num_reqs] draft_tokens: torch.Tensor | None # [num_reqs, num_speculative_steps] idx_mapping: torch.Tensor # [num_reqs] idx_mapping_np: np.ndarray # [num_reqs] - # Records which rows need a deferred postprocess (bool). - need_sampled_mask: np.ndarray # [num_reqs] - # Snapshot of slot generation counters at receive time, used to - # detect requests aborted since then. + need_sampled_mask: np.ndarray | None # [num_reqs] (None => sentinel step) gen_at_receive_np: np.ndarray # [num_reqs] @@ -109,13 +125,10 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: 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 - consumed at step T+pp_size via `get_prev_sampled_outputs`. + """Runs the PP sampled-token broadcast/recv on a side stream. - Uses a dedicated NCCL communicator (sibling of the PP `device_group`) - for the broadcast so it does not serialize on the wire with the - inter-stage hidden-state p2p send/recv ops. + Protocol is data-independent (see module docstring): every step posts + exactly four broadcasts on every rank, preceded by a validation header. """ def __init__( @@ -142,6 +155,11 @@ def __init__( # between PP decodes. self.req_idx_gen_np = np.zeros(max_num_reqs, dtype=np.int32) + # Per-step monotonically increasing sequence number for the broadcast + # wire header. Both sender and receivers increment it once per + # sample_tokens step, in the same order they consume scheduler outputs. + self.step_seq: int = 0 + # Dedicated subgroup for the sampled-token broadcast. self.broadcast_group = get_pp_group().make_sibling_device_group( group_desc="pp_broadcast" @@ -157,6 +175,24 @@ def __init__( def on_req_idx_freed(self, req_idx: int) -> None: self.req_idx_gen_np[req_idx] += 1 + def _wire_header(self, mask_hash: int) -> torch.Tensor: + """Fixed-size broadcast header: [seq_low32, seq_high32, mask_hash].""" + seq = self.step_seq + return torch.tensor( + [seq & 0xFFFFFFFF, (seq >> 32) & 0xFFFFFFFF, mask_hash & 0xFFFFFFFF], + dtype=torch.int32, + device=self.device, + ) + + @staticmethod + def _mask_hash(need_sampled_mask: np.ndarray | None) -> int: + """Deterministic hash of the need-sampled mask for wire validation.""" + if need_sampled_mask is None: + return 0 + # np.packbits -> bytes -> zlib.crc32: deterministic across ranks for the + # same bool array, cheap, and not security-sensitive. + return zlib.crc32(np.packbits(need_sampled_mask).tobytes()) & 0xFFFFFFFF + def get_prev_sampled_outputs( self, ) -> tuple[dict[str, torch.Tensor], torch.Tensor | None, np.ndarray] | None: @@ -174,7 +210,12 @@ def get_prev_sampled_outputs( # Skip requests which did not need sampled output and/or those already # finished. The post_update kernel skips the -1 entries. freed = self.req_idx_gen_np[slot.idx_mapping_np] != slot.gen_at_receive_np - exclude_mask = freed | ~slot.need_sampled_mask + if slot.need_sampled_mask is None: + # Sentinel step: no request needed sampled outputs when this was + # received, so every row is excluded from the deferred postprocess. + exclude_mask = np.ones_like(freed) + else: + exclude_mask = freed | ~slot.need_sampled_mask idx_mapping = slot.idx_mapping idx_mapping_np = slot.idx_mapping_np if exclude_mask.any(): @@ -202,9 +243,7 @@ def receive(self, input_batch: InputBatch) -> bool: requests in the batch.""" assert not self.is_last_rank need_sampled_mask = compute_need_sampled_mask(input_batch) - if need_sampled_mask is None: - # Leave this step's reserved slot as None. - return False + mask_hash = self._mask_hash(need_sampled_mask) # Snapshot the per-slot generation counter so a later free of any of # these RequestStates request indices is detectable at consume time. @@ -213,6 +252,24 @@ def receive(self, input_batch: InputBatch) -> bool: num_reqs = input_batch.num_reqs with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) + header = torch.empty(3, dtype=torch.int32, device=self.device) + torch.distributed.broadcast( + header, src=self.last_rank, group=self.broadcast_group + ) + # Validate sender/receiver step agreement BEFORE consuming payloads. + # Copy to CPU: tiny (3 ints), and a mismatch must raise from here. + expected = self._wire_header(mask_hash) + expected_cpu = expected.cpu() + header_cpu = header.cpu() + if not torch.equal(header_cpu, expected_cpu): + raise RuntimeError( + "PP sampled-token broadcast desync: " + f"sender header={header_cpu.tolist()} " + f"receiver expected={expected_cpu.tolist()} " + f"(step_seq before increment={self.step_seq})" + ) + # Step matched: advance the sequence counter on every rank. + self.step_seq += 1 sampled_tokens = torch.empty( num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device ) @@ -223,27 +280,24 @@ def receive(self, input_batch: InputBatch) -> bool: torch.distributed.broadcast( combined, src=self.last_rank, group=self.broadcast_group ) - draft_tokens = None - if self.max_draft_len > 0: - draft_tokens = torch.empty( - num_reqs, - self.max_draft_len, - dtype=torch.int64, - device=self.device, - ) - torch.distributed.broadcast( - draft_tokens, - src=self.last_rank, - group=self.broadcast_group, - ) + draft_tokens = torch.empty( + num_reqs, + self.max_draft_len, + dtype=torch.int64, + device=self.device, + ) + 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 # later used on the main stream. sampled_tokens.record_stream(self.main_stream) combined.record_stream(self.main_stream) - if draft_tokens is not None: - draft_tokens.record_stream(self.main_stream) + draft_tokens.record_stream(self.main_stream) self.queue[-1] = PendingRecv( event, sampled_tokens, @@ -255,7 +309,9 @@ def receive(self, input_batch: InputBatch) -> bool: need_sampled_mask, gen_at_receive_np, ) - return bool(need_sampled_mask.all()) + # Draft tokens may be a sentinel when there was no real sampled output + # this step; consumers filter via need_sampled_mask at consume time. + return bool(need_sampled_mask is not None and need_sampled_mask.all()) def get_current_received_draft_tokens( self, @@ -274,9 +330,8 @@ def broadcast( input_batch: InputBatch, ) -> bool: assert self.is_last_rank - if compute_need_sampled_mask(input_batch) is None: - # No request needs sampled outputs for a subsequent decode step. - return False + need_sampled_mask = compute_need_sampled_mask(input_batch) + mask_hash = self._mask_hash(need_sampled_mask) assert sampled_token_ids.dtype == torch.int64 @@ -285,6 +340,14 @@ def broadcast( with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) + torch.distributed.broadcast( + self._wire_header(mask_hash), + src=self.last_rank, + group=self.broadcast_group, + ) + # Advance the sequence counter on every rank (sender included) so + # it stays in lockstep with receivers. + self.step_seq += 1 wire_sampled_token_ids = _pad_sampled_tokens_for_broadcast( sampled_token_ids, self.max_sample_len ) @@ -293,7 +356,14 @@ def broadcast( src=self.last_rank, group=self.broadcast_group, ) - combined = torch.stack((num_sampled, num_rejected), dim=0) + if need_sampled_mask is None: + # No request needs sampled outputs this step: send zeroed + # sentinel counts so receivers still get a matching broadcast. + combined = torch.zeros( + 2, input_batch.num_reqs, dtype=torch.int32, device=self.device + ) + else: + combined = torch.stack((num_sampled, num_rejected), dim=0) torch.distributed.broadcast( combined, src=self.last_rank, group=self.broadcast_group ) @@ -302,7 +372,7 @@ def broadcast( return True def broadcast_draft_tokens(self, draft_token_ids: torch.Tensor) -> None: - """Broadcast real post-proposal draft IDs to non-final PP ranks.""" + """Broadcast real post-proposal (or sentinel) draft IDs to non-last PP ranks.""" assert self.is_last_rank if self.max_draft_len == 0: return From fa93533ca4ebca38fd67e57bd839acbe92860ceb Mon Sep 17 00:00:00 2001 From: Niles Jarvis Date: Sat, 22 Aug 2026 22:16:35 +0000 Subject: [PATCH 2/5] PPHandler: harden broadcast wire header (int64 + num_reqs) Wire header fixes on top of the validated PP broadcast protocol: - int64 header instead of int32: mask_hash is an unsigned CRC32 (0..2^32-1); int32 torch.tensor raised RuntimeError on ~50% of distinct masks (values >= 2^31). 64-bit seq also removes the 32-bit split. - Include num_reqs in the header: payload shapes are locally determined per rank ([num_reqs, ...]) and crc32(packbits(mask)) collides across batch sizes (e.g. all-False masks), so a num_reqs mismatch could pass validation and post mismatched-shape payloads, hanging the NCCL FIFO -- the exact wedge class this protocol exists to fail fast on. - Sender increments step_seq before enqueueing the header so sender and receivers advance at the identical protocol point. - record_stream the recv-side header past the blocking .cpu() read. Adds tests/v1/worker/gpu/test_pp_utils_header.py (unittest, 11 tests: dtype/width, high-bit CRC32 roundtrip, seq placement, num_reqs disambiguation, validation equality semantics). --- tests/v1/worker/gpu/test_pp_utils_header.py | 122 ++++++++++++++++++++ vllm/v1/worker/gpu/pp_utils.py | 33 ++++-- 2 files changed, 144 insertions(+), 11 deletions(-) create mode 100755 tests/v1/worker/gpu/test_pp_utils_header.py diff --git a/tests/v1/worker/gpu/test_pp_utils_header.py b/tests/v1/worker/gpu/test_pp_utils_header.py new file mode 100755 index 000000000000..65e60a777bf9 --- /dev/null +++ b/tests/v1/worker/gpu/test_pp_utils_header.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +"""Unit tests for the PPHandler wire header hardening (WEDGE_FIX_GOAL). + +Covers the contract that defends observable behavior: +- header dtype/width: int64 so unsigned CRC32 mask hashes (>= 2**31) survive + (regression: int32 torch.tensor raised RuntimeError on high-bit hashes); +- header includes num_reqs so batch-size disagreement cannot pass validation + with a colliding mask hash and then hang mismatched-shape NCCL payloads; +- seq/mask/num_reqs land in the expected header slots; +- receive-side validation equality semantics behind the fail-fast raise. + +Run: /bin/python tests/v1/worker/gpu/test_pp_utils_header.py +(plain unittest: pytest is not installed in the serving venv) +""" +import os +import sys +import types +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))) + +import numpy as np +import torch + +# Stub heavy vllm imports that pp_utils pulls at module import time. The wire +# helpers are pure tensor code; distributed/platform modules become fakes. +_fake_parallel_state = types.ModuleType("vllm.distributed.parallel_state") +_fake_parallel_state.get_pp_group = lambda: None +_fake_distributed = types.ModuleType("vllm.distributed") +_fake_distributed.parallel_state = _fake_parallel_state +sys.modules.setdefault("vllm.distributed", _fake_distributed) +sys.modules.setdefault("vllm.distributed.parallel_state", _fake_parallel_state) + +from vllm.v1.worker.gpu import pp_utils # noqa: E402 + + +def make_handler(): + """PPHandler without __init__: the wire helpers need only step_seq.""" + h = pp_utils.PPHandler.__new__(pp_utils.PPHandler) + h.step_seq = 0 + h.device = torch.device("cpu") + return h + + +class TestMaskHash(unittest.TestCase): + def test_none_mask_hashes_to_zero(self): + self.assertEqual(pp_utils.PPHandler._mask_hash(None), 0) + + def test_same_mask_same_hash(self): + m = np.array([True, False, True, True] * 8) + self.assertEqual(pp_utils.PPHandler._mask_hash(m), + pp_utils.PPHandler._mask_hash(m.copy())) + + def test_different_batch_sizes_diverge_in_full_key(self): + # CRC32(packbits(mask)) can collide across lengths; num_reqs in the + # header is what disambiguates. Verify the full key separates them. + m4 = np.array([False] * 4) + m8 = np.array([False] * 8) + h = make_handler() + self.assertFalse(torch.equal( + h._wire_header(h._mask_hash(m4), 4), + h._wire_header(h._mask_hash(m8), 8))) + + +class TestWireHeader(unittest.TestCase): + def test_int64_dtype(self): + hdr = make_handler()._wire_header(0, 1) + self.assertEqual(hdr.dtype, torch.int64, + "int32 overflows unsigned CRC32 hashes") + self.assertEqual(hdr.numel(), 3) + + def test_high_bit_mask_hash_roundtrips(self): + high = 0x80000000 # 2**31: raised RuntimeError under int32 + hdr = make_handler()._wire_header(high, 7) + self.assertEqual(int(hdr[1]), high) + self.assertGreater(int(hdr[1]), 0, "mask hash must stay unsigned") + + def test_max_crc32_roundtrips(self): + hdr = make_handler()._wire_header(0xFFFFFFFF, 12) + self.assertEqual(int(hdr[1]), 0xFFFFFFFF) + + def test_seq_and_num_reqs_placement(self): + h = make_handler() + h.step_seq = (1 << 40) + 12345 # beyond 32 bits: int64 must carry it + hdr = h._wire_header(42, 9) + self.assertEqual(int(hdr[0]), h.step_seq) + self.assertEqual(int(hdr[1]), 42) + self.assertEqual(int(hdr[2]), 9) + + def test_every_crc32_value_representable(self): + h = make_handler() + rng = np.random.default_rng(0) + for v in rng.integers(0, 2**32, size=256, dtype=np.uint64): + hdr = h._wire_header(int(v), 1) + self.assertEqual(int(hdr[1]), int(v)) + + +class TestHeaderValidationSemantics(unittest.TestCase): + """Equality semantics behind the receive-side fail-fast desync raise.""" + + def test_equal_headers_match(self): + h = make_handler() + h.step_seq = 99 + self.assertTrue(torch.equal(h._wire_header(7, 4).cpu(), + h._wire_header(7, 4).cpu())) + + def test_mask_mismatch_detected(self): + h = make_handler() + h.step_seq = 99 + self.assertFalse(torch.equal(h._wire_header(7, 4).cpu(), + h._wire_header(8, 4).cpu())) + + def test_num_reqs_mismatch_detected(self): + h = make_handler() + h.step_seq = 99 + self.assertFalse(torch.equal(h._wire_header(7, 4).cpu(), + h._wire_header(7, 5).cpu())) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index c3c011b68ffc..ca52e4e13957 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -175,12 +175,18 @@ def __init__( def on_req_idx_freed(self, req_idx: int) -> None: self.req_idx_gen_np[req_idx] += 1 - def _wire_header(self, mask_hash: int) -> torch.Tensor: - """Fixed-size broadcast header: [seq_low32, seq_high32, mask_hash].""" - seq = self.step_seq + def _wire_header(self, mask_hash: int, num_reqs: int) -> torch.Tensor: + """Fixed-size broadcast header: [seq, mask_hash, num_reqs] as int64. + + int64 throughout: mask_hash is an unsigned CRC32 (0..2^32-1), which + overflows int32 for high-bit values. num_reqs is included because the + payload shapes are locally determined per rank ([num_reqs, ...]); a + batch-size disagreement with a colliding mask hash would otherwise + pass validation and post mismatched-shape payloads, hanging NCCL. + """ return torch.tensor( - [seq & 0xFFFFFFFF, (seq >> 32) & 0xFFFFFFFF, mask_hash & 0xFFFFFFFF], - dtype=torch.int32, + [self.step_seq, mask_hash, num_reqs], + dtype=torch.int64, device=self.device, ) @@ -252,15 +258,18 @@ def receive(self, input_batch: InputBatch) -> bool: num_reqs = input_batch.num_reqs with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) - header = torch.empty(3, dtype=torch.int32, device=self.device) + header = torch.empty(3, dtype=torch.int64, device=self.device) torch.distributed.broadcast( header, src=self.last_rank, group=self.broadcast_group ) # Validate sender/receiver step agreement BEFORE consuming payloads. # Copy to CPU: tiny (3 ints), and a mismatch must raise from here. - expected = self._wire_header(mask_hash) + expected = self._wire_header(mask_hash, num_reqs) expected_cpu = expected.cpu() header_cpu = header.cpu() + # The blocking .cpu() syncs the broadcast stream, so NCCL is done + # reading `header`; keep it alive past the host read anyway. + header.record_stream(self.main_stream) if not torch.equal(header_cpu, expected_cpu): raise RuntimeError( "PP sampled-token broadcast desync: " @@ -270,6 +279,7 @@ def receive(self, input_batch: InputBatch) -> bool: ) # Step matched: advance the sequence counter on every rank. self.step_seq += 1 + del expected_cpu, header_cpu sampled_tokens = torch.empty( num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device ) @@ -340,14 +350,15 @@ def broadcast( with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) + # Capture the header BEFORE incrementing so sender and receivers + # advance the counter at the same protocol point (at send/post). + header = self._wire_header(mask_hash, input_batch.num_reqs) + self.step_seq += 1 torch.distributed.broadcast( - self._wire_header(mask_hash), + header, src=self.last_rank, group=self.broadcast_group, ) - # Advance the sequence counter on every rank (sender included) so - # it stays in lockstep with receivers. - self.step_seq += 1 wire_sampled_token_ids = _pad_sampled_tokens_for_broadcast( sampled_token_ids, self.max_sample_len ) From 1900969c057acc2be8e5f9446c0d72defce0eee5 Mon Sep 17 00:00:00 2001 From: Niles Jarvis Date: Sun, 23 Aug 2026 01:12:36 +0000 Subject: [PATCH 3/5] config: opt-in cap of V2 async in-flight batches at pp_size for PP+DSpark With async scheduling + V2 runner, max_concurrent_batches = pp_size+1. The extra in-flight batch lets a sample_tokens(N+1) RPC race ahead of the last PP rank's sample_tokens(N): flight-recorder dumps show non-last ranks one draft-broadcast sequence ahead (enq N+1) while the last rank is still at N, deadlocking the pp_broadcast group (60s NCCL watchdog, EngineDead). Guarded opt-in (VLLM_WEDGE_AB_CAP_CONCURRENT=1, pp_size>1, method=dspark): cap at pp_size, keeping async_scheduling=True and the V2 PP token path. Default behavior unchanged for every other configuration. Validated: previously-reliable ~60s wedge trigger (max_tokens overflow at temp 1.0) passed 3x consecutively, then a 66-minute sustained stress pass (200K-token prefills, 2-4 concurrent tool-schema streams, 25 overflow rounds) with zero wedges/watchdog events; baseline boots wedged within minutes on every attempt. --- vllm/config/vllm.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index d6e55798cae8..bc2c7b11b78d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -543,6 +543,22 @@ def max_concurrent_batches(self) -> int: pp_size = self.parallel_config.pipeline_parallel_size if self.scheduler_config.async_scheduling: if self.use_v2_model_runner: + # WEDGE A/B (PP8+DSpark investigation), opt-in via env guard: + # cap at pp_size instead of pp_size+1 ONLY for PP>1 + DSpark. + # The +1 admits an extra in-flight batch whose sample_tokens + # RPC can run ahead of the last rank's, racing receivers one + # step ahead on the pp_broadcast group (observed: receivers at + # draft seq N+1 while PP7 still at N). Keeps + # async_scheduling=True and the V2 PP token-propagation path. + # Not committed as default behavior. + import os + if ( + os.environ.get("VLLM_WEDGE_AB_CAP_CONCURRENT", "0") == "1" + and pp_size > 1 + and self.speculative_config is not None + and self.speculative_config.method == "dspark" + ): + return pp_size return pp_size + 1 # V1 Model Runner does not fully support async scheduling with PP. if pp_size <= 1: From 9dc0b8c7be273bdbb170627c3752f7d9168f0fe5 Mon Sep 17 00:00:00 2001 From: Niles Jarvis Date: Sun, 23 Aug 2026 01:21:17 +0000 Subject: [PATCH 4/5] tests: update PP receive test for header protocol; add cap config test - test_pp_receive_posts_sample_metadata_and_real_draft_broadcasts: fill the int64 validation header the way the sender does; assert 4 broadcasts [(3,),(1,6),(2,1),(1,5)], step_seq advance, and the fail-fast desync raise on a stale header. - tests/config/test_max_concurrent_batches_cap.py: 8 cases proving the in-flight cap applies only with VLLM_WEDGE_AB_CAP_CONCURRENT=1 + pp_size>1 + method=dspark, and that PP=1 async (2), non-dspark (9), env-off (9), V1-runner and non-async values are unchanged. --- .../config/test_max_concurrent_batches_cap.py | 86 +++++++++++++++++++ tests/models/test_deepseek_v4_dspark_pp.py | 36 ++++++-- 2 files changed, 116 insertions(+), 6 deletions(-) create mode 100755 tests/config/test_max_concurrent_batches_cap.py diff --git a/tests/config/test_max_concurrent_batches_cap.py b/tests/config/test_max_concurrent_batches_cap.py new file mode 100755 index 000000000000..f33e89c0d3ea --- /dev/null +++ b/tests/config/test_max_concurrent_batches_cap.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Focused test for the gated max_concurrent_batches cap (WEDGE A/B). + +Contract: with async scheduling + V2 runner, the cap to pp_size applies ONLY +when VLLM_WEDGE_AB_CAP_CONCURRENT=1 AND pp_size>1 AND speculative method is +dspark. Every other combination preserves pp_size+1 (or the V1/non-async +values). Plain unittest: pytest is not installed in the serving venv. + +Run: /bin/python tests/config/test_max_concurrent_batches_cap.py +""" +import os +import sys +import unittest +from types import SimpleNamespace +from unittest import mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname( + os.path.abspath(__file__))))) + +from vllm.config.vllm import VllmConfig # noqa: E402 + + +def make_cfg(pp_size, async_sched=True, use_v2=True, method="dspark"): + cfg = VllmConfig.__new__(VllmConfig) + object.__setattr__(cfg, "parallel_config", + SimpleNamespace(pipeline_parallel_size=pp_size)) + object.__setattr__(cfg, "scheduler_config", + SimpleNamespace(async_scheduling=async_sched)) + object.__setattr__( + cfg, "speculative_config", + SimpleNamespace(method=method) if method is not None else None) + # use_v2_model_runner may be a property on the class; bypass via __dict__ + # patching where possible, else patch the type attribute in tests. + return cfg, use_v2 + + +class TestCap(unittest.TestCase): + def _mcb(self, cfg, use_v2): + with mock.patch.object(type(cfg), "use_v2_model_runner", + new_callable=mock.PropertyMock, + return_value=use_v2): + return cfg.max_concurrent_batches + + def test_env_on_pp8_dspark_capped(self): + cfg, v2 = make_cfg(8) + with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}): + self.assertEqual(self._mcb(cfg, v2), 8) + + def test_env_off_pp8_dspark_uncapped(self): + cfg, v2 = make_cfg(8) + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop("VLLM_WEDGE_AB_CAP_CONCURRENT", None) + self.assertEqual(self._mcb(cfg, v2), 9) + + def test_env_on_pp1_uncapped(self): + # PP=1 must keep pp_size+1=2 (async overlap), even with env on. + cfg, v2 = make_cfg(1) + with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}): + self.assertEqual(self._mcb(cfg, v2), 2) + + def test_env_on_non_dspark_uncapped(self): + cfg, v2 = make_cfg(8, method="eagle") + with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}): + self.assertEqual(self._mcb(cfg, v2), 9) + + def test_env_on_no_spec_config_uncapped(self): + cfg, v2 = make_cfg(8, method=None) + with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}): + self.assertEqual(self._mcb(cfg, v2), 9) + + def test_v1_runner_pp_gt1_returns_pp_size(self): + cfg, v2 = make_cfg(8, use_v2=False) + self.assertEqual(self._mcb(cfg, v2), 8) + + def test_v1_runner_pp1_async_returns_2(self): + cfg, v2 = make_cfg(1, use_v2=False) + self.assertEqual(self._mcb(cfg, v2), 2) + + def test_no_async_returns_pp_size(self): + cfg, v2 = make_cfg(8, async_sched=False) + with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}): + self.assertEqual(self._mcb(cfg, v2), 8) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/models/test_deepseek_v4_dspark_pp.py b/tests/models/test_deepseek_v4_dspark_pp.py index e2b4a006a0ba..27c056351930 100644 --- a/tests/models/test_deepseek_v4_dspark_pp.py +++ b/tests/models/test_deepseek_v4_dspark_pp.py @@ -319,17 +319,28 @@ def make_sibling_device_group(self, *, group_desc: str): "compute_need_sampled_mask", lambda _batch: np.array([True]), ) - monkeypatch.setattr( - pp_utils.torch.distributed, - "broadcast", - lambda tensor, **_kwargs: broadcast_shapes.append(tuple(tensor.shape)), - ) + def fake_broadcast(tensor, **_kwargs): + broadcast_shapes.append(tuple(tensor.shape)) + if tensor.numel() == 3 and tensor.dtype == torch.int64: + # Fill the validation header exactly as the sender would: the + # receiver validates [step_seq, mask_hash, num_reqs] before + # consuming payloads (fa93533 wire hardening). + mask_hash = pp_utils.PPHandler._mask_hash(np.array([True])) + tensor.copy_( + torch.tensor( + [handler_ref[0].step_seq, mask_hash, 1], dtype=torch.int64 + ) + ) + monkeypatch.setattr(pp_utils.torch.distributed, "broadcast", fake_broadcast) + + handler_ref: list = [] handler = pp_utils.PPHandler( max_num_reqs=1, num_speculative_steps=5, device=torch.device("cpu"), ) + handler_ref.append(handler) input_batch = SimpleNamespace( idx_mapping=torch.tensor([0]), idx_mapping_np=np.array([0]), @@ -337,7 +348,20 @@ def make_sibling_device_group(self, *, group_desc: str): ) assert handler.receive(input_batch) - assert broadcast_shapes == [(1, 6), (2, 1), (1, 5)] + # Header broadcast [3] precedes the three payload broadcasts. + assert broadcast_shapes == [(3,), (1, 6), (2, 1), (1, 5)] + assert handler.step_seq == 1 # advanced after successful validation + + # A header that disagrees (stale seq) must raise, not hang. + def stale_broadcast(tensor, **_kwargs): + if tensor.numel() == 3 and tensor.dtype == torch.int64: + tensor.copy_(torch.tensor([99, 0, 1], dtype=torch.int64)) + + monkeypatch.setattr(pp_utils.torch.distributed, "broadcast", stale_broadcast) + with pytest.raises(RuntimeError, match="desync"): + handler.receive(input_batch) + + monkeypatch.setattr(pp_utils.torch.distributed, "broadcast", fake_broadcast) received = handler.get_current_received_draft_tokens() assert received is not None received_drafts, ready_event = received From 6ea9ae636e176bf606a4287e251ffeb55c56cf8a Mon Sep 17 00:00:00 2001 From: nilesjarvis Date: Sun, 23 Aug 2026 12:33:32 +0000 Subject: [PATCH 5/5] fix(pp): add stable PP+DSpark fallback Add an opt-in dedicated Gloo transport for the tiny sampled-token protocol, defer draft receives until consumption, and allow DSpark draft-only eager execution while retaining target CUDA graphs. --- tests/distributed/test_comm_ops.py | 32 ++ tests/models/test_deepseek_v4_dspark_pp.py | 364 ++++++++++++++++-- tests/v1/worker/gpu/test_pp_utils_header.py | 110 +++++- vllm/distributed/parallel_state.py | 25 ++ vllm/v1/worker/gpu/model_runner.py | 6 +- vllm/v1/worker/gpu/pp_utils.py | 277 ++++++++++--- .../gpu/spec_decode/dflash/speculator.py | 9 + .../gpu/spec_decode/dspark/speculator.py | 36 ++ vllm/v1/worker/gpu/spec_decode/utils.py | 78 +++- 9 files changed, 841 insertions(+), 96 deletions(-) diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index 19a095c7c6f7..4556f58b60b7 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -238,6 +238,38 @@ def _make_group_for_unit_test( return g +def test_make_sibling_cpu_group_is_distinct_and_preserves_membership( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.distributed import utils as distributed_utils + + calls: list[tuple[list[int], str, str | None, object]] = [] + + def fake_new_group( + ranks: list[int], *, backend: str, group_desc: str | None, timeout: object + ) -> str: + calls.append((ranks, backend, group_desc, timeout)) + return f"group-{len(calls)}" + + timeout = object() + monkeypatch.setattr(torch.distributed, "new_group", fake_new_group) + monkeypatch.setattr( + distributed_utils, "get_cpu_distributed_timeout_or_none", lambda: timeout + ) + + group = _make_group_for_unit_test(rank_in_group=0, world_size=2) + group.rank = 2 + group.group_ranks = [[0, 2], [1, 3]] + + sibling = group.make_sibling_cpu_group(group_desc="control") + + assert sibling == "group-1" + assert calls == [ + ([0, 2], "gloo", "control", timeout), + ([1, 3], "gloo", "control", timeout), + ] + + def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/models/test_deepseek_v4_dspark_pp.py b/tests/models/test_deepseek_v4_dspark_pp.py index 27c056351930..a6ac27b06900 100644 --- a/tests/models/test_deepseek_v4_dspark_pp.py +++ b/tests/models/test_deepseek_v4_dspark_pp.py @@ -6,6 +6,7 @@ from collections import deque from types import SimpleNamespace from typing import cast +from unittest import mock import numpy as np import pytest @@ -222,6 +223,7 @@ def test_sparse_mla_startup_warmup_skip_is_dspark_pp_only( from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( _should_skip_dspark_pp_sparse_mla_warmup, ) + runner = SimpleNamespace( vllm_config=SimpleNamespace( use_v2_model_runner=use_v2, @@ -259,6 +261,35 @@ def test_scheduler_realistic_kernel_warmup_skip_is_dspark_pp_only( assert _should_skip_dspark_pp_kernel_warmup(runner) is expected +@pytest.mark.parametrize( + ("force_eager", "trace_stages", "expected_mode"), + [ + ("0", "0", "FULL_DECODE_ONLY"), + ("1", "0", "NONE"), + ("0", "1", "NONE"), + ], +) +def test_dspark_draft_only_eager_diagnostic_gate( + monkeypatch, + force_eager: str, + trace_stages: str, + expected_mode: str, +): + from vllm.config.compilation import CUDAGraphMode + from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator + from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + + monkeypatch.setenv("VLLM_WEDGE_AB_DSPARK_EAGER", force_eager) + monkeypatch.setenv("VLLM_WEDGE_TRACE_DSPARK_STAGES", trace_stages) + speculator = DSparkSpeculator.__new__(DSparkSpeculator) + requested_mode = CUDAGraphMode.FULL_DECODE_ONLY + + with mock.patch.object(DFlashSpeculator, "init_cudagraph_manager") as parent_init: + speculator.init_cudagraph_manager(requested_mode) + + parent_init.assert_called_once_with(CUDAGraphMode[expected_mode]) + + def test_pp_sampled_broadcast_pads_to_fixed_width(): from vllm.v1.worker.gpu.pp_utils import _pad_sampled_tokens_for_broadcast @@ -278,19 +309,38 @@ def test_pp_sampled_broadcast_pads_to_fixed_width(): ) +@pytest.mark.parametrize("cpu_mode", [False, True]) def test_pp_receive_posts_sample_metadata_and_real_draft_broadcasts( - monkeypatch, + monkeypatch, cpu_mode: bool ) -> None: from vllm.v1.worker.gpu import pp_utils broadcast_shapes: list[tuple[int, ...]] = [] + broadcast_async: list[bool] = [] + + ready_event = object() + + class FakeWork: + def __init__(self) -> None: + self.wait_calls = 0 + + def wait(self) -> None: + self.wait_calls += 1 + + draft_work = FakeWork() class FakeStream: + def __init__(self, name: str) -> None: + self.name = name + self.waited_for: list[FakeStream] = [] + self.record_event_calls = 0 + def wait_stream(self, _stream) -> None: - pass + self.waited_for.append(_stream) def record_event(self): - return object() + self.record_event_calls += 1 + return ready_event class StreamContext: def __enter__(self): @@ -306,31 +356,47 @@ class FakePPGroup: def make_sibling_device_group(self, *, group_desc: str): assert group_desc == "pp_broadcast" - return "sample-group" + return "device-sample-group" + def make_sibling_cpu_group(self, *, group_desc: str): + assert group_desc == "pp_broadcast_cpu" + return "cpu-sample-group" + + main_stream = FakeStream("main") + side_stream = FakeStream("side") + monkeypatch.setenv( + "VLLM_WEDGE_AB_PP_BROADCAST_CPU", + "1" if cpu_mode else "0", + ) monkeypatch.setattr(pp_utils, "get_pp_group", lambda: FakePPGroup()) - monkeypatch.setattr(pp_utils.torch.cuda, "current_stream", lambda _device: object()) - monkeypatch.setattr(pp_utils.torch.cuda, "Stream", lambda _device: FakeStream()) + monkeypatch.setattr( + pp_utils.torch.cuda, "current_stream", lambda _device: main_stream + ) + monkeypatch.setattr(pp_utils.torch.cuda, "Stream", lambda _device: side_stream) monkeypatch.setattr(pp_utils.torch.cuda, "stream", lambda _stream: StreamContext()) monkeypatch.setattr(torch.Tensor, "record_stream", lambda _tensor, _stream: None) monkeypatch.setattr(pp_utils, "_warm_up_broadcast_group", lambda **_kwargs: None) + monkeypatch.setattr( + pp_utils, "_warm_up_cpu_broadcast_group", lambda **_kwargs: None + ) monkeypatch.setattr( pp_utils, "compute_need_sampled_mask", lambda _batch: np.array([True]), ) - def fake_broadcast(tensor, **_kwargs): + + def fake_broadcast(tensor, **kwargs): broadcast_shapes.append(tuple(tensor.shape)) + broadcast_async.append(kwargs.get("async_op", False)) if tensor.numel() == 3 and tensor.dtype == torch.int64: # Fill the validation header exactly as the sender would: the # receiver validates [step_seq, mask_hash, num_reqs] before # consuming payloads (fa93533 wire hardening). mask_hash = pp_utils.PPHandler._mask_hash(np.array([True])) tensor.copy_( - torch.tensor( - [handler_ref[0].step_seq, mask_hash, 1], dtype=torch.int64 - ) + torch.tensor([handler_ref[0].step_seq, mask_hash, 1], dtype=torch.int64) ) + return draft_work if kwargs.get("async_op", False) else None monkeypatch.setattr(pp_utils.torch.distributed, "broadcast", fake_broadcast) @@ -348,8 +414,17 @@ def fake_broadcast(tensor, **_kwargs): ) assert handler.receive(input_batch) + assert handler.broadcast_stream is side_stream + assert handler.broadcast_group == ( + "cpu-sample-group" if cpu_mode else "device-sample-group" + ) + assert side_stream.waited_for == ([] if cpu_mode else [main_stream]) + assert main_stream.record_event_calls == 0 + assert side_stream.record_event_calls == (0 if cpu_mode else 1) # Header broadcast [3] precedes the three payload broadcasts. assert broadcast_shapes == [(3,), (1, 6), (2, 1), (1, 5)] + assert broadcast_async == ([False, False, False, True] if cpu_mode else [False] * 4) + assert draft_work.wait_calls == 0 assert handler.step_seq == 1 # advanced after successful validation # A header that disagrees (stale seq) must raise, not hang. @@ -364,9 +439,19 @@ def stale_broadcast(tensor, **_kwargs): monkeypatch.setattr(pp_utils.torch.distributed, "broadcast", fake_broadcast) received = handler.get_current_received_draft_tokens() assert received is not None - received_drafts, ready_event = received + received_drafts, received_event, cpu_recv = received assert received_drafts.shape == (1, 5) - assert ready_event is handler.queue[-1].event + assert received_event is handler.queue[-1].event + assert (cpu_recv is not None) is cpu_mode + + if cpu_mode: + slot = handler.queue[-1] + assert slot is not None + handler._materialize_cpu_recv(slot) + assert draft_work.wait_calls == 1 + assert slot.cpu_recv is None + assert slot.event is ready_event + assert main_stream.record_event_calls == 1 def test_pp_received_drafts_wait_for_broadcast_before_cpu_copy(monkeypatch) -> None: @@ -414,12 +499,8 @@ def __exit__(self, *_args): req_ids=["ordinary", "req"], has_structured_output_reqs=True, ) - drafts = torch.tensor( - [[1, 2, 3, 4, 5], [11, 12, 13, 14, 15]], dtype=torch.int64 - ) - copied = np.array( - [[1, 2, 3, 4, 5], [11, 12, 13, 14, 15]], dtype=np.int64 - ) + drafts = torch.tensor([[1, 2, 3, 4, 5], [11, 12, 13, 14, 15]], dtype=torch.int64) + copied = np.array([[1, 2, 3, 4, 5], [11, 12, 13, 14, 15]], dtype=np.int64) monkeypatch.setattr(spec_utils, "async_copy_to_np", lambda _drafts: copied) monkeypatch.setattr( @@ -427,7 +508,9 @@ def __exit__(self, *_args): "Event", lambda **_kwargs: FakeCopyEvent(), ) - monkeypatch.setattr(spec_utils.torch.cuda, "stream", lambda _stream: StreamContext()) + monkeypatch.setattr( + spec_utils.torch.cuda, "stream", lambda _stream: StreamContext() + ) monkeypatch.setattr(torch.Tensor, "record_stream", lambda _tensor, _stream: None) spec_utils.DraftTokensHandler.set_draft_tokens( @@ -466,6 +549,109 @@ def __exit__(self, *_args): assert queued_event.synchronized +def test_pp_cpu_drafts_defer_gloo_wait_until_structured_retrieval() -> None: + from vllm.v1.worker.gpu.spec_decode import utils as spec_utils + + class FakeWaitable: + def __init__(self) -> None: + self.wait_calls = 0 + + def wait(self) -> None: + self.wait_calls += 1 + + ready = FakeWaitable() + handler = cast( + spec_utils.DraftTokensHandler, + SimpleNamespace( + req_ids=[], + draft_tokens_np=None, + num_draft_tokens=0, + pending_structured_drafts=deque(), + ), + ) + input_batch = SimpleNamespace( + req_ids=["ordinary", "structured"], + has_structured_output_reqs=True, + ) + drafts = torch.tensor([[1, 2], [11, 12]], dtype=torch.int64) + + spec_utils.DraftTokensHandler.set_draft_tokens( + handler, + input_batch, + drafts, + ready_waitable=ready, + structured_req_ids={"structured"}, + ) + + assert ready.wait_calls == 0 + queued = handler.pending_structured_drafts[0] + assert isinstance(queued[2], spec_utils.DeferredCpuDrafts) + assert queued[3] is None + + result = spec_utils.DraftTokensHandler.get_draft_tokens(handler, {"structured"}) + assert ready.wait_calls == 1 + assert result is not None + assert result.req_ids == ["structured"] + assert result.draft_token_ids == [[11, 12]] + + +def test_pp_cpu_draft_retrieval_does_not_wait_for_unrelated_batch() -> None: + from vllm.v1.worker.gpu.spec_decode import utils as spec_utils + + class FakeWaitable: + def __init__(self) -> None: + self.wait_calls = 0 + + def wait(self) -> None: + self.wait_calls += 1 + + first_ready = FakeWaitable() + unrelated_ready = FakeWaitable() + handler = cast( + spec_utils.DraftTokensHandler, + SimpleNamespace( + req_ids=[], + draft_tokens_np=None, + num_draft_tokens=2, + pending_structured_drafts=deque( + [ + ( + ["wanted"], + [0], + spec_utils.DeferredCpuDrafts( + torch.tensor([[1, 2]], dtype=torch.int64), first_ready + ), + None, + ), + ( + ["unrelated"], + [0], + spec_utils.DeferredCpuDrafts( + torch.tensor([[3, 4]], dtype=torch.int64), + unrelated_ready, + ), + None, + ), + ] + ), + ), + ) + + wanted = spec_utils.DraftTokensHandler.get_draft_tokens(handler, {"wanted"}) + + assert wanted is not None + assert wanted.req_ids == ["wanted"] + assert wanted.draft_token_ids == [[1, 2]] + assert first_ready.wait_calls == 1 + assert unrelated_ready.wait_calls == 0 + assert len(handler.pending_structured_drafts) == 1 + + unrelated = spec_utils.DraftTokensHandler.get_draft_tokens(handler, {"unrelated"}) + assert unrelated is not None + assert unrelated.draft_token_ids == [[3, 4]] + assert unrelated_ready.wait_calls == 1 + + def test_structured_draft_queue_drains_stale_batches_and_keeps_latest() -> None: from vllm.v1.worker.gpu.spec_decode import utils as spec_utils @@ -612,7 +798,8 @@ def synchronize(self) -> None: assert rows.tolist() == [[1, 2], [3, 4]] -def test_pp_draft_sender_orders_buffer_reuse_after_broadcast(monkeypatch) -> None: +@pytest.mark.parametrize("cpu_mode", [False, True]) +def test_pp_draft_sender_uses_selected_transport(monkeypatch, cpu_mode: bool) -> None: from vllm.v1.worker.gpu import pp_utils class FakeStream: @@ -638,6 +825,7 @@ def __exit__(self, *_args): SimpleNamespace( is_last_rank=True, max_draft_len=5, + broadcast_on_cpu=cpu_mode, main_stream=main_stream, broadcast_stream=broadcast_stream, last_rank=7, @@ -656,8 +844,68 @@ def __exit__(self, *_args): pp_utils.PPHandler.broadcast_draft_tokens(handler, drafts) assert sent == [drafts] - assert broadcast_stream.waited_for == [main_stream] - assert main_stream.waited_for == [broadcast_stream] + assert broadcast_stream.waited_for == ([] if cpu_mode else [main_stream]) + assert main_stream.waited_for == ([] if cpu_mode else [broadcast_stream]) + + +@pytest.mark.parametrize("cpu_mode", [False, True]) +def test_pp_sample_sender_uses_selected_transport(monkeypatch, cpu_mode: bool) -> None: + from vllm.v1.worker.gpu import pp_utils + + class FakeStream: + def __init__(self) -> None: + self.waited_for: list[FakeStream] = [] + + def wait_stream(self, stream: "FakeStream") -> None: + self.waited_for.append(stream) + + class StreamContext: + def __enter__(self): + return None + + def __exit__(self, *_args): + return None + + main_stream = FakeStream() + broadcast_stream = FakeStream() + sent: list[tuple[tuple[int, ...], str, torch.device]] = [] + handler = pp_utils.PPHandler.__new__(pp_utils.PPHandler) + handler.is_last_rank = True + handler.max_sample_len = 6 + handler.device = torch.device("cpu") + handler.step_seq = 0 + handler.broadcast_on_cpu = cpu_mode + handler.main_stream = main_stream + handler.broadcast_stream = broadcast_stream + handler.last_rank = 7 + handler.broadcast_group = "cpu-sample-group" if cpu_mode else "device-sample-group" + monkeypatch.setattr( + pp_utils, "compute_need_sampled_mask", lambda _batch: np.array([True]) + ) + monkeypatch.setattr(pp_utils.torch.cuda, "stream", lambda _stream: StreamContext()) + monkeypatch.setattr(torch.Tensor, "record_stream", lambda _tensor, _stream: None) + monkeypatch.setattr( + pp_utils.torch.distributed, + "broadcast", + lambda tensor, *, src, group: sent.append( + (tuple(tensor.shape), group, tensor.device) + ), + ) + + assert pp_utils.PPHandler.broadcast( + handler, + torch.tensor([[7]], dtype=torch.int64), + torch.tensor([1], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + SimpleNamespace(num_reqs=1), + ) + + expected_group = "cpu-sample-group" if cpu_mode else "device-sample-group" + assert broadcast_stream.waited_for == ([] if cpu_mode else [main_stream]) + assert [shape for shape, _, _ in sent] == [(3,), (1, 6), (2, 1)] + assert all(group == expected_group for _, group, _ in sent) + assert all(device == torch.device("cpu") for _, _, device in sent) + assert handler.step_seq == 1 def test_pp_diffusion_draft_mode_fails_fast() -> None: @@ -751,6 +999,30 @@ def fake_broadcast(value, *, src, group) -> None: ] +def test_pp_cpu_broadcast_group_warmup_is_blocking(monkeypatch) -> None: + from vllm.v1.worker.gpu import pp_utils + + events: list[object] = [] + marker = torch.zeros(1, dtype=torch.int32) + + monkeypatch.setattr(pp_utils.torch, "zeros", lambda *_args, **_kwargs: marker) + + def fake_broadcast(value, *, src, group) -> None: + events.append(("broadcast", src, group, value.device.type)) + value.fill_(1) + + monkeypatch.setattr(pp_utils.torch.distributed, "broadcast", fake_broadcast) + + pp_utils._warm_up_cpu_broadcast_group( + group="cpu-sample-group", + last_rank=7, + is_last_rank=False, + ) + + assert events == [("broadcast", 7, "cpu-sample-group", "cpu")] + assert marker.item() == 1 + + def test_pp_handler_eagerly_warms_sibling_group(monkeypatch) -> None: from vllm.v1.worker.gpu import pp_utils @@ -767,6 +1039,7 @@ def make_sibling_device_group(self, *, group_desc: str): return "sample-group" group = FakePPGroup() + monkeypatch.delenv("VLLM_WEDGE_AB_PP_BROADCAST_CPU", raising=False) monkeypatch.setattr(pp_utils, "get_pp_group", lambda: group) monkeypatch.setattr(pp_utils.torch.cuda, "current_stream", lambda _device: object()) monkeypatch.setattr(pp_utils.torch.cuda, "Stream", lambda _device: fake_stream) @@ -774,7 +1047,9 @@ def make_sibling_device_group(self, *, group_desc: str): def fake_warmup(**kwargs) -> None: events.append(("warm", kwargs)) - monkeypatch.setattr(pp_utils, "_warm_up_broadcast_group", fake_warmup, raising=False) + monkeypatch.setattr( + pp_utils, "_warm_up_broadcast_group", fake_warmup, raising=False + ) pp_utils.PPHandler( max_num_reqs=4, @@ -791,3 +1066,46 @@ def fake_warmup(**kwargs) -> None: "last_rank": 7, "is_last_rank": False, } + + +def test_pp_handler_eagerly_warms_dedicated_cpu_group(monkeypatch) -> None: + from vllm.v1.worker.gpu import pp_utils + + events: list[object] = [] + + class FakePPGroup: + is_last_rank = False + last_rank = 7 + world_size = 8 + + def make_sibling_cpu_group(self, *, group_desc: str): + events.append(("make_cpu_group", group_desc)) + return "cpu-sample-group" + + monkeypatch.setenv("VLLM_WEDGE_AB_PP_BROADCAST_CPU", "1") + monkeypatch.setattr(pp_utils, "get_pp_group", lambda: FakePPGroup()) + monkeypatch.setattr(pp_utils.torch.cuda, "current_stream", lambda _device: object()) + monkeypatch.setattr(pp_utils.torch.cuda, "Stream", lambda _device: object()) + monkeypatch.setattr( + pp_utils, + "_warm_up_cpu_broadcast_group", + lambda **kwargs: events.append(("warm_cpu", kwargs)), + ) + + pp_utils.PPHandler( + max_num_reqs=4, + num_speculative_steps=5, + device=torch.device("cpu"), + ) + + assert events == [ + ("make_cpu_group", "pp_broadcast_cpu"), + ( + "warm_cpu", + { + "group": "cpu-sample-group", + "last_rank": 7, + "is_last_rank": False, + }, + ), + ] diff --git a/tests/v1/worker/gpu/test_pp_utils_header.py b/tests/v1/worker/gpu/test_pp_utils_header.py index 65e60a777bf9..a52f9b3d55c9 100755 --- a/tests/v1/worker/gpu/test_pp_utils_header.py +++ b/tests/v1/worker/gpu/test_pp_utils_header.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Unit tests for the PPHandler wire header hardening (WEDGE_FIX_GOAL). +"""Unit tests for the PPHandler wire/header safety gates (WEDGE_FIX_GOAL). Covers the contract that defends observable behavior: - header dtype/width: int64 so unsigned CRC32 mask hashes (>= 2**31) survive @@ -8,17 +8,26 @@ with a colliding mask hash and then hang mismatched-shape NCCL payloads; - seq/mask/num_reqs land in the expected header slots; - receive-side validation equality semantics behind the fail-fast raise. +- dedicated CPU/Gloo transport selection is opt-in and speculative-only. Run: /bin/python tests/v1/worker/gpu/test_pp_utils_header.py (plain unittest: pytest is not installed in the serving venv) """ + import os import sys import types import unittest +from unittest import mock -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))) +sys.path.insert( + 0, + os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ) + ), +) import numpy as np import torch @@ -49,8 +58,9 @@ def test_none_mask_hashes_to_zero(self): def test_same_mask_same_hash(self): m = np.array([True, False, True, True] * 8) - self.assertEqual(pp_utils.PPHandler._mask_hash(m), - pp_utils.PPHandler._mask_hash(m.copy())) + self.assertEqual( + pp_utils.PPHandler._mask_hash(m), pp_utils.PPHandler._mask_hash(m.copy()) + ) def test_different_batch_sizes_diverge_in_full_key(self): # CRC32(packbits(mask)) can collide across lengths; num_reqs in the @@ -58,16 +68,19 @@ def test_different_batch_sizes_diverge_in_full_key(self): m4 = np.array([False] * 4) m8 = np.array([False] * 8) h = make_handler() - self.assertFalse(torch.equal( - h._wire_header(h._mask_hash(m4), 4), - h._wire_header(h._mask_hash(m8), 8))) + self.assertFalse( + torch.equal( + h._wire_header(h._mask_hash(m4), 4), h._wire_header(h._mask_hash(m8), 8) + ) + ) class TestWireHeader(unittest.TestCase): def test_int64_dtype(self): hdr = make_handler()._wire_header(0, 1) - self.assertEqual(hdr.dtype, torch.int64, - "int32 overflows unsigned CRC32 hashes") + self.assertEqual( + hdr.dtype, torch.int64, "int32 overflows unsigned CRC32 hashes" + ) self.assertEqual(hdr.numel(), 3) def test_high_bit_mask_hash_roundtrips(self): @@ -102,20 +115,85 @@ class TestHeaderValidationSemantics(unittest.TestCase): def test_equal_headers_match(self): h = make_handler() h.step_seq = 99 - self.assertTrue(torch.equal(h._wire_header(7, 4).cpu(), - h._wire_header(7, 4).cpu())) + self.assertTrue( + torch.equal(h._wire_header(7, 4).cpu(), h._wire_header(7, 4).cpu()) + ) def test_mask_mismatch_detected(self): h = make_handler() h.step_seq = 99 - self.assertFalse(torch.equal(h._wire_header(7, 4).cpu(), - h._wire_header(8, 4).cpu())) + self.assertFalse( + torch.equal(h._wire_header(7, 4).cpu(), h._wire_header(8, 4).cpu()) + ) def test_num_reqs_mismatch_detected(self): h = make_handler() h.step_seq = 99 - self.assertFalse(torch.equal(h._wire_header(7, 4).cpu(), - h._wire_header(7, 5).cpu())) + self.assertFalse( + torch.equal(h._wire_header(7, 4).cpu(), h._wire_header(7, 5).cpu()) + ) + + +class TestBroadcastCpuGate(unittest.TestCase): + """The PP+speculative CPU transport fallback is narrowly gated.""" + + class FakePPGroup: + is_last_rank = False + last_rank = 7 + world_size = 8 + + def make_sibling_device_group(self, *, group_desc): + if group_desc != "pp_broadcast": + raise AssertionError(group_desc) + return "device-sample-group" + + def make_sibling_cpu_group(self, *, group_desc): + if group_desc != "pp_broadcast_cpu": + raise AssertionError(group_desc) + return "cpu-sample-group" + + def make_initialized_handler(self, env_value, num_speculative_steps=5): + main_stream = object() + side_stream = object() + env = {"VLLM_WEDGE_AB_PP_BROADCAST_CPU": env_value} + with ( + mock.patch.dict(os.environ, env), + mock.patch.object( + pp_utils, "get_pp_group", return_value=self.FakePPGroup() + ), + mock.patch.object( + pp_utils.torch.cuda, "current_stream", return_value=main_stream + ), + mock.patch.object(pp_utils.torch.cuda, "Stream", return_value=side_stream), + mock.patch.object(pp_utils, "_warm_up_broadcast_group"), + mock.patch.object(pp_utils, "_warm_up_cpu_broadcast_group"), + ): + handler = pp_utils.PPHandler( + max_num_reqs=4, + num_speculative_steps=num_speculative_steps, + device=torch.device("cpu"), + ) + return handler, main_stream, side_stream + + def test_gate_on_selects_cpu_group_for_speculative_decode(self): + handler, _, side_stream = self.make_initialized_handler("1") + self.assertTrue(handler.broadcast_on_cpu) + self.assertEqual(handler.broadcast_group, "cpu-sample-group") + self.assertIs(handler.broadcast_stream, side_stream) + + def test_gate_off_preserves_device_group(self): + handler, _, side_stream = self.make_initialized_handler("0") + self.assertFalse(handler.broadcast_on_cpu) + self.assertEqual(handler.broadcast_group, "device-sample-group") + self.assertIs(handler.broadcast_stream, side_stream) + + def test_gate_does_not_change_non_speculative_decode(self): + handler, _, side_stream = self.make_initialized_handler( + "1", num_speculative_steps=0 + ) + self.assertFalse(handler.broadcast_on_cpu) + self.assertEqual(handler.broadcast_group, "device-sample-group") + self.assertIs(handler.broadcast_stream, side_stream) if __name__ == "__main__": diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index a90e8acbcad8..d79d9013b028 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -554,6 +554,31 @@ def make_sibling_device_group(self, group_desc: str | None = None) -> ProcessGro assert sibling is not None return sibling + def make_sibling_cpu_group(self, group_desc: str | None = None) -> ProcessGroup: + """Create a distinct Gloo group with this coordinator's membership. + + This is collective across world ranks, just like + :meth:`make_sibling_device_group`. A separate CPU group is useful when + an ordered control protocol must not share collective sequence space + with the coordinator's normal tensor-dict metadata traffic. + """ + from vllm.distributed.utils import get_cpu_distributed_timeout_or_none + + timeout = get_cpu_distributed_timeout_or_none() + sibling: ProcessGroup | None = None + for ranks in self.group_ranks: + with suppress_stdout(): + pg = torch.distributed.new_group( + ranks, + backend="gloo", + group_desc=group_desc, + timeout=timeout, + ) + if self.rank in ranks: + sibling = pg + assert sibling is not None + return sibling + def create_mq_broadcaster( self, writer_rank=0, external_writer_handle=None, blocking=True ): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a2534439a542..4787636b1b39 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1540,15 +1540,17 @@ def sample_tokens( self.num_speculative_steps > 0 and structured_req_ids and ( - received_drafts := self.pp_handler.get_current_received_draft_tokens() + received_drafts + := self.pp_handler.get_current_received_draft_tokens() ) is not None ): - draft_tokens, draft_ready_event = received_drafts + draft_tokens, draft_ready_event, draft_ready_waitable = received_drafts self.draft_tokens_handler.set_draft_tokens( input_batch, draft_tokens, ready_event=draft_ready_event, + ready_waitable=draft_ready_waitable, structured_req_ids=structured_req_ids, ) # Optimistically update num_computed_tokens for entire batch here. diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index ca52e4e13957..edccc17a238c 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -30,14 +30,19 @@ the DSpark proposal failed/skipped, the last rank sends a sentinel payload (zeroed counts / -1 draft tokens) so receivers still get a matching step. Consumers filter sentinels out via their own need-sampled-mask state. +- An opt-in PP+speculative safety profile moves these tiny control/token + payloads to a dedicated CPU/Gloo group. This isolates them from + activation P2P and other collectives on NCCL without sharing sequence space + with the PP coordinator's normal CPU metadata group. Wire-format is intra-process only (all ranks load the same module), so the change is compatible with any embedding as long as all ranks use this build. """ +import os +import zlib from collections import deque from dataclasses import dataclass -import zlib import numpy as np import torch @@ -92,11 +97,38 @@ def _warm_up_broadcast_group( raise RuntimeError("sampled-token broadcast communicator warmup failed") +def _warm_up_cpu_broadcast_group( + *, + group: torch.distributed.ProcessGroup, + last_rank: int, + is_last_rank: bool, +) -> None: + """Validate the dedicated sampled-token Gloo group at initialization.""" + marker = torch.zeros(1, dtype=torch.int32, device="cpu") + if is_last_rank: + marker.fill_(1) + torch.distributed.broadcast(marker, src=last_rank, group=group) + if marker.item() != 1: + raise RuntimeError("sampled-token CPU broadcast communicator warmup failed") + + +@dataclass +class CpuRecvState: + """Own an asynchronous Gloo receive until deferred consumers need it.""" + + work: torch.distributed.Work | None + + def wait(self) -> None: + if self.work is not None: + self.work.wait() + self.work = None + + @dataclass class PendingRecv: """Per-step slot data for a deferred postprocess on the main stream.""" - event: torch.cuda.Event + event: torch.cuda.Event | None sampled_tokens: torch.Tensor # [num_reqs, max_sample_len] num_sampled: torch.Tensor # [num_reqs] num_rejected: torch.Tensor # [num_reqs] @@ -105,6 +137,7 @@ class PendingRecv: idx_mapping_np: np.ndarray # [num_reqs] need_sampled_mask: np.ndarray | None # [num_reqs] (None => sentinel step) gen_at_receive_np: np.ndarray # [num_reqs] + cpu_recv: CpuRecvState | None = None def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: @@ -125,7 +158,7 @@ def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: class PPHandler: - """Runs the PP sampled-token broadcast/recv on a side stream. + """Runs the PP sampled-token broadcast/recv on a dedicated communicator. Protocol is data-independent (see module docstring): every step posts exactly four broadcasts on every rank, preceded by a validation header. @@ -134,20 +167,29 @@ class PPHandler: def __init__( self, max_num_reqs: int, num_speculative_steps: int, device: torch.device ): - self.is_last_rank = get_pp_group().is_last_rank - self.last_rank = get_pp_group().last_rank + pp_group = get_pp_group() + self.is_last_rank = pp_group.is_last_rank + self.last_rank = pp_group.last_rank self.max_sample_len = num_speculative_steps + 1 self.max_draft_len = num_speculative_steps self.device = device self.main_stream = torch.cuda.current_stream(device) self.broadcast_stream = torch.cuda.Stream(device) + # Correctness fallback for PP+speculative decoding. The tiny wire + # payloads use a dedicated Gloo group instead of a sibling + # NCCL communicator, which can contend cyclically with activation P2P. + # Keep it opt-in while its throughput impact is characterized. + self.broadcast_on_cpu = ( + num_speculative_steps > 0 + and os.environ.get("VLLM_WEDGE_AB_PP_BROADCAST_CPU", "0") == "1" + ) # On non-last ranks, a FIFO with one entry per in-flight step: the entry # pushed by step T's `receive` is consumed pp_size steps later. Pre-seeded # with pp_size None placeholders so the first pp_size consumes are no-ops. # None means no postprocess is pending for that step (broadcast skipped). self.queue: deque[PendingRecv | None] = ( - deque() if self.is_last_rank else deque([None] * get_pp_group().world_size) + deque() if self.is_last_rank else deque([None] * pp_group.world_size) ) # Per req-index generation counter, incremented every time a request @@ -160,22 +202,39 @@ def __init__( # sample_tokens step, in the same order they consume scheduler outputs. self.step_seq: int = 0 - # Dedicated subgroup for the sampled-token broadcast. - self.broadcast_group = get_pp_group().make_sibling_device_group( - group_desc="pp_broadcast" - ) - _warm_up_broadcast_group( - group=self.broadcast_group, - stream=self.broadcast_stream, - device=self.device, - last_rank=self.last_rank, - is_last_rank=self.is_last_rank, - ) + # Never reuse pp_group.cpu_group here: tensor-dict metadata may use it + # independently and interleave a different collective order. + if self.broadcast_on_cpu: + self.broadcast_group = pp_group.make_sibling_cpu_group( + group_desc="pp_broadcast_cpu" + ) + _warm_up_cpu_broadcast_group( + group=self.broadcast_group, + last_rank=self.last_rank, + is_last_rank=self.is_last_rank, + ) + else: + self.broadcast_group = pp_group.make_sibling_device_group( + group_desc="pp_broadcast" + ) + _warm_up_broadcast_group( + group=self.broadcast_group, + stream=self.broadcast_stream, + device=self.device, + last_rank=self.last_rank, + is_last_rank=self.is_last_rank, + ) def on_req_idx_freed(self, req_idx: int) -> None: self.req_idx_gen_np[req_idx] += 1 - def _wire_header(self, mask_hash: int, num_reqs: int) -> torch.Tensor: + def _wire_header( + self, + mask_hash: int, + num_reqs: int, + *, + device: torch.device | str | None = None, + ) -> torch.Tensor: """Fixed-size broadcast header: [seq, mask_hash, num_reqs] as int64. int64 throughout: mask_hash is an unsigned CRC32 (0..2^32-1), which @@ -187,9 +246,24 @@ def _wire_header(self, mask_hash: int, num_reqs: int) -> torch.Tensor: return torch.tensor( [self.step_seq, mask_hash, num_reqs], dtype=torch.int64, - device=self.device, + device=self.device if device is None else device, ) + def _validate_wire_header( + self, header: torch.Tensor, mask_hash: int, num_reqs: int + ) -> None: + """Validate and advance the fixed-size wire header on a CPU tensor.""" + header_cpu = header.cpu() + expected_cpu = self._wire_header(mask_hash, num_reqs, device="cpu") + if not torch.equal(header_cpu, expected_cpu): + raise RuntimeError( + "PP sampled-token broadcast desync: " + f"sender header={header_cpu.tolist()} " + f"receiver expected={expected_cpu.tolist()} " + f"(step_seq before increment={self.step_seq})" + ) + self.step_seq += 1 + @staticmethod def _mask_hash(need_sampled_mask: np.ndarray | None) -> int: """Deterministic hash of the need-sampled mask for wire validation.""" @@ -213,6 +287,8 @@ def get_prev_sampled_outputs( if slot is None: return None + self._materialize_cpu_recv(slot) + # Skip requests which did not need sampled output and/or those already # finished. The post_update kernel skips the -1 entries. freed = self.req_idx_gen_np[slot.idx_mapping_np] != slot.gen_at_receive_np @@ -232,6 +308,7 @@ def get_prev_sampled_outputs( idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np) idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + assert slot.event is not None self.main_stream.wait_event(slot.event) return ( dict( @@ -256,6 +333,82 @@ def receive(self, input_batch: InputBatch) -> bool: gen_at_receive_np = self.req_idx_gen_np[input_batch.idx_mapping_np] num_reqs = input_batch.num_reqs + if self.broadcast_on_cpu: + header = torch.empty(3, dtype=torch.int64, device="cpu") + torch.distributed.broadcast( + header, src=self.last_rank, group=self.broadcast_group + ) + self._validate_wire_header(header, mask_hash, num_reqs) + + sampled_tokens_cpu = torch.empty( + num_reqs, self.max_sample_len, dtype=torch.int64, device="cpu" + ) + combined_cpu = torch.empty(2, num_reqs, dtype=torch.int32, device="cpu") + draft_tokens_cpu = torch.empty( + num_reqs, self.max_draft_len, dtype=torch.int64, device="cpu" + ) + for tensor in (sampled_tokens_cpu, combined_cpu): + torch.distributed.broadcast( + tensor, src=self.last_rank, group=self.broadcast_group + ) + draft_work = torch.distributed.broadcast( + draft_tokens_cpu, + src=self.last_rank, + group=self.broadcast_group, + async_op=True, + ) + assert draft_work is not None + sampled_tokens = sampled_tokens_cpu + num_sampled, num_rejected = combined_cpu.unbind(dim=0) + draft_tokens = draft_tokens_cpu + event = None + cpu_recv = CpuRecvState(draft_work) + else: + sampled_tokens, num_sampled, num_rejected, draft_tokens, event = ( + self._receive_device(mask_hash, num_reqs) + ) + cpu_recv = None + + self.queue[-1] = PendingRecv( + event, + sampled_tokens, + num_sampled, + num_rejected, + draft_tokens, + input_batch.idx_mapping, + input_batch.idx_mapping_np, + need_sampled_mask, + gen_at_receive_np, + cpu_recv, + ) + # Draft tokens may be a sentinel when there was no real sampled output + # this step; consumers filter via need_sampled_mask at consume time. + return bool(need_sampled_mask is not None and need_sampled_mask.all()) + + def _materialize_cpu_recv(self, slot: PendingRecv) -> None: + """Wait for a deferred Gloo draft and move one slot onto the GPU.""" + if slot.cpu_recv is None: + return + slot.cpu_recv.wait() + with torch.cuda.stream(self.main_stream): + slot.sampled_tokens = slot.sampled_tokens.to(self.device) + slot.num_sampled = slot.num_sampled.to(self.device) + slot.num_rejected = slot.num_rejected.to(self.device) + assert slot.draft_tokens is not None + slot.draft_tokens = slot.draft_tokens.to(self.device) + slot.event = self.main_stream.record_event() + slot.cpu_recv = None + + def _receive_device( + self, mask_hash: int, num_reqs: int + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.cuda.Event, + ]: + """Receive one protocol step with the default device communicator.""" with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) header = torch.empty(3, dtype=torch.int64, device=self.device) @@ -264,22 +417,11 @@ def receive(self, input_batch: InputBatch) -> bool: ) # Validate sender/receiver step agreement BEFORE consuming payloads. # Copy to CPU: tiny (3 ints), and a mismatch must raise from here. - expected = self._wire_header(mask_hash, num_reqs) - expected_cpu = expected.cpu() header_cpu = header.cpu() # The blocking .cpu() syncs the broadcast stream, so NCCL is done # reading `header`; keep it alive past the host read anyway. header.record_stream(self.main_stream) - if not torch.equal(header_cpu, expected_cpu): - raise RuntimeError( - "PP sampled-token broadcast desync: " - f"sender header={header_cpu.tolist()} " - f"receiver expected={expected_cpu.tolist()} " - f"(step_seq before increment={self.step_seq})" - ) - # Step matched: advance the sequence counter on every rank. - self.step_seq += 1 - del expected_cpu, header_cpu + self._validate_wire_header(header_cpu, mask_hash, num_reqs) sampled_tokens = torch.empty( num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device ) @@ -308,29 +450,22 @@ def receive(self, input_batch: InputBatch) -> bool: sampled_tokens.record_stream(self.main_stream) combined.record_stream(self.main_stream) draft_tokens.record_stream(self.main_stream) - self.queue[-1] = PendingRecv( - event, + return ( sampled_tokens, num_sampled, num_rejected, draft_tokens, - input_batch.idx_mapping, - input_batch.idx_mapping_np, - need_sampled_mask, - gen_at_receive_np, + event, ) - # Draft tokens may be a sentinel when there was no real sampled output - # this step; consumers filter via need_sampled_mask at consume time. - return bool(need_sampled_mask is not None and need_sampled_mask.all()) def get_current_received_draft_tokens( self, - ) -> tuple[torch.Tensor, torch.cuda.Event] | None: - """Return this step's PP-broadcast drafts and their readiness event.""" + ) -> tuple[torch.Tensor, torch.cuda.Event | None, CpuRecvState | None] | None: + """Return this step's drafts and its CUDA event or CPU wait handle.""" slot = self.queue[-1] if slot is None or slot.draft_tokens is None: return None - return slot.draft_tokens, slot.event + return slot.draft_tokens, slot.event, slot.cpu_recv def broadcast( self, @@ -345,6 +480,17 @@ def broadcast( assert sampled_token_ids.dtype == torch.int64 + if self.broadcast_on_cpu: + self._broadcast_cpu( + sampled_token_ids, + num_sampled, + num_rejected, + input_batch.num_reqs, + mask_hash, + need_sampled_mask is None, + ) + return True + if current_platform.is_xpu(): self.main_stream.synchronize() @@ -382,12 +528,57 @@ def broadcast( tensor.record_stream(self.broadcast_stream) return True + def _broadcast_cpu( + self, + sampled_token_ids: torch.Tensor, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + num_reqs: int, + mask_hash: int, + sentinel_counts: bool, + ) -> None: + """Send the header and sampled-token metadata as blocking CPU tensors.""" + header = self._wire_header(mask_hash, num_reqs, device="cpu") + self.step_seq += 1 + torch.distributed.broadcast( + header, src=self.last_rank, group=self.broadcast_group + ) + + wire_sampled_token_ids = _pad_sampled_tokens_for_broadcast( + sampled_token_ids, self.max_sample_len + ).cpu() + torch.distributed.broadcast( + wire_sampled_token_ids, + src=self.last_rank, + group=self.broadcast_group, + ) + if sentinel_counts: + combined = torch.zeros(2, num_reqs, dtype=torch.int32, device="cpu") + else: + combined = torch.stack((num_sampled, num_rejected), dim=0).cpu() + torch.distributed.broadcast( + combined, src=self.last_rank, group=self.broadcast_group + ) + def broadcast_draft_tokens(self, draft_token_ids: torch.Tensor) -> None: """Broadcast real post-proposal (or sentinel) draft IDs to non-last PP ranks.""" assert self.is_last_rank if self.max_draft_len == 0: return assert draft_token_ids.dtype == torch.int64 + if self.broadcast_on_cpu: + # .cpu() is blocking, so the persistent source buffer is safe to + # reuse as soon as the blocking Gloo broadcast returns. + wire_draft_token_ids = _pad_sampled_tokens_for_broadcast( + draft_token_ids, self.max_draft_len + ).cpu() + torch.distributed.broadcast( + wire_draft_token_ids, + src=self.last_rank, + group=self.broadcast_group, + ) + return + with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) wire_draft_token_ids = _pad_sampled_tokens_for_broadcast( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 333648637bd7..3188c8c974be 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -133,6 +133,10 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: decode_query_len=self.num_query_per_req, ) + def _trace_proposal_stage(self, stage: str) -> None: + """Optional subclass hook for narrowly scoped live diagnostics.""" + del stage + def capture(self) -> None: logger.info("Capturing model for %s speculator...", self._speculator_name) # Reset sampling indices to zero to prevent stale values from prior @@ -346,6 +350,7 @@ def propose( else: hidden_states = last_hidden_states self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens]) + self._trace_proposal_stage("hidden-copy") if dummy_run and skip_attn_for_dummy_run: # Memory profiling path: block_tables / kv_cache_config are not initialized. @@ -400,6 +405,7 @@ def propose( self.max_model_len, self.sample_from_anchor, ) + self._trace_proposal_stage("prepare-inputs") # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph # because the context shape varies per step. During dummy runs the block tables @@ -419,6 +425,7 @@ def propose( self.context_positions[:num_target_tokens], context_slots, ) + self._trace_proposal_stage("context-kv") # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( @@ -466,6 +473,8 @@ def propose( cudagraph_runtime_mode=batch_desc.cg_mode, ) + self._trace_proposal_stage("draft-complete") + return self.draft_tokens[:num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 45dedde3d7db..66c10c3e8aa6 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -23,16 +23,20 @@ backbone forward AND the sequential Markov sampling. """ +import os from typing import Any import torch from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model +logger = init_logger(__name__) + class DSparkSpeculator(DFlashSpeculator): _speculator_name = "DSpark" @@ -40,6 +44,11 @@ class DSparkSpeculator(DFlashSpeculator): def __init__(self, vllm_config: VllmConfig, device: torch.device): super().__init__(vllm_config, device) + self._wedge_trace_stages = ( + os.environ.get("VLLM_WEDGE_TRACE_DSPARK_STAGES", "0") == "1" + ) + self._wedge_proposal_seq = 0 + # Whether to sample from the anchor position. When True, uses anchor-as-first # (N slots, each position predicts the next token). When False, uses 1+N # fill-in block (anchor is a bonus token). @@ -73,6 +82,31 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self._d2t_scatter_index: torch.Tensor | None = None self._draft_scatter_buf: torch.Tensor | None = None + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: + force_eager = ( + os.environ.get("VLLM_WEDGE_AB_DSPARK_EAGER", "0") == "1" + or os.environ.get("VLLM_WEDGE_TRACE_DSPARK_STAGES", "0") == "1" + ) + if force_eager: + logger.warning( + "DSpark draft CUDA graphs are disabled by the wedge diagnostic " + "gate; target-model CUDA graphs remain unchanged." + ) + cudagraph_mode = CUDAGraphMode.NONE + super().init_cudagraph_manager(cudagraph_mode) + + def _trace_proposal_stage(self, stage: str) -> None: + if not self._wedge_trace_stages: + return + torch.cuda.synchronize(self.device) + logger.info( + "[DSPARK WEDGE TRACE] proposal_seq=%d stage=%s synchronized", + self._wedge_proposal_seq, + stage, + ) + if stage == "draft-complete": + self._wedge_proposal_seq += 1 + def load_draft_model( self, target_model: torch.nn.Module, @@ -166,4 +200,6 @@ def _generate_draft( num_tokens_across_dp, cudagraph_runtime_mode, ) + self._trace_proposal_stage("draft-backbone") self._sample_sequential(num_reqs, head_hidden) + self._trace_proposal_stage("draft-sampling") diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 95443445c21f..fe09b99cd7e4 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -3,6 +3,8 @@ # Apache-2.0 Section 4(b): modified by Lime Labs EOOD; upstream notices preserved. # Original overlay modifications only; upstream ownership is not claimed. from collections import deque +from dataclasses import dataclass +from typing import Protocol import numpy as np import torch @@ -12,6 +14,33 @@ from vllm.v1.worker.gpu.input_batch import InputBatch +class Waitable(Protocol): + def wait(self) -> None: ... + + +@dataclass +class DeferredCpuDrafts: + """A Gloo destination whose data becomes readable after `ready.wait()`.""" + + tensor: torch.Tensor + ready: Waitable + + def materialize(self) -> np.ndarray: + self.ready.wait() + return self.tensor.numpy() + + +def _materialize_deferred_drafts( + drafts: np.ndarray | DeferredCpuDrafts, + ready_event: torch.cuda.Event | None, +) -> np.ndarray: + if isinstance(drafts, DeferredCpuDrafts): + return drafts.materialize() + if ready_event is not None: + ready_event.synchronize() + return drafts + + class DraftTokensHandler: def __init__(self, device: torch.device | None = None): self.device = device @@ -21,7 +50,12 @@ def __init__(self, device: torch.device | None = None): self.draft_tokens_np: np.ndarray | None = None self.num_draft_tokens: int = 0 self.pending_structured_drafts: deque[ - tuple[list[str], list[int], np.ndarray, torch.cuda.Event] + tuple[ + list[str], + list[int], + np.ndarray | DeferredCpuDrafts, + torch.cuda.Event | None, + ] ] = deque() def set_draft_tokens( @@ -30,6 +64,7 @@ def set_draft_tokens( draft_tokens: torch.Tensor, *, ready_event: torch.cuda.Event | None = None, + ready_waitable: Waitable | None = None, structured_req_ids: set[str] | None = None, ) -> None: if structured_req_ids is None: @@ -54,6 +89,19 @@ def set_draft_tokens( # the scheduler for this batch. return + if ready_waitable is not None: + assert ready_event is None + assert draft_tokens.device.type == "cpu" + self.pending_structured_drafts.append( + ( + [req_id for _, req_id in structured_rows], + [row for row, _ in structured_rows], + DeferredCpuDrafts(draft_tokens, ready_waitable), + None, + ) + ) + return + # For spec decoding + structured outputs, we must transfer the # draft tokens back to the scheduler for grammar validation. if ready_event is None: @@ -80,9 +128,7 @@ def set_draft_tokens( ) ) - def get_draft_tokens( - self, req_ids: set[str] | None = None - ) -> DraftTokenIds | None: + def get_draft_tokens(self, req_ids: set[str] | None = None) -> DraftTokenIds | None: if self.pending_structured_drafts: # Deferred structured SchedulerOutputs are consumed independently. # Return only the requested rows and preserve unrelated snapshots for @@ -91,15 +137,23 @@ def get_draft_tokens( selected_drafts: dict[str, list[int]] = {} remaining_drafts = deque() while self.pending_structured_drafts: - batch_req_ids, row_indices, draft_tokens_np, copy_event = ( + batch_req_ids, row_indices, draft_tokens, copy_event = ( self.pending_structured_drafts.popleft() ) - copy_event.synchronize() + # With PP, non-last ranks can return from sample_tokens while a + # newer Gloo draft is still in flight. A selective request for + # an older structured batch must not wait for that unrelated + # receive, or it can prevent the worker from accepting the + # pipeline work that lets the last rank produce the draft. + if req_ids is not None and req_ids.isdisjoint(batch_req_ids): + remaining_drafts.append( + (batch_req_ids, row_indices, draft_tokens, copy_event) + ) + continue + draft_tokens_np = _materialize_deferred_drafts(draft_tokens, copy_event) keep_req_ids: list[str] = [] keep_row_indices: list[int] = [] - for req_id, row_index in zip( - batch_req_ids, row_indices, strict=True - ): + for req_id, row_index in zip(batch_req_ids, row_indices, strict=True): if req_ids is None or req_id in req_ids: selected_drafts[req_id] = draft_tokens_np[row_index].tolist() else: @@ -144,17 +198,17 @@ def discard_req_ids(self, req_ids: set[str]) -> None: return remaining_drafts = deque() while self.pending_structured_drafts: - batch_req_ids, row_indices, draft_tokens_np, copy_event = ( + batch_req_ids, row_indices, draft_tokens, copy_event = ( self.pending_structured_drafts.popleft() ) if req_ids.isdisjoint(batch_req_ids): remaining_drafts.append( - (batch_req_ids, row_indices, draft_tokens_np, copy_event) + (batch_req_ids, row_indices, draft_tokens, copy_event) ) continue # The pinned CPU destination cannot be released while the async D2H # copy is still in flight. - copy_event.synchronize() + draft_tokens_np = _materialize_deferred_drafts(draft_tokens, copy_event) keep_indices = [ index for index, req_id in enumerate(batch_req_ids)