diff --git a/tests/worker/test_batched_omni_output.py b/tests/worker/test_batched_omni_output.py index 86e598718..b9b86991a 100644 --- a/tests/worker/test_batched_omni_output.py +++ b/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)] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 7d743603f..9d2d2d171 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/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: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index ae8de37b8..03913aebe 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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()):