[Core] Performance fix for orchestrator bottleneck: separates inter-stage from client outputs (#4527)

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Chenguang ZHENG <645327136@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/worker/test_gpu_ar_model_runner.py

@@ -115,6 +115,7 @@ def _make_async_output_runner(engine_output_type: str = "audio"):
     )
     runner.vllm_config = SimpleNamespace(model_config=model_config)
     runner.model_config = model_config
+    runner._async_chunk = True
     runner.omni_prefix_cache = None
     runner.requests = {"r1": object(), "r2": object()}
     runner.supports_mm_inputs = False
@@ -174,8 +175,10 @@ def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monke
     )
 
     assert output.req_ids == ["r1", "r2"]
-    assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0]]))
-    assert torch.equal(output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
+    assert output.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["hidden"], torch.tensor([[1.0]]))
+    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"]
@@ -217,10 +220,11 @@ def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatc
         query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long),
     )
 
-    assert output.multimodal_outputs is not None
-    assert output.multimodal_outputs[0] == {}
-    assert torch.equal(output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
-    assert output.multimodal_outputs[2] == {}
+    assert output.inter_stage_outputs is not None
+    assert output.inter_stage_outputs[0] is None
+    assert torch.equal(output.inter_stage_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
+    assert output.inter_stage_outputs[2] is None
+    assert output.multimodal_outputs is None
 
 
 def test_process_additional_information_uses_snapshot_request_order(monkeypatch):
@@ -331,10 +335,11 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
         query_start_loc_cpu=torch.tensor([0], dtype=torch.long),
     )
 
-    assert output.multimodal_outputs is not None
-    assert len(output.multimodal_outputs) == 1
-    assert "hidden" not in output.multimodal_outputs[0]
-    assert torch.equal(output.multimodal_outputs[0]["codes.audio"], torch.tensor([[7, 8], [9, 10]], dtype=torch.long))
+    assert output.inter_stage_outputs is not None
+    assert len(output.inter_stage_outputs) == 1
+    assert "hidden" not in output.inter_stage_outputs[0]
+    assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], torch.tensor([[7, 8], [9, 10]], dtype=torch.long))
+    assert output.multimodal_outputs is None
 
 
 def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
@@ -395,6 +400,7 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
     runner.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(engine_output_type="audio"),
     )
+    runner._async_chunk = False
 
     monkeypatch.setattr(
         GPUARModelRunner, "_sample", lambda self, logits, spec_decode_metadata: SimpleNamespace(sampled_token_ids=[])
@@ -428,6 +434,8 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
 
     output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
 
+    assert output.inter_stage_outputs is None
+    assert output.multimodal_outputs is not None
     assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
     assert torch.equal(
         output.multimodal_outputs[1]["hidden"],
@@ -482,5 +490,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )
 
-    assert torch.equal(output.multimodal_outputs[0]["codes.audio"], codes[0:1])
-    assert torch.equal(output.multimodal_outputs[1]["codes.audio"], codes[1:2])
+    assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1])
+    assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2])
+    assert output.multimodal_outputs is None