[Bugfix] Drain the omni connector output on the main thread in the AR runner (#7750)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>

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

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"

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