[Perf][MOSS-TTS] Avoid synchronization in batched index transfers

Construct MTP row offsets and mixed prefill/decode stop indices on the CPU, then enqueue nonblocking device transfers before batched writes. Add CUDA regressions for synchronization, row placement, stop flags, and snapshot ownership.

Signed-off-by: Sy03 <1370724210@qq.com>

tests/worker/test_batched_omni_output.py

@@ -92,3 +92,54 @@ def test_owned_code_batch_and_embeddings_survive_next_graph_output(monkeypatch,
     returned[0][1].fill_(999)
     assert runner.model_intermediate_buffer["r1"]["codes"]["audio"].item() == 0
     assert runner.model_intermediate_buffer["r2"]["codes"]["audio"].item() == 1
+
+
+@pytest.mark.cuda
+def test_mtp_mixed_offsets_do_not_synchronize_cuda(monkeypatch):
+    import vllm_omni.worker.gpu_model_runner as mod
+    from tests.worker.test_omni_gpu_model_runner import _make_runner, _noop_forward_context
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.model.gpu_resident_buffer_keys = {("codes", "audio")}
+    for name in ("talker_mtp_input_ids", "talker_mtp_inputs_embeds", "last_talker_hidden", "text_step"):
+        buffer = getattr(runner, name)
+        buffer.gpu = buffer.gpu.to("cuda")
+    embeds = torch.ones((2, 4), device="cuda")
+    codes = torch.arange(2, device="cuda").view(2, 1)
+    runner.talker_mtp = lambda *args, **kwargs: (embeds, codes)
+    output = torch.full((5, 4), -10.0, device="cuda")
+    previous = torch.cuda.get_sync_debug_mode()
+    try:
+        torch.cuda.set_sync_debug_mode("error")
+        runner._talker_mtp_forward(["r1", "r2"], output, [1, 4])
+    finally:
+        torch.cuda.set_sync_debug_mode(previous)
+    expected = torch.full((5, 4), -10.0)
+    expected[[1, 4]] = 1.0
+    torch.testing.assert_close(output.cpu(), expected, rtol=0, atol=0)
+    codes.fill_(999)
+    assert runner.model_intermediate_buffer["r1"]["codes"]["audio"].item() == 0
+    assert runner.model_intermediate_buffer["r2"]["codes"]["audio"].item() == 1
+
+
+@pytest.mark.cuda
+def test_mixed_stop_flags_do_not_synchronize_cuda():
+    model = object.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    model.n_vq = 12
+    model.audio_pad_token_id = 1024
+    infos = [
+        {"audio_codes": {"current": torch.zeros((1, 12), device="cuda", dtype=torch.long)}},
+        {},
+        {"audio_codes": {"current": torch.full((1, 12), 1024, device="cuda", dtype=torch.long)}},
+    ]
+    hidden = torch.ones((3, 4), device="cuda")
+    previous = torch.cuda.get_sync_debug_mode()
+    try:
+        torch.cuda.set_sync_debug_mode("error")
+        output = model.make_omni_output(hidden, runtime_additional_information=infos)
+    finally:
+        torch.cuda.set_sync_debug_mode(previous)
+    assert model._batch_should_continue.tolist() == [True, True, False]
+    assert [row.shape for row in output.multimodal_outputs["codes"]["audio"]] == [(1, 12), (0, 12), (1, 12)]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -1704,7 +1704,10 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
                 self._batch_should_continue = active
             else:
                 self._batch_should_continue = torch.ones(len(info_dicts), device=hidden.device, dtype=torch.bool)
-                indices = torch.tensor(should_indices, device=hidden.device, dtype=torch.long)
+                # Keep mixed prefill/decode output packing asynchronous.
+                indices = torch.tensor(should_indices, device="cpu", dtype=torch.long).to(
+                    hidden.device, non_blocking=True
+                )
                 self._batch_should_continue.index_copy_(0, indices, active)
 
         if not have_codes:

vllm_omni/worker/gpu_model_runner.py

@@ -1971,7 +1971,11 @@ class OmniGPUModelRunner(GPUModelRunner):
         if start_offsets == list(range(decode_batch_size)):
             inputs_embeds[:decode_batch_size].copy_(req_embeds[:decode_batch_size])
         else:
-            offsets = torch.tensor(start_offsets, device=inputs_embeds.device, dtype=torch.long)
+            # A device tensor constructor waits for the preceding MTP graph.
+            # Enqueue the small host index transfer without synchronizing it.
+            offsets = torch.tensor(start_offsets, device="cpu", dtype=torch.long).to(
+                inputs_embeds.device, non_blocking=True
+            )
             inputs_embeds.index_copy_(0, offsets, req_embeds[:decode_batch_size])
         if code_predictor_codes is not None:
             if out_key in getattr(self.model, "gpu_resident_buffer_keys", set()):