Merge main and preserve dispatch backlog across monitor windows

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

tests/worker/test_fa3_scheduler_metadata_size.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Characterization of ``OmniGPUModelRunner.initialize_metadata_builders``.
 
 Upstream ``FlashAttentionMetadataBuilder`` pre-allocates ``scheduler_metadata``
@@ -19,14 +19,12 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-# The class ``super().initialize_metadata_builders`` dispatches to.
-_UPSTREAM_RUNNER = OmniGPUModelRunner.__mro__[1]
-
 
 def _builder(*, prealloc: int | None, max_num_splits: int) -> SimpleNamespace:
     sm = None if prealloc is None else torch.zeros(prealloc, dtype=torch.int32)
@@ -36,10 +34,14 @@ def _builder(*, prealloc: int | None, max_num_splits: int) -> SimpleNamespace:
 def _make_runner(monkeypatch, *, max_num_seqs: int, builders: list) -> OmniGPUModelRunner:
     # Stub the upstream initialize_metadata_builders; our fake runner already
     # carries the attn_groups the Omni override iterates over.
-    monkeypatch.setattr(_UPSTREAM_RUNNER, "initialize_metadata_builders", lambda self, *a, **k: None)
+    monkeypatch.setattr(GPUModelRunner, "initialize_metadata_builders", lambda self, *a, **k: None)
     runner = object.__new__(OmniGPUModelRunner)
     runner.scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs)
     runner.cache_config = SimpleNamespace(enable_prefix_caching=False)  # skip omni_prefix_cache branch
+    runner.vllm_config = SimpleNamespace(kv_transfer_config=None)
+    runner.model_config = SimpleNamespace()
+    runner.is_pooling_model = False
+    runner.speculative_config = None
     runner.attn_groups = [[SimpleNamespace(metadata_builders=builders)]]
     return runner
 

tests/worker/test_gpu_ar_model_runner.py

@@ -18,6 +18,7 @@ from vllm.sampling_params import SamplingParams
 from vllm.v1.worker import gpu_input_batch
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
+from vllm_omni.core.prefix_cache import ModelCachePolicy
 from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
@@ -37,6 +38,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 def _make_runner(engine_output_type: str | None, downstream_req_ids: set[str]) -> GPUARModelRunner:
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     runner.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(engine_output_type=engine_output_type),
     )
