[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)]