From 5e6b9047c95eca49fa15587afd991b59e5e89fe0 Mon Sep 17 00:00:00 2001 From: Joshna-Medisetty Date: Wed, 23 Sep 2026 11:36:56 -0700 Subject: [PATCH] [Bugfix] Drain the omni connector output on the main thread in the AR runner (#7750) Signed-off-by: Joshna Medisetty --- tests/worker/test_gpu_ar_model_runner.py | 8 ++++---- vllm_omni/deploy/qwen3_omni_moe.yaml | 1 + vllm_omni/worker/gpu_ar_model_runner.py | 11 ++++++++--- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index d09282035..5e07706d2 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -504,7 +504,7 @@ def _make_async_output_runner(engine_output_type: str = "audio"): return runner -def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monkeypatch): +def test_build_omni_output_uses_snapshots_after_accumulation(monkeypatch): runner = _make_async_output_runner() events = [] @@ -519,10 +519,11 @@ def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monke "accumulate_full_payload_output", lambda self, rid, payload, request: events.append(f"accumulate:{rid}"), ) + # The connector drain runs a TP collective and belongs to the caller's thread. monkeypatch.setattr( GPUARModelRunner, "get_omni_connector_output", - lambda self: events.append("connector") or "connector-output", + lambda self: pytest.fail("builder must not drain the connector"), ) output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( @@ -554,8 +555,7 @@ def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monke assert torch.equal(output.inter_stage_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]])) assert output.multimodal_outputs is None assert output.kv_extracted_req_ids == ["r2"] - assert output.omni_connector_output == "connector-output" - assert events == ["accumulate:r1", "accumulate:r2", "connector"] + assert events == ["accumulate:r1", "accumulate:r2"] def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatch): diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index 69f0fc42e..151e90ad6 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -135,5 +135,6 @@ platforms: devices: "4" - stage_id: 2 gpu_memory_utilization: 0.3 + enforce_eager: true max_cudagraph_capture_size: 0 devices: "4" diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 0dd393ac6..0a388d1f9 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1910,8 +1910,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex cudagraph_stats=cudagraph_stats, ) output.kv_extracted_req_ids = kv_extracted_req_ids - with record_function_or_nullcontext("omni_output_builder:get_omni_connector_output"): - output.omni_connector_output = self.get_omni_connector_output() output.routed_experts = routed_experts_lists return output @@ -2126,12 +2124,17 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex multimodal_outputs=multimodal_outputs, ) + # Runs a TP collective, so it must stay on the main thread: the builder + # below runs on the async output thread, which would race execute_model(). + with record_function_or_nullcontext("omni_async_output:get_omni_connector_output"): + omni_connector_output = self.get_omni_connector_output() + def output_builder() -> OmniModelRunnerOutput: if output_tensor_snapshot.async_payload is not None: with record_function_or_nullcontext("omni_async_output:wait_cpu_payload"): output_tensor_snapshot.async_payload.wait() with record_function_or_nullcontext("omni_output_builder:total"): - return self._build_omni_model_runner_output_from_snapshot( + output = self._build_omni_model_runner_output_from_snapshot( scheduler_output=scheduler_output_snapshot, hidden_states=output_tensor_snapshot.hidden_states, staged_hidden_states_cpu=output_tensor_snapshot.staged_hidden_states_cpu, @@ -2151,6 +2154,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex postprocess_already_applied=omni_postprocess_already_applied, prefix_cache_step_id=prefix_cache_step_id, ) + output.omni_connector_output = omni_connector_output + return output if not use_async_omni_output: output = output_builder()