@@ -57,6 +60,8 @@ def _mtp_runner(*, async_scheduling: bool, buffers: dict[str, dict]) -> tuple[GP
         return torch.tensor([[11, 12, 13]], dtype=torch.long)
 
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     runner.use_async_scheduling = async_scheduling
     runner.model = SimpleNamespace(post_sample_talker_mtp=post_sample_talker_mtp)
     runner.requests = {req_id: SimpleNamespace() for req_id in buffers}
@@ -143,6 +148,8 @@ def test_post_sample_talker_mtp_uses_gpu_token_with_async_scheduling() -> None:
 
 def test_post_sample_talker_mtp_skips_shape_validation_without_duplex_rows() -> None:
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     runner.use_async_scheduling = False
     runner.model = SimpleNamespace(
         post_sample_talker_mtp=lambda **_: pytest.fail("hook must not run"),
@@ -178,6 +185,8 @@ def test_post_sample_talker_mtp_rejects_invalid_selected_token_shape(
     error_match: str,
 ) -> None:
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     runner.use_async_scheduling = False
     runner.model = SimpleNamespace(
         post_sample_talker_mtp=lambda **_: pytest.fail("hook must not run"),
@@ -265,6 +274,8 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
 
     received = []
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     runner.input_batch = input_batch
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
@@ -323,6 +334,8 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker():
 
 def test_runner_assisted_full_attention_metadata_request_is_opt_in():
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     runner.model = object()
     runner.scheduler_config = SimpleNamespace(max_num_seqs=16)
 
@@ -377,6 +390,8 @@ def test_runner_assisted_full_attention_metadata_request_and_context_hooks():
             calls.append(("context", {"enabled": enabled, "num_reqs": num_reqs}))
 
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     runner.model = Model()
     runner.scheduler_config = SimpleNamespace(max_num_seqs=8)
 
@@ -466,6 +481,8 @@ def test_omni_async_gpu_model_runner_output_reraises_background_exception():
 
 def _make_async_output_runner(engine_output_type: str = "audio"):
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     model_config = SimpleNamespace(
         engine_output_type=engine_output_type,
         async_chunk=True,
@@ -642,10 +659,11 @@ def test_async_omni_output_guard_requires_safe_conditions():
 
     assert GPUARModelRunner._should_use_async_omni_output(runner)
 
+    # Prefix cache no longer forces the sync output path; leftover mm is
+    # snapshotted at save so the builder can materialize a step late.
     runner.omni_prefix_cache = object()
-    assert not GPUARModelRunner._should_use_async_omni_output(runner)
+    assert GPUARModelRunner._should_use_async_omni_output(runner)
 
-    runner.omni_prefix_cache = None
     runner.model.has_postprocess = True
     assert not GPUARModelRunner._should_use_async_omni_output(runner)
 
@@ -656,6 +674,7 @@ def test_async_omni_output_guard_requires_safe_conditions():
 def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
     runner = _make_async_output_runner(engine_output_type="latent")
     runner.model.omni_pooler_payload_include_hidden = False
+    runner._pooler_payload_include_hidden_flag = False
 
     monkeypatch.setattr(
         GPUARModelRunner,
@@ -699,6 +718,7 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
 def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_only(monkeypatch):
     runner = _make_async_output_runner(engine_output_type="latent")
     runner.model.omni_pooler_payload_include_hidden = False
+    runner._pooler_payload_include_hidden_flag = False
     runner._async_chunk = False
     runner.requests = {"r1": object(), "r2": object(), "r3": object()}
 
@@ -745,6 +765,7 @@ def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monk
     """Thinker mm rows align to hidden_states.shape[0], not padded scheduled count."""
     runner = _make_async_output_runner(engine_output_type="latent")
     runner.model.omni_pooler_payload_include_hidden = False
+    runner._pooler_payload_include_hidden_flag = False
     runner._async_chunk = False
     runner.requests = {"r1": object(), "r2": object(), "r3": object()}
 
@@ -790,6 +811,7 @@ def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monk
 def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
     runner = _make_async_output_runner()
     runner.model.omni_pooler_payload_include_hidden = False
+    runner._pooler_payload_include_hidden_flag = False
 
     payload = GPUARModelRunner._build_omni_async_snapshot_payload(
         runner,
@@ -819,6 +841,8 @@ def test_runner_assisted_full_attention_metadata_refresh_pads_buffers():
             self.commits.append(num_reqs_padded)
 
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     block_table = BlockTable()
     runner.input_batch = SimpleNamespace(
         num_computed_tokens_cpu=torch.tensor([10, 20, 99, 99], dtype=torch.int32),
@@ -846,6 +870,8 @@ def test_runner_assisted_full_attention_metadata_refresh_pads_buffers():
 @pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
 def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
     runner = object.__new__(GPUARModelRunner)
+    runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
+    runner._pooler_payload_include_hidden_flag = True
     runner.execute_model_state = ExecuteModelState(
         SimpleNamespace(
             total_num_scheduled_tokens=3,
@@ -908,12 +934,11 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
     )
     monkeypatch.setattr(GPUARModelRunner, "eplb_step", lambda self: None)
     monkeypatch.setattr(GPUARModelRunner, "_resolve_pooler_payload_req_ids", lambda self, req_ids: ("audio", req_ids))
-    monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: set())
     monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
     monkeypatch.setattr(
         GPUARModelRunner,
-        "_maybe_get_combined_prefix_cache_tensors",
-        lambda *args, **kwargs: (None, None),
+        "_prepare_prefix_cache_pooler_payload_sources",
+        lambda self, **kwargs: (kwargs["staged_hidden_states_cpu"], None, None),
     )
     monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
     monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
@@ -946,8 +971,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch
         lambda self, req_ids: ("latent", req_ids),
     )
     monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
-    monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
-    monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
     monkeypatch.setattr(
         GPUARModelRunner,
         "_prepare_prefix_cache_pooler_payload_sources",
@@ -1095,8 +1118,6 @@ def test_build_omni_output_filters_multimodal_by_partial_downstream_batch(monkey
         lambda self, req_ids: ("latent", ["r1", "r3"]),
     )
     monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
-    monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
-    monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
     monkeypatch.setattr(
         GPUARModelRunner,
         "_prepare_prefix_cache_pooler_payload_sources",
@@ -1172,8 +1193,6 @@ def test_build_omni_output_uses_combined_prefix_cache_mm_payload_for_partial_dow
         lambda self, req_ids: ("latent", ["r1", "r3"]),
     )
     monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
-    monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
-    monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
     monkeypatch.setattr(
         GPUARModelRunner,
         "_prepare_prefix_cache_pooler_payload_sources",

tests/worker/test_min_tokens_stop_ids_guard.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Guard against out-of-vocabulary stop ids in min-tokens masking (#4962).
 
 vLLM's input processor folds the stage tokenizer's EOS id into
@@ -19,7 +19,7 @@ from vllm.v1.sample.logits_processor import (
     MinTokensLogitsProcessor,
 )
 
-from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
+from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids, sanitize_sampling_params_min_tokens_stop_ids
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -28,6 +28,12 @@ CODEC_EOS = 2150
 TEXT_EOS = 151645
 
 
+@pytest.fixture(autouse=True)
+def _disable_vllm_pin_memory_for_cpu(monkeypatch: pytest.MonkeyPatch):
+    """Keep the module-level H2D helper in CPU-only mode."""
+    monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False)
+
+
 def _make_min_tokens_proc(stop_token_ids: list[int], eos_token_id: int | None) -> MinTokensLogitsProcessor:
     """Build a MinTokensLogitsProcessor with one active request, mirroring
     the production path (engine folds the tokenizer EOS via
@@ -113,3 +119,14 @@ def test_guard_handles_multiple_requests():
     assert out[0, CODEC_EOS] == -float("inf")
     assert out[1, 7] == -float("inf")
     assert int((out == -float("inf")).sum()) == 2
+
+
+def test_sampling_params_guard_filters_only_unreachable_min_tokens_ids():
+    params = SamplingParams(min_tokens=2, stop_token_ids=[CODEC_EOS])
+    params.update_from_generation_config({}, TEXT_EOS)
+
+    sanitize_sampling_params_min_tokens_stop_ids(params, TALKER_VOCAB)
+
+    assert params.all_stop_token_ids == {CODEC_EOS}
+    assert params.stop_token_ids == [CODEC_EOS]
+    assert params.eos_token_id == TEXT_EOS

tests/worker/test_model_sampler_extra_args.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for per-request model sampler arguments.
+
+Custom model samplers that set ``model_sampler_wants_extra_args = True``
+(e.g. HunyuanImage3's per-request ``ar_task_mode``, #6088) receive one
+``SamplingParams.extra_args`` entry per batch row, in ``input_batch.req_ids``
+order — the same order ``_build_model_sampler_output_token_ids`` uses. Pure CPU.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.worker.sampling_utils import build_model_sampler_extra_args, call_model_sampler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _state(extra_args):
+    return SimpleNamespace(sampling_params=SimpleNamespace(extra_args=extra_args))
+
+
+def test_row_order_follows_req_ids():
+    input_batch = SimpleNamespace(req_ids=["b", "a"])
+    requests = {
+        "a": _state({"ar_task_mode": "comprehension"}),
+        "b": _state(None),
+    }
+    assert build_model_sampler_extra_args(input_batch, requests) == [None, {"ar_task_mode": "comprehension"}]
+
+
+def test_missing_request_state_yields_none():
+    assert build_model_sampler_extra_args(SimpleNamespace(req_ids=["gone"]), {}) == [None]
+
+
+def test_request_without_sampling_params_yields_none():
+    input_batch = SimpleNamespace(req_ids=["r0"])
+    assert build_model_sampler_extra_args(input_batch, {"r0": SimpleNamespace(sampling_params=None)}) == [None]
+
+
+def test_empty_batch():
+    assert build_model_sampler_extra_args(SimpleNamespace(req_ids=[]), {}) == []
+
+
+def test_opted_in_sampler_receives_per_request_extra_args(mocker):
+    model = SimpleNamespace(model_sampler_wants_extra_args=True)
+    model_sample = mocker.Mock(return_value="sampled")
+    input_batch = SimpleNamespace(req_ids=["a", "b"])
+    logits = torch.empty(0)
+    metadata = SimpleNamespace()
+    requests = {
+        "a": _state(None),
+        "b": _state({"ar_task_mode": "comprehension"}),
+    }
+
+    result = call_model_sampler(
+        model,
+        model_sample,
+        logits,
+        metadata,
+        input_batch=input_batch,
+        requests=requests,
+    )
+
+    assert result == "sampled"
+    model_sample.assert_called_once_with(
+        logits,
+        metadata,
+        per_req_extra_args=[None, {"ar_task_mode": "comprehension"}],
+    )
+
+
+def test_non_opted_in_sampler_keeps_two_argument_contract(mocker):
+    model = SimpleNamespace(model_sampler_wants_extra_args=False)
+    model_sample = mocker.Mock(return_value="sampled")
+    logits = torch.empty(0)
+    metadata = SimpleNamespace()
+
+    result = call_model_sampler(
+        model,
+        model_sample,
+        logits,
+        metadata,
+        input_batch=SimpleNamespace(req_ids=["a"]),
+        requests={"a": _state({"ar_task_mode": "comprehension"})},
+    )
+
+    assert result == "sampled"
+    model_sample.assert_called_once_with(logits, metadata)

tests/worker/test_native_duplex_hooks.py

@@ -2639,7 +2639,7 @@ _REQUIRED_DUPLEX_CALLS = {
     "_sample": "_apply_duplex_sampling",
 }
 
-_MODEL_SAMPLER_CALL = "model_sample"
+_MODEL_SAMPLER_CALL = "call_model_sampler"
 
 
 def _repo_root():

tests/worker/test_omni_connector_mixin.py

@@ -397,7 +397,7 @@ class TestLoadCustomFuncSelection(unittest.TestCase):
     def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self):
         incompatible_paths = [
             "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav",
-            "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit",
+            "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2diffusion",
             "vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow",
             "vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion",
         ]
@@ -1724,3 +1724,16 @@ class TestSendRetry(unittest.TestCase):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+@pytest.mark.parametrize(
+    "fallback,external,expected", [("r", None, "mapped"), (None, "ext", "ext"), (None, None, None)]
+)
+def test_resolve_request_id_requires_internal_or_external_id(fallback, external, expected):
+    host = SimpleNamespace(_request_ids_mapping={"r": "mapped"})
+    request = SimpleNamespace(external_req_id=external)
+    if expected is None:
+        with pytest.raises(ValueError, match="request ID"):
+            OmniConnectorModelRunnerMixin._resolve_external_req_id(host, request, fallback)
+    else:
+        assert OmniConnectorModelRunnerMixin._resolve_external_req_id(host, request, fallback) == expected

tests/worker/test_omni_gpu_model_runner.py

@@ -770,8 +770,10 @@ def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeh
 def test_full_payload_output_accumulation_hook_matrix():
     assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output()
     assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output()
-    assert not _make_full_payload_accumulation_runner(
-        model_stage="code2wav", final_output=True
+    # A stage may publish a user-visible output and still feed another stage
+    # (Qwen3-Omni thinker publishes text while sending audio state to talker).
+    assert _make_full_payload_accumulation_runner(
+        model_stage="thinker", final_output=True
     )._should_accumulate_full_payload_output()
     assert not _make_full_payload_accumulation_runner(
         model_stage="token2audio",

tests/worker/test_omni_v2_profile_dispatch.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+from vllm.compilation.cuda_graph import CUDAGraphMode
+
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("dp_size", [1, 2])
+def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.cudagraph_manager = None
+    runner.dp_size = dp_size
+    runner.dp_rank = 1
+
+    import vllm.v1.worker.gpu.dp_utils as dp_utils
+
+    expected_tokens = object()
+    sync_calls = []
+
+    def sync_padding(*args, **kwargs):
+        sync_calls.append((args, kwargs))
+        return args[1], expected_tokens
+
+    monkeypatch.setattr(
+        dp_utils,
+        "sync_cudagraph_and_dp_padding",
+        sync_padding,
+    )
+
+    batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
+        num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=True, max_query_len=8
+    )
+    assert batch_desc.cg_mode == CUDAGraphMode.NONE
+    if dp_size == 1:
+        assert num_tokens_across_dp is None
+        assert sync_calls == []
+    else:
+        assert num_tokens_across_dp is expected_tokens
+
+    # Non-eager dispatch goes through the cudagraph manager.
+    expected = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=8, num_reqs=1)
+    dispatch_calls = []
+
+    def dispatch(*args, **kwargs):
+        dispatch_calls.append((args, kwargs))
+        return expected
+
+    runner.cudagraph_manager = SimpleNamespace(dispatch=dispatch)
+    runner.dp_size = 1
+    batch_desc, _ = runner._dispatch_batch_descriptor(
+        num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=False, max_query_len=8
+    )
+    assert batch_desc is expected
+    assert dispatch_calls == [((1, 8, 8), {"num_active_loras": 0, "max_query_len": 8})]

tests/worker/test_tts_sampling_seed.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.worker import sampling_utils
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_fixed_seed_contract_separates_stage_sampler_from_talker_mtp_rng() -> None:
+    assert hasattr(sampling_utils, "apply_fixed_seed_to_sampling_params")
+    assert hasattr(sampling_utils, "get_tts_local_seed")
+
+    stage_only = SimpleNamespace(seed=None, extra_args=None)
+    sampling_utils.apply_fixed_seed_to_sampling_params(
+        stage_only,
+        17,
+        seed_talker_mtp=False,
+    )
+    assert stage_only.seed == 17
+    assert sampling_utils.get_tts_local_seed(stage_only) is None
+
+    talker = SimpleNamespace(seed=None, extra_args={"preserved": True})
+    sampling_utils.apply_fixed_seed_to_sampling_params(
+        talker,
+        23,
+        seed_talker_mtp=True,
+    )
+    assert talker.seed == 23
+    assert talker.extra_args == {"preserved": True, "tts_local_seed": 23}
+    assert sampling_utils.get_tts_local_seed(talker) == 23