diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index 4e01ef4c2..2c85c33fe 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -23,8 +23,8 @@ some GPUs. Asynchronous chunk streaming is **enabled by default** within the bundled configuration. For Thinker pipeline parallelism, configure `pipeline_parallel_size` and `devices` -on the Thinker stage and set `async_scheduling: false`. Keep Talker and Code2Wav -at `pipeline_parallel_size: 1`. Audio chunk streaming can remain enabled. +on the Thinker stage. Keep Talker and Code2Wav at `pipeline_parallel_size: 1`. The +default asynchronous scheduling and audio chunk streaming can remain enabled. To explicitly utilize a custom deployment YAML, specify the configuration path: ```bash diff --git a/tests/worker/test_pp_async_token_broadcast.py b/tests/worker/test_pp_async_token_broadcast.py new file mode 100644 index 000000000..d5e938681 --- /dev/null +++ b/tests/worker/test_pp_async_token_broadcast.py @@ -0,0 +1,321 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Async scheduling + pipeline parallelism: the last rank must send what earlier ranks wait for. + +vLLM's ``GPUModelRunner.sample_tokens`` pairs two collectives on the PP device group: +a non-final rank posts an async broadcast *receive* of the newly sampled token ids +(``_pp_receive_prev_sampled_token_ids_to_input_batch``) and the final rank performs +the matching broadcast *send* (``_pp_broadcast_prev_sampled_token_ids``) right after +``_update_states_after_model_execute``. ``GPUARModelRunner.sample_tokens`` overrides the +upstream method and inherited the receive but not the send. With async scheduling and +PP > 1 the first rank never returned from posting that receive: it is the first +collective on the PP group's communicator, NCCL has to create the communicator with every +rank, and the last rank never joined (observed live on 2x A100: rank 0 inside +``torch.distributed.broadcast``, rank 1 waiting for the next intermediate tensors, engine +dead after the RPC timeout). Had the communicator existed, the same missing send would +have blocked the ``_pp_recv_work.wait()`` in ``_prepare_input_ids`` on the next step. + +These tests drive the real Omni override with a recorded ``torch.distributed.broadcast`` +ledger and, separately, a real two-process gloo group so the receive is proven to +complete only because the Omni sampling path sent. +""" + +from __future__ import annotations + +import os +from datetime import timedelta +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from vllm.v1.worker import gpu_model_runner as upstream_runner_module + +from vllm_omni.worker import gpu_ar_model_runner as omni_runner_module +from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, GPUARModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.parallel] + + +class _StopAtBookkeepingError(Exception): + """Raised by the bookkeeping double: everything before it is what we assert on.""" + + +def _pp(rank: int, world_size: int, device_group: object = "pp-group") -> SimpleNamespace: + return SimpleNamespace( + rank=rank, + world_size=world_size, + last_rank=world_size - 1, + is_last_rank=rank == world_size - 1, + is_first_rank=rank == 0, + device_group=device_group, + ) + + +def _input_batch(num_reqs: int) -> SimpleNamespace: + return SimpleNamespace( + num_reqs=num_reqs, + req_ids=[f"req-{i}" for i in range(num_reqs)], + num_tokens_no_spec=np.ones(num_reqs, dtype=np.int64), + is_token_ids=np.zeros((num_reqs, 8), dtype=bool), + sampling_metadata=SimpleNamespace(no_penalties=True, logitsprocs=SimpleNamespace(non_argmax_invariant=[])), + vocab_size=32, + prev_sampled_token_ids=None, + prev_req_id_to_index=None, + ) + + +def _last_rank_runner( + *, + sampled: torch.Tensor, + async_scheduling: bool = True, + broadcast_pp_output: bool = False, + all_chunked_prefill: bool = False, + events: list[str], +) -> GPUARModelRunner: + """A final-rank runner double that reaches the sampler and the state update.""" + num_reqs = sampled.shape[0] + hidden = torch.ones(num_reqs, 8) + runner = object.__new__(GPUARModelRunner) + runner.use_async_scheduling = async_scheduling + runner.broadcast_pp_output = broadcast_pp_output + runner.speculative_config = None + runner.device = torch.device("cpu") + runner.input_batch = _input_batch(num_reqs) + runner.discard_request_mask = SimpleNamespace(np=np.full(num_reqs, all_chunked_prefill, dtype=bool)) + runner.execute_model_state = ExecuteModelState( + scheduler_output=SimpleNamespace(total_num_scheduled_tokens=num_reqs), + logits=torch.zeros(num_reqs, 32), + spec_decode_metadata=None, + spec_decode_common_attn_metadata=None, + hidden_states=hidden, + hidden_states_cpu=None, + sample_hidden_states=hidden, + aux_hidden_states=None, + ec_connector_output=None, + cudagraph_stats=None, + multimodal_outputs={}, + slot_mappings=None, + ) + runner._sample = lambda logits, spec: SimpleNamespace(sampled_token_ids=sampled, logprobs_tensors=None) + runner._update_states_after_model_execute = lambda *args: events.append("state-update") + + def stop(*args): + raise _StopAtBookkeepingError + + runner._bookkeeping_sync = stop + return runner + + +def _non_last_rank_runner(num_reqs: int) -> GPUARModelRunner: + runner = object.__new__(GPUARModelRunner) + runner.use_async_scheduling = True + runner.execute_model_state = None + runner.kv_connector_output = None + runner.device = torch.device("cpu") + runner._pp_recv_work = None # as vLLM's __init__ leaves it before any receive is posted + runner.input_batch = _input_batch(num_reqs) + runner.discard_request_mask = SimpleNamespace(np=np.zeros(num_reqs, dtype=bool)) + runner.requests = {rid: SimpleNamespace(output_token_ids=[]) for rid in runner.input_batch.req_ids} + runner.attach_omni_connector_output = lambda output: output + return runner + + +class _Ledger(list): + """Every torch.distributed.broadcast, plus a shared timeline with the state update.""" + + def __init__(self): + super().__init__() + self.timeline: list[str] = [] + + +@pytest.fixture +def ledger(monkeypatch): + calls = _Ledger() + + class Pending: + def wait(self): + if not any(c["kind"] == "send" for c in calls): + raise RuntimeError("receive has no matching broadcast") + + def broadcast(tensor, src, group, async_op=False): + kind = "receive" if async_op else "send" + calls.append( + dict( + kind=kind, + shape=tuple(tensor.shape), + dtype=tensor.dtype, + src=src, + group=group, + value=tensor.clone(), + ) + ) + calls.timeline.append(kind) + return Pending() if async_op else None + + monkeypatch.setattr(torch.distributed, "broadcast", broadcast) + return calls + + +def _use_pp(monkeypatch, pp): + # The guard lives in the Omni override; the sender/receiver bodies live in vLLM. + monkeypatch.setattr(omni_runner_module, "get_pp_group", lambda: pp) + monkeypatch.setattr(upstream_runner_module, "get_pp_group", lambda: pp) + + +def test_non_last_rank_posts_the_receive_before_the_connector_only_return(monkeypatch, ledger): + _use_pp(monkeypatch, _pp(rank=0, world_size=2)) + runner = _non_last_rank_runner(num_reqs=2) + output = runner.sample_tokens(None) + assert output is not None + assert [c["kind"] for c in ledger] == ["receive"] + assert ledger[0]["src"] == 1 and ledger[0]["group"] == "pp-group" + assert ledger[0]["shape"] == (2, 1) and ledger[0]["dtype"] == torch.int32 + assert runner._pp_recv_work is not None + assert runner.input_batch.prev_sampled_token_ids.shape == (2, 1) + assert runner.input_batch.prev_req_id_to_index == {"req-0": 0, "req-1": 1} + + +def test_last_rank_broadcasts_sampled_tokens_after_the_state_update(monkeypatch, ledger): + pp = _pp(rank=1, world_size=2) + _use_pp(monkeypatch, pp) + sampled = torch.tensor([[4], [7]], dtype=torch.int32) + runner = _last_rank_runner(sampled=sampled, events=ledger.timeline) + with pytest.raises(_StopAtBookkeepingError): + runner.sample_tokens(None) + assert ledger.timeline == ["state-update", "send"], "the send must follow the state update, not precede it" + sends = [c for c in ledger if c["kind"] == "send"] + assert len(sends) == 1 + assert sends[0]["src"] == pp.rank == pp.last_rank + assert sends[0]["group"] == "pp-group" + assert sends[0]["shape"] == (2, 1) and sends[0]["dtype"] == torch.int32 + assert torch.equal(sends[0]["value"], sampled) + + +def test_receive_and_send_form_one_matched_pair(monkeypatch, ledger): + """The negative control from the handoff probe, now expected to complete.""" + _use_pp(monkeypatch, _pp(rank=0, world_size=2)) + first = _non_last_rank_runner(num_reqs=1) + first.sample_tokens(None) + _use_pp(monkeypatch, _pp(rank=1, world_size=2)) + last = _last_rank_runner(sampled=torch.tensor([[4]], dtype=torch.int32), events=[]) + with pytest.raises(_StopAtBookkeepingError): + last.sample_tokens(None) + first._pp_recv_work.wait() + assert [c["kind"] for c in ledger] == ["receive", "send"] + assert ledger[0]["shape"] == ledger[1]["shape"] and ledger[0]["dtype"] == ledger[1]["dtype"] + assert ledger[0]["src"] == ledger[1]["src"] and ledger[0]["group"] == ledger[1]["group"] + + +@pytest.mark.parametrize( + ("async_scheduling", "broadcast_pp_output", "world_size", "reason"), + [ + (False, False, 2, "synchronous scheduling feeds tokens through the scheduler"), + (True, True, 2, "external_launcher already broadcast the logits to every rank"), + (True, False, 1, "a single rank has nobody waiting"), + ], +) +def test_last_rank_does_not_send_outside_the_async_pp_case( + monkeypatch, ledger, async_scheduling, broadcast_pp_output, world_size, reason +): + _use_pp(monkeypatch, _pp(rank=world_size - 1, world_size=world_size)) + runner = _last_rank_runner( + sampled=torch.tensor([[4]], dtype=torch.int32), + async_scheduling=async_scheduling, + broadcast_pp_output=broadcast_pp_output, + events=[], + ) + with pytest.raises(_StopAtBookkeepingError): + runner.sample_tokens(None) + assert ledger == [], reason + + +def test_both_sides_skip_an_all_chunked_prefill_step(monkeypatch, ledger): + """Dummy tokens from unfinished prefill chunks are discarded, so neither side talks.""" + _use_pp(monkeypatch, _pp(rank=0, world_size=2)) + first = _non_last_rank_runner(num_reqs=1) + first.discard_request_mask = SimpleNamespace(np=np.ones(1, dtype=bool)) + first.sample_tokens(None) + assert first._pp_recv_work is None + _use_pp(monkeypatch, _pp(rank=1, world_size=2)) + last = _last_rank_runner(sampled=torch.tensor([[4]], dtype=torch.int32), all_chunked_prefill=True, events=[]) + with pytest.raises(_StopAtBookkeepingError): + last.sample_tokens(None) + assert ledger == [] + + +def test_sender_rejects_multi_token_rows(monkeypatch, ledger): + """The receive buffer is [num_reqs, 1]; the upstream sender guards that contract.""" + _use_pp(monkeypatch, _pp(rank=1, world_size=2)) + runner = _last_rank_runner(sampled=torch.tensor([[4, 5]], dtype=torch.int32), events=[]) + with pytest.raises(AssertionError, match="num_reqs, 1"): + runner.sample_tokens(None) + assert ledger == [] + + +# --- real collective ----------------------------------------------------------------- + + +def _gloo_rank(rank: int, world_size: int, rendezvous: str, result_path: str) -> None: + """Rank 0 runs the Omni receive branch; rank 1 runs the Omni sampling path. + + Both use the real methods bound to doubles; the process group is a real gloo group, + so rank 0's wait can only return because rank 1's ``sample_tokens`` broadcast. + """ + import torch.distributed as dist + + dist.init_process_group( + backend="gloo", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + try: + pp = _pp(rank=rank, world_size=world_size, device_group=dist.group.WORLD) + omni_runner_module.get_pp_group = lambda: pp + upstream_runner_module.get_pp_group = lambda: pp + if rank == world_size - 1: + runner = _last_rank_runner(sampled=torch.tensor([[4], [7]], dtype=torch.int32), events=[]) + try: + runner.sample_tokens(None) + except _StopAtBookkeepingError: + pass + received = None + else: + runner = _non_last_rank_runner(num_reqs=2) + runner.sample_tokens(None) + assert runner._pp_recv_work is not None + # A missing send surfaces here as a gloo timeout instead of a silent hang. + runner._pp_recv_work.wait(timeout=timedelta(seconds=30)) + received = runner.input_batch.prev_sampled_token_ids.tolist() + with open(f"{result_path}.{rank}", "w") as handle: + handle.write(repr(received)) + finally: + dist.destroy_process_group() + + +def test_gloo_pp2_receive_completes_only_because_the_omni_sampling_path_sent(tmp_path): + import torch.multiprocessing as mp + + rendezvous = str(tmp_path / "rendezvous") + result_path = str(tmp_path / "received") + context = mp.get_context("spawn") + processes = [ + context.Process(target=_gloo_rank, args=(rank, 2, rendezvous, result_path), daemon=True) for rank in range(2) + ] + try: + for process in processes: + process.start() + for process in processes: + process.join(timeout=240) + finally: + for process in processes: + if process.is_alive(): + process.kill() + process.join(timeout=30) + exit_codes = [process.exitcode for process in processes] + assert exit_codes == [0, 0], f"rank exit codes {exit_codes}: an unmatched receive times out on rank 0" + with open(f"{result_path}.0") as handle: + assert handle.read() == repr([[4], [7]]) + assert os.path.exists(f"{result_path}.1") diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 6625effe8..06fdfa18b 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -2088,6 +2088,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex sampler_output = self._sample(logits, spec_decode_metadata) self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output) + if self.use_async_scheduling: + pp = get_pp_group() + # Earlier PP ranks posted the matching async receive in the + # connector-only branch above (inherited from vLLM). Without this + # send their next `_prepare_input_ids` waits forever. + # For torchrun external_launcher PP mode with broadcast_pp_output=True, + # PP outputs have been broadcast to all ranks at logits computation, + # so the sampled token ids must not be sent again. + if not self.broadcast_pp_output and pp.world_size > 1 and pp.is_last_rank: + self._pp_broadcast_prev_sampled_token_ids(sampler_output.sampled_token_ids) self._draft_token_ids = None self._draft_token_req_ids = None