[Bugfix] Restore the async PP sampled-token broadcast in the AR runner (#7393)

Signed-off-by: Deep Shah <deep@socratic.co>

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

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")

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