[Bugfix] Route text-only chat as per-request comprehension in HunyuanImage3 AR sampler (#6111)

Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com>

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.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
 """Unit tests for HunyuanImage3 AR sampler logic (stage transitions,
 ratio restriction, comprehension blocking)."""
 
@@ -21,6 +21,10 @@ RATIO_END = 210
 RATIO_OTHER_START = 220
 RATIO_OTHER_END = 223
 TIMESTEP = 224
+CFG = 110
+IMG_TOK = 111
+JOINT_SEP = 112
+EOI = 113
 
 
 class FakeSamplerModel:
@@ -47,10 +51,13 @@ class FakeSamplerModel:
             self._stage_transitions[END_OF_THINK] = [RECAPTION]
             self._stage_transitions[END_OF_RECAPTION] = [ANSWER, BOI, SIZE_TOKEN]
 
-        self._blocked_token_ids: set[int] = set()
-        if is_comprehension:
-            self._blocked_token_ids.update([BOI, SIZE_TOKEN])
-            self._blocked_token_ids.update(self._all_ratio_ids)
+        # Mirrors the model: the blocklist is always built so generation-mode
+        # engines can apply it to per-request comprehension rows (#6088).
+        self._cfg_token_id = CFG
+        self._blocked_token_ids: set[int] = {BOI, EOI, SIZE_TOKEN, CFG, IMG_TOK, TIMESTEP, JOINT_SEP}
+        self._blocked_token_ids.update(self._all_ratio_ids)
+        self._blocked_token_ids_tensor = None
+        self._sampler = None
 
     # Bind the real methods from the model class.
     from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import (
@@ -59,6 +66,8 @@ class FakeSamplerModel:
 
     _get_forced_token = _Real._get_forced_token
     _apply_ratio_restriction = _Real._apply_ratio_restriction
+    _row_is_comprehension = staticmethod(_Real._row_is_comprehension)
+    sample = _Real.sample
 
 
 class FakeTokenEmbedding:
@@ -250,3 +259,114 @@ class TestForceEosAfterRatio:
         assert logits[0, EOS].item() == 0
         non_eos_max = logits[0, :EOS].max().item()
         assert non_eos_max == min_score
+
+
+class _RecordingSampler:
+    """Stub sampler: returns the (mutated) logits so tests can assert on them."""
+
+    def __call__(self, *, logits, sampling_metadata):
+        return logits
+
+
+def _metadata(output_token_ids):
+    from types import SimpleNamespace
+
+    return SimpleNamespace(output_token_ids=output_token_ids)
+
+
+class TestPerRequestTaskMode:
+    """Generation-mode engines must honor per-request comprehension rows (#6088).
+
+    A text chat (I2T/T2T) served by an AR+DiT deployment is marked with
+    extra_args={"ar_task_mode": "comprehension"}; its rows get the stateless
+    image-token blocking instead of stage-transition forcing, while untagged
+    rows keep the generation behavior.
+    """
+
+    VOCAB = 300
+
+    def _run(self, model, histories, per_req_extra_args):
+        logits = torch.zeros(len(histories), self.VOCAB)
+        logits[:, 50] = 7.0  # ordinary text token, must survive comprehension
+        model._sampler = _RecordingSampler()
+        out = model.sample(logits, _metadata(histories), per_req_extra_args=per_req_extra_args)
+        return out
+
+    def test_comprehension_row_skips_forcing_and_blocks_image_tokens(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        min_score = torch.finfo(torch.float32).min
+
+        logits = self._run(
+            model,
+            histories=[[END_OF_THINK], [END_OF_THINK]],
+            per_req_extra_args=[{"ar_task_mode": "comprehension"}, None],
+        )
+
+        # Row 0 (comprehension): no </think> -> <recaption> forcing (ordinary
+        # text token survives, which a force-fill would have erased) and every
+        # DiT scaffold token is masked.
+        assert logits[0, 50].item() == 7.0
+        for tid in (BOI, EOI, SIZE_TOKEN, CFG, IMG_TOK, TIMESTEP, JOINT_SEP, RATIO_START, RATIO_OTHER_START):
+            assert logits[0, tid].item() == min_score, f"token {tid} not blocked"
+
+        # Row 1 (generation): forced to <recaption> exactly as before.
+        assert logits[1, RECAPTION].item() == 0
+        assert logits[1, 50].item() == min_score
+
+    def test_no_extra_args_keeps_generation_behavior(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        min_score = torch.finfo(torch.float32).min
+
+        logits = self._run(model, histories=[[END_OF_THINK]], per_req_extra_args=None)
+
+        assert logits[0, RECAPTION].item() == 0
+        assert logits[0, 50].item() == min_score
+
+    def test_empty_or_foreign_extra_args_keeps_generation_behavior(self):
+        model = FakeSamplerModel(is_comprehension=False)
+
+        logits = self._run(
+            model,
+            histories=[[END_OF_THINK], [END_OF_THINK]],
+            per_req_extra_args=[{}, {"other_key": 1}],
+        )
+
+        assert logits[0, RECAPTION].item() == 0
+        assert logits[1, RECAPTION].item() == 0
+
+    def test_comprehension_row_without_trigger_history(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        min_score = torch.finfo(torch.float32).min
+
+        logits = self._run(
+            model,
+            histories=[[1, 2, 3]],
+            per_req_extra_args=[{"ar_task_mode": "comprehension"}],
+        )
+
+        assert logits[0, 50].item() == 7.0
+        assert logits[0, CFG].item() == min_score
+
+    def test_static_comprehension_engine_blocks_cfg(self):
+        model = FakeSamplerModel(is_comprehension=True)
+        min_score = torch.finfo(torch.float32).min
+
+        logits = self._run(model, histories=[[1, 2]], per_req_extra_args=None)
+
+        for tid in (BOI, EOI, SIZE_TOKEN, CFG, IMG_TOK, TIMESTEP, JOINT_SEP):
+            assert logits[0, tid].item() == min_score
+        assert logits[0, 50].item() == 7.0
+
+    def test_extra_args_shorter_than_batch_defaults_to_generation(self):
+        model = FakeSamplerModel(is_comprehension=False)
+
+        logits = self._run(
+            model,
+            histories=[[END_OF_THINK], [END_OF_THINK]],
+            per_req_extra_args=[{"ar_task_mode": "comprehension"}],
+        )
+        min_score = torch.finfo(torch.float32).min
+
+        assert logits[0, 50].item() == 7.0  # tagged row: comprehension
+        assert logits[1, RECAPTION].item() == 0  # missing entry: generation
+        assert logits[1, 50].item() == min_score

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 Unit tests for OmniOpenAIServingChat sampling params handling.
 
@@ -8,6 +9,7 @@ are correctly applied to the comprehension stage while preserving YAML defaults.
 
 import asyncio
 from types import SimpleNamespace
+from typing import Any, cast
 
 import pytest
 from pytest_mock import MockerFixture
@@ -326,6 +328,64 @@ def test_mixed_consumer_keeps_root_common_args_with_nested_extras(mock_engine_cl
     assert captured["prompt"]["negative_prompt"] == "avoid blur"
 
 
+def test_text_only_request_reaches_engine_with_comprehension_task_mode(mock_engine_client, mocker: MockerFixture):
+    """Exercise the production chat path, not only the tagging helper."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    mock_engine_client.stage_configs = [
+        SimpleNamespace(stage_type="llm", is_comprehension=True),
+        SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+    ]
+    mock_engine_client.default_sampling_params_list = [
+        SamplingParams(),
+        OmniDiffusionSamplingParams(),
+    ]
+    mock_engine_client.output_modalities = ["text", "image"]
+    mock_engine_client.errored = False
+    mock_engine_client.renderer = SimpleNamespace(get_tokenizer=lambda: object())
+    captured: dict[str, object] = {}
+
+    async def results():
+        if False:
+            yield None
+
+    def generate(**kwargs):
+        captured.update(kwargs)
+        return results()
+
+    mock_engine_client.generate = generate
+    serving_chat = build_serving_chat(
+        engine_client=mock_engine_client,
+        models=SimpleNamespace(model_name=lambda _: "test"),
+        online_renderer=SimpleNamespace(validate_chat_template=lambda **_: None),
+        trust_request_chat_template=True,
+    )
+    serving_chat._diffusion_mode = False
+    serving_chat._diffusion_extra_body_params = frozenset()
+    mocker.patch.multiple(
+        serving_chat,
+        _check_model=mocker.AsyncMock(return_value=None),
+        _maybe_get_adapters=mocker.Mock(return_value=None),
+        _effective_chat_template_kwargs=mocker.Mock(return_value={}),
+        _preprocess_chat=mocker.AsyncMock(return_value=([], [{"prompt": "raw"}])),
+        _base_request_id=mocker.Mock(return_value="test"),
+        _log_inputs=mocker.Mock(),
+        chat_completion_full_generator=mocker.AsyncMock(return_value="done"),
+    )
+    request = ChatCompletionRequest(
+        model="test",
+        messages=[{"role": "user", "content": "describe this image"}],
+        modalities=["text"],
+    )
+
+    assert asyncio.run(serving_chat._create_chat_completion(request)) == "done"
+
+    sampling_params_list = cast(list[Any], captured["sampling_params_list"])
+    assert sampling_params_list[0].extra_args == {"ar_task_mode": "comprehension"}
+    assert sampling_params_list[1].extra_args == {}
+    assert captured["output_modalities"] == ["text"]
+
+
 @pytest.fixture
 def mock_request(mocker: MockerFixture):
     """Create a mock request with all OpenAI sampling params set to None."""
@@ -865,3 +925,54 @@ class TestResolveHeightWidth:
         h, w = OmniOpenAIServingChat._resolve_height_width_from_extra_body({"size": "invalid"})
         assert h is None
         assert w is None
+
+
+# Tests for _apply_text_chat_ar_task_mode (#6088)
+
+
+def _tag_request(modalities):
+    return SimpleNamespace(modalities=modalities)
+
+
+def _apply_tag(params, modalities):
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    OmniOpenAIServingChat._apply_text_chat_ar_task_mode(params, _tag_request(modalities))
+    return params
+
+
+@pytest.mark.parametrize("modalities", [[], ["text"]])
+def test_text_only_chat_tags_ar_stage_as_comprehension(modalities):
+    ar = SamplingParams()
+    params = _apply_tag([ar], modalities)
+    assert params[0].extra_args == {"ar_task_mode": "comprehension"}
+
+
+@pytest.mark.parametrize("modalities", [["image"], ["text", "audio"], ["video"]])
+def test_non_text_output_request_is_untouched(modalities):
+    ar = SamplingParams()
+    _apply_tag([ar], modalities)
+    assert ar.extra_args is None
+
+
+def test_diffusion_stage_params_are_untouched():
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    ar = SamplingParams()
+    dit = OmniDiffusionSamplingParams()
+    dit_extra_before = getattr(dit, "extra_args", None)
+    _apply_tag([ar, dit], ["text"])
+    assert ar.extra_args == {"ar_task_mode": "comprehension"}
+    assert getattr(dit, "extra_args", None) == dit_extra_before
+
+
+def test_explicit_caller_ar_task_mode_is_preserved():
+    ar = SamplingParams(extra_args={"ar_task_mode": "generation"})
+    _apply_tag([ar], ["text"])
+    assert ar.extra_args["ar_task_mode"] == "generation"
+
+
+def test_existing_extra_args_are_merged_not_replaced():
+    ar = SamplingParams(extra_args={"custom": 1})
+    _apply_tag([ar], None)
+    assert ar.extra_args == {"custom": 1, "ar_task_mode": "comprehension"}

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():

vllm_omni/entrypoints/openai/serving_chat.py

@@ -818,6 +818,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     if _image_gen_height is not None and _image_gen_width is not None
                     else None
                 )
+                self._apply_text_chat_ar_task_mode(sampling_params_list, request)
                 # Apply user-specified overrides to diffusion stage(s) for image generation
                 for idx, sp in enumerate(sampling_params_list):
                     if idx == comprehension_idx:
@@ -1393,6 +1394,34 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             and delta_message.tool_calls
         )
 
+    @staticmethod
+    def _apply_text_chat_ar_task_mode(
+        sampling_params_list: list[Any],
+        request: ChatCompletionRequest,
+    ) -> None:
+        """Mark AR stages of a text-only chat request as per-request comprehension.
+
+        A generation deployment (e.g. HunyuanImage3 AR+DiT with
+        ``engine_output_type="latent"``) also serves plain chat completions
+        whose answer is text (I2T/T2T). Without a per-request marker the AR
+        model sampler applies its image-generation stage transitions to those
+        requests and leaks DiT scaffold tokens (``<recaption>``/``<answer>``/
+        ``<boi>``/``<img_size_*>``/``<cfg>``) into the text answer (#6088).
+
+        Sets ``extra_args["ar_task_mode"] = "comprehension"`` on every plain
+        ``SamplingParams`` stage when the request output is text-only. Does
+        nothing for image/audio/video-output requests, never overrides an
+        explicit caller-provided ``ar_task_mode``, and models that don't opt
+        into reading extra_args are unaffected.
+        """
+        if set(getattr(request, "modalities", None) or []) - {"text"}:
+            return
+        for sp in sampling_params_list:
+            if isinstance(sp, SamplingParams) and not isinstance(sp, OmniDiffusionSamplingParams):
+                extra_args = dict(getattr(sp, "extra_args", None) or {})
+                extra_args.setdefault("ar_task_mode", "comprehension")
+                sp.extra_args = extra_args
+
     def _build_sampling_params_list_from_request(
         self,
         request: ChatCompletionRequest,

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1448,6 +1448,9 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
 
     prefer_model_sampler = True
     logitsprocs_need_output_token_ids = True
+    # Ask the AR model runner to pass each row's SamplingParams.extra_args to
+    # sample(), so per-request task modes (ar_task_mode) can reach the sampler.
+    model_sampler_wants_extra_args = True
 
     # Siglip2 ViT supports data-parallel encoding (mm_encoder_tp_mode="data"):
     # weights are replicated and the image batch is sharded across TP ranks.
@@ -1571,24 +1574,37 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
             self._all_ratio_ids.update(range(s, e))
 
         # Determine mode: comprehension (I2T/T2T) vs generation (IT2I/T2I).
+        # ``engine_output_type`` is a per-engine static (from the pipeline
+        # registry), but a generation deployment (e.g. the AR+DiT
+        # hunyuan_image_3_moe pipeline, engine_output_type="latent") also
+        # serves plain chat completions whose answer is text (I2T/T2T).
+        # Those requests are marked per-request via
+        # ``SamplingParams.extra_args["ar_task_mode"] == "comprehension"``
+        # by the serving layer and must not go through the image-generation
+        # stage transitions below (#6088).
         engine_output_type = getattr(vllm_config.model_config, "engine_output_type", None)
         self._is_comprehension = engine_output_type in (None, "text")
 
-        # For comprehension mode, block image generation tokens but allow
+        # For comprehension requests, block image generation tokens but allow
         # text structure tokens (<think>, <answer>, etc.) so the model can
         # follow its natural generation pattern. Runtime sampling params
         # decide stop tokens from the active bot_task, matching the official
         # HunyuanImage3 generation path without hard-coded YAML token ids.
-        self._blocked_token_ids: set[int] = set()
-        if self._is_comprehension:
-            self._blocked_token_ids.update(
-                [
-                    self._mrope_boi_token_id,  # <boi>
-                    self._mrope_eoi_token_id,  # <eoi>
-                    self._size_token_id,  # <img_size_*>
-                ]
-            )
-            self._blocked_token_ids.update(self._all_ratio_ids)
+        # <cfg>, <img>, <timestep> and <joint_img_sep> are DiT-side scaffold
+        # tokens that are never valid in a text answer; without them in the
+        # blocklist an I2T request can degenerate into a stream of <cfg>
+        # tokens (#6088).
+        self._cfg_token_id = tokenizer.convert_tokens_to_ids("<cfg>")
+        self._blocked_token_ids: set[int] = {
+            self._mrope_boi_token_id,  # <boi>
+            self._mrope_eoi_token_id,  # <eoi>
+            self._size_token_id,  # <img_size_*>
+            self._cfg_token_id,  # <cfg>
+            self._mrope_img_token_id,  # <img>
+            self._timestep_token_id,  # <timestep>
+            self._mrope_joint_img_sep_token_id,  # <joint_img_sep>
+        }
+        self._blocked_token_ids.update(self._all_ratio_ids)
 
         # For generation mode, build stage transition map.
         # Official logic: </think> → [<recaption>],
@@ -2061,6 +2077,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         self,
         logits: torch.Tensor,
         sampling_metadata: SamplingMetadata,
+        per_req_extra_args: Sequence[Mapping[str, Any] | None] | None = None,
     ) -> SamplerOutput | None:
         if logits is None or logits.numel() == 0:
             return None
@@ -2070,25 +2087,34 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
 
         min_score = torch.finfo(logits.dtype).min
 
+        if self._blocked_token_ids_tensor is None and self._blocked_token_ids:
+            self._blocked_token_ids_tensor = torch.tensor(
+                sorted(self._blocked_token_ids),
+                dtype=torch.long,
+                device=logits.device,
+            )
+
         if self._is_comprehension:
             # Comprehension path is stateless: we only need to mask a fixed
             # set of blocked token ids on every step. Do it in one batched
             # index_fill_ instead of a per-(req, id) Python loop of scalar
             # GPU writes, and skip reading `output_token_ids` from CPU
             # (unused here) so this branch never forces a D2H sync.
-            if self._blocked_token_ids_tensor is None and self._blocked_token_ids:
-                self._blocked_token_ids_tensor = torch.tensor(
-                    sorted(self._blocked_token_ids),
-                    dtype=torch.long,
-                    device=logits.device,
-                )
             if self._blocked_token_ids_tensor is not None:
                 logits.index_fill_(-1, self._blocked_token_ids_tensor, min_score)
             return self._sampler(logits=logits, sampling_metadata=sampling_metadata)
 
-        # Generation path retains the per-request stateful logic for forced
-        # stage-transition tokens and ratio-restriction.
+        # Generation-mode engine. Individual requests can still be
+        # comprehension (I2T/T2T chat served by an AR+DiT deployment); the
+        # serving layer marks those with extra_args["ar_task_mode"] ==
+        # "comprehension", and they get the stateless blocking above instead
+        # of stage-transition forcing (#6088).
         for req_idx in range(logits.shape[0]):
+            if self._row_is_comprehension(per_req_extra_args, req_idx):
+                if self._blocked_token_ids_tensor is not None:
+                    logits[req_idx].index_fill_(-1, self._blocked_token_ids_tensor, min_score)
+                continue
+
             decoded_tokens: list[int] = (
                 sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else []
             )
@@ -2106,6 +2132,24 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
 
         return self._sampler(logits=logits, sampling_metadata=sampling_metadata)
 
+    @staticmethod
+    def _row_is_comprehension(
+        per_req_extra_args: Sequence[Mapping[str, Any] | None] | None,
+        req_idx: int,
+    ) -> bool:
+        """Whether this batch row is a per-request comprehension override.
+
+        Defaults to False (generation) when no extra_args reached the
+        sampler, so offline ``Omni.generate`` image-generation flows keep
+        the stage-transition behavior unchanged.
+        """
+        if per_req_extra_args is None or req_idx >= len(per_req_extra_args):
+            return False
+        extra_args = per_req_extra_args[req_idx]
+        if not extra_args:
+            return False
+        return extra_args.get("ar_task_mode") == "comprehension"
+
     def _get_forced_token(self, decoded_tokens: list[int]) -> int | None:
         """Derive the next forced token from output history (stateless).
 

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -54,7 +54,7 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import (
     OmniConnectorModelRunnerMixin,
     needs_omni_connector,
 )
-from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
+from vllm_omni.worker.sampling_utils import call_model_sampler, sanitize_min_tokens_stop_ids
 
 
 def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]:
@@ -818,7 +818,14 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                     )
                 prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata)
                 self._apply_duplex_sampling(logits, prepared_sampling_metadata)
-                sampler_output = model_sample(logits, prepared_sampling_metadata)
+                sampler_output = call_model_sampler(
+                    self.model,
+                    model_sample,
+                    logits,
+                    prepared_sampling_metadata,
+                    input_batch=self.input_batch,
+                    requests=getattr(self, "requests", None),
+                )
                 if sampler_output is not None:
                     return sampler_output
             return self.sampler(

vllm_omni/worker/gpu_ar_model_runner.py

@@ -59,7 +59,11 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import (
 )
 from vllm_omni.worker.output.payload_build import build_omni_mm_payload
 from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
-from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding, sanitize_min_tokens_stop_ids
+from vllm_omni.worker.sampling_utils import (
+    call_model_sampler,
+    clamp_prompt_ids_to_penalty_padding,
+    sanitize_min_tokens_stop_ids,
+)
 from vllm_omni.worker.sparse_audio import resolve_sparse_mm_routing
 
 logger = init_logger(__name__)
@@ -1315,7 +1319,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                     )
                 prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata)
                 self._apply_duplex_sampling(logits, prepared_sampling_metadata)
-                sampler_output = model_sample(logits, prepared_sampling_metadata)
+                sampler_output = call_model_sampler(
+                    self.model,
+                    model_sample,
+                    logits,
+                    prepared_sampling_metadata,
+                    input_batch=self.input_batch,
+                    requests=getattr(self, "requests", None),
+                )
                 if sampler_output is not None:
                     return sampler_output
                 # Contract: None => fall back to the default sampler (see

vllm_omni/worker/sampling_utils.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Sampling-state guards shared by the GPU and NPU AR model runners."""
+"""Sampling helpers shared by the GPU and NPU AR model runners."""
+
+from typing import Any
 
 import torch
 from vllm.logger import init_logger
@@ -8,7 +10,46 @@ from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsPro
 
 logger = init_logger(__name__)
 
-__all__ = ["clamp_prompt_ids_to_penalty_padding", "sanitize_min_tokens_stop_ids"]
+__all__ = [
+    "build_model_sampler_extra_args",
+    "call_model_sampler",
+    "clamp_prompt_ids_to_penalty_padding",
+    "sanitize_min_tokens_stop_ids",
+]
+
+
+def build_model_sampler_extra_args(input_batch: Any, requests: Any) -> list[dict | None]:
+    """Return per-request ``SamplingParams.extra_args`` in batch-row order."""
+    request_states = requests or {}
+    per_req_extra_args: list[dict | None] = []
+    for req_id in getattr(input_batch, "req_ids", []):
+        state = request_states.get(req_id)
+        params = getattr(state, "sampling_params", None)
+        per_req_extra_args.append(getattr(params, "extra_args", None))
+    return per_req_extra_args
+
+
+def call_model_sampler(
+    model: Any,
+    model_sample: Any,
+    logits: torch.Tensor,
+    sampling_metadata: Any,
+    *,
+    input_batch: Any,
+    requests: Any,
+) -> Any:
+    """Call an opted-in model sampler with per-request extra arguments.
+
+    The opt-in keeps the existing two-argument sampler contract unchanged for
+    every other model while allowing custom samplers to make row-local choices.
+    """
+    if not getattr(model, "model_sampler_wants_extra_args", False):
+        return model_sample(logits, sampling_metadata)
+    return model_sample(
+        logits,
+        sampling_metadata,
+        per_req_extra_args=build_model_sampler_extra_args(input_batch, requests),
+    )
 
 
 def clamp_prompt_ids_to_penalty_padding(prompt_token_ids: torch.Tensor, logits_vocab: int) -> torch.Tensor: