diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 2a2f62c4a82..c5e1b4b7cb7 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -74,11 +74,8 @@ The thinker stage handles multimodal understanding. Key features: from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM -class Qwen3OmniMoeThinkerForConditionalGeneration( - Qwen3MoeForCausalLM, - SupportsMultiModal, - SupportsPP -): + +class Qwen3OmniMoeThinkerForConditionalGeneration(Qwen3MoeForCausalLM, SupportsMultiModal, SupportsPP): """Thinker stage: multimodal understanding → text generation.""" def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -93,10 +90,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( The talker stage converts text embeddings to codec codes: ```python -class Qwen3OmniMoeTalkerForConditionalGeneration( - Qwen3MoeForCausalLM, - SupportsPP -): +class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3MoeForCausalLM, SupportsPP): """Talker stage: text embeddings → RVQ codec codes.""" def __init__(self, vllm_config, talker_config, prefix): @@ -156,8 +150,7 @@ class Qwen3OmniMoeForConditionalGeneration( if self.model_stage == "thinker": # Initialize thinker model thinker_vllm_config = vllm_config.with_hf_config( - config.thinker_config, - architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"] + config.thinker_config, architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"] ) self.thinker = init_vllm_registered_model( vllm_config=thinker_vllm_config, @@ -170,8 +163,7 @@ class Qwen3OmniMoeForConditionalGeneration( elif self.model_stage == "talker": # Initialize talker model talker_vllm_config = vllm_config.with_hf_config( - config.talker_config, - architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"] + config.talker_config, architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"] ) self.talker = init_vllm_registered_model( vllm_config=talker_vllm_config, @@ -184,8 +176,7 @@ class Qwen3OmniMoeForConditionalGeneration( elif self.model_stage == "code2wav": # Initialize code2wav model code2wav_vllm_config = vllm_config.with_hf_config( - config.code2wav_config, - architectures=["Qwen3OmniMoeCode2Wav"] + config.code2wav_config, architectures=["Qwen3OmniMoeCode2Wav"] ) self.code2wav = init_vllm_registered_model( vllm_config=code2wav_vllm_config, @@ -196,8 +187,7 @@ class Qwen3OmniMoeForConditionalGeneration( self.model = self.code2wav else: raise ValueError( - f"Invalid model_stage: {self.model_stage}. " - f"Must be one of: 'thinker', 'talker', 'code2wav'" + f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'talker', 'code2wav'" ) ``` @@ -308,11 +298,10 @@ Register your model in `vllm_omni/model_executor/models/registry.py`: ```python _OMNI_MODELS = { # ... existing models ... - # Your new model "YourModelForConditionalGeneration": ( - "your_model_name", # Module folder name - "your_model", # Module file name (without .py) + "your_model_name", # Module folder name + "your_model", # Module file name (without .py) "YourModelForConditionalGeneration", # Class name ), "YourModelThinkerForConditionalGeneration": ( @@ -410,10 +399,7 @@ The stage transition process follows these steps: else: # Custom transition function (YOUR CODE HERE) return self.custom_process_input_func( - stage_list, - self.engine_input_source, - prompt, - self.requires_multimodal_data + stage_list, self.engine_input_source, prompt, self.requires_multimodal_data ) ``` - If `custom_process_input_func` is configured, it calls that function @@ -506,14 +492,28 @@ thinker_hidden_states = output.multimodal_output["24"] Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function: -| Suffix | Role | Runs when | -| -------- | ------ | ----------- | -| `*_full_payload` | Worker-side payload producer | `async_chunk=false`; accumulates tensors and ships via connector | -| `*_async_chunk` | Scheduler-side streaming producer | `async_chunk=true`; emits per-chunk payloads | -| `*_token_only` | Orchestrator placeholder builder | `async_chunk=false`; allocates downstream prompt slots only | +| Suffix | Registry kind | Runs where | Role | +| -------- | --------------- | ------------ | ------ | +| `*_full_payload` | `producer_full_payload` | Worker (producer-side) | `async_chunk=false`; accumulates tensors and ships via connector | +| `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | `async_chunk=true`; emits per-chunk payloads | +| `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream prompt slots only (both modes) | + +A model that supports **both** `async_chunk=false` and `async_chunk=true` must +implement all three processors for every inter-stage edge — `*_full_payload`, +`*_async_chunk` and `*_token_only` (the `*_token_only` placeholder is required +in both modes). See the +[Stage Input Processor Contract](../../design/feature/async_chunk.md#stage-input-processor-contract) +for the full contract, naming convention, and registry rules. + +Orchestrator-facing builders are normalized to the C1 contract +`(source_outputs, ctx: OrchestratorInputContext)` by +`wrap_orchestrator_processor`. The legacy `*_token_only` shell keeps its +positional signature for call-site compatibility; the actual C1 entry point is +the `build_forward_placeholder` / `build_prewarm_placeholder` pair attached to +the shell (the dual-entry design): ```python -# qwen3_omni.py (Thinker → Talker, non-async path) +# qwen3_omni.py (Thinker → Talker) def thinker2talker_token_only( source_outputs: list[Any], @@ -521,20 +521,33 @@ def thinker2talker_token_only( requires_multimodal_data: bool = False, streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: - """Allocate talker prefill slots; bulk tensors arrive via the connector.""" + """Allocate talker prefill slots; bulk tensors arrive via the connector. + + Legacy positional shell (C2). The orchestrator wraps it once via + ``wrap_orchestrator_processor``, which forwards the C1 + ``OrchestratorInputContext`` fields back into the positional arguments. + """ ... +# Attached C1 dual-entry builders: +# thinker2talker_token_only.build_forward_placeholder(source_outputs, ctx) +# thinker2talker_token_only.build_prewarm_placeholder(*, stage0_prompt, ctx, ...) + + def thinker2talker_full_payload( + *, transfer_manager: Any, pooling_output: dict[str, Any], request: OmniEngineCoreRequest, + is_finished: bool = ..., # optional: the worker retries without it ) -> dict[str, Any] | None: """Pack accumulated thinker hidden states into OmniPayload for stage-1.""" ... def thinker2talker_async_chunk( + *, transfer_manager: Any, multimodal_output: OmniPayload | dict[str, Any], request: OmniEngineCoreRequest, @@ -544,6 +557,11 @@ def thinker2talker_async_chunk( ... ``` +The producer-side builders (`*_full_payload`, `*_async_chunk`) are +keyword-only and never receive an `OrchestratorInputContext`; +`pooling_output` / `multimodal_output` are load-bearing keyword names of the +connector data plane. + Wire these in `pipeline.py`: ```python @@ -560,7 +578,30 @@ StagePipelineConfig( ), ``` -Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is already declared — `_select_processor_funcs()` always prefers the `*_token_only` hook in non-async mode, so the bare function would never run. See `docs/design/rfc_stage_input_processors_refactor.md` for the full contract. +### Registry registration and validation + +The processor registry infers a processor's kind from its name suffix and +structurally validates the signature at startup (it never executes processor +logic). Names that do not follow the suffix convention can register an explicit +kind override once, e.g. at module import time: + +```python +from vllm_omni.model_executor.stage_input_processors import register_processor + +register_processor(f"{_PROC}.thinker2talker_token_only", "placeholder_prompt_builder") +``` + +At startup the runtime resolves each processor through +`resolve_processor(path, expected_kind=...)` — import -> `infer_kind` -> +`validate_processor` -> expected-kind check, returning a `ProcessorSpec` whose +`fn` matches the legacy lookup. Hard violations raise +`ProcessorValidationError`; soft mismatches emit a `RuntimeWarning`. New model +processors should therefore match the suffix convention and the corresponding +signature exactly. + +Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is +already declared — `_select_processor_funcs()` always prefers the `*_token_only` +hook in non-async mode, so the bare function would never run. ## Testing diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 806bec849d1..4754b1f6460 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -238,15 +238,15 @@ Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fu from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer -class YourTTSARStage(nn.Module): +class YourTTSARStage(nn.Module): def __init__(self, config, vllm_config, prefix): - self.layers = nn.ModuleList([ - Qwen3DecoderLayer( - config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}" - ) - for i in range(config.num_hidden_layers) - ]) + self.layers = nn.ModuleList( + [ + Qwen3DecoderLayer(config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}") + for i in range(config.num_hidden_layers) + ] + ) self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size) ``` @@ -257,8 +257,7 @@ See `qwen3_tts_code_predictor_vllm.py` for the full implementation. Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1: ```python -def forward(self, input_ids, positions, intermediate_tensors=None, - inputs_embeds=None, **kwargs) -> OmniOutput: +def forward(self, input_ids, positions, intermediate_tensors=None, inputs_embeds=None, **kwargs) -> OmniOutput: hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds) logits = self.lm_head(hidden_states) @@ -301,8 +300,8 @@ needed). Implement `chunked_decode_streaming()` to support async_chunk streaming ```python # your_model_decoder.py -class YourTTSDecoder(nn.Module): +class YourTTSDecoder(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() # Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.) @@ -310,14 +309,13 @@ class YourTTSDecoder(nn.Module): def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor: return self.decoder(codes) - def chunked_decode_streaming(self, codes, chunk_size=25, - left_context_size=25) -> torch.Tensor: + def chunked_decode_streaming(self, codes, chunk_size=25, left_context_size=25) -> torch.Tensor: """Decode with a sliding context window for smooth chunk boundaries.""" end_index = codes.shape[-1] context_size = 0 if end_index <= chunk_size else left_context_size wav_chunk = self(codes) # Trim left context to avoid duplicate audio - return wav_chunk[..., context_size * self.total_upsample:] + return wav_chunk[..., context_size * self.total_upsample :] ``` ### Step 3: Implement the Unified Model Class @@ -327,8 +325,8 @@ The unified class dispatches to the correct stage based on `model_stage` in the ```python # your_model.py -class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP): +class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.model_stage = vllm_config.model_config.model_stage @@ -411,15 +409,18 @@ Register all stage classes in `vllm_omni/model_executor/models/registry.py`: _OMNI_MODELS = { # (package_name, module_name, class_name) "YourTTSModelForConditionalGeneration": ( - "your_model_name", "your_model", + "your_model_name", + "your_model", "YourTTSModelForConditionalGeneration", ), "YourTTSARStageForConditionalGeneration": ( - "your_model_name", "your_model_ar_stage", + "your_model_name", + "your_model_ar_stage", "YourTTSARStageForConditionalGeneration", ), "YourTTSDecoder": ( - "your_model_name", "your_model_decoder", + "your_model_name", + "your_model_decoder", "YourTTSDecoder", ), } @@ -472,7 +473,7 @@ YOUR_TTS_PIPELINE = PipelineConfig( input_sources=(), owns_tokenizer=True, engine_output_type="latent", - custom_process_next_stage_input_func=f"{_PROC}.ar2decoder", + custom_process_next_stage_input_func=f"{_PROC}.ar2decoder_full_payload", async_chunk_process_next_stage_input_func=f"{_PROC}.ar2decoder_async_chunk", ), StagePipelineConfig( @@ -481,6 +482,7 @@ YOUR_TTS_PIPELINE = PipelineConfig( execution_type=StageExecutionType.LLM_GENERATION, input_sources=(0,), model_arch="YourTTSDecoder", + sync_process_input_func=f"{_PROC}.ar2decoder_token_only", engine_output_type="audio", final_output=True, final_output_type="audio", @@ -521,6 +523,50 @@ Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours See `stage_input_processors/qwen3_tts.py` for the full reference implementation. +### Processor set and naming convention + +Each inter-stage edge provides a **coherent processor set** rather than a single +monolithic function. The suffix convention and the matching registry kind are +load-bearing: + +| Suffix | Registry kind | Runs where | Role | +| -------- | --------------- | ------------ | ------ | +| `*_full_payload` | `producer_full_payload` | Worker (producer-side) | `async_chunk=false`; packs the accumulated Stage 0 output and ships it via the connector | +| `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | `async_chunk=true`; streams one chunk per AR decode step | +| `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream KV slots only; bulk tensors arrive via the connector (both modes) | + +A model that supports **both** `async_chunk=false` and `async_chunk=true` must +implement all three processors per edge — `ar2decoder_full_payload`, +`ar2decoder_async_chunk` and `ar2decoder_token_only`. The `*_token_only` +placeholder builder is wired as the downstream stage's `sync_process_input_func` +and follows the C1 contract `(source_outputs, ctx: OrchestratorInputContext)`: + +```python +from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext + + +def ar2decoder_token_only( + source_outputs: list[Any], + ctx: OrchestratorInputContext, +) -> list[OmniTokensPrompt]: + """Allocate decoder prefill slots; bulk audio codes arrive via the connector.""" + ... +``` + +Names that do not follow the suffix convention can register an explicit kind +override at module import time: + +```python +from vllm_omni.model_executor.stage_input_processors import register_processor + +register_processor(f"{_PROC}.ar2decoder_token_only", "placeholder_prompt_builder") +``` + +For the registry validation rules, legacy-shape adaptation, and the full +contract, see the +[Stage Input Processor Contract](../../design/feature/async_chunk.md#stage-input-processor-contract) +section of the async chunk design document. + ### Data structures Understanding what's available in stage outputs: @@ -532,28 +578,29 @@ Understanding what's available in stage outputs: - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs` - `prompt_token_ids` - original prompt token IDs -### Batch mode (non-streaming) +### Batch mode (non-streaming, full payload) -Collects all Stage 0 outputs and forwards them to Stage 1 in one shot: +Runs on the worker (producer) side for `async_chunk=false`. The connector +invokes it with the exact keyword contract +`(transfer_manager, pooling_output, request, ...)` — note that the name must +carry the `_full_payload` suffix so `infer_kind` resolves it as a +`producer_full_payload` (an unsuffixed `ar2decoder` would be inferred as +`legacy_orchestrator_builder` and skipped by the worker's candidate chain): ```python -def ar2decoder( - stage_list: list[Any], - engine_input_source: list[int], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - source_id = engine_input_source[0] - decoder_inputs = [] - - for output in stage_list[source_id].engine_outputs: - result = output.outputs[0] - codes = result.multimodal_output["audio_codes"].cpu() - decoder_inputs.append( - OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist()) - ) - - return decoder_inputs +def ar2decoder_full_payload( + transfer_manager: Any, + pooling_output: dict[str, Any] | None, + request: Any, + **_: Any, +) -> dict[str, Any] | None: + """Producer: pack the accumulated Stage 0 codec into a connector payload.""" + del transfer_manager + codes = pooling_output.get("codes.audio") if isinstance(pooling_output, dict) else None + if codes is None: + return None + decoder_inputs = OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist()) + return {"engine_inputs": [decoder_inputs]} ``` ### Streaming mode (async_chunk) @@ -563,8 +610,9 @@ have accumulated. The function signature follows the `OmniChunkTransferAdapter` ```python def ar2decoder_async_chunk( + *, transfer_manager: Any, - pooling_output: dict[str, Any] | None, + multimodal_output: dict[str, Any] | None, request: Any, is_finished: bool = False, ) -> dict[str, Any] | None: @@ -573,12 +621,10 @@ def ar2decoder_async_chunk( finished = bool(is_finished or request.is_finished()) # Extract and buffer the latest frame - if isinstance(pooling_output, dict): - frame = extract_frame(pooling_output) + if isinstance(multimodal_output, dict): + frame = extract_frame(multimodal_output) if frame is not None: - transfer_manager.code_prompt_token_ids[request_id].append( - frame.cpu().tolist() - ) + transfer_manager.code_prompt_token_ids[request_id].append(frame.cpu().tolist()) elif not finished: return None @@ -861,11 +907,10 @@ any CUDA allocations. Stream via a per-request generator stored in an instance d ```python class YourModelForCausalLM(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - self._lm: nn.Module | None = None # populated in load_weights() - self._stream_gens: dict[str, Any] = {} # request_key → generator + self._lm: nn.Module | None = None # populated in load_weights() + self._stream_gens: dict[str, Any] = {} # request_key → generator def load_weights(self, weights): # Load self._lm here, after vLLM distributed init diff --git a/docs/design/feature/async_chunk.md b/docs/design/feature/async_chunk.md index 5a74b9f53f6..90caffc43a0 100644 --- a/docs/design/feature/async_chunk.md +++ b/docs/design/feature/async_chunk.md @@ -6,7 +6,8 @@ 2. [Performance](#performance) 3. [Architecture](#architecture) 4. [Configuration](#configuration) -5. [Related Files](#related-files) +5. [Stage Input Processor Contract](#stage-input-processor-contract) +6. [Related Files](#related-files) ## Overview @@ -314,6 +315,166 @@ stages: max_num_seqs: 64 # Enables batched audio generation ``` +## Stage Input Processor Contract + +RFC #4872 standardized the per-model builders that convert one stage's outputs +into the next stage's inputs. Each builder belongs to exactly one *role* +(consumer-side vs producer-side) and is dispatched by the runtime according to +the `async_chunk` mode. The naming convention and the validation rules below +are load-bearing: the processor registry infers a builder's kind from its name +suffix and enforces the matching signature at startup. + +### Naming convention and roles + +| Suffix | Registry kind | Runs where | Role | +| -------- | --------------- | ------------ | ------ | +| `*_full_payload` | `producer_full_payload` | Worker (producer-side) | Packs the accumulated stage output into an `OmniPayload` and ships it through the connector (`FullPayloadProducer`) | +| `*_async_chunk` | `producer_async_chunk` | Scheduler (producer-side) | Streams one chunk per call while `async_chunk` is enabled (`AsyncChunkProducer`) | +| `*_token_only` | `placeholder_prompt_builder` | Orchestrator (consumer-side) | Allocates downstream KV slots only; bulk tensors arrive via the connector (`PlaceholderPromptBuilder`) | +| (no suffix, legacy) | `legacy_orchestrator_builder` | Orchestrator (consumer-side) | Legacy sync builder — a placeholder or diffusion input builder, adapted through `wrap_orchestrator_processor` | + +Diffusion-stage edges use dedicated suffixes +(`ar2diffusion` / `ar2dit` / `thinker2imagegen` / ...) mapped to +`diffusion_input_builder`, and `moss_tts.talker2codec` keeps the legacy +`legacy_multi_source` shape. The canonical naming documentation lives in the +`vllm_omni/model_executor/stage_input_processors/__init__.py` module docstring. + +### Minimal processor set per edge + +A new model that must work in **both** modes implements, for each inter-stage +edge, the three processors below: + +| Mode | Producer-side | Consumer-side (orchestrator) | +| ------------------- | ---------------------------------------- | -------------------------------------------------------- | +| `async_chunk=false` | `*_full_payload` (`FullPayloadProducer`) | `*_token_only` (`PlaceholderPromptBuilder`) | +| `async_chunk=true` | `*_async_chunk` (`AsyncChunkProducer`) | `*_token_only` (prewarm via `build_prewarm_placeholder`) | + +The minimal union set is `*_full_payload`, `*_async_chunk` and `*_token_only`. +The `*_token_only` placeholder builder is required in both modes: in non-async +mode it builds the forward placeholder, and in async mode the orchestrator +reuses it to prewarm the downstream stage. + +### OrchestratorInputContext and the C1 contract + +Every orchestrator-facing builder is invoked under the fixed C1 contract +`(source_outputs, ctx)`: + +```python +from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext + + +def my_token_only( + source_outputs: list[Any], + ctx: OrchestratorInputContext, +) -> list[OmniTokensPrompt]: + """Upstream outputs -> next-stage token prompts (C1 contract).""" + ... +``` + +`OrchestratorInputContext` (defined in `_dispatch.py`) carries `prompt`, +`requires_multimodal_data`, `streaming_context`, `sampling_params`, and the +target-stage config fields `target_model_config` / `next_stage_hf_config` +(populated by the stage client from the stage's own vllm config); it +deliberately has no `model_config` field. + +Processors that already accept `ctx` are used unchanged. Legacy positional +shapes (C0 3-arg, C2 placeholder, C3 diffusion, C4 `moss_tts.talker2codec` +multi-source) are adapted once by `wrap_orchestrator_processor` / +`invoke_orchestrator_processor` (see `_dispatch.py`) and emit a +`DeprecationWarning`. + +### Registry and startup validation + +The registry performs **signature-level structural checks only** — it never +executes processor logic and never loads model weights. Kind inference is +name-driven (suffix-based): + +- `register_processor(path, kind)` — manual kind override for names that do not + follow the suffix convention (escape hatch; overrides are validated eagerly). +- `infer_kind(fn, *, path)` — suffix rules: `_token_only` -> + `placeholder_prompt_builder`, `_full_payload` / `_batch` -> + `producer_full_payload`, `_async_chunk` -> `producer_async_chunk`, diffusion + suffixes -> `diffusion_input_builder`, no suffix -> + `legacy_orchestrator_builder`, `moss_tts.talker2codec` -> + `legacy_multi_source`. +- `validate_processor(fn, *, kind, path, stage_config=None)` — hard violations + raise `ProcessorValidationError`; soft mismatches emit a `RuntimeWarning`. +- `resolve_processor(path, *, expected_kind=None, stage_config=None)` — the + drop-in replacement for legacy `getattr(importlib.import_module(...), ...)` + lookups: import -> infer -> validate -> optional expected-kind check, + returning a `ProcessorSpec` whose `fn` matches the legacy lookup. + +### Dual entry: forward and prewarm placeholders + +The `*_token_only` placeholder builder is exposed through two entry points so +the sync forward path and the async-chunk prewarm path share the same `_common` +length / packing helpers: + +- `build_forward_placeholder(source_outputs, ctx)` — the non-async forward + path. One placeholder `OmniTokensPrompt` per upstream output, sized by + `_common.compute_placeholder_prompt_len(mode="full")`. +- `build_prewarm_placeholder(*, stage0_prompt, ctx, downstream_stage_id)` — + async-chunk mode has no upstream `source_outputs` yet at prewarm time, so the + length is a best-effort estimate from the stage-0 input prompt + (`mode="stage0_only"`, matching the legacy + `adapter.compute_talker_prompt_ids_length`). The connector fixup path replaces + the estimate with the real length once the upstream chunk arrives. + +The dual-entry builder is resolved once at stage-metadata extraction time +(`stage_init_utils._resolve_prewarm_placeholder_builder`) and exposed on the +stage client as `prewarm_placeholder_builder`; the orchestrator's +`_prewarm_async_chunk_stages` reads it from there. + +### Three-state async gate + +The orchestrator only forwards via `process_engine_inputs` when the transition +is not served by the async-chunk data plane: + +```python +if ( + (finished or segment_finished) + and stage_id < req_state.final_stage_id + and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1)) + and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled) +): +``` + +`_stage_receives_async_chunks(stage_id)` reports whether a stage's connector +supplies its runtime inputs. When `async_chunk=true` and the downstream stage +receives async chunks, the orchestrator skips `process_engine_inputs` for that +transition, so the corresponding input processor may be dead for the duration +of the request. `dead_processor_hint` is a pure decision helper (for warnings +and tests only); the runtime is warn-first in the M0 phase. + +### Producer keyword contract + +Producer-side builders never receive an `OrchestratorInputContext`; their +keyword-only parameters are load-bearing parts of the connector data plane and +must not be renamed or made positional: + +```python +def my_full_payload( + *, + transfer_manager: Any, + pooling_output: Any, + request: Any, + is_finished: bool = ..., # optional: the worker retries without it +) -> Any: ... + + +def my_async_chunk( + *, + transfer_manager: Any, + multimodal_output: Any, + request: Any, + is_finished: bool = False, # required: the scheduler always passes it +) -> Any: ... +``` + +`pooling_output` (full payload) and `multimodal_output` (async chunk) are +cross-checked by `validate_processor`; using the wrong keyword name for a kind +is flagged as a structural warning. + ## Related Files - `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`: Chunk processing functions (receive `transfer_manager` as first param) diff --git a/docs/design/index.md b/docs/design/index.md index 986fd19cecd..11a8e4405d1 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -22,6 +22,7 @@ implementation contract; it is not, by itself, a general support claim. - [Disaggregated Inference](feature/disaggregated_inference.md) - [Host Weight Runtime](feature/host_weight_runtime.md) - [Async Chunk](feature/async_chunk.md) +- [Stage Input Processor Contract](feature/async_chunk.md#stage-input-processor-contract) — RFC #4872 `stage_input_processors` refactor - [Async Diffusion Output](feature/async_diffusion_output.md) - [Async Omni Output Materialization](feature/omni_async_output_materialization.md) - [Runner-to-model Prefill/Decode Phase Contract](feature/preprocess_phase_contract.md) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index c7073fee14e..94640dca137 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -575,6 +575,7 @@ def _build( flat_tts_config: bool = False, active_stream_window: int = 0, connector_extra: dict | None = None, + model_config_extra: dict | None = None, custom_process_next_stage_input_func: str | None = None, ): connector = mocker.MagicMock() @@ -618,6 +619,12 @@ def _fake_base_init(self, config): attention_type=tts_attention_type, ) ) + # Both call styles must work: ours-style callers pass the hook via + # ``model_config_extra`` while theirs-style callers pass the explicit + # ``custom_process_next_stage_input_func`` kwarg. Folding them into one + # mapping avoids a duplicate-keyword TypeError when both are supplied. + extra = dict(model_config_extra or {}) + extra.setdefault("custom_process_next_stage_input_func", custom_process_next_stage_input_func) model_config = SimpleNamespace( worker_type=model_mode, max_num_seqs=max_num_seqs, @@ -625,11 +632,11 @@ def _fake_base_init(self, config): hf_config=hf_config, hf_config_name="tts_config" if stage_id == 1 else None, active_stream_window=active_stream_window, - custom_process_next_stage_input_func=custom_process_next_stage_input_func, stage_connector_config={ "name": "SharedMemoryConnector", "extra": connector_extra or {}, }, + **extra, ) scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) adapter = OmniChunkTransferAdapter( @@ -640,6 +647,37 @@ def _fake_base_init(self, config): return _build +def pack(transfer_manager, multimodal_output, request, is_finished=False): + """Arbitrarily named, structurally compatible async-chunk producer.""" + del transfer_manager, multimodal_output, request, is_finished + return None + + +def _incompatible_hook(transfer_manager, request): + """Structurally invalid: no ``multimodal_output`` and no ``is_finished``.""" + del transfer_manager, request + return None + + +def test_adapter_accepts_explicitly_configured_hook_with_arbitrary_name(build_adapter): + """A signature-compatible hook without a ``*_async_chunk`` suffix loads.""" + adapter, _ = build_adapter(model_config_extra={"custom_process_next_stage_input_func": f"{__name__}.pack"}) + + assert adapter.custom_process_next_stage_input_func is pack + + +def test_adapter_rejects_structurally_incompatible_hook(build_adapter): + """A hook missing required parameters fails fast instead of loading.""" + from vllm_omni.model_executor.stage_input_processors import ProcessorValidationError + + with pytest.raises(ProcessorValidationError): + build_adapter( + model_config_extra={ + "custom_process_next_stage_input_func": f"{__name__}._incompatible_hook", + } + ) + + @pytest.mark.parametrize("active_stream_window", [0, 2]) def test_scheduler_chunk_gate_preserves_kv_holding_queue(build_adapter, active_stream_window): adapter, _ = build_adapter(active_stream_window=active_stream_window) @@ -1653,6 +1691,89 @@ def test_send_single_request_personaplex_pending_frame_is_not_segment_boundary( assert sent_payload.meta.is_segment_finished.item() is False +def test_send_single_request_passes_new_token_ids_to_processor(build_adapter, monkeypatch): + """A producer that accepts ``new_token_ids`` receives the step's tokens + through the production ``_send_single_request`` path. The signature cache + (``_processor_accepts_step_tokens``) must be initialized in ``__init__`` + before the base class starts the save thread; without it every send raised + AttributeError that ``_send_single_request_for_generation`` swallowed, + silently dropping the chunk. + """ + adapter, connector = build_adapter(stage_id=1) + request = _req("req-new-tokens", RequestStatus.WAITING, external_req_id="ext-new-tokens") + + seen_kwargs = {} + + def new_tokens_processor(**kwargs): + seen_kwargs.update(kwargs) + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)), + ) + + adapter.custom_process_next_stage_input_func = new_tokens_processor + monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None) + + adapter._send_single_request( + { + "multimodal_output": None, + "request": request, + "is_finished": False, + "is_segment_finished": False, + "new_token_ids": (7, 8, 9), + } + ) + + assert seen_kwargs["new_token_ids"] == (7, 8, 9) + assert connector.put.called + sent_payload = connector.put.call_args.kwargs["data"] + assert sent_payload.meta.finished.item() is False + + +def test_send_single_request_normal_chunk_flushes_processor(build_adapter, monkeypatch): + """A non-terminal chunk with a producer that does **not** accept + ``new_token_ids`` must still reach the wire. Before the cache + re-initialization the ``_accepts_new_token_ids`` probe raised AttributeError + on every call, the processor was never invoked, and the ordinary chunk was + silently dropped at the ``payload_data is None`` guard. + """ + adapter, connector = build_adapter(stage_id=1) + request = _req("req-normal-chunk", RequestStatus.WAITING, external_req_id="ext-normal-chunk") + + processor_calls = [] + + def plain_processor(transfer_manager, multimodal_output, request, is_finished): + processor_calls.append( + { + "transfer_manager": transfer_manager, + "multimodal_output": multimodal_output, + "request": request, + "is_finished": is_finished, + } + ) + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([4, 5], dtype=torch.long)), + ) + + adapter.custom_process_next_stage_input_func = plain_processor + monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None) + + adapter._send_single_request( + { + "multimodal_output": None, + "request": request, + "is_finished": False, + "is_segment_finished": False, + } + ) + + assert len(processor_calls) == 1 + assert processor_calls[0]["is_finished"] is False + assert connector.put.called + sent_payload = connector.put.call_args.kwargs["data"] + assert sent_payload.meta.finished.item() is False + assert sent_payload.meta.is_segment_finished.item() is False + + def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter): from vllm_omni.model_executor.stage_input_processors.personaplex import ( talker2code2wav_async_chunk, diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 5f3d1147145..d6cc196f0cd 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -97,6 +97,8 @@ def test_stage_engine_core_client_builds_payload_sender_info_after_base_init( prompt_transform_func=None, prompt_expand_func=None, custom_process_input_func=None, + sync_process_input_func=None, + prewarm_placeholder_builder=None, ) def fake_base_init(self, config, *_args, **_kwargs): diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 70dc7ca9f88..9ae79f38c31 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -7,10 +7,11 @@ import queue from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import janus import pytest +import torch from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams @@ -100,13 +101,9 @@ def process_inputs(self, **kwargs): class FakePrewarmPool: - """Single-replica pool double that binds requests only on submission.""" - stage_type = "llm" def __init__(self, role: str) -> None: - self.stage_client = SimpleNamespace() - self._bound_request_ids: set[str] = set() self.stage_vllm_config = SimpleNamespace( model_config=SimpleNamespace( max_model_len=64, @@ -114,17 +111,18 @@ def __init__(self, role: str) -> None: ) ) self.submitted: list[Any] = [] + self.stage_client: Any = None # _build_payload_sender_info fallback (returns None → OK) - async def submit_initial(self, request_id, _req_state, request, prompt_text=None): + async def submit_initial(self, _request_id, _req_state, request, prompt_text=None): self.submitted.append(request) - self._bound_request_ids.add(request_id) return 0 - def get_bound_replica_id(self, request_id): - return 0 if request_id in self._bound_request_ids else None + def get_bound_replica_id(self, _request_id): + return 0 - def get_bound_client(self, request_id): - return self.stage_client if self.get_bound_replica_id(request_id) is not None else None + def get_bound_client(self, _request_id: str) -> None: + """No bound client; _build_payload_sender_info falls back to stage_client (None).""" + return None def _request_output(request_id: str) -> RequestOutput: @@ -208,29 +206,17 @@ def decode(self, token_ids, *, skip_special_tokens): @pytest.mark.asyncio -@pytest.mark.parametrize("payload_sender_info", [None, {"host": "10.0.0.2", "zmq_port": 52099}]) -async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> None: - import torch - - from vllm_omni.engine.serialization import deserialize_additional_information - +async def test_async_prewarm_skips_outgoing_only_stage() -> None: orchestrator = object.__new__(Orchestrator) stage0 = FakePrewarmPool("sender") stage1 = FakePrewarmPool("sender") stage2 = FakePrewarmPool("receiver") - if payload_sender_info is not None: - stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info) orchestrator.stage_pools = [stage0, stage1, stage2] orchestrator._emit_tx_edge = lambda **_kwargs: None orchestrator._on_stage_submitted = MagicMock() - max_frames = [256] - prompt = { - "prompt_token_ids": [1, 2], - "additional_information": {"codes": {"ref": torch.arange(6).reshape(3, 2)}, "max_new_frames": max_frames}, - } req_state = OrchestratorRequestState( request_id="req-prewarm", - prompt=prompt, + prompt={"prompt_token_ids": [1, 2]}, sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], final_stage_id=2, ) @@ -240,23 +226,10 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), req_state, ) - max_frames.append(512) - info = deserialize_additional_information(stage2.submitted[0].additional_information) - assert torch.equal(info["codes"]["ref"], torch.arange(6).reshape(3, 2)) - assert info["max_new_frames"] == [256] - assert set(stage2.submitted[0].prompt_token_ids) == {0} - assert prompt["prompt_token_ids"] == [1, 2] assert prewarmed is True assert stage1.submitted == [] - assert stage1.get_bound_client("req-prewarm") is None assert len(stage2.submitted) == 1 - assert stage2.get_bound_client("req-prewarm") is stage2.stage_client - assert stage2.submitted[0].payload_sender_info == payload_sender_info - if payload_sender_info is not None: - stage1.stage_client.get_payload_sender_info.assert_called_once_with() - assert stage2.submitted[0].external_req_id == "req-prewarm" - assert stage2.submitted[0].resumable is True assert 1 not in req_state.stage_submit_ts assert 2 in req_state.stage_submit_ts orchestrator._on_stage_submitted.assert_called_once_with( @@ -303,6 +276,64 @@ async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> Non assert 2 in req_state.stage_submit_ts +async def test_prewarm_uses_registered_build_prewarm_placeholder() -> None: + """Async-chunk prewarm uses the metadata-resolved ``build_prewarm_placeholder``. + + The builder is resolved into stage metadata at init time (see + ``_resolve_prewarm_placeholder_builder``) and read off the stage client — + the orchestrator never re-resolves. The placeholder length is + ``stage0_only`` = the Qwen chat-template scan on the stage-0 prompt. The + synthetic ``[1, 2]`` prompt has no chat-template marker, so the scan yields + 0 and the builder floors it to a 1-token placeholder (best-effort; the + connector fixup path replaces it later). + """ + from vllm_omni.engine.stage_init_utils import _resolve_prewarm_placeholder_builder + + orchestrator = object.__new__(Orchestrator) + stage0 = FakePrewarmPool("sender") + stage1 = FakePrewarmPool("receiver") + # The real qwen3_omni sync placeholder builder exposes build_prewarm_placeholder + # off the resolved fn (see test_placeholder_parity.py). + stage1.stage_client = SimpleNamespace( # type: ignore[attr-defined] + sync_process_input_func=( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only" + ), + prewarm_placeholder_builder=_resolve_prewarm_placeholder_builder( + sync_path=("vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only"), + custom_fn=None, + stage_config=None, + ), + ) + orchestrator.stage_pools = [stage0, stage1] + orchestrator._emit_tx_edge = lambda **_kwargs: None + orchestrator._on_stage_submitted = MagicMock() + req_state = OrchestratorRequestState( + request_id="req-prewarm-resolved", + prompt={"prompt_token_ids": [1, 2]}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=1, + ) + + prewarmed = await orchestrator._prewarm_async_chunk_stages( + "req-prewarm-resolved", + SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), + req_state, + ) + + assert prewarmed is True + assert len(stage1.submitted) == 1 + # build_prewarm_placeholder uses stage0_only = the chat-template scan on the + # stage-0 prompt. The synthetic [1, 2] prompt has no chat-template marker, + # so the scan yields 0 and the builder floors it to a 1-token placeholder. + assert stage1.submitted[0].prompt_token_ids == [0] + orchestrator._on_stage_submitted.assert_called_once_with( + 1, + "req-prewarm-resolved", + 0, + req_state, + ) + + @pytest.mark.asyncio async def test_async_route_forwards_to_outgoing_only_stage() -> None: orchestrator = object.__new__(Orchestrator) @@ -362,3 +393,279 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: orchestrator._cleanup_request_ids.assert_not_awaited() routed = orchestrator.output_async_queue.get_nowait() assert routed.finished is False + + +# --------------------------------------------------------------------------- +# Dead-processor hint must not flag the selected sync hook. +# --------------------------------------------------------------------------- + + +def _prewarm_orchestrator_with_client(client: Any) -> Orchestrator: + orch = object.__new__(Orchestrator) + orch.stage_pools = [ + SimpleNamespace( + stage_client=client, + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + ] + return orch + + +def test_warn_dead_input_processors_skips_selected_sync_hook() -> None: + """The selected sync hook used as the custom hook is NOT dead. + + In non-async mode a stage wiring the same ``*_token_only`` processor as + both ``custom_process_input_func`` and ``sync_process_input_func`` must not + be reported as never-invoked (it is the active forward processor). + """ + from vllm_omni.engine import orchestrator as orch_module + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + path = f"{_sync_hook.__module__}.{_sync_hook.__qualname__}" + + class _Client: + custom_process_input_func = _sync_hook + sync_process_input_func = path + + orch = _prewarm_orchestrator_with_client(_Client()) + orch.async_chunk = False + with patch.object(orch_module.logger, "warning") as mock_warn: + orch._warn_dead_input_processors() + dead_warnings = [call for call in mock_warn.call_args_list if call.args and "never invoked" in call.args[0]] + assert dead_warnings == [] + + +def test_warn_dead_input_processors_flags_distinct_sync_hook() -> None: + """A custom hook overridden by a *different* sync hook is still reported dead.""" + from vllm_omni.engine import orchestrator as orch_module + + def _custom_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + path = f"{_sync_hook.__module__}.{_sync_hook.__qualname__}" + + class _Client: + custom_process_input_func = _custom_hook + sync_process_input_func = path + + orch = _prewarm_orchestrator_with_client(_Client()) + orch.async_chunk = False + with patch.object(orch_module.logger, "warning") as mock_warn: + orch._warn_dead_input_processors() + dead_warnings = [call for call in mock_warn.call_args_list if call.args and "never invoked" in call.args[0]] + assert len(dead_warnings) == 1 + + +# --------------------------------------------------------------------------- +# Prewarm placeholder builder is resolved into stage metadata at init time. +# --------------------------------------------------------------------------- + + +class _FakeProcessorSpec: + """Minimal stand-in for ``ProcessorSpec`` with a controllable ``fn``.""" + + def __init__(self, fn: Any) -> None: + self.path = "fake" + self.kind = "placeholder_prompt_builder" + self.fn = fn + + +def test_prewarm_builder_resolved_once_into_stage_metadata() -> None: + """The dual-entry builder is resolved exactly once, at metadata time. + + ``_resolve_prewarm_placeholder_builder`` reads ``build_prewarm_placeholder`` + off the resolved ``sync_process_input_func`` fn (with the module-level + sibling fallback) so the orchestrator never re-resolves per async request. + """ + from vllm_omni.engine.stage_init_utils import _resolve_prewarm_placeholder_builder + + def _build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] + return {"prompt_token_ids": [0]} + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + _sync_hook.build_prewarm_placeholder = _build_prewarm_placeholder # type: ignore[attr-defined] + + with patch( + "vllm_omni.engine.stage_init_utils.resolve_processor", + return_value=_FakeProcessorSpec(_sync_hook), + ) as mock_resolve: + builder = _resolve_prewarm_placeholder_builder( + sync_path="pkg.mod.fake_token_only", + custom_fn=None, + stage_config=None, + ) + assert builder is _build_prewarm_placeholder + assert mock_resolve.call_count == 1 + + +def test_prewarm_builder_miss_warns_once_then_keeps_inline_estimate() -> None: + """A stage WITHOUT a builder warns once at init and the prewarm keeps the + inline estimate — the orchestrator never re-resolves or re-warns.""" + from vllm_omni.engine import stage_init_utils as siu + from vllm_omni.engine.stage_init_utils import _resolve_prewarm_placeholder_builder + + def _sync_hook(source_outputs, prompt=None, requires_multimodal_data=False): # type: ignore[no-untyped-def] + return [] + + fake_module = SimpleNamespace() + with ( + patch( + "vllm_omni.engine.stage_init_utils.resolve_processor", + return_value=_FakeProcessorSpec(_sync_hook), + ) as mock_resolve, + patch("importlib.import_module", return_value=fake_module) as mock_import, + patch.object(siu.logger, "warning") as mock_warn, + ): + builder = _resolve_prewarm_placeholder_builder( + sync_path="pkg.mod.fake_token_only", + custom_fn=None, + stage_config=None, + ) + assert builder is None + assert mock_resolve.call_count == 1 + assert mock_import.call_count == 1 + inline_estimates = [call for call in mock_warn.call_args_list if call.args and "inline estimate" in call.args[0]] + assert len(inline_estimates) == 1 + + +# --------------------------------------------------------------------------- +# Regression: configured async_chunk_prewarm_prompt_len must survive the +# inline-estimate fallback (amy-why-3459 review of HEAD 22e3a258). +# --------------------------------------------------------------------------- + + +def _override_orchestrator(client: Any, prewarm_len: int = 37) -> Orchestrator: + """Orchestrator whose downstream stage declares a token prefill boundary.""" + orch = object.__new__(Orchestrator) + orch.stage_pools = [ + SimpleNamespace( + stage_client=client, + stage_vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + max_model_len=64, + hf_config=SimpleNamespace(async_chunk_prewarm_prompt_len=prewarm_len), + ), + ), + ) + ] + return orch + + +#: Builder variants shared by the override table below (a None builder is the +#: "no dual-entry builder" miss case; a raising builder exercises the fallback). +def _build_prewarm_placeholder_3(**kwargs): # type: ignore[no-untyped-def] + return [{"prompt_token_ids": [0, 0, 0]}] + + +def _build_prewarm_placeholder_789(**kwargs): # type: ignore[no-untyped-def] + return [{"prompt_token_ids": [7, 8, 9]}] + + +def _broken_build_prewarm_placeholder(**kwargs): # type: ignore[no-untyped-def] + raise RuntimeError("boom") + + +@pytest.mark.parametrize( + "builder,prewarm_len,expected_ids", + [ + # A stage WITHOUT a builder keeps the configured override through the + # inline-estimate fallback (Nemotron's 37-token speaker-prompt prefill). + pytest.param(None, 37, [0] * 37, id="builder-miss-keeps-override"), + # A builder that RAISES falls back to the inline estimate AND keeps it. + pytest.param(_broken_build_prewarm_placeholder, 37, [0] * 37, id="builder-failure-keeps-override"), + # Builder hit must still honor the override (amy-why-3459 review of HEAD ad81ac3a). + pytest.param(_build_prewarm_placeholder_3, 37, [0] * 37, id="builder-success-keeps-override"), + # Without a configured override, a builder hit keeps its own output. + pytest.param(_build_prewarm_placeholder_789, None, [7, 8, 9], id="no-override-keeps-builder-output"), + # A configured length of 1 is not a prefill override: builder output wins. + pytest.param(_build_prewarm_placeholder_3, 1, [0, 0, 0], id="len-one-not-a-prefill-override"), + ], +) +def test_prewarm_prompt_len_override(builder, prewarm_len, expected_ids) -> None: + from vllm_omni.engine import orchestrator as orch_module + + class _Client: + sync_process_input_func = "pkg.mod.fake_token_only" + + def __init__(self) -> None: + # Instance attribute: class-level functions bind self on access. + self.prewarm_placeholder_builder = builder + + if prewarm_len is None: + orch = _prewarm_orchestrator_with_client(_Client()) # no async_chunk_prewarm_prompt_len + else: + orch = _override_orchestrator(_Client(), prewarm_len=prewarm_len) + req_state = OrchestratorRequestState( + request_id="req-override", + prompt={"prompt": "hi"}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + with patch.object(orch_module.logger, "warning"): + base_input = orch._build_prewarm_placeholder_input(0, "req-override", [1, 2], req_state) + + assert base_input["prompt_token_ids"] == expected_ids + assert base_input["multi_modal_data"] is None + assert base_input["mm_processor_kwargs"] is None + + +class _DeepcopySpy: + """Sentinel media object that records any attempted deep-copy.""" + + def __init__(self) -> None: + self.deepcopy_calls = 0 + + def __deepcopy__(self, memo: dict[int, Any]) -> _DeepcopySpy: # type: ignore[type-arg] + self.deepcopy_calls += 1 + return _DeepcopySpy() + + +def test_prewarm_placeholder_input_filters_discarded_media() -> None: + """Discarded prompt fields (media) must never be deep-copied for a prewarm.""" + media = _DeepcopySpy() + orch = _prewarm_orchestrator_with_client(SimpleNamespace(prewarm_placeholder_builder=None)) + req_state = OrchestratorRequestState( + request_id="req-prewarm-media", + prompt={ + "prompt_token_ids": [1, 2], + "multi_modal_data": media, + "mm_processor_kwargs": {"foo": "bar"}, + }, + sampling_params_list=[SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + base_input = orch._build_prewarm_placeholder_input(0, "req-prewarm-media", [1, 2], req_state) + + assert base_input["multi_modal_data"] is None + assert base_input["mm_processor_kwargs"] is None + assert media.deepcopy_calls == 0 + + +def test_prewarm_placeholder_input_shares_additional_information_tensors() -> None: + """``additional_information`` containers are isolated but tensors stay shared.""" + tensor = torch.arange(4, dtype=torch.long) + original = {"t": tensor, "l": [1, 2]} + orch = _prewarm_orchestrator_with_client(SimpleNamespace(prewarm_placeholder_builder=None)) + req_state = OrchestratorRequestState( + request_id="req-prewarm-additional", + prompt={"prompt_token_ids": [1, 2], "additional_information": original}, + sampling_params_list=[SamplingParams(max_tokens=1)], + final_stage_id=0, + ) + + base_input = orch._build_prewarm_placeholder_input(0, "req-prewarm-additional", [1, 2], req_state) + + copied = base_input["additional_information"] + assert copied is not original + assert copied["l"] is not original["l"] + assert copied["t"] is tensor + assert copied["t"].data_ptr() == tensor.data_ptr() diff --git a/tests/engine/test_stage_metadata.py b/tests/engine/test_stage_metadata.py index 6974105691a..a5a57f8367a 100644 --- a/tests/engine/test_stage_metadata.py +++ b/tests/engine/test_stage_metadata.py @@ -4,6 +4,7 @@ import copy import operator from dataclasses import replace +from typing import Any import pytest from vllm.sampling_params import SamplingParams @@ -23,6 +24,7 @@ extract_stage_metadata_from_omni_stage_config, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.stage_input_processors import ProcessorValidationError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -195,3 +197,45 @@ def test_extract_stage_metadata_defaults_missing_engine_input_source(): metadata = extract_stage_metadata(legacy_config) assert metadata.engine_input_source == [] + + +# --------------------------------------------------------------------------- +# A bad custom_process_input_func must FAIL FAST. The config-hook resolution +# path no longer swallows resolve_processor failures and continues with +# custom_process_input_func=None (which silently fell back to +# _default_process_engine_inputs). The worker mixin candidate chain keeps its +# own skip-on-mismatch tolerance; this is the config-hook resolution path. +# --------------------------------------------------------------------------- + + +def _legacy_stage0_config() -> Any: + pipeline, deploy = _metadata_inputs() + return merge_pipeline_deploy( + pipeline, + copy.deepcopy(deploy), + )[0].to_omegaconf() + + +def test_extract_stage_metadata_fails_fast_on_bad_processor_module(): + legacy_config = _legacy_stage0_config() + legacy_config["custom_process_input_func"] = "nonexistent.module.does_not_exist" + with pytest.raises(ImportError): + extract_legacy_stage_metadata(legacy_config) + + +def test_extract_stage_metadata_fails_fast_on_missing_processor_attr(): + legacy_config = _legacy_stage0_config() + legacy_config["custom_process_input_func"] = "operator.does_not_exist" + with pytest.raises(AttributeError): + extract_legacy_stage_metadata(legacy_config) + + +def test_extract_stage_metadata_fails_fast_on_invalid_processor(): + # Resolves to a non-callable symbol -> ProcessorValidationError (rule + # "callable") instead of silently dropping the configured hook. + legacy_config = _legacy_stage0_config() + legacy_config["custom_process_input_func"] = ( + "vllm_omni.model_executor.stage_input_processors._common._IM_START_TOKEN_ID" + ) + with pytest.raises(ProcessorValidationError): + extract_legacy_stage_metadata(legacy_config) diff --git a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py index 64f892717e4..a914ab381b3 100755 --- a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py @@ -1427,10 +1427,14 @@ def test_load_weights_fails_unknown_prefix(self, monkeypatch): class TestStageInputProcessor: + def _common(self): + from vllm_omni.model_executor.stage_input_processors import _common + + return _common + def test_revert_delay_pattern(self): - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _revert_delay_pattern, - ) + c = self._common() + _NUM_CODEBOOKS = 8 # 8 codebooks, 3 real frames -> delayed shape [8, 3+8-1=10] Q, T = 8, 3 @@ -1438,31 +1442,27 @@ def test_revert_delay_pattern(self): for i in range(Q): for t in range(T): delayed[i, i + t] = i * 100 + t # Real codes at shifted positions - result = _revert_delay_pattern(delayed) + result = c.revert_delay_pattern(delayed, expected_codebooks=_NUM_CODEBOOKS, allow_short=False) assert result.shape == (Q, T) for i in range(Q): for t in range(T): assert result[i, t].item() == i * 100 + t def test_revert_delay_pattern_rejects_wrong_codebooks(self): - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _revert_delay_pattern, - ) + c = self._common() # 7 codebooks should be rejected codes = torch.zeros(7, 20) with pytest.raises(ValueError, match="Expected exactly 8 codebook rows"): - _revert_delay_pattern(codes) + c.revert_delay_pattern(codes, expected_codebooks=8, allow_short=False) def test_revert_delay_pattern_rejects_too_few_frames(self): - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _revert_delay_pattern, - ) + c = self._common() # 8 codebooks but only 5 frames (need at least 8) codes = torch.zeros(8, 5) with pytest.raises(ValueError, match="Not enough frames"): - _revert_delay_pattern(codes) + c.revert_delay_pattern(codes, expected_codebooks=8, allow_short=False) def test_talker2code2wav_skips_too_few_frames_without_crashing(self): from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( @@ -1482,9 +1482,8 @@ class TalkerOutput: assert result[0]["prompt_token_ids"] == [] def test_filter_real_code_frames(self): - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _filter_real_code_frames, - ) + c = self._common() + _NUM_REAL_CODES = 1024 # 8 codebooks, 4 frames codes = torch.tensor( @@ -1499,7 +1498,7 @@ def test_filter_real_code_frames(self): [107, 207, 1024, 307], ] ) - result = _filter_real_code_frames(codes) + result = c.filter_real_code_frames(codes, num_real_codes=_NUM_REAL_CODES, layout="codebooks_first") # Frame 2 (column 2) has BOC in all codebooks -> filtered out assert result.shape == (8, 3) assert result[0, 0].item() == 100 diff --git a/tests/model_executor/stage_input_processors/conftest.py b/tests/model_executor/stage_input_processors/conftest.py new file mode 100644 index 00000000000..8be9a5e9ed8 --- /dev/null +++ b/tests/model_executor/stage_input_processors/conftest.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Test-only import fallback: lets pure-logic tests run without real ``vllm``. + +Bypasses ``vllm_omni/__init__.py`` via a synthetic package in ``sys.modules`` +and stubs the ``vllm`` tree with permissive ``_Dummy`` classes. Active only +when real ``vllm`` is unavailable (CI no-op), so it cannot mask genuine +import problems. Scope: pure-logic unit tests only. +""" + +import importlib.abc +import importlib.machinery +import os +import sys +import types + +# --------------------------------------------------------------------------- +# Locate the repo root (this file lives at: +# /tests/model_executor/stage_input_processors/conftest.py) +# --------------------------------------------------------------------------- +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +_VLLM_OMNI_DIR = os.path.join(_REPO_ROOT, "vllm_omni") + + +class _Meta(type): + """Class-level attribute access on a ``_Dummy`` also yields ``_Dummy``.""" + + def __getattr__(cls, name): + return _Dummy + + +class _Dummy(metaclass=_Meta): + """Permissive stand-in for any vllm symbol (import-grade compatibility).""" + + def __init__(self, *args, **kwargs): + pass + + def __call__(self, *args, **kwargs): + return _Dummy() + + def __getattr__(self, name): + return _Dummy + + def __class_getitem__(cls, item): + return cls + + def __mro_entries__(self, bases): + return (_Dummy,) + + def __init_subclass__(cls, **kwargs): + pass + + def __or__(self, other): + return _Dummy + + def __ror__(self, other): + return _Dummy + + +class _StubLoader(importlib.abc.Loader): + """Creates an empty module with a wildcard __getattr__ for any vllm.*.""" + + def create_module(self, spec): + module = types.ModuleType(spec.name) + module.__path__ = [] + module.__getattr__ = lambda attr: _Dummy # type: ignore[method-assign] + return module + + def exec_module(self, module): + pass + + +class _StubFinder(importlib.abc.MetaPathFinder): + """Auto-creates a stub package for ``vllm`` and any ``vllm.*`` submodule.""" + + def find_spec(self, fullname, path=None, target=None): + if fullname == "vllm" or fullname.startswith("vllm."): + return importlib.machinery.ModuleSpec(fullname, _StubLoader(), is_package=True) + return None + + +def _install_shim(): + """Bypass the real vllm_omni/__init__ and stub the vllm package tree.""" + if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + # Synthetic vllm_omni parent package (skips the real __init__). + pkg = types.ModuleType("vllm_omni") + pkg.__path__ = [_VLLM_OMNI_DIR] + pkg.__package__ = "vllm_omni" + pkg._SHIM_ACTIVE = True # type: ignore[attr-defined] + sys.modules.setdefault("vllm_omni", pkg) + # Wildcard vllm stub finder (meta path -> takes precedence for vllm.*). + sys.meta_path.insert(0, _StubFinder()) + + +def _activate_if_needed(): + """Install the import fallback only when real vllm is unavailable.""" + try: + import vllm # noqa: F401 + except ModuleNotFoundError: + _install_shim() + + +_activate_if_needed() diff --git a/tests/model_executor/stage_input_processors/test_common_helpers_golden.py b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py new file mode 100644 index 00000000000..99155faa23f --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_common_helpers_golden.py @@ -0,0 +1,555 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Golden tests for the consolidated ``_common`` stage-input-processor helpers. + +Each case encodes a documented semantic difference (None handling, tensor -> +list, ``ConstantList._x`` unwrapping, tuple handling, codec-frame validity +masks, delay-pattern strictness, ...). Modules whose optional dependencies are +unavailable in the current environment are skipped and exercised on CI. +""" + +import functools +import importlib +from typing import Any + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_PREFIX = "vllm_omni.model_executor.stage_input_processors." + + +def _import(name: str) -> Any: + """Import a stage-input-processor module (or diffusion output_formatter).""" + try: + if name == "diffusion.output_formatter": + return importlib.import_module("vllm_omni.diffusion.output_formatter") + if name == "higgs_audio_v2": + # higgs_audio_v2 keeps its own local _extract_last_frame. + return importlib.import_module(_PREFIX + name) + return importlib.import_module(_PREFIX + name) + except Exception as exc: # pragma: no cover - env dependent + pytest.skip(f"module {name!r} not importable here: {type(exc).__name__}: {exc}") + raise AssertionError("unreachable") + + +def _common_import() -> Any: + """Import the canonical shared helper module.""" + try: + import vllm_omni.model_executor.stage_input_processors._common as c + + return c + except Exception as exc: # pragma: no cover - env dependent + pytest.skip(f"_common not importable here: {type(exc).__name__}: {exc}") + + +class _ConstantList: + """Mimics the ``ConstantList._x`` attribute seen across modules.""" + + def __init__(self, values): + self._x = list(values) + + +def _norm(value: Any) -> Any: + """Normalize a result for comparison: tensors -> nested lists.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + return value + + +def _is_tensor(value: Any) -> bool: + return isinstance(value, torch.Tensor) + + +# =========================================================================== +# _ensure_list +# =========================================================================== + + +_T2 = torch.tensor([1, 2]) +_T22 = torch.tensor([[1, 2], [3, 4]]) +_T24 = torch.tensor([1, 2, 3, 4]) + + +def _ensure_list_cases(): + c = _common_import() + mapping = { + "qwen3_omni": c.ensure_list_unchanged, + "ming_flash_omni": c.ensure_list_strict, + "audex": c.ensure_list_flatten, + "cosyvoice3": c.ensure_list, + "step_audio2": c.ensure_list_preserve_none, + "diffusion.output_formatter": _import("diffusion.output_formatter")._ensure_list, + } + # (module, callable, input, expected, expect_typeerror) + return [ + pytest.param(mapping["qwen3_omni"], [1, 2], [1, 2], False, id="qwen3_omni-list"), + pytest.param(mapping["qwen3_omni"], _ConstantList([1, 2]), [1, 2], False, id="qwen3_omni-constantlist"), + pytest.param(mapping["qwen3_omni"], None, None, False, id="qwen3_omni-none"), + pytest.param(mapping["qwen3_omni"], (1, 2), (1, 2), False, id="qwen3_omni-tuple"), + pytest.param(mapping["qwen3_omni"], _T2, None, False, id="qwen3_omni-tensor"), + pytest.param(mapping["ming_flash_omni"], [1, 2], [1, 2], False, id="ming_flash_omni-list"), + pytest.param( + mapping["ming_flash_omni"], _ConstantList([1, 2]), [1, 2], False, id="ming_flash_omni-constantlist" + ), + pytest.param(mapping["ming_flash_omni"], _T2, [1, 2], False, id="ming_flash_omni-tensor"), + pytest.param(mapping["ming_flash_omni"], (1, 2), [1, 2], False, id="ming_flash_omni-tuple"), + pytest.param(mapping["ming_flash_omni"], None, None, True, id="ming_flash_omni-none-raises"), + pytest.param(mapping["audex"], None, [], False, id="audex-none"), + pytest.param(mapping["audex"], _T24, [1, 2, 3, 4], False, id="audex-tensor"), + pytest.param(mapping["audex"], [1, 2], [1, 2], False, id="audex-list"), + pytest.param(mapping["audex"], (1, 2), [1, 2], False, id="audex-tuple"), + pytest.param(mapping["audex"], 5, None, True, id="audex-scalar-raises"), + pytest.param(mapping["cosyvoice3"], None, [], False, id="cosyvoice3-none"), + pytest.param(mapping["cosyvoice3"], (1, 2), [1, 2], False, id="cosyvoice3-tuple"), + pytest.param(mapping["cosyvoice3"], 5, [5], False, id="cosyvoice3-scalar"), + pytest.param(mapping["cosyvoice3"], _ConstantList([1, 2]), [1, 2], False, id="cosyvoice3-constantlist"), + pytest.param(mapping["cosyvoice3"], [1, 2], [1, 2], False, id="cosyvoice3-list"), + pytest.param(mapping["step_audio2"], _ConstantList([1, 2]), [1, 2], False, id="step_audio2-constantlist"), + pytest.param(mapping["step_audio2"], _T2, [1, 2], False, id="step_audio2-tensor"), + pytest.param(mapping["step_audio2"], (1, 2), [1, 2], False, id="step_audio2-tuple"), + pytest.param(mapping["step_audio2"], 5, [5], False, id="step_audio2-scalar"), + pytest.param(mapping["step_audio2"], None, [None], False, id="step_audio2-none"), + pytest.param(mapping["diffusion.output_formatter"], [1, 2], [1, 2], False, id="output_formatter-list"), + pytest.param(mapping["diffusion.output_formatter"], None, [], False, id="output_formatter-none"), + pytest.param(mapping["diffusion.output_formatter"], 5, [5], False, id="output_formatter-scalar"), + ] + + +@pytest.mark.parametrize("fn,inp,expected,expect_typeerror", _ensure_list_cases()) +def test_ensure_list_golden(fn, inp, expected, expect_typeerror): + if expect_typeerror: + with pytest.raises(TypeError): + fn(inp) + return + result = fn(inp) + if _is_tensor(result): + # qwen3_omni may return a tensor unchanged. + assert expected is None + assert torch.equal(result, _T2) + return + assert _norm(result) == expected + + +def test_output_formatter_ensure_list_wrap_only(): + """output_formatter._ensure_list keeps wrap-only semantics. + + The diffusion formatter is NOT a stage-input helper: a non-list primary + payload is wrapped as ``[x]`` verbatim — no tensor flattening, no dict-key + iteration, no row-wise iteration of an iterable (e.g. a PIL ``Image``). + ``_common.ensure_list`` (the processor flatten) must NOT be reused here. + """ + mod = _import("diffusion.output_formatter") + fn = getattr(mod, "_ensure_list") + assert fn([1, 2]) == [1, 2] + assert fn(None) == [] + assert fn(5) == [5] + # A tensor is wrapped whole, NOT flattened via .tolist(). + t = torch.tensor([1, 2]) + (wrapped,) = fn(t) + assert wrapped is t + # A dict is wrapped whole, NOT iterated over its keys. + d = {"a": 1} + (wrapped_dict,) = fn(d) + assert wrapped_dict is d + + # A row-iterable (PIL.Image-like) is wrapped whole, NOT walked row by row. + class _RowIterable: + def __iter__(self): + yield from [[1, 2], [3, 4]] + + img = _RowIterable() + (wrapped_img,) = fn(img) + assert wrapped_img is img + # Contrast: the canonical processor flatten would have flattened each. + c = _common_import() + assert c.ensure_list(t) == [1, 2] # flattened, not [t] + assert c.ensure_list(d) == ["a"] # dict keys iterated + assert c.ensure_list(img) == [[1, 2], [3, 4]] # rows walked + + +# =========================================================================== +# _extract_last_frame +# =========================================================================== + + +def _codes_audio(tensor): + return {"codes": {"audio": tensor}} + + +def _extract_last_frame_cases(): + """Legacy per-module behavior: one callable + input per table row.""" + c = _common_import() + mapping = { + # higgs_audio_v2 keeps its own local range-filtered helper (RFC refactor + # removed the other per-module helpers; its range filter is not + # expressible through _common.extract_last_codec_frame). + "higgs_audio_v2": _import("higgs_audio_v2")._extract_last_frame, + "fish_speech": functools.partial( + c.extract_last_codec_frame, + key_path=("audio_codes",), + validate="valid_mask", + to_cpu=True, + to_long=True, + ), + "qwen3_tts": functools.partial(c.extract_last_codec_frame, validate="any"), + # voxtral_tts: no validity gate, no cpu/long cast. Under the + # consolidated implementation a 2-D ``codes.audio`` tensor contributes + # its LAST row only (legacy voxtral flattened the whole tensor; the + # refactor narrows this to the last inter-stage frame, matching all + # other modules). + "voxtral_tts": functools.partial(c.extract_last_codec_frame, to_long=False), + } + higgs = mapping["higgs_audio_v2"] + fish = mapping["fish_speech"] + qwen3 = mapping["qwen3_tts"] + voxtral = mapping["voxtral_tts"] + return [ + # higgs_audio_v2: codes.audio key path; range filter [-0, _NUM_REAL_CODES). + pytest.param(higgs, _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False, id="higgs-last-row"), + # >= _NUM_REAL_CODES + pytest.param(higgs, _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 1024]])), None, False, id="higgs-range"), + pytest.param(higgs, _codes_audio(torch.tensor([[0, 1], [2, 3], [-1, 5]])), None, False, id="higgs-negative"), + pytest.param(higgs, _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False, id="higgs-1d"), + pytest.param(higgs, _codes_audio(torch.empty(0, 4)), None, False, id="higgs-empty"), + # 3-D -> ValueError + pytest.param(higgs, _codes_audio(torch.zeros(2, 2, 2)), None, True, id="higgs-3d-raises"), + pytest.param(higgs, {"codes": {}}, None, False, id="higgs-missing-audio"), + # fish_speech: TOP-LEVEL audio_codes key; audio_code_valid mask; cpu+long. + pytest.param( + fish, + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": torch.tensor([1, 1, 1])}, + [4, 5], + False, + id="fish-valid", + ), + pytest.param( + fish, + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": torch.tensor([1, 1, 0])}, + None, + False, + id="fish-invalid-last", + ), + pytest.param( + fish, + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": True}, + [4, 5], + False, + id="fish-valid-true", + ), + pytest.param( + fish, + {"audio_codes": torch.tensor([[0, 1], [2, 3], [4, 5]]), "audio_code_valid": False}, + None, + False, + id="fish-invalid-false", + ), + # no valid key -> frame.any() False + pytest.param(fish, {"audio_codes": torch.tensor([[0, 1], [2, 3], [0, 0]])}, None, False, id="fish-zero-frame"), + pytest.param(fish, {"audio_codes": torch.tensor([1, 2, 3])}, [1, 2, 3], False, id="fish-1d"), + # 1-D semantics: flat frames bypass the validity/zero gate (amy-why-3459 review). + pytest.param(fish, {"audio_codes": torch.zeros(8, dtype=torch.long)}, [0] * 8, False, id="fish-1d-zeros"), + pytest.param( + fish, + {"audio_codes": torch.zeros(8, dtype=torch.long), "audio_code_valid": torch.tensor([False])}, + [0] * 8, + False, + id="fish-1d-zeros-invalid", + ), + pytest.param( + fish, + {"audio_codes": torch.zeros(8, dtype=torch.long), "audio_code_valid": False}, + [0] * 8, + False, + id="fish-1d-zeros-false", + ), + pytest.param(fish, {"audio_codes": torch.zeros(2, 2, 2)}, None, True, id="fish-3d-raises"), + pytest.param(fish, {}, None, False, id="fish-missing-codes"), + # qwen3_tts: codes.audio; frame.any() gate; long (no cpu). + pytest.param(qwen3, _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False, id="qwen3-last-row"), + pytest.param(qwen3, _codes_audio(torch.tensor([[0, 1], [2, 3], [0, 0]])), None, False, id="qwen3-zero-frame"), + pytest.param(qwen3, _codes_audio(torch.tensor([1, 2, 3])), [1, 2, 3], False, id="qwen3-1d"), + # 1-D semantics: flat frames bypass the zero gate (legacy qwen3_tts behavior). + pytest.param(qwen3, _codes_audio(torch.zeros(8, dtype=torch.long)), [0] * 8, False, id="qwen3-1d-zeros"), + pytest.param(qwen3, _codes_audio(torch.zeros(2, 2, 2)), None, True, id="qwen3-3d-raises"), + pytest.param(qwen3, {"codes": {}}, None, False, id="qwen3-missing-audio"), + # voxtral_tts: validated via _common.extract_last_codec_frame(to_long=False). + pytest.param( + voxtral, _codes_audio(torch.tensor([[0, 1], [2, 3], [4, 5]])), [4, 5], False, id="voxtral-last-row" + ), + pytest.param( + voxtral, + {"codes": {"audio": [torch.tensor([1, 2]), torch.tensor([3, 4])]}}, + [3, 4], + False, + id="voxtral-list-audio", + ), + pytest.param(voxtral, {"codes": {}}, None, False, id="voxtral-missing-audio"), + # 3-D per-step frames [B, 1, C] (Voxtral compute_mm_logits L1075): + # normalized via targeted squeeze(1) then the 2-D last-row semantics. + pytest.param( + voxtral, + _codes_audio(torch.arange(8, dtype=torch.float32).reshape(1, 1, 8)), + [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], + False, + id="voxtral-3d-per-step", + ), + # Mirror the real payload: a list whose last element is 3-D (1, 1, C). + pytest.param( + voxtral, + {"codes": {"audio": [torch.ones(1, 1, 4)]}}, + [1.0, 1.0, 1.0, 1.0], + False, + id="voxtral-3d-list-element", + ), + # Non per-step 3-D layouts keep raising ValueError (dim1 != 1). + pytest.param(voxtral, _codes_audio(torch.zeros(1, 2, 3)), None, True, id="voxtral-3d-non-singleton-raises"), + # Generic 3-D (aligned with the fish/qwen3/higgs contracts). + pytest.param(voxtral, _codes_audio(torch.zeros(2, 2, 2)), None, True, id="voxtral-3d-generic-raises"), + ] + + +@pytest.mark.parametrize("fn,payload,expected,expect_error", _extract_last_frame_cases()) +def test_extract_last_frame_golden(fn, payload, expected, expect_error): + if expect_error: + with pytest.raises(ValueError): + fn(payload) + return + assert _norm(fn(payload)) == expected + + +# =========================================================================== +# revert_delay_pattern / filter_real_code_frames +# =========================================================================== +# Behavior is locked by +# ``test_common_revert_delay_pattern_matches_golden`` and +# ``test_common_filter_real_code_frames_matches_golden`` below. + +# =========================================================================== +# _to_cpu_tensor (glm_tts) +# =========================================================================== + + +_T2F = torch.tensor([1.0, 2.0]) + + +@pytest.mark.parametrize( + "inp,expected", + [ + pytest.param(_T2F, _T2F, id="tensor"), + pytest.param([], None, id="empty-list"), + pytest.param([_T2F], _T2F, id="list-tensor"), + pytest.param([1, 2], None, id="list-int"), + pytest.param(None, None, id="none"), + pytest.param(5, None, id="int"), + ], +) +def test_to_cpu_tensor_glm_tts_golden(inp, expected): + result = _common_import().to_cpu_tensor(inp) + if expected is None: + assert result is None + else: + assert torch.equal(result, expected) + + +# =========================================================================== +# _to_token_id_list (dynin_omni; cosyvoice3 recursive variant) +# =========================================================================== + + +@pytest.mark.parametrize( + "inp,expected", + [ + pytest.param(torch.tensor(5), [5], id="scalar-tensor"), + pytest.param(torch.tensor([1, 2, 3]), [1, 2, 3], id="vec-tensor"), + pytest.param(torch.tensor([[1, 2], [3, 4]]), [1, 2], id="mat-tensor"), # 2-D -> first row only + pytest.param([[1, 2], [3, 4]], [1, 2], id="nested-list"), # nested list -> first nested row only + pytest.param([1, 2, 3], [1, 2, 3], id="list"), + pytest.param(None, [], id="none"), + pytest.param(5, [5], id="int"), + ], +) +def test_to_token_id_list_dynin_omni_golden(inp, expected): + assert _common_import().to_token_id_list(inp) == expected + + +def test_to_token_id_list_cosyvoice3_recursive_golden(): + c = _common_import() + fn = functools.partial(c.to_token_id_list, recursive=True) + # cosyvoice3 recursively flattens ALL nesting (dynin takes first row only). + assert fn([[1, 2], [3, 4]]) == [1, 2, 3, 4] + assert fn(None) == [] + assert fn(torch.tensor([[1, 2], [3, 4]])) == [1, 2, 3, 4] + assert fn(5) == [5] + # Per-item normalization flattens a list containing a non-scalar tensor and + # a plain tuple. + assert fn([torch.tensor([[4, 5]])]) == [4, 5] + assert fn((4, 5)) == [4, 5] + + +# =========================================================================== +# placeholder length helper (Qwen chat-template scan). +# =========================================================================== + + +def test_compute_talker_prompt_ids_length_golden(): + c = _common_import() + # im_start markers at 0 and 6; user segment len 6, last assistant +9. + prompt = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] + info = {"ids": {"all": prompt, "prompt": prompt}} + assert c.compute_placeholder_prompt_len(ids_or_prompt=info, mode="full", device="cpu") == 15 + # stage0_only scans the flat stage-0 list (same 15 for the golden prompt). + assert c.compute_placeholder_prompt_len(ids_or_prompt=prompt, mode="stage0_only") == 15 + + +# =========================================================================== +# The shared _common.py helpers must reproduce the golden-locked behaviour. +# =========================================================================== + + +def _shim_active() -> bool: + """True when the test-support import fallback (no real vllm) is active.""" + try: + import vllm_omni # noqa: F401 + + return bool(getattr(vllm_omni, "_SHIM_ACTIVE", False)) + except Exception: + return False + + +def test_common_ensure_list_matches_golden(): + c = _common_import() + assert c.ensure_list([1, 2]) == [1, 2] + assert c.ensure_list((1, 2)) == [1, 2] + assert c.ensure_list(None) == [] + assert c.ensure_list(_ConstantList([1, 2])) == [1, 2] + assert c.ensure_list(5) == [5] + # tolist branch (2-D -> list-of-lists, no flattening). + assert c.ensure_list(torch.tensor([[1, 2], [3, 4]])) == [[1, 2], [3, 4]] + # Named variants. + assert c.ensure_list_unchanged(None) is None + assert c.ensure_list_unchanged((1, 2)) == (1, 2) + assert c.ensure_list_unchanged([1, 2]) == [1, 2] + # Strict: non-iterable scalar raises (ming_flash / audex legacy). + with pytest.raises(TypeError): + c.ensure_list_strict(5) + + +def test_common_to_cpu_tensor_matches_golden(): + c = _common_import() + t = torch.tensor([1.0, 2.0]) + assert torch.equal(c.to_cpu_tensor(t), t) + assert c.to_cpu_tensor([]) is None + assert torch.equal(c.to_cpu_tensor([t]), t) + assert c.to_cpu_tensor([1, 2]) is None + assert c.to_cpu_tensor(None) is None + + +def test_common_to_token_id_list_matches_golden(): + c = _common_import() + # dynin semantics: first nested row only (non-recursive). + assert c.to_token_id_list(torch.tensor(5)) == [5] + assert c.to_token_id_list(torch.tensor([1, 2, 3])) == [1, 2, 3] + assert c.to_token_id_list(torch.tensor([[1, 2], [3, 4]])) == [1, 2] + assert c.to_token_id_list([[1, 2], [3, 4]]) == [1, 2] + assert c.to_token_id_list(None) == [] + assert c.to_token_id_list(5) == [5] + # cosyvoice3 semantics: recursive flatten. + assert c.to_token_id_list([[1, 2], [3, 4]], recursive=True) == [1, 2, 3, 4] + assert c.to_token_id_list(torch.tensor([[1, 2], [3, 4]]), recursive=True) == [1, 2, 3, 4] + # Per-item recursion normalizes a list of non-scalar tensors and a tuple. + assert c.to_token_id_list([torch.tensor([[4, 5]])], recursive=True) == [4, 5] + assert c.to_token_id_list((4, 5), recursive=True) == [4, 5] + assert c.to_token_id_list((torch.tensor(4), torch.tensor([5, 6])), recursive=True) == [4, 5, 6] + + +def test_common_revert_delay_pattern_matches_golden(): + c = _common_import() + inp = torch.tensor( + [ + [10, 1, 2, 3, 4], + [11, 12, 5, 6, 7], + [13, 14, 15, 8, 9], + ] + ) # [Q=3, T=5] + assert c.revert_delay_pattern(inp, allow_short=True).tolist() == [[10, 1, 2], [12, 5, 6], [15, 8, 9]] + # Lenient (v2): T < Q returns input unchanged. + short = torch.tensor([[1], [2]]) + assert torch.equal(c.revert_delay_pattern(short, allow_short=True), short) + # Strict (v3): T < Q raises; wrong codebook count raises. + with pytest.raises(ValueError): + c.revert_delay_pattern(short, allow_short=False) + with pytest.raises(ValueError): + c.revert_delay_pattern(torch.zeros(5, 6), expected_codebooks=3) + + +def test_common_filter_real_code_frames_matches_golden(): + c = _common_import() + nrc = 1024 + # frames-first (v2 layout). + frames_first = torch.tensor( + [ + [0, 1, 2], + [nrc, 5, 6], + [7, 8, 9], + [-1, 0, 1], + ] + ) + assert c.filter_real_code_frames(frames_first, num_real_codes=nrc, layout="frames_first").tolist() == [ + [0, 1, 2], + [7, 8, 9], + ] + # codebooks-first (v3 layout). + codebooks_first = ( + torch.tensor( + [ + [0, 1, 2], + [nrc, 5, 6], + [7, 8, 9], + [-1, 0, 1], + ] + ) + .t() + .contiguous() + ) # [Q=3, frames=4] + assert c.filter_real_code_frames(codebooks_first, num_real_codes=nrc, layout="codebooks_first").tolist() == [ + [0, 7], + [1, 8], + [2, 9], + ] + + +def test_common_compute_placeholder_prompt_len_matches_golden(): + c = _common_import() + prompt = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] + # full mode == the legacy Qwen chat-template scan golden (15). + assert c.compute_placeholder_prompt_len(ids_or_prompt={"ids": {"all": prompt, "prompt": prompt}}, mode="full") == 15 + # stage0_only mode: the same scan on the flat stage-0 list (prewarm) -> 15, + # so the builder and the inline fallback agree. + assert c.compute_placeholder_prompt_len(ids_or_prompt=prompt, mode="stage0_only") == 15 + + +def test_common_pack_placeholder_prompt(): + c = _common_import() + prompt = c.pack_placeholder_prompt(prompt_len=4, voice_metadata={"speaker": 1}) + empty_meta = c.pack_placeholder_prompt(prompt_len=1) + if _shim_active(): + # Under the import fallback, OmniTokensPrompt's real constructor is + # bypassed, so field access returns fallback values. Field assertions + # run where real vllm constructs the model properly. + pytest.skip("OmniTokensPrompt field construction needs real vllm") + # Real vllm (0.27) builds OmniTokensPrompt as a dict subclass (MRO + # ['OmniTokensPrompt', 'dict', 'object']): attribute access is not + # available, so read the packed fields via dict indexing. + assert prompt["prompt_token_ids"] == [0, 0, 0, 0] + assert prompt["additional_information"] == {"speaker": 1} + assert empty_meta["prompt_token_ids"] == [0] + # When additional_information is None the key may be absent from the + # dict, so use .get() for the empty-metadata case. + assert empty_meta.get("additional_information") is None + + +# Keep a reference to the case builder so linters don't flag dead code. +_EXTRA_ALIAS = _ensure_list_cases diff --git a/tests/model_executor/stage_input_processors/test_constants_consistency.py b/tests/model_executor/stage_input_processors/test_constants_consistency.py new file mode 100644 index 00000000000..bc57329cf78 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_constants_consistency.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Constant-consistency guards. + +Verifies that magic special-token ids that were previously duplicated across +processors / pipelines / model implementations now resolve to a single value: + +* qwen3 codec ``4196/4197/4198``: processor ``_QWEN3_CODEC_*`` == model + ``TALKER_CODEC_*`` == HF config defaults (``codec_pad_id`` / + ``codec_bos_id`` / ``codec_eos_token_id``). +* qwen2_5 talker codec end ``8294``: processor ``TALKER_CODEC_END_TOKEN_ID`` + == model ``TALKER_CODEC_EOS_TOKEN_ID`` == pipeline ``stop_token_ids``. +* ``2150``: single-sourced in ``stage_input_processors._constants`` and shared + by the qwen3_omni / qwen3_tts / aura_omni pipelines. + +CPU-only: imports modules and reads constants; never loads model weights. +""" + +import pytest + +from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE +from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import QWEN2_5_OMNI_PIPELINE +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni import TALKER_CODEC_EOS_TOKEN_ID +from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import ( + TALKER_CODEC_BOS_TOKEN_ID, + TALKER_CODEC_PAD_TOKEN_ID, +) +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import ( + TALKER_CODEC_EOS_TOKEN_ID as QWEN3_MODEL_CODEC_EOS, +) +from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSTalkerConfig +from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE +from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID +from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID +from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( + _QWEN3_CODEC_BOS_TOKEN_ID, + _QWEN3_CODEC_EOS_TOKEN_ID, + _QWEN3_CODEC_PAD_TOKEN_ID, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Expected values locked by the baseline commit `731a771`. +QWEN3_CODEC_EXPECTED = (4196, 4197, 4198) # (pad, bos, eos) +QWEN2_5_CODEC_END_EXPECTED = 8294 +QWEN3_CODEC_STOP_EXPECTED = 2150 + +# Processor-side codec ids, single-sourced (pad, bos, eos). +QWEN3_CODEC_PROCESSOR_TUPLE = (_QWEN3_CODEC_PAD_TOKEN_ID, _QWEN3_CODEC_BOS_TOKEN_ID, _QWEN3_CODEC_EOS_TOKEN_ID) + + +def _pipeline_stop_token_ids(pipeline) -> list[int]: + """Extract the (unique) ``stop_token_ids`` across a pipeline's stages.""" + ids: list[int] = [] + for stage in pipeline.stages: + constraints = getattr(stage, "sampling_constraints", None) or {} + stage_ids = constraints.get("stop_token_ids") + if stage_ids: + ids.extend(stage_ids) + return ids + + +def test_qwen3_codec_token_ids_processor_matches_model(): + """4196/4197/4198: processor and model implementation must agree.""" + model = (TALKER_CODEC_PAD_TOKEN_ID, TALKER_CODEC_BOS_TOKEN_ID, QWEN3_MODEL_CODEC_EOS) + assert QWEN3_CODEC_PROCESSOR_TUPLE == model == QWEN3_CODEC_EXPECTED + + +def test_qwen3_codec_token_ids_match_hf_config(): + """4196/4197/4198: processor constants must match the in-repo Qwen3-TTS + HF config defaults (single source of truth).""" + cfg = Qwen3TTSTalkerConfig() + hf = (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) + assert hf == QWEN3_CODEC_EXPECTED + assert QWEN3_CODEC_PROCESSOR_TUPLE == hf + + +def test_qwen3_omni_runtime_talker_config_matches_processor(): + """The transformers ``Qwen3OmniMoeTalkerConfig`` (the config qwen3-omni + actually runs with) must also carry 4196/4197/4198.""" + from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeTalkerConfig, + ) + + cfg = Qwen3OmniMoeTalkerConfig() + assert (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) == QWEN3_CODEC_EXPECTED + + +def test_qwen2_5_stop_token_consistent(): + """8294: processor / model implementation / pipeline must agree.""" + assert TALKER_CODEC_END_TOKEN_ID == QWEN2_5_CODEC_END_EXPECTED + assert TALKER_CODEC_EOS_TOKEN_ID == TALKER_CODEC_END_TOKEN_ID + assert _pipeline_stop_token_ids(QWEN2_5_OMNI_PIPELINE) == [TALKER_CODEC_END_TOKEN_ID] + + +def test_qwen3_codec_stop_token_shared_across_pipelines(): + """2150: the three pipelines must share the single-sourced constant.""" + assert QWEN3_CODEC_EOS_TOKEN_ID == QWEN3_CODEC_STOP_EXPECTED + for pipeline in (QWEN3_OMNI_PIPELINE, QWEN3_TTS_PIPELINE, AURA_OMNI_PIPELINE): + assert _pipeline_stop_token_ids(pipeline) == [QWEN3_CODEC_EOS_TOKEN_ID] diff --git a/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py new file mode 100644 index 00000000000..e5f93635424 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_dispatch_wrap_once.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Orchestrator dispatch wrap-once semantics (table-driven). + +``invoke_orchestrator_processor`` runs per forward, so +``wrap_orchestrator_processor`` must not re-wrap — nor re-emit the legacy +``DeprecationWarning`` — on every forward: the result is memoized per callable. + +Pure-logic, CPU-only (no model loading, no vllm runtime). +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from vllm_omni.model_executor.stage_input_processors._dispatch import ( + OrchestratorInputContext, + invoke_orchestrator_processor, + wrap_orchestrator_processor, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _legacy_c2(source_outputs, prompt=None, requires_multimodal_data=False, streaming_context=None): + """Legacy C2 positional shell (like qwen3_omni.thinker2talker_token_only).""" + return [source_outputs, prompt, requires_multimodal_data, streaming_context] + + +def _legacy_c0(source_outputs, prompt=None, requires_multimodal_data=False): + """Legacy C0 3-arg shell.""" + return [source_outputs, prompt, requires_multimodal_data] + + +def _c1(source_outputs: list[Any], ctx: OrchestratorInputContext): + """Already speaks the C1 ``(source_outputs, ctx)`` contract.""" + return [source_outputs, ctx] + + +def _legacy_c3_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, sampling_params=None): + """C3 shell whose ``sampling_params`` is keyword-only.""" + return [source_outputs, prompt, requires_multimodal_data, sampling_params] + + +def _legacy_c3_required_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, sampling_params): + """C3 shell with a REQUIRED keyword-only ``sampling_params``.""" + return [source_outputs, prompt, requires_multimodal_data, sampling_params] + + +def _legacy_c2_other_name(source_outputs, prompt=None, requires_multimodal_data=False, streaming_ctx=None): + """Four-positional shell with a non-standard streaming param name.""" + return [source_outputs, prompt, requires_multimodal_data, streaming_ctx] + + +def _legacy_c2_keyword_only(source_outputs, prompt=None, requires_multimodal_data=False, *, streaming_context=None): + """C2 shell whose ``streaming_context`` is keyword-only.""" + return [source_outputs, prompt, requires_multimodal_data, streaming_context] + + +def _count_deprecation_warnings(trigger) -> int: + """Run *trigger* under ``simplefilter("always")``; count DeprecationWarnings.""" + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + trigger() + return sum(1 for w in caught if issubclass(w.category, DeprecationWarning)) + + +_INVOKE_CASES = [ + pytest.param( + _legacy_c2, + {"prompt": "p", "requires_multimodal_data": True, "streaming_context": "s"}, + [["out"], "p", True, "s"], + id="c2-streaming-context", + ), + pytest.param( + _legacy_c3_keyword_only, + {"prompt": "p", "requires_multimodal_data": True, "sampling_params": "sp"}, + [["out"], "p", True, "sp"], + id="c3-keyword-only-sampling-params", + ), + pytest.param( + _legacy_c3_required_keyword_only, + {"prompt": "p", "requires_multimodal_data": True, "sampling_params": "sp"}, + [["out"], "p", True, "sp"], + id="c3-required-keyword-only-sampling-params", + ), + pytest.param( + _legacy_c2_other_name, + {"prompt": "p", "requires_multimodal_data": True, "streaming_context": "s"}, + [["out"], "p", True, "s"], + id="c2-four-positional-other-name", + ), + pytest.param( + _legacy_c2_keyword_only, + {"prompt": "p", "requires_multimodal_data": True, "streaming_context": "s"}, + [["out"], "p", True, "s"], + id="c2-keyword-only-streaming-context", + ), +] + + +def test_wrap_legacy_shape_warns_once_then_memoized(): + assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) == 1 + # Subsequent wraps reuse the memoized adapter: no further warning. + assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) == 0 + assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_legacy_c2)) == 0 + + +@pytest.mark.parametrize("fn,ctx_kwargs,expected", _INVOKE_CASES) +def test_invoke_adapts_legacy_shape_and_forwards_context(fn, ctx_kwargs, expected): + ctx = OrchestratorInputContext(**ctx_kwargs) + assert invoke_orchestrator_processor(fn, ["out"], ctx) == expected + + +def test_invoke_does_not_rewarn_on_every_forward(): + ctx = OrchestratorInputContext(prompt="p", requires_multimodal_data=True, streaming_context="s") + + def invoke_many(): + for _ in range(5): + invoke_orchestrator_processor(_legacy_c0, ["out"], ctx) + + # Exactly one DeprecationWarning total across 5 forwards (wrap happens once). + assert _count_deprecation_warnings(invoke_many) == 1 + + +def test_c1_processor_passes_through_without_warning(): + wrapped = wrap_orchestrator_processor(_c1) + assert wrapped is _c1 # already speaks ctx: returned unchanged + ctx = OrchestratorInputContext(prompt="p") + assert invoke_orchestrator_processor(_c1, ["out"], ctx) == [["out"], ctx] + assert _count_deprecation_warnings(lambda: wrap_orchestrator_processor(_c1)) == 0 diff --git a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py index 481408956b0..0691694bb35 100644 --- a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py +++ b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py @@ -46,6 +46,47 @@ def test_thinker2talker_skips_unfinished() -> None: assert thinker2talker_token_only([out]) == [] +def test_thinker2talker_via_process_engine_inputs_uses_configured_talker_len() -> None: + """The stage client's ``_stage_hf_config`` reaches the legacy c2pos + processor as the 4th positional ``next_stage_hf_config`` so the native + talker placeholder is sized to ``talker_init_len``. Before the fix + ``ctx.streaming_context`` (None) was passed instead and the placeholder + fell back to one token. + """ + from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient + + client = object.__new__(StageEngineCoreClient) + client.custom_process_input_func = thinker2talker_token_only + client.requires_multimodal_data = False + client._stage_hf_config = SimpleNamespace(use_native_talker=True, talker_init_len=37) + client.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) + + source_outputs = [_thinker_output(prompt_len=5, generated=[7, 8, 9, 10])] + talker_inputs = client.process_engine_inputs(source_outputs, prompt=None) + + [talker_input] = talker_inputs + assert talker_input["prompt_token_ids"] == [0] * 37 + + +def test_thinker2talker_via_process_engine_inputs_falls_back_without_config() -> None: + """Without ``_stage_hf_config`` the legacy processor keeps its one-token + placeholder fallback, preserving backward-compatible semantics. + """ + from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient + + client = object.__new__(StageEngineCoreClient) + client.custom_process_input_func = thinker2talker_token_only + client.requires_multimodal_data = False + client._stage_hf_config = None + client.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) + + source_outputs = [_thinker_output(prompt_len=5, generated=[7, 8, 9, 10])] + talker_inputs = client.process_engine_inputs(source_outputs, prompt=None) + + [talker_input] = talker_inputs + assert talker_input["prompt_token_ids"] == [0] + + def test_full_payload_trims_prompt_region_rows() -> None: # Rows correspond to timeline steps t=1..T-1; with prompt_len P the NeMo # trim keeps rows for t >= P, i.e. row index P-1 onward. diff --git a/tests/model_executor/stage_input_processors/test_placeholder_parity.py b/tests/model_executor/stage_input_processors/test_placeholder_parity.py new file mode 100644 index 00000000000..d6428f80af2 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_placeholder_parity.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Parity tests: async-chunk prewarm placeholder vs sync forward placeholder. + +``*_token_only`` placeholder builders expose a dual entry — +``build_forward_placeholder(source_outputs, ctx)`` for the sync forward path +and ``build_prewarm_placeholder(*, stage0_prompt, ctx, downstream_stage_id)`` for +the async-chunk prewarm path (no ``source_outputs`` yet). Both must be fed by +the same ``_common`` length / packing helpers so async and sync placeholders stay +consistent, and the streaming fixup helper +(``adapter.compute_talker_prompt_ids_length``) must use the same length helper. + +Length semantics +---------------- +- forward (``mode="full"``) is the Qwen chat-template scan (golden-locked at 15 + for the canonical prompt ``test_common_helpers_golden.py``). +- prewarm (``mode="stage0_only"``) runs the **same** Qwen chat-template scan on + the stage-0 input list (feeding it to both ``all``/``prompt`` roles), so the + prewarm estimate equals the forward placeholder length and the adapter's + streaming-fixup scan (15 == 15 for the golden prompt; 6 == 6 for a single user + segment). The connector fixup path replaces the estimate with the real + length once the upstream chunk arrives. + +These tests are CPU-only (no model loading). They run both without a vllm +runtime (the import fallback is active) and with real vllm, where +``OmniTokensPrompt`` is built as a dict subclass; length capture uses a +``pack_placeholder_prompt`` spy so it works in both environments. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm_omni.model_executor.stage_input_processors import _common +from vllm_omni.model_executor.stage_input_processors import qwen3_omni as q3 +from vllm_omni.model_executor.stage_input_processors._dispatch import OrchestratorInputContext +from vllm_omni.model_executor.stage_input_processors._registry import resolve_processor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +#: Canonical Qwen chat-template prompt used by the golden tests: im_start at 0/6, +#: one user segment of 6 tokens + one trailing assistant segment (+9) -> full 15. +GOLDEN_PROMPT = [151644, 872, 10, 11, 12, 13, 151644, 77091, 20, 21] + +#: Single user segment (no trailing assistant marker): full == stage0_only == 6. +SINGLE_USER_PROMPT = [151644, 872, 10, 11, 12, 13] + + +def _shim_active() -> bool: + """True when the test-support import fallback (no real vllm) is active.""" + try: + import vllm_omni # noqa: F401 + + return bool(getattr(vllm_omni, "_SHIM_ACTIVE", False)) + except Exception: + return False + + +def _source_outputs(prompt_token_ids, output_ids=()): + return [ + SimpleNamespace( + request_id="req-1", + prompt_token_ids=list(prompt_token_ids), + outputs=[SimpleNamespace(cumulative_token_ids=list(output_ids))], + ) + ] + + +def _ctx(prompt=None): + return OrchestratorInputContext( + prompt=prompt, + requires_multimodal_data=False, + streaming_context=None, + ) + + +def _capture_pack(monkeypatch): + """Spy on ``_common.pack_placeholder_prompt``; returns a recorder dict. + + ``pack_placeholder_prompt`` is the single packing chokepoint shared by both + builders, so capturing ``prompt_len`` / ``voice_metadata`` verifies the + builders' length logic without depending on how ``OmniTokensPrompt`` is + constructed (import fallback vs real vllm). + """ + + recorded: dict[str, Any] = {} + + def _spy(*, prompt_len: int, voice_metadata: dict[str, Any] | None = None) -> Any: + recorded["prompt_len"] = prompt_len + recorded["voice_metadata"] = voice_metadata + return "PLACEHOLDER" + + monkeypatch.setattr(_common, "pack_placeholder_prompt", _spy) + return recorded + + +# --------------------------------------------------------------------------- +# Forward builder (sync path) locks the golden full-mode length. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "source,ctx_prompt,expected_len,expected_voice", + [ + # Golden chat-template scan: sum(user segment) + 9 == 15. + pytest.param(_source_outputs(GOLDEN_PROMPT), None, 15, None, id="golden-length"), + pytest.param( + _source_outputs([1, 2], [3]), + {"additional_information": {"speaker": ["ethan"], "language": ["English"]}}, + None, + {"speaker": ["ethan"], "language": ["English"]}, + id="voice-metadata", + ), + ], +) +def test_build_forward_placeholder(monkeypatch, source, ctx_prompt, expected_len, expected_voice): + recorded = _capture_pack(monkeypatch) + out = q3.build_forward_placeholder(source, _ctx(prompt=ctx_prompt)) + assert len(out) == 1 + assert out[0] == "PLACEHOLDER" + if expected_len is not None: + assert recorded["prompt_len"] == expected_len + if expected_voice is not None: + assert recorded["voice_metadata"] == expected_voice + + +# --------------------------------------------------------------------------- +# Prewarm builder (async path) is a stage0-only best-effort estimate. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "stage0_prompt,expected_len", + [ + pytest.param(GOLDEN_PROMPT, 15, id="list-stage0"), + pytest.param(SimpleNamespace(prompt_token_ids=GOLDEN_PROMPT), 15, id="request-like-stage0"), + ], +) +def test_build_prewarm_placeholder(monkeypatch, stage0_prompt, expected_len): + recorded = _capture_pack(monkeypatch) + out = q3.build_prewarm_placeholder(stage0_prompt=stage0_prompt, ctx=_ctx(), downstream_stage_id=1) + assert len(out) == 1 + assert out[0] == "PLACEHOLDER" + # stage0_only runs the chat-template scan on the stage-0 list -> 15 + # (the same number the forward builder and the adapter scan produce). + assert recorded["prompt_len"] == expected_len + # Voice metadata is intentionally not forwarded (matches pre-existing prewarm). + assert recorded["voice_metadata"] is None + + +# --------------------------------------------------------------------------- +# Parity: forward and prewarm must agree where the prompt shape allows. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "prompt,expected_len", + [ + # Single user segment: forward (full) == prewarm (stage0_only) == 6. + pytest.param(SINGLE_USER_PROMPT, 6, id="single-user-segment"), + # Golden chat-template prompt: full == stage0_only scan == 15. The old + # ``len(stage0_prompt)`` == 10 estimate under-reserved KV slots and split + # a request between 10 and 15 depending on resolution success. + pytest.param(GOLDEN_PROMPT, 15, id="golden"), + ], +) +def test_forward_prewarm_parity(monkeypatch, prompt, expected_len): + recorded = _capture_pack(monkeypatch) + q3.build_forward_placeholder(_source_outputs(prompt), _ctx()) + forward_len = recorded["prompt_len"] + recorded.clear() + q3.build_prewarm_placeholder(stage0_prompt=prompt, ctx=_ctx(), downstream_stage_id=1) + assert forward_len == recorded["prompt_len"] == expected_len + + +# --------------------------------------------------------------------------- +# Packed placeholder structure (asserted with real vllm; field reads are +# skipped under the import fallback). +# --------------------------------------------------------------------------- + + +def test_pack_placeholder_prompt_structure(): + prompt = _common.pack_placeholder_prompt(prompt_len=4, voice_metadata={"speaker": ["ethan"]}) + empty_meta = _common.pack_placeholder_prompt(prompt_len=1) + if _shim_active(): + # Under the import fallback, OmniTokensPrompt's real constructor is + # bypassed, so field access returns fallback values. Field assertions + # run where real vllm constructs the model properly. + pytest.skip("OmniTokensPrompt field construction needs real vllm") + assert prompt["prompt_token_ids"] == [0, 0, 0, 0] + assert prompt["additional_information"] == {"speaker": ["ethan"]} + assert empty_meta["prompt_token_ids"] == [0] + assert empty_meta.get("additional_information") is None + + +# --------------------------------------------------------------------------- +# thinker2talker_token_only delegation keeps the public sync entry unchanged. +# --------------------------------------------------------------------------- + + +def test_thinker2talker_token_only_delegates_to_forward_placeholder(monkeypatch): + recorded = _capture_pack(monkeypatch) + prompt = {"additional_information": {"speaker": ["ethan"], "language": ["English"]}} + out = q3.thinker2talker_token_only( + _source_outputs([1, 2], [3]), + prompt=prompt, + requires_multimodal_data=False, + streaming_context=None, + ) + assert len(out) == 1 + assert out[0] == "PLACEHOLDER" + assert recorded["voice_metadata"] == {"speaker": ["ethan"], "language": ["English"]} + + +# --------------------------------------------------------------------------- +# Streaming fixup helper shares the same length helper as the forward builder. +# --------------------------------------------------------------------------- + + +def test_adapter_compute_length_consistent_with_common(): + # Lazy import keeps this module import-light without a vllm runtime. + from vllm_omni.distributed.omni_connectors import adapter + + assert adapter.compute_talker_prompt_ids_length(GOLDEN_PROMPT) == 15 + assert adapter.compute_talker_prompt_ids_length(GOLDEN_PROMPT) == _common.compute_placeholder_prompt_len( + ids_or_prompt={"ids": {"all": GOLDEN_PROMPT, "prompt": GOLDEN_PROMPT}}, + mode="full", + ) + # The prewarm stage0_only scan (what build_prewarm_placeholder and the + # orchestrator inline fallback both use) must equal the adapter scan. + assert adapter.compute_talker_prompt_ids_length(GOLDEN_PROMPT) == _common.compute_placeholder_prompt_len( + ids_or_prompt=GOLDEN_PROMPT, + mode="stage0_only", + ) + assert adapter.compute_talker_prompt_ids_length(SINGLE_USER_PROMPT) == 6 + assert ( + _common.compute_placeholder_prompt_len( + ids_or_prompt=SINGLE_USER_PROMPT, + mode="stage0_only", + ) + == 6 + ) + + +# --------------------------------------------------------------------------- +# Registry surface: the resolved sync fn exposes build_prewarm_placeholder, so +# the orchestrator's prewarm resolution path resolves without a module scan. +# --------------------------------------------------------------------------- + + +def test_resolved_sync_fn_exposes_build_prewarm_placeholder(): + spec = resolve_processor( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only", + expected_kind=None, + ) + assert callable(getattr(spec.fn, "build_prewarm_placeholder", None)) + assert callable(getattr(spec.fn, "build_forward_placeholder", None)) + # The attached builder is the same module-level object used by tests above. + assert spec.fn.build_prewarm_placeholder is q3.build_prewarm_placeholder diff --git a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py index 3ecae3d5688..3272a1a05fd 100644 --- a/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py +++ b/tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.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 """Light coverage for qwen2_5_omni.thinker2talker_full_payload. Covers the finish-reason-aware stop-row trim contract: when the request @@ -85,5 +85,65 @@ def test_missing_hidden_returns_none(): assert thinker2talker_full_payload(transfer_manager=None, pooling_output={}, request=request) is None +def _make_request_with_raw_token_ids(prompt, output, all_ids, status_name=None): + """Build a request without the ``_make_request`` ``list(...)`` coercion so a + ``None`` token-id field survives to the builder unchanged. + """ + status = SimpleNamespace(name=status_name) if status_name else None + return SimpleNamespace( + request_id="r1", + prompt_token_ids=prompt, + output_token_ids=output, + all_token_ids=all_ids, + status=status, + sampling_params=None, + ) + + +def test_none_prompt_and_output_token_ids_do_not_crash(): + """Regression (linyueqian #4290791768): ``None`` token-id fields must be + normalised to ``[]`` by ``ensure_list`` instead of reaching ``len(None)`` / + ``list(None)`` and raising ``TypeError``. + """ + request = _make_request_with_raw_token_ids(None, None, None) + hidden = torch.arange(2 * 4, dtype=torch.float32).reshape(2, 4) + + # No prompt and no output => target_rows == 0 => graceful None (no crash). + payload = thinker2talker_full_payload(transfer_manager=None, pooling_output={"hidden": hidden}, request=request) + assert payload is None + + +def test_none_output_token_ids_with_prompt_builds_payload(): + """``output_token_ids=None`` (a fresh, pre-decode request) must be treated + as an empty list: prompt-only payload is still shipped. + """ + request = _make_request_with_raw_token_ids((1, 2, 3), None, None) + # Tuple prompt is normalised to list; 3 hidden rows match the prompt span. + hidden = torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4) + + payload = thinker2talker_full_payload(transfer_manager=None, pooling_output={"hidden": hidden}, request=request) + + assert payload is not None + assert payload["ids"]["prompt"] == [1, 2, 3] + assert payload["ids"]["output"] == [] + assert payload["embed"]["prefill"].shape[0] == 3 + + +def test_none_all_token_ids_falls_back_to_prompt_plus_output(): + """``all_token_ids=None`` must fall back to ``prompt + output`` rather than + raising on ``list(None)``. + """ + request = _make_request_with_raw_token_ids([1, 2], [10, 11], None) + hidden = torch.arange(4 * 4, dtype=torch.float32).reshape(4, 4) + + payload = thinker2talker_full_payload(transfer_manager=None, pooling_output={"hidden": hidden}, request=request) + + assert payload is not None + assert payload["ids"]["prompt"] == [1, 2] + assert payload["ids"]["output"] == [10, 11] + assert payload["embed"]["prefill"].shape[0] == 2 + assert payload["hidden_states"]["output"].shape[0] == 2 + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 59c02cdec0a..345d95845b4 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -1697,22 +1697,22 @@ def append_and_emit(frame_idx, finished=False): tm.ramp_chunk_count[rid] += 1 return p - # Phase 1 (t=0.0): chunk 0, IC=2 + # t=0.0: chunk 0, IC=2. for i in range(2): append_and_emit(i) - # Phase 2 (t=0.05): build EWMA, emit target=4 at length=6 + # t=0.05: build EWMA, emit target=4 at length=6. clock[0] = 0.05 for i in range(2, 6): append_and_emit(i) - # Phase 3 (t=0.05): greedy target=18 (large buffer), accumulate - # 17 frames while holding (new_frames 1..17 < 18) + # t=0.05: greedy target=18 (large buffer), accumulate 17 frames while + # holding (new_frames 1..17 < 18). for i in range(6, 23): append_and_emit(i) - # Phase 4 (t=100.0): clock jumps, greedy collapses, target drops - # to ramp_floor=6. new_frames=18 > 6 -> bug trigger point. + # t=100.0: clock jumps, greedy collapses, target drops to ramp_floor=6. + # new_frames=18 > 6 -> bug trigger point. clock[0] = 100.0 for i in range(23, total_frames): append_and_emit(i, finished=(i == total_frames - 1)) diff --git a/tests/model_executor/stage_input_processors/test_registry.py b/tests/model_executor/stage_input_processors/test_registry.py new file mode 100644 index 00000000000..25376adbddd --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_registry.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for the stage-input processor registry (table-driven). + +Pure-logic, CPU-only: exercises kind inference, structural validation, +dead-processor hints and ``resolve_processor`` without executing any processor +logic or loading model weights. + +Subject under test: ``vllm_omni.model_executor.stage_input_processors._registry`` +""" + +import importlib +import warnings + +import pytest + +from vllm_omni.model_executor.stage_input_processors import ( + ProcessorKind, + ProcessorSpec, + ProcessorValidationError, + _registry, + dead_processor_hint, + infer_kind, + register_processor, + resolve_processor, + validate_processor, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# --------------------------------------------------------------------------- +# Dummy processors (simple callables used as validation subjects). +# --------------------------------------------------------------------------- + + +def placeholder_a_token_only(source_outputs, _prompt=None, _requires_multimodal_data=False): + return [] + + +def placeholder_b_token_only(prompt, ctx=None): + return [] + + +def diffusion_a_ar2diffusion(source_outputs, _prompt=None, _requires_multimodal_data=False): + return {} + + +def producer_a_full_payload(transfer_manager, pooling_output, request): + return None + + +def producer_b_full_payload(transfer_manager, pooling_output, request, is_finished=False): + return None + + +def producer_c_full_payload(transfer_manager, multimodal_output, request, is_finished=False): + return None + + +def producer_d_full_payload(transfer_manager, pooling_output, request, *, is_finished): + return None + + +def producer_a_async_chunk(transfer_manager, multimodal_output, request, is_finished=False): + return None + + +def producer_b_async_chunk(transfer_manager, multimodal_output, request): + return None + + +def producer_c_async_chunk(transfer_manager, multimodal_output, request, is_finished, /): + return None + + +def producer_kwargs_full_payload(**kwargs): + return None + + +def producer_kwargs_async_chunk(**kwargs): + return None + + +def producer_badname_full_payload(transfer_manager, payload, request, is_finished=False): + return None + + +def producer_badname_async_chunk(transfer_manager, payload, request, is_finished=False): + return None + + +def producer_positional_full_payload(transfer_manager, pooling_output, request, /): + return None + + +def talker2codec_shape(stage_list, engine_input_source, prompt=None, requires_multimodal_data=False): + return [] + + +def _mod_path(fn_name: str) -> str: + """Dotted path to a callable defined in this test module.""" + return f"{__name__}.{fn_name}" + + +# --------------------------------------------------------------------------- +# Kind inference (suffix rules). +# --------------------------------------------------------------------------- + + +class TestInferKind: + @pytest.mark.parametrize( + "path,expected", + [ + ("pkg.mod.foo_token_only", "placeholder_prompt_builder"), + ("pkg.mod.foo_full_payload", "producer_full_payload"), + ("pkg.mod.foo_batch", "producer_full_payload"), + ("pkg.mod.foo_async_chunk", "producer_async_chunk"), + ("pkg.mod.ar2diffusion", "diffusion_input_builder"), + ("pkg.mod.ar2dit", "diffusion_input_builder"), + ("pkg.mod.thinker2imagegen", "diffusion_input_builder"), + ("pkg.mod.legacy_builder", "legacy_orchestrator_builder"), + ("pkg.mod.talker2code2wav", "legacy_orchestrator_builder"), + ("vllm_omni.model_executor.stage_input_processors.moss_tts.talker2codec", "legacy_multi_source"), + ("pkg.moss_tts.talker2codec", "legacy_multi_source"), + ( + "vllm_omni.model_executor.stage_input_processors.moss_tts.talker2codec_delay_async_chunk", + "producer_async_chunk", + ), + ], + ) + def test_suffix_rules(self, path: str, expected: ProcessorKind): + assert infer_kind(None, path=path) == expected + + def test_invalid_path_raises(self): + with pytest.raises(ValueError): + infer_kind(None, path="") + with pytest.raises(ValueError): + infer_kind(None, path="nodots") + + def test_register_processor_override(self): + register_processor("pkg.mod.custom_path", "diffusion_input_builder") + try: + assert infer_kind(None, path="pkg.mod.custom_path") == "diffusion_input_builder" + finally: + _registry._KIND_OVERRIDES.clear() + + def test_register_processor_invalid_kind(self): + with pytest.raises(ValueError): + register_processor("pkg.mod.custom_path", "bogus_kind") # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# Structural validation (table-driven: ok / warn / reject). +# --------------------------------------------------------------------------- + + +def _ok(fn, kind, path): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_processor(fn, kind=kind, path=path) + + +def _warns(fn, kind, path, match): + with pytest.warns(RuntimeWarning, match=match): + validate_processor(fn, kind=kind, path=path) + + +def _rejects(fn, kind, path, rule, match=None): + with pytest.raises(ProcessorValidationError) as excinfo: + validate_processor(fn, kind=kind, path=path) + assert excinfo.value.rule == rule + if match is not None: + assert match in str(excinfo.value) + + +_VALIDATION_CASES = [ + # (fn, kind, path, expected-behaviour, extra) + # full-payload producers + pytest.param( + producer_b_full_payload, + "producer_full_payload", + "pkg.mod.producer_b_full_payload", + "ok", + None, + id="full-payload-ok", + ), + pytest.param( + producer_a_full_payload, + "producer_full_payload", + "pkg.mod.producer_a_full_payload", + "warn", + "is_finished", + id="full-payload-missing-is_finished-warns", + ), + pytest.param( + producer_c_full_payload, + "producer_full_payload", + "pkg.mod.producer_c_full_payload", + "reject", + ("producer_kwargs", None), + id="full-payload-cross-kind-second-param-rejected", + ), + pytest.param( + producer_badname_full_payload, + "producer_full_payload", + "pkg.mod.producer_badname_full_payload", + "reject", + ("producer_kwargs", "pooling_output"), + id="full-payload-bad-param-names-rejected", + ), + pytest.param( + producer_positional_full_payload, + "producer_full_payload", + "pkg.mod.producer_positional_full_payload", + "reject", + ("producer_kwargs", None), + id="full-payload-positional-only-rejected", + ), + pytest.param( + producer_kwargs_full_payload, + "producer_full_payload", + "pkg.mod.producer_kwargs_full_payload", + "ok", + None, + id="full-payload-var-keyword-accepted", + ), + pytest.param( + producer_d_full_payload, + "producer_full_payload", + "pkg.mod.producer_d_full_payload", + "ok", + None, + id="full-payload-required-kw-only-is_finished-accepted", + ), + pytest.param( + producer_b_full_payload, + "producer_full_payload", + "pkg.mod.producer_b_full_payload", + "ok", + None, + id="full-payload-default-is_finished-accepted", + ), + # async-chunk producers + pytest.param( + producer_badname_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_badname_async_chunk", + "reject", + ("producer_kwargs", "multimodal_output"), + id="async-chunk-bad-param-names-rejected", + ), + pytest.param( + producer_kwargs_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_kwargs_async_chunk", + "ok", + None, + id="async-chunk-var-keyword-accepted", + ), + pytest.param( + producer_a_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_a_async_chunk", + "ok", + None, + id="async-chunk-ok", + ), + pytest.param( + producer_b_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_b_async_chunk", + "reject", + ("is_finished_required", "producer_b_async_chunk"), + id="async-chunk-missing-is_finished-hard-fail", + ), + pytest.param( + producer_c_async_chunk, + "producer_async_chunk", + "pkg.mod.producer_c_async_chunk", + "reject", + ("is_finished_positional_only", None), + id="async-chunk-positional-only-is_finished-hard-fail", + ), + # orchestrator-side + pytest.param( + placeholder_a_token_only, + "placeholder_prompt_builder", + "pkg.mod.placeholder_a_token_only", + "ok", + None, + id="placeholder-ok", + ), + pytest.param( + placeholder_b_token_only, + "placeholder_prompt_builder", + "pkg.mod.placeholder_b_token_only", + "ok", + None, + id="placeholder-ctx-contract-ok", + ), + pytest.param( + diffusion_a_ar2diffusion, + "diffusion_input_builder", + "pkg.mod.diffusion_a_ar2diffusion", + "ok", + None, + id="diffusion-ok", + ), + pytest.param( + talker2codec_shape, + "legacy_multi_source", + "vllm_omni.model_executor.stage_input_processors.moss_tts.talker2codec", + "ok", + None, + id="legacy-multi-source-allowlist-ok", + ), + pytest.param( + talker2codec_shape, + "legacy_multi_source", + "pkg.mod.talker2codec", + "reject", + ("legacy_multi_source_allowlist", None), + id="legacy-multi-source-rejects-non-moss", + ), + pytest.param( + object(), + "producer_full_payload", + "pkg.mod.producer_full_payload", + "reject", + ("inspectable_signature", None), + id="uninspectable-callable-fails", + ), +] + + +@pytest.mark.parametrize("fn,kind,path,behaviour,extra", _VALIDATION_CASES) +def test_validate_processor_table(fn, kind, path, behaviour, extra): + if behaviour == "ok": + _ok(fn, kind, path) + elif behaviour == "warn": + _warns(fn, kind, path, extra) + else: + rule, match = extra + _rejects(fn, kind, path, rule, match) + + +# --------------------------------------------------------------------------- +# Dead-processor hint (gate three-state). +# --------------------------------------------------------------------------- + + +# One entry per legacy test method; each entry carries the scenarios the +# original method asserted (keeps the collected-test count unchanged). +_DEAD_CASES = [ + pytest.param( + [ + ("placeholder_prompt_builder", True, True, False, False, True), + ], + id="async-chunk-receiving-downstream-dead", + ), + pytest.param( + [ + ("diffusion_input_builder", True, False, False, False, False), + ], + id="async-chunk-non-receiving-downstream-alive", + ), + pytest.param( + [ + ("diffusion_input_builder", False, False, False, False, False), + ("diffusion_input_builder", False, False, True, False, True), + ], + id="sync-mode-alive-unless-overridden", + ), + pytest.param( + [ + ("placeholder_prompt_builder", False, False, True, True, False), + ("placeholder_prompt_builder", False, False, True, False, True), + ("placeholder_prompt_builder", True, True, True, True, True), + ], + id="sync-mode-selected-sync-hook-is-not-dead", + ), + pytest.param( + [ + ("producer_full_payload", True, True, False, False, False), + ("producer_async_chunk", True, True, False, False, False), + ], + id="producer-never-dead", + ), +] + + +@pytest.mark.parametrize("scenarios", _DEAD_CASES) +def test_dead_processor_hint_table(scenarios): + for kind, async_chunk, downstream, has_sync, custom_is_sync, expected in scenarios: + assert ( + dead_processor_hint( + kind, + async_chunk=async_chunk, + downstream_receives_async_chunks=downstream, + has_sync=has_sync, + custom_is_sync=custom_is_sync, + ) + is expected + ) + + +# --------------------------------------------------------------------------- +# Resolution entry point. +# --------------------------------------------------------------------------- + + +class TestResolveProcessor: + @pytest.mark.parametrize( + "fn_name,expected_kind", + [ + ("placeholder_a_token_only", "placeholder_prompt_builder"), + ("producer_b_full_payload", None), + ("producer_d_full_payload", "producer_full_payload"), + ], + ) + def test_resolve_normal(self, fn_name, expected_kind): + spec = resolve_processor(_mod_path(fn_name), expected_kind=expected_kind) + assert isinstance(spec, ProcessorSpec) + assert spec.kind == (expected_kind or infer_kind(None, path=_mod_path(fn_name))) + assert spec.path == _mod_path(fn_name) + assert spec.fn is getattr(importlib.import_module(__name__), fn_name) + + def test_resolve_drop_in_equivalence(self): + """resolve_processor().fn must be the same object as the legacy lookup.""" + spec = resolve_processor( + _mod_path("producer_b_full_payload"), + expected_kind="producer_full_payload", + ) + legacy = getattr(importlib.import_module(__name__), "producer_b_full_payload") + assert spec.fn is legacy + + def test_expected_kind_mismatch_raises(self): + with pytest.raises(ProcessorValidationError) as excinfo: + resolve_processor( + _mod_path("placeholder_a_token_only"), + expected_kind="producer_full_payload", + ) + assert excinfo.value.rule == "expected_kind" + assert "placeholder_prompt_builder" in str(excinfo.value) + + def test_missing_symbol_raises(self): + with pytest.raises(AttributeError): + resolve_processor(_mod_path("does_not_exist")) + + def test_package_exports(self): + import vllm_omni.model_executor.stage_input_processors as sip + + assert sip.resolve_processor is resolve_processor + assert sip.ProcessorSpec is ProcessorSpec + assert sip.ProcessorValidationError is ProcessorValidationError + assert sip.infer_kind is infer_kind + assert sip.dead_processor_hint is dead_processor_hint + assert sip.register_processor is register_processor diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py index dc248f21179..99df9dcf27b 100644 --- a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.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 from collections import defaultdict from types import SimpleNamespace @@ -213,3 +213,96 @@ def test_none_multimodal_output_not_finished_returns_none(): ) assert payload is None + + +def _compute_mm_logits_payload(batch: int = 1, codebooks: int = 8): + """Mirror voxtral_tts_audio_generation.compute_mm_logits() L1075: + audio_codes [B, C] -> unsqueeze(1) -> [B, 1, C] -> split -> list of (1, 1, C). + """ + audio_codes = torch.arange(batch * codebooks, dtype=torch.float32).reshape(batch, codebooks) + audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) + return {"codes": {"audio": audio_list}} + + +def test_3d_per_step_frame_list_accepted(): + """Voxtral P1 regression: compute_mm_logits 3-D per-step [B, 1, C] + payload must not raise ValueError; the last frame is extracted as 1-D.""" + from vllm_omni.model_executor.stage_input_processors import _common + + transfer_manager = _make_transfer_manager() + request_id = "rid-3d" + request = _req(request_id, finished=False) + + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + multimodal_output=_compute_mm_logits_payload(1, 8), + request=request, + ) + + # The frame is stored as the flat 8-element codec row. + assert transfer_manager.code_prompt_token_ids[request_id][-1] == [float(i) for i in range(8)] + # One frame < chunk_size_at_begin=5: no emission yet. + assert payload is None + # Direct helper call: 1-D float frame, precision preserved (to_long=False). + frame = _common.extract_last_codec_frame( + {"codes": {"audio": torch.zeros(1, 1, 8)}}, + to_long=False, + ) + assert frame is not None + assert frame.shape == (8,) + assert frame.dtype == torch.float32 + + +def test_3d_direct_tensor_accepted(): + """A bare 3-D [1, 1, C] tensor is normalized via squeeze(1) and the + 2-D last-row semantics apply unchanged.""" + from vllm_omni.model_executor.stage_input_processors import _common + + frame = _common.extract_last_codec_frame( + {"codes": {"audio": torch.arange(8, dtype=torch.float32).reshape(1, 1, 8)}}, + to_long=False, + ) + assert frame is not None + assert frame.shape == (8,) + assert frame.dtype == torch.float32 + assert frame.tolist() == [float(i) for i in range(8)] + + +def test_3d_per_step_chunk_emission(): + """Feeding 5 per-step 3-D frames reaches the codec_chunk_frames_at_begin + boundary and emits the standard [ctx_frames, context_length, ...flat] payload.""" + transfer_manager = _make_transfer_manager( + codec_chunk_frames=25, + codec_left_context_frames=25, + codec_chunk_frames_at_begin=5, + ) + request_id = "rid-3d-chunk" + request = _req(request_id, finished=False) + + payload = None + for _ in range(5): + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + multimodal_output=_compute_mm_logits_payload(1, 8), + request=request, + ) + + assert payload is not None + codes = payload.codes.audio + assert int(codes[0].item()) == 0 # ctx_frames + assert int(codes[1].item()) == 5 # context_length == chunk_size_at_begin + assert codes[2:].numel() == 5 * 8 # 5 frames x 8 codebooks + + +def test_3d_non_singleton_frame_dim_raises(): + """Defensive contract: non per-step 3-D layouts (dim1 != 1) still raise + ValueError instead of being silently flattened.""" + transfer_manager = _make_transfer_manager() + request = _req("rid-3d-non-singleton", finished=False) + + with pytest.raises(ValueError): + generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + multimodal_output={"codes": {"audio": torch.zeros(1, 2, 3)}}, + request=request, + ) diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index ea042d4e846..3deac22a5ad 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -545,6 +545,21 @@ def test_finished_load_reqs_flow_to_chunk_ready(self): host.shutdown_omni_connectors() +def pack(transfer_manager, multimodal_output, request, is_finished=False): + """Arbitrarily named, structurally compatible connector payload builder.""" + del transfer_manager, multimodal_output, request, is_finished + return None + + +def pack_full_payload(transfer_manager, multimodal_output, request, is_finished=False): + """Name suffix says ``producer_full_payload`` while the signature follows the + async-chunk contract (``multimodal_output``) -- the exact shape a + name-driven kind inference misclassifies before the structural check. + """ + del transfer_manager, multimodal_output, request, is_finished + return None + + class TestLoadCustomFuncSelection(unittest.TestCase): def test_uses_validator_override_from_public_mixin(self): config = SimpleNamespace( @@ -559,6 +574,43 @@ def test_uses_validator_override_from_public_mixin(self): assert selected_path is None assert func is None + def test_warns_when_explicit_hook_is_skipped(self): + """linyueqian #4290791761: an explicitly configured hook that is dropped + (here: rejected before the structural check) is escalated to a + ``warning`` naming the ``register_processor`` escape hatch, instead of + vanishing into a ``debug``-only line. + """ + config = SimpleNamespace( + async_chunk=True, + custom_process_next_stage_input_func=f"{__name__}.pack_full_payload", + ) + + with patch.object(omni_connector_runtime.logger, "warning") as warn: + selected_path, func = MixinHost._load_custom_func(config) + + assert selected_path is None + assert func is None + warn.assert_called_once() + message = warn.call_args.args[0] + assert "explicitly configured connector payload hook" in message + assert "register_processor" in message + + def test_auto_derived_probe_candidates_stay_at_debug(self): + """Auto-derived ``*_full_payload`` / ``*_batch`` probes are the expected + misses of the scan and must not spam ``warning`` when the sibling + builder does not exist. + """ + config = SimpleNamespace( + async_chunk=False, + custom_process_next_stage_input_func=None, + custom_process_input_func="vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav", + ) + + with patch.object(omni_connector_runtime.logger, "warning") as warn: + MixinHost._load_custom_func(config) + + warn.assert_not_called() + 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", @@ -578,6 +630,50 @@ def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self): assert selected_path != func_path assert func is None or MixinHost._is_connector_payload_builder(func) + def test_accepts_full_payload_pooling_output_hook(self): + """A ``pooling_output`` full-payload hook loads through the payload contract.""" + selected_path, func = MixinHost._load_custom_func( + SimpleNamespace( + async_chunk=False, + custom_process_input_func=None, + custom_process_next_stage_input_func=( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_full_payload" + ), + ) + ) + assert selected_path is not None + assert MixinHost._is_connector_payload_builder(func) + assert MixinHost._connector_payload_kwarg(func) == "pooling_output" + + def test_accepts_async_chunk_multimodal_output_hook(self): + """An async-chunk ``multimodal_output`` hook loads through the payload contract.""" + selected_path, func = MixinHost._load_custom_func( + SimpleNamespace( + async_chunk=True, + custom_process_input_func=None, + custom_process_next_stage_input_func=( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk" + ), + ) + ) + assert selected_path is not None + assert MixinHost._is_connector_payload_builder(func) + assert MixinHost._connector_payload_kwarg(func) == "multimodal_output" + + def test_accepts_explicitly_configured_hook_with_arbitrary_name(self): + """An explicitly configured, structurally compatible hook is not rejected by name.""" + module_path = f"{__name__}" + selected_path, func = MixinHost._load_custom_func( + SimpleNamespace( + async_chunk=True, + custom_process_input_func=None, + custom_process_next_stage_input_func=f"{module_path}.pack", + ) + ) + assert selected_path == f"{module_path}.pack" + assert func is pack + assert MixinHost._is_connector_payload_builder(func) + class TestFullPayloadSendWithCustomFunc(unittest.TestCase): """Test B4: send_full_payload_outputs with full_payload_mode custom process func.""" @@ -646,6 +742,46 @@ def full_payload_func(transfer_manager, pooling_output, request): time.sleep(0.1) host.shutdown_omni_connectors() + def test_full_payload_required_is_finished_kwarg_receives_finished_flag(self): + """Keyword-only REQUIRED ``is_finished`` must bind the worker's best-effort kwarg + (amy-why-3459 review of HEAD ad81ac3a).""" + seen = {} + + def full_payload_func(transfer_manager, pooling_output, request, *, is_finished): + seen["connector"] = transfer_manager.connector + seen["is_finished"] = is_finished + seen["data"] = pooling_output + seen["rid"] = request.request_id if request else None + return {"processed": True, "finished": is_finished} + + host = MixinHost() + host.init_omni_connectors( + model_config=_make_model_config(), + ) + host._omni_connector = MockConnector(stage_id=0) + host._stage_id = 0 + host._custom_process_func = full_payload_func + + req = _make_request("req-1") + req.is_finished = lambda: True + payload = host._build_custom_process_payload( + request_id="req-1", + request=req, + pooling_output={"raw": 100}, + ) + self.assertIsNotNone(payload) + self.assertEqual( + seen, + { + "connector": host._omni_connector, + "is_finished": True, + "data": {"raw": 100}, + "rid": "req-1", + }, + ) + + host.shutdown_omni_connectors() + class TestKVSentReqIdsAccumulation(unittest.TestCase): """Test that kv_sent_req_ids accumulates results from send_kv_cache.""" diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index d8a26fd1161..7660a93ca5d 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1535,6 +1535,7 @@ def merge_pipeline_deploy( stage_type=stage_type, input_sources=list(ps.input_sources), custom_process_input_func=input_proc, + sync_process_input_func=ps.sync_process_input_func, final_output=ps.final_output, final_output_type=ps.final_output_type, worker_type=worker_type, @@ -1564,6 +1565,7 @@ class StageConfig: stage_type: StageType = StageType.LLM input_sources: list[int] = field(default_factory=list) custom_process_input_func: str | None = None + sync_process_input_func: str | None = None final_output: bool = False final_output_type: str | None = None worker_type: str | None = None @@ -1659,6 +1661,11 @@ def to_omegaconf(self) -> Any: if self.custom_process_input_func: config_dict["custom_process_input_func"] = self.custom_process_input_func + # Carry the raw ``sync_process_input_func`` path to the orchestrator so + # async-chunk prewarm can resolve the registered placeholder builder + # (``*_token_only``) without re-reading the YAML. + if self.sync_process_input_func: + config_dict["sync_process_input_func"] = self.sync_process_input_func # Pass through extra YAML fields (default_sampling_params, # output_connectors, input_connectors, tts_args, etc.) diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index fc8ba08097b..50b2b19468a 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -12,6 +12,7 @@ from vllm_omni.diffusion.registry import DiffusionModelRegistry from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniPromptType +from vllm_omni.model_executor.stage_input_processors import _common from vllm_omni.outputs import OmniRequestOutput from vllm_omni.outputs.output_metadata import ( DiffusionMetadata, @@ -214,9 +215,17 @@ def _stream_metadata_from_diffusion_output(diffusion_output: DiffusionOutput) -> def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]: - if isinstance(outputs, list): - return outputs - return [outputs] if outputs is not None else [] + """Wrap a non-list primary payload as ``[x]`` (diffusion formatter semantics). + + ``output_formatter`` is **not** a stage-input helper: the legacy behaviour + was wrap-only — a ``list`` passes through unchanged, a non-list scalar / + tensor / PIL ``Image`` is wrapped verbatim as ``[x]``, and ``None`` becomes + ``[]``. The canonical ``_common.ensure_list`` would flatten tensors, + iterate dict keys, and walk a PIL ``Image`` row by row, so it must not be + reused here; ``_common.ensure_list_wrap_only`` reproduces the exact legacy + semantics. + """ + return _common.ensure_list_wrap_only(outputs) def _primary_payload(postprocess_output: DiffusionPostprocessOutput) -> DiffusionPayloadValue | None: diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index dc3be166fc9..3af7cd2c161 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -193,34 +193,27 @@ def try_recv_via_connector( def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int: """Compute the length of the talker prompt ids. + Delegates to the canonical ``_common.compute_placeholder_prompt_len(mode="full")`` + so the streaming fixup path (``construct_next_stage_streaming_input_prompt``) + and the forward placeholder builder (``build_forward_placeholder``) share one + length helper. Behaviour is golden-locked (``test_common_helpers_golden.py`` + reproduces the 15-token chat-template result for the canonical prompt). + Args: prompt_ids: The prompt ids tensor. Returns: The length of the talker prompt ids. """ - im_start_token_id = 151644 - system_token_id = 8948 - user_token_id = 872 - assistant_token_id = 77091 - im_start_indexes = [i for i in range(len(prompt_ids)) if prompt_ids[i] == im_start_token_id] - im_start_indexes.append(len(prompt_ids)) - sum_user_len = 0 - assistant_len = 0 - for i in range(len(im_start_indexes) - 1): - s = im_start_indexes[i] - e = im_start_indexes[i + 1] - role = prompt_ids[s + 1] - if role == system_token_id: - continue - elif role == user_token_id: - sum_user_len += e - s - elif role == assistant_token_id and i == len(im_start_indexes) - 2: - assistant_len += 9 # 3 + 4 + 1 + 1 - else: - pass - - return sum_user_len + assistant_len + from vllm_omni.model_executor.stage_input_processors import _common + + # A flat prompt list has no ``all``/``prompt`` split, so feed the same list + # to both roles; ``_common``'s terminal index then equals ``len(prompt_ids)`` + # exactly as the legacy inline scan did. + return _common.compute_placeholder_prompt_len( + ids_or_prompt={"ids": {"all": prompt_ids, "prompt": prompt_ids}}, + mode="full", + ) def construct_next_stage_streaming_input_prompt( diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 6667c0e4b1e..ad7a325a61c 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -22,6 +22,10 @@ ) from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec from vllm_omni.distributed.omni_connectors.utils.kv_utils import get_local_tp_rank, get_omni_replica_id +from vllm_omni.model_executor.stage_input_processors import ( + ProcessorValidationError, + resolve_processor, +) from vllm_omni.outputs import OmniConnectorOutput logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin") @@ -609,12 +613,34 @@ def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: ``custom_process_input_func`` (for example ``thinker2talker``), while the connector payload builder lives beside it as ``thinker2talker_full_payload``. In that case, derive the full_payload_mode builder path automatically. + + The loader accepts both supported payload contracts: full-payload + builders name their runner payload ``pooling_output`` while async-chunk + builders name it ``multimodal_output`` (see + :meth:`_connector_payload_kwarg`). For the configured hook the mixin's + structural validator (:meth:`_is_connector_payload_builder`) is the + authoritative acceptance gate. + + Caveat: :func:`resolve_processor` runs the registry's *name-driven* + kind inference before that structural check, so a hook whose name ends + in ``_full_payload`` (for example ``pack_full_payload``) is classified + as a ``producer_full_payload`` builder and validated against *that* + contract. If its signature actually follows a different payload + contract, it is rejected before :meth:`_is_connector_payload_builder` + ever sees it. Such a hook is still reported at ``warning`` level (see + :meth:`_warn_skipped_hook`) rather than silently dropped, and can be + pinned to the intended kind via + ``stage_input_processors.register_processor(path, "producer_async_chunk")``. """ - candidates: list[str] = [] + # Each candidate carries an ``is_explicit`` flag: hooks named in config + # are escalated to ``warning`` when skipped (dropping one disables the + # stage's connector payload transfer), while auto-derived probe + # candidates stay at ``debug`` to avoid noise from the normal scan. + candidates: list[tuple[str, bool]] = [] next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) if isinstance(next_stage_func, str) and next_stage_func: - candidates.append(next_stage_func) + candidates.append((next_stage_func, True)) if not getattr(model_config, "async_chunk", False): input_func = getattr(model_config, "custom_process_input_func", None) @@ -622,37 +648,67 @@ def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: try: module_path, func_name = input_func.rsplit(".", 1) if func_name.endswith("_full_payload") or func_name.endswith("_batch"): - candidates.append(f"{module_path}.{func_name}") + candidates.append((f"{module_path}.{func_name}", False)) else: - candidates.append(f"{module_path}.{func_name}_full_payload") - candidates.append(f"{module_path}.{func_name}_batch") - candidates.append(input_func) + candidates.append((f"{module_path}.{func_name}_full_payload", False)) + candidates.append((f"{module_path}.{func_name}_batch", False)) + candidates.append((input_func, False)) except ValueError: - candidates.append(input_func) + candidates.append((input_func, False)) tried: set[str] = set() - for func_path in candidates: + for func_path, is_explicit in candidates: if func_path in tried: continue tried.add(func_path) try: - module_path, func_name = func_path.rsplit(".", 1) - module = importlib.import_module(module_path) - func = getattr(module, func_name, None) - if callable(func): - if not cls._is_connector_payload_builder(func): - logger.debug( - "Skipping incompatible connector payload hook %s; signature=%s", - func_path, - inspect.signature(func), - ) - continue - return func_path, func + # Resolve + run the registry's structural validation. No + # expected_kind is passed: a caller-configured hook may carry + # any name, and the payload-contract check below is the + # authoritative acceptance gate. + spec = resolve_processor(func_path) + if not cls._is_connector_payload_builder(spec.fn): + try: + detail: Any = inspect.signature(spec.fn) + except (TypeError, ValueError): + detail = "signature unavailable" + cls._warn_skipped_hook(func_path, is_explicit, detail) + continue + return spec.path, spec.fn + except ProcessorValidationError as exc: + # A name-driven kind inference can pick the wrong contract + # (e.g. ``pack_full_payload`` -> ``producer_full_payload``) + # before the structural payload check ever runs. + cls._warn_skipped_hook(func_path, is_explicit, exc) + continue except Exception: logger.warning("Failed to load custom func: %s", func_path, exc_info=True) return None, None + @classmethod + def _warn_skipped_hook(cls, func_path: str, is_explicit: bool, detail: Any) -> None: + """Report a connector payload hook that was skipped during the scan. + + Auto-derived probe candidates (``*_full_payload`` / ``*_batch``) are + logged at ``debug`` because a missing sibling is the expected case. + An *explicitly* configured hook is escalated to ``warning`` since + silently dropping it disables connector payload transfer for the + stage; the message names the ``register_processor`` escape hatch so a + deliberately named hook can be pinned to the correct kind. + """ + if is_explicit: + logger.warning( + "Skipping explicitly configured connector payload hook %s: %s. " + "If this hook is structurally valid, pin its kind via " + "stage_input_processors.register_processor(%r, 'producer_async_chunk').", + func_path, + detail, + func_path, + ) + else: + logger.debug("Skipping incompatible connector payload hook %s: %s", func_path, detail) + @classmethod def _is_connector_payload_builder(cls, func: Any) -> bool: """Whether *func* matches the mixin payload-builder contract.""" diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 300e4fc5f85..c6c79311010 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import copy -import importlib import inspect import threading import time @@ -16,6 +15,11 @@ from vllm.v1.utils import ConstantList from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload +from vllm_omni.model_executor.stage_input_processors import ( + ProcessorValidationError, + resolve_processor, +) +from vllm_omni.model_executor.stage_input_processors._registry import _check_async_chunk, import_symbol from ..adapter import construct_next_stage_streaming_input_prompt from ..connectors.shm_connector import SharedMemoryConnector @@ -31,6 +35,24 @@ _RECLAIM_WARN_INTERVAL = 1000 +def _accepts_async_chunk_contract(fn: Callable[..., Any]) -> bool: + """Whether *fn* structurally matches the async-chunk producer contract.""" + try: + signature = inspect.signature(fn) + except (TypeError, ValueError): + return False + try: + _check_async_chunk( + signature, + path=getattr(fn, "__qualname__", repr(fn)), + kind="producer_async_chunk", + stage_config=None, + ) + except ProcessorValidationError: + return False + return True + + class _SenderGeneration: """Fence one external request generation without blocking cleanup.""" @@ -175,6 +197,10 @@ def __init__(self, vllm_config: Any): ) self.connector = self.create_connector(model_config) self.receives_chunks = stage_receives_chunks(model_config) + # Must exist before super().__init__() starts the save thread, because + # _accepts_new_token_ids() is called from _send_single_request() on the + # save-loop thread before __init__() would otherwise complete. + self._processor_accepts_step_tokens: dict[Callable[..., Any], bool] = {} super().__init__(model_config) self.model_mode = getattr(model_config, "worker_type", None) or "ar" # State specific to Chunk management @@ -187,13 +213,34 @@ def __init__(self, vllm_config: Any): # the request snapshot (other ``_batch`` forms serve other callers). self.custom_process_next_stage_input_batch_func: Callable[..., list[OmniPayloadStruct | None]] | None = None if custom_process_next_stage_input_func: - module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1) - module = importlib.import_module(module_path) - self.custom_process_next_stage_input_func = getattr(module, func_name) - batch_func = getattr(module, f"{func_name}_batch", None) - if getattr(batch_func, "per_step_batch", False) is True: - self.custom_process_next_stage_input_batch_func = batch_func - self._processor_accepts_step_tokens: dict[Callable[..., Any], bool] = {} + # Kind inference is name-driven; accept any matching signature, reject only incompatible hooks. + spec = resolve_processor(custom_process_next_stage_input_func) + if spec.kind == "producer_async_chunk" or _accepts_async_chunk_contract(spec.fn): + self.custom_process_next_stage_input_func = spec.fn + else: + raise ProcessorValidationError( + path=custom_process_next_stage_input_func, + kind="producer_async_chunk", + rule="async_chunk_contract", + message=( + "custom_process_next_stage_input_func does not accept the async-chunk " + "producer call (requires keyword-bindable 'transfer_manager', " + "'multimodal_output', 'request' and a 'is_finished' parameter)" + ), + ) + # Absorb c4871a09b: resolve the same-named ``_batch`` and opt in only + # when it declares per-step batch semantics. The batch form is optional, so a + # missing attribute must not fail construction. Raw ``import_symbol`` (import + + # ``getattr``) is used instead of ``resolve_processor`` because registry kind + # inference classifies any ``*_batch`` name as a full-payload producer and would + # reject the async-chunk batch signature below. + batch_path = f"{custom_process_next_stage_input_func}_batch" + try: + batch_fn = import_symbol(batch_path) + except Exception: + batch_fn = None + if batch_fn is not None and getattr(batch_fn, "per_step_batch", False) is True: + self.custom_process_next_stage_input_batch_func = batch_fn # mapping for request id and chunk id self.put_req_chunk: dict[str, int] = defaultdict(int) self.get_req_chunk: dict[str, int] = defaultdict(int) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 2dc8128f268..41ca7b7265f 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -140,6 +140,32 @@ def _build_terminal_empty_output( ) +def _custom_hook_is_sync_hook(fn: Any, sync_path: str | None) -> bool: + """Whether the custom hook callable *is* the configured sync hook. + + A stage may wire the same ``*_token_only`` processor as both + ``custom_process_input_func`` and ``sync_process_input_func`` (the sync hook + is pre-selected as the active forward processor). In that case the custom + hook is **not** dead — it is the active processor — so + :func:`dead_processor_hint` must not report it. Comparison is by dotted + path first, then by resolving the sync path and identity-checking the + callable. Never raises. + """ + if not callable(fn) or not sync_path: + return False + module = getattr(fn, "__module__", None) + name = getattr(fn, "__qualname__", None) or getattr(fn, "__name__", None) + if module and name and f"{module}.{name}" == sync_path: + return True + try: + from vllm_omni.model_executor.stage_input_processors import resolve_processor + + spec = resolve_processor(sync_path, expected_kind=None) + return spec.fn is fn + except Exception: + return False + + def build_engine_core_request_from_tokens( request_id: str, prompt: dict[str, Any], @@ -371,6 +397,9 @@ def __init__( self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools + # Wire the dead-processor hint at startup (warn-only) so a + # configured-but-never-invoked custom_process_input_func is surfaced. + self._warn_dead_input_processors() self.log_stats = log_stats if log_stats and not collect_stage_metrics: raise ValueError("log_stats=True requires collect_stage_metrics=True") @@ -2171,6 +2200,59 @@ def _stage_receives_async_chunks(self, stage_id: int) -> bool: model_config = getattr(pool.stage_vllm_config, "model_config", None) return stage_receives_chunks(model_config) + def _warn_dead_input_processors(self) -> None: + """Startup (M0, warn-only) wiring of :func:`dead_processor_hint`. + + For every stage whose ``custom_process_input_func`` would never be + invoked under the current async-chunk / sync wiring (mirroring the + ``_route_output`` gate ``(not async_chunk or not + _stage_receives_async_chunks(stage_id))``), log a warning so a + configured-but-dead hook is not silently dropped. This is a pure hint + (warn-first): the orchestrator never enforces it at runtime. + Defensive by design — fake pools / minimal clients in tests or future + connectors may not expose ``stage_client``, so the loop never raises. + """ + from vllm_omni.model_executor.stage_input_processors import dead_processor_hint, infer_kind + + for stage_id, pool in enumerate(self.stage_pools): + try: + client = getattr(pool, "stage_client", None) + fn = getattr(client, "custom_process_input_func", None) + if fn is None: + continue + module = getattr(fn, "__module__", None) + name = getattr(fn, "__qualname__", None) or getattr(fn, "__name__", None) + if not module or not name: + continue + path = f"{module}.{name}" + kind = infer_kind(fn, path=path) + sync_path = getattr(client, "sync_process_input_func", None) + has_sync = bool(sync_path) + # Do not report a hook that *is* the selected sync hook: in + # non-async mode the sync ``*_token_only`` processor is the + # active custom hook, not a dead leftover. + custom_is_sync = _custom_hook_is_sync_hook(fn, sync_path) + dead = dead_processor_hint( + kind, + async_chunk=self.async_chunk, + downstream_receives_async_chunks=self._stage_receives_async_chunks(stage_id), + has_sync=has_sync, + custom_is_sync=custom_is_sync, + ) + if dead: + logger.warning( + "[Orchestrator] stage-%d custom_process_input_func %r is never invoked " + "under the current async-chunk/sync wiring (dead-processor hint)", + stage_id, + path, + ) + except Exception as exc: # pragma: no cover - defensive only + logger.debug( + "[Orchestrator] dead-processor hint lookup for stage-%s failed: %s", + stage_id, + exc, + ) + def _get_stage_input_processor(self, stage_id: int) -> Any: processor = self._stage_input_processors.get(stage_id) if processor is None: @@ -2724,20 +2806,22 @@ async def _forward_to_next_stage_unguarded( if next_client is not None and next_client.custom_process_input_func is not None: _t_ar2d = _time.perf_counter() _fn = next_client.custom_process_input_func - _extra_kwargs: dict[str, Any] = {} - # TODO: replace signature probe with explicit kwarg contract. - try: - import inspect as _inspect + # Dispatch through the shared contract layer, replacing the + # ad-hoc `sampling_params` name probe. + from vllm_omni.model_executor.stage_input_processors._dispatch import ( + OrchestratorInputContext, + invoke_orchestrator_processor, + ) - if "sampling_params" in _inspect.signature(_fn).parameters: - _extra_kwargs["sampling_params"] = params - except (TypeError, ValueError): - pass - diffusion_prompt = _fn( + _ctx = OrchestratorInputContext( + prompt=req_state.prompt, + requires_multimodal_data=requires_multimodal_data, + sampling_params=params, + ) + diffusion_prompt = invoke_orchestrator_processor( + _fn, diffusion_source_outputs, - req_state.prompt, - requires_multimodal_data, - **_extra_kwargs, + _ctx, ) _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000 req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d @@ -3157,51 +3241,16 @@ async def _prewarm_async_chunk_stages( operation="async-chunk prewarm", ) else: - from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length - - try: - next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) - except Exception: - next_prompt_len = max(1, len(prompt_token_ids)) - # A stage may size its own prewarmed placeholder prompt via - # hf_overrides.async_chunk_prewarm_prompt_len (e.g. a talker - # whose engine positions must cover a speaker-prompt prefill; - # the model validates the value and reports the right one on - # mismatch). - prewarm_len = getattr( - getattr(next_pool.stage_vllm_config, "model_config", None), - "hf_config", - None, + # Build the prewarm placeholder through the registered + # placeholder builder's ``build_prewarm_placeholder`` (resolved + # from ``sync_process_input_func``), falling back to the legacy + # inline estimate when no builder is available. + base_input = self._build_prewarm_placeholder_input( + next_stage_id, + request_id, + prompt_token_ids, + req_state, ) - prewarm_len = getattr(prewarm_len, "async_chunk_prewarm_prompt_len", None) - if prewarm_len is not None and int(prewarm_len) > 1: - next_prompt_len = int(prewarm_len) - - original_prompt = req_state.prompt - if isinstance(original_prompt, dict): - from vllm_omni.engine.request_snapshot import copy_request_snapshot - - # The placeholder replaces the prompt and multimodal fields, - # so they are not copied (multimodal inputs can be megabytes). - # additional_information is serialized below, which copies - # tensor bytes; its containers are copied, its tensors shared. - memo: dict[int, Any] = {} - base_input = copy_request_snapshot( - {k: v for k, v in original_prompt.items() if k not in _PREWARM_UNCOPIED_PROMPT_KEYS}, - memo, - ) - if "additional_information" in original_prompt: - base_input["additional_information"] = copy_request_snapshot( - original_prompt["additional_information"], - memo, - shared_types=(torch.Tensor,), - ) - else: - base_input = {} - - base_input["prompt_token_ids"] = [0] * next_prompt_len - base_input["multi_modal_data"] = None - base_input["mm_processor_kwargs"] = None downstream_resumable = bool(getattr(entry_request, "resumable", req_state.streaming.enabled)) # Session-owned and resumable requests keep their downstream # request across turns, so they stay on the chunk-0 path. @@ -3336,6 +3385,108 @@ def _diffusion_submit_kwargs( ) } + @staticmethod + def _apply_prewarm_prompt_len_override(base_input: dict[str, Any], configured_len: int | None) -> dict[str, Any]: + """Apply the configured ``async_chunk_prewarm_prompt_len`` override to base_input.""" + if configured_len is None or int(configured_len) <= 1: + return base_input + patched = dict(base_input) + patched["prompt_token_ids"] = [0] * int(configured_len) + return patched + + def _build_prewarm_placeholder_input( + self, + next_stage_id: int, + request_id: str, + prompt_token_ids: Any, + req_state: OrchestratorRequestState, + ) -> dict[str, Any]: + """Build the async-chunk prewarm ``base_input`` dict for a downstream stage. + + Uses the registered placeholder builder's ``build_prewarm_placeholder`` + when available; otherwise (and on any builder failure) falls back to the + legacy inline estimate (``max(1, compute_talker_prompt_ids_length(...))``) + with a warning — the prewarm path must never raise. + + The configured override is applied after both paths (amy-why-3459 review of HEAD ad81ac3a). + """ + from vllm_omni.model_executor.stage_input_processors import OrchestratorInputContext + + original_prompt = req_state.prompt + if isinstance(original_prompt, dict): + from vllm_omni.engine.request_snapshot import copy_request_snapshot + + # Discarded fields are skipped; additional_information tensors stay shared. + memo: dict[int, Any] = {} + base_input = copy_request_snapshot( + {k: v for k, v in original_prompt.items() if k not in _PREWARM_UNCOPIED_PROMPT_KEYS}, + memo, + ) + if "additional_information" in original_prompt: + base_input["additional_information"] = copy_request_snapshot( + original_prompt["additional_information"], + memo, + shared_types=(torch.Tensor,), + ) + else: + base_input = {} + + # Apply the configured async_chunk_prewarm_prompt_len override on both prewarm paths. + stage_vllm_config = getattr(self.stage_pools[next_stage_id], "stage_vllm_config", None) + model_config = getattr(stage_vllm_config, "model_config", None) + prewarm_len = getattr(getattr(model_config, "hf_config", None), "async_chunk_prewarm_prompt_len", None) + + pool = self.stage_pools[next_stage_id] + client = getattr(pool, "stage_client", None) + prewarm_builder = getattr(client, "prewarm_placeholder_builder", None) + if prewarm_builder is not None: + try: + ctx = OrchestratorInputContext( + prompt=req_state.prompt, + requires_multimodal_data=False, + streaming_context=None, + target_model_config=model_config, + next_stage_hf_config=getattr(model_config, "hf_config", None), + ) + placeholders = prewarm_builder( + stage0_prompt=prompt_token_ids, + ctx=ctx, + downstream_stage_id=next_stage_id, + ) + placeholder = placeholders[0] if placeholders else None + if placeholder is not None: + ph_ids = ( + placeholder["prompt_token_ids"] + if isinstance(placeholder, dict) + else getattr(placeholder, "prompt_token_ids", None) + ) + if ph_ids: + base_input["prompt_token_ids"] = list(ph_ids) + base_input["multi_modal_data"] = None + base_input["mm_processor_kwargs"] = None + return self._apply_prewarm_prompt_len_override(base_input, prewarm_len) + except Exception as exc: + logger.warning( + "[Orchestrator] req=%s: prewarm placeholder builder for " + "stage-%s raised; falling back to inline estimate: %s", + request_id, + next_stage_id, + exc, + ) + + # Legacy inline fallback (warn-and-keep-old behaviour so the prewarm + # never regresses on models without a dual-entry builder). + from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length + + try: + next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids)) + except Exception: + next_prompt_len = max(1, len(prompt_token_ids)) + base_input["prompt_token_ids"] = [0] * next_prompt_len + base_input["multi_modal_data"] = None + base_input["mm_processor_kwargs"] = None + return self._apply_prewarm_prompt_len_override(base_input, prewarm_len) + def _build_kv_sender_info( self, sender_stage_ids: list[int], diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index a8ab3678e9b..48b3ad0ac56 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -9,7 +9,6 @@ from __future__ import annotations -import inspect import os import socket from typing import TYPE_CHECKING, Any @@ -154,6 +153,12 @@ def __init__( self.prompt_expand_func = metadata.prompt_expand_func self.custom_process_input_func = metadata.custom_process_input_func self.async_chunk_prewarm_payload_func = getattr(metadata, "async_chunk_prewarm_payload_func", None) + # Raw ``sync_process_input_func`` path (dead-processor hint) plus the + # already-resolved dual-entry ``build_prewarm_placeholder`` builder. + # Both default to None so older metadata doubles lacking the fields + # (e.g. the upstream async-chunk prewarm tests) keep working. + self.sync_process_input_func = getattr(metadata, "sync_process_input_func", None) + self.prewarm_placeholder_builder = getattr(metadata, "prewarm_placeholder_builder", None) self.engine_outputs: Any = None self.client_addresses = dict(client_addresses or {}) @@ -464,54 +469,32 @@ def process_engine_inputs( and the original prompt. """ if self.custom_process_input_func is not None: - return self._call_custom_process_input(source_outputs, prompt, streaming_context) + # Dispatch through the shared contract layer, replacing the ad-hoc + # arity probe (`len(signature.parameters) >= 4`) + from vllm_omni.model_executor.stage_input_processors._dispatch import ( + OrchestratorInputContext, + invoke_orchestrator_processor, + ) + + ctx = OrchestratorInputContext( + prompt=prompt, + requires_multimodal_data=self.requires_multimodal_data, + streaming_context=streaming_context, + target_model_config=getattr(getattr(self, "vllm_config", None), "model_config", None), + next_stage_hf_config=getattr(self, "_stage_hf_config", None), + ) + return invoke_orchestrator_processor( + self.custom_process_input_func, + source_outputs, + ctx, + ) if not self.engine_input_source: raise ValueError(f"engine_input_source empty for stage {self.stage_id}") + # The default path below is itself a C0-shape processor; it is kept out of + # the contract dispatch to preserve its exact fallback semantics unchanged. return _default_process_engine_inputs(source_outputs, prompt, self.requires_multimodal_data) - def _call_custom_process_input( - self, - source_outputs: list[Any], - prompt: Any, - streaming_context: Any | None, - ) -> list[OmniTokensPrompt]: - """Call a stage input processor with its explicitly requested context.""" - processor = self.custom_process_input_func - assert processor is not None - signature = inspect.signature(processor) - extra_kwargs: dict[str, Any] = {} - if "next_stage_hf_config" in signature.parameters: - # Let a processor size the next stage's prompt from that - # stage's model config (e.g. a talker whose engine positions - # must cover a speaker-prompt prefill). - extra_kwargs["next_stage_hf_config"] = self._stage_hf_config - target_model_config = signature.parameters.get("target_model_config") - if target_model_config is not None: - # The JoyAI bridge needs the Talker tokenizer and model config - # to calculate the exact prompt length. - if target_model_config.kind is not inspect.Parameter.KEYWORD_ONLY: - raise TypeError("target_model_config must be a keyword-only parameter") - extra_kwargs["target_model_config"] = self.vllm_config.model_config - # Match the context parameter by name, including the - # underscore-prefixed spelling some processors use (e.g. - # MiniCPM-o's ``llm2tts(..., _streaming_context)``), so bridge - # state keeps flowing to them. - if "streaming_context" in signature.parameters or "_streaming_context" in signature.parameters: - return processor( - source_outputs, - prompt, - self.requires_multimodal_data, - streaming_context, - **extra_kwargs, - ) - return processor( - source_outputs, - prompt, - self.requires_multimodal_data, - **extra_kwargs, - ) - async def collective_rpc_async( self, method: str, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index e79e7e77098..454907de5e6 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -59,6 +59,7 @@ from vllm_omni.entrypoints.utils import filter_dataclass_kwargs from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams from vllm_omni.inputs.preprocess import build_omni_renderer, omni_renderer_cls +from vllm_omni.model_executor.stage_input_processors import resolve_processor from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig @@ -572,6 +573,13 @@ class StageMetadata: custom_process_input_func: Callable | None model_stage: str | None runtime_cfg: Any + # Raw ``sync_process_input_func`` path (the ``*_token_only`` placeholder + # builder) read by the orchestrator's dead-processor hint. + sync_process_input_func: str | None = None + # Resolved ``build_prewarm_placeholder`` (the dual-entry ``*_token_only`` + # builder), so the orchestrator reads it off the stage client instead of + # re-resolving per async-chunk prewarm request. + prewarm_placeholder_builder: Callable | None = None prompt_transform_func: Callable | None = None prompt_expand_func: Callable | None = None cfg_kv_collect_func: Callable | None = None @@ -605,6 +613,55 @@ def _apply_rocm_attention_backend( engine_args["attention_backend"] = "TRITON_ATTN" +def _resolve_prewarm_placeholder_builder( + *, + sync_path: str | None, + custom_fn: Callable | None, + stage_config: Any, +) -> Callable | None: + """Resolve ``build_prewarm_placeholder`` at metadata-extraction time. + + Order mirrors the legacy orchestrator lookup: the registered + ``sync_process_input_func`` fn first (with a module-level sibling + fallback), then the already-resolved ``custom_process_input_func`` fn. + A missing or failing builder only warns — the async-chunk prewarm keeps + its inline estimate and never raises on builder resolution. + """ + builder: Callable | None = None + if sync_path: + try: + spec = resolve_processor(sync_path, expected_kind=None, stage_config=stage_config) + candidate = getattr(spec.fn, "build_prewarm_placeholder", None) + if not callable(candidate): + # Module-level sibling (covers paths whose fn does not expose + # the attribute directly). + candidate = getattr( + importlib.import_module(sync_path.rsplit(".", 1)[0]), + "build_prewarm_placeholder", + None, + ) + if callable(candidate): + builder = candidate + else: + logger.warning( + "sync_process_input_func=%r has no build_prewarm_placeholder; " + "async-chunk prewarm will use the inline estimate", + sync_path, + ) + except Exception as exc: + logger.warning( + "failed to resolve build_prewarm_placeholder for %r; " + "async-chunk prewarm will use the inline estimate: %s", + sync_path, + exc, + ) + if builder is None and callable(custom_fn): + candidate = getattr(custom_fn, "build_prewarm_placeholder", None) + if callable(candidate): + builder = candidate + return builder + + def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: """Extract metadata through the active production legacy path. @@ -640,8 +697,32 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: custom_process_input_func: Callable | None = None _cpif_path = _get_attr_or_item(stage_config, "custom_process_input_func") if _cpif_path: - mod_path, fn_name = _cpif_path.rsplit(".", 1) - custom_process_input_func = getattr(importlib.import_module(mod_path), fn_name) + # Resolve through the stage-input processor registry. expected_kind is + # None here — orchestrator-side processors carry no suffix contract, so + # we rely on suffix inference only (compatible with existing configs). + # A bad path must **fail fast**: let ProcessorValidationError / ImportError / + # AttributeError propagate instead of silently dropping the configured + # hook and falling back to ``_default_process_engine_inputs``. (The + # worker-side mixin candidate chain, which *probes* for a ``*_full_payload`` + # derivative, keeps its own skip-on-mismatch tolerance.) + custom_process_input_func = resolve_processor( + _cpif_path, + expected_kind=None, + stage_config=stage_config, + ).fn + + sync_process_input_func: str | None = None + _spif_path = _get_attr_or_item(stage_config, "sync_process_input_func") + if _spif_path: + # Keep the raw dotted path (the dead-processor hint reads it); the + # dual-entry ``build_prewarm_placeholder`` is resolved once here. + sync_process_input_func = _spif_path + + prewarm_placeholder_builder = _resolve_prewarm_placeholder_builder( + sync_path=sync_process_input_func, + custom_fn=custom_process_input_func, + stage_config=stage_config, + ) prompt_transform_func: Callable | None = None _ptf_path = _get_attr_or_item(stage_config, "prompt_transform_func") @@ -681,6 +762,8 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: final_output_type=final_output_type, default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, + sync_process_input_func=sync_process_input_func, + prewarm_placeholder_builder=prewarm_placeholder_builder, model_stage=model_stage, runtime_cfg=runtime_cfg, prompt_transform_func=prompt_transform_func, @@ -702,6 +785,8 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: final_output_type=final_output_type, default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, + sync_process_input_func=sync_process_input_func, + prewarm_placeholder_builder=prewarm_placeholder_builder, model_stage=model_stage, runtime_cfg=runtime_cfg, prompt_transform_func=prompt_transform_func, @@ -736,6 +821,18 @@ def extract_stage_metadata_from_omni_stage_config( sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams sampling_params = sampling_params_cls(**(stage_config.model_config.default_sampling_params or {})) custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func) + # Carry the raw sync path (dead-processor hint); resolve the dual-entry + # ``build_prewarm_placeholder`` once here. + sync_process_input_func: str | None = None + _spif_path = getattr(stage_config.stage_pipeline_config, "sync_process_input_func", None) + if _spif_path: + sync_process_input_func = _spif_path + + prewarm_placeholder_builder = _resolve_prewarm_placeholder_builder( + sync_path=sync_process_input_func, + custom_fn=custom_process_input_func, + stage_config=stage_config, + ) if stage_type == "diffusion": return StageMetadata( @@ -749,6 +846,8 @@ def extract_stage_metadata_from_omni_stage_config( final_output_type=stage_config.final_output_type, default_sampling_params=sampling_params, custom_process_input_func=custom_process_input_func, + sync_process_input_func=sync_process_input_func, + prewarm_placeholder_builder=prewarm_placeholder_builder, model_stage=stage_config.model_stage, runtime_cfg=stage_config.runtime_config, prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func), @@ -766,6 +865,8 @@ def extract_stage_metadata_from_omni_stage_config( final_output_type=stage_config.final_output_type, default_sampling_params=sampling_params, custom_process_input_func=custom_process_input_func, + sync_process_input_func=sync_process_input_func, + prewarm_placeholder_builder=prewarm_placeholder_builder, model_stage=stage_config.model_stage, runtime_cfg=stage_config.runtime_config, prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func), diff --git a/vllm_omni/model_executor/models/aura_omni/pipeline.py b/vllm_omni/model_executor/models/aura_omni/pipeline.py index 71f451c001a..667117d1f1a 100644 --- a/vllm_omni/model_executor/models/aura_omni/pipeline.py +++ b/vllm_omni/model_executor/models/aura_omni/pipeline.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 """AURA Omni pipeline topology. Semantic modules: @@ -18,6 +18,10 @@ StagePipelineConfig, ) +# Talker stop token single-sourced in the shared constants module (also used +# by qwen3_omni / qwen3_tts pipelines). +from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID + _AURA_PROC = "vllm_omni.model_executor.stage_input_processors.aura_omni" _QWEN3_TTS_PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" @@ -66,7 +70,7 @@ async_chunk_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_async_chunk", sampling_constraints={ "detokenize": False, - "stop_token_ids": [2150], + "stop_token_ids": [QWEN3_CODEC_EOS_TOKEN_ID], }, ), StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index 2036c42aa07..8880c40d3c6 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.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 """Qwen2.5-Omni pipeline topology (frozen). Stage 0: Thinker — multimodal understanding + text generation @@ -13,6 +13,9 @@ StagePipelineConfig, ) +# Talker codec end token id (8294) single-sourced in the stage input processor. +from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID + _PROC = "vllm_omni.model_executor.stage_input_processors.qwen2_5_omni" QWEN2_5_OMNI_PIPELINE = PipelineConfig( @@ -45,7 +48,7 @@ custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", sampling_constraints={ "detokenize": True, - "stop_token_ids": [8294], + "stop_token_ids": [TALKER_CODEC_END_TOKEN_ID], }, requires_full_payload_input=True, ), diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index 22dbcfac32c..bec1b4b9870 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -44,6 +44,10 @@ Qwen2_5OmniThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges + +# The talker codec end token id is single-sourced in the stage input processor +# (``TALKER_CODEC_END_TOKEN_ID``). +from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import TALKER_CODEC_END_TOKEN_ID from vllm_omni.platforms import current_omni_platform # `get_llm_pos_ids_for_vision` was removed from upstream vLLM (commit @@ -51,7 +55,9 @@ # OmniMRotaryEmbedding. Bind it at module level to keep call sites unchanged. get_llm_pos_ids_for_vision = OmniMRotaryEmbedding._get_llm_pos_ids_for_vision -TALKER_CODEC_EOS_TOKEN_ID = 8294 +# Alias kept for backwards compatibility; the single source is +# ``TALKER_CODEC_END_TOKEN_ID`` in ``stage_input_processors.qwen2_5_omni``. +TALKER_CODEC_EOS_TOKEN_ID = TALKER_CODEC_END_TOKEN_ID TALKER_CODEC_BOS_TOKEN_ID = 8293 diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index a32462d3293..f1936254fcf 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -17,6 +17,10 @@ pipeline_cfg_resolver, ) +# Talker stop token single-sourced in the shared constants module (also used +# by the qwen3_tts / aura_omni pipelines). +from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID + _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_omni" QWEN3_OMNI_PIPELINE = PipelineConfig( @@ -68,7 +72,7 @@ supports_duplex_mrv2=True, sampling_constraints={ "detokenize": False, - "stop_token_ids": [2150], + "stop_token_ids": [QWEN3_CODEC_EOS_TOKEN_ID], }, requires_full_payload_input=True, ), diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index b0537ead0ff..42be2acaf61 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -65,6 +65,16 @@ from vllm_omni.model_executor.models.qwen3_omni.tool_speech import spoken_span_end from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_decode_windows + +# Codec token ids are single-sourced in the stage input processor +# (``_QWEN3_CODEC_*``); the model keeps the ``TALKER_CODEC_*`` aliases for +# backwards compatibility. At runtime the talker actually reads them from +# ``config.talker_config.codec_{pad,bos,eos}``. +from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( + _QWEN3_CODEC_BOS_TOKEN_ID, + _QWEN3_CODEC_EOS_TOKEN_ID, + _QWEN3_CODEC_PAD_TOKEN_ID, +) from vllm_omni.platforms import current_omni_platform from vllm_omni.utils.device_copy import to_device_nonblocking @@ -96,10 +106,13 @@ def resolve_default_speaker(speaker_token_ids: dict[str, int]) -> str: TTS_EOS_TOKEN_ID = 151673 # (end of dialogue) TTS_BOS_SINGLE_TOKEN_ID = 151674 # -# Talker codec tokens (talker vocabulary, used for RVQ code generation) -TALKER_CODEC_PAD_TOKEN_ID = 4196 # Padding token -TALKER_CODEC_BOS_TOKEN_ID = 4197 # Beginning of speech -TALKER_CODEC_EOS_TOKEN_ID = 4198 # End of speech +# Talker codec tokens (talker vocabulary, used for RVQ code generation). +# Single-sourced from ``stage_input_processors.qwen3_omni`` +# (``_QWEN3_CODEC_*``); at runtime the talker reads these ids from +# ``config.talker_config.codec_{pad,bos,eos}`` instead. +TALKER_CODEC_PAD_TOKEN_ID = _QWEN3_CODEC_PAD_TOKEN_ID # Padding token +TALKER_CODEC_BOS_TOKEN_ID = _QWEN3_CODEC_BOS_TOKEN_ID # Beginning of speech +TALKER_CODEC_EOS_TOKEN_ID = _QWEN3_CODEC_EOS_TOKEN_ID # End of speech TALKER_CODEC_NOTHINK_ID = 4203 # No-think mode TALKER_CODEC_THINK_BOS_ID = 4204 # Think mode start TALKER_CODEC_THINK_EOS_ID = 4205 # Think mode end diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index ea3e2fcfe36..a2d258796d5 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -81,6 +81,15 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config + # The codec token-id consistency guard is deliberately not run at module + # import (it pulls in transformers / the qwen3-tts config). Invoke it + # explicitly at model-load startup so the processor constants stay in + # sync with the HF configs this model actually runs with. + from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( + _assert_codec_token_ids_consistent, + ) + + _assert_codec_token_ids_consistent() rope_params = getattr(talker_config.text_config, "rope_scaling", None) if rope_params is None: rope_params = getattr(talker_config.text_config, "rope_parameters", None) or {} diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index e1644f24215..f6f3cb712f6 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -10,6 +10,10 @@ StageExecutionType, StagePipelineConfig, ) + +# Talker stop token single-sourced in the shared constants module (also used +# by qwen3_omni / aura_omni pipelines). +from vllm_omni.model_executor.stage_input_processors._constants import QWEN3_CODEC_EOS_TOKEN_ID from vllm_omni.outputs.output_modality import ( TensorAccumulationStrategy, register_key_accumulation_strategy, @@ -47,7 +51,7 @@ custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", sampling_constraints={ "detokenize": False, - "stop_token_ids": [2150], + "stop_token_ids": [QWEN3_CODEC_EOS_TOKEN_ID], }, ), StagePipelineConfig( diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py index e69de29bb2d..9451ca067a8 100644 --- a/vllm_omni/model_executor/stage_input_processors/__init__.py +++ b/vllm_omni/model_executor/stage_input_processors/__init__.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Stage input processors: per-model builders for inter-stage transitions. + +Naming convention: ``*_full_payload``/``*_async_chunk`` are producer-side +(worker) builders (``pooling_output``/``multimodal_output`` load-bearing); +``*_token_only`` and no-suffix names are consumer-side (orchestrator) builders, +normalized through ``wrap_orchestrator_processor``. + +This package re-exports only the dispatch contract and registry — importing it +never pulls in model-specific code. +""" + +from vllm_omni.model_executor.stage_input_processors._dispatch import ( + DiffusionInputBuilder, + OrchestratorInputContext, + PlaceholderPromptBuilder, + invoke_orchestrator_processor, + wrap_orchestrator_processor, +) +from vllm_omni.model_executor.stage_input_processors._registry import ( + ProcessorKind, + ProcessorSpec, + ProcessorValidationError, + dead_processor_hint, + infer_kind, + register_processor, + resolve_processor, + validate_processor, +) + +__all__ = [ + "OrchestratorInputContext", + "PlaceholderPromptBuilder", + "DiffusionInputBuilder", + "invoke_orchestrator_processor", + "wrap_orchestrator_processor", + "ProcessorKind", + "ProcessorSpec", + "ProcessorValidationError", + "resolve_processor", + "register_processor", + "infer_kind", + "validate_processor", + "dead_processor_hint", +] diff --git a/vllm_omni/model_executor/stage_input_processors/_common.py b/vllm_omni/model_executor/stage_input_processors/_common.py new file mode 100644 index 00000000000..f97853aeefc --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/_common.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Canonical shared helpers; semantics golden-locked in ``test_common_helpers_golden.py``.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch + +from vllm_omni.inputs.data import OmniTokensPrompt + +# =========================================================================== +# ensure_list +# =========================================================================== + + +def _unwrap_constant_list(x: Any) -> Any: + """Unwrap a ``ConstantList``-like object exposing ``_x``.""" + if hasattr(x, "_x"): + return x._x + return x + + +def ensure_list(x: Any) -> list[Any]: + """Canonical ``ensure_list`` (union of legacy variants). + + ``ConstantList`` unwrapped; tensor ``.tolist()``; ``None`` -> ``[]``; + other iterable via ``list(x)``; non-iterable scalar wrapped as ``[x]``. + """ + x = _unwrap_constant_list(x) + if isinstance(x, torch.Tensor): + return x.detach().cpu().tolist() + if x is None: + return [] + if isinstance(x, (list, tuple)): + return list(x) + try: + return list(x) + except TypeError: + return [x] + + +def ensure_list_unchanged(x: Any) -> Any: + """Legacy ``qwen3_omni._ensure_list``: non-``list`` (incl. ``None``) unchanged.""" + x = _unwrap_constant_list(x) + if not isinstance(x, list): + return x + return list(x) + + +def ensure_list_strict(x: Any) -> list[Any]: + """Strict variant (``ming_flash_omni``): no ``None`` case; non-iterable raises.""" + x = _unwrap_constant_list(x) + if isinstance(x, torch.Tensor): + return x.detach().cpu().tolist() + if isinstance(x, (list, tuple)): + return list(x) + return list(x) + + +def ensure_list_flatten(x: Any) -> list[Any]: + """Flattening variant (``audex``): tensors via ``reshape(-1).tolist()``.""" + x = _unwrap_constant_list(x) + if x is None: + return [] + if isinstance(x, torch.Tensor): + return x.detach().cpu().reshape(-1).tolist() + if isinstance(x, (list, tuple)): + return list(x) + return list(x) + + +def ensure_list_preserve_none(x: Any) -> list[Any]: + """Preserve-``None`` variant (``step_audio2``): non-iterable wrapped as ``[x]``.""" + if hasattr(x, "_x"): + return list(x._x) + if isinstance(x, torch.Tensor): + return x.detach().cpu().tolist() + if isinstance(x, (list, tuple)): + return list(x) + return list(x) if hasattr(x, "__iter__") else [x] + + +def ensure_list_wrap_only(x: Any) -> list[Any]: + """Wrap-only variant (``diffusion.output_formatter``): non-list wrapped whole. + + Unlike :func:`ensure_list`: no tensor flattening, no iterable iteration. + """ + if isinstance(x, list): + return x + return [x] if x is not None else [] + + +# =========================================================================== +# to_cpu_tensor +# =========================================================================== + + +def to_cpu_tensor(value: Any) -> torch.Tensor | None: + """CPU tensor when possible; else ``None`` (``glm_tts``).""" + if isinstance(value, list): + if not value: + return None + value = value[0] + if isinstance(value, torch.Tensor): + return value.detach().cpu() + return None + + +# =========================================================================== +# to_token_id_list +# =========================================================================== + + +def to_token_id_list(value: Any, *, recursive: bool = False) -> list[int]: + """Flatten a token-ish value to ``list[int]``. + + ``recursive=False`` matches ``dynin_omni`` (first row only); + ``recursive=True`` matches ``cosyvoice3`` (per-item flatten). + """ + if recursive: + return _to_token_id_list_recursive(value) + if isinstance(value, torch.Tensor): + value = value.detach().to("cpu") + if value.ndim == 0: + return [int(value.item())] + if value.ndim > 1: + value = value[0] + return [int(x) for x in value.tolist()] + if isinstance(value, list): + if not value: + return [] + if isinstance(value[0], list): + return [int(x) for x in value[0]] + return [int(x) for x in value] + if value is None: + return [] + return [int(value)] + + +def _to_token_id_list_recursive(value: Any) -> list[int]: + """Recursive token-id flattening: tensors reshaped, scalars ``int()``.""" + if value is None: + return [] + if isinstance(value, torch.Tensor): + value = value.detach().to("cpu").reshape(-1).tolist() + out: list[int] = [] + for item in ensure_list(value): + if isinstance(item, torch.Tensor): + out.extend(_to_token_id_list_recursive(item)) + continue + if isinstance(item, (list, tuple)): + out.extend(_to_token_id_list_recursive(item)) + continue + out.append(int(item)) + return out + + +# =========================================================================== +# revert_delay_pattern (higgs_audio_v2 lenient vs higgs_audio_v3 strict) +# =========================================================================== + + +def revert_delay_pattern( + audio_codes_qt: torch.Tensor, + *, + expected_codebooks: int | None = None, + allow_short: bool = False, +) -> torch.Tensor: + """Reverse the MusicGen-style delay pattern ``[Q, T]`` -> ``[Q, seq_len]``. + + ``allow_short`` (higgs v2) vs strict raise (v3); ``expected_codebooks`` + row-count check (v3 only). + """ + if audio_codes_qt.ndim != 2: + raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}") + q, t = audio_codes_qt.shape + if expected_codebooks is not None and q != expected_codebooks: + raise ValueError(f"Expected exactly {expected_codebooks} codebook rows, got {q}. Input shape: [{q}, {t}]") + if t < q: + if allow_short: + return audio_codes_qt + raise ValueError(f"Not enough frames to revert delay pattern: T={t} < Q={q}") + seq_len = t - q + 1 + out_l = [audio_codes_qt[i : i + 1, i : seq_len + i] for i in range(q)] + return torch.cat(out_l, dim=0) + + +# =========================================================================== +# filter_real_code_frames (higgs v2 [frames,Q] vs v3 [Q,frames]) +# =========================================================================== + + +def filter_real_code_frames( + audio_codes: torch.Tensor, + *, + num_real_codes: int, + layout: str = "frames_first", +) -> torch.Tensor: + """Keep frames with all codebook values in ``[0, num_real_codes)``. + + ``frames_first`` (v2) vs ``codebooks_first`` (v3, transposed) layouts. + """ + if audio_codes.numel() == 0: + return audio_codes + if layout == "frames_first": + if audio_codes.ndim != 2: + raise ValueError(f"expected [num_frames, num_codebooks] audio_codes; got shape {tuple(audio_codes.shape)}") + valid = (audio_codes >= 0).all(dim=1) & (audio_codes < num_real_codes).all(dim=1) + return audio_codes[valid] + if layout == "codebooks_first": + frames = audio_codes.t() + valid = (frames >= 0).all(dim=1) & (frames < num_real_codes).all(dim=1) + return frames[valid].t().contiguous() + raise ValueError(f"unknown layout: {layout!r}") + + +# =========================================================================== +# extract_last_codec_frame +# =========================================================================== + + +def extract_last_codec_frame( + payload: Any, + *, + key_path: tuple[str, ...] = ("codes", "audio"), + validate: str | None = None, + to_cpu: bool = False, + to_long: bool = True, +) -> torch.Tensor | None: + """Extract the last inter-stage codec frame (4 legacy variants merged). + + ``key_path`` codes/audio vs audio_codes; ``validate`` None/"any"/"valid_mask"; + ``to_cpu``/``to_long`` (fish_speech). Validity gate applies only to 2-D + frames; 1-D frames return directly; result is always 1-D. 3-D per-step + frames ``[B, 1, C]`` (Voxtral ``compute_mm_logits``) are normalized to + ``[B, C]`` first; any other 3-D layout still raises ``ValueError``. + """ + if not isinstance(payload, dict): + return None + audio_codes: Any = payload + for key in key_path: + audio_codes = audio_codes.get(key) if isinstance(audio_codes, dict) else None + if audio_codes is None: + return None + if isinstance(audio_codes, list): + if not audio_codes: + return None + audio_codes = audio_codes[-1] + if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: + return None + if audio_codes.ndim == 3 and audio_codes.shape[1] == 1: + # Voxtral TTS: compute_mm_logits() emits per-step [B, 1, C] codes + # (payload list elements of shape (1, 1, C)); the legacy voxtral + # helper flattened them, so 3-D per-step input is a supported + # layout, not a contract violation. Targeted squeeze(1) reduces + # it to [B, C] so the shared 2-D last-row semantics (including + # any validity gate) apply unchanged. Non-singleton frame dims + # (e.g. (2, 2, 2)) intentionally fall through to the ValueError + # below, preserving the golden-locked fish/qwen3/higgs contracts. + audio_codes = audio_codes.squeeze(1) + if audio_codes.ndim == 2: + frame: torch.Tensor = audio_codes[-1] + if frame.numel() == 0: + return None + if validate in ("any", "valid_mask"): + if validate == "valid_mask": + valid = payload.get("audio_code_valid") + if isinstance(valid, torch.Tensor) and valid.numel() > 0: + is_valid = bool(valid.reshape(-1)[-1].item()) + elif valid is not None: + is_valid = bool(valid) + else: + is_valid = bool(frame.any().item()) + else: + is_valid = bool(frame.any().item()) + if not is_valid: + return None + elif audio_codes.ndim == 1: + # Legacy semantics: a flat (1-D) codec frame is returned directly + # without the validity/zero gate (Fish Speech / Qwen3-TTS). + frame = audio_codes + else: + raise ValueError(f"unexpected audio_codes shape: {tuple(audio_codes.shape)}") + frame = frame.flatten() + if to_cpu: + frame = frame.detach().cpu() + if to_long: + frame = frame.to(torch.long) + return frame + + +# =========================================================================== +# Placeholder prompt length + packing +# =========================================================================== + +# Qwen chat-template sentinel ids (single source within this module; the +# codec token ids themselves are centralized separately in ``_constants``). +_IM_START_TOKEN_ID = 151644 +_SYSTEM_TOKEN_ID = 8948 +_USER_TOKEN_ID = 872 +_ASSISTANT_TOKEN_ID = 77091 +_ASSISTANT_TAIL_LEN = 9 # 3 + 4 + 1 + 1 + + +def _chat_template_prompt_len( + all_ids: Sequence[int], + prompt_ids: Sequence[int], + device: torch.device | str, +) -> int: + """Qwen chat-template scan: ``sum(user segments) + 9`` (assistant tail). + + Shared by both modes so sync forward and async prewarm agree. + """ + thinker_sequences = torch.tensor(all_ids, dtype=torch.long, device=device).unsqueeze(0) # [1, T] + input_ids = torch.tensor(prompt_ids, dtype=torch.long, device=device).unsqueeze(0) # [1, T] + + im_start_indexes = torch.cat( + [ + torch.nonzero(input_ids[0] == _IM_START_TOKEN_ID).squeeze(1), + torch.tensor( + [thinker_sequences.shape[-1]], + device=input_ids.device, + dtype=input_ids.dtype, + ), + ], + dim=0, + ) + + sum_user_len = 0 + assistant_len = 0 + for i in range(len(im_start_indexes) - 1): + s = int(im_start_indexes[i].item()) + e = int(im_start_indexes[i + 1].item()) + role = int(input_ids[0, s + 1].item()) + if role == _SYSTEM_TOKEN_ID: + continue + if role == _USER_TOKEN_ID: + sum_user_len += e - s + elif role == _ASSISTANT_TOKEN_ID and i == len(im_start_indexes) - 2: + assistant_len += _ASSISTANT_TAIL_LEN + return sum_user_len + assistant_len + + +def compute_placeholder_prompt_len( + *, + ids_or_prompt: Any, + mode: str = "full", + device: torch.device | str = "cpu", +) -> int: + """Compute the downstream placeholder prompt length. + + ``mode="full"``: ``ids.all``/``ids.prompt`` scan (``sum(user) + 9``). + ``mode="stage0_only"``: prewarm estimate feeding stage-0 ids to both roles, + matching ``adapter.compute_talker_prompt_ids_length`` so builder and inline + fallback agree. + """ + if mode == "stage0_only": + prompt = list(ids_or_prompt) + return _chat_template_prompt_len(prompt, prompt, device) + if mode != "full": + raise ValueError(f"unknown mode: {mode!r}") + + ids = (ids_or_prompt or {}).get("ids", {}) + return _chat_template_prompt_len(ids["all"], ids["prompt"], device) + + +def pack_placeholder_prompt( + *, + prompt_len: int, + voice_metadata: dict[str, Any] | None = None, +) -> OmniTokensPrompt: + """Pack a KV-slot placeholder prompt: ``[0] * prompt_len`` + voice meta.""" + return OmniTokensPrompt( + prompt_token_ids=[0] * max(1, int(prompt_len)), + additional_information=(voice_metadata if voice_metadata else None), # type: ignore[typeddict-item] + ) diff --git a/vllm_omni/model_executor/stage_input_processors/_constants.py b/vllm_omni/model_executor/stage_input_processors/_constants.py new file mode 100644 index 00000000000..5af5b949ff7 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/_constants.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Canonical cross-model constant definitions. + +Special-token ids shared across model pipelines / processors / model +implementations live here so no module hard-codes a magic number that could +drift out of sync. This module is deliberately **dependency-free** (no torch / +vllm imports) so pure-config modules can import it without pulling in the +runtime stack. + +**Scope note:** only constants genuinely shared across more than one module +belong here. Model-specific codec ids keep their module-top home with the +processor as source and the HF config as the runtime single source of truth. +""" + +# Qwen3-family talker codec stop (EOS) token id; used as the AR talker +# stage's ``stop_token_ids`` across the qwen3_omni / qwen3_tts / aura_omni +# pipelines. +QWEN3_CODEC_EOS_TOKEN_ID = 2150 diff --git a/vllm_omni/model_executor/stage_input_processors/_dispatch.py b/vllm_omni/model_executor/stage_input_processors/_dispatch.py new file mode 100644 index 00000000000..6934e4ed30b --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/_dispatch.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Consumer-side (orchestrator) input-dispatch contract layer. + +Collapses the two legacy runtime probes (arity in ``process_engine_inputs``, +sampling_params-name in ``_forward_to_next_stage``) into one normalization: + +* ``OrchestratorInputContext`` — the fixed transition context passed to every + consumer-side builder; ``PlaceholderPromptBuilder`` / ``DiffusionInputBuilder`` + are the two C1-compatible callable roles ``(source_outputs, ctx)``. +* ``wrap_orchestrator_processor`` passes through ``ctx``-native processors and + adapts legacy positional shapes (C0 3-arg; C2 ``streaming_context`` placeholder + incl. ``_streaming_context``/four-positional; C3 ``sampling_params`` diffusion; + C4 ``moss_tts.talker2codec``) with a one-time ``DeprecationWarning``. + +Producer roles are out of scope: their ``pooling_output``/``multimodal_output`` +kwargs are load-bearing and stay as-is (validated in ``_registry``). +""" + +from __future__ import annotations + +import inspect +import warnings +from dataclasses import dataclass +from typing import Any, Protocol + +from vllm_omni.inputs.data import OmniTokensPrompt + +__all__ = [ + "OrchestratorInputContext", + "PlaceholderPromptBuilder", + "DiffusionInputBuilder", + "wrap_orchestrator_processor", + "invoke_orchestrator_processor", +] + + +@dataclass(frozen=True) +class OrchestratorInputContext: + """Fixed transition context passed to orchestrator-facing processors. + + ``target_model_config`` / ``next_stage_hf_config`` are populated by the + stage client from the stage's own vllm config so that legacy processors + (e.g. JoyAI ``joyai_action_to_tts``, Nemotron ``thinker2talker_token_only``) + can size next-stage prompts from the target model configuration. + """ + + prompt: Any | None = None + requires_multimodal_data: bool = False + streaming_context: Any | None = None + sampling_params: Any | None = None + target_model_config: Any | None = None + next_stage_hf_config: Any | None = None + + +class PlaceholderPromptBuilder(Protocol): + """C2-style sync builder: upstream outputs -> next-stage token prompts.""" + + def __call__( + self, + source_outputs: list[Any], + ctx: OrchestratorInputContext, + ) -> list[OmniTokensPrompt]: ... + + +class DiffusionInputBuilder(Protocol): + """Diffusion-stage input builder: upstream outputs -> diffusion payload(s).""" + + def __call__( + self, + source_outputs: list[Any], + ctx: OrchestratorInputContext, + ) -> dict | list[dict] | None: ... + + +# --------------------------------------------------------------------------- +# Signature inspection helpers +# --------------------------------------------------------------------------- + + +def _is_orchestrator_context_annotation(annotation: Any) -> bool: + """Whether ``annotation`` names ``OrchestratorInputContext`` (class, string, or dotted).""" + if annotation is inspect.Parameter.empty: + return False + if annotation is OrchestratorInputContext: + return True + if isinstance(annotation, str): + text = annotation.strip().strip("'\"") + return text == "OrchestratorInputContext" or text.endswith("OrchestratorInputContext") + return False + + +def _accepts_orchestrator_input_context(fn: Any) -> bool: + """Whether ``fn`` already speaks the ``ctx`` contract (param name or annotation).""" + try: + signature = inspect.signature(fn) + except (TypeError, ValueError): + return False + for name, param in signature.parameters.items(): + if name == "ctx": + return True + if _is_orchestrator_context_annotation(param.annotation): + return True + return False + + +def _legacy_signature(fn: Any) -> inspect.Signature | None: + """Best-effort ``inspect.signature`` probe for a callable.""" + try: + return inspect.signature(fn) + except (TypeError, ValueError): + return None + + +def _positional_names(signature: inspect.Signature) -> list[str]: + """Names of positional (or positional-or-keyword) parameters.""" + return [ + name + for name, param in signature.parameters.items() + if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + ] + + +def _legacy_shape(signature: inspect.Signature | None) -> str: + """Classify a legacy positional shape (full signature, keyword-only included). + + ``sampling_params`` -> C3 (shells too); ``>= 4`` positional -> C2 streaming. + """ + if signature is None: + return "c0" + params = signature.parameters + pos_names = _positional_names(signature) + if len(pos_names) >= 2 and pos_names[0] == "stage_list" and pos_names[1] == "engine_input_source": + return "moss" + if "sampling_params" in pos_names: + return "c3" + if len(pos_names) >= 4: + return "c2pos" + if "sampling_params" in params: + return "c3" + if "streaming_context" in params or "_streaming_context" in params: + return "c2kw" + return "c0" + + +def _ctx_kwarg_names(signature: inspect.Signature | None) -> tuple[str, ...]: + """Keyword-only ctx config params a legacy signature accepts. + + Excludes ``sampling_params``/``streaming_context`` (injected by the C3/C2 + adapters — forwarding again would raise a duplicate-kwarg ``TypeError``). + """ + if signature is None: + return () + return tuple( + name + for name, param in signature.parameters.items() + if param.kind is inspect.Parameter.KEYWORD_ONLY and name in ("target_model_config", "next_stage_hf_config") + ) + + +def _warn_legacy_contract(fn: Any) -> None: + """Emit a ``DeprecationWarning`` for a legacy positional processor shape.""" + name = getattr(fn, "__qualname__", None) or getattr(fn, "__name__", None) or repr(fn) + warnings.warn( + f"stage-input processor {name!r} uses the legacy positional contract; " + "migrate it to the OrchestratorInputContext contract.", + DeprecationWarning, + stacklevel=2, + ) + + +# --------------------------------------------------------------------------- +# Legacy adapters +# --------------------------------------------------------------------------- + + +def _legacy_adapter(fn: Any, signature: inspect.Signature | None, shape: str) -> Any: + """Build the C1 adapter for a legacy *shape*; resolve and validate once. + + Call shape fixed at wrap time; dry-run ``bind`` fails fast on broken + legacy calls at configuration time. + """ + pos_names = _positional_names(signature) if signature is not None else [] + ctx_kwarg_names = _ctx_kwarg_names(signature) + c2pos_fourth = ( + "next_stage_hf_config" if len(pos_names) > 3 and pos_names[3] == "next_stage_hf_config" else "streaming_context" + ) + if signature is not None: + if shape == "c0": + signature.bind(None, None, None, **{name: None for name in ctx_kwarg_names}) + elif shape == "c2pos": + signature.bind(None, None, None, None, **{name: None for name in ctx_kwarg_names}) + elif shape == "c3": + signature.bind(None, None, None, sampling_params=None) + elif shape == "c2kw": + signature.bind(None, None, None, streaming_context=None) + elif shape == "moss": + signature.bind(None, None) + + def _adapted(source_outputs: list[Any], ctx: OrchestratorInputContext) -> Any: + if shape == "moss": + return fn(source_outputs, ctx.prompt) + if shape == "c3": + return fn( + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + sampling_params=ctx.sampling_params, + ) + if shape == "c2pos": + return fn( + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + getattr(ctx, c2pos_fourth), + **{name: getattr(ctx, name) for name in ctx_kwarg_names}, + ) + if shape == "c2kw": + return fn( + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + streaming_context=ctx.streaming_context, + ) + # C0 3-arg legacy (and the uninspectable / builtin fallback). + return fn( + source_outputs, + ctx.prompt, + ctx.requires_multimodal_data, + **{name: getattr(ctx, name) for name in ctx_kwarg_names}, + ) + + _warn_legacy_contract(fn) + return _adapted + + +#: Wrap memo: the signature probe, the call-shape resolution and the +#: ``DeprecationWarning`` happen once per processor, not per forward. +_WRAP_CACHE: dict[Any, Any] = {} + + +def wrap_orchestrator_processor(fn: Any) -> PlaceholderPromptBuilder | DiffusionInputBuilder: + """Return a C1-compatible ``(source_outputs, ctx)`` callable for ``fn``. + + ``ctx``-native processors pass through unchanged; legacy shapes are adapted + per ``_legacy_shape`` with a one-time ``DeprecationWarning``. Result is + memoized in ``_WRAP_CACHE`` so the wrap (and warning) happens once per + processor, not per forward. + """ + if not callable(fn): + raise TypeError(f"stage-input processor must be callable, got {fn!r}") + + try: + cached = _WRAP_CACHE.get(fn) + except TypeError: # pragma: no cover - unhashable callable (rare) + cached = None + if cached is not None: + return cached # type: ignore[return-value] + + if _accepts_orchestrator_input_context(fn): + wrapped: Any = fn + else: + signature = _legacy_signature(fn) + wrapped = _legacy_adapter(fn, signature, _legacy_shape(signature)) + + try: + _WRAP_CACHE[fn] = wrapped + except TypeError: # pragma: no cover - unhashable callable (rare) + pass + return wrapped # type: ignore[return-value] + + +def invoke_orchestrator_processor( + fn: Any, + source_outputs: list[Any], + ctx: OrchestratorInputContext, +) -> Any: + """Invoke an orchestrator-facing processor under ``(source_outputs, ctx)``.""" + return wrap_orchestrator_processor(fn)(source_outputs, ctx) diff --git a/vllm_omni/model_executor/stage_input_processors/_registry.py b/vllm_omni/model_executor/stage_input_processors/_registry.py new file mode 100644 index 00000000000..d66e28a0489 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/_registry.py @@ -0,0 +1,535 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Stage-input processor registry: signature-level structural validation. + +Serves the three importlib resolution points (orchestrator-side +``custom_process_input_func``/``sync_process_input_func``, producer-side +``*_full_payload``, producer-side ``*_async_chunk``). Checks are purely +signature-based (``inspect.signature``); behaviour is covered by the golden / +parity tests. + +Producer kwargs are contract: full-payload keeps ``pooling_output`` +(``is_finished`` optional — the worker retries without it); async-chunk keeps +``multimodal_output`` and **must** accept ``is_finished`` (the scheduler always +passes it). + +Kind inference is name-driven (suffix-based); ``resolve_processor`` returns a +:class:`ProcessorSpec` whose ``fn`` matches the legacy importlib lookup. +""" + +from __future__ import annotations + +import importlib +import inspect +import warnings +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any, Literal, NoReturn, get_args + +ProcessorKind = Literal[ + "placeholder_prompt_builder", + "diffusion_input_builder", + "legacy_orchestrator_builder", + "producer_full_payload", + "producer_async_chunk", + "legacy_multi_source", +] + +#: Suffixes that identify a producer-side full-payload builder. +_FULL_PAYLOAD_SUFFIXES: tuple[str, ...] = ("_full_payload", "_batch") +#: Suffix that identifies a producer-side async-chunk builder. +_ASYNC_CHUNK_SUFFIX: str = "_async_chunk" +#: Suffix that identifies a consumer-side placeholder builder. +_TOKEN_ONLY_SUFFIX: str = "_token_only" + +#: Diffusion-stage input-builder name stems (suffix match). +_DIFFUSION_SUFFIXES: tuple[str, ...] = ( + "ar2diffusion", + "ar2dit", + "thinker2imagegen", + "ar2acoustic", + "ar_to_dit", + "llm2audio_vae", + "thinker2token2wav", +) + +#: Module hint for the only remaining legacy multi-source processor family. +_MOSS_MULTI_SOURCE_MODULE_HINT: str = "moss_tts" +#: Exact function names in ``moss_tts`` with the multi-source shape. +_MOSS_MULTI_SOURCE_NAMES: frozenset[str] = frozenset({"talker2codec"}) + +#: Producer kinds are never orchestrator-side, so the dead-processor hint never +#: applies to them. +_PRODUCER_KINDS: frozenset[str] = frozenset({"producer_full_payload", "producer_async_chunk"}) + +#: All valid :class:`ProcessorKind` values (used to validate manual overrides). +_PROCESSOR_KINDS: frozenset[str] = frozenset(get_args(ProcessorKind)) + +#: Soft-check scalar return annotations that contradict a stage-input role. +_SCALAR_RETURN_ANNOTATIONS: frozenset[Any] = frozenset({int, str, bool, float}) + +__all__ = [ + "ProcessorKind", + "ProcessorSpec", + "ProcessorValidationError", + "import_symbol", + "infer_kind", + "validate_processor", + "resolve_processor", + "dead_processor_hint", + "register_processor", +] + + +@dataclass(frozen=True) +class ProcessorSpec: + """Resolved and validated stage-input processor (path, kind, fn).""" + + path: str + kind: ProcessorKind + fn: Any + + +class ProcessorValidationError(Exception): + """A stage-input processor failed structural validation.""" + + def __init__( + self, + *, + stage_id: int | None = None, + path: str | None = None, + kind: ProcessorKind | None = None, + rule: str | None = None, + message: str | None = None, + ) -> None: + self.stage_id = stage_id + self.path = path + self.kind = kind + self.rule = rule + self.message = message + super().__init__(self._format()) + + def _format(self) -> str: + detail = ( + ", ".join( + f"{name}={value}" + for name, value in ( + ("stage_id", self.stage_id), + ("path", repr(self.path) if self.path else None), + ("kind", repr(self.kind) if self.kind else None), + ("rule", self.rule), + ) + if value is not None + ) + or "processor validation failed" + ) + return f"{detail}: {self.message}" if self.message else detail + + +# --------------------------------------------------------------------------- +# Small path / config helpers +# --------------------------------------------------------------------------- + + +def _fn_name(path: str) -> str: + return path.rsplit(".", 1)[-1] + + +def _module_part(path: str) -> str: + return path.rsplit(".", 1)[0] + + +def _stage_id_of(stage_config: Any) -> int | None: + """Best-effort stage id extraction from an arbitrary config object.""" + sid = getattr(stage_config, "stage_id", None) + if sid is None: + return None + try: + return int(sid) + except (TypeError, ValueError): + return None + + +# --------------------------------------------------------------------------- +# Symbol resolution +# --------------------------------------------------------------------------- + + +def import_symbol(path: str) -> Any: + """Import a dotted ``module.attr`` symbol; missing module/attr propagate.""" + if not isinstance(path, str) or not path: + raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") + module_path, attr_name = path.rsplit(".", 1) + module = importlib.import_module(module_path) + return getattr(module, attr_name) + + +# --------------------------------------------------------------------------- +# Kind inference (name / suffix driven) +# --------------------------------------------------------------------------- + +#: Manual kind overrides registered via :func:`register_processor`. +_KIND_OVERRIDES: dict[str, ProcessorKind] = {} + + +def register_processor(path: str, kind: ProcessorKind) -> None: + """Register a manual kind override (escape hatch for non-suffix names).""" + if not isinstance(path, str) or not path: + raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") + if kind not in _PROCESSOR_KINDS: + raise ValueError(f"unknown processor kind {kind!r}") + _KIND_OVERRIDES[path] = kind + + +def infer_kind(fn: Any, *, path: str) -> ProcessorKind: + """Infer the processor kind from the *suffix* of ``path`` (name-driven).""" + del fn # inference is name-based; kept for a symmetric API + if not isinstance(path, str) or not path or "." not in path: + raise ValueError(f"processor path must be a non-empty dotted string, got {path!r}") + + override = _KIND_OVERRIDES.get(path) + if override is not None: + return override + + fn_name = _fn_name(path) + + if _MOSS_MULTI_SOURCE_MODULE_HINT in _module_part(path) and fn_name in _MOSS_MULTI_SOURCE_NAMES: + return "legacy_multi_source" + + if fn_name.endswith(_TOKEN_ONLY_SUFFIX): + return "placeholder_prompt_builder" + if fn_name.endswith(_FULL_PAYLOAD_SUFFIXES): + return "producer_full_payload" + if fn_name.endswith(_ASYNC_CHUNK_SUFFIX): + return "producer_async_chunk" + for suffix in _DIFFUSION_SUFFIXES: + if fn_name.endswith(suffix): + return "diffusion_input_builder" + return "legacy_orchestrator_builder" + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +def _warn(path: str, kind: ProcessorKind | None, rule: str, message: str) -> None: + warnings.warn( + f"stage-input processor {path!r} (kind={kind!r}, rule={rule}): {message}", RuntimeWarning, stacklevel=3 + ) + + +def _fail( + *, + path: str, + kind: ProcessorKind | None, + rule: str, + stage_config: Any, + message: str, +) -> NoReturn: + raise ProcessorValidationError( + stage_id=_stage_id_of(stage_config), + path=path, + kind=kind, + rule=rule, + message=message, + ) + + +def _check_orchestrator_first_param( + signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """First param must be ``source_outputs`` or the signature must carry ``ctx``.""" + params = signature.parameters + names = [ + name + for name, param in params.items() + if param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) + ] + first = names[0] if names else None + if first != "source_outputs" and "ctx" not in params: + _warn( + path, + kind, + "orchestrator_first_param", + f"expected first parameter 'source_outputs' or a 'ctx' parameter, got first={first!r}", + ) + + +def _has_var_keyword(params: Mapping[str, inspect.Parameter]) -> bool: + """Whether the signature accepts arbitrary keywords (``**kwargs``).""" + return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()) + + +def _check_is_finished( + params: Mapping[str, inspect.Parameter], + *, + path: str, + kind: ProcessorKind, + stage_config: Any, + required: bool, +) -> None: + """Validate the ``is_finished`` kwarg contract (hard only for async-chunk).""" + if _has_var_keyword(params): + return # **kwargs accepts is_finished + if "is_finished" not in params: + if required: + _fail( + path=path, + kind=kind, + rule="is_finished_required", + stage_config=stage_config, + message="async-chunk producer must declare 'is_finished' (the scheduler always passes it)", + ) + _warn( + path, + kind, + "is_finished_optional", + "full-payload producer does not declare is_finished; worker retries without it (best-effort)", + ) + return + is_finished_param = params["is_finished"] + if is_finished_param.kind == inspect.Parameter.POSITIONAL_ONLY: + if required: + _fail( + path=path, + kind=kind, + rule="is_finished_positional_only", + stage_config=stage_config, + message="async-chunk producer declares 'is_finished' positional-only; " + "it must be keyword-only or have a default", + ) + _warn( + path, + kind, + "is_finished_positional_only", + "full-payload producer declares is_finished positional-only; worker retries without it", + ) + return + if required and ( + is_finished_param.kind != inspect.Parameter.KEYWORD_ONLY + and is_finished_param.default is inspect.Parameter.empty + ): + _warn( + path, + kind, + "is_finished_default", + "async-chunk producer declares a required positional-or-keyword 'is_finished'; " + "prefer keyword-only or a default", + ) + + +#: Sentinel value for :func:`_check_producer_keyword_bind` probes (the value +#: itself is never inspected or passed to the processor). +_BIND_PLACEHOLDER: Any = object() + + +def _check_producer_keyword_bind( + signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, + stage_config: Any, + required: tuple[str, ...], + what: str, +) -> None: + """Bind the worker keyword call to *signature* via ``inspect.Signature.bind``. + + Producers that declare the right names but reject the keyword call shape + fail at configuration time instead of when the worker passes + ``pooling_output=`` / ``multimodal_output=``. + """ + try: + signature.bind(**{name: _BIND_PLACEHOLDER for name in required}) + except TypeError as exc: + _fail( + path=path, + kind=kind, + rule="producer_kwargs", + stage_config=stage_config, + message=(f"{what} producer does not accept the worker keyword call ({', '.join(required)}): {exc}"), + ) + + +def _check_full_payload( + signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """Full-payload producer: keyword bind, then soft is_finished.""" + params = signature.parameters + required: tuple[str, ...] = ("transfer_manager", "pooling_output", "request") + is_finished_param = params.get("is_finished") + if is_finished_param is not None and is_finished_param.kind in ( + inspect.Parameter.KEYWORD_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + required = required + ("is_finished",) + _check_producer_keyword_bind( + signature, + path=path, + kind=kind, + stage_config=stage_config, + required=required, + what="full-payload", + ) + _check_is_finished(params, path=path, kind=kind, stage_config=stage_config, required=False) + + +def _check_async_chunk( + signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """Async-chunk producer: hard is_finished, then keyword bind.""" + params = signature.parameters + _check_is_finished(params, path=path, kind=kind, stage_config=stage_config, required=True) + _check_producer_keyword_bind( + signature, + path=path, + kind=kind, + stage_config=stage_config, + required=("transfer_manager", "multimodal_output", "request", "is_finished"), + what="async-chunk", + ) + + +def _check_legacy_multi_source( + _signature: inspect.Signature, + *, + path: str, + kind: ProcessorKind, + stage_config: Any, +) -> None: + """Legacy multi-source processors are allowlisted to ``moss_tts.*``.""" + if _MOSS_MULTI_SOURCE_MODULE_HINT not in _module_part(path): + _fail( + path=path, + kind=kind, + rule="legacy_multi_source_allowlist", + stage_config=stage_config, + message=f"legacy multi-source processors are allowlisted to moss_tts.*, got {path!r}", + ) + + +#: Per-kind structural checks run in order by :func:`validate_processor`. +#: ``legacy_orchestrator_builder`` is deliberately check-free: the orchestrator +#: normalizes legacy shapes through ``wrap_orchestrator_processor``. +_VALIDATORS: dict[ProcessorKind, tuple[Callable[..., None], ...]] = { + "placeholder_prompt_builder": (_check_orchestrator_first_param,), + "diffusion_input_builder": (_check_orchestrator_first_param,), + "producer_full_payload": (_check_full_payload,), + "producer_async_chunk": (_check_async_chunk,), + "legacy_multi_source": (_check_legacy_multi_source,), + "legacy_orchestrator_builder": (), +} + + +def validate_processor(fn: Any, *, kind: ProcessorKind, path: str, stage_config: Any = None) -> None: + """Structurally validate *fn* against the expected *kind* (signature only). + + Hard violations raise; soft mismatches warn; checks come from ``_VALIDATORS``. + """ + try: + signature = inspect.signature(fn) + except (TypeError, ValueError) as exc: + _fail( + path=path, + kind=kind, + rule="inspectable_signature", + stage_config=stage_config, + message=f"cannot inspect signature: {exc}", + ) + + validators = _VALIDATORS.get(kind) + if validators is None: # pragma: no cover - defensive against future literal additions + raise ValueError(f"unknown processor kind {kind!r}") + for check in validators: + check(signature, path=path, kind=kind, stage_config=stage_config) + + # Soft return-annotation check: mismatches warn, never fail. + ann = signature.return_annotation + if ann is not inspect.Signature.empty and not isinstance(ann, str) and ann in _SCALAR_RETURN_ANNOTATIONS: + _warn( + path, + kind, + "return_annotation", + f"return annotation {ann!r} is a scalar, unlikely a stage-input processor result", + ) + + +# --------------------------------------------------------------------------- +# Resolution entry point +# --------------------------------------------------------------------------- + + +def resolve_processor( + path: str, + *, + expected_kind: ProcessorKind | None = None, + stage_config: Any = None, +) -> ProcessorSpec: + """Resolve *path* to a validated :class:`ProcessorSpec`. + + Flow: ``import_symbol`` -> ``infer_kind`` -> ``validate_processor`` -> + optional ``expected_kind`` check; ``spec.fn`` matches the legacy lookup. + """ + fn = import_symbol(path) + if not callable(fn): + _fail( + path=path, + kind=expected_kind, + rule="callable", + stage_config=stage_config, + message=f"resolved symbol is not callable: {fn!r}", + ) + + kind = infer_kind(fn, path=path) + validate_processor(fn, kind=kind, path=path, stage_config=stage_config) + + if expected_kind is not None and kind != expected_kind: + _fail( + path=path, + kind=kind, + rule="expected_kind", + stage_config=stage_config, + message=f"inferred kind {kind!r} does not match expected_kind {expected_kind!r}", + ) + + return ProcessorSpec(path=path, kind=kind, fn=fn) + + +# --------------------------------------------------------------------------- +# Dead-processor hint (gate three-state) +# --------------------------------------------------------------------------- + + +def dead_processor_hint( + kind: ProcessorKind, + *, + async_chunk: bool, + downstream_receives_async_chunks: bool, + has_sync: bool, + custom_is_sync: bool = False, +) -> bool: + """Whether an orchestrator-side custom hook is never invoked (gate three-state). + + Mirrors ``engine/orchestrator._route_output``: async-chunk transitions skip + forward processing when the downstream stage receives chunks; otherwise a + differing ``sync_process_input_func`` overrides the custom hook. Pure + decision helper (warn-first; the orchestrator does not enforce it). + """ + if kind in _PRODUCER_KINDS: + return False + if async_chunk: + return bool(downstream_receives_async_chunks) + return bool(has_sync) and not custom_is_sync diff --git a/vllm_omni/model_executor/stage_input_processors/audex.py b/vllm_omni/model_executor/stage_input_processors/audex.py index 011ec904306..979d240dffe 100644 --- a/vllm_omni/model_executor/stage_input_processors/audex.py +++ b/vllm_omni/model_executor/stage_input_processors/audex.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 """Audex thinker → code2wav stage input processors. The thinker is a plain token-autoregressive LM whose sampled stream contains @@ -21,6 +21,7 @@ from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt logger = init_logger(__name__) @@ -43,14 +44,6 @@ _STATE_KEY = "_audex_async_state" -def _ensure_list(value: Any) -> list[int]: - if value is None: - return [] - if isinstance(value, torch.Tensor): - return value.reshape(-1).tolist() - return list(value) - - def _connector_extra(transfer_manager: Any) -> dict[str, Any]: connector = getattr(transfer_manager, "connector", None) raw_cfg = getattr(connector, "config", {}) or {} @@ -85,7 +78,7 @@ def _finished_request_codes( cfg = _connector_extra(transfer_manager) codec_offset = int(cfg.get(offset_key, default_offset)) codec_size = int(cfg.get(size_key, default_size)) - output_token_ids = _ensure_list(getattr(request, "output_token_ids", [])) + output_token_ids = _common.ensure_list_flatten(getattr(request, "output_token_ids", [])) codes = _codec_frames(output_token_ids, codec_offset, codec_size) request_id = getattr(request, "external_req_id", None) or getattr(request, "request_id", None) return codes, request_id @@ -159,7 +152,7 @@ def _extract_new_codes( codec_size: int, ) -> list[int]: """Return codec frame ids from tokens sampled since the last call.""" - output_token_ids = _ensure_list(getattr(request, "output_token_ids", [])) + output_token_ids = _common.ensure_list_flatten(getattr(request, "output_token_ids", [])) new_tokens = output_token_ids[int(state.get("seen_len", 0)) :] state["seen_len"] = len(output_token_ids) return _codec_frames(new_tokens, codec_offset, codec_size) @@ -348,8 +341,8 @@ def thinker2code2wav_token_only( if not source_output.finished: continue output = source_output.outputs[0] - prompt_ids = _ensure_list(source_output.prompt_token_ids) - raw_output_ids = _ensure_list(output.cumulative_token_ids) + prompt_ids = _common.ensure_list_flatten(source_output.prompt_token_ids) + raw_output_ids = _common.ensure_list_flatten(output.cumulative_token_ids) if raw_output_ids[: len(prompt_ids)] == prompt_ids: raw_output_ids = raw_output_ids[len(prompt_ids) :] engine_inputs.append(OmniTokensPrompt(prompt_token_ids=raw_output_ids)) diff --git a/vllm_omni/model_executor/stage_input_processors/audio8_tts.py b/vllm_omni/model_executor/stage_input_processors/audio8_tts.py index 4b453c13ab7..b286da0ee7d 100644 --- a/vllm_omni/model_executor/stage_input_processors/audio8_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/audio8_tts.py @@ -12,6 +12,7 @@ MetaStruct, OmniPayloadStruct, ) +from vllm_omni.model_executor.stage_input_processors import _common #: Chunking defaults; ~21.5 codec frames per second, so 25 frames ~= 1.16 s. DEFAULT_CHUNK_FRAMES = 25 @@ -49,28 +50,16 @@ def extract_last_frame(multimodal_output: Mapping[str, Any]) -> torch.Tensor | N Prefill steps emit an all-zero placeholder frame, which must not be forwarded to the decoder. - """ - audio_codes = multimodal_output.get("audio_codes") - if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: - return None - if audio_codes.ndim == 1: - return audio_codes.to(device="cpu", dtype=torch.long).reshape(-1) - if audio_codes.ndim != 2: - raise ValueError(f"Invalid audio_codes shape for Audio8 TTS async_chunk: {tuple(audio_codes.shape)}") - frame = audio_codes[-1] - if frame.numel() == 0: - return None - valid = multimodal_output.get("audio_code_valid") - if isinstance(valid, torch.Tensor) and valid.numel() > 0: - is_valid = bool(valid.reshape(-1)[-1].item()) - elif valid is not None: - is_valid = bool(valid) - else: - is_valid = bool(frame.any().item()) - if not is_valid: - return None - return frame.to(device="cpu", dtype=torch.long).reshape(-1) + Delegates to the canonical ``_common.extract_last_codec_frame``. + """ + return _common.extract_last_codec_frame( + multimodal_output, + key_path=("audio_codes",), + validate="valid_mask", + to_cpu=True, + to_long=True, + ) def slow_ar_to_codec_decoder_async_chunk( diff --git a/vllm_omni/model_executor/stage_input_processors/auk.py b/vllm_omni/model_executor/stage_input_processors/auk.py index 7c417aee171..17c196543a1 100644 --- a/vllm_omni/model_executor/stage_input_processors/auk.py +++ b/vllm_omni/model_executor/stage_input_processors/auk.py @@ -76,7 +76,14 @@ def _extract_text_cond(ar_output: Any) -> Any: def _to_cpu_tensor(text_cond: Any) -> torch.Tensor: - """Coerce the emitted condition to a single 2-D CPU tensor.""" + """Coerce the emitted condition to a CPU tensor. + + Restores the AuK-wide semantics that the shared ``_common.to_cpu_tensor`` + (glm_tts-narrow) does not provide: a list of segments is recursively + converted and concatenated along ``dim=0``, an empty list raises + ``ValueError``, and numpy / array-like inputs are converted via + ``torch.as_tensor``. The 2-D shape check lives at the call site. + """ if isinstance(text_cond, list): if not text_cond: raise ValueError("AuK encoder emitted an empty text condition list") @@ -87,10 +94,7 @@ def _to_cpu_tensor(text_cond: Any) -> torch.Tensor: else: # numpy array or any array-like the connector round-tripped. tensor = torch.as_tensor(text_cond) - tensor = tensor.detach().cpu() - if tensor.ndim != 2: - raise ValueError(f"AuK text condition must be [tokens, hidden]; got shape {tuple(tensor.shape)}") - return tensor + return tensor.detach().cpu() def encoder2dit( @@ -119,6 +123,10 @@ def encoder2dit( "stage 1 cannot run without the fused thinker hidden states" ) prompt_embeds = _to_cpu_tensor(text_cond) + if prompt_embeds is None: + raise ValueError("AuK encoder emitted no text condition tensor for stage 1") + if prompt_embeds.ndim != 2: + raise ValueError(f"AuK text condition must be [tokens, hidden]; got shape {tuple(prompt_embeds.shape)}") original = _as_dict(prompt) original_knobs = original.get("additional_information") or {} diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index e8123fe8751..cc1d8652e44 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -18,6 +18,7 @@ from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.cosyvoice3.utils import unpad_prompt_conditioning +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) @@ -36,38 +37,6 @@ def _build_prompt_embed_struct(prompt_payload: dict[str, Any]) -> EmbeddingsStru ) -def _ensure_list(x: Any) -> list[Any]: - if hasattr(x, "_x"): - return list(x._x) - if isinstance(x, list): - return list(x) - if isinstance(x, tuple): - return list(x) - if x is None: - return [] - try: - return list(x) - except TypeError: - return [x] - - -def _to_token_id_list(value: Any) -> list[int]: - if value is None: - return [] - if isinstance(value, torch.Tensor): - value = value.detach().to("cpu").reshape(-1).tolist() - token_ids: list[int] = [] - for item in _ensure_list(value): - if isinstance(item, torch.Tensor): - token_ids.extend(_to_token_id_list(item)) - continue - if isinstance(item, (list, tuple)): - token_ids.extend(_to_token_id_list(item)) - continue - token_ids.append(int(item)) - return token_ids - - def _strip_prompt_prefix(output_ids: list[Any], prefix_ids: list[Any]) -> list[Any]: if prefix_ids and len(output_ids) >= len(prefix_ids) and output_ids[: len(prefix_ids)] == prefix_ids: return output_ids[len(prefix_ids) :] @@ -80,17 +49,7 @@ def _prompt_speech_token_ids(multi_modal_data: dict[str, Any]) -> list[int]: embed = multi_modal_data.get("embed") if isinstance(embed, dict): speech_token = embed.get("speech_token") - return _to_token_id_list(speech_token) - - -def _to_cpu_tensor(x: Any) -> torch.Tensor | None: - if isinstance(x, list): - if not x: - return None - x = x[0] - if isinstance(x, torch.Tensor): - return x.detach().cpu() - return None + return _common.to_token_id_list(speech_token, recursive=True) def talker2code2wav_async_chunk( @@ -128,7 +87,7 @@ def talker2code2wav_async_chunk( prompt_payload = {} cond_keys = ("speech_token", "speech_feat", "embedding", "speech_token_len") for key in cond_keys: - value = _to_cpu_tensor(info_embed.get(key)) + value = _common.to_cpu_tensor(info_embed.get(key)) if value is not None: prompt_payload[key] = value if isinstance(multimodal_output, Mapping): @@ -138,7 +97,7 @@ def talker2code2wav_async_chunk( for key in cond_keys: if key in prompt_payload: continue - value = _to_cpu_tensor(mm_embed.get(key)) + value = _common.to_cpu_tensor(mm_embed.get(key)) if value is not None: prompt_payload[key] = value # Drop any right-padding carried from batched talker emission so @@ -186,7 +145,7 @@ def talker2code2wav_async_chunk( return None with nullcontext(): - output_token_ids = _ensure_list(getattr(request, "output_token_ids", [])) + output_token_ids = _common.ensure_list(getattr(request, "output_token_ids", [])) seen_len = int(state.get("seen_len", 0)) new_tokens = output_token_ids[seen_len:] if seen_len < len(output_token_ids) else [] state["seen_len"] = len(output_token_ids) @@ -318,8 +277,8 @@ def text2flow_token_only( if not source_output.finished: continue output = source_output.outputs[0] - prefix_ids = _ensure_list(source_output.prompt_token_ids) - raw_output_ids = _ensure_list(output.cumulative_token_ids) + prefix_ids = _common.ensure_list(source_output.prompt_token_ids) + raw_output_ids = _common.ensure_list(output.cumulative_token_ids) output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids) multi_modal_data = output.multimodal_output if multi_modal_data is None: diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index cfcb277eabb..977797a0844 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py @@ -13,6 +13,7 @@ MetaStruct, OmniPayloadStruct, ) +from vllm_omni.model_executor.stage_input_processors import _common def _get_connector_extra(transfer_manager: Any) -> dict[str, Any]: @@ -95,28 +96,6 @@ def _select_backlog_chunk_size( return backlog_chunk_size -def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | None: - """Extract the last frame of audio codes from the multimodal output.""" - audio_codes = multimodal_output.get("audio_codes") - if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: - return None - if audio_codes.ndim == 2: - frame = audio_codes[-1] - valid = multimodal_output.get("audio_code_valid") - if isinstance(valid, torch.Tensor) and valid.numel() > 0: - is_valid = bool(valid.reshape(-1)[-1].item()) - elif valid is not None: - is_valid = bool(valid) - else: - is_valid = bool(frame.any().item()) - if frame.numel() == 0 or not is_valid: - return None - return frame.to(device="cpu", dtype=torch.long).reshape(-1) - if audio_codes.ndim == 1: - return audio_codes.to(device="cpu", dtype=torch.long).reshape(-1) - raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}") - - def slow_ar_to_dac_decoder_async_chunk( transfer_manager: Any, multimodal_output: dict[str, Any] | None, @@ -134,7 +113,13 @@ def slow_ar_to_dac_decoder_async_chunk( cfg = _get_connector_extra(transfer_manager) if isinstance(multimodal_output, Mapping): - frame = _extract_last_frame(multimodal_output) + frame = _common.extract_last_codec_frame( + multimodal_output, + key_path=("audio_codes",), + validate="valid_mask", + to_cpu=True, + to_long=True, + ) if frame is not None: transfer_manager.code_prompt_token_ids[request_id].append(frame.detach()) elif not finished: diff --git a/vllm_omni/model_executor/stage_input_processors/glm_tts.py b/vllm_omni/model_executor/stage_input_processors/glm_tts.py index af037aee77c..d3cc21dabec 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/glm_tts.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 """Stage input processor for GLM-TTS: AR → DiT Pipeline. Supports both sync (non-streaming) and async_chunk (streaming) modes. @@ -18,6 +18,7 @@ OmniPayloadStruct, ) from vllm_omni.engine.serialization import deserialize_additional_information +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) @@ -37,7 +38,7 @@ def _copy_voice_clone_payload( prompt = src.get("prompt_token") if prompt is not None: if to_cpu: - prompt = _to_cpu_tensor(prompt) + prompt = _common.to_cpu_tensor(prompt) if prompt is not None: dst["prompt_speech_token"] = prompt @@ -47,7 +48,7 @@ def _copy_voice_clone_payload( val = src.get(key) if val is not None: if to_cpu: - val = _to_cpu_tensor(val) + val = _common.to_cpu_tensor(val) if val is not None: dst[key] = val @@ -169,18 +170,6 @@ def _extract_last_speech_token(pooling_output: dict[str, Any]) -> int | None: return token_val -def _to_cpu_tensor(value: Any) -> torch.Tensor | None: - """Convert value to CPU tensor if possible.""" - if isinstance(value, torch.Tensor): - return value.detach().cpu() - if isinstance(value, list): - if not value: - return None - if isinstance(value[0], torch.Tensor): - return value[0].detach().cpu() - return None - - # --------------------------------------------------------------------------- # Async streaming processor: emit speech token chunks as AR produces them # --------------------------------------------------------------------------- diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py index 53f60ba4ba4..dbb5a10afd4 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py @@ -29,6 +29,7 @@ OmniPayload, OmniPayloadStruct, ) +from vllm_omni.model_executor.stage_input_processors import _common __all__ = ["talker2code2wav", "talker2code2wav_async_chunk"] @@ -48,30 +49,6 @@ _NUM_REAL_CODES = _AUDIO_STREAM_BOS_ID # codes in [0, 1023] are real -def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: - """Reverse the delay-pattern shift, mirroring upstream's - ``boson_multimodal.model.higgs_audio.utils.revert_delay_pattern``. - - Input shape: ``[num_codebooks, seq_len + num_codebooks - 1]``. - Output shape: ``[num_codebooks, seq_len]``. - - For each codebook ``i``, slice rows ``[i:i+1]`` from columns - ``[i : seq_len + i]`` to remove ``i`` leading BOS pads and ``Q-1-i`` - trailing PAD entries. - """ - if audio_codes_qt.ndim != 2: - raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}") - q, t = audio_codes_qt.shape - if t < q: - # Not enough frames to revert delay pattern; return as-is. - return audio_codes_qt - seq_len = t - q + 1 - out_l = [] - for i in range(q): - out_l.append(audio_codes_qt[i : i + 1, i : seq_len + i]) - return torch.cat(out_l, dim=0) - - def talker2code2wav( source_outputs: list[Any], prompt: Any = None, @@ -121,7 +98,7 @@ def talker2code2wav( # the upstream ``HiggsAudioServeEngine.generate`` decode line: # vq_code = revert_delay_pattern(out).clip(0, real-1)[:, 1:-1] codes_qt = audio_codes.transpose(0, 1).contiguous().cpu() - codes_qt = _revert_delay_pattern(codes_qt) + codes_qt = _common.revert_delay_pattern(codes_qt, allow_short=True) codes_qt = codes_qt.clamp_(min=0, max=_NUM_REAL_CODES - 1) if codes_qt.shape[-1] >= 3: codes_qt = codes_qt[:, 1:-1] diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py index 8a7c66121aa..a4322601fb2 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py @@ -34,6 +34,7 @@ OmniPayloadStruct, ) from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common __all__ = ["talker2code2wav", "talker2code2wav_async_chunk"] @@ -60,40 +61,25 @@ def _empty_code2wav_prompt() -> Any: def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: - """Reverse the MusicGen-style delay pattern. + """Reverse the MusicGen-style delay pattern (higgs-v3 strict variant). - Input shape: [num_codebooks, seq_len + num_codebooks - 1]. - Output shape: [num_codebooks, seq_len]. - - For each codebook i, extract delayed[i, i : i + seq_len] to remove - the i leading BOC pads and Q-1-i trailing EOC entries. + Delegates to the canonical ``_common.revert_delay_pattern`` with the + strict codebook-count and short-input checks. """ - if audio_codes_qt.ndim != 2: - raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}") - q, t = audio_codes_qt.shape - if q != _NUM_CODEBOOKS: - raise ValueError(f"Expected exactly {_NUM_CODEBOOKS} codebook rows, got {q}. Input shape: [{q}, {t}]") - if t < q: - raise ValueError(f"Not enough frames to revert delay pattern: T={t} < Q={q}") - seq_len = t - q + 1 - out_l = [] - for i in range(q): - out_l.append(audio_codes_qt[i : i + 1, i : seq_len + i]) - return torch.cat(out_l, dim=0) + return _common.revert_delay_pattern( + audio_codes_qt, + expected_codebooks=_NUM_CODEBOOKS, + allow_short=False, + ) def _filter_real_code_frames(audio_codes_qt: torch.Tensor) -> torch.Tensor: """Keep only frames where ALL codebook values are in [0, 1023]. - Input shape: [num_codebooks, num_frames]. - Called AFTER delay pattern reversal. + Delegates to the canonical ``_common.filter_real_code_frames`` with the + higgs-v3 ``[num_codebooks, num_frames]`` layout. """ - if audio_codes_qt.numel() == 0: - return audio_codes_qt - # Transpose to [num_frames, num_codebooks] for per-frame filtering - frames = audio_codes_qt.t() - valid = (frames >= 0).all(dim=1) & (frames < _NUM_REAL_CODES).all(dim=1) - return frames[valid].t().contiguous() + return _common.filter_real_code_frames(audio_codes_qt, num_real_codes=_NUM_REAL_CODES, layout="codebooks_first") def _whole_utterance_codec_codes(audio_rows: torch.Tensor) -> torch.Tensor: @@ -342,7 +328,7 @@ def _flush_async_chunk( codes_qt = torch.tensor(window_rows, dtype=torch.long).t().contiguous() try: - de_delayed = _revert_delay_pattern(codes_qt) + de_delayed = _common.revert_delay_pattern(codes_qt, expected_codebooks=_NUM_CODEBOOKS, allow_short=False) except ValueError: logger.warning( "async_chunk: insufficient frames for delay pattern reversal " diff --git a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py index da498811429..f85c440eecb 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py @@ -14,6 +14,7 @@ from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common from vllm_omni.model_executor.stage_input_processors.ming_talker import ( build_ming_talker_inputs, ) @@ -208,14 +209,11 @@ def _resolve_image_patch_token_id(stage: Any) -> int: def _ensure_list(x) -> list[int]: - """Convert ConstantList / tensor-like to plain list.""" - if hasattr(x, "_x"): - return list(x._x) - if isinstance(x, list): - return x - if hasattr(x, "tolist"): - return x.tolist() - return list(x) + """Convert ConstantList / tensor-like to plain list. + + Delegates to the canonical ``_common.ensure_list_strict``. + """ + return _common.ensure_list_strict(x) def _slice_patch_hidden( @@ -246,7 +244,7 @@ def _slice_patch_hidden( logger.warning("[thinker2imagegen] %s: missing final_hidden_states (keys=%s)", tag, list(mm_out.keys())) return None - prompt_ids = _ensure_list(thinker_output.prompt_token_ids) + prompt_ids = _common.ensure_list_strict(thinker_output.prompt_token_ids) prompt_ids_t = torch.tensor(prompt_ids, dtype=torch.long, device=full_hidden.device) patch_indices = (prompt_ids_t == image_patch_token_id).nonzero(as_tuple=False).squeeze(-1) total_patches = int(patch_indices.numel()) diff --git a/vllm_omni/model_executor/stage_input_processors/personaplex.py b/vllm_omni/model_executor/stage_input_processors/personaplex.py index 71f33ba3387..8a4efdefc00 100644 --- a/vllm_omni/model_executor/stage_input_processors/personaplex.py +++ b/vllm_omni/model_executor/stage_input_processors/personaplex.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 """Talker -> Code2Wav input processors for PersonaPlex. The talker (stage 0) emits, per frame, the ``dep_q`` depformer audio codes under @@ -130,8 +130,6 @@ def _codes_from(src: Any) -> torch.Tensor | None: getattr(request, "additional_information_cpu", None), pooling_output, # Temporary compatibility shim for the two producer call contracts. - # Remove after https://github.com/vllm-project/vllm-omni/issues/4872 - # provides validated full-payload and async-chunk processor protocols. kwargs.get("multimodal_output"), ): audio = _codes_from(source) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 42a9b19e3d7..bb0bf5ff707 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import logging import torch @@ -12,6 +15,7 @@ to_dict, ) from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors import _common logger = logging.getLogger(__name__) @@ -241,19 +245,13 @@ def thinker2talker_full_payload( ) return None - def _ensure_list(x): - if x is None: - return [] - if hasattr(x, "_x"): - # vLLM wraps cached token-id lists in ConstantList-like objects. - return list(x._x) - if isinstance(x, list): - return list(x) - return list(x) - - prompt_token_ids = _ensure_list(getattr(request, "prompt_token_ids", None)) - output_token_ids = _ensure_list(getattr(request, "output_token_ids", None)) - all_token_ids = _ensure_list(getattr(request, "all_token_ids", None) or []) + # Use ``ensure_list`` (``None`` -> ``[]``), not ``ensure_list_unchanged`` + # (``None`` -> ``None``): the downstream arithmetic below indexes and + # takes ``len(...)`` of these lists, so a ``None`` would raise instead of + # degrading gracefully. This matches the pre-refactor behaviour. + prompt_token_ids = _common.ensure_list(getattr(request, "prompt_token_ids", None)) + output_token_ids = _common.ensure_list(getattr(request, "output_token_ids", None)) + all_token_ids = _common.ensure_list(getattr(request, "all_token_ids", None)) if not all_token_ids: all_token_ids = list(prompt_token_ids) + list(output_token_ids) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 3ff4b5e2813..b64a48a68f0 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -24,6 +24,8 @@ from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.qwen3_omni.tool_speech import TOOL_CALL_START_TOKEN_ID +from vllm_omni.model_executor.stage_input_processors import _common +from vllm_omni.model_executor.stage_input_processors._dispatch import OrchestratorInputContext from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import compute_ramp_emit, parse_chunk_ramp from vllm_omni.model_executor.stage_input_processors.tts_utils import ( extract_language_from_prompt, @@ -48,6 +50,45 @@ _QWEN3_CODEC_EOS_TOKEN_ID = 4198 +def _assert_codec_token_ids_consistent() -> None: + """Fail-fast guard: codec token ids must not drift. + + These ids are model-specific special tokens. The HF config classes are + the single source of truth at runtime -- the talker reads + ``codec_pad_id`` / ``codec_bos_id`` / ``codec_eos_token_id`` off + ``talker_config``. The processor keeps its own module-top constants + (used by ``talker2code2wav_*`` bookkeeping) and verifies them against + both the in-repo Qwen3-TTS config and the transformers + ``Qwen3OmniMoeTalkerConfig`` (the config qwen3-omni actually runs with). + """ + from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeTalkerConfig, + ) + + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( + Qwen3TTSTalkerConfig, + ) + + expected = ( + _QWEN3_CODEC_PAD_TOKEN_ID, + _QWEN3_CODEC_BOS_TOKEN_ID, + _QWEN3_CODEC_EOS_TOKEN_ID, + ) + for name, cfg in ( + ("qwen3_tts", Qwen3TTSTalkerConfig()), + ("qwen3_omni_moe", Qwen3OmniMoeTalkerConfig()), + ): + actual = (cfg.codec_pad_id, cfg.codec_bos_id, cfg.codec_eos_token_id) + assert actual == expected, ( + f"{name} codec token ids {actual} drift from the stage_input_processors.qwen3_omni constants {expected}" + ) + + +# The codec-token consistency guard is intentionally **not** run at import time +# (it pulls in transformers / the in-repo qwen3-tts config and must not break +# light imports). It is invoked explicitly at model-load startup +# (qwen3_omni_moe_talker) so the processor constants stay in sync with the HF +# configs the model actually runs with. def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None: """Fetch layer tensor with tolerant key lookup (str/int/semantic).""" if not isinstance(layers, dict): @@ -88,7 +129,7 @@ def _get_accept_hidden_layer_index(transfer_manager: Any) -> int: def _is_final_token(request: Any) -> bool | None: """Whether the token sampled last ends the Thinker's turn (``None``: unknown).""" - output_token_ids = _ensure_list(getattr(request, "output_token_ids", None) or []) + output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", None) or []) count = getattr(request, "output_token_count", None) count = len(output_token_ids) if count is None else count token = getattr(request, "last_output_token_id", None) @@ -138,7 +179,7 @@ def _sampled_token_is_tool_call(request: Any) -> bool: def _output_reaches_tool_call(request: Any) -> bool: """True once this step's Thinker sample is ```` or after it.""" - output_token_ids = _ensure_list(getattr(request, "output_token_ids", None) or []) + output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", None) or []) if TOOL_CALL_START_TOKEN_ID in output_token_ids: return True return _sampled_token_is_tool_call(request) @@ -160,57 +201,11 @@ def _sampled_token_embed(thinker_embed: Mapping[str, Any], request: Any) -> torc return sampled.detach().cpu() -def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | str = "cuda") -> int: - im_start_token_id = 151644 - system_token_id = 8948 - user_token_id = 872 - assistant_token_id = 77091 - - ids = info.get("ids", {}) - thinker_sequences = torch.tensor(ids["all"], dtype=torch.long, device=device).unsqueeze(0) # [1, T] - - input_ids = torch.tensor(ids["prompt"], dtype=torch.long, device=device).unsqueeze(0) # [1, T] - - im_start_indexes = torch.cat( - [ - torch.nonzero(input_ids[0] == im_start_token_id).squeeze(1), - torch.tensor([thinker_sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype), - ], - dim=0, - ) - - sum_user_len = 0 - assistant_len = 0 - for i in range(len(im_start_indexes) - 1): - s = int(im_start_indexes[i].item()) - e = int(im_start_indexes[i + 1].item()) - role = int(input_ids[0, s + 1].item()) - if role == system_token_id: - continue - elif role == user_token_id: - sum_user_len += e - s - elif role == assistant_token_id and i == len(im_start_indexes) - 2: - assistant_len += 9 # 3 + 4 + 1 + 1 - else: - pass - - return sum_user_len + assistant_len - - # ========================= # Common helpers # ========================= -def _ensure_list(x): - """Convert ConstantList / tensor-like to Python list.""" - if hasattr(x, "_x"): - return list(x._x) - elif not isinstance(x, list): - return x - return list(x) - - def _as_tensor_or_none(value: Any) -> torch.Tensor | None: if isinstance(value, torch.Tensor): return value.detach().cpu() @@ -299,7 +294,7 @@ def _construct_thinker2talker_streaming_input_async_chunk( request_id = request.external_req_id output_token_ids = request.output_token_ids # Convert ConstantList to regular list for OmniSerializer serialization - output_token_ids = _ensure_list(output_token_ids) + output_token_ids = _common.ensure_list_unchanged(output_token_ids) speaker = extract_speaker_from_request(request) language = extract_language_from_request(request) finished = torch.tensor(is_finished, dtype=torch.bool) @@ -329,8 +324,8 @@ def _construct_thinker2talker_streaming_input_async_chunk( embed=EmbeddingsStruct(prefill=emb_cpu), hidden_states=HiddenStatesStruct(output=hid_cpu), ids=IdsStruct( - all=_ensure_list(request.all_token_ids[-new_prompt_len - 1 :]), - prompt=_ensure_list(request.prompt_token_ids[-new_prompt_len:]), + all=_common.ensure_list_unchanged(request.all_token_ids[-new_prompt_len - 1 :]), + prompt=_common.ensure_list_unchanged(request.prompt_token_ids[-new_prompt_len:]), ), speaker=speaker, language=language, @@ -524,8 +519,8 @@ def _maybe_cpu(t: Any) -> torch.Tensor | None: return t.detach().cpu() if isinstance(t, torch.Tensor) else None if chunk_id == 0: - all_token_ids = _ensure_list(request.all_token_ids) - prompt_token_ids = _ensure_list(request.prompt_token_ids) + all_token_ids = _common.ensure_list_unchanged(request.all_token_ids) + prompt_token_ids = _common.ensure_list_unchanged(request.prompt_token_ids) payload = OmniPayloadStruct( embed=EmbeddingsStruct( prefill=thinker_emb.detach().cpu(), @@ -677,10 +672,10 @@ def thinker2talker_full_payload( ) return None - prompt_token_ids = _ensure_list(getattr(request, "prompt_token_ids", []) or []) - all_token_ids = _ensure_list(getattr(request, "all_token_ids", None) or []) + prompt_token_ids = _common.ensure_list_unchanged(getattr(request, "prompt_token_ids", []) or []) + all_token_ids = _common.ensure_list_unchanged(getattr(request, "all_token_ids", None) or []) if not all_token_ids: - output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or []) + output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", []) or []) all_token_ids = list(prompt_token_ids) + list(output_token_ids) # Drop the terminal stop-token row only when more than one row was @@ -728,32 +723,24 @@ def thinker2talker_full_payload( return payload -def thinker2talker_token_only( +def build_forward_placeholder( source_outputs: list[Any], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, - streaming_context: Any | None = None, + ctx: OrchestratorInputContext, ) -> list[OmniTokensPrompt]: - """Orchestrator-side placeholder builder for Stage-1 (Talker) when - ``async_chunk=False``. - - After the communication-layer refactor, this function only allocates a - placeholder ``prompt_token_ids`` of the correct length so the scheduler can - reserve KV-cache slots. It does **not** forward bulk tensors. + """Sync forward placeholder: one KV-sized ``OmniTokensPrompt`` per output. - Bulk talker conditioning is sent through the connector. Speaker and - language are also copied from the original prompt so they survive when - Stage-0 request metadata is unavailable to the connector payload. - - ``prompt`` / ``requires_multimodal_data`` are kept for call-site signature - compatibility with other orchestrator input processors; they are unused. + Mirrors the legacy ``thinker2talker_token_only`` path; speaker/language are + copied from the prompt so they survive without Stage-0 request metadata. + Length/packing delegated to ``_common.compute_placeholder_prompt_len(mode="full")`` + and ``_common.pack_placeholder_prompt``. """ talker_inputs: list[OmniTokensPrompt] = [] + streaming_context = ctx.streaming_context for i, thinker_output in enumerate(source_outputs): output = thinker_output.outputs[0] req_id = str(getattr(thinker_output, "request_id", f"idx-{i}")) - prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids) - output_ids = _ensure_list(output.cumulative_token_ids) + prompt_token_ids = _common.ensure_list_unchanged(thinker_output.prompt_token_ids) + output_ids = _common.ensure_list_unchanged(output.cumulative_token_ids) is_streaming_session = bool(getattr(streaming_context, "enabled", False)) if is_streaming_session: prompt_token_ids, output_ids = _get_streaming_talker_tokens( @@ -767,25 +754,80 @@ def thinker2talker_token_only( thinker_sequences = prompt_token_ids + output_ids thinker_input_ids = prompt_token_ids info_for_len: OmniPayload = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}} - prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu") + prompt_len = _common.compute_placeholder_prompt_len( + ids_or_prompt=info_for_len, + mode="full", + device="cpu", + ) # Keep this fallback until the connector reliably preserves voice metadata. additional_information = to_dict( OmniPayloadStruct( - speaker=extract_speaker_from_prompt(prompt, index=i), - language=extract_language_from_prompt(prompt, index=i), + speaker=extract_speaker_from_prompt(ctx.prompt, index=i), + language=extract_language_from_prompt(ctx.prompt, index=i), ) ) talker_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[0] * prompt_len, - additional_information=additional_information or None, - multi_modal_data=None, - mm_processor_kwargs=None, + _common.pack_placeholder_prompt( + prompt_len=prompt_len, + voice_metadata=additional_information or None, ) ) return talker_inputs +def build_prewarm_placeholder( + *, + stage0_prompt: Any, + ctx: OrchestratorInputContext, + downstream_stage_id: int, +) -> list[OmniTokensPrompt]: + """Best-effort async-chunk prewarm estimate (no ``source_outputs`` yet). + + Length from ``_common.compute_placeholder_prompt_len(mode="stage0_only")`` + (the same Qwen scan ``adapter.compute_talker_prompt_ids_length`` runs), so + prewarm matches sync-forward and the inline fallback; the connector fixup + path replaces it once the upstream chunk arrives. Voice metadata is not + forwarded (matching pre-existing prewarm behaviour); per-edge refinement + is out of scope. + """ + del ctx, downstream_stage_id # Only the stage-0 input estimate is used. + if hasattr(stage0_prompt, "prompt_token_ids"): + stage0_prompt = stage0_prompt.prompt_token_ids + prompt_len = max( + 1, + _common.compute_placeholder_prompt_len( + ids_or_prompt=stage0_prompt, + mode="stage0_only", + ), + ) + return [_common.pack_placeholder_prompt(prompt_len=prompt_len, voice_metadata=None)] + + +def thinker2talker_token_only( + source_outputs: list[Any], + prompt: OmniTokensPrompt | TextPrompt | None = None, + requires_multimodal_data: bool = False, + streaming_context: Any | None = None, +) -> list[OmniTokensPrompt]: + """Stage-1 sync placeholder (``async_chunk=False``): KV-sized slot only. + + Thin delegation to :func:`build_forward_placeholder`; legacy params are kept + for call-site compatibility and folded into an ``OrchestratorInputContext``. + """ + ctx = OrchestratorInputContext( + prompt=prompt, + requires_multimodal_data=requires_multimodal_data, + streaming_context=streaming_context, + ) + return build_forward_placeholder(source_outputs, ctx) + + +# Expose the dual-entry builders on the resolved ``*_token_only`` fn so the +# orchestrator's ``sync_process_input_func`` lookup finds them without a scan. +thinker2talker_token_only.build_forward_placeholder = build_forward_placeholder # type: ignore[attr-defined] +thinker2talker_token_only.build_prewarm_placeholder = build_prewarm_placeholder # type: ignore[attr-defined] + + # ========================= # Talker -> Code2Wav # ========================= @@ -805,9 +847,7 @@ def talker2code2wav_async_chunk( request: OmniEngineCoreRequest, is_finished: bool = False, ) -> OmniPayloadStruct | None: - """ - Multimodal output version. - """ + """Async-chunk producer: multimodal output version.""" request_id = request.external_req_id code_predictor_codes = None frame_valid = None @@ -970,7 +1010,7 @@ def talker2code2wav_full_payload( ) return None - output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or []) + output_token_ids = _common.ensure_list_unchanged(getattr(request, "output_token_ids", []) or []) raw_shape = tuple(code_predictor_codes.shape) code_predictor_codes, codec_stats = _extract_qwen3_full_payload_codec_rows( code_predictor_codes.to(torch.long), diff --git a/vllm_omni/model_executor/stage_input_processors/step_audio2.py b/vllm_omni/model_executor/stage_input_processors/step_audio2.py index c0c0932b27a..8424985cfce 100644 --- a/vllm_omni/model_executor/stage_input_processors/step_audio2.py +++ b/vllm_omni/model_executor/stage_input_processors/step_audio2.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage input processor for Step-Audio2: Thinker → Token2Wav transition.""" from typing import Any @@ -14,21 +15,11 @@ DEFAULT_STREAM_CONFIG, DEFAULT_TOKEN_CONFIG, ) +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) -def _ensure_list(x): - """Convert ConstantList / tensor / iterable to Python list.""" - if hasattr(x, "_x"): - return list(x._x) - elif isinstance(x, torch.Tensor): - return x.tolist() - elif not isinstance(x, list): - return list(x) if hasattr(x, "__iter__") else [x] - return list(x) - - # ========================= # Async chunk processor # ========================= @@ -77,8 +68,8 @@ def thinker2token2wav_async_chunk( # Only look at decode (generated) tokens — the prompt may contain # historical audio tokens from prior conversation turns. - all_token_ids = _ensure_list(request.all_token_ids) - prompt_len = len(_ensure_list(request.prompt_token_ids)) + all_token_ids = _common.ensure_list_preserve_none(request.all_token_ids) + prompt_len = len(_common.ensure_list_preserve_none(request.prompt_token_ids)) generated_ids = all_token_ids[prompt_len:] # Extract audio tokens and convert to 0-based IDs for Token2Wav diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py index af2cb0a257b..d7dcfc18de4 100644 --- a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Mapping from typing import Any @@ -10,35 +13,11 @@ OmniPayload, OmniPayloadStruct, ) +from vllm_omni.model_executor.stage_input_processors import _common logger = init_logger(__name__) -def _codec_audio_tensors(multimodal_output: OmniPayload | dict[str, Any]) -> list[torch.Tensor] | None: - """Inter-stage codec frames from ``codes.audio``.""" - if not isinstance(multimodal_output, Mapping): - return None - codes = multimodal_output.get("codes") - if not isinstance(codes, Mapping) or "audio" not in codes: - return None - audio = codes["audio"] - if isinstance(audio, torch.Tensor): - return [audio] - if isinstance(audio, list) and audio: - return audio - return None - - -def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None: - audio_tensors = _codec_audio_tensors(multimodal_output) - if not audio_tensors: - return None - frame = audio_tensors[-1] - if not isinstance(frame, torch.Tensor) or frame.numel() == 0: - return None - return frame.flatten() - - def generator2tokenizer_async_chunk( transfer_manager: Any, multimodal_output: OmniPayload | dict[str, Any], @@ -49,7 +28,7 @@ def generator2tokenizer_async_chunk( finished = bool(is_finished or request.is_finished()) if isinstance(multimodal_output, Mapping): - frame = _extract_last_frame(multimodal_output) + frame = _common.extract_last_codec_frame(multimodal_output, to_long=False) if frame is not None: codec_codes = frame.cpu().tolist() transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 3f94c16a725..5d75227848e 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -23,5 +23,5 @@ class OmniConnectorModelRunnerMixin(_OmniConnectorPayloadTransportMixin): """Unified data-plane communication interface for model runners. Runtime ownership and payload transport live under - ``distributed.omni_connectors``. + "distributed.omni_connectors". """