From 16734bbdf3850bc7407bafb3484de613a73efd84 Mon Sep 17 00:00:00 2001 From: DanaerLee Date: Sun, 20 Sep 2026 23:23:53 +0800 Subject: [PATCH] [Bugfix] Route text-only chat as per-request comprehension in HunyuanImage3 AR sampler (#6111) Signed-off-by: MrlixiangWE Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com> --- .../test_hunyuan_image3_sampler.py | 130 +++++++++++++++++- .../test_serving_chat_sampling_params.py | 111 +++++++++++++++ tests/worker/test_model_sampler_extra_args.py | 91 ++++++++++++ tests/worker/test_native_duplex_hooks.py | 2 +- vllm_omni/entrypoints/openai/serving_chat.py | 29 ++++ .../models/hunyuan_image3/hunyuan_image3.py | 82 ++++++++--- .../npu/worker/npu_ar_model_runner.py | 11 +- vllm_omni/worker/gpu_ar_model_runner.py | 15 +- vllm_omni/worker/sampling_utils.py | 45 +++++- 9 files changed, 485 insertions(+), 31 deletions(-) create mode 100644 tests/worker/test_model_sampler_extra_args.py diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py index 34f86347a..eeb14421f 100644 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py +++ b/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 -> 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 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 diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index f81e61cf1..8b2586bd5 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/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"} diff --git a/tests/worker/test_model_sampler_extra_args.py b/tests/worker/test_model_sampler_extra_args.py new file mode 100644 index 000000000..3c996331e --- /dev/null +++ b/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) diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index 388e334cc..1b1bf4b7e 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/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(): diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 13b50b72f..6926253c5 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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 (````/````/ + ````/````/````) 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, diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 0e48412d5..02106a2cc 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/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 (, , 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, # - self._mrope_eoi_token_id, # - self._size_token_id, # - ] - ) - self._blocked_token_ids.update(self._all_ratio_ids) + # , , and 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 + # tokens (#6088). + self._cfg_token_id = tokenizer.convert_tokens_to_ids("") + self._blocked_token_ids: set[int] = { + self._mrope_boi_token_id, # + self._mrope_eoi_token_id, # + self._size_token_id, # + self._cfg_token_id, # + self._mrope_img_token_id, # + self._timestep_token_id, # + self._mrope_joint_img_sep_token_id, # + } + self._blocked_token_ids.update(self._all_ratio_ids) # For generation mode, build stage transition map. # Official logic: → [], @@ -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). diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 18353971c..aa01b7a8e 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/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( diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b918b57f1..0dd393ac6 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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 diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py index 39ecae469..8bcea3e6d 100644 --- a/vllm_omni/worker/sampling_utils.py +++ b/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: