[Model] Default MOSS Local 1.5 to adaptive CUDA MRV2 with MPS (#8438)

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

tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py

@@ -318,26 +318,45 @@ def test_failed_output_handoff_retains_decoded_samples(cuda):
     fast.close()
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("enabled,complete", [(False, False), (True, False), (True, True), (True, None)])
+def test_gate_waits_only_for_enabled_inflight_decode(mocker, enabled, complete):
+    # A fixed GPU sleep cannot prove that an event is still pending when the
+    # host reaches gate(), particularly across CUDA and ROCm implementations.
+    fast = object.__new__(MossFirstChunkFastPath)
+    fast._device = torch.device("cuda")
+    fast._gate_main = enabled
+    event = None if complete is None else mocker.Mock()
+    if event is not None:
+        event.query.return_value = complete
+    fast._inflight = event
+    stream = mocker.Mock()
+    current_stream = mocker.patch("torch.cuda.current_stream", return_value=stream)
+    fast.gate()
+    if enabled and complete is False:
+        current_stream.assert_called_once_with(fast._device)
+        stream.wait_event.assert_called_once_with(event)
+    else:
+        current_stream.assert_not_called()
+        stream.wait_event.assert_not_called()
+
+
 @pytest.mark.cuda
-def test_gate_orders_main_stream_after_inflight_fast_decode(cuda):
+def test_gate_hands_inflight_gpu_output_to_main_stream(cuda):
     fast, _, _ = _fast(cuda)
     side = torch.cuda.Stream()
+    output = torch.zeros(16, device=cuda)
+    side.wait_stream(torch.cuda.current_stream())
     with torch.cuda.stream(side):
-        torch.cuda._sleep(200_000_000)
+        output.fill_(7)
         pending = torch.cuda.Event()
         pending.record(side)
     fast._inflight = pending
-    fast.gate()  # disabled: no wait
-    unordered = torch.cuda.Event()
-    unordered.record()
     fast._gate_main = True
     fast.gate()
-    ordered = torch.cuda.Event()
-    ordered.record()
-    unordered.synchronize()
-    assert not pending.query() and not ordered.query()
-    ordered.synchronize()
-    assert pending.query()
+    received = output.clone()
+    torch.testing.assert_close(received, torch.full_like(received, 7))
+    fast.close()
 
 
 def _cpu_fast(monkeypatch, *, scheduler=None):
@@ -545,3 +564,41 @@ def test_platform_failure_wakes_pending_handoff(monkeypatch):
     assert fast._closed
     with pytest.raises(RuntimeError, match="decode failed"):
         fast.order_after(0)
+
+
+@pytest.mark.cpu
+@pytest.mark.tts
+@pytest.mark.parametrize("limit", [1, 2])
+def test_congested_admission_preserves_routes_and_slots_then_recovers(monkeypatch, limit):
+    import sys
+    from functools import partial
+
+    helpers = sys.modules[__name__]
+    from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import MossFirstChunkFastPath
+
+    monkeypatch.setattr(helpers, "MossFirstChunkFastPath", partial(MossFirstChunkFastPath, max_active_streams=limit))
+    fast, session, sink, _ = helpers._cpu_fast(monkeypatch)
+    calls = []
+
+    class Sink:
+        def prepare(self, ids):
+            calls.append(ids)
+            return sink.prepare(ids)
+
+    fast._sink = Sink()
+    thread = threading.Thread()
+    monkeypatch.setattr(thread, "is_alive", lambda: True)
+    fast._thread = thread
+    slots = {f"busy{i}": session.acquire() for i in range(limit)}
+    before_slots, before_free = slots.copy(), session.free.copy()
+    codes = torch.tensor([1, 2])
+    assert not fast.submit("r", "r", codes, slots)
+    assert slots == before_slots and session.free == before_free
+    assert not calls and not fast._decoded and not fast._handoffs and fast._jobs.empty()
+    released = slots.pop("busy0")
+    session.release(released)
+    assert fast.submit("r", "r", codes, slots)
+    codes.fill_(99)
+    job = fast._jobs.get_nowait()
+    assert job.slot == slots["r"] and calls == [["r"]]
+    assert job.codes.tolist() == [[1], [2]]

tests/model_executor/models/moss_tts/test_local_compile_config.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Deployment sampler settings override the legacy environment switch."""
+
+import pytest
+from transformers import GPT2Config
+
+from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local_depth as depth
+
+pytestmark = [pytest.mark.core_model, pytest.mark.tts, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    "configured,env,expected", [(None, "0", False), (None, "1", True), (True, "0", True), (False, "1", False)]
+)
+def test_sampler_compile_setting_and_idempotent_setup(monkeypatch, configured, env, expected):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", env)
+    monkeypatch.setattr(depth.current_omni_platform, "supports_torch_inductor", lambda: True)
+    calls = []
+
+    def compile_fn(fn, **kwargs):
+        calls.append((fn, kwargs))
+        return fn
+
+    monkeypatch.setattr(depth.torch, "compile", compile_fn)
+    model = depth.MossTTSLocalDepthTransformer(
+        GPT2Config(n_embd=80, n_head=1, n_inner=160), compile_audio_sampler=configured
+    )
+    model.setup_compile()
+    model.setup_compile()
+    assert len(calls) == 1 + expected
+    assert (model._compiled_audio_sampler is depth._sample_token) is expected
+    sampler_calls = [kwargs for fn, kwargs in calls if fn is depth._sample_token]
+    assert len(sampler_calls) == int(expected)
+    if expected:
+        assert sampler_calls[0] == {"fullgraph": True, "dynamic": True, "options": {"fallback_random": True}}
+
+
+def test_sampler_compile_setting_respects_platform_support(monkeypatch):
+    monkeypatch.setattr(depth.current_omni_platform, "supports_torch_inductor", lambda: False)
+
+    def forbidden(*args, **kwargs):
+        raise AssertionError("compile called on unsupported platform")
+
+    monkeypatch.setattr(depth.torch, "compile", forbidden)
+    model = depth.MossTTSLocalDepthTransformer(GPT2Config(n_embd=80, n_head=1, n_inner=160), compile_audio_sampler=True)
+    model.setup_compile()
+    assert model._compiled_audio_sampler is None

tests/model_executor/models/moss_tts/test_local_model_state.py

@@ -40,14 +40,18 @@ def _model(device):
     model.audio_assistant_slot_token_id, model.im_end_token_id = 2, 3
     model.talker_mtp_output_key = ("audio_codes", "current")
     model.talker_mtp_graph_safe = False
+    model.talker_mtp_accepts_per_row_generators = True
     model.gpu_resident_buffer_keys = {("hidden_states", "last"), ("audio_codes", "current")}
     model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
     model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).to(dtype) / 8
     model.force_stop = False
 
-    def frame(hidden, *_args, generator=None, temperature=None, top_k=None, top_p=None, **_kwargs):
+    def frame(hidden, *_args, generator=None, generators=None, temperature=None, top_k=None, top_p=None, **_kwargs):
         assert (temperature, top_k, top_p) == (1.7, 25, 0.8)
-        noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator)
+        if generators is not None:
+            noise = torch.cat([torch.randint(0, 8, (1, 2), device=device, generator=g) for g in generators])
+        else:
+            noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator)
         codes = (hidden[:, :2].mul(16).long() + noise) % 8
         return torch.full((hidden.shape[0],), not model.force_stop, device=device), codes
 
@@ -85,11 +89,11 @@ def _admit(state, slot, name, seed):
     state.intermediate_buffer.buffers[slot]["codes"] = {"ref": torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])}
 
 
-def _batch(device, slots, counts):
+def _batch(device, slots, counts, index_dtype=torch.int32):
     starts = np.array([0, *np.cumsum(counts)], dtype=np.int32)
     return SimpleNamespace(
         idx_mapping_np=np.array(slots),
-        idx_mapping=torch.tensor(slots, device=device),
+        idx_mapping=torch.tensor(slots, device=device, dtype=index_dtype),
         num_reqs=len(slots),
         num_tokens=int(starts[-1]),
         num_scheduled_tokens=np.array(counts),
@@ -122,7 +126,8 @@ def _step(state, batch, req_states, dispatcher=None):
 
 @pytest.mark.parametrize("seed", [None, 17])
 @pytest.mark.parametrize("batch_prefill", [False, True])
-def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill):
+@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64])
+def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill, index_dtype):
     reference, candidate = (_state(cls, device) for cls in (OmniModelState, MossLocalModelState))
     candidate._batch_prefill = batch_prefill
     for state in (reference, candidate):
@@ -154,7 +159,7 @@ def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, see
         for state in (reference, candidate):
             state.model.force_stop = index == 3
             torch.manual_seed(91 + index)
-            outputs.append(_step(state, _batch(device, slots, counts), req_states))
+            outputs.append(_step(state, _batch(device, slots, counts, index_dtype), req_states))
         a, b = outputs
         torch.testing.assert_close(a[0], b[0], rtol=0, atol=0)
         torch.testing.assert_close(a[2], b[2], rtol=0, atol=0)
@@ -190,6 +195,23 @@ def test_slot_honors_external_stopping_control(device):
         assert output[1]["codes"]["audio"][0].tolist() == [[8, 8]]
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("state_cls", [OmniModelState, MossLocalModelState])
+def test_local_seeded_mtp_keeps_per_row_generators_in_one_batch(state_cls):
+    state = _state(state_cls, torch.device("cpu"))
+    generators = [torch.Generator().manual_seed(seed) for seed in (17, 29)]
+    inputs = [torch.zeros(2, dtype=torch.long), *[torch.ones(2, 4) for _ in range(3)]]
+    frame = MagicMock(wraps=state.model.local_transformer.generate_frame)
+    state.model.local_transformer.generate_frame = frame
+
+    state._call_mtp_with_sampling(*inputs, buffers=[{}, {}], req_ids=["first", "second"], generators=generators)
+
+    frame.assert_called_once()
+    assert frame.call_args.args[0].shape[0] == 2
+    assert frame.call_args.kwargs["generators"] == generators
+    assert frame.call_args.kwargs["generator"] is None
+
+
 @pytest.mark.cuda
 def test_slot_mtp_graph_padding_reorder_and_owned_output():
     if not torch.cuda.is_available():

tests/model_executor/models/moss_tts/test_local_prefill_batch.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Reference batching preserves coordinates and asynchronous source ownership."""
 
-from types import SimpleNamespace
+from dataclasses import dataclass
 
 import numpy as np
 import pytest
@@ -15,6 +15,12 @@ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import Mo
 pytestmark = pytest.mark.core_model
 
 
+@dataclass
+class _PrefillPositions:
+    prompt_len: np.ndarray
+    num_computed_tokens: np.ndarray
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("flat", [False, True])
 def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mocker, flat):
@@ -28,7 +34,8 @@ def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mock
             "ref_offset": offset,
         }
     batch = _batch(torch.device("cpu"), [3, 0], [2, 1])
-    req = SimpleNamespace(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5))
+    req = _PrefillPositions(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5))
+    req.num_computed_tokens[[3, 0]] = [1, 2]
     embeds = state._static_inputs_embeds[:3]
     expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[[1, 2, 2]])
     embed = mocker.spy(state.model.model.embed_tokens, "forward")
@@ -40,6 +47,65 @@ def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mock
     assert state.intermediate_buffer.buffers[0]["ref_offset"] == 3
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("batch_prefill", [False, True])
+@pytest.mark.parametrize("cached,count", [(2, 1), (1, 2)])
+def test_cached_prefix_prefill_uses_absolute_reference_position(batch_prefill, cached, count):
+    state = _state(MossLocalModelState, torch.device("cpu"))
+    state._batch_prefill = batch_prefill
+    codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])
+    state.intermediate_buffer.buffers[0] = {"req_id": "prefix-hit", "codes": {"ref": codes}}
+    batch = _batch(torch.device("cpu"), [0], [count])
+    req = _PrefillPositions(prompt_len=np.full(5, 4), num_computed_tokens=np.full(5, cached))
+    embeds = state._static_inputs_embeds[:count]
+    expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[cached : cached + count])
+    with torch.inference_mode():
+        state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req)
+    torch.testing.assert_close(embeds, expected, rtol=0, atol=0)
+    assert state.intermediate_buffer.buffers[0]["ref_offset"] == cached + count
+
+
+@pytest.mark.cpu
+def test_scalar_prefill_uses_cached_prompt_position():
+    state = _state(MossLocalModelState, torch.device("cpu"))
+    codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])
+    ids = torch.tensor([2])
+    expected = state.model.embed_input_ids(ids) + state.model._audio_embed(codes[2:3])
+    with torch.inference_mode():
+        _, actual, updates = state.model.preprocess(
+            ids, None, codes={"ref": codes}, _omni_is_prefill=True, _omni_num_computed_tokens=2
+        )
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    assert updates["ref_offset"] == 3
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("batch_prefill", [False, True])
+@pytest.mark.parametrize("omit_trailing_pad", [False, True])
+def test_prefill_chunk_crosses_reference_end(batch_prefill, omit_trailing_pad):
+    state = _state(MossLocalModelState, torch.device("cpu"))
+    state._batch_prefill = batch_prefill
+    # codes.ref uses prompt coordinates, including PAD rows for text tokens.
+    codes = torch.tensor([[8, 8], [1, 2], [3, 4], [5, 6], [8, 8], [8, 8]])
+    if omit_trailing_pad:
+        codes = codes[:4]
+    state.model._audio_embed = lambda rows: rows[:, :1].masked_fill(rows[:, :1] == 8, 0).expand(-1, 4) / 8
+    state.intermediate_buffer.buffers[0] = {"req_id": "reference-boundary", "codes": {"ref": codes}}
+    req = _PrefillPositions(prompt_len=np.full(5, 6), num_computed_tokens=np.zeros(5))
+    outputs = []
+    with torch.inference_mode():
+        for start, count in [(0, 2), (2, 4)]:
+            req.num_computed_tokens[0] = start
+            batch = _batch(torch.device("cpu"), [0], [count])
+            embeds = state._static_inputs_embeds[:count]
+            state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req)
+            outputs.append(embeds.clone())
+    expected = state.model.embed_input_ids(torch.full((6,), 2))
+    expected[1:4] += torch.tensor([1, 3, 5]).reshape(-1, 1) / 8
+    torch.testing.assert_close(torch.cat(outputs), expected, rtol=0, atol=0)
+    assert state.intermediate_buffer.buffers[0]["ref_offset"] == 6
+
+
 @pytest.mark.cuda
 def test_pending_reference_uploads_survive_staging_reuse_and_overflow():
     if not torch.cuda.is_available():