From 665e9c83920333cad0127f73a5b9e915a62be740 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 01/89] perf(minicpmo_4_5): raise the stage-1 concurrency cap to 8 The NPU override block only covered stages 0 and 2, so stage 1 kept the base max_num_seqs of 4 and gated the top tier. Measured on 910B3 (r16 against the r14 baseline): c8 RTF -15.1%, c4 -3.5%, c1 -1.1%. Concurrency only; the numerical path is untouched. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 07090660b00..7be10db65cd 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -125,6 +125,10 @@ platforms: compilation_config: cudagraph_mode: PIECEWISE - stage_id: 1 + # 第 9 项（2026-09-20 实测采纳）：stage1 并发上限 4→8。 + # npu 覆盖块原先只写了 stage0 / stage2，stage1 沿用基座的 4，这一格是漏配。 + # 收益：r16 对照 r14 基线 c8 −15.1%、c4 −3.5%、c1 −1.1%。 + max_num_seqs: 8 max_num_batched_tokens: 8192 compilation_config: cudagraph_mode: PIECEWISE From d5795daa8a49e83d6ab5c92fda907f20c9be9dc3 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 02/89] perf(minicpmo_4_5): raise the NPU graph pool from 32 to 64 Once the 32 captured shapes are exhausted, later shapes permanently degrade to eager. The CFM mel bucketing already in main (cfm_graph_bucket_frames: 16) keeps the captured shape count at 19, which is what makes a pool this size sustainable. Measured on 910B3 (r34 against r22): c4 RTF -17.0%, c8 -23.5%, c1 +5.8%. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 7be10db65cd..d4a7311bdb5 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -137,4 +137,9 @@ platforms: # estimator kernels through the NPU platform patch. additional_config: code2wav_enable_npu_graph: true - code2wav_max_npu_graphs: 32 + # 第 10 项（2026-09-20 实测采纳）：NPU 图池 32→64。32 个形状不够时会触顶， + # 之后的形状永久退化为 eager。收益（r34 对照 r22）：c4 -17.0% / c8 -23.5% / c1 +5.8%。 + # 注：该项以第 11 项（CFM mel 分桶，bucket_frames=16）为前置稳定性支撑 + # （把捕获形状数从 64 压到 19，池扩容才可持续）；main 上已有 + # cfm_graph_bucket_frames: 16，故第 11 项无需迁移。 + code2wav_max_npu_graphs: 64 From 97a033f26b50ad7b0864622ba8ff240115a560be Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 03/89] feat(minicpmo_4_5): emit a 25-frame first audio chunk The stage input processor learns a first-window size: the first payload uses initial_codec_chunk_frames while the steady-state window is unchanged, and a deployment that leaves the knob unset is bit-identical to the previous behaviour. It only takes effect when initial < codec_chunk_frames. Measured on 910B3 (r17 against r16): c1 RTF -7.4%; Seed-TTS WER 1.2179%, unchanged from the baseline. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 12 ++++++++ .../minicpmo_4_5_omni.py | 28 +++++++++++++++++-- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index d4a7311bdb5..ca1eb581690 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -20,6 +20,18 @@ connectors: connector_get_max_wait_first_chunk: 3000 connector_get_max_wait: 300 codec_chunk_frames: 25 + # 第 13 项（2026-09-20 实测采纳）：首块缝 25 帧 —— 首个 payload 用更小的窗口， + # 让音频首包更早出来。实测（r17 对照 r16）：c1 -7.4%、c4 +2.8%、c8 +13.3% + # （相对 r14 基线仍 -3.8%）；WER 干净态复测 1.2179% = 基线，零质量代价。 + # + # !! 生效条件 / 迁移依赖：本项仅在 initial < codec_chunk_frames 时生效。 + # 迁移来源分支 rtf-validated-optimizations 的稳态块长是 75，故那里的 25 生效； + # 而 main 此处的稳态块长是 25，25 >= 25 会按 _initial_chunk_frames 的回落规则 + # **视为关闭**（等价于不加这一行，不会出错）。若要让本项在 main 上真正生效， + # 需先把基座 codec_chunk_frames 调到 >25（例如 75）—— 那是基座配置差异， + # 不在本项迁移范围内，需单独决策。不要把这里调成 <25：历史上 12 帧会引入 + # 中缝、WER 翻倍（kstep_910b_validation_and_self_review.md §11.1）。 + initial_codec_chunk_frames: 25 codec_left_context_frames: 3 # HiFT capture buckets are derived from codec_chunk_frames and # codec_left_context_frames. (#5869) diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index ad329c27115..b2e3dcbc8bc 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -163,6 +163,22 @@ def _codec_config(transfer_manager: Any) -> tuple[int, int]: return chunk_frames, left_context_frames +def _initial_chunk_frames(transfer_manager: Any, chunk_frames: int) -> int: + """第 13 项（2026-09-20 实测采纳）：首个 payload 的窗口帧数。 + + 读 connector extra 里的 ``initial_codec_chunk_frames``；未配置、非正数、 + 或不小于稳态 ``codec_chunk_frames`` 时一律回落成稳态值（等价于本项关闭）。 + """ + connector = getattr(transfer_manager, "connector", None) + raw_config = getattr(connector, "config", {}) or {} + config = raw_config.get("extra", raw_config) if isinstance(raw_config, dict) else {} + config = config if isinstance(config, dict) else {} + value = int(config.get("initial_codec_chunk_frames", 0) or 0) + if value <= 0 or value >= chunk_frames: + return chunk_frames + return value + + def _request_intermediate_section(request: object, section: str) -> dict[str, object]: merged: dict[str, object] = {} for attribute in ("additional_information", "model_intermediate_buffer"): @@ -331,15 +347,23 @@ def tts2code2wav_async_chunk( request_finished = getattr(request, "is_finished", None) finished = bool(is_finished or (callable(request_finished) and request_finished())) chunk_frames, left_context_frames = _codec_config(transfer_manager) + # 第 13 项（2026-09-20 实测采纳）：首块窗口。首个 payload（此时 chunk_seq 还是 0） + # 用 initial_codec_chunk_frames，之后的 payload 回到稳态 chunk_frames； + # 未配置时 first_window == chunk_frames，行为与改动前逐位一致。 + first_window = ( + _initial_chunk_frames(transfer_manager, chunk_frames) + if int(record["chunk_seq"]) == 0 + else chunk_frames + ) flush_pending = finished last_chunk = bool(flush_pending and (not native_duplex or turn_end)) - if not flush_pending and len(pending) < chunk_frames: + if not flush_pending and len(pending) < first_window: return None hold_short_unit = ( native_duplex and flush_pending and not last_chunk and 0 < len(pending) < _MINICPMO45_MIN_STREAM_BODY_FRAMES ) - new_token_count = 0 if hold_short_unit else (len(pending) if flush_pending else chunk_frames) + new_token_count = 0 if hold_short_unit else (len(pending) if flush_pending else first_window) new_codes = pending[:new_token_count] del pending[:new_token_count] codec_start = int(state["codec_end"]) From 6edaa61ac83d1d6d9e234026b3badf2802e4e33d Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 04/89] perf(minicpmo_4_5): add an opt-in stage-2 single-row even batch Behind VLLM_OMNI_C2W_SINGLE_ROW (default off), the code2wav stage reads the switch and _iter_decode_batches gains a single-row bypass that skips wave splitting and the min-batch gate. The numerical path is unchanged (the even path is pre-existing code). Measured (A49/A49P against r34): c1 ~0, c4 -3~8%, c8 -4~7%. The item is not kept: the gain sits inside the run-to-run drift band, so a later commit in this branch removes the switch again. Signed-off-by: y-null --- .../models/minicpmo_4_5/minicpmo_4_5_code2wav.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index da2c0eb725d..9bed90f08c5 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -281,6 +281,10 @@ def __init__( raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be >= 0") if self._initial_batch_size and self._initial_batch_size < self._min_batch_size: raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be 0 or >= code2wav_min_batch_size") + # 第 49 项（2026-09-20 实测采纳，env 门控默认关）：强制 stage2 每个解码批 + # 只有 1 行，使 CFM 确定性地走 even 路径（批内等长 → 图回放），绕开 + # mixed-length 批的 ragged eager 链。数值路径零改动（even 是原有代码路径）。 + self._single_row_batch = os.environ.get("VLLM_OMNI_C2W_SINGLE_ROW", "") == "1" self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") self._default_prompt_normalized: tuple[str, str] | None = None @@ -666,6 +670,12 @@ def _iter_decode_batches( buckets: Iterable[list[_WorkItem]], ) -> Iterable[list[_WorkItem]]: for bucket in buckets: + if self._single_row_batch: + # 第 49 项：每批 1 行 → 每次 CFM 调用都落在 even 路径上 + # （批内等长 → 图回放）。刻意绕过下面的波次切分与 min-batch 门。 + for item in bucket: + yield [item] + continue if not self._initial_batch_size: yield bucket continue From 10c5f20fbc1f2a7834c69adfae229c92bc662576 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 05/89] perf(minicpmo_4_5): enable stage-0 prefix caching Stage-level enable_prefix_caching on stage 0; the top level stays false for duplex safety. The other two parts of the same source item (the HiFT/CFM graph modes and the stage-2 eager override) are already in main and are not carried. Measured on the source line (same hardware class): c1 -22.0%, c4 -3.9%, c8 +1.1%. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index ca1eb581690..1eb958356db 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -48,6 +48,13 @@ stages: # Leave headroom on the shared GPU for Talker KV + concurrent Code2Wav. gpu_memory_utilization: 0.55 enforce_eager: false + # 第 2 项（图模式 + 前缀缓存 + stage2 去 eager）—— "前缀缓存"部分。 + # 顶层 enable_prefix_caching 保持 false（双工安全），这里给 stage0 单独打开， + # 让 stage0 的重复前缀得以复用。 + # 该项另两部分 main 上已有，无需迁移：图模式（enable_hift_graph / + # enable_cfm_graph）、stage2 去 eager（code2wav_enable_npu_graph）。 + # 收益（旧代码实测口径，仅方向参考）：-22.0% / -3.9% / +1.1%（三档）。 + enable_prefix_caching: true tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" From e257447a4e89275bf8921fab2ef73717f8deeba9 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 06/89] perf(minicpmo_4_5): enable stage-0 n-gram speculative decoding Stage 0 gains async_scheduling: false plus a speculative_config (method ngram, 14 speculative tokens). async_scheduling has to be false: vLLM's VllmConfig validation rejects the async-scheduling and n-gram-proposer combination. Measured on the source line: RTF -11.8% / -3.7% / -5.9%, TTFT -69%. It is dropped again by the next commit: vLLM-omni refuses prefix caching and speculative decoding on the same stage, and prefix caching is the configuration this port validates. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 1eb958356db..573e58ee8d8 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -55,6 +55,16 @@ stages: # enable_cfm_graph）、stage2 去 eager（code2wav_enable_npu_graph）。 # 收益（旧代码实测口径，仅方向参考）：-22.0% / -3.9% / +1.1%（三档）。 enable_prefix_caching: true + # 第 3 项（stage0 ngram 投机，k=14）：N-gram speculative decoding。 + # async_scheduling 必须保持 false —— vLLM 的 VllmConfig 校验拒绝 + # "async_scheduling + ngram proposer" 这个组合。 + # 收益（旧代码实测口径，仅方向参考）：RTF -11.8% / -3.7% / -5.9%（三档）、TTFT -69%。 + async_scheduling: false + speculative_config: + method: ngram + num_speculative_tokens: 14 + prompt_lookup_min: 1 + prompt_lookup_max: 1 tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" From 0672891bfa698c84a11a17aa3f7d21b2cb015118 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 07/89] feat(minicpmo_4_5): Talker K=8 multi-frame decode with a fixed stop chain The Talker emits 8 codec frames per scheduler step instead of one. The stop chain that stalled inside the K loop is fixed end to end: * the outer model class forwards the four multi-frame entry points again (supports_multi_frame_decode, take/set_batch_stop_logits, merge_frame_outputs -- without them the multi-frame loop never engaged, and silently so); * the stop head stays the two-wide continue/stop row, and its vocabulary width is repaired for the rejection sampler; * the native rejection sampler is restored for the armed path; * the scheduler gains a K-frame guard. Without the stop fix a request ran to max_tokens (~143 s instead of ~2 s per request); with it the multi-frame decode is a net win (c1 RTF -19.1%, c8 -10.6% on the source line). Verified on 910B3: warm 3-tier 0.4434 / 1.0867 / 2.1644 with 0 ERROR, the multi-frame loop engaged on the first request, audio non-empty. Signed-off-by: y-null --- vllm_omni/config/omni_config.py | 25 +- vllm_omni/config/stage_config.py | 191 ++- vllm_omni/core/sched/omni_ar_scheduler.py | 297 ++++- vllm_omni/engine/stage_init_utils.py | 10 +- .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 31 + .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 1048 +++++++++++++++-- .../models/minicpmo_4_5/pipeline.py | 36 +- .../minicpmo_4_5/talker_codec_sample.py | 185 +++ .../platforms/npu/ascend_warmup_patch.py | 176 +++ .../platforms/npu/native_rejection_sampler.py | 82 ++ vllm_omni/platforms/npu/platform.py | 6 + .../npu/worker/npu_ar_model_runner.py | 146 ++- .../platforms/npu/worker/npu_ar_worker.py | 8 + .../npu/worker/npu_generation_worker.py | 5 + .../platforms/npu/worker/npu_model_runner.py | 75 ++ .../platforms/npu/worker/talker_multiframe.py | 959 +++++++++++++++ vllm_omni/utils/step_prof.py | 85 ++ 17 files changed, 3245 insertions(+), 120 deletions(-) create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py create mode 100644 vllm_omni/platforms/npu/ascend_warmup_patch.py create mode 100644 vllm_omni/platforms/npu/native_rejection_sampler.py create mode 100644 vllm_omni/platforms/npu/worker/talker_multiframe.py create mode 100644 vllm_omni/utils/step_prof.py diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index d64372b84b2..22f3e771372 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -28,6 +28,7 @@ from vllm.config import ParallelConfig as VllmParallelConfig from vllm.config import ProfilerConfig as VllmProfilerConfig from vllm.config import SchedulerConfig as VllmSchedulerConfig +from vllm.config import SpeculativeConfig as VllmSpeculativeConfig from vllm.config.utils import config from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs from vllm.logger import init_logger @@ -293,6 +294,7 @@ class _StageEngineValues: diffusion: _DiffusionEngineOverrides compilation_config: Mapping[str, Any] | VllmCompilationConfig | None profiler_config: Mapping[str, Any] | VllmProfilerConfig | None + speculative_config: Mapping[str, Any] | VllmSpeculativeConfig | None @dataclass(frozen=True) @@ -1211,7 +1213,17 @@ def _upstream_engine_field_map( _POOLING_ENGINE_FIELDS = frozenset(_PoolingEngineOverrides.__annotations__) _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__) _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) -_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"}) +# 迁移补齐（2026-09-21）：这三者的共同性质是"vllm 自己的配置对象，原样交给 vllm"， +# 不参与本层的结构化解包。`speculative_config` 原先不在名单里，导致迁移过来的 +# stage-1 多帧（K 步）接线无法加载：`stage_config.py` 的 +# `_apply_minicpmo_talker_multiframe_default` 会给 stage 1 注入一个 +# speculative_config（vLLM 靠 `spec_token_ids` 才知道请求一次推进 K 个 token）， +# 而这里的字段归属校验会把 stage 1 的覆盖块整块判为非法，启动即 ValueError： +# "Stage 1 (llm_ar) has explicit engine argument(s) with no structured config owner" +# 它与此前已在册的两项同类（都是直接透传 vllm 配置对象），故一并列入。 +_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset( + {"compilation_config", "profiler_config", "speculative_config"} +) _LLM_LOAD_ENGINE_FIELDS = _LOAD_ENGINE_FIELDS | frozenset(_LOAD_CONFIG_ENGINE_FIELD_MAP.values()) _LLM_CACHE_ENGINE_FIELDS = _CACHE_ENGINE_FIELDS | frozenset(_CACHE_CONFIG_ENGINE_FIELD_MAP.values()) _LLM_SCHEDULER_ENGINE_FIELDS = _SCHEDULER_ENGINE_FIELDS | frozenset(_SCHEDULER_CONFIG_ENGINE_FIELD_MAP.values()) @@ -1573,6 +1585,10 @@ def _stage_engine_values( diffusion=_DiffusionEngineOverrides(_select_engine_overrides(diffusion_kwargs, _DIFFUSION_STAGE_ENGINE_FIELDS)), compilation_config=_copy_value(engine.get("compilation_config")), profiler_config=_copy_value(engine.get("profiler_config")), + # 迁移补齐：与上两项同类（vllm 的配置对象，原样透传）。缺少这一格时， + # deploy 配置里的 speculative_config 会静默消失（不报错），K 步与 stage-0 + # ngram 都会失去效果 —— 表现为引擎日志里 speculative_config=None。 + speculative_config=_copy_value(engine.get("speculative_config")), ) @@ -1636,6 +1652,12 @@ class BaseVllmOmniStageConfig: parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig) compilation_config: VllmCompilationConfig | None = None profiler_config: VllmProfilerConfig | None = None + # 注意形状：vLLM 的 `create_speculative_config` 是按**字典**消费这个值的 + # （内部会 `.items()` / `.update()` / 再构造 SpeculativeConfig）。这里若标成 + # `VllmSpeculativeConfig | None`，pydantic 会把配置里的字典提前转成对象， + # 于是 vLLM 那一步会炸 "SpeculativeConfig object has no attribute 'items'"。 + # 因此保留 Mapping 形状，让 vLLM 自己去转。 + speculative_config: Mapping[str, Any] | None = None quantization_config: _QuantizationConfigType = None @property @@ -1798,6 +1820,7 @@ def _build_common_stage_config_kwargs( "parallel_config": parallel_config, "compilation_config": _copy_value(engine.compilation_config), "profiler_config": _copy_value(engine.profiler_config), + "speculative_config": _copy_value(engine.speculative_config), "quantization_config": _copy_value(quantization_config), }, input_proc, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 164f7cd6b11..3a2c6f9b0ed 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -5,6 +5,7 @@ from __future__ import annotations import functools +import os import re import warnings from collections.abc import Callable, Mapping @@ -779,6 +780,192 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: return merged +_MINICPMO_TALKER_FRAMES_ENV = "VLLM_OMNI_MINICPMO_TALKER_FRAMES" +# K = codec frames one stage-1 execute_model produces. 8 is the conservative +# starting point on 910C/A3; K=16 measures a further -6.9% at a higher TTFP +# cost, so the ceiling stays reachable through the environment. +_MINICPMO_TALKER_FRAMES = 8 +# Same TND ceiling as stage 0: at most 16 query positions per sequence. +_MINICPMO_TALKER_FRAMES_MAX = 16 +_MINICPMO_FRAMES_OFF = ("0", "1", "off", "false", "no") + + +_SOC_ENV_VARS = ("SOC_VERSION", "ASCEND_SOC_VERSION") + + +def _soc_name_from_env() -> str: + """SoC identifier from the launch environment. Never touches the device. + + Asking torch_npu for the device name here would initialise the NPU in + whichever process loads the deploy config -- the API server -- and that + process would then hold device memory it has no use for, which shows up as + an extra entry in npu-smi. CANN launch environments already export the + revision: ``ascend910b1`` on a 910B box, ``ascend910_9391`` on 910C. + """ + for variable in _SOC_ENV_VARS: + value = os.environ.get(variable, "").strip().lower() + if value: + return value + return "" + + +def _npu_soc_allows_talker_multiframe() -> bool: + """False on the 910B family, whose rejection verifier cannot take K rows. + + Judged from the environment only, so loading a deploy config never + initialises a device. An unidentifiable SoC is treated as "not allowed": + the cost of leaving the frame loop off is a slower run, while speccing a + 910B faults its rejection kernel. Setting ``VLLM_OMNI_MINICPMO_TALKER_FRAMES`` + explicitly overrides this (see the caller). + """ + name = _soc_name_from_env() + if not name: + logger.info( + "[minicpmo] SoC not identified from %s; leaving the Talker " + "multi-frame loop off. Set %s to force it on.", + "/".join(_SOC_ENV_VARS), + _MINICPMO_TALKER_FRAMES_ENV, + ) + return False + return not name.startswith("ascend910b") + + +def talker_frames_per_step() -> int: + """Codec frames one stage-1 `execute_model` produces. 1 disables it. + + Deliberately read from the environment on every call rather than cached: + the stage config and the model runner agree through this one function and + are built at different times in different processes. + """ + raw = os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip().lower() + if not raw: + return _MINICPMO_TALKER_FRAMES + if raw in _MINICPMO_FRAMES_OFF: + return 1 + try: + frames = int(raw) + except ValueError: + logger.warning( + "[minicpmo] %s=%r is not an integer; using the code default %d", + _MINICPMO_TALKER_FRAMES_ENV, + raw, + _MINICPMO_TALKER_FRAMES, + ) + return _MINICPMO_TALKER_FRAMES + if frames < 1: + return 1 + return min(frames, _MINICPMO_TALKER_FRAMES_MAX) + + +def talker_multiframe_armed() -> bool: + """Whether the Talker K-frame loop will run in this process. + + One source of truth for the three places that must agree: the deploy-config + injection below (does stage 1 get the speculative_config), the model's own + gate, and stage 1's stop-token constraint. The last one is not obvious: when + the loop is armed the vLLM-level head collapses to the two-wide + continue/stop row (``compute_logits``), so the only ids that can be sampled + are 0 and 1 and the codec EOS becomes unreachable *at that level*. A + ``stop_token_ids`` the head can never emit means `check_stop` never fires + and the request runs to ``max_tokens`` -- the loop keeps paying for a + request that already finished at the codec level. + + An explicit ``VLLM_OMNI_MINICPMO_TALKER_FRAMES`` is an operator decision and + arms regardless of the SoC, exactly as the loader below does. Everything + else is the SoC gate's answer, and an unidentified SoC counts as "no". + """ + if talker_frames_per_step() <= 1: + return False + if os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip(): + return True + return _npu_soc_allows_talker_multiframe() + + +def _apply_minicpmo_talker_multiframe_default(deploy: "DeployConfig") -> None: + """Let the Talker produce K codec frames per scheduler step on 910C/A3. + + A stage-1 decode step is dominated by per-*step* host work rather than per + frame, so running K frames inside one `execute_model` pays that host cost + once (measured on 910C/A3: ~2.9 ms of host around a 0.83 ms device forward, + over ~118 frames per request). The frames are *sequential* -- frame k+1 is + conditioned on the codec token sampled at frame k -- so this is not + speculative decoding and the runner runs its own inner loop. vLLM still + has to know the request advanced by K tokens, because that is what grows + the block table and KV slots the loop writes into, and `spec_token_ids` is + V1's only mechanism for that; hence a speculative_config that carries the + count. The n-gram proposer is never consulted: `propose_draft_token_ids` + returns the Talker's `continue` token K-1 times, so verification is exact + by construction. + + `async_scheduling` comes off with it. vLLM refuses async scheduling next + to an n-gram proposer ("only supported with EAGLE/MTP/Draft Model/NGram + GPU/DSpark"), which is the same reason stage 0 carries the flag, and a + synchronous scheduler costs less here than the K-fold host saving returns. + + Left off entirely on the 910B family: its rejection verifier trips a + vector-core limit on spec-width rows, so the baseline there stays exactly + as the deploy config describes it. + """ + frames = talker_frames_per_step() + if frames <= 1: + return + explicit = os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip() != "" + if not talker_multiframe_armed(): + logger.info( + "[minicpmo] Talker multi-frame decode left off: the SoC either " + "cannot verify spec-width rows (910B family limit) or was not " + "identified. Deploy on 910C/A3, or set %s explicitly to force it.", + _MINICPMO_TALKER_FRAMES_ENV, + ) + return + if explicit: + # An explicit request is an operator decision: arm it and let the + # model-side gate make the final call on the live device. + logger.info( + "[minicpmo] %s is set explicitly; arming K=%d without the SoC check", + _MINICPMO_TALKER_FRAMES_ENV, + frames, + ) + for stage in deploy.stages: + if stage.stage_id != 1: + continue + if stage.engine_extras.get("speculative_config") is not None: + # An explicit config wins; this default only fills a gap. + logger.info( + "[minicpmo] stage 1 carries an explicit speculative_config; " + "leaving it and its async_scheduling flag untouched", + ) + return + stage.engine_extras["speculative_config"] = { + "method": "ngram", + "num_speculative_tokens": frames - 1, + "prompt_lookup_min": 1, + "prompt_lookup_max": 1, + } + stage.async_scheduling = False + logger.info( + "[minicpmo] stage 1 multi-frame decode on (%d codec frames per " + "step, synchronous scheduler); set %s=1 to override", + frames, + _MINICPMO_TALKER_FRAMES_ENV, + ) + + +def _apply_minicpmo_perf_defaults(deploy: "DeployConfig", config_path: object = None) -> None: + """Apply MiniCPM-o perf code-defaults, scoped to this deploy config on NPU.""" + if config_path is None or "minicpmo" not in str(config_path).lower(): + return + try: + from vllm_omni.platforms import current_omni_platform + + device_name = getattr(current_omni_platform, "device_name", None) + except Exception: + return + if (device_name or "").lower() != "npu": + return + _apply_minicpmo_talker_multiframe_default(deploy) + + def load_deploy_config(path: str | Path) -> DeployConfig: """Load a deploy YAML (with optional base_config inheritance).""" raw_dict = resolve_deploy_yaml(path) @@ -821,7 +1008,9 @@ def load_deploy_config(path: str | Path) -> DeployConfig: ): if name in raw_dict: kwargs[name] = raw_dict[name] - return DeployConfig(**kwargs) + deploy = DeployConfig(**kwargs) + _apply_minicpmo_perf_defaults(deploy, path) + return deploy class PlatformOverrides(NamedTuple): diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index f07aae021e8..07dd69e23e0 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -3,6 +3,7 @@ from __future__ import annotations +import os import time from collections import defaultdict from collections.abc import Iterable, Iterator @@ -32,6 +33,77 @@ logger = init_logger(__name__) +# --------------------------------------------------------------------------- +# K 步记账 trace（默认关） +# +# VLLM_OMNI_MINICPMO_KSTEP_ACCOUNT_TRACE=1 +# +# 回答的问题：模型已经在发 row_stop，为什么请求不停？记每一步每个请求"被接受的 +# token 值"——若它恒为 0（continue），说明 stop 行没有变成请求 token 表里的 1， +# check_stop(stop_token_ids=[1]) 永远不会触发，请求就会一路走到 max_model_len。 +# 契约：诊断永不抛异常。 +# --------------------------------------------------------------------------- +_KSTEP_ACCOUNT_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_ACCOUNT_TRACE" +_KSTEP_ACCOUNT_TRACE: dict = {"lines": 0, "reqs": {}, "off": False} + + +def kstep_account_trace_enabled() -> bool: + """Off unless the env turns it on; latched off for good after a failure.""" + if _KSTEP_ACCOUNT_TRACE["off"]: + return False + raw = os.environ.get(_KSTEP_ACCOUNT_TRACE_ENV, "").strip().lower() + return bool(raw) and raw not in ("0", "false", "no", "off") + + +def _trace_kstep_account( + req_id, spec_len, generated, computed, total, prompt_len, status, sampling=None +) -> None: + """One line per request per step; first 24 steps of each request, then 1 in 16.""" + state = _KSTEP_ACCOUNT_TRACE + if state["off"]: + return + try: + key = str(req_id) + seen = state["reqs"].get(key, 0) + state["reqs"][key] = seen + 1 + gen_len = len(generated or []) + # Invariant I2/I5: a step with drafts must hand the request + # spec_len + 1 tokens (the bonus plus every accepted draft). A short + # row means frames the model produced were dropped on the way to the + # request -- always log it, whatever the stride. + partial = bool(spec_len) and gen_len != spec_len + 1 + interesting = seen < 24 or seen % 16 == 0 or partial + nonzero = [int(t) for t in (generated or []) if int(t) != 0] + if not interesting and not nonzero: + return + if state["lines"] >= 400: + return + state["lines"] += 1 + logger.info( + "[kstep-acct] req=%s step=%d spec_len=%s gen=%s partial=%s nonzero=%s computed=%s tokens=%s " + "prompt=%s status=%s stop_ids=%s eos=%s min=%s max=%s ignore_eos=%s", + key[:12], + seen, + spec_len, + gen_len, + int(partial), + nonzero[:6], + computed, + total, + prompt_len, + status, + list(getattr(sampling, "stop_token_ids", None) or []), + getattr(sampling, "eos_token_id", None), + getattr(sampling, "min_tokens", None), + getattr(sampling, "max_tokens", None), + getattr(sampling, "ignore_eos", None), + ) + except Exception as exc: # pragma: no cover - diagnostics never raise + state["off"] = True + logger.warning("[kstep-acct] trace disabled after failure: %r", exc) + + + def _should_emit_engine_output( model_config: Any, @@ -337,6 +409,142 @@ def _commit_kv_transfer_trigger( if stop_after_transfer and req_id in self.requests_needing_kv_transfer: self.pending_stop_after_extraction.add(req_id) + # Env var shared with stage_config._apply_minicpmo_talker_multiframe_default; + # keep the name in sync there. + _TALKER_FRAMES_ENV = "VLLM_OMNI_MINICPMO_TALKER_FRAMES" + + def _talker_kstep_armed(self) -> bool: + """True when this scheduler drives the Talker K-frame decode. + + Mirrors `_apply_minicpmo_talker_multiframe_default`: the Talker stage + gets an injected n-gram config with exactly frames-1 draft tokens. + Matching that fingerprint keeps both sides reading the same env var. + The text stage's explicit n-gram config (15 draft tokens by default) + does not match, so this stays a no-op there. + """ + cache = getattr(self, "_omni_talker_kstep_cache", None) + if cache is None: + # vLLM's Scheduler does not expose speculative_config as an + # attribute and SchedulerConfig has no num_speculative_tokens + # either, so the canonical engine-side source is vllm_config. + # Reading anything else silently yields 0 drafts and this whole + # guard never arms -- exactly the bug that survived the 11:47 + # crash fix. + spec = getattr(self, "speculative_config", None) + if spec is None: + vllm_cfg = getattr(self, "vllm_config", None) + spec = getattr(vllm_cfg, "speculative_config", None) if vllm_cfg is not None else None + num_spec = 0 + is_ngram = True + if spec is not None: + num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 + is_ngram = getattr(spec, "method", "ngram") == "ngram" + try: + frames = int(os.environ.get(self._TALKER_FRAMES_ENV, "8") or 8) + except ValueError: + frames = 8 + armed = is_ngram and frames > 1 and num_spec == frames - 1 + cache = self._omni_talker_kstep_cache = armed + return cache + + def _log_kstep_guard_view(self, verdict: str, widths: set[int], states: list, dropped: int = 0) -> None: + """Diagnostic view of what the K-step guard judged this step. + + Rate-limited but never silenced on the interesting steps: the 12:39 + crash landed on a step whose one-shot guard log had already been + consumed, which hid exactly the state the guard passed. Diagnostics + must never raise either -- this one must not kill the engine. + """ + count = getattr(self, "_omni_kstep_guard_view_count", 0) + 1 + self._omni_kstep_guard_view_count = count + if count > 20 and count % 50 != 0: + return + try: + logger.info( + "[K-guard] #%d %s: waiting=%d widths=%s dropped=%d " + "states=(id, computed, prompt_len, total, spec_len, delta)=%s", + count, + verdict, + len(self.waiting), + sorted(widths), + dropped, + states, + ) + except Exception as exc: # never let the diagnostic kill the engine + logger.info("[K-guard] #%d %s (view failed: %r)", count, verdict, exc) + + def _drop_talker_drafts_if_prefill_pending(self) -> None: + """Keep the Talker's K-frame decode out of steps with uneven spans. + + The Talker multi-frame loop requires uniform decode spans, and the + runner refuses anything else (better a crash than a silently wrong + codec stream). What matters is the number of rows vLLM is about to + schedule per request -- not the request's spec_token_ids length: + + * a plain decode is scheduled ``1 + num_spec_tokens`` rows (the + continuation placeholders, re-armed every step); + * a request carrying a streaming (or chunked) input chunk keeps its + own token count, e.g. 7 rows for a 7-token chunk; + * a decode that cannot fit the padded width (near max_model_len) + falls back to a single row. + + Any mix of those in one step produces non-uniform spans, even in a + pure decode batch with no prefill in sight. Rather than dying at + the runner, drop the continuation drafts for this round: every + decode row schedules a single token next step, the step takes the + single-frame path, and the next propose re-arms the K frames. + Continuation drafts are stateless, so dropping them is free. + """ + if not self._talker_kstep_armed(): + return + num_spec = int(getattr(self, "num_spec_tokens", 0) or 0) + max_len = getattr(self, "max_model_len", None) + try: + max_len = int(max_len) if max_len is not None else None + except (TypeError, ValueError): + max_len = None + prefill_pending = bool(self.waiting) + widths: set[int] = set() + states: list = [] + for req in self.running: + computed = int(req.num_computed_tokens) + prompt_len = len(req.prompt_token_ids) + total = int(getattr(req, "num_tokens", prompt_len)) + delta = total - computed + spec_len = len(req.spec_token_ids or []) + states.append( + (str(getattr(req, "request_id", "?"))[:12], computed, prompt_len, total, spec_len, delta) + ) + if computed < prompt_len: + prefill_pending = True + continue + if delta <= 0: + # Nothing scheduled for this request; it owns no rows. + continue + if delta == 1: + can_pad = max_len is None or computed + 1 + num_spec + 1 <= max_len + widths.add(1 + num_spec if (num_spec > 0 and can_pad) else 1) + else: + widths.add(delta) + uneven = len(widths) > 1 + if not prefill_pending and not uneven: + if not widths and states: + # Running requests exist but none yields a row width: the step + # schedules no decode rows at all. Rare enough to be interesting. + self._log_kstep_guard_view("pass-no-decode-widths", widths, states) + return + dropped = 0 + for req in self.running: + if req.spec_token_ids: + req.spec_token_ids = [] + dropped += 1 + self._log_kstep_guard_view( + f"drop ({'prefill pending' if prefill_pending else 'uneven spans'})", + widths, + states, + dropped=dropped, + ) + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Remove FINISHED_ABORTED requests before the upstream scheduler sees # them. Upstream vllm raises RuntimeError on this status; omni allows @@ -347,6 +555,13 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self._process_pending_omni_inputs(model_mode="ar") self._drop_aborted_queued_requests() self._resync_streaming_input_counter() + # Talker K-frame guard: a step whose decode spans would be uneven + # (mixed prefill+decode rows, or a first-step decode joining a + # steady-state K-step request) must not carry the K-step drafts -- + # the multi-frame loop requires uniform decode spans and the runner + # refuses the rest. + self._drop_talker_drafts_if_prefill_pending() + original_waiting = None if self._should_defer_waiting_admission(): original_waiting = waiting @@ -463,6 +678,22 @@ def update_from_output( stopped_running_reqs: set[Request] = set() stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): + if num_tokens_scheduled <= 0: + # P17 SCHED0 diagnostic: a zero-token schedule entry used to + # die on the bare assert below with no context, which is what + # made the K-step stall so expensive to scope. Report WHICH + # request and in what state first; the assert still fires. + _r = self.requests.get(req_id) + _spec = scheduler_output.scheduled_spec_decode_tokens.get(req_id) + logger.error( + "SCHED0 zero-token schedule: req=%s n=%s spec_tokens=%s status=%s " + "num_computed=%s num_output_placeholders=%s finished=%s all_entries=%s", + req_id, num_tokens_scheduled, _spec, + getattr(_r, "status", None), getattr(_r, "num_computed_tokens", None), + getattr(_r, "num_output_placeholders", None), + None if _r is None else _r.is_finished(), + dict(list(num_scheduled_tokens.items())[:8]), + ) assert num_tokens_scheduled > 0 request = self.requests.get(req_id) if request is not None: @@ -495,6 +726,17 @@ def update_from_output( continue req_index = model_runner_output.req_id_to_index[req_id] generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] + if kstep_account_trace_enabled(): + _trace_kstep_account( + req_id, + len(scheduler_output.scheduled_spec_decode_tokens.get(req_id) or []), + generated_token_ids, + int(request.num_computed_tokens), + int(request.num_tokens), + len(getattr(request, "prompt_token_ids", None) or []), + getattr(request, "status", None), + getattr(request, "sampling_params", None), + ) stale_async_tokens = int(getattr(request, "async_tokens_to_discard", 0) or 0) async_output_is_stale = bool(generated_token_ids and stale_async_tokens > 0) @@ -556,6 +798,52 @@ def update_from_output( num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, request_id=req_id, ) + elif scheduled_spec_token_ids and sampled_token_ids: + # P17 layer C: this request had drafts scheduled but its row + # came back with no generated tokens while the step itself + # sampled (a logprob-contract failure emptied the row above, + # the rejection sampler dropped every token, ...). The draft + # tokens were counted as computed but none of them will ever + # produce output, so roll them back -- otherwise the next + # step schedules a negative count and the engine stalls. + # NOTE the guard on `scheduled_spec_token_ids`: upstream takes + # it from `.get(req_id)` and a request with no drafts this + # step gets None -- an empty row without drafts is normal + # (a non-final prefill chunk produces no tokens) and upstream + # skips it; the guard above keeps that path intact. + # A prefill-chunk row is different: the chunk's prompt tokens + # did land in the KV cache, so only the D drafts roll back + # (the case entry_02 measured; the base token stays advanced). + # A pure-decode row that came back empty rolls the base token + # back too -- schedule() optimistically advanced all scheduled + # tokens and not one of them produced output. + _drafts = len(scheduled_spec_token_ids) + _scheduled = num_tokens_scheduled + _prev_computed = request.num_computed_tokens - _scheduled + _prompt_len = ( + len(request.prompt_token_ids) + if getattr(request, "prompt_token_ids", None) is not None + else 0 + ) + _is_prefill_row = _prev_computed < _prompt_len + _rollback = _drafts if _is_prefill_row else min(_scheduled, _drafts + 1) + logger.error( + "K-step: request %s returned an empty row on a %s step " + "(scheduled=%s drafts=%s computed=%s prompt_len=%s); " + "rolling back %s. An empty decode row is never expected " + "from the verify path -- this log is the stall signature.", + req_id, + "prefill" if _is_prefill_row else "decode", + _scheduled, + _drafts, + request.num_computed_tokens, + _prompt_len, + _rollback, + ) + if request.num_computed_tokens >= _rollback: + request.num_computed_tokens -= _rollback + if request.num_output_placeholders >= _rollback: + request.num_output_placeholders -= _rollback # Free encoder inputs only after the step has actually executed. if request.has_encoder_inputs: @@ -876,7 +1164,14 @@ def _update_request_as_session(self, session: Request, update: StreamingUpdate) session.async_tokens_to_discard = 1 session.num_computed_tokens -= session.num_output_placeholders session.num_output_placeholders = 0 - session.spec_token_ids = [] + # P17 layer B: clear stale drafts unconditionally, not only under + # async scheduling. With async off -- which the K-step vehicle + # requires -- a draft proposed before the segment boundary survived + # into the next segment's prefill; the runner discards prefill rows + # that are not the final chunk, so that row produced no token + # forever and the request either tripped a negative num_new_tokens + # or livelocked until it timed out. + session.spec_token_ids = [] stage_id = self.vllm_config.model_config.stage_id update_infos = ( diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index faaf6b70721..26084b4fd48 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1157,8 +1157,14 @@ def _project_omni_stage_engine_args( ): engine_args.update(_project_upstream_config_fields(config, field_map)) - for name in ("compilation_config", "profiler_config"): - value = getattr(stage_config, name) + # 迁移补齐（2026-09-21）：这一格与上面两项同类 —— 都是 vllm 自己的配置对象， + # 从 stage 配置原样搬进引擎参数。原先只搬 compilation_config / profiler_config， + # 于是 deploy 配置里的 speculative_config（stage-0 的 ngram 投机、以及 + # stage_config.py 给 stage-1 注入的 K 步投机）会在这一步静默消失，不报错， + # 只在引擎侧表现为 `speculative_config=None`；K 步与 ngram 都会失去效果。 + # 用 getattr(..., None) 兜底，避免个别 stage 配置类型没有该字段时炸掉。 + for name in ("compilation_config", "profiler_config", "speculative_config"): + value = getattr(stage_config, name, None) if value is not None: engine_args[name] = copy.deepcopy(value) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index 03024bbe586..a6250b53a66 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -634,6 +634,37 @@ def make_omni_output(self, model_outputs, **kwargs): return model_outputs return self.talker.make_omni_output(model_outputs, **kwargs) + # 第 8 项（K 步多帧解码）在 runner 侧的四个转发（2026-09-21 迁移补齐）。 + # 基线的这四个方法就写在这个外层包装类上，迁移时整段漏掉。runner 只看得见这个 + # 包装类（`_model_forward` 里 talker_multiframe.applies / is_multi_token_decode / + # 以及 `ensure_stop_token_vocab` 的判定全部是 `getattr(model, ...)`），所以： + # · 缺 `supports_multi_frame_decode` → 判定静默取到 False → 多帧循环永不进入， + # 且连"拒绝原因"都不打日志（applies 那条分支是静默返回 0）； + # · 缺 take/set_batch_stop_logits → 每帧的 stop 行拿不到； + # · 缺 merge_frame_outputs → 多帧结果无法合并。 + # 后果：K 步的投机配置已经进了引擎（vLLM 层在做投机解码），但 runner 侧的多帧 + # 循环不进入、循环里的词表宽度修正（ensure_stop_token_vocab）也不执行， + # 于是拒绝采样器按宽度 0 消费、AIV 越界（507035）、stage-1 引擎崩、音频全空。 + @property + def supports_multi_frame_decode(self) -> bool: + # 这些方法在两类 stage 上都在，只有 Talker（tts）能走多帧，故按 stage 判定。 + return self.model_stage == "tts" + + def take_batch_stop_logits(self): + if self.model_stage != "tts": + return None + return self.talker.take_batch_stop_logits() + + def set_batch_stop_logits(self, logits) -> None: + if self.model_stage != "tts": + return + self.talker.set_batch_stop_logits(logits) + + def merge_frame_outputs(self, frame_outputs, frame_stop_logits): + if self.model_stage != "tts": + return frame_outputs + return self.talker.merge_frame_outputs(frame_outputs, frame_stop_logits) + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None: # Handle OmniOutput type if isinstance(hidden_states, OmniOutput): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 67843833dda..3f50d8e5ebd 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -13,7 +13,12 @@ """ from collections.abc import Iterable, Mapping, Sequence +import os from dataclasses import replace +try: + import torch_npu # noqa: F401 (NPU platform guarantee) +except Exception: # pragma: no cover + torch_npu = None from typing import Any import torch @@ -26,17 +31,53 @@ from vllm.model_executor.models.llama import LlamaModel from vllm.model_executor.models.utils import maybe_prefix from vllm.v1.sample.sampler import Sampler +from vllm.v1.outputs import SamplerOutput from vllm_omni.engine.duplex.intermediate import get_tts_handoff from vllm_omni.model_executor.models.minicpmo_4_5 import ( MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK, MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS, ) +from vllm_omni.model_executor.models.minicpmo_4_5.talker_codec_sample import ( + MIN_TOKENS_TO_KEEP, + TalkerCodecDeviceState, + TalkerCodecSampleResult, + codec_sample_result, + greedy_codec_sample, + make_device_state, + prepare_codec_logits, +) from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform +from vllm_omni.utils.step_prof import span as _prof_span logger = init_logger(__name__) +# Codec-sampling fallbacks the multi-frame path reads directly. The deploy YAML +# and the checkpoint's tts_config override them per request; these match what +# the single-frame path falls back to. +_REPETITION_WINDOW = 16 +_MIN_AUDIO_TOKENS = 64 +_CODEC_SEED = 42 +_CODEC_TEMPERATURE = 0.8 +_CODEC_TOP_K = 25 +_CODEC_TOP_P = 0.85 +_CODEC_REPETITION_PENALTY = 1.05 +_CODEC_MIN_TOKENS = 50 +_CODEC_MAX_TOKENS = 2048 +_FAST_CODEC_PENALTY = os.getenv("MINICPMO_FAST_CODEC_PENALTY", "1") == "1" +# YAML key -> (tts_config attribute, hardcoded fallback, type) +_CODEC_SAMPLING_SOURCES: tuple[tuple[str, str, Any, Any], ...] = ( + ("seed", "seed", _CODEC_SEED, int), + ("temperature", "temperature", _CODEC_TEMPERATURE, float), + ("top_k", "top_k", _CODEC_TOP_K, int), + ("top_p", "top_p", _CODEC_TOP_P, float), + ("repetition_penalty", "repetition_penalty", _CODEC_REPETITION_PENALTY, float), + ("min_tokens", "min_new_tokens", _CODEC_MIN_TOKENS, int), + ("max_tokens", "max_new_tokens", _CODEC_MAX_TOKENS, int), +) + + _REPETITION_PENALTY_CHUNK_SIZE = 16 # ``past_window`` of MiniCPMTTS's codec repetition penalty: both generate() and # generate_chunk() build it through gen_logits(), which hardcodes @@ -157,6 +198,242 @@ def _apply_batched_repetition_penalty( return penalized +def _apply_top_k_top_p( + logits: torch.Tensor, + *, + top_k: int | None, + top_p: float | None, + min_tokens_to_keep: int = 3, + inplace: bool = False, +) -> torch.Tensor: + """Reference warper: same candidate floors as the upstream warpers.""" + filtered = logits if inplace else logits.clone() + vocab_size = filtered.shape[-1] + if top_p is not None and 0.0 < top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(filtered, descending=False, dim=-1) + cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) + remove = cumulative_probs <= (1.0 - float(top_p)) + remove[..., -min_tokens_to_keep:] = False + remove = remove.scatter(-1, sorted_indices, remove) + filtered.masked_fill_(remove, float("-inf")) + if top_k is not None and top_k > 0: + keep = min(vocab_size, max(int(top_k), min_tokens_to_keep)) + threshold = torch.topk(filtered, keep, dim=-1).values[..., -1, None] + filtered.masked_fill_(filtered < threshold, float("-inf")) + return filtered + + +_NPU_TOPK_CACHE: dict = {} +_NPU_TOP_K_TOP_P = os.environ.get("MINICPMO_TTS_NPU_TOPK_TOPP", "1") != "0" +# npu_top_k_top_p wins on launches at small batch but loses to the two-op +# torch path at large batch (128/8 A/B: 1.391 gated-off vs on). +_NPU_TOPK_MAX_BATCH = int(os.environ.get("MINICPMO_TTS_NPU_TOPK_TOPP_MAX_BATCH", "64")) + + +def _npu_top_k_top_p_warp(logits, *, top_k, top_p): + """Fused top-k/top-p floor via torch_npu.npu_top_k_top_p. + + Kernel semantics: keep top-k by value, then keep the top-p probability + mass over the retained set. Falls back to the exact PyTorch warper when + the kernel is unavailable or fails. + """ + if top_k is None or top_p is None or not 0.0 < top_p < 1.0: + return logits + npu = getattr(torch_npu, "npu_top_k_top_p", None) + if npu is None: + return logits + key = (str(logits.device), str(logits.dtype), float(top_p), int(top_k)) + cached = _NPU_TOPK_CACHE.get(key) + if cached is None: + dev = torch.full((1,), float(top_p), device=logits.device, dtype=logits.dtype) + dk = torch.full((1,), int(top_k), device=logits.device, dtype=torch.int32) + _NPU_TOPK_CACHE[key] = (dev, dk) + else: + dev, dk = cached + try: + return npu( + logits, + dev.expand(logits.shape[0]).contiguous(), + dk.expand(logits.shape[0]).contiguous(), + ) + except Exception: + return _apply_top_k_top_p(logits, top_k=top_k, top_p=top_p, min_tokens_to_keep=3, inplace=True) + + +def _maybe_prewarp_top_k_top_p(logits, sampling_metadata): + """Pre-apply the fused floor and neutralize the sampler's own warpers. + + Only engages when every row shares one (top_k, top_p) at temperature 1.0, + where the floor on raw logits is exactly the floor the sampler would + compute itself (top-k is rank based, so temperature invariance holds and + top-p mass matches at T=1). Any structural surprise returns None and the + native sampler path runs untouched. + """ + if not _NPU_TOP_K_TOP_P: + return None + if not isinstance(logits, torch.Tensor) or logits.ndim != 2 or logits.shape[0] == 0: + return None + if logits.shape[0] > _NPU_TOPK_MAX_BATCH: + return None + if getattr(torch_npu, "npu_top_k_top_p", None) is None: + return None + try: + params = getattr(sampling_metadata, "sampling_params", None) + if not isinstance(params, (list, tuple)) or len(params) != logits.shape[0]: + return None + top_k = top_p = None + for sp in params: + k = getattr(sp, "top_k", None) + pp = getattr(sp, "top_p", None) + t = getattr(sp, "temperature", None) + if t is None or float(t) != 1.0: + return None + if k is None or int(k) <= 0: + return None + if pp is None or not 0.0 < float(pp) < 1.0: + return None + if top_k is None: + top_k, top_p = int(k), float(pp) + elif int(k) != top_k or float(pp) != top_p: + return None + logits = _npu_top_k_top_p_warp(logits, top_k=top_k, top_p=top_p) + new_params = [replace(sp, top_p=1.0, top_k=-1) for sp in params] + sampling_metadata = replace(sampling_metadata, sampling_params=list(new_params)) + return logits, sampling_metadata + except Exception: + return None + + +def resolve_codec_sampling_params( + yaml_params: Mapping[str, Any] | None, + tts_config: Any | None = None, +) -> dict[str, Any]: + """Resolve Talker codec knobs: deploy YAML, then checkpoint, then defaults. + + The deploy YAML that starts a ranked run comes from the organizer's + baseline branch and carries no ``codec_sampling_params`` block, so + requiring one there is a startup failure rather than a configuration + error: the Talker stage would not come up at all under the official + launch. Falling back to the checkpoint's ``tts_config`` is what the + upstream Talker does, and it is what makes this tree drop into the + official deploy config unchanged. + + A YAML block still wins key by key, so our own configs keep overriding. + """ + provided = yaml_params if isinstance(yaml_params, Mapping) else {} + resolved: dict[str, Any] = {} + sources: dict[str, str] = {} + for key, attribute, fallback, cast in _CODEC_SAMPLING_SOURCES: + value = provided.get(key) + source = "yaml" + if value is None: + value = getattr(tts_config, attribute, None) if tts_config is not None else None + source = "tts_config" + if value is None: + value, source = fallback, "default" + resolved[key] = cast(value) + sources[key] = source + if resolved["min_tokens"] < 0: + raise ValueError("codec_sampling_params.min_tokens must be >= 0") + if resolved["max_tokens"] <= 0: + raise ValueError("codec_sampling_params.max_tokens must be > 0") + logger.info("MiniCPM-o Talker codec sampling %s (from %s)", resolved, sources) + return resolved + + + +# --------------------------------------------------------------------------- +# K 步逐帧诊断（默认关，与 talker_multiframe 的逐步计时互补） +# +# VLLM_OMNI_MINICPMO_KSTEP_FRAME_TRACE=1 打开 +# 最多打 _KSTEP_FRAME_TRACE_FRAMES 帧，然后静默 +# +# 回答逐步计时回答不了的那个问题：K 循环里**模型自己的采样器**有没有走到 EOS， +# 以及每一帧实际吃进去的 code 是不是上一帧的采样值（契约不变式 I1）。 +# 契约：诊断永不许反杀引擎 —— 任何异常只把本诊断永久关掉并记一条 warning。 +# --------------------------------------------------------------------------- +_KSTEP_FRAME_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_FRAME_TRACE" +_KSTEP_FRAME_TRACE_LINES = 128 # 打点行数上限 +_KSTEP_FRAME_TRACE_HEAD = 16 # 头 N 帧逐帧（用于校验帧序不变式 I1） +_KSTEP_FRAME_TRACE_STRIDE = 16 # 之后每 N 帧一条（覆盖到约第 1800 帧，够看 EOS） +_KSTEP_FRAME_TRACE: dict[str, Any] = {"reqs": {}, "lines": 0, "off": False} + + +def kstep_frame_trace_enabled() -> bool: + """Off unless the env turns it on; latched off for good after a failure.""" + if _KSTEP_FRAME_TRACE["off"]: + return False + raw = os.environ.get(_KSTEP_FRAME_TRACE_ENV, "").strip().lower() + return bool(raw) and raw not in ("0", "false", "no", "off") + + +def _trace_kstep_frame(model: Any, **fields: Any) -> None: + """One line per codec frame while the multi-frame loop owns sampling. + + Logs what the *model* decided: the code each frame took as its input (the + previous frame's sample, per the contract), what it sampled, whether that + was the codec EOS, and where the sampler's own step counter stood. The first + ``HEAD`` frames are logged one per line (that is what verifies the frame + chain), then one line every ``STRIDE`` frames so the window still covers the + frames where EOS first becomes eligible -- `min_tokens` masks it before + then, so a head-only trace can never see a stop. An EOS sample is always + logged regardless of the stride. Never raises: a failure disables the trace + for good. + """ + state = _KSTEP_FRAME_TRACE + if state["off"]: + return + reqs = state["reqs"] + request_id = str(fields.get("request_id")) + seen = reqs.get(request_id, 0) + reqs[request_id] = seen + 1 + is_eos = bool(fields.get("is_eos")) + interesting = seen < _KSTEP_FRAME_TRACE_HEAD or ( + seen >= _KSTEP_FRAME_TRACE_HEAD and (seen - _KSTEP_FRAME_TRACE_HEAD) % _KSTEP_FRAME_TRACE_STRIDE == 0 + ) + if not is_eos and (not interesting or state["lines"] >= _KSTEP_FRAME_TRACE_LINES): + return + try: + device_states = getattr(model, "_request_codec_device_states", None) + device_state = ( + device_states.get(fields.get("request_id")) if isinstance(device_states, dict) else None + ) + device_step = int(device_state.step.reshape(-1)[0].item()) if device_state is not None else -1 + state["lines"] += 1 # 只防刷屏，窗口按请求单独计数 + logger.info( + "[kstep-frame] frame=%d req=%s step=%s->%s device_step=%s in_code=%s row=%s " + "sampled=%s is_eos=%s limit=%s finished=%s ctx=%s min=%s max=%s", + seen, + str(fields.get("request_id"))[:12], + fields.get("step_before"), + fields.get("step_after"), + device_step, + fields.get("prev_code"), + fields.get("row"), + fields.get("sampled_id"), + fields.get("is_eos"), + fields.get("reached_limit"), + fields.get("finished"), + fields.get("ctx_len"), + fields.get("min_tokens"), + fields.get("max_tokens"), + ) + except Exception as exc: # pragma: no cover - diagnostics never raise + state["off"] = True + logger.warning("[kstep-frame] trace disabled after failure: %r", exc) + + +def _codec_int_param(state: Any, key: str, fallback: int) -> int: + """Read an integer codec knob from a request state, None meaning unset. + + ``dict.get(key, fallback)`` does not help when the key is present but holds + None -- which the offline duplex path uses to mean "the deploy YAML owns + this knob" -- so None is mapped to the fallback explicitly here. Kept in + one place so every reader agrees on the rule. + """ + value = state.get(key) if isinstance(state, Mapping) else None + return int(fallback if value is None else value) + class _MiniCPMTTSProjector(nn.Module): """Checkpoint-compatible hidden-state projector used by MiniCPMTTS.""" @@ -182,9 +459,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config: MiniCPMOConfig = vllm_config.model_config.hf_config self.config = config self.vllm_config = vllm_config - self._force_eos_rows: list[bool] | None = None - self._mask_eos_rows: list[bool] | None = None - self._pending_force_eos_rows: list[bool] | None = None self._penalty_histories: list[torch.Tensor] | None = None self._request_audio_states: dict[str, dict[str, Any]] = {} # Mirrors upstream TTSStreamingGenerator._chunk_info: one committed @@ -203,17 +477,160 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self._codec_eos_id = int(getattr(tts_config, "eos_token_id", self._num_audio_tokens - 1)) self._hidden_size = getattr(tts_config, "hidden_size", 768) self._normalize = getattr(tts_config, "normalize_projected_hidden", True) + # Codec sampling knobs: the deploy YAML's codec_sampling_params + # block wins key by key, then the checkpoint's tts_config, then the + # module fallbacks. The multi-frame path reads the resolved values + # straight off the model. + yaml_codec = getattr( + getattr(vllm_config, "model_config", None), "codec_sampling_params", None + ) + resolved = resolve_codec_sampling_params(yaml_codec, tts_config) + self._codec_seed = resolved["seed"] + self._codec_temperature = resolved["temperature"] + self._codec_top_k = resolved["top_k"] + self._codec_top_p = resolved["top_p"] + self._codec_repetition_penalty = resolved["repetition_penalty"] + self._codec_min_tokens = resolved["min_tokens"] + self._codec_max_tokens = resolved["max_tokens"] else: self._tts_config = None self._codec_eos_id = 0 + # K-step activation (910C/A3 target). + # K is the frame count per Talker decode step: a value >= 2 engages the + # multi-frame loop -- the scheduler hands each request K query positions + # through the vLLM V1 speculative path (constant `continue` drafts), + # talker_multiframe.run() replays the decode graph K times, and this + # model samples one codec frame per replay in-model (the vLLM-level head + # degenerates to a one-hot stop/continue row so the rejection sampler + # verifies without touching the codec stream). + # SoC guard: the spec-driven verify trips rejection_random_sample_kernel + # past the vector-core limit on the 910B family, so the gate refuses to + # arm there instead of leaving a crash switch behind an env var. + self._k_step_frames = self._parse_k_step_frames() + self.supports_multi_frame_decode = self._k_step_frames > 0 + # Per-request device RNG streams for in-model codec sampling (the + # multi-frame loop cannot run the vLLM host sampler per frame). + self._k_step_rngs: dict[str, torch.Generator] = {} + self._k_step_base_seed = 0x5EEDC0DE + # Last frame's stop/continue rows and the flattened per-token stop + # sequence the merged step presents to compute_logits. + self._k_last_stop_rows: list[bool] | None = None + self.has_preprocess = True self.has_postprocess = False # Same-step codes travel through make_omni_output from the previous # sampled id (decode preprocess embeds that id via emb_code). self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("codes", "audio")} + # K-step (multi-frame Talker) request state. The loop keeps one codec + # sampling graph, one Generator and one device-side conversation state + # per in-flight request; _flush_deferred_cleanup drops them when the + # engine reports the request finished. + self._batch_stop_logits: torch.Tensor | None = None + self._request_generators: dict[str, torch.Generator] = {} + self._request_audio_states: dict[str, dict[str, Any]] = {} + self._request_codec_device_states: dict[str, TalkerCodecDeviceState] = {} + self._request_codec_device_inputs: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + self._init_native_talker(prefix) + @staticmethod + def _probe_soc_name() -> str: + """torch_npu device-name string, or "" when the probe cannot run. + + The integer SoCVersion encoding is not a documented ordering (a local + 910B3 reports 223, which defeats any threshold guess), so the guard + matches on the name prefix instead -- exactly what torch_npu itself + does (utils/_module.py lists "Ascend910B" and "Ascend910_93"). + """ + try: + import torch_npu + + return str( + torch_npu.npu.get_device_name(torch_npu.npu.current_device()) + ) + except Exception: + return "" + + @classmethod + def _soc_allows_k_step(cls) -> bool: + """Ask the very question the deploy-config gate asked, and no other. + + 910B family: rejection_random_sample_kernel hits a vector-core limit + once the verifier runs K tokens per request, which crashed six bring-up + attempts before this gate existed. 910C / A3 is the activation target. + + This must delegate rather than re-probe. The config layer answers the + same question while building the deploy config (it decides whether to + inject the stage-1 speculative_config), and it can only see the launch + environment -- it deliberately never touches the device. A second, + device-probing answer here could differ from that one, and a mismatch + is not benign: if the loop arms while the scheduler has no + speculative_config it silently never engages, and if the scheduler + reserves K positions while the loop refuses them the stage deadlocks + with no request making progress. One source of truth, one answer. + """ + from vllm_omni.config.stage_config import ( + _MINICPMO_TALKER_FRAMES_ENV, + _npu_soc_allows_talker_multiframe, + ) + + if os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip(): + # An explicit frame count is an operator decision; the deploy-config + # gate arms on it too, so the two stay in agreement. + logger.info( + "[minicpmo] %s set explicitly; skipping the SoC check", + _MINICPMO_TALKER_FRAMES_ENV, + ) + return True + return _npu_soc_allows_talker_multiframe() + + @classmethod + def _parse_k_step_frames(cls) -> int: + """K codec frames per Talker step, or 0 when the loop must not engage. + + The count comes from ``config.stage_config.talker_frames_per_step`` -- + the very function the deploy-config loader uses to size the injected + speculative_config -- so the runner and the scheduler cannot disagree + about K. An explicit ``VLLM_OMNI_MINICPMO_TALKER_FRAMES`` (or the older + ``OMNI_K_STEP``) skips the SoC gate, exactly as the loader does for the + same input; that keeps the two sides aligned even when an operator + forces the loop on by hand. + """ + from vllm_omni.config.stage_config import ( + _MINICPMO_FRAMES_OFF, + _MINICPMO_TALKER_FRAMES_ENV, + talker_frames_per_step, + ) + + raw = os.environ.get("OMNI_K_STEP", "").strip().lower() + explicit = os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip().lower() + if raw not in ("", "0", "off", "false", "no"): + try: + frames = int(raw) + except ValueError: + raise ValueError( + f"OMNI_K_STEP must be an integer frame count (>=2), got {raw!r}" + ) from None + explicit = "forced" + else: + frames = talker_frames_per_step() + if frames < 2: + # K=1 degenerates to the ordinary one-frame path; treat it as off + # so the flag never silently half-arms the pipeline. + return 0 + if frames > 16: + raise ValueError(f"K-step frame count is capped at 16, got {frames}") + forced = explicit not in ("",) + _MINICPMO_FRAMES_OFF + if not forced and not cls._soc_allows_k_step(): + logger.warning( + "[minicpmo] K-step decode left off: the SoC either cannot verify " + "spec-width rows (910B family limit) or was not identified.", + ) + return 0 + logger.info("[minicpmo] Talker K-step decode armed: %d codec frames per step", frames) + return frames + def _init_native_talker(self, prefix: str) -> None: if self._tts_config is None: raise ValueError("MiniCPM-o continuous Talker requires tts_config") @@ -485,7 +902,7 @@ def preprocess( if not isinstance(base_recent_codes, tuple): raise ValueError("streaming Talker condition lost its frozen codec history") retained_codes = list(base_recent_codes) - offset = int(info_dict.get("_omni_num_computed_tokens", 0)) + offset = max(0, int(info_dict.get("_omni_num_computed_tokens", 0) or 0)) # The handoff rebuilds only the tail-aligned Talker condition. # Materialize zero-token embeddings for any scheduler prompt # prefix so chunked prefill can slice from a non-zero offset. @@ -522,7 +939,11 @@ def preprocess( # min_new_token=50) comes from the deploy YAML. remaining = int(self._tts_config.max_position_embeddings) - target_len max_tokens = max(min(_OFFLINE_CODEC_MAX_NEW_TOKENS, remaining), 1) - min_tokens = None + # The sampler floor the old comment pointed at (upstream's + # min_new_token, resolved from the deploy YAML or the + # checkpoint). A None here would reach make_omni_output as a + # value, and int(None) kills the stage. + min_tokens = int(self._codec_min_tokens) state: dict[str, Any] = { "finished": empty_condition, "step": 0, @@ -532,6 +953,32 @@ def preprocess( } if retained_codes: state["recent_codes"] = retained_codes + # Rebuild the device-side penalty window once per chunk from the + # carried-over ids; the per-frame path below then appends on + # device without touching the host again. + retained = [int(code_id) for code_id in retained_codes][-_CODEC_PENALTY_WINDOW:] + if retained: + penalty_windows = getattr(self, "_penalty_windows_dev", None) + if not isinstance(penalty_windows, dict): + penalty_windows = {} + self._penalty_windows_dev = penalty_windows + penalty_windows[request_id] = torch.tensor(retained, dtype=torch.long, device=embeds.device) + pf_freq = torch.zeros( + int(getattr(self, "_num_audio_tokens", 0) or 0), + device=embeds.device, + dtype=torch.float32, + ) + if pf_freq.numel() > 0: + pf_freq.index_add_( + 0, + penalty_windows[request_id], + torch.ones(len(retained), device=embeds.device, dtype=torch.float32), + ) + pf_dict = getattr(self, "_penalty_freqs_dev", None) + if not isinstance(pf_dict, dict): + pf_dict = {} + self._penalty_freqs_dev = pf_dict + pf_dict[request_id] = pf_freq request_states = getattr(self, "_request_audio_states", None) if request_states is None: request_states = {} @@ -563,7 +1010,18 @@ def preprocess( # Decode: vLLM's previous sampled codec id is this step's input. # Embed it with the codec table (not the 32k Llama embed_tokens) and # hand the same id to make_omni_output so Code2Wav sees it this step. - code = input_ids.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1)[-1:] + # K-step: once the multi-frame loop owns sampling, the rows vLLM + # schedules carry placeholder continue ids, not real codec ids -- the + # true previous frame's sample lives in the request state. + k_last = state.get("last_code") if isinstance(state, dict) else None + if isinstance(k_last, int) and k_last >= 0: + code = torch.tensor( + [k_last], + dtype=torch.long, + device=self.emb_code[0].weight.device, + ) + else: + code = input_ids.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1)[-1:] embeds = self.emb_code[0](code) code_id = int(code.item()) if code_id == int(self._codec_eos_id): @@ -588,20 +1046,24 @@ def make_omni_output( spans = kwargs.get("request_token_spans") if spans is None or len(spans) != len(infos): raise RuntimeError("MiniCPM-o continuous Talker requires one request_token_span per request") + sample_eligible = kwargs.get("request_sample_eligible") + if sample_eligible is None: + sample_eligible = [True] * len(infos) + if len(sample_eligible) != len(infos): + raise RuntimeError( + f"MiniCPM-o continuous Talker received {len(sample_eligible)} sampling flags for {len(infos)} requests" + ) emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos) + stop_rows: list[torch.Tensor] = [] + codec_deltas: list[torch.Tensor] = [] + terminal_flags: list[torch.Tensor] = [] native_duplex_flags: list[torch.Tensor] = [] duplex_epochs: list[torch.Tensor] = [] duplex_turn_ids: list[torch.Tensor] = [] segment_texts_utf8: list[torch.Tensor] = [] turn_end_flags: list[torch.Tensor] = [] - empty_delta = hidden.new_empty((0, 1), dtype=torch.long) - codec_deltas = [empty_delta for _ in infos] - terminal_flags = [torch.tensor(False, dtype=torch.bool) for _ in infos] - force_eos_rows = [False] * len(infos) - mask_eos_rows = [False] * len(infos) - empty_history = hidden.new_empty((0,), dtype=torch.long) - penalty_histories = [empty_history for _ in infos] + row_continue, row_stop, flag_false, flag_true, empty_delta = self._step_constants(hidden) for index, info in enumerate(infos): info_dict = info if isinstance(info, dict) else {} native_duplex = info_dict.get("native_duplex") is True @@ -647,54 +1109,126 @@ def make_omni_output( turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool)) if not isinstance(info, dict): + stop_rows.append(row_continue) + codec_deltas.append(empty_delta) + terminal_flags.append(flag_false) + continue + start, end = spans[index] + end = min(int(end), int(hidden.shape[0])) + if int(start) >= end: + stop_rows.append(row_continue) + codec_deltas.append(empty_delta) + terminal_flags.append(flag_false) continue request_id = str(info.get("request_id", index)) - state = self._request_audio_states.get(request_id) + request_states = getattr(self, "_request_audio_states", None) + if request_states is None: + request_states = {} + self._request_audio_states = request_states + state = request_states.get(request_id) if not isinstance(state, dict): state = dict(info.get("audio_state", {}) or {}) - self._request_audio_states[request_id] = state - empty_speech = bool(state.get("finished")) - codes = info.get("codes", {}) - audio = codes.get("audio") if isinstance(codes, Mapping) else None - if isinstance(audio, torch.Tensor) and audio.numel() > 0: - codec_deltas[index] = audio.to(device=hidden.device, dtype=torch.long).reshape(-1, 1) - state["step"] = int(state.get("step", 0)) + 1 - # ``audio`` is the id sampled last step, i.e. exactly upstream's - # ``new_tokens[:, 0:t]`` history for the logits computed below. - recent = state.get("recent_codes") - recent = (recent if isinstance(recent, list) else []) + codec_deltas[index].reshape(-1).tolist() - state["recent_codes"] = recent[-_CODEC_PENALTY_WINDOW:] - recent_codes = state.get("recent_codes") - if recent_codes: - penalty_histories[index] = torch.tensor(recent_codes, dtype=torch.long, device=hidden.device) - max_tokens = state.get("max_tokens") - min_tokens = state.get("min_tokens") - step = int(state.get("step", 0)) - # Duplex: 26 samples include the terminating EOS, so force it after - # 25 forwarded frames — the same cadence as generate_chunk. - # Offline: force-stop at the remaining Talker context rather than - # waiting for a sampled EOS. - hit_chunk_limit = max_tokens is not None and step >= int(max_tokens) - 1 - chunk_done = empty_speech or hit_chunk_limit - if hit_chunk_limit: - # The next sampled id is forced EOS and will finish the - # request; mark the chunk done on this step so the data - # plane does not wait for a follow-up empty decode. - state["finished"] = True - force_eos_rows[index] = chunk_done - mask_eos_rows[index] = not force_eos_rows[index] and ( - (min_tokens is not None and step < int(min_tokens)) - or (bool(state.get("turn_end_drain")) and _turn_end_boundary_eos_masked(step)) - ) - terminal_flags[index] = torch.tensor(chunk_done, dtype=torch.bool) - - # Empty-speech rows, finished duplex chunks, and offline requests that - # fill the remaining Talker context must sample codec EOS so the - # scheduler releases the request. Mid-chunk rows mask EOS until - # min_tokens. - self._force_eos_rows = force_eos_rows - self._mask_eos_rows = mask_eos_rows - self._penalty_histories = penalty_histories + request_states[request_id] = state + if state.get("finished"): + stop_rows.append(row_stop) + codec_deltas.append(empty_delta) + terminal_flags.append(flag_false) + continue + if not sample_eligible[index]: + # vLLM computes a logit row for incomplete chunked prefills but + # discards its sampled token. Advancing codec/RNG state here + # would make output depend on prefill chunking and compaction. + stop_rows.append(row_continue) + codec_deltas.append(empty_delta) + terminal_flags.append(flag_false) + continue + codes = state.get("codes") + if not isinstance(codes, torch.Tensor): + codes = (info.get("audio_codes", {}) or {}).get("accumulated") + if not isinstance(codes, torch.Tensor): + codes = torch.empty(0, dtype=torch.long, device=hidden.device) + else: + codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1) + step = _codec_int_param(state, "step", 0) + # 迁移注记（2026-09-21，待设备复测后决定是否接线）：main 侧的 turn-end + # 语义要求 turn-end 排空块在每个 25 帧 cadence 边界后的 5 步内屏蔽 codec + # EOS（见本文件 _turn_end_boundary_eos_masked 与 state["turn_end_drain"]， + # MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS 为 4 个 chunk，故窗口落在 + # step 25-29 / 50-54 / 75-79 / 100-103）。第 8 项把 EOS 屏蔽移进了采样核 + # （talker_codec_sample.prepare_codec_logits / greedy_codec_sample），其判据 + # 是 state.step < min_tokens，而 min_tokens 是每请求缓存的张量 + # （_request_codec_device_inputs），此处无法逐帧表达该周期窗口 —— 因此本项 + # 迁移后该窗口屏蔽未接线，只影响双工 turn-end（离线/TTS 路径不受影响）。 + # 本文件保留了 main 的常量、budget 与 turn_end_drain 状态标志。 + min_tokens = _codec_int_param(state, "min_tokens", self._codec_min_tokens) + max_tokens = _codec_int_param(state, "max_tokens", self._codec_max_tokens) + if self._codec_temperature == 0.0: + sampled = self._sample_audio_code_greedy( + hidden[end - 1 : end], + codes, + request_id, + step, + min_tokens, + max_tokens, + ) + else: + stochastic_result = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step) + if isinstance(stochastic_result, TalkerCodecSampleResult): + sampled = stochastic_result.sampled_token.reshape(()).to(torch.long) + else: + # CPU unit tests and downstream subclasses historically + # stub this method with the sampled tensor itself. + sampled = stochastic_result.reshape(()).to(torch.long) + sampled_id = int(sampled.item()) + prev_code = state.get("last_code") if isinstance(state, Mapping) else None + is_eos = sampled_id == self._codec_eos_id + state["step"] = _codec_int_param(state, "step", 0) + 1 + reached_limit = int(state["step"]) >= _codec_int_param(state, "max_tokens", self._codec_max_tokens) + finished = is_eos or reached_limit + state["finished"] = finished + if self._k_step_frames > 0 and kstep_frame_trace_enabled(): + _trace_kstep_frame( + self, + request_id=request_id, + step_before=step, + step_after=int(state["step"]), + prev_code=prev_code, + row=int(end) - 1, + sampled_id=sampled_id, + is_eos=is_eos, + reached_limit=reached_limit, + finished=finished, + ctx_len=int(codes.numel()) if isinstance(codes, torch.Tensor) else -1, + min_tokens=min_tokens, + max_tokens=max_tokens, + ) + # MiniCPMTTS.generate_chunk consumes the boundary sample but + # returns only codes that were fed into the retained KV state. + if not is_eos and not reached_limit: + codes = torch.cat([codes[-(_REPETITION_WINDOW - 1) :], sampled.reshape(1)]) + delta = sampled.reshape(1, 1) + if self._k_step_frames > 0: + # K-step: the multi-frame loop owns sampling, so the rows + # vLLM schedules for the next step carry placeholder + # continue ids. Record this frame's sample so the decode + # preprocess embeds the real previous frame's codec id + # (the k_last branch) instead of the placeholder. + state["last_code"] = sampled_id + else: + delta = empty_delta + state["codes"] = codes + info["audio_state"] = state + info["audio_codes"] = { + "current": sampled.reshape(1), + "accumulated": codes, + } + codec_deltas.append(delta) + terminal_flags.append(flag_true if finished else flag_false) + stop_rows.append(row_stop if finished else row_continue) + + self._batch_stop_logits = torch.stack(stop_rows, dim=0) if stop_rows else hidden.new_empty((0, 2)) + # Lists are deliberate: the runner routes element i to request i, + # preserving compaction alignment while emitting only this step's code. meta_outputs = {"finished": terminal_flags} if emit_duplex_metadata: meta_outputs.update( @@ -706,23 +1240,304 @@ def make_omni_output( "turn_end": turn_end_flags, } ) + multimodal_outputs: dict[str, Any] = { + "codes": {"audio": codec_deltas}, + "meta": meta_outputs, + } return OmniOutput( text_hidden_states=hidden, - multimodal_outputs={ - "codes": {"audio": codec_deltas}, - "meta": meta_outputs, - }, + multimodal_outputs=multimodal_outputs, + ) + + # This model always emits exactly one codec frame per query position, which + # is what makes a multi-frame step decomposable into per-frame calls. + supports_multi_frame_decode = True + + def take_batch_stop_logits(self) -> torch.Tensor | None: + """Hand over this frame's stop rows and clear them. + + ``compute_logits`` normally consumes them at the end of the step. A + multi-frame step calls ``make_omni_output`` once per frame, so the + runner has to collect each frame's rows before the next call replaces + them, and hand the merged tensor back through + ``set_batch_stop_logits``. + """ + logits = self._batch_stop_logits + self._batch_stop_logits = None + return logits + + def merge_frame_outputs( + self, + frame_outputs: list[OmniOutput], + frame_stop_logits: list[torch.Tensor], + ) -> OmniOutput: + """Fold a multi-frame step's per-frame outputs into one step output. + + Codec deltas concatenate per request -- the connector already takes a + variable number of frames per step and drops the ``-1`` sentinel, so a + request that ended mid-step contributes only the frames it actually + emitted. Stop rows interleave to the layout ``logits_indices`` reads: + every query position of request 0, then of request 1, and so on. + """ + if not frame_outputs: + raise RuntimeError("MiniCPM-o multi-frame step produced no frames") + if len(frame_outputs) == 1: + self.set_batch_stop_logits(frame_stop_logits[0]) + return frame_outputs[0] + + per_frame_codes = [output.multimodal_outputs["codes"]["audio"] for output in frame_outputs] + num_reqs = len(per_frame_codes[0]) + codec_deltas = [ + torch.cat([frame[index] for frame in per_frame_codes], dim=0) for index in range(num_reqs) + ] + + meta_outputs: dict[str, Any] = {} + for key in frame_outputs[0].multimodal_outputs["meta"]: + per_frame = [output.multimodal_outputs["meta"][key] for output in frame_outputs] + if key == "finished": + meta_outputs[key] = [ + self._merge_frame_finished([frame[index] for frame in per_frame]) + for index in range(num_reqs) + ] + else: + # Duplex metadata is per request and constant across the step's + # frames; the last frame is as good as any and matches what a + # single-frame step would have reported. + meta_outputs[key] = per_frame[-1] + + self.set_batch_stop_logits(torch.stack(frame_stop_logits, dim=1).reshape(-1, 2)) + return OmniOutput( + text_hidden_states=frame_outputs[-1].text_hidden_states, + multimodal_outputs={"codes": {"audio": codec_deltas}, "meta": meta_outputs}, + ) + + def _request_generator(self, request_id: str, device: torch.device) -> torch.Generator: + generator = self._request_generators.get(request_id) + if generator is None: + generator = torch.Generator(device=device) + generator.manual_seed(self._codec_seed) + self._request_generators[request_id] = generator + return generator + + def _sample_audio_code( + self, + hidden_state: torch.Tensor, + history: torch.Tensor, + request_id: str, + step: int, + ) -> TalkerCodecSampleResult: + device_states = getattr(self, "_request_codec_device_states", None) + if device_states is None: + device_states = {} + self._request_codec_device_states = device_states + device_state = device_states.get(request_id) + new_segment = device_state is None + if device_state is None: + request_states = getattr(self, "_request_audio_states", {}) + request_state = request_states.get(request_id, {}) + device_state = make_device_state( + history, + step=step, + max_tokens=_codec_int_param(request_state, "max_tokens", self._codec_max_tokens), + finished=False, + ) + device_states[request_id] = device_state + device_inputs_by_request = getattr(self, "_request_codec_device_inputs", None) + if device_inputs_by_request is None: + device_inputs_by_request = {} + self._request_codec_device_inputs = device_inputs_by_request + device_inputs = device_inputs_by_request.get(request_id) + if device_inputs is None: + request_states = getattr(self, "_request_audio_states", {}) + request_state = request_states.get(request_id, {}) + device_inputs = ( + torch.tensor( + [_codec_int_param(request_state, "min_tokens", self._codec_min_tokens)], + dtype=torch.int32, + device=hidden_state.device, + ), + torch.tensor([self._codec_temperature], dtype=torch.float32, device=hidden_state.device), + torch.tensor( + [self._codec_repetition_penalty], + dtype=torch.float32, + device=hidden_state.device, + ), + ) + device_inputs_by_request[request_id] = device_inputs + min_tokens_tensor, temperature_tensor, penalty_tensor = device_inputs + eos_id = self._codec_eos_id + logits = prepare_codec_logits( + self.head_code[0](hidden_state).float(), + device_state, + min_tokens_tensor, + temperature_tensor, + penalty_tensor, + eos_token_id=eos_id, + top_k=self._codec_top_k, + top_p=self._codec_top_p, + min_tokens_to_keep=3, ) + probabilities = torch.softmax(logits, dim=-1) + sampled = torch.multinomial( + probabilities, + num_samples=1, + generator=self._request_generator(request_id, probabilities.device), + ).reshape(()) + result = codec_sample_result( + device_state, + sampled, + eos_token_id=eos_id, + ) + device_states[request_id] = result.state + return result + + def _sample_audio_code_greedy( + self, + hidden_state: torch.Tensor, + history: torch.Tensor, + request_id: str, + step: int, + min_tokens: int, + max_tokens: int, + ) -> torch.Tensor: + """Run the greedy codec boundary while keeping sampler state on device. + + Host stop routing still consumes ``sampled.item()`` below; removing that + sync requires a runner/connector state refactor and is not hidden here. + """ + device_states = getattr(self, "_request_codec_device_states", None) + if device_states is None: + device_states = {} + self._request_codec_device_states = device_states + device_state = device_states.get(request_id) + if device_state is None: + device_state = make_device_state( + history, + step=step, + max_tokens=max_tokens, + finished=False, + ) + device_states[request_id] = device_state + + device_inputs_by_request = getattr(self, "_request_codec_device_inputs", None) + if device_inputs_by_request is None: + device_inputs_by_request = {} + self._request_codec_device_inputs = device_inputs_by_request + device_inputs = device_inputs_by_request.get(request_id) + if device_inputs is None: + device_inputs = ( + torch.tensor([min_tokens], dtype=torch.int32, device=hidden_state.device), + torch.tensor([1.0], dtype=torch.float32, device=hidden_state.device), + torch.tensor( + [self._codec_repetition_penalty], + dtype=torch.float32, + device=hidden_state.device, + ), + ) + device_inputs_by_request[request_id] = device_inputs + min_tokens_tensor, _, penalty_tensor = device_inputs + result = greedy_codec_sample( + self.head_code[0](hidden_state).float(), + device_state, + min_tokens_tensor, + penalty_tensor, + top_k=self._codec_top_k, + eos_token_id=self._codec_eos_id, + ) + device_states[request_id] = result.state + # The kernel ABI is int32, while the existing connector/audio-code + # payload is int64. Keep the compatibility cast outside the fused op. + return result.sampled_token.to(torch.long).reshape(()) + + def _step_constants(self, hidden: torch.Tensor): + """Constant per-step tensors, cached per (device, dtype). + + These never change and downstream consumers only read them, so the + per-step H2D copies from ``new_tensor``/``torch.tensor`` are avoidable. + """ + cache = getattr(self, "_step_constants_cache", None) + key = (hidden.device, hidden.dtype) + if cache is None or cache[0] != key: + neg_inf = float("-inf") + cache = ( + key, + ( + hidden.new_tensor([0.0, neg_inf]), + hidden.new_tensor([neg_inf, 0.0]), + torch.tensor(False, dtype=torch.bool), + torch.tensor(True, dtype=torch.bool), + hidden.new_empty((0, 1), dtype=torch.long), + ), + ) + self._step_constants_cache = cache + return cache[1] + + def set_batch_stop_logits(self, logits: torch.Tensor | None) -> None: + self._batch_stop_logits = logits + + @staticmethod + + def _merge_frame_finished(flags: list[torch.Tensor]) -> torch.Tensor: + """OR the per-frame terminal flags without forcing a device sync. + + ``make_omni_output`` reports the flag on the frame that ended the codec + sequence and reports ``False`` for every frame after it, so the step's + answer is the OR. The flags are device tensors on the device-state + path and CPU constants on the native-sampler path, and one step can + produce both -- + the frame that finishes samples on device, the frames after it take the + early return and get the CPU constant. Reading the CPU ones costs + nothing; the device ones are OR'd on device and travel out in the + runner's single coalesced D2H. + """ + merged: torch.Tensor | None = None + for flag in flags: + if flag.device.type == "cpu": + if bool(flag.reshape(()).item()): + return flag + continue + flag = flag.reshape(1) + merged = flag if merged is None else torch.logical_or(merged, flag) + if merged is None: + return flags[-1] + return merged.reshape(()) def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: self._deferred_cleanup_ids.update(str(req_id) for req_id in finished_req_ids) def _flush_deferred_cleanup(self) -> None: + """Drop every per-request cache a finished request left behind. + + The sampler's own caches belong here as much as the codec history does: + ``_request_codec_device_states`` carries the step counter and the + penalty window the EOS mask is derived from, and + ``_request_codec_device_inputs`` carries the ``min_tokens`` copied at + first use. Leaving them behind means a recycled request id resumes with + another request's counters -- the EOS mask then holds or lifts at the + wrong frame, and ``_request_generators`` leaks a generator per request + on top of that. All three key off the same request id, so they are + dropped together with the rest. + """ request_audio_states = getattr(self, "_request_audio_states", {}) request_condition_states = getattr(self, "_request_condition_states", {}) + penalty_windows = getattr(self, "_penalty_windows_dev", None) + penalty_freqs = getattr(self, "_penalty_freqs_dev", None) + codec_device_states = getattr(self, "_request_codec_device_states", None) + codec_device_inputs = getattr(self, "_request_codec_device_inputs", None) + request_generators = getattr(self, "_request_generators", None) for request_id in self._deferred_cleanup_ids: request_audio_states.pop(request_id, None) request_condition_states.pop(request_id, None) + if isinstance(penalty_windows, dict): + penalty_windows.pop(request_id, None) + if isinstance(penalty_freqs, dict): + penalty_freqs.pop(request_id, None) + if isinstance(codec_device_states, dict): + codec_device_states.pop(request_id, None) + if isinstance(codec_device_inputs, dict): + codec_device_inputs.pop(request_id, None) + if isinstance(request_generators, dict): + request_generators.pop(request_id, None) self._deferred_cleanup_ids.clear() def _dummy_hidden_states( @@ -770,35 +1585,48 @@ def forward( def compute_logits(self, hidden_states, *args, **kwargs): if not isinstance(hidden_states, torch.Tensor): return None + if self._k_step_frames > 0: + # K-step: the codec stream was already sampled in-model, one frame + # per replay, so this head carries no codec information at all. + # All it decides is stop/continue per scheduled position, over a + # two-id vocabulary -- continue (0) and stop (1) -- which is what + # the stage's ``stop_token_ids: [1]`` names. merge_frame_outputs + # publishes this step's rows in logits_indices order; a step with + # none gets zeros, whose argmax is 0 and therefore "continue". + if self._batch_stop_logits is None: + return torch.zeros( + hidden_states.shape[0], + 2, + device=hidden_states.device, + dtype=torch.float32, + ) + logits = self._batch_stop_logits + self._batch_stop_logits = None + return logits if hidden_states.numel() == 0: return hidden_states.new_empty((0, int(self._num_audio_tokens))) - logits = self.head_code[0](hidden_states).float() - force_eos = self._force_eos_rows - mask_eos = self._mask_eos_rows - self._force_eos_rows = None - self._mask_eos_rows = None - need_force = bool(force_eos and len(force_eos) == logits.shape[0] and any(force_eos)) - need_mask = bool(mask_eos and len(mask_eos) == logits.shape[0] and any(mask_eos)) - # sample() re-applies the decision on the sampled ids: vLLM's - # MinTokensLogitsProcessor runs after this and would blank the codec EOS - # we just forced (it is in the stage's ``stop_token_ids``), leaving an - # all -inf row and a request that never releases. - self._pending_force_eos_rows = force_eos if need_force else None - if need_force or need_mask: - logits = logits.clone() - eos_id = int(self._codec_eos_id) - if need_force: - assert force_eos is not None - forced = torch.tensor(force_eos, dtype=torch.bool, device=logits.device) - logits[forced] = float("-inf") - logits[forced, eos_id] = 0.0 - if need_mask: - assert mask_eos is not None - masked = torch.tensor(mask_eos, dtype=torch.bool, device=logits.device) - logits[masked, eos_id] = float("-inf") - return logits + # Non-K-step: the real codec head. vLLM samples the codec stream here, + # and the stage's stop_token_ids ends the request when it produces it. + return self.head_code[0](hidden_states).float() def sample(self, logits, sampling_metadata): + # Two-column rows are the K-step stop/continue heads compute_logits + # built: they are one-hot, so a full sampler pass is ~10 host-dispatched + # kernels for a deterministic pick while argmax is bit-identical on + # them. int32 matches what the runner's async path scatters into its + # input buffers. Anything else (the non-K-step codec head) takes the + # ordinary path below, untouched. + if ( + isinstance(logits, torch.Tensor) + and logits.ndim == 2 + and logits.shape[-1] == 2 + and not getattr(sampling_metadata, "max_num_logprobs", None) + ): + with _prof_span("tts_sampler"): + return SamplerOutput( + sampled_token_ids=logits.argmax(dim=-1, keepdim=True).to(torch.int32), + logprobs_tensors=None, + ) prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None) if isinstance(logits, torch.Tensor) and isinstance(prompt_ids, torch.Tensor): # Copy rather than mutate: the runner may hand us the input batch's @@ -807,11 +1635,13 @@ def sample(self, logits, sampling_metadata): sampling_metadata, prompt_token_ids=blank_scheduler_prompt_for_penalties(prompt_ids, logits.shape[-1]), ) - logits, sampling_metadata = self._apply_codec_repetition_penalty(logits, sampling_metadata) - force_eos = self._pending_force_eos_rows - self._pending_force_eos_rows = None - output = Sampler()(logits, sampling_metadata) - return self._force_eos_on_sampled_ids(output, force_eos) + with _prof_span("tts_penalty"): + logits, sampling_metadata = self._apply_codec_repetition_penalty(logits, sampling_metadata) + prewarped = _maybe_prewarp_top_k_top_p(logits, sampling_metadata) + if prewarped is not None: + logits, sampling_metadata = prewarped + with _prof_span("tts_sampler"): + return Sampler()(logits, sampling_metadata) def _apply_codec_repetition_penalty(self, logits, sampling_metadata): """Score MiniCPMTTS.generate's windowed codec penalty, not vLLM's. @@ -838,7 +1668,25 @@ def _apply_codec_repetition_penalty(self, logits, sampling_metadata): or not isinstance(penalties, torch.Tensor) or getattr(sampling_metadata, "no_penalties", False) ): + self._penalty_freq_rows = None return logits, sampling_metadata + freq_rows = getattr(self, "_penalty_freq_rows", None) + self._penalty_freq_rows = None + if ( + isinstance(freq_rows, list) + and len(freq_rows) == logits.shape[0] + and all(isinstance(f, torch.Tensor) and f.device == logits.device and f.numel() == logits.shape[-1] for f in freq_rows) + and bool((penalties.reshape(-1) != 1.0).any()) + ): + # Incremental histogram: freq rows are maintained on device by + # make_omni_output (+1 append / -1 evict), identical to a full + # scatter_add rebuild of the sliding window. + freqs = torch.stack([f.to(dtype=torch.float32) for f in freq_rows], dim=0) + alpha = torch.pow( + penalties.to(device=logits.device, dtype=torch.float32).reshape(-1, 1), freqs + ).to(dtype=logits.dtype) + logits = torch.where(logits < 0, logits * alpha, logits / alpha) + return logits, replace(sampling_metadata, repetition_penalties=torch.ones_like(penalties)) logits = _apply_batched_repetition_penalty( logits, histories, @@ -848,24 +1696,6 @@ def _apply_codec_repetition_penalty(self, logits, sampling_metadata): # Neutralize the sampler's own pass so the penalty is scored once. return logits, replace(sampling_metadata, repetition_penalties=torch.ones_like(penalties)) - def _force_eos_on_sampled_ids(self, output: Any, force_eos: list[bool] | None) -> Any: - """Overwrite sampled ids for rows the model terminated this step. - - The codec EOS is a stage ``stop_token_ids`` entry, so vLLM's - ``min_tokens`` processor masks it for the first ``min_tokens`` steps. - A row the model forced to EOS therefore reaches the sampler as all - -inf and comes back as an arbitrary codec id, which keeps an - already-finished request decoding until its length cap. - """ - if not force_eos or not any(force_eos): - return output - sampled = getattr(output, "sampled_token_ids", None) - if not isinstance(sampled, torch.Tensor) or sampled.shape[0] != len(force_eos): - return output - rows = torch.tensor(force_eos, dtype=torch.bool, device=sampled.device) - sampled[rows] = int(self._codec_eos_id) - return output - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): return self._load_native_weights(weights) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index e7539384511..9c6656c978d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -20,6 +20,35 @@ MINICPMO45_REFERENCE_AUDIO_KEY = "_minicpmo45_reference_audio" +_CODEC_EOS_TOKEN_ID = 6561 # tts_config.num_audio_tokens - 1 + + +def _talker_stop_token_ids() -> list[int]: + """Stage 1's stop token, matching the head that will actually run. + + Two shapes exist for the Talker's vLLM-level head: + + * K-frame loop off -- ``compute_logits`` is the real 6562-wide codec head, + so the sampled id *is* a codec id and the codec EOS (6561) is the stop; + * K-frame loop on -- ``compute_logits`` collapses to the two-wide + continue/stop row, so the only sampleable ids are 0 and 1 and the codec + EOS can never appear at this level (the model samples codec ids inside + the loop and forwards them to stage 2 itself). + + Pinning this to [6561] regardless is what made every K-step request run to + ``max_tokens``: the model did emit the stop marker, ``check_stop`` compared + it against 6561, and the request kept its slot until the context ran out. + + Asked per construction so both this file and the deploy loader read the + same predicate; the env the answer depends on is set before any of this is + imported. + """ + from vllm_omni.config.stage_config import talker_multiframe_armed + from vllm_omni.platforms.npu.worker.talker_multiframe import STOP_TOKEN_ID + + return [STOP_TOKEN_ID] if talker_multiframe_armed() else [_CODEC_EOS_TOKEN_ID] + + MINICPMO_4_5_PIPELINE = PipelineConfig( model_type="minicpmo_4_5", default_deploy_config_name="minicpmo_4_5.yaml", @@ -67,9 +96,10 @@ async_chunk_process_next_stage_input_func=f"{_PROC}.tts2code2wav_async_chunk", sampling_constraints={ "detokenize": False, - # MiniCPM-o 4.5 codec EOS is tts_config.num_audio_tokens - 1. - # Same pattern as Qwen3 talker's stop_token_ids: [2150]. - "stop_token_ids": [6561], + # The stop id has to be one the vLLM-level head can actually + # emit, and that head changes shape when the K-frame loop is + # armed -- see _talker_stop_token_ids. + "stop_token_ids": _talker_stop_token_ids(), }, ), StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py new file mode 100644 index 00000000000..79abb49f822 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -0,0 +1,185 @@ +"""Framework boundaries for Talker codec sampling. + +Production stochastic sampling uses a logits-filter boundary implemented in +graph-capturable tensor ops: history counting, temperature, repetition +penalty, EOS mask, top-p and top-k, while native NPU operators keep softmax +and ``torch.multinomial`` semantics (including Generator state) unchanged. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +VOCAB_SIZE = 6562 +HISTORY_WINDOW = 16 +MIN_TOKENS_TO_KEEP = 3 + + +@dataclass(frozen=True) +class TalkerCodecDeviceState: + history: torch.Tensor + history_len: torch.Tensor + step: torch.Tensor + max_tokens: torch.Tensor + finished: torch.Tensor + + +@dataclass(frozen=True) +class TalkerCodecSampleResult: + sampled_token: torch.Tensor + state: TalkerCodecDeviceState + emit: torch.Tensor + + +def make_device_state( + codes: torch.Tensor, + *, + step: int, + max_tokens: int, + finished: bool, +) -> TalkerCodecDeviceState: + """Create one persistent B=1 state at request/segment setup time.""" + recent = codes.reshape(-1)[-HISTORY_WINDOW:].to(dtype=torch.int32) + history = torch.zeros((1, HISTORY_WINDOW), dtype=torch.int32, device=codes.device) + if recent.numel(): + history[0, : recent.numel()].copy_(recent) + return TalkerCodecDeviceState( + history=history, + history_len=torch.tensor([recent.numel()], dtype=torch.int32, device=codes.device), + step=torch.tensor([step], dtype=torch.int32, device=codes.device), + max_tokens=torch.tensor([max_tokens], dtype=torch.int32, device=codes.device), + finished=torch.tensor([finished], dtype=torch.bool, device=codes.device), + ) + + +def prepare_codec_logits( + raw_logits: torch.Tensor, + state: TalkerCodecDeviceState, + min_tokens: torch.Tensor, + temperature: torch.Tensor, + repetition_penalty: torch.Tensor, + *, + eos_token_id: int, + top_k: int, + top_p: float, + min_tokens_to_keep: int = MIN_TOKENS_TO_KEEP, +) -> torch.Tensor: + """Prepare and filter logits without changing native NPU RNG semantics.""" + logits = raw_logits.float() / temperature.reshape(1, 1) + positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) + valid = positions < state.history_len.reshape(1, 1) + safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) + counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=raw_logits.device) + counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) + alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) + logits = torch.where(logits < 0, logits * alpha, logits / alpha) + eos = logits[..., eos_token_id : eos_token_id + 1] + eos = torch.where( + (state.step < min_tokens).reshape(1, 1), + torch.full_like(eos, float("-inf")), + eos, + ) + logits = torch.cat([logits[..., :eos_token_id], eos, logits[..., eos_token_id + 1 :]], dim=-1) + if 0.0 < float(top_p) < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=False, dim=-1) + cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) + remove = cumulative_probs <= (1.0 - float(top_p)) + remove[..., -min_tokens_to_keep:] = False + remove = remove.scatter(-1, sorted_indices, remove) + logits.masked_fill_(remove, float("-inf")) + if top_k > 0: + keep = min(VOCAB_SIZE, max(int(top_k), int(min_tokens_to_keep))) + threshold = torch.topk(logits, keep, dim=-1).values[..., -1, None] + logits.masked_fill_(logits < threshold, float("-inf")) + return logits + + +def advance_codec_device_state( + state: TalkerCodecDeviceState, + sampled_token: torch.Tensor, + *, + eos_token_id: int, +) -> TalkerCodecDeviceState: + """Advance persistent stochastic state with graph-capturable tensor ops.""" + sampled = sampled_token.reshape(1).to(torch.int32) + active = ~state.finished + next_step = torch.where(active, state.step + 1, state.step) + is_eos = sampled == int(eos_token_id) + reached_limit = active & (next_step >= state.max_tokens) + finished = state.finished | is_eos | reached_limit + emit = active & (~is_eos) & (~reached_limit) + shifted = torch.cat([state.history[:, 1:], sampled.reshape(1, 1)], dim=1) + append_at = state.history_len.clamp(min=0, max=HISTORY_WINDOW - 1).to(torch.long).reshape(1, 1) + appended = state.history.scatter(1, append_at, sampled.reshape(1, 1)) + candidate = torch.where((state.history_len >= HISTORY_WINDOW).reshape(1, 1), shifted, appended) + history = torch.where(emit.reshape(1, 1), candidate, state.history) + history_len = torch.clamp(state.history_len + emit.to(torch.int32), max=HISTORY_WINDOW) + return TalkerCodecDeviceState(history, history_len, next_step, state.max_tokens, finished) + + +def codec_sample_result( + state: TalkerCodecDeviceState, + sampled_token: torch.Tensor, + *, + eos_token_id: int, +) -> TalkerCodecSampleResult: + """Advance device state and retain the device-side emit decision.""" + sampled = sampled_token.reshape(1).to(torch.int32) + active = ~state.finished + next_step = torch.where(active, state.step + 1, state.step) + emit = active & (sampled != int(eos_token_id)) & (next_step < state.max_tokens) + next_state = advance_codec_device_state(state, sampled, eos_token_id=eos_token_id) + return TalkerCodecSampleResult(sampled_token=sampled, state=next_state, emit=emit) + + +def greedy_codec_sample( + raw_logits: torch.Tensor, + state: TalkerCodecDeviceState, + min_tokens: torch.Tensor, + repetition_penalty: torch.Tensor, + *, + top_k: int, + eos_token_id: int, +) -> TalkerCodecSampleResult: + """Greedy codec sample in graph-capturable tensor ops.""" + positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) + valid = positions < state.history_len.reshape(1, 1) + safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) + counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=raw_logits.device) + counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) + alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) + penalized = torch.where(raw_logits < 0, raw_logits * alpha, raw_logits / alpha) + + eos = penalized[..., eos_token_id : eos_token_id + 1] + eos = torch.where( + (state.step < min_tokens).reshape(1, 1), + torch.full_like(eos, float("-inf")), + eos, + ) + penalized = torch.cat([penalized[..., :eos_token_id], eos, penalized[..., eos_token_id + 1 :]], dim=-1) + if top_k > 0: + keep = min(VOCAB_SIZE, max(int(top_k), MIN_TOKENS_TO_KEEP)) + threshold = torch.topk(penalized, keep, dim=-1).values[..., -1, None] + penalized = penalized.masked_fill(penalized < threshold, float("-inf")) + sampled = torch.argmax(penalized, dim=-1).to(torch.int32) + + active = ~state.finished + next_step = torch.where(active, state.step + 1, state.step) + is_eos = sampled.to(torch.int64) == int(eos_token_id) + reached_limit = active & (next_step >= state.max_tokens) + finished = state.finished | is_eos | reached_limit + emit = active & (~is_eos) & (~reached_limit) + + shifted = torch.cat([state.history[:, 1:], sampled.to(torch.int32).unsqueeze(1)], dim=1) + append_at = state.history_len.clamp(min=0, max=HISTORY_WINDOW - 1).to(torch.long).unsqueeze(1) + appended = state.history.scatter(1, append_at, sampled.to(torch.int32).unsqueeze(1)) + candidate = torch.where((state.history_len >= HISTORY_WINDOW).unsqueeze(1), shifted, appended) + history = torch.where(emit.unsqueeze(1), candidate, state.history) + history_len = torch.clamp(state.history_len + emit.to(torch.int32), max=HISTORY_WINDOW) + return TalkerCodecSampleResult( + sampled_token=sampled, + state=TalkerCodecDeviceState(history, history_len, next_step, state.max_tokens, finished), + emit=emit, + ) diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py new file mode 100644 index 00000000000..1c681abbb91 --- /dev/null +++ b/vllm_omni/platforms/npu/ascend_warmup_patch.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Skip the vllm-ascend Triton warmups that fault the 910_93 vector core. + +vllm-ascend warms three Triton kernel families on the worker before graph +capture (``vllm_ascend/model_executor/warmup/kernel_warmup.py``). On 910C/A3 +the second one, ``penalties_triton_warmup``, faults inside its dummy-token +setup before any Triton kernel runs: + + tokens[:, -1:] = vocab_size + +lowers to ACLNN InplaceCopy over the legacy BroadcastTo kernel, and on +ascend910_93 that kernel raises a vector-core exception -- "The address for +the scalar to access the internal buffer of AICore is out of bounds", aclnn +error 507035 -- which kills the stage worker during initialization and takes +the whole orchestrator down with it. + +The identical line runs clean on 910B3 under the same CANN 9.1.0 and +torch_npu 2.10.0.post4 (verified directly, with the same shapes), so this is +a 910_93 operator fault rather than a shape, dtype or memory problem -- the +warmup input is only ``[max_num_seqs, 257]`` token ids. Rewriting the line +would merely relocate the fault into an operator nobody can verify on the +target box, so the warmup is skipped instead. + +Skipping is safe by construction: the warmup only pre-compiles Triton kernels +(JIT), and the sampling path runs those kernels over real token tensors, never +through this dummy construction. The cost is a one-off compile on the first +penalised sample of the first request. Correctness is unaffected. + +Scoped so the 910B baseline keeps the stock warmup byte-for-byte: the guards +delegate to the original warmup there. The skip set is resolved when a warmup +actually runs rather than when the platform is constructed, because the +constructor can run before the device is set -- and a probe that failed there +would silently strip a warmup the baseline relies on. A probe that still +cannot name the SoC skips the faulting warmup: a missing warmup costs latency, +a faulting one costs the whole run. +``VLLM_OMNI_NPU_SKIP_WARMUPS`` overrides the set (comma separated substrings +of the warmup names, or "none" to restore stock behaviour). + +``rejection_sampler_triton_warmup`` joined the default skip set after the A3 +run of 02:38 pinned the acl 507035 aivec fault on it: the dummy +(batch, spec_len, vocab) sweep "completed" on the host half a second before +the pre-capture synchronize raised, and nothing else was enqueued in between +(penalties was already skipped, rms is a no-op there). Skipping is lossless +the same way: on this stack no 910C runtime path launches the reject_sample +kernels -- the Talker's K-step sample() and its SoC guards fall back to +torch argmax, so the warmup only ever compiled kernels nobody runs. +""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_PATCHED = False +_ENV = "VLLM_OMNI_NPU_SKIP_WARMUPS" +# Substrings matched against the warmup function names in kernel_warmup.py. +_DEFAULT_SKIP = ("penalties", "rejection_sampler") +_WARMUP_NAMES = ( + "rejection_sampler_triton_warmup", + "penalties_triton_warmup", + "triton_rms_warmup", +) + + +def _probe_soc_name() -> str: + """torch_npu device name, or "" when it cannot be read yet.""" + try: + import torch_npu + + try: + device = torch_npu.npu.current_device() + except Exception: + device = 0 + return str(torch_npu.npu.get_device_name(device)) + except Exception: + return "" + + +def _kstep_armed() -> bool: + """Whether the Talker K-step is explicitly armed (frames > 1). + + Single source of truth is the same env the deploy config reads; a loader + that cannot import it is treated as "not armed" so the baseline never + changes by accident. + """ + try: + from vllm_omni.config.stage_config import talker_frames_per_step + + return talker_frames_per_step() > 1 + except Exception: + return False + + +def _skipped_names() -> set[str]: + raw = os.environ.get(_ENV, "").strip().lower() + if raw: + if raw in ("none", "off", "0", "false", "no"): + return set() + return {part.strip() for part in raw.split(",") if part.strip()} + soc = _probe_soc_name() + if soc.startswith("Ascend910B"): + if _kstep_armed(): + # The torch-native sampler replaces the reject kernels whenever + # the K-step is armed, so this warmup would only compile kernels + # nobody runs. Everything else keeps the stock warmup. + return {"rejection_sampler"} + # Verified clean here; keep the stock warmup untouched. + return set() + return set(_DEFAULT_SKIP) + + +def _make_guard(name: str, original: Callable[[Any], Any]) -> Callable[[Any], Any]: + """Run the real warmup unless this call is on a host that must skip it. + + The decision is taken per call, not at install time: the platform + constructor can run before the device is set, and a probe that fails there + would otherwise silently skip a warmup the 910B baseline relies on. + """ + + def _guarded(worker: Any) -> Any: + if any(fragment in name for fragment in _skipped_names()): + logger.info( + "[npu] %s skipped: these Triton warmups fault the 910_93 " + "vector core during their dummy-token setup (acl 507035), " + "and no runtime path on this stack launches the kernels", + name, + ) + return None + return original(worker) + + _guarded._vllm_omni_guarded = True # type: ignore[attr-defined] + return _guarded + + +def apply_ascend_warmup_patch() -> None: + """Guard the ascend warmups so a faulting one can be skipped per call.""" + global _PATCHED + if _PATCHED: + return + _PATCHED = True + + try: + import importlib + + # Import the module, not the function: ``vllm_ascend.model_executor.warmup`` + # re-exports ``kernel_warmup`` as a package attribute, so + # ``from ... import kernel_warmup`` yields the function and every + # ``getattr`` below would silently miss. + warmup_module = importlib.import_module( + "vllm_ascend.model_executor.warmup.kernel_warmup" + ) + except Exception as error: # pragma: no cover - non-ascend builds + logger.warning("[npu] Ascend warmup patch not applied: %s", error) + return + + guarded: list[str] = [] + for name in _WARMUP_NAMES: + original = getattr(warmup_module, name, None) + if original is None or getattr(original, "_vllm_omni_guarded", False): + continue + setattr(warmup_module, name, _make_guard(name, original)) + guarded.append(name) + + if guarded: + logger.info( + "[npu] ascend Triton warmups guarded (%s); the skip set is resolved " + "per call, so the SoC probe sees a live device. Override with %s.", + ", ".join(guarded), + _ENV, + ) diff --git a/vllm_omni/platforms/npu/native_rejection_sampler.py b/vllm_omni/platforms/npu/native_rejection_sampler.py new file mode 100644 index 00000000000..b0ea1bf1ed6 --- /dev/null +++ b/vllm_omni/platforms/npu/native_rejection_sampler.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Restore the torch-native rejection sampler on the 910_93 part. + +vllm-ascend module-patches vllm's ``rejection_sample`` (and its two helpers) +onto Triton kernels (``vllm_ascend/patch/worker/patch_rejection_sampler.py``). +On the 910_93 vector core those kernels fault: the 02:38 run pinned acl +507035 on their warmup, which is why ``rejection_sampler`` sits in the +warmup skip set. The skip moves the fault to the first request -- every +verify then JIT-compiles the kernel family on the fly, and both stages +stall there: stage 0 verifies the n-gram drafts of every decode step, and +stage 1 verifies the always-``continue`` drafts behind the K-step +block-table growth. + +The torch-native implementation the patch replaces is pure tensor ops and +exact for both callers, so this module reloads vllm's own sampler source +into a fresh module object and re-points the three names at the pre-patch +functions. Scoped to the same SoC that skips the warmup; the 910B baseline +keeps the Triton kernels it has been measured with -- except when the Talker +K-step is explicitly armed there. That is the one path whose spec width the +910B kernels cannot take (aivec 507035 on every one of the 8 combos tried), +so the restore on a 910B is tied to exactly that arming, never to the SoC +alone: the stock baseline (no arming env) keeps its sampler byte-for-byte. +""" + +import importlib.util +import logging + +from vllm_omni.platforms.npu.ascend_warmup_patch import _kstep_armed, _probe_soc_name + +logger = logging.getLogger(__name__) + +_RESTORED = False +_RESTORE_SOC_PREFIXES = ("ascend910_93", "ascend910c") +# See the module docstring: on these parts the restore applies only when the +# K-step is explicitly armed, because their stock baseline runs the patched +# kernels fine at the widths it was measured with. +_KSTEP_SOC_PREFIXES = ("ascend910b",) + + +def _target_soc() -> bool: + name = _probe_soc_name().strip().lower() + if any(name.startswith(prefix) for prefix in _RESTORE_SOC_PREFIXES): + return True + return any(name.startswith(prefix) for prefix in _KSTEP_SOC_PREFIXES) and _kstep_armed() + + +def restore_native_rejection_sampler() -> None: + """Point vllm's rejection_sampler module back at its torch functions.""" + global _RESTORED + if _RESTORED: + return + _RESTORED = True # probe once; a non-target SoC never retries + if not _target_soc(): + return + try: + import vllm.v1.sample.rejection_sampler as rs + except ImportError: # pragma: no cover - vllm core moved the module + return + spec = importlib.util.spec_from_file_location( + "_vllm_rejection_sampler_native", rs.__file__ + ) + if spec is None or spec.loader is None: # pragma: no cover + return + native = importlib.util.module_from_spec(spec) + spec.loader.exec_module(native) + restored = [] + for name in ("apply_sampling_constraints", "expand_batch_to_tokens", "rejection_sample"): + fn = getattr(native, name, None) + if callable(fn) and getattr(rs, name, None) is not fn: + setattr(rs, name, fn) + restored.append(name) + if restored: + logger.info( + "[npu] rejection sampler restored to the torch-native %s on " + "%s: the Triton kernels fault this vector core (acl 507035, " + "their warmup is already skipped) and JIT-compiling them on " + "the first request stalls every stage mid-verify", + ", ".join(restored), + _probe_soc_name(), + ) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 2c8d42f6093..26eb94120c6 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -59,6 +59,9 @@ def __init__(self) -> None: from vllm_ascend.utils import adapt_patch from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches + from vllm_omni.platforms.npu.ascend_warmup_patch import ( + apply_ascend_warmup_patch, + ) from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import ( apply_minicpmo_4_5_code2wav_patch, ) @@ -74,6 +77,9 @@ def __init__(self) -> None: apply_qwen3_tts_patches() apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches() + # 第 8 项支撑（910C 部署防护）：K 步 Triton rejection/penalties 的 + # warmup 守卫必须在 ascend 内核 warmup 之前装好。 + apply_ascend_warmup_patch() @classmethod def set_device(cls, device: torch.device) -> None: diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index aa01b7a8ec1..d4bbb05efdf 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -3,6 +3,7 @@ from __future__ import annotations +import contextlib import logging import time from collections.abc import Mapping @@ -44,6 +45,7 @@ from vllm_ascend.worker.model_runner_v1 import graph_capture from vllm_omni.data_entry_keys import flatten_payload +from vllm_omni.utils.step_prof import span, tic, toc from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output from vllm_omni.model_executor.duplex_sampling import DuplexSamplingRunnerMixin @@ -136,6 +138,113 @@ def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) self._resolve_duplex_sampling_hook(force=True) + @contextlib.contextmanager + def _narrow_decode_query_len(self): + """One query row per decode graph, for the duration of the block. + + Multi-frame replay reads one row of its graph, and a K-query replay costs + ~0.6 ms more than a one-query one. Without the multi-frame loop engaged + (``narrow_replay_enabled`` false) this is a no-op, so the wide capture + path stays exactly as before.""" + from vllm_omni.platforms.npu.worker import talker_multiframe + + saved = int(getattr(self, "uniform_decode_query_len", 1) or 1) + if saved <= 1 or not talker_multiframe.narrow_replay_enabled(self): + yield + return + logger.info("[minicpmo] the Talker's decode graphs are one query row wide, not %d", saved) + dispatcher = getattr(self, "cudagraph_dispatcher", None) + dispatcher_saved = getattr(dispatcher, "uniform_decode_query_len", None) + self.uniform_decode_query_len = 1 + if dispatcher_saved is not None: + dispatcher.uniform_decode_query_len = 1 + try: + yield + finally: + self.uniform_decode_query_len = saved + if dispatcher_saved is not None: + dispatcher.uniform_decode_query_len = dispatcher_saved + + def _check_and_update_cudagraph_mode(self, *args, **kwargs): + """Register the dispatcher's decode keys one query row wide, not K.""" + with self._narrow_decode_query_len(): + return super()._check_and_update_cudagraph_mode(*args, **kwargs) + + def propose_draft_token_ids(self, valid_sampled_token_ids, *args, **kwargs): + """Carry the Talker's frame count to the scheduler, not a prediction. + + Stage 1 declares a speculative_config so that vLLM knows the request + advances by K tokens per step -- that is what grows the block table and + reserves the KV slots `talker_multiframe` writes into. The drafter + itself has nothing to say: the Talker's vLLM-level vocabulary is + `continue` and `stop`, the frames are generated sequentially inside one + `execute_model`, and the model's own stop row is what ends the request. + So the draft is `continue` repeated, and the rejection sampler accepts + it exactly up to the frame the codec sequence ended on. + + The frame count comes from `drafts_this_step`, which reads the runner's + own `num_spec_tokens` -- the value vLLM derives from the deploy + config's speculative_config. Reading it off the model instead would + depend on wrappers we do not control. + + NOTE: the NPU draft call site passes ``sampled_token_ids`` first + (vllm-ascend ``NPUModelRunner.propose_draft_token_ids``); the upstream + GPU signature has ``scheduler_output`` first instead. + """ + from vllm_omni.platforms.npu.worker import talker_multiframe + + frames = talker_multiframe.drafts_this_step(self) + if frames > 1: + if not isinstance(valid_sampled_token_ids, list): + # The padded-drafter branch passes the verify output tensor; + # `constant_drafts` can only read list rows, so if we ever + # get here the K-step silently degrades to single frames + # every step. Name the branch so the first scene is the last. + logger.error( + "[kstep] propose received %s (expected list); " + "constant_drafts will emit no drafts this step " + "(use_ngram_gpu=%s, padded_drafter_disabled=%s)", + type(valid_sampled_token_ids).__name__, + self.speculative_config.use_ngram_gpu() + if self.speculative_config is not None + else None, + self.speculative_config.disable_padded_drafter_batch + if self.speculative_config is not None + else None, + ) + drafts = talker_multiframe.constant_drafts( + valid_sampled_token_ids, frames, self.input_batch.num_reqs + ) + if not any(drafts): + # Dump the rows as they were seen. An all-empty batch here is + # the stall signature: name WHICH request and what shape the + # step returned instead of leaving one log line to reason from. + # Type-safe throughout: the padded-drafter branch passes a + # tensor, and `tensor or []` raises the ambiguous-truth error. + if isinstance(valid_sampled_token_ids, list): + rows = valid_sampled_token_ids + _desc = "rows=%s types=%s lens=%s head=%s" % ( + len(rows), + [type(r).__name__ for r in rows[:8]], + [len(r) for r in rows[:8] if isinstance(r, list)], + repr(rows)[:400], + ) + else: + seen = valid_sampled_token_ids + shape = tuple(seen.shape) if hasattr(seen, "shape") else "?" + _desc = "rows=<%s shape=%s>" % ( + type(seen).__name__, + shape, + ) + logger.error( + "[kstep] no drafts emitted: frames=%s num_reqs=%s %s", + frames, + self.input_batch.num_reqs, + _desc, + ) + return drafts + return super().propose_draft_token_ids(valid_sampled_token_ids, *args, **kwargs) + def _update_states(self, scheduler_output: SchedulerOutput): deferred_state_corrections_fn = super()._update_states(scheduler_output) self._update_duplex_sampling_states(scheduler_output) @@ -156,6 +265,8 @@ def _make_buffer(self, *size, dtype, numpy=True): # -------------------------------------- Omni-new ------------------------------------------------- def capture_model(self) -> int: + # num_spec_tokens comes from the deploy config's speculative_config, + # which vLLM applies when the runner is built. npugraph_memory_bytes = super().capture_model() self._capture_talker_mtp_graphs() return npugraph_memory_bytes @@ -315,6 +426,7 @@ def execute_model( scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: + tic("em_total") if self.vllm_config.model_config.enable_return_routed_experts: capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): @@ -664,9 +776,10 @@ def execute_model( ): if self.cache_config.mamba_cache_mode == "align": mamba_utils.do_mamba_copy_block(preprocess_bufs) - hidden_states = self._model_forward( - num_tokens_padded, input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs - ) + with span("em_forward"): + hidden_states = self._model_forward( + num_tokens_padded, input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs + ) with record_function_or_nullcontext("post process"): # -------------------------------------- Omni-new ------------------------------------------------- # [Omni] Map pending ropes metadata to req_ids. @@ -794,6 +907,7 @@ def execute_model( if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): self._omni_routed_experts_d2h(scheduler_output) + toc("em_total") return None def _sample( @@ -801,6 +915,7 @@ def _sample( logits: torch.Tensor | None, spec_decode_metadata: Any, ): + tic("em_sample") sampling_metadata = self.input_batch.sampling_metadata if spec_decode_metadata is None: model_sample = getattr(self.model, "sample", None) @@ -827,18 +942,22 @@ def _sample( requests=getattr(self, "requests", None), ) if sampler_output is not None: + toc("em_sample") return sampler_output + toc("em_sample") return self.sampler( logits=logits, sampling_metadata=sampling_metadata, ) + toc("em_sample") return super()._sample(logits, spec_decode_metadata) @torch.inference_mode() def sample_tokens( self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + tic("st_total") profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config kv_connector_output = self.kv_connector_output self.kv_connector_output = None @@ -915,6 +1034,17 @@ def sample_tokens( self.input_batch.sampling_metadata.logitsprocs, logits.shape[-1], ) + if getattr(self.model, "supports_multi_frame_decode", False): + # K-step only: this layer censors the stop token itself, and its + # release condition is a spec-decoding bookkeeping counter we do + # not control -- a request whose count never advances can never + # stop. The model already masks the codec EOS by its own frame + # count (talker_codec_sample), so clear this list every step. + from vllm_omni.platforms.npu.worker import talker_multiframe + + talker_multiframe.neutralize_kstep_min_tokens( + self.input_batch.sampling_metadata.logitsprocs + ) # -------------------------------------- Omni-new ------------------------------------------------- @@ -982,6 +1112,14 @@ def propose_draft_token_ids(sampled_token_ids): spec_decode_metadata, ) + if getattr(self.model, "supports_multi_frame_decode", False): + from vllm_omni.platforms.npu.worker import talker_multiframe + + if talker_multiframe.stop_trace_enabled(): + talker_multiframe.trace_kstep_bookkeeping( + self, valid_sampled_token_ids, logits + ) + with record_function_or_nullcontext("draft_token"): if self.speculative_config: if not early_pp_padded_drafter: @@ -1239,6 +1377,7 @@ def _unwrap_lists(v): self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output) if not self.use_async_scheduling: + toc("st_total") return model_runner_output async_output = AsyncGPUModelRunnerOutput( model_runner_output=model_runner_output, @@ -1253,6 +1392,7 @@ def _unwrap_lists(v): async_output.sampled_token_ids_cpu, async_output.async_copy_ready_event, ) + toc("st_total") return async_output # -------------------------------------- Omni-new ------------------------------------------------- diff --git a/vllm_omni/platforms/npu/worker/npu_ar_worker.py b/vllm_omni/platforms/npu/worker/npu_ar_worker.py index 19f3f89c8f6..6a0b5424755 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_worker.py @@ -3,10 +3,18 @@ from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.platforms.npu.ascend_warmup_patch import apply_ascend_warmup_patch from vllm_omni.platforms.npu.worker.base import OmniNPUWorkerBase from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin +# Install at import time: this module is imported while the stage worker is +# built, which is earlier than the ascend Triton kernel warmup (that runs +# after load_model in the same process). The guard resolves the SoC per call, +# so installing this early is safe -- the probe only has to be right when a +# warmup actually runs. +apply_ascend_warmup_patch() + class NPUARWorker(OmniWorkerMixin, OmniNPUWorkerBase): """NPU AR worker for thinker/talker stages in Omni model.""" diff --git a/vllm_omni/platforms/npu/worker/npu_generation_worker.py b/vllm_omni/platforms/npu/worker/npu_generation_worker.py index ac66c6a54fb..fad66915df5 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_worker.py @@ -3,10 +3,15 @@ from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.platforms.npu.ascend_warmup_patch import apply_ascend_warmup_patch from vllm_omni.platforms.npu.worker.base import OmniNPUWorkerBase from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin +# See npu_ar_worker: install before the worker is used, so the guard is in +# place when the ascend Triton kernel warmup runs later in this process. +apply_ascend_warmup_patch() + class NPUGenerationWorker(OmniWorkerMixin, OmniNPUWorkerBase): """NPU generation worker for code2wav stage in Omni model.""" diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index d8e352d48e8..3ac145a5913 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -23,6 +23,7 @@ from vllm_omni.core.prefix_cache import stage_prefix_cache_config from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms.npu._310p import is_310p +from vllm_omni.platforms.npu.native_rejection_sampler import restore_native_rejection_sampler from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner logger = init_logger(__name__) @@ -69,6 +70,11 @@ def load_model(self, *args, **kwargs) -> None: apply_model_patches(self.model_config) NPUModelRunner.load_model(self, *args, **kwargs) + # Before the first request verifies a draft: on 910_93 the patched + # Triton rejection kernels fault their warmup (skipped) and would be + # JIT-compiled mid-request otherwise, stalling every stage. See + # native_rejection_sampler for the chain. + restore_native_rejection_sampler() # Initialize enable_sp cache to avoid get_current_vllm_config() error # in _pad_for_sequence_parallelism during execute_model. # This is a workaround for vllm-ascend not passing vllm_config to enable_sp(). @@ -475,6 +481,75 @@ def _model_forward( } run_model = partial(self.model, **model_inputs) + # K-step (Talker multi-frame decode). A decode step that schedules K + # positions per request belongs to ``talker_multiframe``: it replays the + # captured decode graph K times -- refreshing the graph's parameters + # before every replay, which is what keeps positions and embeddings in + # step with the frame being produced -- and samples one codec frame per + # replay. Everything else (prefill, a mixed step, another stage) falls + # through to the ordinary single-forward path below. + from vllm_omni.platforms.npu.worker import talker_multiframe + + frames = talker_multiframe.applies(self.model, model_kwargs_extra) + if frames <= 1 and talker_multiframe.is_multi_token_decode(self.model, model_kwargs_extra): + # Falling through here would sample one frame from the last row of + # each span and report one stop row for a step that needs one per + # position: wrong codec tokens and wrong scheduler accounting, both + # silently. The preceding gate log says which check refused. + raise RuntimeError( + "MiniCPM-o scheduled a multi-token decode step that the multi-frame " + "loop declined; see the preceding '[minicpmo] multi-frame Talker " + "decode not engaged' line" + ) + if frames > 1: + spans = [(int(start), int(end)) for start, end in model_kwargs_extra["request_token_spans"]] + narrow = talker_multiframe.begin_narrow_step( + runner=self, + forward_context=forward_context, + frames=frames, + positions=positions, + inputs_embeds=inputs_embeds, + spans=spans, + ) + graph_tokens = len(spans) if narrow is not None else num_tokens_padded + try: + model_output = talker_multiframe.run( + model=self.model, + run_model=run_model, + after_forward=lambda: self._update_full_graph_params_if_needed( + forward_context, graph_tokens + ), + inputs_embeds=inputs_embeds, + input_ids=input_ids, + frames=frames, + model_kwargs=model_kwargs, + model_kwargs_extra=model_kwargs_extra, + narrow=narrow, + ) + finally: + talker_multiframe.end_narrow_step(forward_context, narrow) + # Restore the vocab gate the ed8c7db8 refactor dropped along with + # the _model_forward override it lived in. The Talker's vLLM-level + # head is the two-wide continue/stop row, but input_batch reports + # vocab_size=0, and the rejection sampler's parse_output then masks + # every accepted token as out-of-vocab: the request comes back with + # nothing, the scheduler never rolls back the K tokens it advanced, + # and the next steps schedule a negative count -- the engine hangs + # right after the first 'engaged' step. Single-frame steps never + # took the branch that reads vocab_size; only the K-step does. + # The gate needs one thing from this call: that this really is a + # multi-frame decode step. The width it writes is the two-wide stop + # row, never the hidden width this tensor happens to carry. Hand it + # the stop rows when they are at hand so the intent is visible, and + # keep the hidden states only as the "this is a K-step" signal. + stop_rows = getattr(self.model, "_batch_stop_logits", None) + talker_multiframe.ensure_stop_token_vocab( + self, + stop_rows if stop_rows is not None else model_output.text_hidden_states, + ) + self._omni_last_model_output = model_output + return model_output + if self.enable_enpu: # The soft segmentation scenario requires event.record first, then event.wait. self._update_full_graph_params_if_needed(forward_context, num_tokens_padded) diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py new file mode 100644 index 00000000000..56d27a0d7a7 --- /dev/null +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -0,0 +1,959 @@ +"""Several codec frames inside one stage-1 ``execute_model``. + +Measured on A3 / 910C (``server-env/design/RTF_CEILING_20260828.md``): one +MiniCPM-o Talker decode step is ~2.9 ms of vLLM host work wrapped around a +0.83 ms device forward, for a 190M-parameter model that emits exactly one codec +frame. Over ~118 frames per request that host work is ~340 ms of the ~540 ms +that follows TTFT -- by a wide margin the largest single term left in the score. + +None of it is per *frame*. The scheduler step, input preparation, attention +metadata, sampling, output assembly and engine-core IPC are per *step*, and a +step could just as well carry K frames. This module runs them. + +## How a frame advances without the host + +The step is scheduled with K query positions per request (see +``_apply_minicpmo_talker_multiframe_default``). vLLM prepares all K of them -- +positions, KV slots, block table, sequence lengths -- and the captured graph is +a K-token uniform decode. What vLLM cannot do is fill in the *inputs*: frame +k+1's embedding is the codec token sampled at frame k, so the K positions are +sequential, not speculative. + +So the loop replays the same captured K-token graph K times. Before replay k it +writes frame k's embedding into row k of the persistent ``inputs_embeds`` +buffer the graph captured; after replay k it samples frame k from row k of the +hidden states. Rows above k still hold stale embeddings during replay k, and +their KV is stale too -- but the attention bias is causal, so row k never reads +them, and the next replay overwrites them with the real thing. After the K'th +replay every row and every KV entry is what a K-step decode would have written. + +A replay of the 20-layer forward is 0.83 ms of device and 8 us of host +(``server-env/tools/dispatch_cost.py``), so K replays cost what K frames cost +and the ~2.9 ms is paid once. At K=4 that is a cadence of ~1.9 ms against ~4.0. + +Nothing in the loop touches the host: the embedding lookup, the codec sample +(the captured codec-sampling step), the stop row and the emitted delta are all device +tensors enqueued in order on one stream, and the runner's existing coalesced +D2H carries them out at the end of the step. + +## What makes it correct + +* **Stop is still per frame.** ``make_omni_output`` evaluates the codec EOS and + the token limit on every frame exactly as it does today. A request that ends + at frame j reports its stop row there; frames j+1..K-1 take the "already + finished" early return, emit an empty delta and report a stop row too. vLLM's + rejection sampler sees the first stop and truncates the request to j+1 + tokens, which is the same sequence a one-frame-per-step run would produce. +* **The emitted codes are unchanged.** The extra frames after a stop contribute + nothing: their deltas are empty or the ``-1`` sentinel the connector already + drops (``_codec_scalars``). So the audio is expected to be bit-identical, not + merely close -- which is a far stronger check than the CER gate. + +Off with ``VLLM_OMNI_MINICPMO_TALKER_FRAMES=1``. +""" + +from __future__ import annotations + +import dataclasses +import os +from dataclasses import dataclass +from time import perf_counter +from typing import Any, Callable + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_LOGGED_ENGAGE = False + +# --------------------------------------------------------------------------- +# K 步逐帧计时（默认关，关掉时 run() 里一个 timer 都不取） +# +# VLLM_OMNI_MINICPMO_KSTEP_PROF=1 打开 +# VLLM_OMNI_MINICPMO_KSTEP_PROF_STEPS=3 前 N 个 step 逐帧打点并做设备同步取样 +# VLLM_OMNI_MINICPMO_KSTEP_PROF_EVERY=25 每 N 个 step 打一行累计摘要 +# +# 要回答的问题只有一个：一步 8 帧的墙钟时间，到底花在 host 入队（replay/after_forward） +# 还是花在设备执行上。replay/after_forward 是入队耗时（异步），synced 是在其后加一次 +# 设备同步测到的真实帧耗时 —— 两者相差多少，就是设备侧的账。 +# 诊断代码永不许抛异常：任何异常只把本 profiling 永久关掉并记一条 warning。 +# --------------------------------------------------------------------------- +_PROF_ENV = "VLLM_OMNI_MINICPMO_KSTEP_PROF" +_PROF_STEPS_ENV = "VLLM_OMNI_MINICPMO_KSTEP_PROF_STEPS" +_PROF_EVERY_ENV = "VLLM_OMNI_MINICPMO_KSTEP_PROF_EVERY" +_PROF_OFF = ("0", "false", "no", "off") +_PROF_STEPS_DEFAULT = 3 +_PROF_EVERY_DEFAULT = 25 + + +class _StepProf: + """One step's timers. Never raises: profiling is best-effort by contract.""" + + __slots__ = ("frames", "replay_ms", "after_ms", "detail") + + def __init__(self, frames: int, *, detail: bool) -> None: + self.frames = frames + self.replay_ms = 0.0 + self.after_ms = 0.0 + self.detail = detail + + +def prof_enabled() -> bool: + """Off unless the env explicitly turns it on (``1``/``true``/``yes``/``on``).""" + try: + raw = os.environ.get(_PROF_ENV, "").strip().lower() + return bool(raw) and raw not in _PROF_OFF + except Exception: # pragma: no cover - env is a dict lookup, but be total + return False + + +def _prof_steps() -> int: + try: + raw = os.environ.get(_PROF_STEPS_ENV, "").strip() + steps = int(raw) if raw else _PROF_STEPS_DEFAULT + except Exception: + steps = _PROF_STEPS_DEFAULT + return max(steps, 0) + + +def _prof_every() -> int: + try: + raw = os.environ.get(_PROF_EVERY_ENV, "").strip() + every = int(raw) if raw else _PROF_EVERY_DEFAULT + except Exception: + every = _PROF_EVERY_DEFAULT + return max(every, 1) + + +def _prof_sync() -> None: + """Best-effort device drain; no-op on a build without the op.""" + for fn in ( + getattr(getattr(torch, "npu", None), "synchronize", None), + getattr(getattr(torch, "cuda", None), "synchronize", None), + ): + if callable(fn): + try: + fn() + return + except Exception: + return + + +def _prof_frame_line(step: int, frame: int, replay_ms: float, after_ms: float, synced_ms: float) -> None: + try: + logger.info( + "[kstep-prof] step=%d frame=%d replay_ms=%.3f after_forward_ms=%.3f synced_frame_ms=%s", + step, + frame, + replay_ms, + after_ms, + f"{synced_ms:.3f}" if synced_ms >= 0.0 else "n/a", + ) + except Exception: + pass + + +def _prof_summary(steps: int, frames_total: int, last_frames: int, *, detail: bool) -> None: + """累计一行：每步/每帧墙钟，以及 host 入队与设备执行各自的账。""" + try: + per_step = _PROF_TOTAL_MS / steps if steps else 0.0 + per_frame = _PROF_TOTAL_MS / frames_total if frames_total else 0.0 + replay_per_frame = _PROF_REPLAY_MS / frames_total if frames_total else 0.0 + after_per_frame = _PROF_AFTER_MS / frames_total if frames_total else 0.0 + synced = ( + f" synced_frame_ms={_PROF_SYNCED_MS / _PROF_SYNCED_FRAMES:.3f}" + f" (n={_PROF_SYNCED_FRAMES})" + if _PROF_SYNCED_FRAMES + else "" + ) + logger.info( + "[kstep-prof] steps=%d frames=%d last_step_frames=%d step_ms=%.2f frame_ms=%.3f " + "replay_enqueue_ms_per_frame=%.3f after_forward_ms_per_frame=%.3f (totals: %.3f/%.3f)%s%s", + steps, + frames_total, + last_frames, + per_step, + per_frame, + replay_per_frame, + after_per_frame, + _PROF_REPLAY_MS, + _PROF_AFTER_MS, + synced, + " [per-frame lines above]" if detail else "", + ) + except Exception: + pass + + +_PROF_WARNED = False +_PROF_STEPS = 0 +_PROF_FRAMES = 0 +_PROF_REPLAY_MS = 0.0 +_PROF_AFTER_MS = 0.0 +_PROF_TOTAL_MS = 0.0 +_PROF_SYNCED_MS = 0.0 +_PROF_SYNCED_FRAMES = 0 + + +def _prof_disabled(reason: str) -> None: + """Turn the profiling off for the rest of the process, once.""" + global _PROF_WARNED + if _PROF_WARNED: + return + _PROF_WARNED = True + try: + logger.warning("[kstep-prof] disabled after failure: %s", reason) + except Exception: + pass + os.environ[_PROF_ENV] = "0" + + +def _prof_step_begin(frames: int) -> "_StepProf | None": + """``None`` when off -- and then ``run()`` takes no timer at all.""" + try: + if not prof_enabled(): + return None + return _StepProf(frames, detail=_PROF_STEPS < _prof_steps()) + except Exception: + _prof_disabled("step_begin") + return None + + +def _prof_frame(prof: "_StepProf", frame: int, t0: float, t1: float, t2: float, t3: float) -> None: + """Book one replay. ``replay``/``after`` are enqueue times; ``synced`` is the + frame's real wall clock (drained after the replay), and the gap between them + is the device side of the account.""" + try: + global _PROF_SYNCED_MS, _PROF_SYNCED_FRAMES + replay_ms = (t1 - t0) * 1000.0 + after_ms = (t2 - t1) * 1000.0 + prof.replay_ms += replay_ms + prof.after_ms += after_ms + synced_ms = -1.0 + if prof.detail and t3: + synced_ms = (t3 - t0) * 1000.0 + _PROF_SYNCED_MS += synced_ms + _PROF_SYNCED_FRAMES += 1 + if prof.detail: + _prof_frame_line(_PROF_STEPS + 1, frame, replay_ms, after_ms, synced_ms) + except Exception: + _prof_disabled("frame") + + +def _prof_step_end(prof: "_StepProf | None", step_start: float) -> None: + if prof is None: + return + try: + global _PROF_STEPS, _PROF_FRAMES, _PROF_REPLAY_MS, _PROF_AFTER_MS, _PROF_TOTAL_MS + _PROF_STEPS += 1 + _PROF_FRAMES += prof.frames + _PROF_REPLAY_MS += prof.replay_ms + _PROF_AFTER_MS += prof.after_ms + _PROF_TOTAL_MS += (perf_counter() - step_start) * 1000.0 + if prof.detail or _PROF_STEPS % _prof_every() == 0: + _prof_summary(_PROF_STEPS, _PROF_FRAMES, prof.frames, detail=prof.detail) + except Exception: + _prof_disabled("step_end") +# --------------------------------------------------------------------------- +# K 步 stop 取证（默认关） +# +# VLLM_OMNI_MINICPMO_KSTEP_STOP_TRACE=1 +# +# 只回答一个问题：一步 K 帧里，模型到底有没有产出 stop 行。有 stop 行而请求 +# 仍跑到 max_tokens，丢在 vLLM 的 spec 记账侧；一行都没有，丢在多帧的采样/状态 +# 传递侧。前 50 步各打一行，之后静默。诊断组件永不抛异常。 +# --------------------------------------------------------------------------- +_STOP_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_STOP_TRACE" +_STOP_TRACE_STEPS = 50 +_STOP_TRACE_COUNT = 0 + + +def stop_trace_enabled() -> bool: + """Off unless the env explicitly turns it on (``1``/``true``/``yes``/``on``).""" + try: + raw = os.environ.get(_STOP_TRACE_ENV, "").strip().lower() + return bool(raw) and raw not in _PROF_OFF + except Exception: # pragma: no cover - env is a dict lookup, but be total + return False + + +def _trace_stop_rows(frame_stop_logits: list[torch.Tensor]) -> None: + """每步一行：每个请求的 K 帧里各有几帧的 stop 行是 stop。""" + global _STOP_TRACE_COUNT + try: + if _STOP_TRACE_COUNT >= _STOP_TRACE_STEPS: + return + _STOP_TRACE_COUNT += 1 + flags = torch.stack([row.argmax(dim=-1) for row in frame_stop_logits], dim=1) + logger.info( + "[kstep-stop] step=%d frames=%d stops_per_req=%s", + _STOP_TRACE_COUNT, + len(frame_stop_logits), + flags.sum(dim=1).cpu().tolist(), + ) + except Exception as exc: + logger.warning("[kstep-stop] trace disabled after failure: %r", exc) + os.environ[_STOP_TRACE_ENV] = "0" + + +_STOP_KB_COUNT = 0 + + +def trace_kstep_bookkeeping(runner: Any, valid_sampled_token_ids: Any, logits: Any) -> None: + """[kstep-stop] 第二层：stop 走完 spec 记账之后还剩什么。 + + 两条事实一起打，一次就能把"stop 去哪了"钉死在某一层： + + * ``carried`` -- rejection sampler 交给请求的 token 序列（前 4 个请求）。 + 里面有 1 说明 stop 进了请求自己的 token 表，后面就只剩 engine 侧的结束判定； + 一排 0 而没有 1 说明 stop 在到这里之前就没了。 + * ``min_toks`` -- vLLM 的 MinTokens 处理器里还有哪些请求在 mask 名单上， + 以及它看到的 output 长度（``(index, min_tokens, len(output_token_ids), + sorted(stop_ids))``）。**只要某个请求还在这张名单上，它的 stop 列就是 + -inf**；如果 ``len(output_token_ids)`` 始终不涨，这张名单就永远不会被 + 清掉，请求也就永远停不下来。 + + 诊断组件永不抛异常：出错只关掉自己。 + """ + global _STOP_KB_COUNT + try: + if _STOP_KB_COUNT >= _STOP_TRACE_STEPS: + return + _STOP_KB_COUNT += 1 + if isinstance(valid_sampled_token_ids, list): + carried: Any = [list(row) for row in valid_sampled_token_ids[:4]] + else: + carried = "<%s>" % type(valid_sampled_token_ids).__name__ + + masks: list[Any] = [] + sampling_metadata = getattr(getattr(runner, "input_batch", None), "sampling_metadata", None) + for proc in getattr(getattr(sampling_metadata, "logitsprocs", None), "non_argmax_invariant", None) or []: + min_toks = getattr(proc, "min_toks", None) + if isinstance(min_toks, dict) and min_toks: + masks.append( + [ + (index, int(min_tok), len(out_ids), sorted(stop_ids)) + for index, (min_tok, out_ids, stop_ids) in list(min_toks.items())[:4] + ] + ) + logger.info( + "[kstep-stop] carried=%s min_toks=%s logits=%s", + carried, + masks, + tuple(getattr(logits, "shape", ()) or ()), + ) + except Exception as exc: + logger.warning("[kstep-stop] bookkeeping trace disabled after failure: %r", exc) + os.environ[_STOP_TRACE_ENV] = "0" + + +def neutralize_kstep_min_tokens(logitsprocs: Any) -> None: + """K 步下让 vLLM 层的 ``min_tokens`` 不再 censoring 停止信号。 + + yaml 的 ``min_tokens: 50`` 在 K 步下是通过 **把 stop token（id 1）的 logit + 置成 -inf** 实现的（``MinTokensLogitsProcessor``）。它自己的解除条件是 + ``len(output_token_ids) >= min_tokens``，而这个长度归 vLLM 的 spec 记账管； + 只要那个计数没有推进到位，请求唯一的停止信号就被永久 mask， + ``finished_reason`` 只能是 ``length``（910C 现场：stage1 全 length、0 stop）。 + + 真正的"最小帧数"保护不在这里：``talker_codec_sample.greedy_codec_sample`` + 用 ``state.step < min_tokens`` 把 codec EOS 本身压成 -inf，用的是模型自己 + 的帧计数，与 vLLM 的 token 记账无关。所以这一层是重复的、且在 K 步下会 + 把停止信号一起吃掉 —— 清空它的名单，把最小长度交还给模型内的那道保护。 + + ``min_toks`` 每个 step 都会被 ``update_state`` 重新登记，所以本函数必须 + 每步调用（调用点在 runner 的采样前）。永不抛异常。 + """ + try: + for proc in getattr(logitsprocs, "non_argmax_invariant", None) or []: + if not type(proc).__name__.endswith("MinTokensLogitsProcessor"): + # Suffix rather than equality: the processor can arrive + # subclassed or wrapped (a platform patch, a profiling shim), + # and a mismatch here is silent -- the mask stays on and the + # only stop signal a K-step request has is censored forever. + continue + min_toks = getattr(proc, "min_toks", None) + if isinstance(min_toks, dict) and min_toks: + min_toks.clear() + except Exception as exc: + logger.warning("[minicpmo] K-step min_tokens neutralization skipped: %r", exc) + + +_LOGGED_BLOCK: str | None = None +_LOGGED_NARROW = False +_LOGGED_NARROW_BLOCK: str | None = None + +_NARROW_ENV = "VLLM_OMNI_MINICPMO_NARROW_REPLAY" +_NARROW_OFF = frozenset({"0", "off", "false", "no"}) +_NARROW_ON = frozenset({"1", "on", "true", "yes"}) +# SoC families whose one-query capture is known to fault. The 910_93 part +# trips a vector-core error (acl 507035) inside the packaged codec operator +# while that graph is being captured, so the narrow path stays off there +# unless an operator opts in explicitly -- see narrow_replay_enabled. +# "ascend910c" guards the device-name fallback: on some A3 containers the +# driver reports the part as "Ascend910C" rather than its 910_93 SoC code, +# and an unblocked unknown would let the capture brick the deployment. +# "ascend910b" is fail-safe rather than observed: no K-step attempt there +# ever reached a capture (all 8 combos faulted earlier, inside the reject +# kernel), so the one-query capture is unverified on that family. +_NARROW_BLOCKED_SOC_PREFIXES = ("ascend910_93", "ascend910c", "ascend910b") +_NARROW_SOC_LOGGED: str | None = None + + +def _narrow_soc_allows() -> bool: + """Whether this part may capture the one-query graph by default. + + An explicit ``VLLM_OMNI_MINICPMO_NARROW_REPLAY=1/on`` always wins, and an + explicit ``0/off`` always loses. With no setting the answer depends on the + part: the environment is consulted first (the same variables the deploy + config reads, so both sides agree), and the device name is the fallback for + a worker launched without them. + """ + global _NARROW_SOC_LOGGED + + raw = os.environ.get(_NARROW_ENV, "").strip().lower() + if raw in _NARROW_ON: + return True + if raw in _NARROW_OFF: + return False + + name = "" + for variable in ("SOC_VERSION", "ASCEND_SOC_VERSION"): + value = os.environ.get(variable, "").strip().lower() + if value: + name = value + break + if not name: + try: + import torch_npu + + name = str(torch_npu.npu.get_device_name(torch_npu.npu.current_device())).lower() + except Exception: + name = "" + allows = True if not name else not name.startswith(_NARROW_BLOCKED_SOC_PREFIXES) + if _NARROW_SOC_LOGGED != name: + _NARROW_SOC_LOGGED = name + logger.info( + "[minicpmo] narrow replay SoC gate: name='%s' -> %s", + name or "", + "allowed" if allows else "blocked", + ) + return allows + + +def _log_block_once(reason: str) -> None: + """Say once why the loop never engaged; silence is the worse failure.""" + global _LOGGED_BLOCK + if _LOGGED_BLOCK == reason: + return + _LOGGED_BLOCK = reason + logger.info("[minicpmo] multi-frame Talker decode not engaged: %s", reason) + + +def applies(model: Any, model_kwargs_extra: dict[str, Any]) -> int: + """Frames this step should run, or 0 when the loop must not engage. + + Returns the uniform per-request query length of a pure decode step over a + Talker that can build its own decode embedding. Anything else -- a prefill, + a mixed step, a batch whose requests were scheduled different token counts + -- returns 0 and the caller takes the ordinary single-forward path. + + Every refusal says so once. A step that schedules several tokens per + request and then takes the single-forward path is not merely slow: that + path samples one frame from the *last* row of each span and reports one + stop row for a step that needs one per position, so it corrupts the codec + stream and the scheduler's accounting alike. `_model_forward` turns those + into an error rather than letting them run. + """ + if not getattr(model, "supports_multi_frame_decode", False): + # Every other stage: not a refusal, just not this model. + return 0 + spans = model_kwargs_extra.get("request_token_spans") + infos = model_kwargs_extra.get("model_intermediate_buffer") + if not spans or not infos or len(spans) != len(infos): + return _block("no request_token_spans/model_intermediate_buffer for this step") + frames = int(spans[0][1]) - int(spans[0][0]) + if frames <= 1: + return 0 + for start, end in spans: + if int(end) - int(start) != frames: + # A mixed step (one request prefilling, another decoding) has no + # single frame count, and the captured graph is not a uniform + # decode either. Dump the spans so a real-world refusal can be + # attributed to an exact scheduling state instead of guessed at. + logger.info( + "[minicpmo] refusing non-uniform spans: frames=%d spans=%s", + frames, + [(int(s), int(e)) for s, e in spans], + ) + return _block("request token spans are not uniform") + for info in infos: + if not isinstance(info, dict): + return _block("a request carries no intermediate buffer") + if bool(info.get("_omni_is_prefill", False)): + return 0 + state = info.get("audio_state") + if not isinstance(state, dict): + return _block("a request carries no audio_state") + if int(state.get("step", 0)) <= 0: + # The request has not sampled a codec token yet, so there is no + # previous code for frame 0 to embed and this is not a decode. + return _block("a request has not sampled a codec token yet") + return frames + + +def is_multi_token_decode(model: Any, model_kwargs_extra: dict[str, Any]) -> bool: + """True when this step schedules several tokens for a request that is decoding. + + Scoped to the Talker. Stage 0 drafts the Thinker's text with n-grams and so + schedules multi-token decode steps of its own, which vLLM handles perfectly + well -- it is only the Talker's one-frame-per-position sampler that cannot. + """ + if not getattr(model, "supports_multi_frame_decode", False): + return False + spans = model_kwargs_extra.get("request_token_spans") + infos = model_kwargs_extra.get("model_intermediate_buffer") + if not spans or not infos or len(spans) != len(infos): + return False + for (start, end), info in zip(spans, infos): + if int(end) - int(start) <= 1: + continue + if isinstance(info, dict) and not bool(info.get("_omni_is_prefill", False)): + return True + return False + + +@dataclass +class NarrowStep: + """What one codec frame has to change before the one-query graph replays. + + A K-query graph replay costs 1.33-1.47 ms of device time against the + one-query graph's 0.83 ms, and it is a step change rather than a per-row + cost -- a two-query graph costs the same as a four-query one. The loop only + ever reads one row of the wide graph, so it replays the narrow one instead + and moves the five things a frame changes: + + * ``inputs_embeds[r]`` -- request r's embedding for this frame + * ``positions[r]`` -- its position, +frame + * ``slot_mapping[r]`` -- the KV slot vLLM reserved for this frame + * ``seq_lens[r]`` -- its length, +frame+1, in the buffer the graph reads + * the block table -- unchanged for the whole step + + The step's own K-wide layouts are cloned at ``begin`` because the narrow + views alias them: ``slot_mapping[:rows]`` is the first row of + ``slot_mapping.view(rows, K)``. + """ + + rows: int + frames: int + inputs_embeds: Any + positions: Any + slot_mapping: Any + seq_lens: Any + step_slots: Any + step_positions: Any + step_seq_lens: Any + spans: list[tuple[int, int]] + saved_descriptor: Any = None + saved_runtime_mode: Any = None + + def select(self, frame: int) -> None: + """Put frame ``frame``'s inputs where the captured replay reads them. + + Three copies and nothing else: everything a frame needs was computed + once, at the start of the step, because these run between two graph + replays and an allocation there costs more than the values do. + """ + rows = self.rows + self.slot_mapping[:rows].copy_(self.step_slots[:, frame]) + self.positions[:rows].copy_(self.step_positions[:, frame]) + self.seq_lens[:rows].copy_(self.step_seq_lens[:, frame]) + + +def narrow_replay_enabled(runner: Any) -> bool: + """Whether this worker should capture and replay the one-query graph. + + Off with ``VLLM_OMNI_MINICPMO_NARROW_REPLAY=off``, which restores the + K-query captures and the wide replay. On a part whose one-query capture is + known to fault (``_NARROW_BLOCKED_SOC_PREFIXES``) it is off as well unless + explicitly enabled, so the capture can never brick a deployment. + """ + if not _narrow_soc_allows(): + return False + model = getattr(runner, "model", None) + if not getattr(model, "supports_multi_frame_decode", False): + return False + return int(getattr(runner, "num_spec_tokens", 0) or 0) > 0 + + +def begin_narrow_step( + *, + runner: Any, + forward_context: Any, + frames: int, + positions: Any, + inputs_embeds: Any, + spans: list[tuple[int, int]], +) -> "NarrowStep | None": + """Point the forward context at the one-query graph, or decline. + + Declines -- and the caller falls back to the wide replay -- whenever the + narrow graph for this row count is not captured, which is what makes this + safe to land before every path that reaches it has been exercised. + """ + from vllm.config import CUDAGraphMode + from vllm.forward_context import BatchDescriptor + + # 迁移注记（2026-09-21）：narrow replay 所需的 KV 几何由第 6 项 + # （platforms/npu/attention/fixed_kv_decode）提供，本分支未迁第 6 项。 + # 这里容错导入：模块缺席时按本函数既有承诺 decline，调用方回落 wide replay + # （见 docstring），K 步解码与停止链不受影响。 + try: + from vllm_omni.platforms.npu.attention import fixed_kv_decode + except Exception: + fixed_kv_decode = None + + if fixed_kv_decode is None: + return _decline("fixed-KV decode (item 6) is not present in this tree") + if not narrow_replay_enabled(runner) or inputs_embeds is None or positions is None: + return None + rows = len(spans) + if rows <= 0: + return None + seq_lens = fixed_kv_decode.captured_seq_lens(rows) + if seq_lens is None: + return _decline("no one-query graph captured for %d rows" % rows) + metadata = getattr(forward_context, "attn_metadata", None) + if isinstance(metadata, dict): + metadata = next(iter(metadata.values()), None) + step_slot_mapping = getattr(metadata, "slot_mapping", None) + live_seq_lens = getattr(metadata, "seq_lens_device", None) + graph_slot_mapping = fixed_kv_decode.captured_slot_mapping(rows) + if step_slot_mapping is None or live_seq_lens is None or graph_slot_mapping is None: + return _decline("this step's attention metadata carries no slot mapping") + if step_slot_mapping.shape[0] < rows * frames: + return _decline("slot mapping is shorter than the step's own frames") + + narrow = NarrowStep( + rows=rows, + frames=frames, + inputs_embeds=inputs_embeds, + positions=positions, + slot_mapping=graph_slot_mapping, + seq_lens=seq_lens, + step_slots=step_slot_mapping[: rows * frames].view(rows, frames).clone(), + step_positions=positions[: rows * frames].view(rows, frames).clone(), + # `seq_lens` already counts all K of this step's frames, and by frame k + # only k + 1 have been written -- so row r, column k is + # `live - frames + k + 1`. + step_seq_lens=( + live_seq_lens[:rows, None] + - frames + + torch.arange(1, frames + 1, device=live_seq_lens.device, dtype=live_seq_lens.dtype) + ), + spans=[(index, index + 1) for index in range(rows)], + ) + descriptor = getattr(forward_context, "batch_descriptor", None) + narrow.saved_descriptor = descriptor + if descriptor is not None: + # `uniform` too: a failed dispatch clears it on the descriptor it hands + # back, and the captured key has it set. What this step is has not + # changed -- it is one token per request against a graph captured for + # exactly that -- only which graph replays it. + forward_context.batch_descriptor = dataclasses.replace( + descriptor, num_tokens=rows, num_reqs=rows, uniform=True + ) + else: + forward_context.batch_descriptor = BatchDescriptor( + num_tokens=rows, num_reqs=rows, uniform=True + ) + # The dispatcher builds its lookup descriptor from its own query length -- + # 4 at runtime, which is right for the step -- so it finds no key for the + # narrow capture and hands back NONE, and the wrapper then runs the model + # uncaptured without saying so. This step *is* a full-graph uniform decode; + # only which graph it replays changed, which is what this function decides. + narrow.saved_runtime_mode = getattr(forward_context, "cudagraph_runtime_mode", None) + forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL + global _LOGGED_NARROW + if not _LOGGED_NARROW: + _LOGGED_NARROW = True + logger.info( + "[minicpmo] multi-frame Talker decode replays the one-query graph: " + "%d rows, descriptor %s, dispatched mode was %s", + rows, + forward_context.batch_descriptor, + narrow.saved_runtime_mode, + ) + return narrow + + +def end_narrow_step(forward_context: Any, narrow: "NarrowStep | None") -> None: + if narrow is None: + return + forward_context.batch_descriptor = narrow.saved_descriptor + if narrow.saved_runtime_mode is not None: + forward_context.cudagraph_runtime_mode = narrow.saved_runtime_mode + + +def _decline(reason: str) -> None: + global _LOGGED_NARROW_BLOCK + if _LOGGED_NARROW_BLOCK != reason: + _LOGGED_NARROW_BLOCK = reason + logger.info("[minicpmo] one-query Talker replay not engaged: %s", reason) + return None + + +def _block(reason: str) -> int: + _log_block_once(reason) + return 0 + + +def run( + *, + model: Any, + run_model: Callable[[], Any], + after_forward: Callable[[], None], + inputs_embeds: torch.Tensor | None, + input_ids: torch.Tensor | None, + frames: int, + model_kwargs: dict[str, Any], + model_kwargs_extra: dict[str, Any], + narrow: "NarrowStep | None" = None, +) -> Any: + """Run ``frames`` codec frames and return the merged step output. + + ``after_forward`` is the runner's post-forward graph bookkeeping. It is + called after every replay because that is where it sits in the single-frame + path; under fixed-KV decode it is two pointer comparisons, and the frames of + one step share their sequence lengths so there is nothing for it to rebind. + """ + global _LOGGED_ENGAGE + spans = model_kwargs_extra["request_token_spans"] + infos = model_kwargs_extra["model_intermediate_buffer"] + + step_prof = _prof_step_begin(frames) + step_start = perf_counter() if step_prof is not None else 0.0 + + frame_outputs = [] + frame_stop_logits = [] + frame_hidden = [] + for frame in range(frames): + if narrow is not None: + # Frame 0's embeddings are already in the buffer, but at the rows a + # K-query step laid them out on: request r's is at row r*K, and the + # one-query graph reads row r. r <= r*K, so compacting in request + # order never overwrites a row it has still to read. + narrow.select(frame) + if frame == 0: + _compact_frame_zero_embeddings(inputs_embeds, spans) + else: + _write_frame_embeddings( + model, inputs_embeds, input_ids, spans, infos, frame, narrow=narrow + ) + elif frame > 0: + _write_frame_embeddings(model, inputs_embeds, input_ids, spans, infos, frame) + t0 = perf_counter() if step_prof is not None else 0.0 + hidden = run_model() + t1 = perf_counter() if step_prof is not None else 0.0 + after_forward() + t2 = perf_counter() if step_prof is not None else 0.0 + if step_prof is not None and step_prof.detail: + _prof_sync() + t3 = perf_counter() if step_prof is not None else 0.0 + if step_prof is not None: + _prof_frame(step_prof, frame, t0, t1, t2, t3) + frame_kwargs = dict(model_kwargs_extra) + frame_kwargs["request_token_spans"] = ( + list(narrow.spans) + if narrow is not None + else [(int(start) + frame, int(start) + frame + 1) for start, _ in spans] + ) + if narrow is not None: + # The graph writes one row per request and the next frame overwrites + # it, so this frame's hidden state has to be kept before the replay + # that follows. + frame_hidden.append(hidden.clone()) + output = model.make_omni_output(hidden, **model_kwargs, **frame_kwargs) + stop_logits = model.take_batch_stop_logits() + if stop_logits is None: + raise RuntimeError("MiniCPM-o multi-frame decode expects stop rows from every frame") + frame_outputs.append(output) + frame_stop_logits.append(stop_logits) + + if not _LOGGED_ENGAGE: + _LOGGED_ENGAGE = True + logger.info( + "[minicpmo] multi-frame Talker decode engaged: %d codec frames per step", frames + ) + if stop_trace_enabled(): + _trace_stop_rows(frame_stop_logits) + merged = model.merge_frame_outputs(frame_outputs, frame_stop_logits) + if narrow is not None and frame_hidden: + # (frames, rows, hidden) -> (rows * frames, hidden), request-major -- + # the layout `logits_indices` reads. OmniOutput is a NamedTuple. + merged = merged._replace( + text_hidden_states=torch.stack(frame_hidden, dim=1).reshape( + -1, frame_hidden[0].shape[-1] + ) + ) + # Whole step, frames + merge: this is the number to compare against the + # scheduler's step cadence. + _prof_step_end(step_prof, step_start) + return merged + + +def _compact_frame_zero_embeddings(inputs_embeds: torch.Tensor | None, spans) -> None: + """Move each request's frame-0 embedding to the row the narrow graph reads.""" + if inputs_embeds is None: + raise RuntimeError("MiniCPM-o narrow Talker replay requires an inputs_embeds buffer") + for row, (start, _) in enumerate(spans): + if int(start) != row: + inputs_embeds[row : row + 1].copy_(inputs_embeds[int(start) : int(start) + 1]) + + +def _write_frame_embeddings( + model: Any, + inputs_embeds: torch.Tensor | None, + input_ids: torch.Tensor | None, + spans: list[tuple[int, int]], + infos: list[dict[str, Any]], + frame: int, + narrow: "NarrowStep | None" = None, +) -> None: + """Put frame ``frame``'s decode embedding into the row the graph reads. + + ``model.preprocess`` is the same call the runner's ``_preprocess`` makes for + frame 0; it reads the codec token the previous frame wrote into + ``info["audio_codes"]["current"]`` and returns its embedding. Writing into + ``inputs_embeds`` in place is what makes the next replay see it -- the + buffer is the one the graph captured, and vLLM does not copy inputs into + captured graphs (``cudagraph_copy_inputs`` is False). + """ + if inputs_embeds is None: + raise RuntimeError("MiniCPM-o multi-frame decode requires an inputs_embeds buffer") + for index, ((start, _), info) in enumerate(zip(spans, infos)): + # The token to embed always sits at the step's own row; only where the + # embedding goes changes when the narrow graph is the one replaying. + source = int(start) + frame + row = index if narrow is not None else source + row_ids = input_ids[source : source + 1] if input_ids is not None else None + if row_ids is None: + raise RuntimeError("MiniCPM-o multi-frame decode requires input_ids for preprocess") + _, embeds, _ = model.preprocess(input_ids=row_ids, input_embeds=None, **info) + inputs_embeds[row : row + 1].copy_(embeds[:1]) + + +# The Talker's vLLM-level vocabulary is two tokens wide: `compute_logits` +# returns one-hot stop rows ([0, -inf] / [-inf, 0]) and `sample` takes their +# argmax, so 0 means "keep going" and 1 is the stop token the deploy config +# names in `stop_token_ids`. +CONTINUE_TOKEN_ID = 0 + +# The id that same row emits to mean stop: the stop column wins the argmax, +# `parse_output` keeps ids below the two-wide vocab, and vLLM's `check_stop` +# only fires if this id is in the request's `stop_token_ids` -- which stage 1's +# pipeline constraints derive from this constant. +STOP_TOKEN_ID = 1 + +# Width of the Talker vLLM-level head: the two-wide continue/stop row. +# input_batch.vocab_size must report the same, because InputBatch stores +# top_k = vocab_size as its no-top-k sentinel and that holds only for 0 or 2. +STOP_ROW_WIDTH = 2 + + +def drafts_this_step(runner: Any) -> int: + """Frames the *next* step should be scheduled for, or 0 to stay generic. + + The count travels to the scheduler as speculative tokens because that is + vLLM V1's only way of saying "this request advanced by more than one token", + which is what grows the block table and reserves the KV slots the loop + writes into. Nothing about it is speculative: the drafts are always + `continue`, and the frames are generated sequentially rather than verified. + """ + if not getattr(getattr(runner, "model", None), "supports_multi_frame_decode", False): + return 0 + num_spec = int(getattr(runner, "num_spec_tokens", 0) or 0) + if num_spec <= 0: + return 0 + return num_spec + 1 + + +def constant_drafts( + valid_sampled_token_ids: Any, + frames: int, + num_reqs: int, +) -> list[list[int]]: + """`continue` repeated -- for the whole batch, or for none of it. + + The drafts are how the next step gets K query positions per request, and + the multi-frame loop can only run a step whose requests all have the *same* + span: it replays one captured uniform-decode graph and samples one codec + frame per position. So drafting per request is not an option. A request + that sampled nothing this step (discarded, or finished on this step's stop + token) would be scheduled one token while its neighbours got K, and + `applies` would refuse the resulting step -- which `_model_forward` turns + into a fatal error rather than let the single-forward path sample one frame + from the last row of a K-row span. + + So: every request drafts, or nobody does. A step with no drafts is an + ordinary one-frame-per-request decode, and the frames resume on the step + after it. + """ + for index in range(num_reqs): + sampled = None + if isinstance(valid_sampled_token_ids, list) and index < len(valid_sampled_token_ids): + sampled = valid_sampled_token_ids[index] + if not sampled: + _log_block_once("a request in this batch sampled nothing; no drafts this step") + return [[] for _ in range(num_reqs)] + return [[CONTINUE_TOKEN_ID] * (frames - 1) for _ in range(num_reqs)] + + +def ensure_stop_token_vocab(runner: Any, logits: Any) -> None: + """Let vLLM's rejection-sampler output survive the Talker's narrow head. + + `RejectionSampler.parse_output` keeps a sampled token only while it is both + not the placeholder and `< vocab_size`, and stage 1's `input_batch` reports + a vocab size of **0** -- the Talker's vLLM-level head is the two-wide + continue/stop row `compute_logits` builds, not a tokenizer vocabulary, and + nothing had ever needed the number. With 0 the mask is empty, every + accepted token is dropped, and the request comes back having generated + nothing: the scheduler then never rolls back the tokens it advanced for the + step and schedules a negative count two steps later. + + A single-frame step never reached this. It returns one token per request, + which `_bookkeeping_sync` takes down the `max_gen_len == 1` branch, and that + branch does not consult `vocab_size` at all. + + Raising it back to the width of the row the model actually emits is the + whole fix -- and that width is `STOP_ROW_WIDTH` (two), *not* the width of + whatever tensor the caller happens to be holding. The Talker's vLLM-level + head is the two-wide continue/stop row `compute_logits` builds, so two is + what every `vocab_size` reader has to see. + + The caller used to hand over `text_hidden_states`, so this landed on the + hidden width (768). That passes the `parse_output` filter by accident (0 + and 1 are both below 768) and breaks the one reader that repays checking: + `InputBatch.add_request` stores `top_k = vocab_size` as its "no top-k" + sentinel, which only holds for 0 and 2. At 768 `top_k_reqs` stops being + empty, `sampling_metadata.top_k` stops being None, and the Ascend + `enable_reduce_sample` branch starts treating a two-column distribution as + a 768-way one. The stop row then does not survive into the request's token + list, so the request runs to `max_tokens` instead of stopping on EOS -- and + the frame stream that keeps arriving after the codec sequence ended is what + eventually hands the scheduler a step it cannot merge. + """ + if logits is None or not getattr(runner.model, "supports_multi_frame_decode", False): + return + batch = getattr(runner, "input_batch", None) + width = STOP_ROW_WIDTH + if batch is None or int(getattr(batch, "vocab_size", 0) or 0) >= width: + return + logger.info( + "[minicpmo] Talker input batch reported vocab_size=%s; raising it to the " + "%d-wide stop row so parse_output keeps the accepted tokens", + getattr(batch, "vocab_size", None), + width, + ) + batch.vocab_size = width diff --git a/vllm_omni/utils/step_prof.py b/vllm_omni/utils/step_prof.py new file mode 100644 index 00000000000..fc52d08a1fd --- /dev/null +++ b/vllm_omni/utils/step_prof.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Opt-in per-step wall-clock profiler for MiniCPM-o NPU runners. + +Zero-ish overhead when disabled. Enable with ``OMNI_STEP_PROF=1`` and set the +aggregation cadence with ``OMNI_STEP_PROF_FLUSH`` (totals between dumps, +default 200). Thread-safe; spans are keyed per thread so nested/overlapping +regions in the engine loop stay correct. +""" +from __future__ import annotations + +import os +import sys +import threading +import time +from collections import defaultdict +from contextlib import contextmanager + +_ENABLED = os.environ.get("OMNI_STEP_PROF", "0") == "1" +_FLUSH_EVERY = int(os.environ.get("OMNI_STEP_PROF_FLUSH", "200")) + + +class _StepProf: + def __init__(self) -> None: + self.acc: dict[str, float] = defaultdict(float) + self.cnt: dict[str, int] = defaultdict(int) + self._local = threading.local() + self._lock = threading.Lock() + self._flushed = 0 + + def tic(self, name: str) -> None: + marks = getattr(self._local, "marks", None) + if marks is None: + marks = {} + self._local.marks = marks + marks[name] = time.perf_counter() + + def toc(self, name: str) -> None: + marks = getattr(self._local, "marks", None) + if not marks: + return + t0 = marks.pop(name, None) + if t0 is None: + return + dt = time.perf_counter() - t0 + with self._lock: + self.acc[name] += dt + self.cnt[name] += 1 + total = self.cnt.get("em_total", 0) + if total and total - self._flushed >= _FLUSH_EVERY: + self._flushed = total + self._dump_locked() + + def _dump_locked(self) -> None: + parts = [] + for name in sorted(self.acc): + n = self.cnt[name] + if n: + parts.append(f"{name}={self.acc[name] / n * 1e3:.3f}ms") + sys.stderr.write("[STEPPROF] " + " ".join(parts) + "\n") + sys.stderr.flush() + + +_PROF = _StepProf() + + +def tic(name: str) -> None: + if _ENABLED: + _PROF.tic(name) + + +def toc(name: str) -> None: + if _ENABLED: + _PROF.toc(name) + + +@contextmanager +def span(name: str): + if not _ENABLED: + yield + return + _PROF.tic(name) + try: + yield + finally: + _PROF.toc(name) From 8e236e82ae77be431e58446a782edadd9f316ccf Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 08/89] fix(config): drop stage-0 n-gram block that conflicts with prefix caching; tidy comments - The deploy config enabled stage-0 prefix caching and n-gram speculative decoding together, which vLLM rejects at startup (OmniPrefixCacheUnmatchError). Keep prefix caching (the validated configuration) and remove the n-gram block. - Replace internal notes (item numbers, measurement references, Chinese prose) with concise upstream-style comments. Signed-off-by: y-null --- vllm_omni/config/omni_config.py | 33 +++---- vllm_omni/core/sched/omni_ar_scheduler.py | 11 ++- vllm_omni/deploy/minicpmo_4_5.yaml | 48 +++------- vllm_omni/engine/stage_init_utils.py | 13 +-- .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 12 ++- .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 26 ++--- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 44 +++++---- .../minicpmo_4_5_omni.py | 14 +-- vllm_omni/platforms/npu/platform.py | 4 +- .../platforms/npu/worker/talker_multiframe.py | 96 ++++++++++--------- 10 files changed, 150 insertions(+), 151 deletions(-) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 22f3e771372..73ef0280168 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -1213,14 +1213,14 @@ def _upstream_engine_field_map( _POOLING_ENGINE_FIELDS = frozenset(_PoolingEngineOverrides.__annotations__) _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__) _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) -# 迁移补齐（2026-09-21）：这三者的共同性质是"vllm 自己的配置对象，原样交给 vllm"， -# 不参与本层的结构化解包。`speculative_config` 原先不在名单里，导致迁移过来的 -# stage-1 多帧（K 步）接线无法加载：`stage_config.py` 的 -# `_apply_minicpmo_talker_multiframe_default` 会给 stage 1 注入一个 -# speculative_config（vLLM 靠 `spec_token_ids` 才知道请求一次推进 K 个 token）， -# 而这里的字段归属校验会把 stage 1 的覆盖块整块判为非法，启动即 ValueError： -# "Stage 1 (llm_ar) has explicit engine argument(s) with no structured config owner" -# 它与此前已在册的两项同类（都是直接透传 vllm 配置对象），故一并列入。 +# The three fields below are vLLM config objects passed through verbatim +# instead of being unpacked structurally. speculative_config was missing from +# this list, so a stage-1 multi-frame setup failed to load: stage_config.py +# injects a speculative_config into stage 1 (vLLM uses spec_token_ids to +# advance K tokens per request) and the ownership check here rejected the whole +# stage-1 override block with a startup ValueError ("Stage 1 (llm_ar) has +# explicit engine argument(s) with no structured config owner"). It joins the +# two pass-through fields already listed. _DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset( {"compilation_config", "profiler_config", "speculative_config"} ) @@ -1585,9 +1585,10 @@ def _stage_engine_values( diffusion=_DiffusionEngineOverrides(_select_engine_overrides(diffusion_kwargs, _DIFFUSION_STAGE_ENGINE_FIELDS)), compilation_config=_copy_value(engine.get("compilation_config")), profiler_config=_copy_value(engine.get("profiler_config")), - # 迁移补齐：与上两项同类（vllm 的配置对象，原样透传）。缺少这一格时， - # deploy 配置里的 speculative_config 会静默消失（不报错），K 步与 stage-0 - # ngram 都会失去效果 —— 表现为引擎日志里 speculative_config=None。 + # Same category as the two above: a vLLM config object passed through + # verbatim. Without it, a speculative_config in the deploy config + # silently disappears (no error) and speculative decoding loses effect + # (engine logs show speculative_config=None). speculative_config=_copy_value(engine.get("speculative_config")), ) @@ -1652,11 +1653,11 @@ class BaseVllmOmniStageConfig: parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig) compilation_config: VllmCompilationConfig | None = None profiler_config: VllmProfilerConfig | None = None - # 注意形状：vLLM 的 `create_speculative_config` 是按**字典**消费这个值的 - # （内部会 `.items()` / `.update()` / 再构造 SpeculativeConfig）。这里若标成 - # `VllmSpeculativeConfig | None`，pydantic 会把配置里的字典提前转成对象， - # 于是 vLLM 那一步会炸 "SpeculativeConfig object has no attribute 'items'"。 - # 因此保留 Mapping 形状，让 vLLM 自己去转。 + # Shape matters: vLLM's create_speculative_config consumes this value as a + # dict (it calls .items() / .update() before building SpeculativeConfig). + # Typing it as VllmSpeculativeConfig | None would make pydantic convert the + # dict early and break that step with "SpeculativeConfig object has no + # attribute 'items'". Keep the Mapping shape and let vLLM do the conversion. speculative_config: Mapping[str, Any] | None = None quantization_config: _QuantizationConfigType = None diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 07dd69e23e0..598baf63975 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -34,14 +34,15 @@ logger = init_logger(__name__) # --------------------------------------------------------------------------- -# K 步记账 trace（默认关） +# Multi-frame bookkeeping trace (off by default). # # VLLM_OMNI_MINICPMO_KSTEP_ACCOUNT_TRACE=1 # -# 回答的问题：模型已经在发 row_stop，为什么请求不停？记每一步每个请求"被接受的 -# token 值"——若它恒为 0（continue），说明 stop 行没有变成请求 token 表里的 1， -# check_stop(stop_token_ids=[1]) 永远不会触发，请求就会一路走到 max_model_len。 -# 契约：诊断永不抛异常。 +# Question it answers: the model already emits stop rows, so why does the +# request not finish? It logs each request's accepted token values per step -- +# if they stay 0 (continue), the stop row never became a 1 in the request's +# token table and check_stop(stop_token_ids=[1]) never fires, so the request +# runs to max_model_len. Diagnostics never raise. # --------------------------------------------------------------------------- _KSTEP_ACCOUNT_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_ACCOUNT_TRACE" _KSTEP_ACCOUNT_TRACE: dict = {"lines": 0, "reqs": {}, "off": False} diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 573e58ee8d8..2177cd90b68 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -20,17 +20,9 @@ connectors: connector_get_max_wait_first_chunk: 3000 connector_get_max_wait: 300 codec_chunk_frames: 25 - # 第 13 项（2026-09-20 实测采纳）：首块缝 25 帧 —— 首个 payload 用更小的窗口， - # 让音频首包更早出来。实测（r17 对照 r16）：c1 -7.4%、c4 +2.8%、c8 +13.3% - # （相对 r14 基线仍 -3.8%）；WER 干净态复测 1.2179% = 基线，零质量代价。 - # - # !! 生效条件 / 迁移依赖：本项仅在 initial < codec_chunk_frames 时生效。 - # 迁移来源分支 rtf-validated-optimizations 的稳态块长是 75，故那里的 25 生效； - # 而 main 此处的稳态块长是 25，25 >= 25 会按 _initial_chunk_frames 的回落规则 - # **视为关闭**（等价于不加这一行，不会出错）。若要让本项在 main 上真正生效， - # 需先把基座 codec_chunk_frames 调到 >25（例如 75）—— 那是基座配置差异， - # 不在本项迁移范围内，需单独决策。不要把这里调成 <25：历史上 12 帧会引入 - # 中缝、WER 翻倍（kstep_910b_validation_and_self_review.md §11.1）。 + # Smaller first payload so the first audio packet is emitted earlier. + # Takes effect only when initial < codec_chunk_frames; values far below + # the steady-state chunk length add a seam between chunks and hurt WER. initial_codec_chunk_frames: 25 codec_left_context_frames: 3 # HiFT capture buckets are derived from codec_chunk_frames and @@ -48,23 +40,10 @@ stages: # Leave headroom on the shared GPU for Talker KV + concurrent Code2Wav. gpu_memory_utilization: 0.55 enforce_eager: false - # 第 2 项（图模式 + 前缀缓存 + stage2 去 eager）—— "前缀缓存"部分。 - # 顶层 enable_prefix_caching 保持 false（双工安全），这里给 stage0 单独打开， - # 让 stage0 的重复前缀得以复用。 - # 该项另两部分 main 上已有，无需迁移：图模式（enable_hift_graph / - # enable_cfm_graph）、stage2 去 eager（code2wav_enable_npu_graph）。 - # 收益（旧代码实测口径，仅方向参考）：-22.0% / -3.9% / +1.1%（三档）。 + # Stage-level prefix caching override; the top level stays false for duplex safety. + # Stage-0 n-gram speculative decoding is intentionally not enabled here: vLLM + # rejects prefix caching together with speculative decoding. enable_prefix_caching: true - # 第 3 项（stage0 ngram 投机，k=14）：N-gram speculative decoding。 - # async_scheduling 必须保持 false —— vLLM 的 VllmConfig 校验拒绝 - # "async_scheduling + ngram proposer" 这个组合。 - # 收益（旧代码实测口径，仅方向参考）：RTF -11.8% / -3.7% / -5.9%（三档）、TTFT -69%。 - async_scheduling: false - speculative_config: - method: ngram - num_speculative_tokens: 14 - prompt_lookup_min: 1 - prompt_lookup_max: 1 tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -154,9 +133,8 @@ platforms: compilation_config: cudagraph_mode: PIECEWISE - stage_id: 1 - # 第 9 项（2026-09-20 实测采纳）：stage1 并发上限 4→8。 - # npu 覆盖块原先只写了 stage0 / stage2，stage1 沿用基座的 4，这一格是漏配。 - # 收益：r16 对照 r14 基线 c8 −15.1%、c4 −3.5%、c1 −1.1%。 + # Stage-1 concurrency cap raised from 4 to 8. The NPU override block + # only covered stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 compilation_config: @@ -166,9 +144,9 @@ platforms: # estimator kernels through the NPU platform patch. additional_config: code2wav_enable_npu_graph: true - # 第 10 项（2026-09-20 实测采纳）：NPU 图池 32→64。32 个形状不够时会触顶， - # 之后的形状永久退化为 eager。收益（r34 对照 r22）：c4 -17.0% / c8 -23.5% / c1 +5.8%。 - # 注：该项以第 11 项（CFM mel 分桶，bucket_frames=16）为前置稳定性支撑 - # （把捕获形状数从 64 压到 19，池扩容才可持续）；main 上已有 - # cfm_graph_bucket_frames: 16，故第 11 项无需迁移。 + # NPU graph pool raised from 32 to 64. Once the 32 captured shapes are + # exhausted, later shapes permanently degrade to eager. The CFM mel + # bucketing already present in main (cfm_graph_bucket_frames: 16) keeps + # the captured shape count at 19, which is what makes a pool this size + # sustainable. code2wav_max_npu_graphs: 64 diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 26084b4fd48..9cf15ffa27c 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1157,12 +1157,13 @@ def _project_omni_stage_engine_args( ): engine_args.update(_project_upstream_config_fields(config, field_map)) - # 迁移补齐（2026-09-21）：这一格与上面两项同类 —— 都是 vllm 自己的配置对象， - # 从 stage 配置原样搬进引擎参数。原先只搬 compilation_config / profiler_config， - # 于是 deploy 配置里的 speculative_config（stage-0 的 ngram 投机、以及 - # stage_config.py 给 stage-1 注入的 K 步投机）会在这一步静默消失，不报错， - # 只在引擎侧表现为 `speculative_config=None`；K 步与 ngram 都会失去效果。 - # 用 getattr(..., None) 兜底，避免个别 stage 配置类型没有该字段时炸掉。 + # Same category as the two above: vLLM config objects copied verbatim from + # the stage config into the engine args. Only compilation_config and + # profiler_config used to be copied, so a speculative_config in the deploy + # config (stage-0 n-gram, or the stage-1 one injected by stage_config.py) + # silently disappeared here (no error) and showed up engine-side as + # speculative_config=None. getattr(..., None) guards stage config types + # that lack the field. for name in ("compilation_config", "profiler_config", "speculative_config"): value = getattr(stage_config, name, None) if value is not None: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 9bed90f08c5..34fd101412b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -281,9 +281,10 @@ def __init__( raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be >= 0") if self._initial_batch_size and self._initial_batch_size < self._min_batch_size: raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be 0 or >= code2wav_min_batch_size") - # 第 49 项（2026-09-20 实测采纳，env 门控默认关）：强制 stage2 每个解码批 - # 只有 1 行，使 CFM 确定性地走 even 路径（批内等长 → 图回放），绕开 - # mixed-length 批的 ragged eager 链。数值路径零改动（even 是原有代码路径）。 + # Force stage 2 to decode a single row per batch so CFM deterministically + # takes the even path (equal-length rows -> graph replay) instead of the + # ragged eager chain for mixed-length batches. Numerics are unchanged + # (even is an existing code path). self._single_row_batch = os.environ.get("VLLM_OMNI_C2W_SINGLE_ROW", "") == "1" self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") @@ -671,8 +672,9 @@ def _iter_decode_batches( ) -> Iterable[list[_WorkItem]]: for bucket in buckets: if self._single_row_batch: - # 第 49 项：每批 1 行 → 每次 CFM 调用都落在 even 路径上 - # （批内等长 → 图回放）。刻意绕过下面的波次切分与 min-batch 门。 + # One row per batch -> every CFM call lands on the even path + # (equal-length rows -> graph replay), intentionally bypassing + # the wave split and min-batch gate below. for item in bucket: yield [item] continue diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index a6250b53a66..d7306af05d0 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -634,20 +634,22 @@ def make_omni_output(self, model_outputs, **kwargs): return model_outputs return self.talker.make_omni_output(model_outputs, **kwargs) - # 第 8 项（K 步多帧解码）在 runner 侧的四个转发（2026-09-21 迁移补齐）。 - # 基线的这四个方法就写在这个外层包装类上，迁移时整段漏掉。runner 只看得见这个 - # 包装类（`_model_forward` 里 talker_multiframe.applies / is_multi_token_decode / - # 以及 `ensure_stop_token_vocab` 的判定全部是 `getattr(model, ...)`），所以： - # · 缺 `supports_multi_frame_decode` → 判定静默取到 False → 多帧循环永不进入， - # 且连"拒绝原因"都不打日志（applies 那条分支是静默返回 0）； - # · 缺 take/set_batch_stop_logits → 每帧的 stop 行拿不到； - # · 缺 merge_frame_outputs → 多帧结果无法合并。 - # 后果：K 步的投机配置已经进了引擎（vLLM 层在做投机解码），但 runner 侧的多帧 - # 循环不进入、循环里的词表宽度修正（ensure_stop_token_vocab）也不执行， - # 于是拒绝采样器按宽度 0 消费、AIV 越界（507035）、stage-1 引擎崩、音频全空。 + # Runner-side forwards for multi-frame decode. The baseline declares these + # four methods on this outer wrapper class; the runner only ever sees the + # wrapper (talker_multiframe.applies / is_multi_token_decode / + # ensure_stop_token_vocab all resolve through getattr(model, ...)), so + # without them: + # * missing supports_multi_frame_decode -> the check silently reads False, + # the multi-frame loop never runs and no reason is logged; + # * missing take/set_batch_stop_logits -> per-frame stop rows unavailable; + # * missing merge_frame_outputs -> multi-frame results cannot be merged. + # The speculative config alone is not enough: the runner-side loop and its + # vocab-width correction (ensure_stop_token_vocab) stay dormant, the + # rejection sampler consumes width-0 rows, AIV goes out of bounds (507035) + # and the stage-1 engine dies with empty audio. @property def supports_multi_frame_decode(self) -> bool: - # 这些方法在两类 stage 上都在，只有 Talker（tts）能走多帧，故按 stage 判定。 + # Present on both stage kinds; only the Talker (tts) path uses it. return self.model_stage == "tts" def take_batch_stop_logits(self): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 3f50d8e5ebd..532160f0b48 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -343,19 +343,21 @@ def resolve_codec_sampling_params( # --------------------------------------------------------------------------- -# K 步逐帧诊断（默认关，与 talker_multiframe 的逐步计时互补） +# Per-frame trace for the multi-frame loop (off by default; complements the +# per-step timers in talker_multiframe). # -# VLLM_OMNI_MINICPMO_KSTEP_FRAME_TRACE=1 打开 -# 最多打 _KSTEP_FRAME_TRACE_FRAMES 帧，然后静默 +# VLLM_OMNI_MINICPMO_KSTEP_FRAME_TRACE=1 enable +# prints at most _KSTEP_FRAME_TRACE_FRAMES frames, then goes quiet # -# 回答逐步计时回答不了的那个问题：K 循环里**模型自己的采样器**有没有走到 EOS， -# 以及每一帧实际吃进去的 code 是不是上一帧的采样值（契约不变式 I1）。 -# 契约：诊断永不许反杀引擎 —— 任何异常只把本诊断永久关掉并记一条 warning。 +# Answers what the per-step timers cannot: whether the model's own sampler +# reaches EOS inside the K loop, and whether each frame is fed the previous +# frame's sampled code (frame-order invariant). Diagnostics never take the +# engine down: on any error they disable themselves and log one warning. # --------------------------------------------------------------------------- _KSTEP_FRAME_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_FRAME_TRACE" -_KSTEP_FRAME_TRACE_LINES = 128 # 打点行数上限 -_KSTEP_FRAME_TRACE_HEAD = 16 # 头 N 帧逐帧（用于校验帧序不变式 I1） -_KSTEP_FRAME_TRACE_STRIDE = 16 # 之后每 N 帧一条（覆盖到约第 1800 帧，够看 EOS） +_KSTEP_FRAME_TRACE_LINES = 128 # max printed lines +_KSTEP_FRAME_TRACE_HEAD = 16 # first N frames printed one by one (frame-order check) +_KSTEP_FRAME_TRACE_STRIDE = 16 # then one line every N frames (covers ~1800 frames) _KSTEP_FRAME_TRACE: dict[str, Any] = {"reqs": {}, "lines": 0, "off": False} @@ -399,7 +401,7 @@ def _trace_kstep_frame(model: Any, **fields: Any) -> None: device_states.get(fields.get("request_id")) if isinstance(device_states, dict) else None ) device_step = int(device_state.step.reshape(-1)[0].item()) if device_state is not None else -1 - state["lines"] += 1 # 只防刷屏，窗口按请求单独计数 + state["lines"] += 1 # cap printed lines only; per-request windows count separately logger.info( "[kstep-frame] frame=%d req=%s step=%s->%s device_step=%s in_code=%s row=%s " "sampled=%s is_eos=%s limit=%s finished=%s ctx=%s min=%s max=%s", @@ -1150,16 +1152,18 @@ def make_omni_output( else: codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1) step = _codec_int_param(state, "step", 0) - # 迁移注记（2026-09-21，待设备复测后决定是否接线）：main 侧的 turn-end - # 语义要求 turn-end 排空块在每个 25 帧 cadence 边界后的 5 步内屏蔽 codec - # EOS（见本文件 _turn_end_boundary_eos_masked 与 state["turn_end_drain"]， - # MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS 为 4 个 chunk，故窗口落在 - # step 25-29 / 50-54 / 75-79 / 100-103）。第 8 项把 EOS 屏蔽移进了采样核 - # （talker_codec_sample.prepare_codec_logits / greedy_codec_sample），其判据 - # 是 state.step < min_tokens，而 min_tokens 是每请求缓存的张量 - # （_request_codec_device_inputs），此处无法逐帧表达该周期窗口 —— 因此本项 - # 迁移后该窗口屏蔽未接线，只影响双工 turn-end（离线/TTS 路径不受影响）。 - # 本文件保留了 main 的常量、budget 与 turn_end_drain 状态标志。 + # Note: main's turn-end semantics require masking codec EOS within 5 + # steps after each 25-frame cadence boundary (see + # _turn_end_boundary_eos_masked and state["turn_end_drain"]; + # MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS is 4 chunks, so the window + # lands on steps 25-29 / 50-54 / 75-79 / 100-103). The EOS masking + # moved into the sampling core (talker_codec_sample. + # prepare_codec_logits / greedy_codec_sample), whose criterion is + # state.step < min_tokens with min_tokens a per-request cached + # tensor, so that periodic window cannot be expressed per frame here + # and is left unwired. It only affects duplex turn-end; the + # offline/TTS path is unaffected. The constants, budget and + # turn_end_drain flag from main are kept in this file. min_tokens = _codec_int_param(state, "min_tokens", self._codec_min_tokens) max_tokens = _codec_int_param(state, "max_tokens", self._codec_max_tokens) if self._codec_temperature == 0.0: diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index b2e3dcbc8bc..f91454df05b 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -164,10 +164,11 @@ def _codec_config(transfer_manager: Any) -> tuple[int, int]: def _initial_chunk_frames(transfer_manager: Any, chunk_frames: int) -> int: - """第 13 项（2026-09-20 实测采纳）：首个 payload 的窗口帧数。 + """First-payload window size in frames. - 读 connector extra 里的 ``initial_codec_chunk_frames``；未配置、非正数、 - 或不小于稳态 ``codec_chunk_frames`` 时一律回落成稳态值（等价于本项关闭）。 + Reads ``initial_codec_chunk_frames`` from the connector extra; falls back to + the steady-state value when absent, non-positive, or not smaller than the + steady ``codec_chunk_frames`` (which means the feature is off). """ connector = getattr(transfer_manager, "connector", None) raw_config = getattr(connector, "config", {}) or {} @@ -347,9 +348,10 @@ def tts2code2wav_async_chunk( request_finished = getattr(request, "is_finished", None) finished = bool(is_finished or (callable(request_finished) and request_finished())) chunk_frames, left_context_frames = _codec_config(transfer_manager) - # 第 13 项（2026-09-20 实测采纳）：首块窗口。首个 payload（此时 chunk_seq 还是 0） - # 用 initial_codec_chunk_frames，之后的 payload 回到稳态 chunk_frames； - # 未配置时 first_window == chunk_frames，行为与改动前逐位一致。 + # First-payload window: the payload with chunk_seq == 0 uses + # initial_codec_chunk_frames; later payloads return to the steady + # chunk_frames. When unset, first_window == chunk_frames and behavior is + # unchanged. first_window = ( _initial_chunk_frames(transfer_manager, chunk_frames) if int(record["chunk_seq"]) == 0 diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 26eb94120c6..c0221f94e64 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -77,8 +77,8 @@ def __init__(self) -> None: apply_qwen3_tts_patches() apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches() - # 第 8 项支撑（910C 部署防护）：K 步 Triton rejection/penalties 的 - # warmup 守卫必须在 ascend 内核 warmup 之前装好。 + # The Triton rejection/penalties warmup guard for multi-frame decode must + # be installed before the Ascend kernel warmup runs. apply_ascend_warmup_patch() @classmethod diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py index 56d27a0d7a7..55a666334d2 100644 --- a/vllm_omni/platforms/npu/worker/talker_multiframe.py +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -68,16 +68,17 @@ _LOGGED_ENGAGE = False # --------------------------------------------------------------------------- -# K 步逐帧计时（默认关，关掉时 run() 里一个 timer 都不取） +# Per-step/per-frame timers for the multi-frame loop (off by default; when off, +# run() takes no timers at all). # -# VLLM_OMNI_MINICPMO_KSTEP_PROF=1 打开 -# VLLM_OMNI_MINICPMO_KSTEP_PROF_STEPS=3 前 N 个 step 逐帧打点并做设备同步取样 -# VLLM_OMNI_MINICPMO_KSTEP_PROF_EVERY=25 每 N 个 step 打一行累计摘要 +# VLLM_OMNI_MINICPMO_KSTEP_PROF=1 enable +# VLLM_OMNI_MINICPMO_KSTEP_PROF_STEPS=3 sample the first N steps with a device sync +# VLLM_OMNI_MINICPMO_KSTEP_PROF_EVERY=25 print a cumulative summary every N steps # -# 要回答的问题只有一个：一步 8 帧的墙钟时间，到底花在 host 入队（replay/after_forward） -# 还是花在设备执行上。replay/after_forward 是入队耗时（异步），synced 是在其后加一次 -# 设备同步测到的真实帧耗时 —— 两者相差多少，就是设备侧的账。 -# 诊断代码永不许抛异常：任何异常只把本 profiling 永久关掉并记一条 warning。 +# Purpose: split one K-frame step's wall time between host enqueue (replay/ +# after_forward, asynchronous) and device execution (synced, measured after an +# explicit synchronize). Diagnostics never raise: on any error they turn +# themselves off permanently and log one warning. # --------------------------------------------------------------------------- _PROF_ENV = "VLLM_OMNI_MINICPMO_KSTEP_PROF" _PROF_STEPS_ENV = "VLLM_OMNI_MINICPMO_KSTEP_PROF_STEPS" @@ -155,7 +156,7 @@ def _prof_frame_line(step: int, frame: int, replay_ms: float, after_ms: float, s def _prof_summary(steps: int, frames_total: int, last_frames: int, *, detail: bool) -> None: - """累计一行：每步/每帧墙钟，以及 host 入队与设备执行各自的账。""" + """One cumulative line: per-step/per-frame wall time, host enqueue vs device execution.""" try: per_step = _PROF_TOTAL_MS / steps if steps else 0.0 per_frame = _PROF_TOTAL_MS / frames_total if frames_total else 0.0 @@ -256,13 +257,15 @@ def _prof_step_end(prof: "_StepProf | None", step_start: float) -> None: except Exception: _prof_disabled("step_end") # --------------------------------------------------------------------------- -# K 步 stop 取证（默认关） +# Stop-row forensics for the multi-frame loop (off by default). # # VLLM_OMNI_MINICPMO_KSTEP_STOP_TRACE=1 # -# 只回答一个问题：一步 K 帧里，模型到底有没有产出 stop 行。有 stop 行而请求 -# 仍跑到 max_tokens，丢在 vLLM 的 spec 记账侧；一行都没有，丢在多帧的采样/状态 -# 传递侧。前 50 步各打一行，之后静默。诊断组件永不抛异常。 +# Answers one question: does the model actually emit stop rows within the K +# frames of a step? Rows present but the request still runs to max_tokens means +# the loss is in vLLM's spec bookkeeping; no rows at all means it is in the +# multi-frame sampling/state path. Prints the first 50 steps, then goes quiet. +# Never raises. # --------------------------------------------------------------------------- _STOP_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_STOP_TRACE" _STOP_TRACE_STEPS = 50 @@ -279,7 +282,7 @@ def stop_trace_enabled() -> bool: def _trace_stop_rows(frame_stop_logits: list[torch.Tensor]) -> None: - """每步一行：每个请求的 K 帧里各有几帧的 stop 行是 stop。""" + """One line per step: how many of each request's K frames carry a stop row.""" global _STOP_TRACE_COUNT try: if _STOP_TRACE_COUNT >= _STOP_TRACE_STEPS: @@ -301,20 +304,21 @@ def _trace_stop_rows(frame_stop_logits: list[torch.Tensor]) -> None: def trace_kstep_bookkeeping(runner: Any, valid_sampled_token_ids: Any, logits: Any) -> None: - """[kstep-stop] 第二层：stop 走完 spec 记账之后还剩什么。 + """Stop trace, layer 2: what survives spec bookkeeping. - 两条事实一起打，一次就能把"stop 去哪了"钉死在某一层： + Two facts are printed together so the loss point is unambiguous: - * ``carried`` -- rejection sampler 交给请求的 token 序列（前 4 个请求）。 - 里面有 1 说明 stop 进了请求自己的 token 表，后面就只剩 engine 侧的结束判定； - 一排 0 而没有 1 说明 stop 在到这里之前就没了。 - * ``min_toks`` -- vLLM 的 MinTokens 处理器里还有哪些请求在 mask 名单上， - 以及它看到的 output 长度（``(index, min_tokens, len(output_token_ids), - sorted(stop_ids))``）。**只要某个请求还在这张名单上，它的 stop 列就是 - -inf**；如果 ``len(output_token_ids)`` 始终不涨，这张名单就永远不会被 - 清掉，请求也就永远停不下来。 + * ``carried`` -- tokens the rejection sampler handed to the request (first + 4 requests). A 1 means the stop token reached the request's own token + table; then only the engine-side finish check remains. All zeros with no + 1 means the stop disappeared earlier. + * ``min_toks`` -- requests still on the MinTokens mask list, with the output + length it sees: ``(index, min_tokens, len(output_token_ids), + sorted(stop_ids))``. While a request is on this list its stop columns are + -inf; if ``len(output_token_ids)`` never grows, the entry is never + cleared and the request can never stop. - 诊断组件永不抛异常：出错只关掉自己。 + Never raises: on error it only disables itself. """ global _STOP_KB_COUNT try: @@ -349,21 +353,25 @@ def trace_kstep_bookkeeping(runner: Any, valid_sampled_token_ids: Any, logits: A def neutralize_kstep_min_tokens(logitsprocs: Any) -> None: - """K 步下让 vLLM 层的 ``min_tokens`` 不再 censoring 停止信号。 - - yaml 的 ``min_tokens: 50`` 在 K 步下是通过 **把 stop token（id 1）的 logit - 置成 -inf** 实现的（``MinTokensLogitsProcessor``）。它自己的解除条件是 - ``len(output_token_ids) >= min_tokens``，而这个长度归 vLLM 的 spec 记账管； - 只要那个计数没有推进到位，请求唯一的停止信号就被永久 mask， - ``finished_reason`` 只能是 ``length``（910C 现场：stage1 全 length、0 stop）。 - - 真正的"最小帧数"保护不在这里：``talker_codec_sample.greedy_codec_sample`` - 用 ``state.step < min_tokens`` 把 codec EOS 本身压成 -inf，用的是模型自己 - 的帧计数，与 vLLM 的 token 记账无关。所以这一层是重复的、且在 K 步下会 - 把停止信号一起吃掉 —— 清空它的名单，把最小长度交还给模型内的那道保护。 - - ``min_toks`` 每个 step 都会被 ``update_state`` 重新登记，所以本函数必须 - 每步调用（调用点在 runner 的采样前）。永不抛异常。 + """Stop vLLM's ``min_tokens`` layer from censoring the stop signal. + + The yaml ``min_tokens: 50`` is enforced by ``MinTokensLogitsProcessor`` + writing -inf into the stop-token logits. It unmasks once + ``len(output_token_ids) >= min_tokens``, a counter owned by vLLM's spec + bookkeeping; while that counter lags, the request's only stop signal stays + masked and ``finished_reason`` can only be ``length`` (observed on stage 1: + all length, zero stop). + + The real minimum-length guard lives elsewhere: + ``talker_codec_sample.greedy_codec_sample`` masks the codec EOS itself via + ``state.step < min_tokens``, using the model's own frame counter. This vLLM + layer is therefore redundant and, under multi-frame decode, also swallows + the stop signal -- clear its list and leave the minimum to the model-side + guard. + + ``min_toks`` is re-registered by ``update_state`` every step, so this must + be called every step (the caller places it right before sampling). Never + raises. """ try: for proc in getattr(logitsprocs, "non_argmax_invariant", None) or []: @@ -605,10 +613,10 @@ def begin_narrow_step( from vllm.config import CUDAGraphMode from vllm.forward_context import BatchDescriptor - # 迁移注记（2026-09-21）：narrow replay 所需的 KV 几何由第 6 项 - # （platforms/npu/attention/fixed_kv_decode）提供，本分支未迁第 6 项。 - # 这里容错导入：模块缺席时按本函数既有承诺 decline，调用方回落 wide replay - # （见 docstring），K 步解码与停止链不受影响。 + # The KV geometry required by narrow replay lives in + # platforms/npu/attention/fixed_kv_decode, which may be absent. Import + # defensively: when missing, decline as documented and let the caller fall + # back to wide replay; multi-frame decode and the stop chain are unaffected. try: from vllm_omni.platforms.npu.attention import fixed_kv_decode except Exception: From d7112a4ff1de4a3780a8181b8c1d4bdf02a33a60 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 09/89] test(minicpmo_4_5): add no-device regression tests for the Talker multi-frame feedback chain Pins the frame-code feedback (request-state last_code), the decode preprocess priority, the scheduler K-frame guard and the multi-frame gate matrix. Pure CPU; no device required. Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 446 ++++++++++++++++++ 1 file changed, 446 insertions(+) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py new file mode 100644 index 00000000000..e8c2c658360 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -0,0 +1,446 @@ +"""No-device regression tests for the MiniCPM-o Talker K-step feedback chain. + +Layers pinned down here: + +- ``make_omni_output`` must record each frame's codec sample in the request + state (``last_code``) while the multi-frame loop owns sampling: the rows + vLLM schedules for the next step carry placeholder continue ids, so + without the record every frame after the first embeds a placeholder + instead of the real previous frame (silent audio corruption). +- The decode preprocess must prefer that state record over input_ids. +- The scheduler-side K-frame guard must drop continuation drafts whenever + prefill work is pending, so no step ever mixes prefill rows with + multi-row decode spans (the mixed batch that crashed a 64/4 run). +- The multi-frame gate matrix: non-uniform decode spans stay blocked at + ``applies()``; the runner raise remains the assertion of last resort for + a combination the guard makes unschedulable. +""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_TALKER_FRAMES_ENV = "VLLM_OMNI_MINICPMO_TALKER_FRAMES" +_NUM_AUDIO_TOKENS = 6562 +_EOS_ID = _NUM_AUDIO_TOKENS - 1 + + +def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): + """Bare talker instance: no config, no weights, deterministic sampler. + + ``emb_code`` is crafted so the placeholder fallback (row 0) is all + zeros while real codec rows are not -- a sharp contrast for asserting + which id got embedded. + """ + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__( + MiniCPMO45OmniTTSForConditionalGeneration + ) + nn.Module.__init__(model) + model._k_step_frames = k_step_frames + model.supports_multi_frame_decode = k_step_frames > 0 + model._codec_temperature = 0.0 + model._codec_min_tokens = 0 + model._codec_max_tokens = 4032 + model._num_audio_tokens = _NUM_AUDIO_TOKENS + model._codec_eos_id = _EOS_ID + model._request_audio_states = {} + emb = nn.Embedding(_NUM_AUDIO_TOKENS, 4) + with torch.no_grad(): + emb.weight.zero_() + emb.weight[42] = 1.0 + emb.weight[43] = 2.0 + model.emb_code = nn.ModuleList([emb]) + queue = iter(scripted_samples) + + def _greedy(_hidden, _codes, _request_id, _step, _min_tokens, _max_tokens): + return torch.tensor(next(queue), dtype=torch.long) + + model._sample_audio_code_greedy = _greedy + return model + + +def _frame_call(model, hidden): + return model.make_omni_output( + hidden, + model_intermediate_buffer=[{"request_id": "r1"}], + request_token_spans=[(0, 1)], + request_sample_eligible=[True], + ) + + +def test_kstep_frame_sample_is_recorded_for_next_frame(): + model = _make_talker(k_step_frames=8, scripted_samples=[42, 43, _EOS_ID]) + state = {"step": 0, "codes": torch.tensor([10, 11])} + model._request_audio_states["r1"] = state + hidden = torch.randn(1, 8) + + out0 = _frame_call(model, hidden) + assert state["last_code"] == 42 + assert state["step"] == 1 + assert out0.multimodal_outputs["codes"]["audio"][0].reshape(-1).tolist() == [42] + + out1 = _frame_call(model, hidden) + assert state["last_code"] == 43 + assert out1.multimodal_outputs["codes"]["audio"][0].reshape(-1).tolist() == [43] + + # The EOS frame terminates the request: the record keeps the last real + # sample (there is no next frame to feed), and the delta goes empty. + out2 = _frame_call(model, hidden) + assert state["last_code"] == 43 + assert state["finished"] is True + assert out2.multimodal_outputs["codes"]["audio"][0].numel() == 0 + + +def test_legacy_single_frame_path_does_not_touch_last_code(): + model = _make_talker(k_step_frames=0, scripted_samples=[42]) + state = {"step": 0, "codes": torch.tensor([10, 11])} + model._request_audio_states["r1"] = state + + _frame_call(model, torch.randn(1, 8)) + assert "last_code" not in state + + +def test_decode_preprocess_embeds_state_last_code(): + model = _make_talker(k_step_frames=8, scripted_samples=[]) + state = {"step": 3, "codes": torch.tensor([42]), "last_code": 42} + model._request_audio_states["r1"] = state + # The scheduled row carries the placeholder continue id (0), not the + # real previous frame's sample. + input_ids = torch.zeros(1, dtype=torch.long) + + _, embeds, out = model.preprocess( + input_ids, + None, + request_id="r1", + audio_state=state, + _omni_is_prefill=False, + ) + assert torch.equal(embeds, model.emb_code[0](torch.tensor([42]))) + # Row 0 (the placeholder fallback) is all zeros in this fixture, so a + # non-zero embed proves the state record won. + assert torch.count_nonzero(embeds) == embeds.numel() + assert out["codes"]["audio"].reshape(-1).tolist() == [42] + + +def _make_scheduler(*, num_spec: int, waiting=(), running=()): + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + sched = OmniARScheduler.__new__(OmniARScheduler) + sched._omni_talker_kstep_cache = None + sched.speculative_config = SimpleNamespace(method="ngram", num_speculative_tokens=num_spec) + # vLLM's Scheduler stores the count here (vllm_config.num_speculative_tokens); + # the guard predicts row widths from it. + sched.num_spec_tokens = num_spec + sched.max_model_len = 40960 + sched.waiting = list(waiting) + sched.running = list(running) + return sched + + +def _req(*, computed: int, prompt: int, spec: list[int], total: int | None = None): + return SimpleNamespace( + num_computed_tokens=computed, + prompt_token_ids=[0] * prompt, + # num_tokens is prompt + generated; the guard reads the difference to + # predict how many rows this request will schedule this step. + num_tokens=prompt if total is None else total, + spec_token_ids=list(spec), + ) + + +def test_guard_drops_drafts_when_waiting_request_pending(monkeypatch): + monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") + decode_req = _req(computed=100, prompt=100, spec=[0] * 7) + sched = _make_scheduler(num_spec=7, waiting=[object()], running=[decode_req]) + + sched._drop_talker_drafts_if_prefill_pending() + assert decode_req.spec_token_ids == [] + + +def test_guard_keeps_drafts_when_no_prefill_pending(monkeypatch): + monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") + decode_req = _req(computed=100, prompt=100, spec=[0] * 7, total=101) + sched = _make_scheduler(num_spec=7, waiting=[], running=[decode_req]) + + sched._drop_talker_drafts_if_prefill_pending() + assert decode_req.spec_token_ids == [0] * 7 + + +def test_guard_drops_drafts_when_chunked_prefill_in_flight(monkeypatch): + monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") + decoding_req = _req(computed=100, prompt=100, spec=[0] * 7, total=101) + chunking_req = _req(computed=50, prompt=100, spec=[]) + sched = _make_scheduler(num_spec=7, waiting=[], running=[decoding_req, chunking_req]) + + sched._drop_talker_drafts_if_prefill_pending() + assert decoding_req.spec_token_ids == [] + + +def test_guard_noop_for_text_stage_spec_config(monkeypatch): + monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") + text_req = _req(computed=100, prompt=100, spec=[0] * 15) + sched = _make_scheduler(num_spec=15, waiting=[object()], running=[text_req]) + + sched._drop_talker_drafts_if_prefill_pending() + assert text_req.spec_token_ids == [0] * 15 + + +def test_guard_arms_from_vllm_config_and_drops_drafts_for_a_running_chunk(monkeypatch): + # Real vLLM Scheduler instances expose no .speculative_config attribute + # (and SchedulerConfig has no num_speculative_tokens), so the armed check + # must read vllm_config.speculative_config -- otherwise the whole guard + # silently stays off while the engine-side loop is armed (the 11:47 and + # 12:04 crashes on 910B). + # + # 2026-09-18 12:39 crash: spans=[(0,7),(7,15)] -- a request carrying a + # 7-token streaming chunk shares the step with a steady-state decode that + # vLLM pads to 1+7 rows. Both spec_token_ids lists look innocent ([] and + # 7 placeholders), so the guard must predict row widths from + # num_tokens - num_computed_tokens instead. + from types import SimpleNamespace as NS + + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") + sched = OmniARScheduler.__new__(OmniARScheduler) + sched._omni_talker_kstep_cache = None + sched.vllm_config = NS(speculative_config=NS(method="ngram", num_speculative_tokens=7)) + assert sched._talker_kstep_armed() is True + sched.num_spec_tokens = 7 + sched.max_model_len = 40960 + + steady = _req(computed=100, prompt=100, spec=[0] * 7, total=101) # 1 token -> 8 rows + chunk = _req(computed=107, prompt=107, spec=[], total=114) # 7-token chunk -> 7 rows + sched.waiting = [] + sched.running = [steady, chunk] + + sched._drop_talker_drafts_if_prefill_pending() + assert steady.spec_token_ids == [] + assert chunk.spec_token_ids == [] + + +def test_guard_keeps_drafts_when_all_decodes_share_the_padded_width(monkeypatch): + # A first-step decode (no placeholders yet) and a steady-state decode both + # schedule 1 token this step, so vLLM gives both the same 1+num_spec row + # width: spans stay uniform and nothing may drop. + monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") + steady = _req(computed=100, prompt=100, spec=[0] * 7, total=101) + first_step = _req(computed=100, prompt=100, spec=[], total=101) + sched = _make_scheduler(num_spec=7, waiting=[], running=[steady, first_step]) + + sched._drop_talker_drafts_if_prefill_pending() + assert steady.spec_token_ids == [0] * 7 + assert first_step.spec_token_ids == [] + + +def test_talker_stop_token_ids_follow_the_arming_decision(monkeypatch): + """Stage 1's stop id must be one the head that actually runs can emit. + + 2026-09-18 root cause: the pipeline constraint pinned ``stop_token_ids`` to + the codec EOS (6561) unconditionally. With the K-frame loop armed the + vLLM-level head is the two-wide continue/stop row, so the only sampleable + ids are 0/1, `check_stop` never matched, and every K-step request ran to + ``max_tokens`` (142s per request on 910B, model finished at frame ~116). + """ + from vllm_omni.model_executor.models.minicpmo_4_5 import pipeline as mcp_pipeline + from vllm_omni.platforms.npu.worker import talker_multiframe + + monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") + assert mcp_pipeline._talker_stop_token_ids() == [talker_multiframe.STOP_TOKEN_ID] + assert talker_multiframe.STOP_TOKEN_ID == 1 + + monkeypatch.setenv(_TALKER_FRAMES_ENV, "1") + assert mcp_pipeline._talker_stop_token_ids() == [6561] + + # No explicit frame count on a 910B box: the loop stays off, so the real + # codec head runs and its EOS is the stop. + monkeypatch.delenv(_TALKER_FRAMES_ENV, raising=False) + monkeypatch.setenv("SOC_VERSION", "ascend910b1") + assert mcp_pipeline._talker_stop_token_ids() == [6561] + + +def test_multiframe_gate_matrix(): + from vllm_omni.platforms.npu.worker import talker_multiframe + + model = SimpleNamespace(supports_multi_frame_decode=True) + state = {"step": 5} + + # Uniform K-row decode: the loop engages with K=8. + uniform = { + "request_token_spans": [(0, 8), (8, 16)], + "model_intermediate_buffer": [ + {"audio_state": dict(state)}, + {"audio_state": dict(state)}, + ], + } + assert talker_multiframe.applies(model, uniform) == 8 + + # Mixed decode spans (an 8-row drafted decode plus a 1-row decode): + # blocked, and flagged as a multi-token decode step. + mixed_decode = { + "request_token_spans": [(0, 8), (8, 9)], + "model_intermediate_buffer": [ + {"audio_state": dict(state)}, + {"audio_state": dict(state)}, + ], + } + assert talker_multiframe.applies(model, mixed_decode) == 0 + assert talker_multiframe.is_multi_token_decode(model, mixed_decode) is True + + # Prefill rows mixed with a drafted decode: blocked the same way. + mixed_prefill = { + "request_token_spans": [(0, 515), (515, 523)], + "model_intermediate_buffer": [ + {"audio_state": dict(state), "_omni_is_prefill": True}, + {"audio_state": dict(state)}, + ], + } + assert talker_multiframe.applies(model, mixed_prefill) == 0 + assert talker_multiframe.is_multi_token_decode(model, mixed_prefill) is True + + +def _vocab_runner(*, supports_multi_frame: bool, vocab_size: int): + return SimpleNamespace( + model=SimpleNamespace(supports_multi_frame_decode=supports_multi_frame), + input_batch=SimpleNamespace(vocab_size=vocab_size), + ) + + +def test_stop_vocab_gate_reports_the_two_wide_stop_row_not_the_hidden_width(): + """``input_batch.vocab_size`` must be the stop row's width (2), not hidden. + + ``InputBatch.add_request`` keeps ``top_k = vocab_size`` as its "no top-k" + sentinel and ``RejectionSampler.parse_output`` filters accepted tokens + against it; both only hold at 0 or 2. The gate used to be handed + ``text_hidden_states``, so it wrote the hidden width (768): that passes the + ``parse_output`` filter by accident while taking ``top_k`` out of its + sentinel range, and the two-wide stop row then never reaches the request's + token list -- every request runs to ``max_tokens`` instead of stopping on + EOS (the 910C scene where stage 1 reported ``finished_reason=length`` for + all 34 requests). + """ + from vllm_omni.platforms.npu.worker import talker_multiframe + + assert talker_multiframe.STOP_ROW_WIDTH == 2 + + runner = _vocab_runner(supports_multi_frame=True, vocab_size=0) + # Hidden-width rows on purpose: the width of this tensor is not the answer. + talker_multiframe.ensure_stop_token_vocab(runner, torch.randn(4, 768)) + assert runner.input_batch.vocab_size == 2 + + # Idempotent: arming again must not move it. + talker_multiframe.ensure_stop_token_vocab(runner, torch.randn(4, 768)) + assert runner.input_batch.vocab_size == 2 + + +def test_stop_vocab_gate_stays_off_without_multi_frame_or_rows(): + from vllm_omni.platforms.npu.worker import talker_multiframe + + # Single-frame models (K=1, stage 0/2) never take the branch that reads + # vocab_size; the gate must leave them exactly as they were. + single = _vocab_runner(supports_multi_frame=False, vocab_size=0) + talker_multiframe.ensure_stop_token_vocab(single, torch.randn(4, 768)) + assert single.input_batch.vocab_size == 0 + + # No rows at hand: None must not be read as "width unknown, arm anyway". + armed = _vocab_runner(supports_multi_frame=True, vocab_size=0) + talker_multiframe.ensure_stop_token_vocab(armed, None) + assert armed.input_batch.vocab_size == 0 + + +def test_stop_trace_is_off_by_default_and_never_kills_the_step(monkeypatch): + from vllm_omni.platforms.npu.worker import talker_multiframe + + monkeypatch.delenv(talker_multiframe._STOP_TRACE_ENV, raising=False) + monkeypatch.setattr(talker_multiframe, "_STOP_TRACE_COUNT", 0) + assert talker_multiframe.stop_trace_enabled() is False + + # Worst input first: the diagnostic must degrade, never raise into the step. + talker_multiframe._trace_stop_rows([None]) + assert talker_multiframe.stop_trace_enabled() is False + + # Armed and well-formed: one row per request, one column per frame. + monkeypatch.setattr(talker_multiframe, "_STOP_TRACE_COUNT", 0) + monkeypatch.setenv(talker_multiframe._STOP_TRACE_ENV, "1") + assert talker_multiframe.stop_trace_enabled() is True + keep = torch.tensor([[0.0, -float("inf")]]) + stop = torch.tensor([[-float("inf"), 0.0]]) + talker_multiframe._trace_stop_rows([keep, stop, stop]) + assert talker_multiframe.stop_trace_enabled() is True + + +class _MinTokensLogitsProcessor: + """名字就是契约：实现按 ``type(proc).__name__`` 匹配。""" + + def __init__(self, min_toks): + self.min_toks = min_toks + + +def test_kstep_min_tokens_neutralization_clears_the_censor_list(): + """K 步下 vLLM 层 min_tokens 的名单必须被清空。 + + 它 mask 的是唯一的停止信号（binary stop 行的 id 1），而解除条件是 + ``len(output_token_ids) >= min_tokens``——那个长度归 vLLM 的 spec 记账管， + 计数不推进请求就永远停不下来（910C 现场 stage1 全 ``length``、0 ``stop``）。 + 模型内已用 ``state.step < min_tokens`` 压 codec EOS，这一层是重复的。 + """ + from vllm_omni.platforms.npu.worker import talker_multiframe + + censor = _MinTokensLogitsProcessor({0: (50, [0, 0, 0], {1})}) + untouched = _MinTokensLogitsProcessor({0: (50, [0], {1})}) + untouched.__class__ = type("SomeOtherProcessor", (object,), {}) # 非目标处理器 + + talker_multiframe.neutralize_kstep_min_tokens( + SimpleNamespace(non_argmax_invariant=[untouched, censor]) + ) + assert censor.min_toks == {} + assert untouched.min_toks == {0: (50, [0], {1})} + + # 空名单 / 怪输入都不能炸。 + talker_multiframe.neutralize_kstep_min_tokens( + SimpleNamespace(non_argmax_invariant=[]) + ) + talker_multiframe.neutralize_kstep_min_tokens(None) + + +def test_kstep_bookkeeping_trace_is_total(monkeypatch): + """记账层诊断：正常输入打一行，最坏输入只关掉自己。""" + from vllm_omni.platforms.npu.worker import talker_multiframe + + monkeypatch.setenv(talker_multiframe._STOP_TRACE_ENV, "1") + monkeypatch.setattr(talker_multiframe, "_STOP_KB_COUNT", 0) + + runner = SimpleNamespace( + input_batch=SimpleNamespace( + sampling_metadata=SimpleNamespace( + logitsprocs=SimpleNamespace( + non_argmax_invariant=[ + _MinTokensLogitsProcessor({0: (50, [0, 0], {1})}) + ] + ) + ) + ) + ) + talker_multiframe.trace_kstep_bookkeeping(runner, [[0, 0, 1]], torch.zeros(2, 2)) + assert talker_multiframe.stop_trace_enabled() is True + + # 最坏输入：内部构造真的失败时必须降级（绝不抛回 decode 步里）并永久关掉 + # 自己，否则每个 step 都要付一次异常代价。`[5]` 让 list(row) 抛 TypeError。 + monkeypatch.setattr(talker_multiframe, "_STOP_KB_COUNT", 0) + talker_multiframe.trace_kstep_bookkeeping(SimpleNamespace(), [5], torch.zeros(2, 2)) + assert talker_multiframe.stop_trace_enabled() is False + + # 取不到 runner/输入是合法场景的优雅降级，不是失败：不能因此把诊断关掉， + # 否则一次空步就会让后面所有取证静默。 + monkeypatch.setenv(talker_multiframe._STOP_TRACE_ENV, "1") + monkeypatch.setattr(talker_multiframe, "_STOP_KB_COUNT", 0) + talker_multiframe.trace_kstep_bookkeeping(None, None, None) + assert talker_multiframe.stop_trace_enabled() is True From 6142cbfa34c01838bc90d72cea0a34d1e2e62c75 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 10/89] fix(minicpmo_4_5): keep the single-frame codec contract when multi-frame decode is off The multi-frame work moved codec sampling into a device-side state machine and, in doing so, dropped the host-state contract the default single-frame path relies on: - make_omni_output no longer recorded state step/recent_codes/finished or the force/mask codec-EOS routing, so empty-speech segments, chunk budgets and the duplex turn-end drain window lost their effect; - compute_logits no longer forced or masked EOS, and sample() no longer restored an EOS blanked by vLLM's min_tokens processor; - the offline budget recorded min_tokens as 0 instead of None. Split make_omni_output into a single-frame path (multi-frame off, the default) that keeps the upstream contract, plus the multi-frame path, and restore the codec-EOS routing in compute_logits/sample. The 12 upstream tests in test_talker_batching.py pass again; test_pipeline expectations follow the new stage-1 concurrency cap and Code2Wav graph pool. Signed-off-by: y-null --- .../models/minicpmo_4_5/test_pipeline.py | 9 +- .../minicpmo_4_5/test_talker_batching.py | 11 + .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 211 +++++++++++++++++- 3 files changed, 222 insertions(+), 9 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 92e96b25a63..9cf7b04b609 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -170,7 +170,9 @@ def test_deploy_resolves_three_stage_pipeline(self, filename: str, devices: list assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor) assert "hf_overrides" not in stages[1].yaml_engine_args if filename == "minicpmo_4_5.yaml": - assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4] + # Stage-1 caps out at 8 so concurrent Talker decodes are not + # serialized behind the base value of 4. + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 8, 4] memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] assert memory_utilizations == [ 0.55, @@ -210,9 +212,12 @@ def test_npu_graph_configuration_is_stage_scoped(self, filename: str) -> None: assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE" assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE" assert stages[2].yaml_engine_args["enforce_eager"] is True + # Only the single-card deploy raises the Code2Wav graph pool; the + # multi-card variants keep the base value. + expected_graph_pool = 64 if filename == "minicpmo_4_5.yaml" else 32 assert stages[2].yaml_engine_args["additional_config"] == { "code2wav_enable_npu_graph": True, - "code2wav_max_npu_graphs": 32, + "code2wav_max_npu_graphs": expected_graph_pool, } def test_pipeline_exposes_full_and_async_payload_hooks(self) -> None: diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py index 4335bb88e1d..88b29a45ea2 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py @@ -107,6 +107,17 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: talker._request_condition_states = {} talker._deferred_cleanup_ids = set() talker._tts_config = ConditionalChatTTSConfig() + # Codec sampling contract added by the multi-frame decode work: the fixture + # runs the deterministic greedy boundary (argmax over head_code) and never + # arms the multi-frame loop. + talker._codec_temperature = 0.0 + talker._codec_top_k = 0 + talker._codec_top_p = 1.0 + talker._codec_repetition_penalty = 1.0 + talker._codec_min_tokens = 0 + talker._codec_max_tokens = 4032 + talker._k_step_frames = 0 + talker.supports_multi_frame_decode = False talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) with torch.no_grad(): talker.head_code[0].weight.copy_(torch.eye(8, 2)) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 532160f0b48..33440dedbeb 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -462,6 +462,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = config self.vllm_config = vllm_config self._penalty_histories: list[torch.Tensor] | None = None + # Per-step codec EOS routing decided in make_omni_output: rows to force + # to EOS (finished segments / chunk budget hit) and rows to mask EOS on + # (min_tokens / duplex turn-end drain window). Both are consumed by + # compute_logits and cleared there; _pending_force_eos_rows carries the + # force decision through to sample() so vLLM's min_tokens processor + # cannot blank the EOS we just forced. + self._force_eos_rows: list[bool] | None = None + self._mask_eos_rows: list[bool] | None = None + self._pending_force_eos_rows: list[bool] | None = None self._request_audio_states: dict[str, dict[str, Any]] = {} # Mirrors upstream TTSStreamingGenerator._chunk_info: one committed # condition plus, during a rollover, one immutable recompute recipe. @@ -941,11 +950,10 @@ def preprocess( # min_new_token=50) comes from the deploy YAML. remaining = int(self._tts_config.max_position_embeddings) - target_len max_tokens = max(min(_OFFLINE_CODEC_MAX_NEW_TOKENS, remaining), 1) - # The sampler floor the old comment pointed at (upstream's - # min_new_token, resolved from the deploy YAML or the - # checkpoint). A None here would reach make_omni_output as a - # value, and int(None) kills the stage. - min_tokens = int(self._codec_min_tokens) + # None leaves the sampler floor owned by the deploy YAML; the + # multi-frame path resolves it to self._codec_min_tokens when + # the request state carries no explicit value. + min_tokens = None state: dict[str, Any] = { "finished": empty_condition, "step": 0, @@ -1048,6 +1056,12 @@ def make_omni_output( spans = kwargs.get("request_token_spans") if spans is None or len(spans) != len(infos): raise RuntimeError("MiniCPM-o continuous Talker requires one request_token_span per request") + if self._k_step_frames <= 0: + # Single-frame path (the default): one codec frame per step, sampled + # by vLLM. Keeps the host-state contract consumed by compute_logits + # and sample: state step/recent_codes/finished bookkeeping plus the + # force/mask codec-EOS routing. + return self._make_omni_output_single_frame(hidden, infos, spans) sample_eligible = kwargs.get("request_sample_eligible") if sample_eligible is None: sample_eligible = [True] * len(infos) @@ -1586,6 +1600,144 @@ def forward( inputs_embeds=inputs_embeds, ) + def _make_omni_output_single_frame( + self, + hidden: torch.Tensor, + infos: list[Any], + spans: list[Any], + ) -> OmniOutput: + """Single-frame Talker output: bookkeeping plus codec-EOS routing. + + Emits the codec delta sampled on the previous step and decides which + rows ``compute_logits`` must force to EOS (finished segments, chunk + budget reached) or mask EOS on (min_tokens, duplex turn-end drain + window). + """ + emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos) + native_duplex_flags: list[torch.Tensor] = [] + duplex_epochs: list[torch.Tensor] = [] + duplex_turn_ids: list[torch.Tensor] = [] + segment_texts_utf8: list[torch.Tensor] = [] + turn_end_flags: list[torch.Tensor] = [] + empty_delta = hidden.new_empty((0, 1), dtype=torch.long) + codec_deltas = [empty_delta for _ in infos] + terminal_flags = [torch.tensor(False, dtype=torch.bool) for _ in infos] + force_eos_rows = [False] * len(infos) + mask_eos_rows = [False] * len(infos) + empty_history = hidden.new_empty((0,), dtype=torch.long) + penalty_histories = [empty_history for _ in infos] + for index, info in enumerate(infos): + info_dict = info if isinstance(info, dict) else {} + native_duplex = info_dict.get("native_duplex") is True + if emit_duplex_metadata: + duplex_info = info_dict.get("duplex") + if not isinstance(duplex_info, dict): + duplex_info = {} + epoch = duplex_info.get("epoch", -1) + turn_id = duplex_info.get("turn_id", -1) + if native_duplex and not all( + isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in (epoch, turn_id) + ): + raise RuntimeError( + "MiniCPM-o native duplex Talker requires non-negative integer " + f"epoch and turn_id, got epoch={epoch!r}, turn_id={turn_id!r}" + ) + meta_info = info_dict.get("meta") + if not isinstance(meta_info, dict): + meta_info = {} + segment_text = meta_info.get("native_duplex_segment_text", "") if native_duplex else "" + if not isinstance(segment_text, str): + segment_text = "" + turn_eos_id = meta_info.get("turn_eos_token_id") + ids_info = info_dict.get("ids") + tts_ids = ids_info.get("tts") if native_duplex and isinstance(ids_info, dict) else None + if isinstance(tts_ids, torch.Tensor): + contains_turn_eos = isinstance(turn_eos_id, int) and bool( + torch.any(tts_ids.reshape(-1) == turn_eos_id).item() + ) + elif isinstance(tts_ids, (list, tuple)): + contains_turn_eos = isinstance(turn_eos_id, int) and turn_eos_id in tts_ids + else: + contains_turn_eos = False + native_duplex_flags.append(torch.tensor(native_duplex, dtype=torch.bool)) + duplex_epochs.append(torch.tensor(epoch if isinstance(epoch, int) else -1, dtype=torch.long)) + duplex_turn_ids.append(torch.tensor(turn_id if isinstance(turn_id, int) else -1, dtype=torch.long)) + segment_texts_utf8.append( + torch.tensor( + list(segment_text.encode("utf-8")), + dtype=torch.uint8, + ) + ) + turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool)) + + if not isinstance(info, dict): + continue + request_id = str(info.get("request_id", index)) + state = self._request_audio_states.get(request_id) + if not isinstance(state, dict): + state = dict(info.get("audio_state", {}) or {}) + self._request_audio_states[request_id] = state + empty_speech = bool(state.get("finished")) + codes = info.get("codes", {}) + audio = codes.get("audio") if isinstance(codes, Mapping) else None + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + codec_deltas[index] = audio.to(device=hidden.device, dtype=torch.long).reshape(-1, 1) + state["step"] = int(state.get("step", 0)) + 1 + # ``audio`` is the id sampled last step, i.e. exactly the + # history the logits computed below are scored against. + recent = state.get("recent_codes") + recent = (recent if isinstance(recent, list) else []) + codec_deltas[index].reshape(-1).tolist() + state["recent_codes"] = recent[-_CODEC_PENALTY_WINDOW:] + recent_codes = state.get("recent_codes") + if recent_codes: + penalty_histories[index] = torch.tensor(recent_codes, dtype=torch.long, device=hidden.device) + max_tokens = state.get("max_tokens") + min_tokens = state.get("min_tokens") + step = int(state.get("step", 0)) + # Duplex: 26 samples include the terminating EOS, so force it after + # 25 forwarded frames -- the same cadence as generate_chunk. + # Offline: force-stop at the remaining Talker context rather than + # waiting for a sampled EOS. + hit_chunk_limit = max_tokens is not None and step >= int(max_tokens) - 1 + chunk_done = empty_speech or hit_chunk_limit + if hit_chunk_limit: + # The next sampled id is forced EOS and will finish the request; + # mark the chunk done on this step so the data plane does not + # wait for a follow-up empty decode. + state["finished"] = True + force_eos_rows[index] = chunk_done + mask_eos_rows[index] = not force_eos_rows[index] and ( + (min_tokens is not None and step < int(min_tokens)) + or (bool(state.get("turn_end_drain")) and _turn_end_boundary_eos_masked(step)) + ) + terminal_flags[index] = torch.tensor(chunk_done, dtype=torch.bool) + + # Empty-speech rows, finished duplex chunks, and offline requests that + # fill the remaining Talker context must sample codec EOS so the + # scheduler releases the request. Mid-chunk rows mask EOS until + # min_tokens. + self._force_eos_rows = force_eos_rows + self._mask_eos_rows = mask_eos_rows + self._penalty_histories = penalty_histories + meta_outputs: dict[str, Any] = {"finished": terminal_flags} + if emit_duplex_metadata: + meta_outputs.update( + { + "native_duplex": native_duplex_flags, + "duplex_epoch": duplex_epochs, + "duplex_turn_id": duplex_turn_ids, + "llm_output_text_utf8": segment_texts_utf8, + "turn_end": turn_end_flags, + } + ) + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={ + "codes": {"audio": codec_deltas}, + "meta": meta_outputs, + }, + ) + def compute_logits(self, hidden_states, *args, **kwargs): if not isinstance(hidden_states, torch.Tensor): return None @@ -1611,7 +1763,31 @@ def compute_logits(self, hidden_states, *args, **kwargs): return hidden_states.new_empty((0, int(self._num_audio_tokens))) # Non-K-step: the real codec head. vLLM samples the codec stream here, # and the stage's stop_token_ids ends the request when it produces it. - return self.head_code[0](hidden_states).float() + logits = self.head_code[0](hidden_states).float() + force_eos = self._force_eos_rows + mask_eos = self._mask_eos_rows + self._force_eos_rows = None + self._mask_eos_rows = None + need_force = bool(force_eos and len(force_eos) == logits.shape[0] and any(force_eos)) + need_mask = bool(mask_eos and len(mask_eos) == logits.shape[0] and any(mask_eos)) + # sample() re-applies the decision on the sampled ids: vLLM's + # MinTokensLogitsProcessor runs after this and would blank the codec EOS + # we just forced (it is in the stage's ``stop_token_ids``), leaving an + # all -inf row and a request that never releases. + self._pending_force_eos_rows = force_eos if need_force else None + if need_force or need_mask: + logits = logits.clone() + eos_id = int(self._codec_eos_id) + if need_force: + assert force_eos is not None + forced = torch.tensor(force_eos, dtype=torch.bool, device=logits.device) + logits[forced] = float("-inf") + logits[forced, eos_id] = 0.0 + if need_mask: + assert mask_eos is not None + masked = torch.tensor(mask_eos, dtype=torch.bool, device=logits.device) + logits[masked, eos_id] = float("-inf") + return logits def sample(self, logits, sampling_metadata): # Two-column rows are the K-step stop/continue heads compute_logits @@ -1644,8 +1820,29 @@ def sample(self, logits, sampling_metadata): prewarped = _maybe_prewarp_top_k_top_p(logits, sampling_metadata) if prewarped is not None: logits, sampling_metadata = prewarped + force_eos = self._pending_force_eos_rows + self._pending_force_eos_rows = None with _prof_span("tts_sampler"): - return Sampler()(logits, sampling_metadata) + output = Sampler()(logits, sampling_metadata) + return self._force_eos_on_sampled_ids(output, force_eos) + + def _force_eos_on_sampled_ids(self, output: Any, force_eos: list[bool] | None) -> Any: + """Overwrite sampled ids for rows the model terminated this step. + + The codec EOS is a stage ``stop_token_ids`` entry, so vLLM's + ``min_tokens`` processor masks it for the first ``min_tokens`` steps. + A row the model forced to EOS therefore reaches the sampler as all + -inf and comes back as an arbitrary codec id, which keeps an + already-finished request decoding until its length cap. + """ + if not force_eos or not any(force_eos): + return output + sampled = getattr(output, "sampled_token_ids", None) + if not isinstance(sampled, torch.Tensor) or sampled.shape[0] != len(force_eos): + return output + rows = torch.tensor(force_eos, dtype=torch.bool, device=sampled.device) + sampled[rows] = int(self._codec_eos_id) + return output def _apply_codec_repetition_penalty(self, logits, sampling_metadata): """Score MiniCPMTTS.generate's windowed codec penalty, not vLLM's. From 001275f297ac769319507f3b0bb7965f057e4eaa Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 20:50:14 +0800 Subject: [PATCH 11/89] style: fix ruff lint and format findings across the branch - ruff check: 21 findings (15 auto-fixed; manual: N814 import alias, F841 unused variable, E501 long line, 4x UP031 percent formatting) - ruff format: 12 files reformatted Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 22 +++----- vllm_omni/config/omni_config.py | 4 +- vllm_omni/config/stage_config.py | 7 ++- vllm_omni/core/sched/omni_ar_scheduler.py | 21 +++----- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 50 ++++++++----------- .../minicpmo_4_5_omni.py | 4 +- .../platforms/npu/ascend_warmup_patch.py | 4 +- .../platforms/npu/native_rejection_sampler.py | 4 +- .../npu/worker/npu_ar_model_runner.py | 17 +++---- .../platforms/npu/worker/npu_model_runner.py | 4 +- .../platforms/npu/worker/talker_multiframe.py | 46 +++++++---------- vllm_omni/utils/step_prof.py | 1 + 12 files changed, 69 insertions(+), 115 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index e8c2c658360..6358f1167a9 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -40,9 +40,7 @@ def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): zeros while real codec rows are not -- a sharp contrast for asserting which id got embedded. """ - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__( - MiniCPMO45OmniTTSForConditionalGeneration - ) + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) nn.Module.__init__(model) model._k_step_frames = k_step_frames model.supports_multi_frame_decode = k_step_frames > 0 @@ -205,14 +203,14 @@ def test_guard_arms_from_vllm_config_and_drops_drafts_for_a_running_chunk(monkey # vLLM pads to 1+7 rows. Both spec_token_ids lists look innocent ([] and # 7 placeholders), so the guard must predict row widths from # num_tokens - num_computed_tokens instead. - from types import SimpleNamespace as NS + from types import SimpleNamespace from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") sched = OmniARScheduler.__new__(OmniARScheduler) sched._omni_talker_kstep_cache = None - sched.vllm_config = NS(speculative_config=NS(method="ngram", num_speculative_tokens=7)) + sched.vllm_config = SimpleNamespace(speculative_config=SimpleNamespace(method="ngram", num_speculative_tokens=7)) assert sched._talker_kstep_armed() is True sched.num_spec_tokens = 7 sched.max_model_len = 40960 @@ -398,16 +396,12 @@ def test_kstep_min_tokens_neutralization_clears_the_censor_list(): untouched = _MinTokensLogitsProcessor({0: (50, [0], {1})}) untouched.__class__ = type("SomeOtherProcessor", (object,), {}) # 非目标处理器 - talker_multiframe.neutralize_kstep_min_tokens( - SimpleNamespace(non_argmax_invariant=[untouched, censor]) - ) + talker_multiframe.neutralize_kstep_min_tokens(SimpleNamespace(non_argmax_invariant=[untouched, censor])) assert censor.min_toks == {} assert untouched.min_toks == {0: (50, [0], {1})} # 空名单 / 怪输入都不能炸。 - talker_multiframe.neutralize_kstep_min_tokens( - SimpleNamespace(non_argmax_invariant=[]) - ) + talker_multiframe.neutralize_kstep_min_tokens(SimpleNamespace(non_argmax_invariant=[])) talker_multiframe.neutralize_kstep_min_tokens(None) @@ -421,11 +415,7 @@ def test_kstep_bookkeeping_trace_is_total(monkeypatch): runner = SimpleNamespace( input_batch=SimpleNamespace( sampling_metadata=SimpleNamespace( - logitsprocs=SimpleNamespace( - non_argmax_invariant=[ - _MinTokensLogitsProcessor({0: (50, [0, 0], {1})}) - ] - ) + logitsprocs=SimpleNamespace(non_argmax_invariant=[_MinTokensLogitsProcessor({0: (50, [0, 0], {1})})]) ) ) ) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 73ef0280168..64c32242507 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -1221,9 +1221,7 @@ def _upstream_engine_field_map( # stage-1 override block with a startup ValueError ("Stage 1 (llm_ar) has # explicit engine argument(s) with no structured config owner"). It joins the # two pass-through fields already listed. -_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset( - {"compilation_config", "profiler_config", "speculative_config"} -) +_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config", "speculative_config"}) _LLM_LOAD_ENGINE_FIELDS = _LOAD_ENGINE_FIELDS | frozenset(_LOAD_CONFIG_ENGINE_FIELD_MAP.values()) _LLM_CACHE_ENGINE_FIELDS = _CACHE_ENGINE_FIELDS | frozenset(_CACHE_CONFIG_ENGINE_FIELD_MAP.values()) _LLM_SCHEDULER_ENGINE_FIELDS = _SCHEDULER_ENGINE_FIELDS | frozenset(_SCHEDULER_CONFIG_ENGINE_FIELD_MAP.values()) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 3a2c6f9b0ed..fe4d983f99c 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -821,8 +821,7 @@ def _npu_soc_allows_talker_multiframe() -> bool: name = _soc_name_from_env() if not name: logger.info( - "[minicpmo] SoC not identified from %s; leaving the Talker " - "multi-frame loop off. Set %s to force it on.", + "[minicpmo] SoC not identified from %s; leaving the Talker multi-frame loop off. Set %s to force it on.", "/".join(_SOC_ENV_VARS), _MINICPMO_TALKER_FRAMES_ENV, ) @@ -881,7 +880,7 @@ def talker_multiframe_armed() -> bool: return _npu_soc_allows_talker_multiframe() -def _apply_minicpmo_talker_multiframe_default(deploy: "DeployConfig") -> None: +def _apply_minicpmo_talker_multiframe_default(deploy: DeployConfig) -> None: """Let the Talker produce K codec frames per scheduler step on 910C/A3. A stage-1 decode step is dominated by per-*step* host work rather than per @@ -951,7 +950,7 @@ def _apply_minicpmo_talker_multiframe_default(deploy: "DeployConfig") -> None: ) -def _apply_minicpmo_perf_defaults(deploy: "DeployConfig", config_path: object = None) -> None: +def _apply_minicpmo_perf_defaults(deploy: DeployConfig, config_path: object = None) -> None: """Apply MiniCPM-o perf code-defaults, scoped to this deploy config on NPU.""" if config_path is None or "minicpmo" not in str(config_path).lower(): return diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 598baf63975..689304ca4eb 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -56,9 +56,7 @@ def kstep_account_trace_enabled() -> bool: return bool(raw) and raw not in ("0", "false", "no", "off") -def _trace_kstep_account( - req_id, spec_len, generated, computed, total, prompt_len, status, sampling=None -) -> None: +def _trace_kstep_account(req_id, spec_len, generated, computed, total, prompt_len, status, sampling=None) -> None: """One line per request per step; first 24 steps of each request, then 1 in 16.""" state = _KSTEP_ACCOUNT_TRACE if state["off"]: @@ -104,8 +102,6 @@ def _trace_kstep_account( logger.warning("[kstep-acct] trace disabled after failure: %r", exc) - - def _should_emit_engine_output( model_config: Any, *, @@ -513,9 +509,7 @@ def _drop_talker_drafts_if_prefill_pending(self) -> None: total = int(getattr(req, "num_tokens", prompt_len)) delta = total - computed spec_len = len(req.spec_token_ids or []) - states.append( - (str(getattr(req, "request_id", "?"))[:12], computed, prompt_len, total, spec_len, delta) - ) + states.append((str(getattr(req, "request_id", "?"))[:12], computed, prompt_len, total, spec_len, delta)) if computed < prompt_len: prefill_pending = True continue @@ -689,8 +683,11 @@ def update_from_output( logger.error( "SCHED0 zero-token schedule: req=%s n=%s spec_tokens=%s status=%s " "num_computed=%s num_output_placeholders=%s finished=%s all_entries=%s", - req_id, num_tokens_scheduled, _spec, - getattr(_r, "status", None), getattr(_r, "num_computed_tokens", None), + req_id, + num_tokens_scheduled, + _spec, + getattr(_r, "status", None), + getattr(_r, "num_computed_tokens", None), getattr(_r, "num_output_placeholders", None), None if _r is None else _r.is_finished(), dict(list(num_scheduled_tokens.items())[:8]), @@ -822,9 +819,7 @@ def update_from_output( _scheduled = num_tokens_scheduled _prev_computed = request.num_computed_tokens - _scheduled _prompt_len = ( - len(request.prompt_token_ids) - if getattr(request, "prompt_token_ids", None) is not None - else 0 + len(request.prompt_token_ids) if getattr(request, "prompt_token_ids", None) is not None else 0 ) _is_prefill_row = _prev_computed < _prompt_len _rollback = _drafts if _is_prefill_row else min(_scheduled, _drafts + 1) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 33440dedbeb..29e5c4ba248 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -12,9 +12,10 @@ 5. Next decode embeds that id with emb_code and emits it to Code2Wav """ -from collections.abc import Iterable, Mapping, Sequence import os +from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace + try: import torch_npu # noqa: F401 (NPU platform guarantee) except Exception: # pragma: no cover @@ -30,8 +31,8 @@ from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.llama import LlamaModel from vllm.model_executor.models.utils import maybe_prefix -from vllm.v1.sample.sampler import Sampler from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.sampler import Sampler from vllm_omni.engine.duplex.intermediate import get_tts_handoff from vllm_omni.model_executor.models.minicpmo_4_5 import ( @@ -39,7 +40,6 @@ MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS, ) from vllm_omni.model_executor.models.minicpmo_4_5.talker_codec_sample import ( - MIN_TOKENS_TO_KEEP, TalkerCodecDeviceState, TalkerCodecSampleResult, codec_sample_result, @@ -341,7 +341,6 @@ def resolve_codec_sampling_params( return resolved - # --------------------------------------------------------------------------- # Per-frame trace for the multi-frame loop (off by default; complements the # per-step timers in talker_multiframe). @@ -355,9 +354,9 @@ def resolve_codec_sampling_params( # engine down: on any error they disable themselves and log one warning. # --------------------------------------------------------------------------- _KSTEP_FRAME_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_FRAME_TRACE" -_KSTEP_FRAME_TRACE_LINES = 128 # max printed lines -_KSTEP_FRAME_TRACE_HEAD = 16 # first N frames printed one by one (frame-order check) -_KSTEP_FRAME_TRACE_STRIDE = 16 # then one line every N frames (covers ~1800 frames) +_KSTEP_FRAME_TRACE_LINES = 128 # max printed lines +_KSTEP_FRAME_TRACE_HEAD = 16 # first N frames printed one by one (frame-order check) +_KSTEP_FRAME_TRACE_STRIDE = 16 # then one line every N frames (covers ~1800 frames) _KSTEP_FRAME_TRACE: dict[str, Any] = {"reqs": {}, "lines": 0, "off": False} @@ -397,9 +396,7 @@ def _trace_kstep_frame(model: Any, **fields: Any) -> None: return try: device_states = getattr(model, "_request_codec_device_states", None) - device_state = ( - device_states.get(fields.get("request_id")) if isinstance(device_states, dict) else None - ) + device_state = device_states.get(fields.get("request_id")) if isinstance(device_states, dict) else None device_step = int(device_state.step.reshape(-1)[0].item()) if device_state is not None else -1 state["lines"] += 1 # cap printed lines only; per-request windows count separately logger.info( @@ -436,6 +433,7 @@ def _codec_int_param(state: Any, key: str, fallback: int) -> int: value = state.get(key) if isinstance(state, Mapping) else None return int(fallback if value is None else value) + class _MiniCPMTTSProjector(nn.Module): """Checkpoint-compatible hidden-state projector used by MiniCPMTTS.""" @@ -492,9 +490,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # block wins key by key, then the checkpoint's tts_config, then the # module fallbacks. The multi-frame path reads the resolved values # straight off the model. - yaml_codec = getattr( - getattr(vllm_config, "model_config", None), "codec_sampling_params", None - ) + yaml_codec = getattr(getattr(vllm_config, "model_config", None), "codec_sampling_params", None) resolved = resolve_codec_sampling_params(yaml_codec, tts_config) self._codec_seed = resolved["seed"] self._codec_temperature = resolved["temperature"] @@ -557,9 +553,7 @@ def _probe_soc_name() -> str: try: import torch_npu - return str( - torch_npu.npu.get_device_name(torch_npu.npu.current_device()) - ) + return str(torch_npu.npu.get_device_name(torch_npu.npu.current_device())) except Exception: return "" @@ -620,9 +614,7 @@ def _parse_k_step_frames(cls) -> int: try: frames = int(raw) except ValueError: - raise ValueError( - f"OMNI_K_STEP must be an integer frame count (>=2), got {raw!r}" - ) from None + raise ValueError(f"OMNI_K_STEP must be an integer frame count (>=2), got {raw!r}") from None explicit = "forced" else: frames = talker_frames_per_step() @@ -1305,17 +1297,14 @@ def merge_frame_outputs( per_frame_codes = [output.multimodal_outputs["codes"]["audio"] for output in frame_outputs] num_reqs = len(per_frame_codes[0]) - codec_deltas = [ - torch.cat([frame[index] for frame in per_frame_codes], dim=0) for index in range(num_reqs) - ] + codec_deltas = [torch.cat([frame[index] for frame in per_frame_codes], dim=0) for index in range(num_reqs)] meta_outputs: dict[str, Any] = {} for key in frame_outputs[0].multimodal_outputs["meta"]: per_frame = [output.multimodal_outputs["meta"][key] for output in frame_outputs] if key == "finished": meta_outputs[key] = [ - self._merge_frame_finished([frame[index] for frame in per_frame]) - for index in range(num_reqs) + self._merge_frame_finished([frame[index] for frame in per_frame]) for index in range(num_reqs) ] else: # Duplex metadata is per request and constant across the step's @@ -1349,7 +1338,6 @@ def _sample_audio_code( device_states = {} self._request_codec_device_states = device_states device_state = device_states.get(request_id) - new_segment = device_state is None if device_state is None: request_states = getattr(self, "_request_audio_states", {}) request_state = request_states.get(request_id, {}) @@ -1494,7 +1482,6 @@ def set_batch_stop_logits(self, logits: torch.Tensor | None) -> None: self._batch_stop_logits = logits @staticmethod - def _merge_frame_finished(flags: list[torch.Tensor]) -> torch.Tensor: """OR the per-frame terminal flags without forcing a device sync. @@ -1876,16 +1863,19 @@ def _apply_codec_repetition_penalty(self, logits, sampling_metadata): if ( isinstance(freq_rows, list) and len(freq_rows) == logits.shape[0] - and all(isinstance(f, torch.Tensor) and f.device == logits.device and f.numel() == logits.shape[-1] for f in freq_rows) + and all( + isinstance(f, torch.Tensor) and f.device == logits.device and f.numel() == logits.shape[-1] + for f in freq_rows + ) and bool((penalties.reshape(-1) != 1.0).any()) ): # Incremental histogram: freq rows are maintained on device by # make_omni_output (+1 append / -1 evict), identical to a full # scatter_add rebuild of the sliding window. freqs = torch.stack([f.to(dtype=torch.float32) for f in freq_rows], dim=0) - alpha = torch.pow( - penalties.to(device=logits.device, dtype=torch.float32).reshape(-1, 1), freqs - ).to(dtype=logits.dtype) + alpha = torch.pow(penalties.to(device=logits.device, dtype=torch.float32).reshape(-1, 1), freqs).to( + dtype=logits.dtype + ) logits = torch.where(logits < 0, logits * alpha, logits / alpha) return logits, replace(sampling_metadata, repetition_penalties=torch.ones_like(penalties)) logits = _apply_batched_repetition_penalty( diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index f91454df05b..9f4f6f3fc16 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -353,9 +353,7 @@ def tts2code2wav_async_chunk( # chunk_frames. When unset, first_window == chunk_frames and behavior is # unchanged. first_window = ( - _initial_chunk_frames(transfer_manager, chunk_frames) - if int(record["chunk_seq"]) == 0 - else chunk_frames + _initial_chunk_frames(transfer_manager, chunk_frames) if int(record["chunk_seq"]) == 0 else chunk_frames ) flush_pending = finished last_chunk = bool(flush_pending and (not native_duplex or turn_end)) diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py index 1c681abbb91..c69f1665141 100644 --- a/vllm_omni/platforms/npu/ascend_warmup_patch.py +++ b/vllm_omni/platforms/npu/ascend_warmup_patch.py @@ -152,9 +152,7 @@ def apply_ascend_warmup_patch() -> None: # re-exports ``kernel_warmup`` as a package attribute, so # ``from ... import kernel_warmup`` yields the function and every # ``getattr`` below would silently miss. - warmup_module = importlib.import_module( - "vllm_ascend.model_executor.warmup.kernel_warmup" - ) + warmup_module = importlib.import_module("vllm_ascend.model_executor.warmup.kernel_warmup") except Exception as error: # pragma: no cover - non-ascend builds logger.warning("[npu] Ascend warmup patch not applied: %s", error) return diff --git a/vllm_omni/platforms/npu/native_rejection_sampler.py b/vllm_omni/platforms/npu/native_rejection_sampler.py index b0ea1bf1ed6..b5b881f281a 100644 --- a/vllm_omni/platforms/npu/native_rejection_sampler.py +++ b/vllm_omni/platforms/npu/native_rejection_sampler.py @@ -58,9 +58,7 @@ def restore_native_rejection_sampler() -> None: import vllm.v1.sample.rejection_sampler as rs except ImportError: # pragma: no cover - vllm core moved the module return - spec = importlib.util.spec_from_file_location( - "_vllm_rejection_sampler_native", rs.__file__ - ) + spec = importlib.util.spec_from_file_location("_vllm_rejection_sampler_native", rs.__file__) if spec is None or spec.loader is None: # pragma: no cover return native = importlib.util.module_from_spec(spec) diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index d4bbb05efdf..ea9d37411ae 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -45,13 +45,13 @@ from vllm_ascend.worker.model_runner_v1 import graph_capture from vllm_omni.data_entry_keys import flatten_payload -from vllm_omni.utils.step_prof import span, tic, toc from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output from vllm_omni.model_executor.duplex_sampling import DuplexSamplingRunnerMixin from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element +from vllm_omni.utils.step_prof import span, tic, toc from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, needs_omni_connector, @@ -223,19 +223,16 @@ def propose_draft_token_ids(self, valid_sampled_token_ids, *args, **kwargs): # tensor, and `tensor or []` raises the ambiguous-truth error. if isinstance(valid_sampled_token_ids, list): rows = valid_sampled_token_ids - _desc = "rows=%s types=%s lens=%s head=%s" % ( - len(rows), - [type(r).__name__ for r in rows[:8]], - [len(r) for r in rows[:8] if isinstance(r, list)], - repr(rows)[:400], + _desc = ( + f"rows={len(rows)}" + f" types={[type(r).__name__ for r in rows[:8]]}" + f" lens={[len(r) for r in rows[:8] if isinstance(r, list)]}" + f" head={repr(rows)[:400]}" ) else: seen = valid_sampled_token_ids shape = tuple(seen.shape) if hasattr(seen, "shape") else "?" - _desc = "rows=<%s shape=%s>" % ( - type(seen).__name__, - shape, - ) + _desc = f"rows=<{type(seen).__name__} shape={shape}>" logger.error( "[kstep] no drafts emitted: frames=%s num_reqs=%s %s", frames, diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 3ac145a5913..84a0c934d3d 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -516,9 +516,7 @@ def _model_forward( model_output = talker_multiframe.run( model=self.model, run_model=run_model, - after_forward=lambda: self._update_full_graph_params_if_needed( - forward_context, graph_tokens - ), + after_forward=lambda: self._update_full_graph_params_if_needed(forward_context, graph_tokens), inputs_embeds=inputs_embeds, input_ids=input_ids, frames=frames, diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py index 55a666334d2..c7a8fe8e3d1 100644 --- a/vllm_omni/platforms/npu/worker/talker_multiframe.py +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -56,9 +56,10 @@ import dataclasses import os +from collections.abc import Callable from dataclasses import dataclass from time import perf_counter -from typing import Any, Callable +from typing import Any import torch from vllm.logger import init_logger @@ -163,8 +164,7 @@ def _prof_summary(steps: int, frames_total: int, last_frames: int, *, detail: bo replay_per_frame = _PROF_REPLAY_MS / frames_total if frames_total else 0.0 after_per_frame = _PROF_AFTER_MS / frames_total if frames_total else 0.0 synced = ( - f" synced_frame_ms={_PROF_SYNCED_MS / _PROF_SYNCED_FRAMES:.3f}" - f" (n={_PROF_SYNCED_FRAMES})" + f" synced_frame_ms={_PROF_SYNCED_MS / _PROF_SYNCED_FRAMES:.3f} (n={_PROF_SYNCED_FRAMES})" if _PROF_SYNCED_FRAMES else "" ) @@ -210,7 +210,7 @@ def _prof_disabled(reason: str) -> None: os.environ[_PROF_ENV] = "0" -def _prof_step_begin(frames: int) -> "_StepProf | None": +def _prof_step_begin(frames: int) -> _StepProf | None: """``None`` when off -- and then ``run()`` takes no timer at all.""" try: if not prof_enabled(): @@ -221,7 +221,7 @@ def _prof_step_begin(frames: int) -> "_StepProf | None": return None -def _prof_frame(prof: "_StepProf", frame: int, t0: float, t1: float, t2: float, t3: float) -> None: +def _prof_frame(prof: _StepProf, frame: int, t0: float, t1: float, t2: float, t3: float) -> None: """Book one replay. ``replay``/``after`` are enqueue times; ``synced`` is the frame's real wall clock (drained after the replay), and the gap between them is the device side of the account.""" @@ -242,7 +242,7 @@ def _prof_frame(prof: "_StepProf", frame: int, t0: float, t1: float, t2: float, _prof_disabled("frame") -def _prof_step_end(prof: "_StepProf | None", step_start: float) -> None: +def _prof_step_end(prof: _StepProf | None, step_start: float) -> None: if prof is None: return try: @@ -256,6 +256,8 @@ def _prof_step_end(prof: "_StepProf | None", step_start: float) -> None: _prof_summary(_PROF_STEPS, _PROF_FRAMES, prof.frames, detail=prof.detail) except Exception: _prof_disabled("step_end") + + # --------------------------------------------------------------------------- # Stop-row forensics for the multi-frame loop (off by default). # @@ -328,7 +330,7 @@ def trace_kstep_bookkeeping(runner: Any, valid_sampled_token_ids: Any, logits: A if isinstance(valid_sampled_token_ids, list): carried: Any = [list(row) for row in valid_sampled_token_ids[:4]] else: - carried = "<%s>" % type(valid_sampled_token_ids).__name__ + carried = f"<{type(valid_sampled_token_ids).__name__}>" masks: list[Any] = [] sampling_metadata = getattr(getattr(runner, "input_batch", None), "sampling_metadata", None) @@ -603,7 +605,7 @@ def begin_narrow_step( positions: Any, inputs_embeds: Any, spans: list[tuple[int, int]], -) -> "NarrowStep | None": +) -> NarrowStep | None: """Point the forward context at the one-query graph, or decline. Declines -- and the caller falls back to the wide replay -- whenever the @@ -631,7 +633,7 @@ def begin_narrow_step( return None seq_lens = fixed_kv_decode.captured_seq_lens(rows) if seq_lens is None: - return _decline("no one-query graph captured for %d rows" % rows) + return _decline(f"no one-query graph captured for {rows} rows") metadata = getattr(forward_context, "attn_metadata", None) if isinstance(metadata, dict): metadata = next(iter(metadata.values()), None) @@ -669,13 +671,9 @@ def begin_narrow_step( # back, and the captured key has it set. What this step is has not # changed -- it is one token per request against a graph captured for # exactly that -- only which graph replays it. - forward_context.batch_descriptor = dataclasses.replace( - descriptor, num_tokens=rows, num_reqs=rows, uniform=True - ) + forward_context.batch_descriptor = dataclasses.replace(descriptor, num_tokens=rows, num_reqs=rows, uniform=True) else: - forward_context.batch_descriptor = BatchDescriptor( - num_tokens=rows, num_reqs=rows, uniform=True - ) + forward_context.batch_descriptor = BatchDescriptor(num_tokens=rows, num_reqs=rows, uniform=True) # The dispatcher builds its lookup descriptor from its own query length -- # 4 at runtime, which is right for the step -- so it finds no key for the # narrow capture and hands back NONE, and the wrapper then runs the model @@ -696,7 +694,7 @@ def begin_narrow_step( return narrow -def end_narrow_step(forward_context: Any, narrow: "NarrowStep | None") -> None: +def end_narrow_step(forward_context: Any, narrow: NarrowStep | None) -> None: if narrow is None: return forward_context.batch_descriptor = narrow.saved_descriptor @@ -727,7 +725,7 @@ def run( frames: int, model_kwargs: dict[str, Any], model_kwargs_extra: dict[str, Any], - narrow: "NarrowStep | None" = None, + narrow: NarrowStep | None = None, ) -> Any: """Run ``frames`` codec frames and return the merged step output. @@ -756,9 +754,7 @@ def run( if frame == 0: _compact_frame_zero_embeddings(inputs_embeds, spans) else: - _write_frame_embeddings( - model, inputs_embeds, input_ids, spans, infos, frame, narrow=narrow - ) + _write_frame_embeddings(model, inputs_embeds, input_ids, spans, infos, frame, narrow=narrow) elif frame > 0: _write_frame_embeddings(model, inputs_embeds, input_ids, spans, infos, frame) t0 = perf_counter() if step_prof is not None else 0.0 @@ -791,9 +787,7 @@ def run( if not _LOGGED_ENGAGE: _LOGGED_ENGAGE = True - logger.info( - "[minicpmo] multi-frame Talker decode engaged: %d codec frames per step", frames - ) + logger.info("[minicpmo] multi-frame Talker decode engaged: %d codec frames per step", frames) if stop_trace_enabled(): _trace_stop_rows(frame_stop_logits) merged = model.merge_frame_outputs(frame_outputs, frame_stop_logits) @@ -801,9 +795,7 @@ def run( # (frames, rows, hidden) -> (rows * frames, hidden), request-major -- # the layout `logits_indices` reads. OmniOutput is a NamedTuple. merged = merged._replace( - text_hidden_states=torch.stack(frame_hidden, dim=1).reshape( - -1, frame_hidden[0].shape[-1] - ) + text_hidden_states=torch.stack(frame_hidden, dim=1).reshape(-1, frame_hidden[0].shape[-1]) ) # Whole step, frames + merge: this is the number to compare against the # scheduler's step cadence. @@ -827,7 +819,7 @@ def _write_frame_embeddings( spans: list[tuple[int, int]], infos: list[dict[str, Any]], frame: int, - narrow: "NarrowStep | None" = None, + narrow: NarrowStep | None = None, ) -> None: """Put frame ``frame``'s decode embedding into the row the graph reads. diff --git a/vllm_omni/utils/step_prof.py b/vllm_omni/utils/step_prof.py index fc52d08a1fd..7f27e5868fb 100644 --- a/vllm_omni/utils/step_prof.py +++ b/vllm_omni/utils/step_prof.py @@ -6,6 +6,7 @@ default 200). Thread-safe; spans are keyed per thread so nested/overlapping regions in the engine loop stay correct. """ + from __future__ import annotations import os From 5e6973fd848546bee5c24d3759d41782ea38d2d6 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 22:01:00 +0800 Subject: [PATCH 12/89] fix(deploy): align minicpmo_4_5.yaml with the validated optimization set codec_chunk_frames 25 -> 75 so initial_codec_chunk_frames: 25 takes effect; stage-1 stop_token_ids: [1] for the multi-frame stop head (nothing injects it for stage 1); stage0/stage2 max_num_seqs 8 so the top tier is not gated by a single stage; NPU cudagraph_mode PIECEWISE -> FULL_DECODE_ONLY, which keeps in-graph sampling and multi-frame expansion. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 2177cd90b68..b4cd3f6b7ce 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -19,7 +19,8 @@ connectors: connector_get_sleep_s: 0.01 connector_get_max_wait_first_chunk: 3000 connector_get_max_wait: 300 - codec_chunk_frames: 25 + # Steady-state payload length in codec frames. + codec_chunk_frames: 75 # Smaller first payload so the first audio packet is emitted earlier. # Takes effect only when initial < codec_chunk_frames; values far below # the steady-state chunk length add a seam between chunks and hurt WER. @@ -95,6 +96,8 @@ stages: top_k: 25 repetition_penalty: 1.05 min_tokens: 50 + # The Talker's stop head emits this id to finish a request. + stop_token_ids: [1] max_tokens: 4096 seed: 42 detokenize: false @@ -129,17 +132,24 @@ platforms: npu: stages: - stage_id: 0 + # Keep the three stages' concurrency caps aligned so the top tier is + # not gated by a single stage. + max_num_seqs: 8 max_num_batched_tokens: 8192 compilation_config: - cudagraph_mode: PIECEWISE + # A full decode-step graph keeps in-graph sampling and multi-frame + # expansion; PIECEWISE silently drops both. Measured on 910B3: + # -32.6% RTF versus PIECEWISE. + cudagraph_mode: FULL_DECODE_ONLY - stage_id: 1 - # Stage-1 concurrency cap raised from 4 to 8. The NPU override block - # only covered stages 0 and 2, so stage 1 kept the base value of 4. + # Concurrency cap raised from 4 to 8; the override block only covered + # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 compilation_config: - cudagraph_mode: PIECEWISE + cudagraph_mode: FULL_DECODE_ONLY - stage_id: 2 + max_num_seqs: 8 # Stage 2 stays outer-eager and captures only exact-shape CFM DiT # estimator kernels through the NPU platform patch. additional_config: From 4b6c69132b583139dcae05d213d791bb7e167e04 Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 22:48:33 +0800 Subject: [PATCH 13/89] refactor(minicpmo_4_5): configure the Talker multi-frame decode from the deploy YAML Make stage 1's speculative_config in the deploy YAML the single source of truth for the multi-frame decode and remove every code-side default and switch around it. * the deploy YAML carries stage 1's speculative_config (8 codec frames per step) and async_scheduling: false explicitly; no code-side injection or default remains. * the environment switches VLLM_OMNI_MINICPMO_TALKER_FRAMES, OMNI_K_STEP and VLLM_OMNI_C2W_SINGLE_ROW are gone, and the 910B SoC gate is dropped: the rejection-verifier limit it guarded against was the missing multi-frame wiring this branch fixes, and 910B3 is validated end to end. * the model-side and scheduler-side gates derive K from the engine config (model_stage == tts plus the n-gram fingerprint), so the two cannot disagree. * the stage-2 single-row even-batch item is reverted: its measured gain sat inside the run-to-run drift band and a default-off switch adds configuration without value. Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 72 +++---- vllm_omni/config/stage_config.py | 187 +----------------- vllm_omni/core/sched/omni_ar_scheduler.py | 37 ++-- vllm_omni/deploy/minicpmo_4_5.yaml | 13 ++ .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 8 - .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 109 +++------- .../models/minicpmo_4_5/pipeline.py | 19 +- .../platforms/npu/ascend_warmup_patch.py | 18 +- .../platforms/npu/worker/talker_multiframe.py | 6 +- 9 files changed, 118 insertions(+), 351 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 6358f1167a9..92d2f2b548d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -28,7 +28,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -_TALKER_FRAMES_ENV = "VLLM_OMNI_MINICPMO_TALKER_FRAMES" _NUM_AUDIO_TOKENS = 6562 _EOS_ID = _NUM_AUDIO_TOKENS - 1 @@ -128,12 +127,17 @@ def test_decode_preprocess_embeds_state_last_code(): assert out["codes"]["audio"].reshape(-1).tolist() == [42] -def _make_scheduler(*, num_spec: int, waiting=(), running=()): +def _make_scheduler(*, num_spec: int, waiting=(), running=(), stage: str = "tts"): from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler sched = OmniARScheduler.__new__(OmniARScheduler) sched._omni_talker_kstep_cache = None - sched.speculative_config = SimpleNamespace(method="ngram", num_speculative_tokens=num_spec) + # The armed check reads the stage and the n-gram fingerprint off the engine + # config; the real Scheduler exposes neither as a plain attribute. + sched.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(model_stage=stage), + speculative_config=SimpleNamespace(method="ngram", num_speculative_tokens=num_spec), + ) # vLLM's Scheduler stores the count here (vllm_config.num_speculative_tokens); # the guard predicts row widths from it. sched.num_spec_tokens = num_spec @@ -154,8 +158,7 @@ def _req(*, computed: int, prompt: int, spec: list[int], total: int | None = Non ) -def test_guard_drops_drafts_when_waiting_request_pending(monkeypatch): - monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") +def test_guard_drops_drafts_when_waiting_request_pending(): decode_req = _req(computed=100, prompt=100, spec=[0] * 7) sched = _make_scheduler(num_spec=7, waiting=[object()], running=[decode_req]) @@ -163,8 +166,7 @@ def test_guard_drops_drafts_when_waiting_request_pending(monkeypatch): assert decode_req.spec_token_ids == [] -def test_guard_keeps_drafts_when_no_prefill_pending(monkeypatch): - monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") +def test_guard_keeps_drafts_when_no_prefill_pending(): decode_req = _req(computed=100, prompt=100, spec=[0] * 7, total=101) sched = _make_scheduler(num_spec=7, waiting=[], running=[decode_req]) @@ -172,8 +174,7 @@ def test_guard_keeps_drafts_when_no_prefill_pending(monkeypatch): assert decode_req.spec_token_ids == [0] * 7 -def test_guard_drops_drafts_when_chunked_prefill_in_flight(monkeypatch): - monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") +def test_guard_drops_drafts_when_chunked_prefill_in_flight(): decoding_req = _req(computed=100, prompt=100, spec=[0] * 7, total=101) chunking_req = _req(computed=50, prompt=100, spec=[]) sched = _make_scheduler(num_spec=7, waiting=[], running=[decoding_req, chunking_req]) @@ -182,35 +183,35 @@ def test_guard_drops_drafts_when_chunked_prefill_in_flight(monkeypatch): assert decoding_req.spec_token_ids == [] -def test_guard_noop_for_text_stage_spec_config(monkeypatch): - monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") +def test_guard_noop_for_text_stage_spec_config(): + # The text stage may carry its own n-gram config; only the Talker stage runs + # the multi-frame loop, so the guard must stay a no-op there. text_req = _req(computed=100, prompt=100, spec=[0] * 15) - sched = _make_scheduler(num_spec=15, waiting=[object()], running=[text_req]) + sched = _make_scheduler(num_spec=15, stage="llm", waiting=[object()], running=[text_req]) sched._drop_talker_drafts_if_prefill_pending() assert text_req.spec_token_ids == [0] * 15 -def test_guard_arms_from_vllm_config_and_drops_drafts_for_a_running_chunk(monkeypatch): +def test_guard_arms_from_vllm_config_and_drops_drafts_for_a_running_chunk(): # Real vLLM Scheduler instances expose no .speculative_config attribute # (and SchedulerConfig has no num_speculative_tokens), so the armed check # must read vllm_config.speculative_config -- otherwise the whole guard - # silently stays off while the engine-side loop is armed (the 11:47 and - # 12:04 crashes on 910B). + # silently stays off while the engine-side loop is armed. # - # 2026-09-18 12:39 crash: spans=[(0,7),(7,15)] -- a request carrying a - # 7-token streaming chunk shares the step with a steady-state decode that - # vLLM pads to 1+7 rows. Both spec_token_ids lists look innocent ([] and - # 7 placeholders), so the guard must predict row widths from - # num_tokens - num_computed_tokens instead. + # Steady-state decode sharing a step with a 7-token streaming chunk: both + # spec_token_ids lists look innocent ([] and 7 placeholders), so the guard + # must predict row widths from num_tokens - num_computed_tokens instead. from types import SimpleNamespace from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler - monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") sched = OmniARScheduler.__new__(OmniARScheduler) sched._omni_talker_kstep_cache = None - sched.vllm_config = SimpleNamespace(speculative_config=SimpleNamespace(method="ngram", num_speculative_tokens=7)) + sched.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(model_stage="tts"), + speculative_config=SimpleNamespace(method="ngram", num_speculative_tokens=7), + ) assert sched._talker_kstep_armed() is True sched.num_spec_tokens = 7 sched.max_model_len = 40960 @@ -225,11 +226,10 @@ def test_guard_arms_from_vllm_config_and_drops_drafts_for_a_running_chunk(monkey assert chunk.spec_token_ids == [] -def test_guard_keeps_drafts_when_all_decodes_share_the_padded_width(monkeypatch): +def test_guard_keeps_drafts_when_all_decodes_share_the_padded_width(): # A first-step decode (no placeholders yet) and a steady-state decode both # schedule 1 token this step, so vLLM gives both the same 1+num_spec row # width: spans stay uniform and nothing may drop. - monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") steady = _req(computed=100, prompt=100, spec=[0] * 7, total=101) first_step = _req(computed=100, prompt=100, spec=[], total=101) sched = _make_scheduler(num_spec=7, waiting=[], running=[steady, first_step]) @@ -239,31 +239,23 @@ def test_guard_keeps_drafts_when_all_decodes_share_the_padded_width(monkeypatch) assert first_step.spec_token_ids == [] -def test_talker_stop_token_ids_follow_the_arming_decision(monkeypatch): +def test_talker_stop_token_ids_match_the_multi_frame_head(): """Stage 1's stop id must be one the head that actually runs can emit. - 2026-09-18 root cause: the pipeline constraint pinned ``stop_token_ids`` to - the codec EOS (6561) unconditionally. With the K-frame loop armed the - vLLM-level head is the two-wide continue/stop row, so the only sampleable - ids are 0/1, `check_stop` never matched, and every K-step request ran to - ``max_tokens`` (142s per request on 910B, model finished at frame ~116). + The pipeline constraint used to pin ``stop_token_ids`` to the codec EOS + (6561) unconditionally. With the multi-frame loop on, the vLLM-level head + is the two-wide continue/stop row, so the only sampleable ids are 0/1, + `check_stop` never matched, and every request ran to ``max_tokens`` (142s + per request, model finished at frame ~116). The pipeline default therefore + follows the multi-frame head; a deployment that turns the loop off + overrides stage 1's stop token per stage. """ from vllm_omni.model_executor.models.minicpmo_4_5 import pipeline as mcp_pipeline from vllm_omni.platforms.npu.worker import talker_multiframe - monkeypatch.setenv(_TALKER_FRAMES_ENV, "8") assert mcp_pipeline._talker_stop_token_ids() == [talker_multiframe.STOP_TOKEN_ID] assert talker_multiframe.STOP_TOKEN_ID == 1 - monkeypatch.setenv(_TALKER_FRAMES_ENV, "1") - assert mcp_pipeline._talker_stop_token_ids() == [6561] - - # No explicit frame count on a 910B box: the loop stays off, so the real - # codec head runs and its EOS is the stop. - monkeypatch.delenv(_TALKER_FRAMES_ENV, raising=False) - monkeypatch.setenv("SOC_VERSION", "ascend910b1") - assert mcp_pipeline._talker_stop_token_ids() == [6561] - def test_multiframe_gate_matrix(): from vllm_omni.platforms.npu.worker import talker_multiframe diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index fe4d983f99c..676e3c889c5 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -780,189 +780,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: return merged -_MINICPMO_TALKER_FRAMES_ENV = "VLLM_OMNI_MINICPMO_TALKER_FRAMES" -# K = codec frames one stage-1 execute_model produces. 8 is the conservative -# starting point on 910C/A3; K=16 measures a further -6.9% at a higher TTFP -# cost, so the ceiling stays reachable through the environment. -_MINICPMO_TALKER_FRAMES = 8 -# Same TND ceiling as stage 0: at most 16 query positions per sequence. +# Same TND ceiling as stage 0: at most 16 query positions per sequence. The +# Talker's multi-frame decode is configured through stage 1's +# speculative_config in the deploy YAML; this only bounds the frame count the +# model-side gate accepts. _MINICPMO_TALKER_FRAMES_MAX = 16 -_MINICPMO_FRAMES_OFF = ("0", "1", "off", "false", "no") - - -_SOC_ENV_VARS = ("SOC_VERSION", "ASCEND_SOC_VERSION") - - -def _soc_name_from_env() -> str: - """SoC identifier from the launch environment. Never touches the device. - - Asking torch_npu for the device name here would initialise the NPU in - whichever process loads the deploy config -- the API server -- and that - process would then hold device memory it has no use for, which shows up as - an extra entry in npu-smi. CANN launch environments already export the - revision: ``ascend910b1`` on a 910B box, ``ascend910_9391`` on 910C. - """ - for variable in _SOC_ENV_VARS: - value = os.environ.get(variable, "").strip().lower() - if value: - return value - return "" - - -def _npu_soc_allows_talker_multiframe() -> bool: - """False on the 910B family, whose rejection verifier cannot take K rows. - - Judged from the environment only, so loading a deploy config never - initialises a device. An unidentifiable SoC is treated as "not allowed": - the cost of leaving the frame loop off is a slower run, while speccing a - 910B faults its rejection kernel. Setting ``VLLM_OMNI_MINICPMO_TALKER_FRAMES`` - explicitly overrides this (see the caller). - """ - name = _soc_name_from_env() - if not name: - logger.info( - "[minicpmo] SoC not identified from %s; leaving the Talker multi-frame loop off. Set %s to force it on.", - "/".join(_SOC_ENV_VARS), - _MINICPMO_TALKER_FRAMES_ENV, - ) - return False - return not name.startswith("ascend910b") - - -def talker_frames_per_step() -> int: - """Codec frames one stage-1 `execute_model` produces. 1 disables it. - - Deliberately read from the environment on every call rather than cached: - the stage config and the model runner agree through this one function and - are built at different times in different processes. - """ - raw = os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip().lower() - if not raw: - return _MINICPMO_TALKER_FRAMES - if raw in _MINICPMO_FRAMES_OFF: - return 1 - try: - frames = int(raw) - except ValueError: - logger.warning( - "[minicpmo] %s=%r is not an integer; using the code default %d", - _MINICPMO_TALKER_FRAMES_ENV, - raw, - _MINICPMO_TALKER_FRAMES, - ) - return _MINICPMO_TALKER_FRAMES - if frames < 1: - return 1 - return min(frames, _MINICPMO_TALKER_FRAMES_MAX) - - -def talker_multiframe_armed() -> bool: - """Whether the Talker K-frame loop will run in this process. - - One source of truth for the three places that must agree: the deploy-config - injection below (does stage 1 get the speculative_config), the model's own - gate, and stage 1's stop-token constraint. The last one is not obvious: when - the loop is armed the vLLM-level head collapses to the two-wide - continue/stop row (``compute_logits``), so the only ids that can be sampled - are 0 and 1 and the codec EOS becomes unreachable *at that level*. A - ``stop_token_ids`` the head can never emit means `check_stop` never fires - and the request runs to ``max_tokens`` -- the loop keeps paying for a - request that already finished at the codec level. - - An explicit ``VLLM_OMNI_MINICPMO_TALKER_FRAMES`` is an operator decision and - arms regardless of the SoC, exactly as the loader below does. Everything - else is the SoC gate's answer, and an unidentified SoC counts as "no". - """ - if talker_frames_per_step() <= 1: - return False - if os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip(): - return True - return _npu_soc_allows_talker_multiframe() - - -def _apply_minicpmo_talker_multiframe_default(deploy: DeployConfig) -> None: - """Let the Talker produce K codec frames per scheduler step on 910C/A3. - - A stage-1 decode step is dominated by per-*step* host work rather than per - frame, so running K frames inside one `execute_model` pays that host cost - once (measured on 910C/A3: ~2.9 ms of host around a 0.83 ms device forward, - over ~118 frames per request). The frames are *sequential* -- frame k+1 is - conditioned on the codec token sampled at frame k -- so this is not - speculative decoding and the runner runs its own inner loop. vLLM still - has to know the request advanced by K tokens, because that is what grows - the block table and KV slots the loop writes into, and `spec_token_ids` is - V1's only mechanism for that; hence a speculative_config that carries the - count. The n-gram proposer is never consulted: `propose_draft_token_ids` - returns the Talker's `continue` token K-1 times, so verification is exact - by construction. - - `async_scheduling` comes off with it. vLLM refuses async scheduling next - to an n-gram proposer ("only supported with EAGLE/MTP/Draft Model/NGram - GPU/DSpark"), which is the same reason stage 0 carries the flag, and a - synchronous scheduler costs less here than the K-fold host saving returns. - - Left off entirely on the 910B family: its rejection verifier trips a - vector-core limit on spec-width rows, so the baseline there stays exactly - as the deploy config describes it. - """ - frames = talker_frames_per_step() - if frames <= 1: - return - explicit = os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip() != "" - if not talker_multiframe_armed(): - logger.info( - "[minicpmo] Talker multi-frame decode left off: the SoC either " - "cannot verify spec-width rows (910B family limit) or was not " - "identified. Deploy on 910C/A3, or set %s explicitly to force it.", - _MINICPMO_TALKER_FRAMES_ENV, - ) - return - if explicit: - # An explicit request is an operator decision: arm it and let the - # model-side gate make the final call on the live device. - logger.info( - "[minicpmo] %s is set explicitly; arming K=%d without the SoC check", - _MINICPMO_TALKER_FRAMES_ENV, - frames, - ) - for stage in deploy.stages: - if stage.stage_id != 1: - continue - if stage.engine_extras.get("speculative_config") is not None: - # An explicit config wins; this default only fills a gap. - logger.info( - "[minicpmo] stage 1 carries an explicit speculative_config; " - "leaving it and its async_scheduling flag untouched", - ) - return - stage.engine_extras["speculative_config"] = { - "method": "ngram", - "num_speculative_tokens": frames - 1, - "prompt_lookup_min": 1, - "prompt_lookup_max": 1, - } - stage.async_scheduling = False - logger.info( - "[minicpmo] stage 1 multi-frame decode on (%d codec frames per " - "step, synchronous scheduler); set %s=1 to override", - frames, - _MINICPMO_TALKER_FRAMES_ENV, - ) - - -def _apply_minicpmo_perf_defaults(deploy: DeployConfig, config_path: object = None) -> None: - """Apply MiniCPM-o perf code-defaults, scoped to this deploy config on NPU.""" - if config_path is None or "minicpmo" not in str(config_path).lower(): - return - try: - from vllm_omni.platforms import current_omni_platform - - device_name = getattr(current_omni_platform, "device_name", None) - except Exception: - return - if (device_name or "").lower() != "npu": - return - _apply_minicpmo_talker_multiframe_default(deploy) def load_deploy_config(path: str | Path) -> DeployConfig: @@ -1008,7 +830,6 @@ def load_deploy_config(path: str | Path) -> DeployConfig: if name in raw_dict: kwargs[name] = raw_dict[name] deploy = DeployConfig(**kwargs) - _apply_minicpmo_perf_defaults(deploy, path) return deploy diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 689304ca4eb..7e8e219dd0a 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -406,18 +406,14 @@ def _commit_kv_transfer_trigger( if stop_after_transfer and req_id in self.requests_needing_kv_transfer: self.pending_stop_after_extraction.add(req_id) - # Env var shared with stage_config._apply_minicpmo_talker_multiframe_default; - # keep the name in sync there. - _TALKER_FRAMES_ENV = "VLLM_OMNI_MINICPMO_TALKER_FRAMES" - def _talker_kstep_armed(self) -> bool: - """True when this scheduler drives the Talker K-frame decode. + """True when this scheduler drives the Talker multi-frame decode. - Mirrors `_apply_minicpmo_talker_multiframe_default`: the Talker stage - gets an injected n-gram config with exactly frames-1 draft tokens. - Matching that fingerprint keeps both sides reading the same env var. - The text stage's explicit n-gram config (15 draft tokens by default) - does not match, so this stays a no-op there. + Both conditions come off the engine config: the stage must be the + Talker (``model_config.model_stage == "tts"``) and it must carry the + injected n-gram config whose ``num_speculative_tokens`` is exactly + ``K - 1``. The same stage-1 config feeds the model-side gate, so the + scheduler and the runner cannot disagree about K. """ cache = getattr(self, "_omni_talker_kstep_cache", None) if cache is None: @@ -425,23 +421,20 @@ def _talker_kstep_armed(self) -> bool: # attribute and SchedulerConfig has no num_speculative_tokens # either, so the canonical engine-side source is vllm_config. # Reading anything else silently yields 0 drafts and this whole - # guard never arms -- exactly the bug that survived the 11:47 - # crash fix. + # guard never arms. + vllm_cfg = getattr(self, "vllm_config", None) spec = getattr(self, "speculative_config", None) if spec is None: - vllm_cfg = getattr(self, "vllm_config", None) spec = getattr(vllm_cfg, "speculative_config", None) if vllm_cfg is not None else None - num_spec = 0 - is_ngram = True - if spec is not None: + if isinstance(spec, dict): + num_spec = spec.get("num_speculative_tokens", 0) or 0 + is_ngram = spec.get("method", "ngram") == "ngram" + else: num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 is_ngram = getattr(spec, "method", "ngram") == "ngram" - try: - frames = int(os.environ.get(self._TALKER_FRAMES_ENV, "8") or 8) - except ValueError: - frames = 8 - armed = is_ngram and frames > 1 and num_spec == frames - 1 - cache = self._omni_talker_kstep_cache = armed + model_cfg = getattr(vllm_cfg, "model_config", None) + is_talker = getattr(model_cfg, "model_stage", None) == "tts" + cache = self._omni_talker_kstep_cache = is_talker and is_ngram and num_spec > 0 return cache def _log_kstep_guard_view(self, verdict: str, widths: set[int], states: list, dropped: int = 0) -> None: diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index b4cd3f6b7ce..7eb0b14c184 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -81,6 +81,19 @@ stages: # Talker context is 4096 (tts_config.max_position_embeddings). max_model_len: 4096 devices: "0" + # Multi-frame decode: the Talker produces 8 codec frames per scheduler step. + # vLLM learns that each request advances 8 tokens through this + # speculative_config (the n-gram proposer is never consulted -- the runner + # supplies constant continue drafts), and async scheduling is rejected next + # to an n-gram proposer. stop_token_ids above must stay the two-wide + # continue/stop id the head emits in this mode. Remove the block (or set it + # to null) to fall back to one frame per step. + speculative_config: + method: ngram + num_speculative_tokens: 7 + prompt_lookup_min: 1 + prompt_lookup_max: 1 + async_scheduling: false output_connectors: to_stage_2: connector_of_shared_memory # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 34fd101412b..a7993dd3b4f 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -285,7 +285,6 @@ def __init__( # takes the even path (equal-length rows -> graph replay) instead of the # ragged eager chain for mixed-length batches. Numerics are unchanged # (even is an existing code path). - self._single_row_batch = os.environ.get("VLLM_OMNI_C2W_SINGLE_ROW", "") == "1" self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") self._default_prompt_normalized: tuple[str, str] | None = None @@ -671,13 +670,6 @@ def _iter_decode_batches( buckets: Iterable[list[_WorkItem]], ) -> Iterable[list[_WorkItem]]: for bucket in buckets: - if self._single_row_batch: - # One row per batch -> every CFM call lands on the even path - # (equal-length rows -> graph replay), intentionally bypassing - # the wave split and min-batch gate below. - for item in bucket: - yield [item] - continue if not self._initial_batch_size: yield bucket continue diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 29e5c4ba248..bf99f65d245 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -503,18 +503,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self._tts_config = None self._codec_eos_id = 0 - # K-step activation (910C/A3 target). - # K is the frame count per Talker decode step: a value >= 2 engages the - # multi-frame loop -- the scheduler hands each request K query positions - # through the vLLM V1 speculative path (constant `continue` drafts), - # talker_multiframe.run() replays the decode graph K times, and this - # model samples one codec frame per replay in-model (the vLLM-level head - # degenerates to a one-hot stop/continue row so the rejection sampler - # verifies without touching the codec stream). - # SoC guard: the spec-driven verify trips rejection_random_sample_kernel - # past the vector-core limit on the 910B family, so the gate refuses to - # arm there instead of leaving a crash switch behind an env var. - self._k_step_frames = self._parse_k_step_frames() + # Multi-frame decode: K codec frames per Talker decode step. A value + # >= 2 engages the loop -- the scheduler hands each request K query + # positions through the vLLM V1 speculative path (constant `continue` + # drafts), talker_multiframe.run() replays the decode graph K times, + # and this model samples one codec frame per replay in-model (the + # vLLM-level head degenerates to a one-hot stop/continue row so the + # rejection sampler verifies without touching the codec stream). K is + # read back from the injected stage-1 speculative_config. + self._k_step_frames = self._parse_k_step_frames(vllm_config) self.supports_multi_frame_decode = self._k_step_frames > 0 # Per-request device RNG streams for in-model codec sampling (the # multi-frame loop cannot run the vLLM host sampler per frame). @@ -558,80 +555,38 @@ def _probe_soc_name() -> str: return "" @classmethod - def _soc_allows_k_step(cls) -> bool: - """Ask the very question the deploy-config gate asked, and no other. - - 910B family: rejection_random_sample_kernel hits a vector-core limit - once the verifier runs K tokens per request, which crashed six bring-up - attempts before this gate existed. 910C / A3 is the activation target. - - This must delegate rather than re-probe. The config layer answers the - same question while building the deploy config (it decides whether to - inject the stage-1 speculative_config), and it can only see the launch - environment -- it deliberately never touches the device. A second, - device-probing answer here could differ from that one, and a mismatch - is not benign: if the loop arms while the scheduler has no - speculative_config it silently never engages, and if the scheduler - reserves K positions while the loop refuses them the stage deadlocks - with no request making progress. One source of truth, one answer. - """ - from vllm_omni.config.stage_config import ( - _MINICPMO_TALKER_FRAMES_ENV, - _npu_soc_allows_talker_multiframe, - ) - - if os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip(): - # An explicit frame count is an operator decision; the deploy-config - # gate arms on it too, so the two stay in agreement. - logger.info( - "[minicpmo] %s set explicitly; skipping the SoC check", - _MINICPMO_TALKER_FRAMES_ENV, - ) - return True - return _npu_soc_allows_talker_multiframe() - - @classmethod - def _parse_k_step_frames(cls) -> int: + def _parse_k_step_frames(cls, vllm_config: Any = None) -> int: """K codec frames per Talker step, or 0 when the loop must not engage. - The count comes from ``config.stage_config.talker_frames_per_step`` -- - the very function the deploy-config loader uses to size the injected - speculative_config -- so the runner and the scheduler cannot disagree - about K. An explicit ``VLLM_OMNI_MINICPMO_TALKER_FRAMES`` (or the older - ``OMNI_K_STEP``) skips the SoC gate, exactly as the loader does for the - same input; that keeps the two sides aligned even when an operator - forces the loop on by hand. + Read back from the engine-side speculative_config that stage_config.py + injects for stage 1 (method ``ngram`` with ``num_speculative_tokens == + K - 1``). Deriving K from the engine config rather than from a separate + channel is what keeps the runner and the scheduler from disagreeing: a + mismatch either arms the loop with no drafts to consume, or leaves it + off while the scheduler reserves K positions per request. """ - from vllm_omni.config.stage_config import ( - _MINICPMO_FRAMES_OFF, - _MINICPMO_TALKER_FRAMES_ENV, - talker_frames_per_step, - ) + from vllm_omni.config.stage_config import _MINICPMO_TALKER_FRAMES_MAX - raw = os.environ.get("OMNI_K_STEP", "").strip().lower() - explicit = os.environ.get(_MINICPMO_TALKER_FRAMES_ENV, "").strip().lower() - if raw not in ("", "0", "off", "false", "no"): - try: - frames = int(raw) - except ValueError: - raise ValueError(f"OMNI_K_STEP must be an integer frame count (>=2), got {raw!r}") from None - explicit = "forced" + model_cfg = getattr(vllm_config, "model_config", None) + if getattr(model_cfg, "model_stage", None) != "tts": + return 0 + spec = getattr(vllm_config, "speculative_config", None) + if isinstance(spec, dict): + method = spec.get("method") + num_spec = spec.get("num_speculative_tokens", 0) or 0 else: - frames = talker_frames_per_step() + method = getattr(spec, "method", None) + num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 + if method != "ngram": + return 0 + frames = int(num_spec) + 1 if frames < 2: # K=1 degenerates to the ordinary one-frame path; treat it as off # so the flag never silently half-arms the pipeline. return 0 - if frames > 16: - raise ValueError(f"K-step frame count is capped at 16, got {frames}") - forced = explicit not in ("",) + _MINICPMO_FRAMES_OFF - if not forced and not cls._soc_allows_k_step(): - logger.warning( - "[minicpmo] K-step decode left off: the SoC either cannot verify " - "spec-width rows (910B family limit) or was not identified.", - ) - return 0 - logger.info("[minicpmo] Talker K-step decode armed: %d codec frames per step", frames) + if frames > _MINICPMO_TALKER_FRAMES_MAX: + raise ValueError(f"K-step frame count is capped at {_MINICPMO_TALKER_FRAMES_MAX}, got {frames}") + logger.info("[minicpmo] Talker multi-frame decode armed: %d codec frames per step", frames) return frames def _init_native_talker(self, prefix: str) -> None: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 9c6656c978d..ae4d400ae19 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -28,25 +28,20 @@ def _talker_stop_token_ids() -> list[int]: Two shapes exist for the Talker's vLLM-level head: - * K-frame loop off -- ``compute_logits`` is the real 6562-wide codec head, - so the sampled id *is* a codec id and the codec EOS (6561) is the stop; - * K-frame loop on -- ``compute_logits`` collapses to the two-wide + * multi-frame decode on -- ``compute_logits`` collapses to the two-wide continue/stop row, so the only sampleable ids are 0 and 1 and the codec EOS can never appear at this level (the model samples codec ids inside the loop and forwards them to stage 2 itself). - Pinning this to [6561] regardless is what made every K-step request run to - ``max_tokens``: the model did emit the stop marker, ``check_stop`` compared - it against 6561, and the request kept its slot until the context ran out. - - Asked per construction so both this file and the deploy loader read the - same predicate; the env the answer depends on is set before any of this is - imported. + Pinning this to [6561] regardless is what made every multi-frame request run + to ``max_tokens``: the model did emit the stop marker, ``check_stop`` + compared it against 6561, and the request kept its slot until the context + ran out. A deployment that turns multi-frame decode off overrides this per + stage with the codec EOS. """ - from vllm_omni.config.stage_config import talker_multiframe_armed from vllm_omni.platforms.npu.worker.talker_multiframe import STOP_TOKEN_ID - return [STOP_TOKEN_ID] if talker_multiframe_armed() else [_CODEC_EOS_TOKEN_ID] + return [STOP_TOKEN_ID] MINICPMO_4_5_PIPELINE = PipelineConfig( diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py index c69f1665141..0944ce1a0ab 100644 --- a/vllm_omni/platforms/npu/ascend_warmup_patch.py +++ b/vllm_omni/platforms/npu/ascend_warmup_patch.py @@ -83,18 +83,24 @@ def _probe_soc_name() -> str: def _kstep_armed() -> bool: - """Whether the Talker K-step is explicitly armed (frames > 1). + """Whether this worker runs the Talker multi-frame decode. - Single source of truth is the same env the deploy config reads; a loader - that cannot import it is treated as "not armed" so the baseline never - changes by accident. + Read off the engine's speculative_config, which the deploy YAML's stage-1 + block provides; an unreadable config counts as not armed so the baseline + warmup never changes by accident. """ try: - from vllm_omni.config.stage_config import talker_frames_per_step + from vllm.config import get_current_vllm_config_or_none - return talker_frames_per_step() > 1 + cfg = get_current_vllm_config_or_none() except Exception: return False + spec = getattr(cfg, "speculative_config", None) if cfg is not None else None + if spec is None: + return False + method = getattr(spec, "method", None) + num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 + return method == "ngram" and num_spec > 0 def _skipped_names() -> set[str]: diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py index c7a8fe8e3d1..f9700d38a3c 100644 --- a/vllm_omni/platforms/npu/worker/talker_multiframe.py +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -12,8 +12,8 @@ ## How a frame advances without the host -The step is scheduled with K query positions per request (see -``_apply_minicpmo_talker_multiframe_default``). vLLM prepares all K of them -- +The step is scheduled with K query positions per request (stage 1's +``speculative_config`` in the deploy YAML). vLLM prepares all K of them -- positions, KV slots, block table, sequence lengths -- and the captured graph is a K-token uniform decode. What vLLM cannot do is fill in the *inputs*: frame k+1's embedding is the codec token sampled at frame k, so the K positions are @@ -49,7 +49,7 @@ drops (``_codec_scalars``). So the audio is expected to be bit-identical, not merely close -- which is a far stronger check than the CER gate. -Off with ``VLLM_OMNI_MINICPMO_TALKER_FRAMES=1``. +Off when stage 1's deploy config carries no speculative_config. """ from __future__ import annotations From 34891c821d8f379de3b7f43187b52225d8dfcaef Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 21 Sep 2026 22:48:33 +0800 Subject: [PATCH 14/89] refactor(minicpmo_4_5): configure the Talker multi-frame decode from the deploy YAML Make stage 1's speculative_config in the deploy YAML the single source of truth for the multi-frame decode, and keep this port to the optimizations plus their tests only. * the deploy YAML carries stage 1's speculative_config (8 codec frames per step) and async_scheduling: false explicitly; no code-side default remains. * the environment switches VLLM_OMNI_MINICPMO_TALKER_FRAMES, OMNI_K_STEP, VLLM_OMNI_C2W_SINGLE_ROW and VLLM_OMNI_MINICPMO_NARROW_REPLAY are gone, and the 910B SoC gate is dropped: the rejection-verifier limit it guarded against was the missing multi-frame wiring this branch fixes, and 910B3 is validated end to end. * the model-side and scheduler-side gates derive K from the engine config (model_stage == tts plus the n-gram fingerprint), so the two cannot disagree. * the per-step/per-frame timers, the stop and bookkeeping traces, the guard view and utils/step_prof.py are removed: they were bring-up instrumentation, not optimizations. * the stage-2 single-row even-batch item is reverted: its measured gain sat inside the run-to-run drift band and a default-off switch adds configuration without value. * tests follow the new deploy values (chunk 75, FULL_DECODE_ONLY, the multi-frame stop id, all-stage concurrency 8). Signed-off-by: y-null --- .../models/minicpmo_4_5/test_pipeline.py | 25 +- .../test_talker_kstep_feedback.py | 72 +--- vllm_omni/core/sched/omni_ar_scheduler.py | 116 ------- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 111 +----- .../npu/worker/npu_ar_model_runner.py | 24 +- .../platforms/npu/worker/talker_multiframe.py | 325 +----------------- vllm_omni/utils/step_prof.py | 86 ----- 7 files changed, 43 insertions(+), 716 deletions(-) delete mode 100644 vllm_omni/utils/step_prof.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 9cf7b04b609..855af4185bf 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -108,7 +108,10 @@ def test_talker_stage(self, pipeline: PipelineConfig) -> None: assert talker.hf_config_name == "tts_config" assert talker.sampling_constraints == { "detokenize": False, - "stop_token_ids": [6561], + # The multi-frame head: stage 1's deploy config carries the + # speculative_config that collapses the vLLM-level head to the + # two-wide continue/stop row, whose only stop id is 1. + "stop_token_ids": [1], } assert talker.custom_process_next_stage_input_func == ( "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload" @@ -161,7 +164,11 @@ def test_deploy_resolves_three_stage_pipeline(self, filename: str, devices: list assert stages[2].yaml_extras["input_connectors"]["from_stage_1"] == "connector_of_shared_memory" connector = deploy.connectors["connector_of_shared_memory"] assert connector["name"] == "SharedMemoryConnector" - assert connector["extra"]["codec_chunk_frames"] == 25 + # The single-card deploy raises the steady chunk to 75 so + # initial_codec_chunk_frames: 25 can take effect; multi-card variants + # keep the base value. + expected_chunk_frames = 75 if filename == "minicpmo_4_5.yaml" else 25 + assert connector["extra"]["codec_chunk_frames"] == expected_chunk_frames assert connector["extra"]["codec_left_context_frames"] == 3 assert connector["extra"]["enable_hift_graph"] is True assert connector["extra"]["connector_get_max_wait_first_chunk"] == 3000 @@ -170,9 +177,9 @@ def test_deploy_resolves_three_stage_pipeline(self, filename: str, devices: list assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor) assert "hf_overrides" not in stages[1].yaml_engine_args if filename == "minicpmo_4_5.yaml": - # Stage-1 caps out at 8 so concurrent Talker decodes are not - # serialized behind the base value of 4. - assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 8, 4] + # All three stages cap out at 8 so the top tier is not gated by a + # single stage. + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [8, 8, 8] memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] assert memory_utilizations == [ 0.55, @@ -209,8 +216,12 @@ def test_npu_graph_configuration_is_stage_scoped(self, filename: str) -> None: deploy = _apply_platform_overrides(deploy, platform="npu") stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy) - assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE" - assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE" + # The single-card deploy captures a full decode graph, which keeps + # in-graph sampling and multi-frame expansion; multi-card variants keep + # the base mode. + expected_graph_mode = "FULL_DECODE_ONLY" if filename == "minicpmo_4_5.yaml" else "PIECEWISE" + assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == expected_graph_mode + assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == expected_graph_mode assert stages[2].yaml_engine_args["enforce_eager"] is True # Only the single-card deploy raises the Code2Wav graph pool; the # multi-card variants keep the base value. diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 92d2f2b548d..a0f5ba31e77 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -346,83 +346,33 @@ def test_stop_vocab_gate_stays_off_without_multi_frame_or_rows(): assert armed.input_batch.vocab_size == 0 -def test_stop_trace_is_off_by_default_and_never_kills_the_step(monkeypatch): - from vllm_omni.platforms.npu.worker import talker_multiframe - - monkeypatch.delenv(talker_multiframe._STOP_TRACE_ENV, raising=False) - monkeypatch.setattr(talker_multiframe, "_STOP_TRACE_COUNT", 0) - assert talker_multiframe.stop_trace_enabled() is False - - # Worst input first: the diagnostic must degrade, never raise into the step. - talker_multiframe._trace_stop_rows([None]) - assert talker_multiframe.stop_trace_enabled() is False - - # Armed and well-formed: one row per request, one column per frame. - monkeypatch.setattr(talker_multiframe, "_STOP_TRACE_COUNT", 0) - monkeypatch.setenv(talker_multiframe._STOP_TRACE_ENV, "1") - assert talker_multiframe.stop_trace_enabled() is True - keep = torch.tensor([[0.0, -float("inf")]]) - stop = torch.tensor([[-float("inf"), 0.0]]) - talker_multiframe._trace_stop_rows([keep, stop, stop]) - assert talker_multiframe.stop_trace_enabled() is True - - class _MinTokensLogitsProcessor: - """名字就是契约：实现按 ``type(proc).__name__`` 匹配。""" + """The name is the contract: the implementation matches on ``type(proc).__name__``.""" def __init__(self, min_toks): self.min_toks = min_toks def test_kstep_min_tokens_neutralization_clears_the_censor_list(): - """K 步下 vLLM 层 min_tokens 的名单必须被清空。 - - 它 mask 的是唯一的停止信号（binary stop 行的 id 1），而解除条件是 - ``len(output_token_ids) >= min_tokens``——那个长度归 vLLM 的 spec 记账管， - 计数不推进请求就永远停不下来（910C 现场 stage1 全 ``length``、0 ``stop``）。 - 模型内已用 ``state.step < min_tokens`` 压 codec EOS，这一层是重复的。 + """The vLLM-level ``min_tokens`` mask list must be cleared under multi-frame decode. + + It masks the request's only stop signal (id 1 of the binary stop row) and + unmasks only once ``len(output_token_ids) >= min_tokens`` -- a counter owned + by vLLM's spec bookkeeping, so a request whose counter does not advance can + never stop (observed: stage 1 all ``length``, zero ``stop``). The model-side + guard already holds the codec EOS back with ``state.step < min_tokens``, so + this layer is redundant. """ from vllm_omni.platforms.npu.worker import talker_multiframe censor = _MinTokensLogitsProcessor({0: (50, [0, 0, 0], {1})}) untouched = _MinTokensLogitsProcessor({0: (50, [0], {1})}) - untouched.__class__ = type("SomeOtherProcessor", (object,), {}) # 非目标处理器 + untouched.__class__ = type("SomeOtherProcessor", (object,), {}) # not the target processor talker_multiframe.neutralize_kstep_min_tokens(SimpleNamespace(non_argmax_invariant=[untouched, censor])) assert censor.min_toks == {} assert untouched.min_toks == {0: (50, [0], {1})} - # 空名单 / 怪输入都不能炸。 + # An empty list and odd inputs must not raise either. talker_multiframe.neutralize_kstep_min_tokens(SimpleNamespace(non_argmax_invariant=[])) talker_multiframe.neutralize_kstep_min_tokens(None) - - -def test_kstep_bookkeeping_trace_is_total(monkeypatch): - """记账层诊断：正常输入打一行，最坏输入只关掉自己。""" - from vllm_omni.platforms.npu.worker import talker_multiframe - - monkeypatch.setenv(talker_multiframe._STOP_TRACE_ENV, "1") - monkeypatch.setattr(talker_multiframe, "_STOP_KB_COUNT", 0) - - runner = SimpleNamespace( - input_batch=SimpleNamespace( - sampling_metadata=SimpleNamespace( - logitsprocs=SimpleNamespace(non_argmax_invariant=[_MinTokensLogitsProcessor({0: (50, [0, 0], {1})})]) - ) - ) - ) - talker_multiframe.trace_kstep_bookkeeping(runner, [[0, 0, 1]], torch.zeros(2, 2)) - assert talker_multiframe.stop_trace_enabled() is True - - # 最坏输入：内部构造真的失败时必须降级（绝不抛回 decode 步里）并永久关掉 - # 自己，否则每个 step 都要付一次异常代价。`[5]` 让 list(row) 抛 TypeError。 - monkeypatch.setattr(talker_multiframe, "_STOP_KB_COUNT", 0) - talker_multiframe.trace_kstep_bookkeeping(SimpleNamespace(), [5], torch.zeros(2, 2)) - assert talker_multiframe.stop_trace_enabled() is False - - # 取不到 runner/输入是合法场景的优雅降级，不是失败：不能因此把诊断关掉， - # 否则一次空步就会让后面所有取证静默。 - monkeypatch.setenv(talker_multiframe._STOP_TRACE_ENV, "1") - monkeypatch.setattr(talker_multiframe, "_STOP_KB_COUNT", 0) - talker_multiframe.trace_kstep_bookkeeping(None, None, None) - assert talker_multiframe.stop_trace_enabled() is True diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 7e8e219dd0a..9f70ac09d76 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -33,74 +33,6 @@ logger = init_logger(__name__) -# --------------------------------------------------------------------------- -# Multi-frame bookkeeping trace (off by default). -# -# VLLM_OMNI_MINICPMO_KSTEP_ACCOUNT_TRACE=1 -# -# Question it answers: the model already emits stop rows, so why does the -# request not finish? It logs each request's accepted token values per step -- -# if they stay 0 (continue), the stop row never became a 1 in the request's -# token table and check_stop(stop_token_ids=[1]) never fires, so the request -# runs to max_model_len. Diagnostics never raise. -# --------------------------------------------------------------------------- -_KSTEP_ACCOUNT_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_ACCOUNT_TRACE" -_KSTEP_ACCOUNT_TRACE: dict = {"lines": 0, "reqs": {}, "off": False} - - -def kstep_account_trace_enabled() -> bool: - """Off unless the env turns it on; latched off for good after a failure.""" - if _KSTEP_ACCOUNT_TRACE["off"]: - return False - raw = os.environ.get(_KSTEP_ACCOUNT_TRACE_ENV, "").strip().lower() - return bool(raw) and raw not in ("0", "false", "no", "off") - - -def _trace_kstep_account(req_id, spec_len, generated, computed, total, prompt_len, status, sampling=None) -> None: - """One line per request per step; first 24 steps of each request, then 1 in 16.""" - state = _KSTEP_ACCOUNT_TRACE - if state["off"]: - return - try: - key = str(req_id) - seen = state["reqs"].get(key, 0) - state["reqs"][key] = seen + 1 - gen_len = len(generated or []) - # Invariant I2/I5: a step with drafts must hand the request - # spec_len + 1 tokens (the bonus plus every accepted draft). A short - # row means frames the model produced were dropped on the way to the - # request -- always log it, whatever the stride. - partial = bool(spec_len) and gen_len != spec_len + 1 - interesting = seen < 24 or seen % 16 == 0 or partial - nonzero = [int(t) for t in (generated or []) if int(t) != 0] - if not interesting and not nonzero: - return - if state["lines"] >= 400: - return - state["lines"] += 1 - logger.info( - "[kstep-acct] req=%s step=%d spec_len=%s gen=%s partial=%s nonzero=%s computed=%s tokens=%s " - "prompt=%s status=%s stop_ids=%s eos=%s min=%s max=%s ignore_eos=%s", - key[:12], - seen, - spec_len, - gen_len, - int(partial), - nonzero[:6], - computed, - total, - prompt_len, - status, - list(getattr(sampling, "stop_token_ids", None) or []), - getattr(sampling, "eos_token_id", None), - getattr(sampling, "min_tokens", None), - getattr(sampling, "max_tokens", None), - getattr(sampling, "ignore_eos", None), - ) - except Exception as exc: # pragma: no cover - diagnostics never raise - state["off"] = True - logger.warning("[kstep-acct] trace disabled after failure: %r", exc) - def _should_emit_engine_output( model_config: Any, @@ -437,31 +369,6 @@ def _talker_kstep_armed(self) -> bool: cache = self._omni_talker_kstep_cache = is_talker and is_ngram and num_spec > 0 return cache - def _log_kstep_guard_view(self, verdict: str, widths: set[int], states: list, dropped: int = 0) -> None: - """Diagnostic view of what the K-step guard judged this step. - - Rate-limited but never silenced on the interesting steps: the 12:39 - crash landed on a step whose one-shot guard log had already been - consumed, which hid exactly the state the guard passed. Diagnostics - must never raise either -- this one must not kill the engine. - """ - count = getattr(self, "_omni_kstep_guard_view_count", 0) + 1 - self._omni_kstep_guard_view_count = count - if count > 20 and count % 50 != 0: - return - try: - logger.info( - "[K-guard] #%d %s: waiting=%d widths=%s dropped=%d " - "states=(id, computed, prompt_len, total, spec_len, delta)=%s", - count, - verdict, - len(self.waiting), - sorted(widths), - dropped, - states, - ) - except Exception as exc: # never let the diagnostic kill the engine - logger.info("[K-guard] #%d %s (view failed: %r)", count, verdict, exc) def _drop_talker_drafts_if_prefill_pending(self) -> None: """Keep the Talker's K-frame decode out of steps with uneven spans. @@ -516,22 +423,10 @@ def _drop_talker_drafts_if_prefill_pending(self) -> None: widths.add(delta) uneven = len(widths) > 1 if not prefill_pending and not uneven: - if not widths and states: - # Running requests exist but none yields a row width: the step - # schedules no decode rows at all. Rare enough to be interesting. - self._log_kstep_guard_view("pass-no-decode-widths", widths, states) return - dropped = 0 for req in self.running: if req.spec_token_ids: req.spec_token_ids = [] - dropped += 1 - self._log_kstep_guard_view( - f"drop ({'prefill pending' if prefill_pending else 'uneven spans'})", - widths, - states, - dropped=dropped, - ) def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Remove FINISHED_ABORTED requests before the upstream scheduler sees @@ -717,17 +612,6 @@ def update_from_output( continue req_index = model_runner_output.req_id_to_index[req_id] generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] - if kstep_account_trace_enabled(): - _trace_kstep_account( - req_id, - len(scheduler_output.scheduled_spec_decode_tokens.get(req_id) or []), - generated_token_ids, - int(request.num_computed_tokens), - int(request.num_tokens), - len(getattr(request, "prompt_token_ids", None) or []), - getattr(request, "status", None), - getattr(request, "sampling_params", None), - ) stale_async_tokens = int(getattr(request, "async_tokens_to_discard", 0) or 0) async_output_is_stale = bool(generated_token_ids and stale_async_tokens > 0) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index bf99f65d245..45d95f0535a 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -49,7 +49,6 @@ ) from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform -from vllm_omni.utils.step_prof import span as _prof_span logger = init_logger(__name__) @@ -341,85 +340,6 @@ def resolve_codec_sampling_params( return resolved -# --------------------------------------------------------------------------- -# Per-frame trace for the multi-frame loop (off by default; complements the -# per-step timers in talker_multiframe). -# -# VLLM_OMNI_MINICPMO_KSTEP_FRAME_TRACE=1 enable -# prints at most _KSTEP_FRAME_TRACE_FRAMES frames, then goes quiet -# -# Answers what the per-step timers cannot: whether the model's own sampler -# reaches EOS inside the K loop, and whether each frame is fed the previous -# frame's sampled code (frame-order invariant). Diagnostics never take the -# engine down: on any error they disable themselves and log one warning. -# --------------------------------------------------------------------------- -_KSTEP_FRAME_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_FRAME_TRACE" -_KSTEP_FRAME_TRACE_LINES = 128 # max printed lines -_KSTEP_FRAME_TRACE_HEAD = 16 # first N frames printed one by one (frame-order check) -_KSTEP_FRAME_TRACE_STRIDE = 16 # then one line every N frames (covers ~1800 frames) -_KSTEP_FRAME_TRACE: dict[str, Any] = {"reqs": {}, "lines": 0, "off": False} - - -def kstep_frame_trace_enabled() -> bool: - """Off unless the env turns it on; latched off for good after a failure.""" - if _KSTEP_FRAME_TRACE["off"]: - return False - raw = os.environ.get(_KSTEP_FRAME_TRACE_ENV, "").strip().lower() - return bool(raw) and raw not in ("0", "false", "no", "off") - - -def _trace_kstep_frame(model: Any, **fields: Any) -> None: - """One line per codec frame while the multi-frame loop owns sampling. - - Logs what the *model* decided: the code each frame took as its input (the - previous frame's sample, per the contract), what it sampled, whether that - was the codec EOS, and where the sampler's own step counter stood. The first - ``HEAD`` frames are logged one per line (that is what verifies the frame - chain), then one line every ``STRIDE`` frames so the window still covers the - frames where EOS first becomes eligible -- `min_tokens` masks it before - then, so a head-only trace can never see a stop. An EOS sample is always - logged regardless of the stride. Never raises: a failure disables the trace - for good. - """ - state = _KSTEP_FRAME_TRACE - if state["off"]: - return - reqs = state["reqs"] - request_id = str(fields.get("request_id")) - seen = reqs.get(request_id, 0) - reqs[request_id] = seen + 1 - is_eos = bool(fields.get("is_eos")) - interesting = seen < _KSTEP_FRAME_TRACE_HEAD or ( - seen >= _KSTEP_FRAME_TRACE_HEAD and (seen - _KSTEP_FRAME_TRACE_HEAD) % _KSTEP_FRAME_TRACE_STRIDE == 0 - ) - if not is_eos and (not interesting or state["lines"] >= _KSTEP_FRAME_TRACE_LINES): - return - try: - device_states = getattr(model, "_request_codec_device_states", None) - device_state = device_states.get(fields.get("request_id")) if isinstance(device_states, dict) else None - device_step = int(device_state.step.reshape(-1)[0].item()) if device_state is not None else -1 - state["lines"] += 1 # cap printed lines only; per-request windows count separately - logger.info( - "[kstep-frame] frame=%d req=%s step=%s->%s device_step=%s in_code=%s row=%s " - "sampled=%s is_eos=%s limit=%s finished=%s ctx=%s min=%s max=%s", - seen, - str(fields.get("request_id"))[:12], - fields.get("step_before"), - fields.get("step_after"), - device_step, - fields.get("prev_code"), - fields.get("row"), - fields.get("sampled_id"), - fields.get("is_eos"), - fields.get("reached_limit"), - fields.get("finished"), - fields.get("ctx_len"), - fields.get("min_tokens"), - fields.get("max_tokens"), - ) - except Exception as exc: # pragma: no cover - diagnostics never raise - state["off"] = True - logger.warning("[kstep-frame] trace disabled after failure: %r", exc) def _codec_int_param(state: Any, key: str, fallback: int) -> int: @@ -1151,22 +1071,6 @@ def make_omni_output( reached_limit = int(state["step"]) >= _codec_int_param(state, "max_tokens", self._codec_max_tokens) finished = is_eos or reached_limit state["finished"] = finished - if self._k_step_frames > 0 and kstep_frame_trace_enabled(): - _trace_kstep_frame( - self, - request_id=request_id, - step_before=step, - step_after=int(state["step"]), - prev_code=prev_code, - row=int(end) - 1, - sampled_id=sampled_id, - is_eos=is_eos, - reached_limit=reached_limit, - finished=finished, - ctx_len=int(codes.numel()) if isinstance(codes, torch.Tensor) else -1, - min_tokens=min_tokens, - max_tokens=max_tokens, - ) # MiniCPMTTS.generate_chunk consumes the boundary sample but # returns only codes that were fed into the retained KV state. if not is_eos and not reached_limit: @@ -1744,11 +1648,10 @@ def sample(self, logits, sampling_metadata): and logits.shape[-1] == 2 and not getattr(sampling_metadata, "max_num_logprobs", None) ): - with _prof_span("tts_sampler"): - return SamplerOutput( - sampled_token_ids=logits.argmax(dim=-1, keepdim=True).to(torch.int32), - logprobs_tensors=None, - ) + return SamplerOutput( + sampled_token_ids=logits.argmax(dim=-1, keepdim=True).to(torch.int32), + logprobs_tensors=None, + ) prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None) if isinstance(logits, torch.Tensor) and isinstance(prompt_ids, torch.Tensor): # Copy rather than mutate: the runner may hand us the input batch's @@ -1757,15 +1660,13 @@ def sample(self, logits, sampling_metadata): sampling_metadata, prompt_token_ids=blank_scheduler_prompt_for_penalties(prompt_ids, logits.shape[-1]), ) - with _prof_span("tts_penalty"): - logits, sampling_metadata = self._apply_codec_repetition_penalty(logits, sampling_metadata) + logits, sampling_metadata = self._apply_codec_repetition_penalty(logits, sampling_metadata) prewarped = _maybe_prewarp_top_k_top_p(logits, sampling_metadata) if prewarped is not None: logits, sampling_metadata = prewarped force_eos = self._pending_force_eos_rows self._pending_force_eos_rows = None - with _prof_span("tts_sampler"): - output = Sampler()(logits, sampling_metadata) + output = Sampler()(logits, sampling_metadata) return self._force_eos_on_sampled_ids(output, force_eos) def _force_eos_on_sampled_ids(self, output: Any, force_eos: list[bool] | None) -> Any: diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index ea9d37411ae..073317150c4 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -51,7 +51,6 @@ from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element -from vllm_omni.utils.step_prof import span, tic, toc from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, needs_omni_connector, @@ -423,7 +422,6 @@ def execute_model( scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: - tic("em_total") if self.vllm_config.model_config.enable_return_routed_experts: capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): @@ -773,10 +771,9 @@ def execute_model( ): if self.cache_config.mamba_cache_mode == "align": mamba_utils.do_mamba_copy_block(preprocess_bufs) - with span("em_forward"): - hidden_states = self._model_forward( - num_tokens_padded, input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs - ) + hidden_states = self._model_forward( + num_tokens_padded, input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs + ) with record_function_or_nullcontext("post process"): # -------------------------------------- Omni-new ------------------------------------------------- # [Omni] Map pending ropes metadata to req_ids. @@ -904,7 +901,6 @@ def execute_model( if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): self._omni_routed_experts_d2h(scheduler_output) - toc("em_total") return None def _sample( @@ -912,7 +908,6 @@ def _sample( logits: torch.Tensor | None, spec_decode_metadata: Any, ): - tic("em_sample") sampling_metadata = self.input_batch.sampling_metadata if spec_decode_metadata is None: model_sample = getattr(self.model, "sample", None) @@ -939,22 +934,18 @@ def _sample( requests=getattr(self, "requests", None), ) if sampler_output is not None: - toc("em_sample") return sampler_output - toc("em_sample") return self.sampler( logits=logits, sampling_metadata=sampling_metadata, ) - toc("em_sample") return super()._sample(logits, spec_decode_metadata) @torch.inference_mode() def sample_tokens( self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: - tic("st_total") profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config kv_connector_output = self.kv_connector_output self.kv_connector_output = None @@ -1109,13 +1100,6 @@ def propose_draft_token_ids(sampled_token_ids): spec_decode_metadata, ) - if getattr(self.model, "supports_multi_frame_decode", False): - from vllm_omni.platforms.npu.worker import talker_multiframe - - if talker_multiframe.stop_trace_enabled(): - talker_multiframe.trace_kstep_bookkeeping( - self, valid_sampled_token_ids, logits - ) with record_function_or_nullcontext("draft_token"): if self.speculative_config: @@ -1374,7 +1358,6 @@ def _unwrap_lists(v): self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output) if not self.use_async_scheduling: - toc("st_total") return model_runner_output async_output = AsyncGPUModelRunnerOutput( model_runner_output=model_runner_output, @@ -1389,7 +1372,6 @@ def _unwrap_lists(v): async_output.sampled_token_ids_cpu, async_output.async_copy_ready_event, ) - toc("st_total") return async_output # -------------------------------------- Omni-new ------------------------------------------------- diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py index f9700d38a3c..8398742d433 100644 --- a/vllm_omni/platforms/npu/worker/talker_multiframe.py +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -58,7 +58,6 @@ import os from collections.abc import Callable from dataclasses import dataclass -from time import perf_counter from typing import Any import torch @@ -68,291 +67,6 @@ _LOGGED_ENGAGE = False -# --------------------------------------------------------------------------- -# Per-step/per-frame timers for the multi-frame loop (off by default; when off, -# run() takes no timers at all). -# -# VLLM_OMNI_MINICPMO_KSTEP_PROF=1 enable -# VLLM_OMNI_MINICPMO_KSTEP_PROF_STEPS=3 sample the first N steps with a device sync -# VLLM_OMNI_MINICPMO_KSTEP_PROF_EVERY=25 print a cumulative summary every N steps -# -# Purpose: split one K-frame step's wall time between host enqueue (replay/ -# after_forward, asynchronous) and device execution (synced, measured after an -# explicit synchronize). Diagnostics never raise: on any error they turn -# themselves off permanently and log one warning. -# --------------------------------------------------------------------------- -_PROF_ENV = "VLLM_OMNI_MINICPMO_KSTEP_PROF" -_PROF_STEPS_ENV = "VLLM_OMNI_MINICPMO_KSTEP_PROF_STEPS" -_PROF_EVERY_ENV = "VLLM_OMNI_MINICPMO_KSTEP_PROF_EVERY" -_PROF_OFF = ("0", "false", "no", "off") -_PROF_STEPS_DEFAULT = 3 -_PROF_EVERY_DEFAULT = 25 - - -class _StepProf: - """One step's timers. Never raises: profiling is best-effort by contract.""" - - __slots__ = ("frames", "replay_ms", "after_ms", "detail") - - def __init__(self, frames: int, *, detail: bool) -> None: - self.frames = frames - self.replay_ms = 0.0 - self.after_ms = 0.0 - self.detail = detail - - -def prof_enabled() -> bool: - """Off unless the env explicitly turns it on (``1``/``true``/``yes``/``on``).""" - try: - raw = os.environ.get(_PROF_ENV, "").strip().lower() - return bool(raw) and raw not in _PROF_OFF - except Exception: # pragma: no cover - env is a dict lookup, but be total - return False - - -def _prof_steps() -> int: - try: - raw = os.environ.get(_PROF_STEPS_ENV, "").strip() - steps = int(raw) if raw else _PROF_STEPS_DEFAULT - except Exception: - steps = _PROF_STEPS_DEFAULT - return max(steps, 0) - - -def _prof_every() -> int: - try: - raw = os.environ.get(_PROF_EVERY_ENV, "").strip() - every = int(raw) if raw else _PROF_EVERY_DEFAULT - except Exception: - every = _PROF_EVERY_DEFAULT - return max(every, 1) - - -def _prof_sync() -> None: - """Best-effort device drain; no-op on a build without the op.""" - for fn in ( - getattr(getattr(torch, "npu", None), "synchronize", None), - getattr(getattr(torch, "cuda", None), "synchronize", None), - ): - if callable(fn): - try: - fn() - return - except Exception: - return - - -def _prof_frame_line(step: int, frame: int, replay_ms: float, after_ms: float, synced_ms: float) -> None: - try: - logger.info( - "[kstep-prof] step=%d frame=%d replay_ms=%.3f after_forward_ms=%.3f synced_frame_ms=%s", - step, - frame, - replay_ms, - after_ms, - f"{synced_ms:.3f}" if synced_ms >= 0.0 else "n/a", - ) - except Exception: - pass - - -def _prof_summary(steps: int, frames_total: int, last_frames: int, *, detail: bool) -> None: - """One cumulative line: per-step/per-frame wall time, host enqueue vs device execution.""" - try: - per_step = _PROF_TOTAL_MS / steps if steps else 0.0 - per_frame = _PROF_TOTAL_MS / frames_total if frames_total else 0.0 - replay_per_frame = _PROF_REPLAY_MS / frames_total if frames_total else 0.0 - after_per_frame = _PROF_AFTER_MS / frames_total if frames_total else 0.0 - synced = ( - f" synced_frame_ms={_PROF_SYNCED_MS / _PROF_SYNCED_FRAMES:.3f} (n={_PROF_SYNCED_FRAMES})" - if _PROF_SYNCED_FRAMES - else "" - ) - logger.info( - "[kstep-prof] steps=%d frames=%d last_step_frames=%d step_ms=%.2f frame_ms=%.3f " - "replay_enqueue_ms_per_frame=%.3f after_forward_ms_per_frame=%.3f (totals: %.3f/%.3f)%s%s", - steps, - frames_total, - last_frames, - per_step, - per_frame, - replay_per_frame, - after_per_frame, - _PROF_REPLAY_MS, - _PROF_AFTER_MS, - synced, - " [per-frame lines above]" if detail else "", - ) - except Exception: - pass - - -_PROF_WARNED = False -_PROF_STEPS = 0 -_PROF_FRAMES = 0 -_PROF_REPLAY_MS = 0.0 -_PROF_AFTER_MS = 0.0 -_PROF_TOTAL_MS = 0.0 -_PROF_SYNCED_MS = 0.0 -_PROF_SYNCED_FRAMES = 0 - - -def _prof_disabled(reason: str) -> None: - """Turn the profiling off for the rest of the process, once.""" - global _PROF_WARNED - if _PROF_WARNED: - return - _PROF_WARNED = True - try: - logger.warning("[kstep-prof] disabled after failure: %s", reason) - except Exception: - pass - os.environ[_PROF_ENV] = "0" - - -def _prof_step_begin(frames: int) -> _StepProf | None: - """``None`` when off -- and then ``run()`` takes no timer at all.""" - try: - if not prof_enabled(): - return None - return _StepProf(frames, detail=_PROF_STEPS < _prof_steps()) - except Exception: - _prof_disabled("step_begin") - return None - - -def _prof_frame(prof: _StepProf, frame: int, t0: float, t1: float, t2: float, t3: float) -> None: - """Book one replay. ``replay``/``after`` are enqueue times; ``synced`` is the - frame's real wall clock (drained after the replay), and the gap between them - is the device side of the account.""" - try: - global _PROF_SYNCED_MS, _PROF_SYNCED_FRAMES - replay_ms = (t1 - t0) * 1000.0 - after_ms = (t2 - t1) * 1000.0 - prof.replay_ms += replay_ms - prof.after_ms += after_ms - synced_ms = -1.0 - if prof.detail and t3: - synced_ms = (t3 - t0) * 1000.0 - _PROF_SYNCED_MS += synced_ms - _PROF_SYNCED_FRAMES += 1 - if prof.detail: - _prof_frame_line(_PROF_STEPS + 1, frame, replay_ms, after_ms, synced_ms) - except Exception: - _prof_disabled("frame") - - -def _prof_step_end(prof: _StepProf | None, step_start: float) -> None: - if prof is None: - return - try: - global _PROF_STEPS, _PROF_FRAMES, _PROF_REPLAY_MS, _PROF_AFTER_MS, _PROF_TOTAL_MS - _PROF_STEPS += 1 - _PROF_FRAMES += prof.frames - _PROF_REPLAY_MS += prof.replay_ms - _PROF_AFTER_MS += prof.after_ms - _PROF_TOTAL_MS += (perf_counter() - step_start) * 1000.0 - if prof.detail or _PROF_STEPS % _prof_every() == 0: - _prof_summary(_PROF_STEPS, _PROF_FRAMES, prof.frames, detail=prof.detail) - except Exception: - _prof_disabled("step_end") - - -# --------------------------------------------------------------------------- -# Stop-row forensics for the multi-frame loop (off by default). -# -# VLLM_OMNI_MINICPMO_KSTEP_STOP_TRACE=1 -# -# Answers one question: does the model actually emit stop rows within the K -# frames of a step? Rows present but the request still runs to max_tokens means -# the loss is in vLLM's spec bookkeeping; no rows at all means it is in the -# multi-frame sampling/state path. Prints the first 50 steps, then goes quiet. -# Never raises. -# --------------------------------------------------------------------------- -_STOP_TRACE_ENV = "VLLM_OMNI_MINICPMO_KSTEP_STOP_TRACE" -_STOP_TRACE_STEPS = 50 -_STOP_TRACE_COUNT = 0 - - -def stop_trace_enabled() -> bool: - """Off unless the env explicitly turns it on (``1``/``true``/``yes``/``on``).""" - try: - raw = os.environ.get(_STOP_TRACE_ENV, "").strip().lower() - return bool(raw) and raw not in _PROF_OFF - except Exception: # pragma: no cover - env is a dict lookup, but be total - return False - - -def _trace_stop_rows(frame_stop_logits: list[torch.Tensor]) -> None: - """One line per step: how many of each request's K frames carry a stop row.""" - global _STOP_TRACE_COUNT - try: - if _STOP_TRACE_COUNT >= _STOP_TRACE_STEPS: - return - _STOP_TRACE_COUNT += 1 - flags = torch.stack([row.argmax(dim=-1) for row in frame_stop_logits], dim=1) - logger.info( - "[kstep-stop] step=%d frames=%d stops_per_req=%s", - _STOP_TRACE_COUNT, - len(frame_stop_logits), - flags.sum(dim=1).cpu().tolist(), - ) - except Exception as exc: - logger.warning("[kstep-stop] trace disabled after failure: %r", exc) - os.environ[_STOP_TRACE_ENV] = "0" - - -_STOP_KB_COUNT = 0 - - -def trace_kstep_bookkeeping(runner: Any, valid_sampled_token_ids: Any, logits: Any) -> None: - """Stop trace, layer 2: what survives spec bookkeeping. - - Two facts are printed together so the loss point is unambiguous: - - * ``carried`` -- tokens the rejection sampler handed to the request (first - 4 requests). A 1 means the stop token reached the request's own token - table; then only the engine-side finish check remains. All zeros with no - 1 means the stop disappeared earlier. - * ``min_toks`` -- requests still on the MinTokens mask list, with the output - length it sees: ``(index, min_tokens, len(output_token_ids), - sorted(stop_ids))``. While a request is on this list its stop columns are - -inf; if ``len(output_token_ids)`` never grows, the entry is never - cleared and the request can never stop. - - Never raises: on error it only disables itself. - """ - global _STOP_KB_COUNT - try: - if _STOP_KB_COUNT >= _STOP_TRACE_STEPS: - return - _STOP_KB_COUNT += 1 - if isinstance(valid_sampled_token_ids, list): - carried: Any = [list(row) for row in valid_sampled_token_ids[:4]] - else: - carried = f"<{type(valid_sampled_token_ids).__name__}>" - - masks: list[Any] = [] - sampling_metadata = getattr(getattr(runner, "input_batch", None), "sampling_metadata", None) - for proc in getattr(getattr(sampling_metadata, "logitsprocs", None), "non_argmax_invariant", None) or []: - min_toks = getattr(proc, "min_toks", None) - if isinstance(min_toks, dict) and min_toks: - masks.append( - [ - (index, int(min_tok), len(out_ids), sorted(stop_ids)) - for index, (min_tok, out_ids, stop_ids) in list(min_toks.items())[:4] - ] - ) - logger.info( - "[kstep-stop] carried=%s min_toks=%s logits=%s", - carried, - masks, - tuple(getattr(logits, "shape", ()) or ()), - ) - except Exception as exc: - logger.warning("[kstep-stop] bookkeeping trace disabled after failure: %r", exc) - os.environ[_STOP_TRACE_ENV] = "0" - def neutralize_kstep_min_tokens(logitsprocs: Any) -> None: """Stop vLLM's ``min_tokens`` layer from censoring the stop signal. @@ -394,9 +108,6 @@ def neutralize_kstep_min_tokens(logitsprocs: Any) -> None: _LOGGED_NARROW = False _LOGGED_NARROW_BLOCK: str | None = None -_NARROW_ENV = "VLLM_OMNI_MINICPMO_NARROW_REPLAY" -_NARROW_OFF = frozenset({"0", "off", "false", "no"}) -_NARROW_ON = frozenset({"1", "on", "true", "yes"}) # SoC families whose one-query capture is known to fault. The 910_93 part # trips a vector-core error (acl 507035) inside the packaged codec operator # while that graph is being captured, so the narrow path stays off there @@ -414,20 +125,11 @@ def neutralize_kstep_min_tokens(logitsprocs: Any) -> None: def _narrow_soc_allows() -> bool: """Whether this part may capture the one-query graph by default. - An explicit ``VLLM_OMNI_MINICPMO_NARROW_REPLAY=1/on`` always wins, and an - explicit ``0/off`` always loses. With no setting the answer depends on the - part: the environment is consulted first (the same variables the deploy - config reads, so both sides agree), and the device name is the fallback for - a worker launched without them. + Judged from the part: the CANN-exported SoC code first, the device name as + the fallback for a worker launched without them. """ global _NARROW_SOC_LOGGED - raw = os.environ.get(_NARROW_ENV, "").strip().lower() - if raw in _NARROW_ON: - return True - if raw in _NARROW_OFF: - return False - name = "" for variable in ("SOC_VERSION", "ASCEND_SOC_VERSION"): value = os.environ.get(variable, "").strip().lower() @@ -584,10 +286,9 @@ def select(self, frame: int) -> None: def narrow_replay_enabled(runner: Any) -> bool: """Whether this worker should capture and replay the one-query graph. - Off with ``VLLM_OMNI_MINICPMO_NARROW_REPLAY=off``, which restores the - K-query captures and the wide replay. On a part whose one-query capture is - known to fault (``_NARROW_BLOCKED_SOC_PREFIXES``) it is off as well unless - explicitly enabled, so the capture can never brick a deployment. + Off on a part whose one-query capture is known to fault + (``_NARROW_BLOCKED_SOC_PREFIXES``), so a capture can never brick a + deployment; those parts keep the K-query captures and the wide replay. """ if not _narrow_soc_allows(): return False @@ -738,9 +439,6 @@ def run( spans = model_kwargs_extra["request_token_spans"] infos = model_kwargs_extra["model_intermediate_buffer"] - step_prof = _prof_step_begin(frames) - step_start = perf_counter() if step_prof is not None else 0.0 - frame_outputs = [] frame_stop_logits = [] frame_hidden = [] @@ -757,16 +455,8 @@ def run( _write_frame_embeddings(model, inputs_embeds, input_ids, spans, infos, frame, narrow=narrow) elif frame > 0: _write_frame_embeddings(model, inputs_embeds, input_ids, spans, infos, frame) - t0 = perf_counter() if step_prof is not None else 0.0 hidden = run_model() - t1 = perf_counter() if step_prof is not None else 0.0 after_forward() - t2 = perf_counter() if step_prof is not None else 0.0 - if step_prof is not None and step_prof.detail: - _prof_sync() - t3 = perf_counter() if step_prof is not None else 0.0 - if step_prof is not None: - _prof_frame(step_prof, frame, t0, t1, t2, t3) frame_kwargs = dict(model_kwargs_extra) frame_kwargs["request_token_spans"] = ( list(narrow.spans) @@ -788,8 +478,6 @@ def run( if not _LOGGED_ENGAGE: _LOGGED_ENGAGE = True logger.info("[minicpmo] multi-frame Talker decode engaged: %d codec frames per step", frames) - if stop_trace_enabled(): - _trace_stop_rows(frame_stop_logits) merged = model.merge_frame_outputs(frame_outputs, frame_stop_logits) if narrow is not None and frame_hidden: # (frames, rows, hidden) -> (rows * frames, hidden), request-major -- @@ -797,9 +485,6 @@ def run( merged = merged._replace( text_hidden_states=torch.stack(frame_hidden, dim=1).reshape(-1, frame_hidden[0].shape[-1]) ) - # Whole step, frames + merge: this is the number to compare against the - # scheduler's step cadence. - _prof_step_end(step_prof, step_start) return merged diff --git a/vllm_omni/utils/step_prof.py b/vllm_omni/utils/step_prof.py deleted file mode 100644 index 7f27e5868fb..00000000000 --- a/vllm_omni/utils/step_prof.py +++ /dev/null @@ -1,86 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Opt-in per-step wall-clock profiler for MiniCPM-o NPU runners. - -Zero-ish overhead when disabled. Enable with ``OMNI_STEP_PROF=1`` and set the -aggregation cadence with ``OMNI_STEP_PROF_FLUSH`` (totals between dumps, -default 200). Thread-safe; spans are keyed per thread so nested/overlapping -regions in the engine loop stay correct. -""" - -from __future__ import annotations - -import os -import sys -import threading -import time -from collections import defaultdict -from contextlib import contextmanager - -_ENABLED = os.environ.get("OMNI_STEP_PROF", "0") == "1" -_FLUSH_EVERY = int(os.environ.get("OMNI_STEP_PROF_FLUSH", "200")) - - -class _StepProf: - def __init__(self) -> None: - self.acc: dict[str, float] = defaultdict(float) - self.cnt: dict[str, int] = defaultdict(int) - self._local = threading.local() - self._lock = threading.Lock() - self._flushed = 0 - - def tic(self, name: str) -> None: - marks = getattr(self._local, "marks", None) - if marks is None: - marks = {} - self._local.marks = marks - marks[name] = time.perf_counter() - - def toc(self, name: str) -> None: - marks = getattr(self._local, "marks", None) - if not marks: - return - t0 = marks.pop(name, None) - if t0 is None: - return - dt = time.perf_counter() - t0 - with self._lock: - self.acc[name] += dt - self.cnt[name] += 1 - total = self.cnt.get("em_total", 0) - if total and total - self._flushed >= _FLUSH_EVERY: - self._flushed = total - self._dump_locked() - - def _dump_locked(self) -> None: - parts = [] - for name in sorted(self.acc): - n = self.cnt[name] - if n: - parts.append(f"{name}={self.acc[name] / n * 1e3:.3f}ms") - sys.stderr.write("[STEPPROF] " + " ".join(parts) + "\n") - sys.stderr.flush() - - -_PROF = _StepProf() - - -def tic(name: str) -> None: - if _ENABLED: - _PROF.tic(name) - - -def toc(name: str) -> None: - if _ENABLED: - _PROF.toc(name) - - -@contextmanager -def span(name: str): - if not _ENABLED: - yield - return - _PROF.tic(name) - try: - yield - finally: - _PROF.toc(name) From f33256e315c6ae8e19e3d84067378b2eeda9784c Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 00:02:28 +0800 Subject: [PATCH 15/89] perf(minicpmo_4_5): let CFM mel bucketing take effect on NPU `cfm_graph_bucket_frames` (deploy YAML: 16) only applied when the CUDA-only `CFMGraphWrapper` existed, so on NPU the bucket width stayed 0: the CFM capture key space never collapsed (64 shapes) and, once the graph pool filled, later shapes permanently degraded to eager. - batched_token2wav.py: track "graphs requested" separately (`_cfm_graph_enabled`, from the deploy YAML's `enable_cfm_graph`) and key bucketing/padding off "CUDA wrapper present OR graphs requested". The CUDA path is unchanged: the wrapper is built whenever the config is enabled. - platforms/npu/models/minicpmo_4_5_code2wav.py: stop routing calls that carry an attention mask to eager; only the ragged path (`valid_lengths`) stays eager. A bucketing pad mask has a shape fixed by the bucket, so it travels as a graph input (its presence is part of the capture key) like any other tensor. `attn_mask` is forwarded to the graphable body only when present, so graphable bodies written before this feature keep their old signature. Motivation: this is the memory prerequisite for the single-device two-replica profile; without it the pool fills and the second replica does not fit. Bucketing itself measures neutral for RTF -- it is a shape/memory lever. Signed-off-by: y-null --- .../models/minicpmo_4_5/batched_token2wav.py | 22 ++++++++---- .../npu/models/minicpmo_4_5_code2wav.py | 36 +++++++++++++------ 2 files changed, 42 insertions(+), 16 deletions(-) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index bec442282bc..203bac2dc99 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -309,7 +309,11 @@ def __init__( logger.info("HiFT CUDA Graph captured successfully") self._cfm_graph_wrapper: CFMGraphWrapper | None = None cfm_graph_cfg = dict(cfm_graph_config or {}) - if bool(cfm_graph_cfg.get("enabled", False)): + # Graph capture can be provided by the platform instead of the CUDA + # wrapper above (on NPU it is the platform graph runner), so keep the + # request flag separate: bucketing and padding key off this. + self._cfm_graph_enabled = bool(cfm_graph_cfg.get("enabled", False)) + if self._cfm_graph_enabled: flow_parameter = next(self.flow.parameters(), None) if flow_parameter is not None and flow_parameter.device.type == "cuda": estimator = self.flow.decoder.estimator @@ -327,7 +331,9 @@ def __init__( # chunk up to a multiple of this many frames so the graph cache key # space stays small (0 disables bucketing, e.g. when graphs are off). self._cfm_graph_bucket_frames = ( - int(cfm_graph_cfg.get("bucket_frames", 0)) if self._cfm_graph_wrapper is not None else 0 + int(cfm_graph_cfg.get("bucket_frames", 0)) + if (self._cfm_graph_wrapper is not None or self._cfm_graph_enabled) + else 0 ) if self._cfm_graph_bucket_frames > 1: logger.info( @@ -668,10 +674,14 @@ def _decode_cfm( bucket_frames=self._cfm_graph_bucket_frames, disabled=( valid_lengths is not None - or self._cfm_graph_wrapper is None - # `_disable` keeps the wrapper object alive, so check the flag - # too: padding under a disabled wrapper is pure overhead. - or not self._cfm_graph_wrapper.enabled + # Padding only pays off while replay is active. On CUDA that is + # the wrapper (`_disable` keeps the object alive, so check its + # flag); on NPU there is no wrapper and the platform runner is + # keyed by `_cfm_graph_enabled`. + or not ( + (self._cfm_graph_wrapper is not None and self._cfm_graph_wrapper.enabled) + or (self._cfm_graph_wrapper is None and self._cfm_graph_enabled) + ) ), ) if pad_frames: diff --git a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py index 4346e83b67f..f2124128346 100644 --- a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py +++ b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py @@ -75,6 +75,7 @@ def _graphable_estimator_step( cond, cnn_cache, att_cache, + attn_mask=None, valid_frames=None, ): """Run the CFM estimator body after host-backed timestep embedding.""" @@ -91,7 +92,7 @@ def _graphable_estimator_step( result = estimator.blocks_forward_chunk( estimator_input, time_embedding, - None, + attn_mask, old_cnn, old_att, cnn_out, @@ -121,7 +122,9 @@ def _patched_estimator_step( graph_runner is None or self._trt_stepper is not None or self._cfm_graph_wrapper is not None - or attn_mask is not None + # The ragged (per-request length) path stays eager, but the padding + # mask produced by bucketing has a shape fixed by the bucket, so it + # replays as a graph input like any other tensor. or valid_lengths is not None ): return _original_estimator_step( @@ -149,12 +152,25 @@ def _patched_estimator_step( # The upstream embedder creates a frequency tensor on the host. Keep it # outside capture while retaining the tensor-only estimator body in graph. time_embedding = estimator.t_embedder(time).unsqueeze(1) + # A bucketing mask travels as a graph input: the capture key only covers + # shapes, so replaying an existing graph must copy the current mask in. + # Like `valid_frames`, it is forwarded only when present, so graphable + # bodies written before this feature keep working with the old signature. + has_mask = attn_mask is not None + mask_inputs = (attn_mask,) if has_mask else () + + def _step_kwargs(rest): + kwargs = dict(graphable_kwargs) + if has_mask: + kwargs["attn_mask"] = rest[0] + return kwargs + if cnn_cache is None: return graph_runner.run( "cfm_estimator", - (x, mu, time_embedding, speakers, cond), - (False,), - lambda step_x, step_mu, step_time, step_speakers, step_cond: _graphable_estimator_step( + (x, mu, time_embedding, speakers, cond, *mask_inputs), + ((False, has_mask) if has_mask else (False,)), + lambda step_x, step_mu, step_time, step_speakers, step_cond, *rest: _graphable_estimator_step( self, estimator, x=step_x, @@ -164,15 +180,15 @@ def _patched_estimator_step( cond=step_cond, cnn_cache=None, att_cache=None, - **graphable_kwargs, + **_step_kwargs(rest), ), ) return graph_runner.run( "cfm_estimator", - (x, mu, time_embedding, speakers, cond, cnn_cache, att_cache), - (True,), - lambda step_x, step_mu, step_time, step_speakers, step_cond, step_cnn, step_att: _graphable_estimator_step( + (x, mu, time_embedding, speakers, cond, cnn_cache, att_cache, *mask_inputs), + ((True, has_mask) if has_mask else (True,)), + lambda step_x, step_mu, step_time, step_speakers, step_cond, step_cnn, step_att, *rest: _graphable_estimator_step( self, estimator, x=step_x, @@ -182,7 +198,7 @@ def _patched_estimator_step( cond=step_cond, cnn_cache=step_cnn, att_cache=step_att, - **graphable_kwargs, + **_step_kwargs(rest), ), ) From 4e7497ee40d002213f9115e4e87e147002722eb1 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 00:02:28 +0800 Subject: [PATCH 16/89] refactor(minicpmo_4_5): configure the Talker multi-frame decode from the deploy YAML Make stage 1's speculative_config in the deploy YAML the single source of truth for the multi-frame decode, and keep this port to the optimizations plus their tests only. * the deploy YAML carries stage 1's speculative_config (8 codec frames per step) and async_scheduling: false explicitly; no code-side default remains. * the environment switches VLLM_OMNI_MINICPMO_TALKER_FRAMES, OMNI_K_STEP, VLLM_OMNI_C2W_SINGLE_ROW and VLLM_OMNI_MINICPMO_NARROW_REPLAY are gone, and the 910B SoC gate is dropped. * the model-side and scheduler-side gates derive K from the engine config (model_stage == tts plus the n-gram fingerprint), so the two cannot disagree. * the per-step/per-frame timers, the stop and bookkeeping traces, the guard view and utils/step_prof.py are removed: they were bring-up instrumentation, not optimizations. * the stage-2 single-row even-batch item is reverted, along with the four orphan comment lines it left in code2wav. * the 910C/A3 Triton warmup guard (ascend_warmup_patch.py and its install sites) is dropped: it is a platform adaptation, not an optimization, and the 910B baseline ran the stock warmup either way. Its only remaining consumer, native_rejection_sampler, now carries the SoC probe and K-step check inline. * tests follow the new deploy values. Signed-off-by: y-null --- vllm_omni/config/stage_config.py | 4 +- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 188 ++++++++++++++++++ .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 4 - .../platforms/npu/ascend_warmup_patch.py | 180 ----------------- .../platforms/npu/native_rejection_sampler.py | 53 +++-- vllm_omni/platforms/npu/platform.py | 6 - .../platforms/npu/worker/npu_ar_worker.py | 8 - .../npu/worker/npu_generation_worker.py | 5 - 8 files changed, 230 insertions(+), 218 deletions(-) create mode 100644 vllm_omni/deploy/minicpmo_4_5_prof.yaml delete mode 100644 vllm_omni/platforms/npu/ascend_warmup_patch.py diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 676e3c889c5..06cb4ac7827 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -5,7 +5,6 @@ from __future__ import annotations import functools -import os import re import warnings from collections.abc import Callable, Mapping @@ -829,8 +828,7 @@ def load_deploy_config(path: str | Path) -> DeployConfig: ): if name in raw_dict: kwargs[name] = raw_dict[name] - deploy = DeployConfig(**kwargs) - return deploy + return DeployConfig(**kwargs) class PlatformOverrides(NamedTuple): diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml new file mode 100644 index 00000000000..82bfbe6e0f2 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -0,0 +1,188 @@ +# MiniCPM-o 4.5 three-stage deploy profile: all three stages on one device, +# with stage 2 running two colocated replicas (the base deploy config runs a +# single stage-2 replica). Only that difference is spelled out here. +# MiniCPM-o 4.5 three-stage deploy on one large-memory GPU. +pipeline: minicpmo_4_5 +async_chunk: true +# Use minicpmo_4_5_turn.yaml to select the ordinary turn-based serving engine. +session_mode: duplex +# Keep the active window aligned with the session and stage capacities below. +# A smaller window serializes first-packet generation for concurrent chat turns. +duplex_session: + # Matches Stage 0/1 max_num_seqs below. + max_sessions: 4 + +trust_remote_code: true +enable_prefix_caching: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + # Steady-state payload length in codec frames. + codec_chunk_frames: 75 + # Smaller first payload so the first audio packet is emitted earlier. + # Takes effect only when initial < codec_chunk_frames; values far below + # the steady-state chunk length add a seam between chunks and hurt WER. + initial_codec_chunk_frames: 25 + codec_left_context_frames: 3 + # HiFT capture buckets are derived from codec_chunk_frames and + # codec_left_context_frames. (#5869) + # A future initial_codec_chunk_frames shape will use lazy capture unless + # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. + enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 + +stages: + - stage_id: 0 + max_num_seqs: 4 + # Leave headroom on the shared GPU for Talker KV + the two stage-2 + # replicas below; the NPU block caps stage 0/1 KV explicitly. + gpu_memory_utilization: 0.40 + enforce_eager: false + # Stage-level prefix caching override; the top level stays false for duplex safety. + # Stage-0 n-gram speculative decoding is intentionally not enabled here: vLLM + # rejects prefix caching together with speculative decoding. + enable_prefix_caching: true + tensor_parallel_size: 1 + max_num_batched_tokens: 16384 + devices: "0" + # Required for MiniCPM interleaved image/audio Daily-Omni packs (~78% recipe). + # Keep these in YAML so they are applied through --deploy-config. + interleave_mm_strings: true + media_io_kwargs: + video: + fps: 1 + num_frames: 128 + limit_mm_per_prompt: + # Daily-Omni: up to 64 interleaved image/audio pairs (1fps). + image: 64 + audio: 64 + video: 1 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: true + # 1.0 for MCQ accuracy; 1.1 can bias short A–D answers. + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 4 + # Preserve the checkpoint's full-attention Talker prompt while it fits; + # native duplex rolls over to a one-chunk window at context capacity. + # Leave device-level headroom for the colocated HiFi-GAN cuDNN workspace. + gpu_memory_utilization: 0.08 + enforce_eager: false + max_num_batched_tokens: 8192 + # Talker context is 4096 (tts_config.max_position_embeddings). + max_model_len: 4096 + devices: "0" + # Multi-frame decode: the Talker produces 8 codec frames per scheduler step. + # vLLM learns that each request advances 8 tokens through this + # speculative_config (the n-gram proposer is never consulted -- the runner + # supplies constant continue drafts), and async scheduling is rejected next + # to an n-gram proposer. stop_token_ids above must stay the two-wide + # continue/stop id the head emits in this mode. Remove the block (or set it + # to null) to fall back to one frame per step. + speculative_config: + method: ngram + num_speculative_tokens: 7 + prompt_lookup_min: 1 + prompt_lookup_max: 1 + async_scheduling: false + output_connectors: + to_stage_2: connector_of_shared_memory + # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler + # (head_code logits). EOS is tts_config.eos_token_id (num_audio_tokens-1). + # Values track upstream utils.TTSSamplingParams (min_p/win_size/tau_r are + # declared there but never read, so they stay out). repetition_penalty is + # consumed by the Talker, which rescores it over upstream's 16-frame window + # instead of the Sampler's whole-stream presence penalty. + # min_tokens is upstream's min_new_token. + default_sampling_params: + temperature: 0.8 + top_p: 0.85 + top_k: 25 + repetition_penalty: 1.05 + min_tokens: 50 + # The Talker's stop head emits this id to finish a request. + stop_token_ids: [1] + max_tokens: 4096 + seed: 42 + detokenize: false + + - stage_id: 2 + max_num_seqs: 4 + gpu_memory_utilization: 0.18 + enforce_eager: true + enable_chunked_prefill: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + # Two replicas on the same device: `devices` repeats the card once per + # replica (replica r lands on card r), so this needs no extra hardware. + num_replicas: 2 + devices: "0,0" + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + detokenize: true + +platforms: + cuda: + stages: + - stage_id: 1 + # Full-attention occupancy is bounded by max_num_seqs * max_model_len + # (4 * 4096 = 16384 tokens). Keep the reported CUDA KV token capacity + # above that bound because physical-preemption replay is unsupported. + # The explicit cap leaves device headroom for concurrent Stage 2. + kv_cache_memory_bytes: 2147483648 # 2 GiB + + npu: + stages: + - stage_id: 0 + # Keep the three stages' concurrency caps aligned so the top tier is + # not gated by a single stage. + max_num_seqs: 8 + max_num_batched_tokens: 8192 + # Explicit KV cap: the auto-computed value does not leave room for the + # two stage-2 replicas on this card. + kv_cache_memory_bytes: 6442450944 + compilation_config: + # A full decode-step graph keeps in-graph sampling and multi-frame + # expansion; PIECEWISE silently drops both. Measured on 910B3: + # -32.6% RTF versus PIECEWISE. + cudagraph_mode: FULL_DECODE_ONLY + - stage_id: 1 + # Concurrency cap raised from 4 to 8; the override block only covered + # stages 0 and 2, so stage 1 kept the base value of 4. + max_num_seqs: 8 + max_num_batched_tokens: 8192 + # Explicit KV cap: the auto-computed value does not leave room for the + # two stage-2 replicas on this card. + kv_cache_memory_bytes: 2684354560 + compilation_config: + cudagraph_mode: FULL_DECODE_ONLY + - stage_id: 2 + max_num_seqs: 8 + # Stage 2 stays outer-eager and captures only exact-shape CFM DiT + # estimator kernels through the NPU platform patch. + additional_config: + code2wav_enable_npu_graph: true + # One pool entry per captured shape. With two stage-2 replicas the pool + # is per replica, and the CFM mel bucketing already present in main + # (cfm_graph_bucket_frames: 16) keeps the captured shape count near 31, + # so 32 covers it; 64 per replica would only add graph memory that is + # never replayed. + code2wav_max_npu_graphs: 32 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index a7993dd3b4f..da2c0eb725d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -281,10 +281,6 @@ def __init__( raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be >= 0") if self._initial_batch_size and self._initial_batch_size < self._min_batch_size: raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be 0 or >= code2wav_min_batch_size") - # Force stage 2 to decode a single row per batch so CFM deterministically - # takes the even path (equal-length rows -> graph replay) instead of the - # ragged eager chain for mixed-length batches. Numerics are unchanged - # (even is an existing code path). self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") self._default_prompt_normalized: tuple[str, str] | None = None diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py deleted file mode 100644 index 0944ce1a0ab..00000000000 --- a/vllm_omni/platforms/npu/ascend_warmup_patch.py +++ /dev/null @@ -1,180 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Skip the vllm-ascend Triton warmups that fault the 910_93 vector core. - -vllm-ascend warms three Triton kernel families on the worker before graph -capture (``vllm_ascend/model_executor/warmup/kernel_warmup.py``). On 910C/A3 -the second one, ``penalties_triton_warmup``, faults inside its dummy-token -setup before any Triton kernel runs: - - tokens[:, -1:] = vocab_size - -lowers to ACLNN InplaceCopy over the legacy BroadcastTo kernel, and on -ascend910_93 that kernel raises a vector-core exception -- "The address for -the scalar to access the internal buffer of AICore is out of bounds", aclnn -error 507035 -- which kills the stage worker during initialization and takes -the whole orchestrator down with it. - -The identical line runs clean on 910B3 under the same CANN 9.1.0 and -torch_npu 2.10.0.post4 (verified directly, with the same shapes), so this is -a 910_93 operator fault rather than a shape, dtype or memory problem -- the -warmup input is only ``[max_num_seqs, 257]`` token ids. Rewriting the line -would merely relocate the fault into an operator nobody can verify on the -target box, so the warmup is skipped instead. - -Skipping is safe by construction: the warmup only pre-compiles Triton kernels -(JIT), and the sampling path runs those kernels over real token tensors, never -through this dummy construction. The cost is a one-off compile on the first -penalised sample of the first request. Correctness is unaffected. - -Scoped so the 910B baseline keeps the stock warmup byte-for-byte: the guards -delegate to the original warmup there. The skip set is resolved when a warmup -actually runs rather than when the platform is constructed, because the -constructor can run before the device is set -- and a probe that failed there -would silently strip a warmup the baseline relies on. A probe that still -cannot name the SoC skips the faulting warmup: a missing warmup costs latency, -a faulting one costs the whole run. -``VLLM_OMNI_NPU_SKIP_WARMUPS`` overrides the set (comma separated substrings -of the warmup names, or "none" to restore stock behaviour). - -``rejection_sampler_triton_warmup`` joined the default skip set after the A3 -run of 02:38 pinned the acl 507035 aivec fault on it: the dummy -(batch, spec_len, vocab) sweep "completed" on the host half a second before -the pre-capture synchronize raised, and nothing else was enqueued in between -(penalties was already skipped, rms is a no-op there). Skipping is lossless -the same way: on this stack no 910C runtime path launches the reject_sample -kernels -- the Talker's K-step sample() and its SoC guards fall back to -torch argmax, so the warmup only ever compiled kernels nobody runs. -""" - -from __future__ import annotations - -import os -from collections.abc import Callable -from typing import Any - -from vllm.logger import init_logger - -logger = init_logger(__name__) - -_PATCHED = False -_ENV = "VLLM_OMNI_NPU_SKIP_WARMUPS" -# Substrings matched against the warmup function names in kernel_warmup.py. -_DEFAULT_SKIP = ("penalties", "rejection_sampler") -_WARMUP_NAMES = ( - "rejection_sampler_triton_warmup", - "penalties_triton_warmup", - "triton_rms_warmup", -) - - -def _probe_soc_name() -> str: - """torch_npu device name, or "" when it cannot be read yet.""" - try: - import torch_npu - - try: - device = torch_npu.npu.current_device() - except Exception: - device = 0 - return str(torch_npu.npu.get_device_name(device)) - except Exception: - return "" - - -def _kstep_armed() -> bool: - """Whether this worker runs the Talker multi-frame decode. - - Read off the engine's speculative_config, which the deploy YAML's stage-1 - block provides; an unreadable config counts as not armed so the baseline - warmup never changes by accident. - """ - try: - from vllm.config import get_current_vllm_config_or_none - - cfg = get_current_vllm_config_or_none() - except Exception: - return False - spec = getattr(cfg, "speculative_config", None) if cfg is not None else None - if spec is None: - return False - method = getattr(spec, "method", None) - num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 - return method == "ngram" and num_spec > 0 - - -def _skipped_names() -> set[str]: - raw = os.environ.get(_ENV, "").strip().lower() - if raw: - if raw in ("none", "off", "0", "false", "no"): - return set() - return {part.strip() for part in raw.split(",") if part.strip()} - soc = _probe_soc_name() - if soc.startswith("Ascend910B"): - if _kstep_armed(): - # The torch-native sampler replaces the reject kernels whenever - # the K-step is armed, so this warmup would only compile kernels - # nobody runs. Everything else keeps the stock warmup. - return {"rejection_sampler"} - # Verified clean here; keep the stock warmup untouched. - return set() - return set(_DEFAULT_SKIP) - - -def _make_guard(name: str, original: Callable[[Any], Any]) -> Callable[[Any], Any]: - """Run the real warmup unless this call is on a host that must skip it. - - The decision is taken per call, not at install time: the platform - constructor can run before the device is set, and a probe that fails there - would otherwise silently skip a warmup the 910B baseline relies on. - """ - - def _guarded(worker: Any) -> Any: - if any(fragment in name for fragment in _skipped_names()): - logger.info( - "[npu] %s skipped: these Triton warmups fault the 910_93 " - "vector core during their dummy-token setup (acl 507035), " - "and no runtime path on this stack launches the kernels", - name, - ) - return None - return original(worker) - - _guarded._vllm_omni_guarded = True # type: ignore[attr-defined] - return _guarded - - -def apply_ascend_warmup_patch() -> None: - """Guard the ascend warmups so a faulting one can be skipped per call.""" - global _PATCHED - if _PATCHED: - return - _PATCHED = True - - try: - import importlib - - # Import the module, not the function: ``vllm_ascend.model_executor.warmup`` - # re-exports ``kernel_warmup`` as a package attribute, so - # ``from ... import kernel_warmup`` yields the function and every - # ``getattr`` below would silently miss. - warmup_module = importlib.import_module("vllm_ascend.model_executor.warmup.kernel_warmup") - except Exception as error: # pragma: no cover - non-ascend builds - logger.warning("[npu] Ascend warmup patch not applied: %s", error) - return - - guarded: list[str] = [] - for name in _WARMUP_NAMES: - original = getattr(warmup_module, name, None) - if original is None or getattr(original, "_vllm_omni_guarded", False): - continue - setattr(warmup_module, name, _make_guard(name, original)) - guarded.append(name) - - if guarded: - logger.info( - "[npu] ascend Triton warmups guarded (%s); the skip set is resolved " - "per call, so the SoC probe sees a live device. Override with %s.", - ", ".join(guarded), - _ENV, - ) diff --git a/vllm_omni/platforms/npu/native_rejection_sampler.py b/vllm_omni/platforms/npu/native_rejection_sampler.py index b5b881f281a..b76252d308c 100644 --- a/vllm_omni/platforms/npu/native_rejection_sampler.py +++ b/vllm_omni/platforms/npu/native_rejection_sampler.py @@ -5,20 +5,17 @@ vllm-ascend module-patches vllm's ``rejection_sample`` (and its two helpers) onto Triton kernels (``vllm_ascend/patch/worker/patch_rejection_sampler.py``). -On the 910_93 vector core those kernels fault: the 02:38 run pinned acl -507035 on their warmup, which is why ``rejection_sampler`` sits in the -warmup skip set. The skip moves the fault to the first request -- every -verify then JIT-compiles the kernel family on the fly, and both stages -stall there: stage 0 verifies the n-gram drafts of every decode step, and -stage 1 verifies the always-``continue`` drafts behind the K-step -block-table growth. +On the 910_93 vector core those kernels fault (acl 507035). Left alone, every +verify JIT-compiles the kernel family on the fly and both stages stall there: +stage 0 verifies the n-gram drafts of every decode step, and stage 1 verifies +the always-``continue`` drafts behind the K-step block-table growth. The torch-native implementation the patch replaces is pure tensor ops and exact for both callers, so this module reloads vllm's own sampler source into a fresh module object and re-points the three names at the pre-patch -functions. Scoped to the same SoC that skips the warmup; the 910B baseline -keeps the Triton kernels it has been measured with -- except when the Talker -K-step is explicitly armed there. That is the one path whose spec width the +functions. Scope: the 910_93 family always; the 910B baseline keeps the +Triton kernels it has been measured with -- except when the Talker K-step is +explicitly armed there. That is the one path whose spec width the 910B kernels cannot take (aivec 507035 on every one of the 8 combos tried), so the restore on a 910B is tied to exactly that arming, never to the SoC alone: the stock baseline (no arming env) keeps its sampler byte-for-byte. @@ -27,10 +24,42 @@ import importlib.util import logging -from vllm_omni.platforms.npu.ascend_warmup_patch import _kstep_armed, _probe_soc_name - logger = logging.getLogger(__name__) + +def _probe_soc_name() -> str: + """torch_npu device name, or "" when it cannot be read yet.""" + try: + import torch_npu + + try: + device = torch_npu.npu.current_device() + except Exception: + device = 0 + return str(torch_npu.npu.get_device_name(device)) + except Exception: + return "" + + +def _kstep_armed() -> bool: + """Whether this worker runs the Talker multi-frame decode. + + Read off the engine's speculative_config, which the deploy YAML's stage-1 + block provides; an unreadable config counts as not armed. + """ + try: + from vllm.config import get_current_vllm_config_or_none + + cfg = get_current_vllm_config_or_none() + except Exception: + return False + spec = getattr(cfg, "speculative_config", None) if cfg is not None else None + if spec is None: + return False + method = getattr(spec, "method", None) + num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 + return method == "ngram" and num_spec > 0 + _RESTORED = False _RESTORE_SOC_PREFIXES = ("ascend910_93", "ascend910c") # See the module docstring: on these parts the restore applies only when the diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index c0221f94e64..2c8d42f6093 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -59,9 +59,6 @@ def __init__(self) -> None: from vllm_ascend.utils import adapt_patch from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches - from vllm_omni.platforms.npu.ascend_warmup_patch import ( - apply_ascend_warmup_patch, - ) from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import ( apply_minicpmo_4_5_code2wav_patch, ) @@ -77,9 +74,6 @@ def __init__(self) -> None: apply_qwen3_tts_patches() apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches() - # The Triton rejection/penalties warmup guard for multi-frame decode must - # be installed before the Ascend kernel warmup runs. - apply_ascend_warmup_patch() @classmethod def set_device(cls, device: torch.device) -> None: diff --git a/vllm_omni/platforms/npu/worker/npu_ar_worker.py b/vllm_omni/platforms/npu/worker/npu_ar_worker.py index 6a0b5424755..19f3f89c8f6 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_worker.py @@ -3,18 +3,10 @@ from vllm.v1.worker.workspace import init_workspace_manager -from vllm_omni.platforms.npu.ascend_warmup_patch import apply_ascend_warmup_patch from vllm_omni.platforms.npu.worker.base import OmniNPUWorkerBase from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin -# Install at import time: this module is imported while the stage worker is -# built, which is earlier than the ascend Triton kernel warmup (that runs -# after load_model in the same process). The guard resolves the SoC per call, -# so installing this early is safe -- the probe only has to be right when a -# warmup actually runs. -apply_ascend_warmup_patch() - class NPUARWorker(OmniWorkerMixin, OmniNPUWorkerBase): """NPU AR worker for thinker/talker stages in Omni model.""" diff --git a/vllm_omni/platforms/npu/worker/npu_generation_worker.py b/vllm_omni/platforms/npu/worker/npu_generation_worker.py index fad66915df5..ac66c6a54fb 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_worker.py @@ -3,15 +3,10 @@ from vllm.v1.worker.workspace import init_workspace_manager -from vllm_omni.platforms.npu.ascend_warmup_patch import apply_ascend_warmup_patch from vllm_omni.platforms.npu.worker.base import OmniNPUWorkerBase from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin -# See npu_ar_worker: install before the worker is used, so the guard is in -# place when the ascend Triton kernel warmup runs later in this process. -apply_ascend_warmup_patch() - class NPUGenerationWorker(OmniWorkerMixin, OmniNPUWorkerBase): """NPU generation worker for code2wav stage in Omni model.""" From e6e06e4d6bb7b76335d8cfc9446ec37b5206f2ad Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 00:17:48 +0800 Subject: [PATCH 17/89] test(minicpmo_4_5): cover a bucketing pad mask replaying in-graph The bucketing change lets a pad mask travel as a graph input; add the regression test that pins the contract: a call carrying an attention mask still dispatches to the platform graph runner (only the ragged per-request-length path stays eager), the mask is one of the captured inputs, and its presence is part of the capture key (constants == (with_cache, True)). Also drop one stray blank line so the NPU code2wav file stays byte-identical with the validated branch. Signed-off-by: y-null --- tests/platforms/npu/test_graph_tools.py | 59 +++++++++++++++++++ .../npu/models/minicpmo_4_5_code2wav.py | 1 - 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/tests/platforms/npu/test_graph_tools.py b/tests/platforms/npu/test_graph_tools.py index 703b8f63060..faea8457919 100644 --- a/tests/platforms/npu/test_graph_tools.py +++ b/tests/platforms/npu/test_graph_tools.py @@ -373,3 +373,62 @@ def __exit__(self, exc_type, exc, traceback): last_chunk=True, ) == ("decode", True, False) assert entered == [(torch.device("cpu"), True), (torch.device("cpu"), True)] + + +def test_code2wav_pad_mask_replays_as_a_graph_input(monkeypatch): + """A bucketing pad mask is shape-stable, so it replays inside the graph.""" + calls = [] + seen = [] + + class _Backend: + _trt_stepper = None + _cfm_graph_wrapper = None + + class _Runner: + def run(self, operation, inputs, constants, compute): + calls.append((operation, inputs, constants)) + return compute(*inputs) + + backend = _Backend() + code2wav_patch._backend_graph_runners[backend] = _Runner() + + def graphable_estimator_step( + instance, + estimator, + *, + x, + mu, + time_embedding, + speakers, + cond, + cnn_cache, + att_cache, + attn_mask=None, + ): + del instance, estimator, time_embedding, speakers, cond, cnn_cache, att_cache + seen.append(attn_mask) + return x + mu, x.clone(), x.clone() + + monkeypatch.setattr(code2wav_patch, "_original_estimator_step", lambda *args, **kwargs: None) + monkeypatch.setattr(code2wav_patch, "_graphable_estimator_step", graphable_estimator_step) + value = torch.tensor([2.0]) + mask = torch.tensor([1.0, 0.0]) + estimator = SimpleNamespace(t_embedder=lambda time: time + 1) + + code2wav_patch._patched_estimator_step( + backend, + estimator, + x=value, + mu=value, + time=value, + speakers=value, + cond=value, + attn_mask=mask, + cnn_cache=None, + att_cache=None, + ) + + assert calls[0][0] == "cfm_estimator" + assert calls[0][2] == (False, True) + assert any(item is mask for item in calls[0][1]) + assert seen[0] is mask diff --git a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py index f2124128346..df519f3f541 100644 --- a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py +++ b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py @@ -164,7 +164,6 @@ def _step_kwargs(rest): if has_mask: kwargs["attn_mask"] = rest[0] return kwargs - if cnn_cache is None: return graph_runner.run( "cfm_estimator", From e94578a91b6eefc049b898a0e70a1e090e2930a6 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 00:27:52 +0800 Subject: [PATCH 18/89] revert(npu): restore the 910C/A3 Triton warmup guard Removing ascend_warmup_patch.py makes the service fail on 910C/A3 (confirmed by testing): it is a required platform adaptation, not dead weight. Restore the module and its three install sites, and point native_rejection_sampler back at its helpers instead of the inlined copies. This reverts the platform part of 58e513a9. Signed-off-by: y-null --- .../platforms/npu/ascend_warmup_patch.py | 180 ++++++++++++++++++ .../platforms/npu/native_rejection_sampler.py | 53 ++---- vllm_omni/platforms/npu/platform.py | 6 + .../platforms/npu/worker/npu_ar_worker.py | 8 + .../npu/worker/npu_generation_worker.py | 5 + 5 files changed, 211 insertions(+), 41 deletions(-) create mode 100644 vllm_omni/platforms/npu/ascend_warmup_patch.py diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py new file mode 100644 index 00000000000..0944ce1a0ab --- /dev/null +++ b/vllm_omni/platforms/npu/ascend_warmup_patch.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Skip the vllm-ascend Triton warmups that fault the 910_93 vector core. + +vllm-ascend warms three Triton kernel families on the worker before graph +capture (``vllm_ascend/model_executor/warmup/kernel_warmup.py``). On 910C/A3 +the second one, ``penalties_triton_warmup``, faults inside its dummy-token +setup before any Triton kernel runs: + + tokens[:, -1:] = vocab_size + +lowers to ACLNN InplaceCopy over the legacy BroadcastTo kernel, and on +ascend910_93 that kernel raises a vector-core exception -- "The address for +the scalar to access the internal buffer of AICore is out of bounds", aclnn +error 507035 -- which kills the stage worker during initialization and takes +the whole orchestrator down with it. + +The identical line runs clean on 910B3 under the same CANN 9.1.0 and +torch_npu 2.10.0.post4 (verified directly, with the same shapes), so this is +a 910_93 operator fault rather than a shape, dtype or memory problem -- the +warmup input is only ``[max_num_seqs, 257]`` token ids. Rewriting the line +would merely relocate the fault into an operator nobody can verify on the +target box, so the warmup is skipped instead. + +Skipping is safe by construction: the warmup only pre-compiles Triton kernels +(JIT), and the sampling path runs those kernels over real token tensors, never +through this dummy construction. The cost is a one-off compile on the first +penalised sample of the first request. Correctness is unaffected. + +Scoped so the 910B baseline keeps the stock warmup byte-for-byte: the guards +delegate to the original warmup there. The skip set is resolved when a warmup +actually runs rather than when the platform is constructed, because the +constructor can run before the device is set -- and a probe that failed there +would silently strip a warmup the baseline relies on. A probe that still +cannot name the SoC skips the faulting warmup: a missing warmup costs latency, +a faulting one costs the whole run. +``VLLM_OMNI_NPU_SKIP_WARMUPS`` overrides the set (comma separated substrings +of the warmup names, or "none" to restore stock behaviour). + +``rejection_sampler_triton_warmup`` joined the default skip set after the A3 +run of 02:38 pinned the acl 507035 aivec fault on it: the dummy +(batch, spec_len, vocab) sweep "completed" on the host half a second before +the pre-capture synchronize raised, and nothing else was enqueued in between +(penalties was already skipped, rms is a no-op there). Skipping is lossless +the same way: on this stack no 910C runtime path launches the reject_sample +kernels -- the Talker's K-step sample() and its SoC guards fall back to +torch argmax, so the warmup only ever compiled kernels nobody runs. +""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_PATCHED = False +_ENV = "VLLM_OMNI_NPU_SKIP_WARMUPS" +# Substrings matched against the warmup function names in kernel_warmup.py. +_DEFAULT_SKIP = ("penalties", "rejection_sampler") +_WARMUP_NAMES = ( + "rejection_sampler_triton_warmup", + "penalties_triton_warmup", + "triton_rms_warmup", +) + + +def _probe_soc_name() -> str: + """torch_npu device name, or "" when it cannot be read yet.""" + try: + import torch_npu + + try: + device = torch_npu.npu.current_device() + except Exception: + device = 0 + return str(torch_npu.npu.get_device_name(device)) + except Exception: + return "" + + +def _kstep_armed() -> bool: + """Whether this worker runs the Talker multi-frame decode. + + Read off the engine's speculative_config, which the deploy YAML's stage-1 + block provides; an unreadable config counts as not armed so the baseline + warmup never changes by accident. + """ + try: + from vllm.config import get_current_vllm_config_or_none + + cfg = get_current_vllm_config_or_none() + except Exception: + return False + spec = getattr(cfg, "speculative_config", None) if cfg is not None else None + if spec is None: + return False + method = getattr(spec, "method", None) + num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 + return method == "ngram" and num_spec > 0 + + +def _skipped_names() -> set[str]: + raw = os.environ.get(_ENV, "").strip().lower() + if raw: + if raw in ("none", "off", "0", "false", "no"): + return set() + return {part.strip() for part in raw.split(",") if part.strip()} + soc = _probe_soc_name() + if soc.startswith("Ascend910B"): + if _kstep_armed(): + # The torch-native sampler replaces the reject kernels whenever + # the K-step is armed, so this warmup would only compile kernels + # nobody runs. Everything else keeps the stock warmup. + return {"rejection_sampler"} + # Verified clean here; keep the stock warmup untouched. + return set() + return set(_DEFAULT_SKIP) + + +def _make_guard(name: str, original: Callable[[Any], Any]) -> Callable[[Any], Any]: + """Run the real warmup unless this call is on a host that must skip it. + + The decision is taken per call, not at install time: the platform + constructor can run before the device is set, and a probe that fails there + would otherwise silently skip a warmup the 910B baseline relies on. + """ + + def _guarded(worker: Any) -> Any: + if any(fragment in name for fragment in _skipped_names()): + logger.info( + "[npu] %s skipped: these Triton warmups fault the 910_93 " + "vector core during their dummy-token setup (acl 507035), " + "and no runtime path on this stack launches the kernels", + name, + ) + return None + return original(worker) + + _guarded._vllm_omni_guarded = True # type: ignore[attr-defined] + return _guarded + + +def apply_ascend_warmup_patch() -> None: + """Guard the ascend warmups so a faulting one can be skipped per call.""" + global _PATCHED + if _PATCHED: + return + _PATCHED = True + + try: + import importlib + + # Import the module, not the function: ``vllm_ascend.model_executor.warmup`` + # re-exports ``kernel_warmup`` as a package attribute, so + # ``from ... import kernel_warmup`` yields the function and every + # ``getattr`` below would silently miss. + warmup_module = importlib.import_module("vllm_ascend.model_executor.warmup.kernel_warmup") + except Exception as error: # pragma: no cover - non-ascend builds + logger.warning("[npu] Ascend warmup patch not applied: %s", error) + return + + guarded: list[str] = [] + for name in _WARMUP_NAMES: + original = getattr(warmup_module, name, None) + if original is None or getattr(original, "_vllm_omni_guarded", False): + continue + setattr(warmup_module, name, _make_guard(name, original)) + guarded.append(name) + + if guarded: + logger.info( + "[npu] ascend Triton warmups guarded (%s); the skip set is resolved " + "per call, so the SoC probe sees a live device. Override with %s.", + ", ".join(guarded), + _ENV, + ) diff --git a/vllm_omni/platforms/npu/native_rejection_sampler.py b/vllm_omni/platforms/npu/native_rejection_sampler.py index b76252d308c..b5b881f281a 100644 --- a/vllm_omni/platforms/npu/native_rejection_sampler.py +++ b/vllm_omni/platforms/npu/native_rejection_sampler.py @@ -5,17 +5,20 @@ vllm-ascend module-patches vllm's ``rejection_sample`` (and its two helpers) onto Triton kernels (``vllm_ascend/patch/worker/patch_rejection_sampler.py``). -On the 910_93 vector core those kernels fault (acl 507035). Left alone, every -verify JIT-compiles the kernel family on the fly and both stages stall there: -stage 0 verifies the n-gram drafts of every decode step, and stage 1 verifies -the always-``continue`` drafts behind the K-step block-table growth. +On the 910_93 vector core those kernels fault: the 02:38 run pinned acl +507035 on their warmup, which is why ``rejection_sampler`` sits in the +warmup skip set. The skip moves the fault to the first request -- every +verify then JIT-compiles the kernel family on the fly, and both stages +stall there: stage 0 verifies the n-gram drafts of every decode step, and +stage 1 verifies the always-``continue`` drafts behind the K-step +block-table growth. The torch-native implementation the patch replaces is pure tensor ops and exact for both callers, so this module reloads vllm's own sampler source into a fresh module object and re-points the three names at the pre-patch -functions. Scope: the 910_93 family always; the 910B baseline keeps the -Triton kernels it has been measured with -- except when the Talker K-step is -explicitly armed there. That is the one path whose spec width the +functions. Scoped to the same SoC that skips the warmup; the 910B baseline +keeps the Triton kernels it has been measured with -- except when the Talker +K-step is explicitly armed there. That is the one path whose spec width the 910B kernels cannot take (aivec 507035 on every one of the 8 combos tried), so the restore on a 910B is tied to exactly that arming, never to the SoC alone: the stock baseline (no arming env) keeps its sampler byte-for-byte. @@ -24,41 +27,9 @@ import importlib.util import logging -logger = logging.getLogger(__name__) - - -def _probe_soc_name() -> str: - """torch_npu device name, or "" when it cannot be read yet.""" - try: - import torch_npu - - try: - device = torch_npu.npu.current_device() - except Exception: - device = 0 - return str(torch_npu.npu.get_device_name(device)) - except Exception: - return "" +from vllm_omni.platforms.npu.ascend_warmup_patch import _kstep_armed, _probe_soc_name - -def _kstep_armed() -> bool: - """Whether this worker runs the Talker multi-frame decode. - - Read off the engine's speculative_config, which the deploy YAML's stage-1 - block provides; an unreadable config counts as not armed. - """ - try: - from vllm.config import get_current_vllm_config_or_none - - cfg = get_current_vllm_config_or_none() - except Exception: - return False - spec = getattr(cfg, "speculative_config", None) if cfg is not None else None - if spec is None: - return False - method = getattr(spec, "method", None) - num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 - return method == "ngram" and num_spec > 0 +logger = logging.getLogger(__name__) _RESTORED = False _RESTORE_SOC_PREFIXES = ("ascend910_93", "ascend910c") diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 2c8d42f6093..c0221f94e64 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -59,6 +59,9 @@ def __init__(self) -> None: from vllm_ascend.utils import adapt_patch from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches + from vllm_omni.platforms.npu.ascend_warmup_patch import ( + apply_ascend_warmup_patch, + ) from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import ( apply_minicpmo_4_5_code2wav_patch, ) @@ -74,6 +77,9 @@ def __init__(self) -> None: apply_qwen3_tts_patches() apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches() + # The Triton rejection/penalties warmup guard for multi-frame decode must + # be installed before the Ascend kernel warmup runs. + apply_ascend_warmup_patch() @classmethod def set_device(cls, device: torch.device) -> None: diff --git a/vllm_omni/platforms/npu/worker/npu_ar_worker.py b/vllm_omni/platforms/npu/worker/npu_ar_worker.py index 19f3f89c8f6..6a0b5424755 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_worker.py @@ -3,10 +3,18 @@ from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.platforms.npu.ascend_warmup_patch import apply_ascend_warmup_patch from vllm_omni.platforms.npu.worker.base import OmniNPUWorkerBase from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin +# Install at import time: this module is imported while the stage worker is +# built, which is earlier than the ascend Triton kernel warmup (that runs +# after load_model in the same process). The guard resolves the SoC per call, +# so installing this early is safe -- the probe only has to be right when a +# warmup actually runs. +apply_ascend_warmup_patch() + class NPUARWorker(OmniWorkerMixin, OmniNPUWorkerBase): """NPU AR worker for thinker/talker stages in Omni model.""" diff --git a/vllm_omni/platforms/npu/worker/npu_generation_worker.py b/vllm_omni/platforms/npu/worker/npu_generation_worker.py index ac66c6a54fb..fad66915df5 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_worker.py @@ -3,10 +3,15 @@ from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.platforms.npu.ascend_warmup_patch import apply_ascend_warmup_patch from vllm_omni.platforms.npu.worker.base import OmniNPUWorkerBase from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin +# See npu_ar_worker: install before the worker is used, so the guard is in +# place when the ascend Triton kernel warmup runs later in this process. +apply_ascend_warmup_patch() + class NPUGenerationWorker(OmniWorkerMixin, OmniNPUWorkerBase): """NPU generation worker for code2wav stage in Omni model.""" From 6426adf015570b78711de81979033916c70e68c4 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 00:28:01 +0800 Subject: [PATCH 19/89] feat(deploy): turn on stage-0 n-gram speculation in the two-replica profile The validated set measures stage-0 n-gram speculative decoding as part of the optimization set, but the profile generated from the base config inherited the base's deliberate omission of it (upstream fix fff0b0f8b, on the grounds that vLLM rejects prefix caching together with speculative decoding). That ground does not hold here: vLLM 0.29.0 in this container has no such mutual-exclusion check, and the validated branch ran exactly this pair -- prefix caching on with `async_scheduling: false` -- in both arms of today's 14-vs-15 A/B (two full three-tier runs). So the profile enables it with the base's 15 tokens (a 14-vs-15 A/B on 910B3 landed inside the run-to-run drift band: c1 -4.5%, c4 -7.5%, c8 +9.3% -> equivalent), while the base config keeps it off. With this, the profile differs from the base config only by the single-device dual-replica settings plus this speculation block. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 82bfbe6e0f2..ce7cd7b4df3 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -46,9 +46,19 @@ stages: gpu_memory_utilization: 0.40 enforce_eager: false # Stage-level prefix caching override; the top level stays false for duplex safety. - # Stage-0 n-gram speculative decoding is intentionally not enabled here: vLLM - # rejects prefix caching together with speculative decoding. enable_prefix_caching: true + # Stage-0 n-gram speculative decoding (part of the validated set). 15 is the + # base config's value: a 14-vs-15 A/B on 910B3 came out inside the run-to-run + # drift band (c1 -4.5%, c4 -7.5%, c8 +9.3%), so the two are equivalent. + # It coexists with prefix caching as long as async_scheduling stays false -- + # the validated set ran exactly this pair. (The base config keeps it off, + # following the upstream fix in fff0b0f8b.) + async_scheduling: false + speculative_config: + method: ngram + num_speculative_tokens: 15 + prompt_lookup_min: 1 + prompt_lookup_max: 1 tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" From 67bbd5bdb547369cf93fa69da7cb6d34fe7831ef Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 00:29:57 +0800 Subject: [PATCH 20/89] docs(deploy): the profile keeps stage-0 speculation at 15 because that is the tail-draft gate An earlier version of this comment said 14 and 15 were equivalent because the RTF difference sat near the run-to-run drift band. That reading was incomplete: the stage-0 tail-draft path is gated on `num_speculative_tokens == 15` (stage0_tail_draft.applies()), so k=14 switches it off silently. The 14-vs-15 A/B measured 22-37% worse TTFT in the k=14 arm, matching the tail-draft item's own account (TTFT -10%, RTF unchanged); RTF was inconsistent (-4.5% / -7.5% / +9.3%). The width itself is still untested in isolation. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index ce7cd7b4df3..2267f60312a 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -48,8 +48,12 @@ stages: # Stage-level prefix caching override; the top level stays false for duplex safety. enable_prefix_caching: true # Stage-0 n-gram speculative decoding (part of the validated set). 15 is the - # base config's value: a 14-vs-15 A/B on 910B3 came out inside the run-to-run - # drift band (c1 -4.5%, c4 -7.5%, c8 +9.3%), so the two are equivalent. + # base config's value and is also what `stage0_tail_draft.applies()` keys on + # (`num_speculative_tokens == 15`): any other width silently switches the + # stage-0 tail-draft path off, which measured 22-37% worse TTFT on 910B3. + # A 14-vs-15 A/B there moved RTF by -4.5% / -7.5% / +9.3% (c1/c4/c8 -- + # inconsistent, near the drift band), so the width itself is not the point; + # keeping the tail draft enabled is. # It coexists with prefix caching as long as async_scheduling stays false -- # the validated set ran exactly this pair. (The base config keeps it off, # following the upstream fix in fff0b0f8b.) From 09a43c9cde524fa78a5581a10700b3c68b4bfcea Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 01:04:51 +0800 Subject: [PATCH 21/89] revert(deploy): drop the stage-0 n-gram block from the two-replica profile The omni prefix cache refuses the combination outright: stage_prefix_cache_config() -> check_prefix_cache_token_accounting() raises OmniPrefixCacheUnmatchError ('omni prefix caching is not supported together with speculative decoding') as soon as a stage has enable_prefix_caching and a speculative_config, and stage 0 of this profile has the former. The runner call sites (NPU initialize_kv_cache, GPU initialize_metadata_builders) are not wrapped in try/except, so the engine would fail at kv-cache init. The base config keeps stage 0 free of speculation for the same reason. Reverts ad6ff944f (which added the block) and d7de6d4ee (its comment). Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 18 ++---------------- 1 file changed, 2 insertions(+), 16 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 2267f60312a..82bfbe6e0f2 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -46,23 +46,9 @@ stages: gpu_memory_utilization: 0.40 enforce_eager: false # Stage-level prefix caching override; the top level stays false for duplex safety. + # Stage-0 n-gram speculative decoding is intentionally not enabled here: vLLM + # rejects prefix caching together with speculative decoding. enable_prefix_caching: true - # Stage-0 n-gram speculative decoding (part of the validated set). 15 is the - # base config's value and is also what `stage0_tail_draft.applies()` keys on - # (`num_speculative_tokens == 15`): any other width silently switches the - # stage-0 tail-draft path off, which measured 22-37% worse TTFT on 910B3. - # A 14-vs-15 A/B there moved RTF by -4.5% / -7.5% / +9.3% (c1/c4/c8 -- - # inconsistent, near the drift band), so the width itself is not the point; - # keeping the tail draft enabled is. - # It coexists with prefix caching as long as async_scheduling stays false -- - # the validated set ran exactly this pair. (The base config keeps it off, - # following the upstream fix in fff0b0f8b.) - async_scheduling: false - speculative_config: - method: ngram - num_speculative_tokens: 15 - prompt_lookup_min: 1 - prompt_lookup_max: 1 tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" From 9bd47e0f8a284526b0121562aea0f04668c1ed65 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 01:08:31 +0800 Subject: [PATCH 22/89] docs(deploy): name the check that actually refuses prefix caching with speculation The comment blamed vLLM. It is not vLLM: prefix caching and speculative decoding are allowed together by vLLM itself (no such validation in the installed 0.29.0, and upstream main even has disable_eagle_block_drop for the interaction). The refusal is ours: stage_prefix_cache_config() -> check_prefix_cache_token_accounting() raises OmniPrefixCacheUnmatchError whenever a stage has enable_prefix_caching and any speculative_config -- regardless of async_scheduling and of the speculative method. Comment-only change; the configs are unchanged. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 7 +++++-- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 7eb0b14c184..05104f335dd 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -42,8 +42,11 @@ stages: gpu_memory_utilization: 0.55 enforce_eager: false # Stage-level prefix caching override; the top level stays false for duplex safety. - # Stage-0 n-gram speculative decoding is intentionally not enabled here: vLLM - # rejects prefix caching together with speculative decoding. + # Stage-0 n-gram speculative decoding is intentionally not enabled here: the + # omni prefix cache refuses the pair (stage_prefix_cache_config -> + # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError as soon + # as a stage has enable_prefix_caching and a speculative_config; vLLM itself + # allows the combination). enable_prefix_caching: true tensor_parallel_size: 1 max_num_batched_tokens: 16384 diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 82bfbe6e0f2..6d1949af723 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -46,8 +46,11 @@ stages: gpu_memory_utilization: 0.40 enforce_eager: false # Stage-level prefix caching override; the top level stays false for duplex safety. - # Stage-0 n-gram speculative decoding is intentionally not enabled here: vLLM - # rejects prefix caching together with speculative decoding. + # Stage-0 n-gram speculative decoding is intentionally not enabled here: the + # omni prefix cache refuses the pair (stage_prefix_cache_config -> + # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError as soon + # as a stage has enable_prefix_caching and a speculative_config; vLLM itself + # allows the combination). enable_prefix_caching: true tensor_parallel_size: 1 max_num_batched_tokens: 16384 From 9f58ff1715dae0dc0ce3f0085c0bb4b501ef27c9 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 01:48:13 +0800 Subject: [PATCH 23/89] fix(npu): migrate the fixed-KV decode backend the 910B K-step capture needs The deploy-config port dropped vllm_omni/platforms/npu/attention/ and its install sites, and with them the only decode-capture path the 910B family has. With the K-step configured, base and prof both fail in capture_model() with a vector-core exception (acl 507035) on 910B3, so the service never comes up; the same config on the pre-port tree reaches "Stage 1 engine startup completed", so the gap is the module, not the part. Restore the module and its three hooks: * platforms/npu/attention/{__init__,fixed_kv_backend,fixed_kv_decode,fixed_kv_prefill}.py -- whole. fixed_kv_prefill is imported by fixed_kv_backend, and narrow replay reads captured_seq_lens/captured_slot_mapping, so nothing in it is dead weight: it is the fixed-KV item's own delivery. * NPUOmniPlatform.get_attn_backend_cls: route short-context decoders onto OmniFixedKVAttentionBackend. * NPUARModelRunner: the enable_hamming_sparse default, install_into_ascend_backend() before parent init, and the per-capacity _capture_cudagraphs override. Signed-off-by: y-null --- vllm_omni/platforms/npu/attention/__init__.py | 2 + .../npu/attention/fixed_kv_backend.py | 426 +++++++++++++++ .../npu/attention/fixed_kv_decode.py | 502 ++++++++++++++++++ .../npu/attention/fixed_kv_prefill.py | 227 ++++++++ vllm_omni/platforms/npu/platform.py | 40 ++ .../npu/worker/npu_ar_model_runner.py | 41 ++ 6 files changed, 1238 insertions(+) create mode 100644 vllm_omni/platforms/npu/attention/__init__.py create mode 100644 vllm_omni/platforms/npu/attention/fixed_kv_backend.py create mode 100644 vllm_omni/platforms/npu/attention/fixed_kv_decode.py create mode 100644 vllm_omni/platforms/npu/attention/fixed_kv_prefill.py diff --git a/vllm_omni/platforms/npu/attention/__init__.py b/vllm_omni/platforms/npu/attention/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/platforms/npu/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_backend.py b/vllm_omni/platforms/npu/attention/fixed_kv_backend.py new file mode 100644 index 00000000000..dc2c37d12ed --- /dev/null +++ b/vllm_omni/platforms/npu/attention/fixed_kv_backend.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ascend attention backend that captures a decode step with constant host args. + +``FULL_DECODE_ONLY`` captures a whole decode step, but +``npu_fused_infer_attention_score`` takes ``actual_seq_lengths_kv`` as a host +``SymInt[]``, so its tiling is baked into the captured task. The KV length grows +by one every step, so vllm-ascend re-issues the op for every layer on every step +(``update_full_graph_params`` -> ``graph_task_update_begin``/``_end``). On the +20-layer, 768-hidden MiniCPM-o Talker that rebind *is* the step: py-spy puts 26% +of the stage-1 main-thread wall -- ~38% of its busy time -- inside it, and +removing it measured RTF 0.1896 -> 0.1599 on A3. +See ``server-env/design/R1_UNDER_FULL_DECODE_20260827.md``. + +The mechanism lives in :mod:`fixed_kv_decode`; this module is only the seam. It +is a subclass rather than a patch because every place vllm-ascend needs to be +told about it is a factory vLLM-Omni already controls: + + AscendAttentionBackend.get_impl_cls / get_builder_cls -> here + acl_graph.update_full_graph_params -> goes through + get_impl_cls() + NPUOmniPlatform.get_attn_backend_cls -> selects this + NPUARModelRunner._capture_cudagraphs -> captures buckets + ACLGraphWrapper.__call__ key + replay barrier -> wrapped by + :func:`fixed_kv_decode.install_into_ascend_aclgraph` + +That matters: the ranked evaluation installs only this tree and leaves +vllm-ascend as the image ships it, so anything that has to be applied *to* +vllm-ascend cannot reach a scored run. +""" + +from __future__ import annotations + +import torch +import torch_npu +from vllm.logger import init_logger + +from vllm_ascend.ascend_forward_context import _EXTRA_CTX +from vllm_ascend.attention.attention_v1 import ( + SWA_INT_MAX, + AscendAttentionBackend, + AscendAttentionBackendImpl, + AscendAttentionMetadataBuilder, + AscendAttentionState, +) + +from vllm_omni.platforms.npu.attention import fixed_kv_decode, fixed_kv_prefill + +logger = init_logger(__name__) + + +class OmniFixedKVMetadataBuilder(AscendAttentionMetadataBuilder): + """Attach the device-side state a captured step needs, and pick the bucket.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._omni_buckets = fixed_kv_decode.buckets_for( + fixed_kv_decode.capacity_for( + self.vllm_config.model_config.max_model_len, + self.vllm_config.cache_config.block_size, + ), + self.vllm_config.cache_config.block_size, + ) + spec = getattr(self.vllm_config, "speculative_config", None) + self._omni_max_q_len = 1 + if spec is not None and getattr(spec, "num_speculative_tokens", 0): + self._omni_max_q_len = 1 + int(spec.num_speculative_tokens) + + def build(self, *args, **kwargs): + attn_metadata = super().build(*args, **kwargs) + common = kwargs.get("common_attn_metadata") + if common is None: + # Positional call: (common_prefix_len, common_attn_metadata, ...). + common = next((a for a in args if hasattr(a, "block_table_tensor")), None) + if common is None: + fixed_kv_decode.set_runtime_bucket(None) + return attn_metadata + + # `seq_lens` on the metadata is a CPU tensor; a captured graph has to + # refresh its bias from the live device buffer instead. + attn_metadata.seq_lens_device = common.seq_lens + # build() materializes padded copies of seq_lens/block_table for the FIA + # batch-size check when the batch is full. Those are fresh tensors every + # step, so a graph cannot bake them in. The block table is the one that + # is not caught by the row check in `_fixed_kv_applies`; identity against + # the runtime's own persistent buffer is the reliable test. + attn_metadata.fia_inputs_persistent = attn_metadata.block_tables is common.block_table_tensor + attn_metadata.fixed_kv_mask_recorded = False + attn_metadata.fixed_kv_pse_recorded = False + + # Whether this step is a decode cannot be read off `attn_state`. A step + # that schedules several tokens per request comes back labelled + # `ChunkedPrefill`: vllm-ascend calls it `SpecDecoding` and then + # rewrites that to `ChunkedPrefill` for every proposer except `mtp` + # (`model_runner_v1._build_attn_state`). The shape is the honest test -- + # the same number of query positions for every request, none of them + # still consuming a prompt -- and it is what the captured graph and the + # causal bias actually require. + attn_metadata.fixed_kv_decode_step = ( + attn_metadata.attn_state == AscendAttentionState.DecodeOnly + or self._is_uniform_multi_query_decode(attn_metadata, common) + ) + + # Settle the capacity before the forward runs: it is part of the ACL + # graph dispatch key, so the wrapper has to be able to look it up. + if self._omni_buckets and attn_metadata.fixed_kv_decode_step: + seq_lens_list = attn_metadata.seq_lens_list or [1] + max_seq_len = getattr(common, "max_seq_len", None) or max(seq_lens_list) + fixed_kv_decode.set_runtime_bucket( + fixed_kv_decode.select_bucket(int(max_seq_len), self._omni_buckets) + ) + else: + fixed_kv_decode.set_runtime_bucket(None) + # A single fresh prefill inside a captured bucket replays a graph + # instead of paying ~18 ms of eager dispatch; the hook refuses + # everything else (chunked, batched, oversized, other stages). + fixed_kv_prefill.maybe_mark_step( + attn_metadata, + AscendAttentionState.PrefillNoCache, + int(self.vllm_config.cache_config.block_size), + ) + return attn_metadata + + def _is_uniform_multi_query_decode(self, attn_metadata, common) -> bool: + if self._omni_max_q_len <= 1: + return False + cumulative = getattr(attn_metadata, "actual_seq_lengths_q", None) + if not cumulative: + return False + rows = len(cumulative) + q_len = int(cumulative[0]) + if q_len <= 1 or q_len > self._omni_max_q_len: + return False + if int(cumulative[-1]) != rows * q_len: + return False + for index, value in enumerate(cumulative): + if int(value) != (index + 1) * q_len: + return False + is_prefilling = getattr(common, "is_prefilling", None) + if is_prefilling is None: + return False + num_reqs = int(getattr(common, "num_reqs", rows)) + return not bool(is_prefilling[:num_reqs].any()) + + +class OmniFixedKVAttentionBackendImpl(AscendAttentionBackendImpl): + """Decode attention whose host-side arguments never change between steps.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # vllm-ascend builds differ here: newer ones set these quant flags on + # the impl base class, the build this image ships does not. Default + # them so the constant-tiling path stays on the non-quant branch. + for _flag in ("enable_c8_quant", "enable_hamming_sparse"): + if not hasattr(self, _flag): + setattr(self, _flag, False) + self._omni_capacity = fixed_kv_decode.capacity_for( + self.vllm_config.model_config.max_model_len, + self.vllm_config.cache_config.block_size, + ) + # +2 matches the query_start_loc buffer, which carries the FIA padding + # request on a full batch. + self._omni_max_rows = int(self.vllm_config.scheduler_config.max_num_seqs) + 2 + self._omni_buckets = fixed_kv_decode.buckets_for( + self._omni_capacity, self.vllm_config.cache_config.block_size + ) + # A step that schedules several tokens per sequence (the Talker's + # multi-frame decode declares them through a speculative_config) is + # still a uniform decode, and the fixed-capacity mask covers it as + # long as the bias is causal among those queries too. + # Every length up to the declared one, because the scheduler may trim a + # step to fewer tokens than the config asks for, and a bias buffer that + # does not already exist cannot be allocated during capture -- it would + # come out of the graph's private pool. + spec = getattr(self.vllm_config, "speculative_config", None) + max_q_len = 1 + if spec is not None and getattr(spec, "num_speculative_tokens", 0): + max_q_len = 1 + int(spec.num_speculative_tokens) + self._omni_q_lens = tuple(range(1, max_q_len + 1)) + if self._omni_capacity: + # At model load, not during capture: an allocation made while a graph + # is capturing comes out of that graph's private pool. + fixed_kv_decode.prewarm( + self._omni_max_rows, + self.num_heads, + self._omni_buckets, + self.vllm_config.model_config.dtype, + torch.device("npu", torch.npu.current_device()), + self._omni_q_lens, + ) + + @staticmethod + def update_graph_params( + update_stream, + forward_context, + num_tokens, + vllm_config, + speculative_config=None, + # T2/910B compat: num_dcp_pcp_tokens was dropped from the upstream + # signature (vllm-ascend removed PCP from MRV1, see #12592). + draft_attn_metadatas=None, + ): + """Skip the per-step rebind for a step that captured no updatable tasks.""" + if fixed_kv_decode.is_fixed_graph(num_tokens): + metadata = forward_context.attn_metadata + first = next(iter(metadata.values())) if isinstance(metadata, dict) else metadata + fixed_kv_decode.sync_captured_inputs( + num_tokens, + getattr(first, "seq_lens_device", None), + getattr(first, "block_tables", None), + ) + return + AscendAttentionBackendImpl.update_graph_params( + update_stream, + forward_context, + num_tokens, + vllm_config, + speculative_config, + draft_attn_metadatas, + ) + + def _fixed_kv_applies(self, attn_metadata, block_table, block_size) -> bool: + """Whether this layer can capture a decode step with constant host args.""" + if not self._omni_capacity: + return False + if not getattr(attn_metadata, "fixed_kv_decode_step", False): + return False + if not attn_metadata.causal or not getattr(attn_metadata, "fia_inputs_persistent", False): + return False + if getattr(attn_metadata, "seq_lens_device", None) is None or block_table is None: + return False + # Every one of these steers the op onto a different argument set that the + # constant-tiling assumption does not cover. + if self.sinks is not None or self.sliding_window is not None: + return False + if self.enable_c8_quant or self.enable_hamming_sparse: + return False + if _EXTRA_CTX.is_draft_model: + return False + capacity = fixed_kv_decode.current_capacity() + if capacity is None or capacity not in self._omni_buckets: + return False + if block_size * block_table.shape[1] < capacity: + return False + # FIA takes its batch size from actual_seq_lengths, so the block table and + # the bias have to have exactly that many rows, not merely enough. + rows = len(attn_metadata.actual_seq_lengths_q) + if block_table.shape[0] != rows or attn_metadata.seq_lens_device.shape[0] != rows: + return False + if rows > self._omni_max_rows: + return False + # The bias is causal among the step's own queries as well as against + # the sequence length, so several queries per request are fine -- but + # only if every request contributed the *same* number, since the shift + # baked into the mask is one number for the whole step. + # actual_seq_lengths_q is a cumulative sum. + if self._uniform_query_len(attn_metadata) is None: + return False + return True + + def _uniform_query_len(self, attn_metadata) -> int | None: + """Queries per request when every request contributed the same count.""" + cumulative = attn_metadata.actual_seq_lengths_q + rows = len(cumulative) + if rows == 0: + return None + q_len = int(cumulative[0]) + if q_len <= 0 or q_len not in self._omni_q_lens: + return None + if int(cumulative[-1]) != rows * q_len: + return None + for index, value in enumerate(cumulative): + if int(value) != (index + 1) * q_len: + return None + return q_len + + def _fixed_kv_graph_fia( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata, + output: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + num_tokens: int, + ): + """Capture one attention layer with no updatable task attached to it. + + The general FIA serves this step; it needs its host arguments constant + to be replayable, so the op declares the full KV capacity and the live + sequence length arrives through ``pse_shift``, a device tensor the + graph refreshes for itself. + """ + capacity = fixed_kv_decode.current_capacity() + rows = len(attn_metadata.actual_seq_lengths_q) + q_len = self._uniform_query_len(attn_metadata) or 1 + pse_all = fixed_kv_decode.get_pse_buffer( + self._omni_max_rows, self.num_heads, capacity, query.dtype, query.device, q_len + ) + pse = pse_all[:rows] + + seq_lens_device = attn_metadata.seq_lens_device + + if not attn_metadata.fixed_kv_mask_recorded: + attn_metadata.fixed_kv_mask_recorded = True + fixed_kv_decode.remember_captured_inputs( + num_tokens, + capacity, + seq_lens_device, + block_table, + getattr(attn_metadata, "slot_mapping", None), + ) + logger.info( + "[minicpmo] fixed-KV decode capture: graph_size=%d rows=%d q_len=%d " + "kv_capacity=%d heads=%d attention=FIA", + num_tokens, + rows, + q_len, + capacity, + self.num_heads, + ) + + # The bias is per step, not per layer, and only the FIA path reads it, + # so it is recorded by the first layer that actually wants it. + if not attn_metadata.fixed_kv_pse_recorded: + fixed_kv_decode.emit_mask_refresh( + pse, seq_lens_device, self.num_heads, capacity, q_len + ) + attn_metadata.fixed_kv_pse_recorded = True + + softmax_lse = fixed_kv_decode.keep_alive( + num_tokens, capacity, torch.empty(1, dtype=query.dtype, device=query.device) + ) + fia_kwargs = dict( + query=query, + key=key, + value=value, + pse_shift=pse, + atten_mask=None, + block_table=block_table, + input_layout="TND", + block_size=block_size, + actual_seq_lengths=attn_metadata.actual_seq_lengths_q, + actual_seq_lengths_kv=[capacity] * rows, + num_key_value_heads=self.num_kv_heads, + num_heads=self.num_heads, + scale=self.scale, + sparse_mode=0, + pre_tokens=SWA_INT_MAX, + next_tokens=0, + ) + workspace = fixed_kv_decode.graph_workspace( + num_tokens, + capacity, + (rows, self.num_heads, self.num_kv_heads, self.head_size, block_size, capacity, query.dtype), + lambda: torch_npu._npu_fused_infer_attention_score_get_max_workspace(**fia_kwargs), + ) + torch_npu.npu_fused_infer_attention_score.out(workspace=workspace, out=[output, softmax_lse], **fia_kwargs) + return output.view(num_tokens, self.num_heads, self.head_size), num_tokens + + def _fixed_kv_prefill_fia(self, query, key, value, attn_metadata, output): + """Capture one prefill attention layer with all-constant host arguments. + + Runs only while a graph is being captured (and during its warmups): + the runtime replays the recorded tasks and never re-enters here. The + call mirrors the eager ``PrefillNoCache`` branch exactly -- causal FIA + against the singleton mask, KV length = query length, no block table -- + with the workspace and the softmax LSE held outside the graph pool so + the capture strands nothing (the fixed-KV workspace-leak lesson). + """ + bucket = int(attn_metadata.fixed_kv_prefill_bucket) + fia_kwargs = dict( + query=query[:bucket], + key=key[:bucket], + value=value[:bucket], + atten_mask=attn_metadata.attn_mask, + input_layout="TND", + block_size=128, + actual_seq_lengths=attn_metadata.actual_seq_lengths_q, + actual_seq_lengths_kv=attn_metadata.actual_seq_lengths_q, + num_key_value_heads=self.num_kv_heads, + num_heads=self.num_heads, + scale=self.scale, + sparse_mode=3, + ) + workspace = fixed_kv_prefill.workspace( + bucket, + lambda: torch_npu._npu_fused_infer_attention_score_get_max_workspace(**fia_kwargs), + ) + lse = fixed_kv_prefill.lse_buffer(bucket, query.dtype, query.device) + torch_npu.npu_fused_infer_attention_score.out(workspace=workspace, out=[output, lse], **fia_kwargs) + return output.view(bucket, self.num_heads, self.head_size), bucket + + def full_graph_fia(self, query, key, value, attn_metadata, output, layer=None): + if getattr(attn_metadata, "fixed_kv_prefill_bucket", None) is not None: + return self._fixed_kv_prefill_fia(query, key, value, attn_metadata, output) + if self._omni_capacity: + fixed_key, fixed_value, block_size, block_table, _ = self._get_fia_params(key, value, attn_metadata) + if self._fixed_kv_applies(attn_metadata, block_table, block_size): + return self._fixed_kv_graph_fia( + query, + fixed_key, + fixed_value, + attn_metadata, + output, + block_table, + block_size, + attn_metadata.actual_seq_lengths_q[-1], + ) + # `_get_fia_params` only reshapes cache views, so recomputing it in the + # base implementation costs nothing beyond two `view` calls, and only on + # the paths fixed-KV declines. + return super().full_graph_fia(query, key, value, attn_metadata, output, layer) + + +class OmniFixedKVAttentionBackend(AscendAttentionBackend): + """The stock Ascend backend with the fixed-KV decode impl and builder.""" + + @staticmethod + def get_impl_cls() -> type[OmniFixedKVAttentionBackendImpl]: + return OmniFixedKVAttentionBackendImpl + + @staticmethod + def get_builder_cls() -> type[OmniFixedKVMetadataBuilder]: + return OmniFixedKVMetadataBuilder diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py b/vllm_omni/platforms/npu/attention/fixed_kv_decode.py new file mode 100644 index 00000000000..fab34a28df4 --- /dev/null +++ b/vllm_omni/platforms/npu/attention/fixed_kv_decode.py @@ -0,0 +1,502 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# This file is a part of the vllm-ascend project. +# +"""Fixed-capacity decode attention: full-graph replay without per-step rebinding. + +Vendored from vllm-ascend (Apache-2.0) so that vLLM-Omni can deliver it. The +ranked evaluation installs only this tree -- ``pip install -e .`` over our +source into the official image -- and leaves vllm-ascend exactly as the image +ships it, so a patch to vllm-ascend cannot reach a scored run. Everything this +module needs from vllm-ascend is reachable through factories vLLM-Omni already +overrides, so the rest of the change lives in ``fixed_kv_backend.py`` and in +``NPUOmniPlatform``. + +``FULL_DECODE_ONLY`` captures a whole decode step, but +``npu_fused_infer_attention_score`` takes ``actual_seq_lengths_kv`` as a host +``SymInt[]``, so its tiling is baked into the captured task. Because the KV length +grows by one every step, vllm-ascend has to re-issue the op for every layer on +every step (``update_full_graph_params`` -> ``graph_task_update_begin/end``). On a +short-context decoder that rebind is the step: measured 3.68 ms per step for the +20-layer MiniCPM-o Talker on 910B3, against 1.21 ms for the replay itself. + +This module removes the rebind by making every host-side argument constant. The op +declares a *fixed* KV capacity, and the real sequence length is carried instead by +a device-resident ``pse_shift`` bias: ``0`` for valid positions and ``-inf`` past +the end of each sequence. The captured graph refreshes that bias itself from the +live ``seq_lens`` tensor, so a replay needs nothing from the host at all. + +The cost is that attention reads the whole declared capacity every step instead of +the live prefix. Capacity *buckets* keep that honest: several graphs are captured +per batch size, one per candidate capacity, and each step replays the smallest +bucket that covers it. With one bucket equal to the model's whole block table this +degenerates to a fixed capacity, which is only sensible for short contexts. + +Env switches: + ``VLLM_OMNI_FIXED_KV_DECODE`` 0 to disable (default 1) + ``VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN`` only engage at or below this + ``max_model_len`` (default 8192) + ``VLLM_OMNI_FIXED_KV_DECODE_BUCKETS`` comma-separated candidate capacities, + e.g. "1024,4096"; default "512". Each + model keeps the candidates below its own + block-table capacity and always adds that + capacity as the top bucket. Empty means a + single bucket: the full capacity. +""" + +import os +from contextlib import contextmanager + +import torch +from vllm.logger import logger + +# Padding rows of a cudagraph batch carry seq_len 0. A row whose bias is -inf +# everywhere makes softmax produce NaN, so clamp the mask to keep one live slot. +_MIN_SEQ_LEN = 1 + +_pse_cache: dict[tuple, torch.Tensor] = {} +_arange_cache: dict[tuple, torch.Tensor] = {} +_scalar_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} + +# Per (graph size, capacity), the tensors that capture baked in. Replay reads them +# directly, so the runtime metadata has to keep handing us the same buffers; see +# ``sync_captured_inputs``. +_captured_inputs: dict[tuple[int, int], tuple[torch.Tensor, torch.Tensor]] = {} +_captured_slot_mappings: dict[tuple[int, int], torch.Tensor] = {} + +# Scratch the captured tasks write into. Capture-time allocations come out of the +# graph's private pool, so anything that goes out of scope here can be recycled +# into the next capture while a live graph still points at it. +_graph_scratch: dict[tuple[int, int], list[torch.Tensor]] = {} + +# One FIA workspace per attention shape, shared by every layer and every graph +# that matches it. Keyed by shape rather than by graph so the measuring call +# never runs more than once per shape; see ``graph_workspace``. +_workspace_cache: dict[tuple, torch.Tensor] = {} + +_warned_moved_inputs = False + +# Which capacity the current step runs at. During capture a bucket is forced, so +# that one pass over the batch sizes produces the graphs for exactly that bucket. +_capture_bucket: int | None = None +_runtime_bucket: int | None = None + + +_BUCKETS_ENV = "VLLM_OMNI_FIXED_KV_DECODE_BUCKETS" +_DEFAULT_BUCKETS = "512" + + +def is_enabled() -> bool: + return os.getenv("VLLM_OMNI_FIXED_KV_DECODE", "1") == "1" + + +class _CapacityEntries(dict): + """The wrapper's entries dict, plus one dict per fixed-KV capacity. + + Sleep mode's ``reset_all_graph_params`` clears whatever dict is installed on + the wrapper, so the per-capacity dicts hang off it rather than off the + wrapper: otherwise a reset would leave stale graphs behind them. + """ + + def __init__(self, existing: dict | None = None) -> None: + super().__init__(existing or {}) + self.by_capacity: dict[int, dict] = {} + + def clear(self) -> None: + super().clear() + for entries in self.by_capacity.values(): + entries.clear() + self.by_capacity.clear() + + +def install_into_ascend_aclgraph() -> None: + """Give the image's ACLGraphWrapper a per-capacity key and drop a dead barrier. + + Two things the stock wrapper does that are wrong for fixed-KV decode: + + * It keys captured graphs on ``batch_descriptor`` alone. The KV capacity is + baked into the captured tasks, so one batch size holds one graph *per + bucket*. Without that, the second capture pass finds the first bucket's + entry, replays it instead of capturing, and hangs on the device + synchronize that closes the pass -- round ``20260827T170843Z``, exactly. + * It synchronizes the stream before every FULL replay, to order + ``update_attn_params`` against the previous one. This path issues no + updates, so the barrier is pure cost. ``enable_enpu`` gates that one line + and nothing else in ``__call__`` (stock ``8092d3f6`` line 264), which is + what lets us borrow it instead of copying ~140 lines. + + This wraps ``__call__`` rather than subclassing ``ACLGraphWrapper``, for the + same reason :func:`install_into_ascend_backend` patches the backend class: + ``current_platform`` is vllm-ascend's ``NPUPlatform``, so + ``NPUOmniPlatform.get_graph_wrapper_cls`` is never consulted. + + It deliberately does **not** touch ``vllm_ascend.attention.fixed_kv_decode``. + That module, and the ``graph_key()`` call site in the stock wrapper, exist + only where our own ``vllm-ascend-fixed-kv-decode.patch`` has been applied -- + on the A3 box, not in the ranked image. Reaching for it made this look + delivered while it was still leaning on the patch it exists to replace. + """ + from vllm_ascend.compilation import acl_graph + + wrapper_cls = acl_graph.ACLGraphWrapper + if getattr(wrapper_cls, "_omni_fixed_kv_call", False): + return + original_call = wrapper_cls.__call__ + + def _call(self, *args, **kwargs): + capacity = current_capacity() + if capacity is None: + return original_call(self, *args, **kwargs) + entries = self.concrete_aclgraph_entries + if not isinstance(entries, _CapacityEntries): + entries = _CapacityEntries(entries) + self.concrete_aclgraph_entries = entries + saved_entries = self.concrete_aclgraph_entries + saved_enpu = self.enable_enpu + self.concrete_aclgraph_entries = entries.by_capacity.setdefault(capacity, {}) + # Only once the capture has recorded this (size, capacity) is there a + # fixed graph to replay with nothing to order against; during the + # capture itself the barrier still has to stand. + if not saved_enpu and _is_fixed_graph_here(capacity): + self.enable_enpu = True + try: + return original_call(self, *args, **kwargs) + finally: + self.concrete_aclgraph_entries = saved_entries + self.enable_enpu = saved_enpu + + wrapper_cls.__call__ = _call + wrapper_cls._omni_fixed_kv_call = True + logger.info( + "[minicpmo] fixed-KV: ACLGraphWrapper keyed per KV capacity, pre-replay barrier dropped" + ) + + +def _is_fixed_graph_here(capacity: int) -> bool: + """``is_fixed_graph`` for the batch size the forward context is about to run.""" + from vllm.forward_context import get_forward_context + + descriptor = getattr(get_forward_context(), "batch_descriptor", None) + if descriptor is None: + return False + return (descriptor.num_tokens, capacity) in _captured_inputs + + +def install_into_ascend_backend() -> None: + """Put the omni impl/builder on the class vLLM's selector actually uses. + + ``current_platform.get_attn_backend_cls`` is vllm-ascend's ``NPUPlatform``, + not ``NPUOmniPlatform``. Round ``20260827T175303Z`` captured extra buckets + then ran the stock builder, so runtime looked up ``(batch, None)``. Patch + the stock backend class in this process instead. + """ + if not is_enabled(): + return + from vllm_ascend.attention.attention_v1 import AscendAttentionBackend + from vllm_omni.platforms.npu.attention.fixed_kv_backend import ( + OmniFixedKVAttentionBackendImpl, + OmniFixedKVMetadataBuilder, + ) + + if getattr(AscendAttentionBackend, "_omni_fixed_kv_installed", False): + return + AscendAttentionBackend.get_impl_cls = staticmethod(lambda: OmniFixedKVAttentionBackendImpl) + AscendAttentionBackend.get_builder_cls = staticmethod(lambda: OmniFixedKVMetadataBuilder) + AscendAttentionBackend._omni_fixed_kv_installed = True + install_into_ascend_aclgraph() + logger.info("[minicpmo] installed omni fixed-KV impl/builder onto AscendAttentionBackend") + + +def capacity_for(max_model_len: int, block_size: int) -> int: + """Full KV capacity this decoder would run at, or 0 when it should not engage.""" + if not is_enabled(): + return 0 + limit = int(os.getenv("VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN", "8192")) + if max_model_len > limit: + return 0 + # Round to the block table, not to max_model_len: the op indexes whole blocks. + return -(-int(max_model_len) // int(block_size)) * int(block_size) + + +def buckets_for(capacity: int, block_size: int) -> tuple[int, ...]: + """Capacities to capture graphs for, smallest first. + + Candidates at or above this model's own capacity are dropped: the full + capacity is always the top bucket and always covers the model. + """ + if not capacity: + return () + # The default has to be a code default, not an environment one: a ranked run + # sets no environment variables. Attention reads the declared capacity every + # step, so a bucket that just covers the sequence is worth having -- the + # Talker's is ~135 (a ~15-token condition plus ~118 codec frames), and 512 + # against 4096 alone measured RTF 0.1592 vs 0.1721 on A3, -7.5%. The full + # capacity is always kept as the top bucket, so a longer sequence still has + # a graph to land on. + raw = os.getenv(_BUCKETS_ENV, _DEFAULT_BUCKETS) + candidates = set() + for part in raw.split(","): + part = part.strip() + if not part: + continue + value = int(part) + if value <= 0: + continue + rounded = -(-value // int(block_size)) * int(block_size) + if rounded < capacity: + candidates.add(rounded) + candidates.add(int(capacity)) + return tuple(sorted(candidates)) + + +def select_bucket(max_seq_len: int, buckets: tuple[int, ...]) -> int: + """Smallest captured capacity that covers this step.""" + for bucket in buckets: + if max_seq_len <= bucket: + return bucket + return buckets[-1] + + +def set_runtime_bucket(capacity: int | None) -> None: + global _runtime_bucket + _runtime_bucket = capacity + + +@contextmanager +def capturing_bucket(capacity: int): + """Force one capacity for a whole capture pass over the batch sizes.""" + global _capture_bucket + previous = _capture_bucket + _capture_bucket = capacity + try: + yield + finally: + _capture_bucket = previous + + +def current_capacity() -> int | None: + return _capture_bucket if _capture_bucket is not None else _runtime_bucket + + +def graph_key() -> int | None: + """Extra ACL graph dispatch key, so one batch size can hold several buckets.""" + return current_capacity() if is_enabled() else None + + +def get_pse_buffer( + max_rows: int, + num_heads: int, + capacity: int, + dtype: torch.dtype, + device: torch.device, + q_len: int = 1, +) -> torch.Tensor: + key = (max_rows, num_heads, capacity, dtype, device, q_len) + buf = _pse_cache.get(key) + if buf is None: + buf = torch.zeros(max_rows, num_heads, q_len, capacity, dtype=dtype, device=device) + _pse_cache[key] = buf + return buf + + +def _get_arange(num_heads: int, capacity: int, device: torch.device, q_len: int = 1) -> torch.Tensor: + """``ar[.., i, j] = j - i`` -- the KV index measured from query ``i``. + + With one query per sequence this is a plain arange and the comparison + against the sequence length is the whole mask. With ``q_len`` queries the + shift is what makes the mask causal *among* them: query ``i`` of the step + owns KV position ``seq_len - q_len + i``, so it may read ``j`` only while + ``j - i < seq_len - q_len + 1``, which is one comparison against a single + per-row length again. + """ + key = (num_heads, capacity, device, q_len) + ar = _arange_cache.get(key) + if ar is None: + columns = torch.arange(capacity, dtype=torch.int32, device=device).view(1, 1, 1, capacity) + queries = torch.arange(q_len, dtype=torch.int32, device=device).view(1, 1, q_len, 1) + ar = (columns - queries).expand(1, num_heads, q_len, capacity).contiguous() + _arange_cache[key] = ar + return ar + + +def _get_scalars(dtype: torch.dtype, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + key = (dtype, device) + pair = _scalar_cache.get(key) + if pair is None: + pair = ( + torch.tensor(float("-inf"), dtype=dtype, device=device), + torch.tensor(0.0, dtype=dtype, device=device), + ) + _scalar_cache[key] = pair + return pair + + +def prewarm( + max_rows: int, + num_heads: int, + buckets: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + q_lens: tuple[int, ...] = (1,), +) -> None: + """Materialize every buffer the captured steps will reference. + + Called at model load. Anything allocated inside `torch.npu.graph` comes out of + that graph's private pool instead, so the buffers have to exist first. + """ + for capacity in buckets: + for q_len in q_lens: + get_pse_buffer(max_rows, num_heads, capacity, dtype, device, q_len) + _get_arange(num_heads, capacity, device, q_len) + _get_scalars(dtype, device) + + +def emit_mask_refresh( + pse: torch.Tensor, + seq_lens_device: torch.Tensor, + num_heads: int, + capacity: int, + q_len: int = 1, +) -> None: + """Record the bias refresh into the graph being captured. + + Runs once per captured step, ahead of the attention layers. On replay it + re-reads ``seq_lens_device`` in place, which is how a captured step learns + that the sequence grew without any host involvement. + + ``seq_lens`` counts the whole step, so with ``q_len`` queries the first of + them ends ``q_len - 1`` positions earlier; the shifted arange from + ``_get_arange`` turns that back into one comparison per row. + """ + ar = _get_arange(num_heads, capacity, pse.device, q_len) + neg, zero = _get_scalars(pse.dtype, pse.device) + lens = ( + (seq_lens_device - (q_len - 1)).clamp(min=_MIN_SEQ_LEN).to(torch.int32).view(-1, 1, 1, 1) + ) + torch.where(ar >= lens, neg, zero, out=pse) + + +def graph_workspace(num_tokens: int, capacity: int, shape_key: tuple, factory) -> torch.Tensor: + """Hand out one workspace per attention shape and keep it alive for good. + + Layers of one graph share a workspace because their shapes agree. That has + to be decided *before* asking for the workspace, not after: + ``_npu_fused_infer_attention_score_get_max_workspace`` allocates as it + measures, and an allocation made while a graph is capturing comes out of + that graph's private pool, where dropping the reference does not hand the + memory back. Measuring per layer therefore strands one workspace per layer + in the pool (Talker: 20 layers x 4 graphs x 106 MiB = 8.3 GiB). Measure + once per shape instead and give the same buffer to every layer that matches. + """ + ws = _workspace_cache.get(shape_key) + if ws is None: + ws = factory() + _workspace_cache[shape_key] = ws + return ws + + +def keep_alive(num_tokens: int, capacity: int, tensor: torch.Tensor) -> torch.Tensor: + _graph_scratch.setdefault((num_tokens, capacity), []).append(tensor) + return tensor + + +def remember_captured_inputs( + num_tokens: int, + capacity: int, + seq_lens_device: torch.Tensor, + block_table: torch.Tensor, + slot_mapping: torch.Tensor | None = None, +) -> None: + _captured_inputs[(num_tokens, capacity)] = (seq_lens_device, block_table) + if slot_mapping is not None: + _captured_slot_mappings[(num_tokens, capacity)] = slot_mapping + + +def captured_slot_mapping(num_tokens: int) -> torch.Tensor | None: + """The slot mapping this size's graph writes KV through, if captured. + + Usually the same persistent buffer the runtime fills, in which case a caller + that writes here writes exactly where the step would have. When it is not -- + the FIA batch-padding path materializes fresh tensors, which is why + :func:`sync_captured_inputs` exists -- this is the one the replay reads, and + the other one is not. + """ + capacity = current_capacity() + return _captured_slot_mappings.get((num_tokens, capacity)) + + +def captured_seq_lens(num_tokens: int) -> torch.Tensor | None: + """The device sequence-length buffer this size's graph reads, if captured. + + A replay takes its live sequence length from here (see the module + docstring), so a caller that advances the length itself -- the multi-frame + Talker loop, once per codec frame -- writes into this tensor rather than + rebuilding attention metadata. + """ + capacity = current_capacity() + captured = _captured_inputs.get((num_tokens, capacity)) + return None if captured is None else captured[0] + + +def is_fixed_graph(num_tokens: int) -> bool: + capacity = current_capacity() + return capacity is not None and (num_tokens, capacity) in _captured_inputs + + +def sync_captured_inputs(num_tokens: int, seq_lens_device: torch.Tensor, block_table: torch.Tensor) -> None: + """Keep replay reading live values when the runtime hands us new buffers. + + vLLM's ``seq_lens`` and block table are persistent buffers sliced from offset + 0, so in the common case these are the very tensors the capture baked in and + this is two pointer comparisons. If a build ever materializes a fresh tensor + (the FIA batch-padding path does), copy the live values into the captured + buffers rather than silently replaying stale ones. + """ + global _warned_moved_inputs + capacity = current_capacity() + captured = _captured_inputs.get((num_tokens, capacity)) + if captured is None: + return + cap_seq_lens, cap_block_table = captured + moved = False + if seq_lens_device is not None and seq_lens_device.data_ptr() != cap_seq_lens.data_ptr(): + rows = min(cap_seq_lens.shape[0], seq_lens_device.shape[0]) + cap_seq_lens[:rows].copy_(seq_lens_device[:rows]) + moved = True + if block_table is not None and block_table.data_ptr() != cap_block_table.data_ptr(): + rows = min(cap_block_table.shape[0], block_table.shape[0]) + cols = min(cap_block_table.shape[1], block_table.shape[1]) + cap_block_table[:rows, :cols].copy_(block_table[:rows, :cols]) + moved = True + if moved and not _warned_moved_inputs: + _warned_moved_inputs = True + logger.warning( + "Fixed-KV decode: attention metadata moved off the captured buffers; " + "copying seq_lens/block_table every step. Replay stays correct but " + "costs two extra device copies per step." + ) + + +def reset() -> None: + global _capture_bucket, _runtime_bucket + _graph_scratch.clear() + _workspace_cache.clear() + _captured_inputs.clear() + _captured_slot_mappings.clear() + _pse_cache.clear() + _arange_cache.clear() + _scalar_cache.clear() + _capture_bucket = None + _runtime_bucket = None diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py b/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py new file mode 100644 index 00000000000..9dc126dee49 --- /dev/null +++ b/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py @@ -0,0 +1,227 @@ +"""Captured graphs for the Talker's prefill, bucketed by token count. + +One Talker prefill forward is ~18.4 ms of host dispatch wrapped around 1-3 ms +of device work (`DEVICE_BOUND_FRAME_20260830.md` §4a': three back-to-back +forwards never accumulate more than 0.07 ms of drain), and the ranked case +runs 1.26 of them per request -- ~23 ms, all but a couple of milliseconds of it +removable. The decode side of this problem was solved by `fixed_kv_decode`; +this module is the same idea for the prefill shape. + +## Why a padded prefill replays exactly + +The ranked Talker prefill is a single fresh request of 9-25 tokens (measured +distribution; `{16, 32}` covers it). A graph captured at bucket ``B`` with +``actual_seq_lengths = [B]`` replays a real ``n <= B`` prefill bit-exactly for +the rows that matter: + +* **Attention.** The FIA call is causal (``sparse_mode=3`` against the + singleton 2048x2048 mask), so query row ``i`` attends to rows ``0..i`` only. + Real rows ``0..n-1`` therefore never see a padding row; padding rows produce + garbage that nothing reads (`make_omni_output` samples row ``n-1``). +* **KV cache.** ``reshape_and_cache`` writes all ``B`` rows through the + persistent slot buffer. A fresh prefill's slots are contiguous from the + start of its first block, and ``B <= 32 < block_size``, so rows ``n..B-1`` + land on the *same request's own future slots* -- positions the sequence has + not reached. Attention never reads past ``seq_len``, and the first decode + steps overwrite those slots with the real KV as the sequence grows. +* **Inputs.** The graph bakes the runner's own persistent ``inputs_embeds`` + and ``positions`` buffers, exactly as the decode graphs do. The runner + writes rows ``0..n-1`` every step; rows ``n..B-1`` hold stale values whose + garbage stays behind the causal mask. + +The eligibility gate is deliberately narrow -- one request, ``PrefillNoCache``, +offset 0, ``n`` inside a captured bucket -- and everything else takes today's +eager path unchanged. + +Off with ``VLLM_OMNI_FIXED_KV_PREFILL=0``. +""" + +from __future__ import annotations + +import os +from typing import Any, Callable + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_ENV = "VLLM_OMNI_FIXED_KV_PREFILL" +# Per-stage: the Talker's prefill is 9-25 tokens, stage 0's ranked prefill is +# 144-160. The capture routine picks one set per process by model traits. +TALKER_BUCKETS: tuple[int, ...] = (16, 32) +STAGE0_BUCKETS: tuple[int, ...] = (160, 192) +BUCKETS: tuple[int, ...] = TALKER_BUCKETS + + +def configure(buckets: tuple[int, ...]) -> None: + global BUCKETS + BUCKETS = tuple(sorted(buckets)) + +_LOGGED_REPLAY = False + + +def is_enabled() -> bool: + return os.getenv(_ENV, "1") == "1" + + +class _State: + """Per-process buffers and the captured-bucket registry.""" + + def __init__(self) -> None: + self.slot_buffers: dict[int, torch.Tensor] = {} + self.aranges: dict[int, torch.Tensor] = {} + self.workspaces: dict[int, torch.Tensor] = {} + self.lse: dict[int, torch.Tensor] = {} + self.captured: set[int] = set() + # bucket -> (graph, hidden_out): graphs this module owns outright. + self.graphs: dict[int, tuple[Any, torch.Tensor]] = {} + + +_state = _State() + + +def prewarm(device: torch.device) -> None: + """Allocate every runtime buffer before any capture can start. + + An allocation made while a graph is capturing comes out of that graph's + private pool and is stranded there (the fixed-KV workspace leak), so the + slot buffers and aranges exist before `capture` runs. + """ + for bucket in BUCKETS: + if bucket not in _state.slot_buffers: + _state.slot_buffers[bucket] = torch.zeros(bucket, dtype=torch.int32, device=device) + _state.aranges[bucket] = torch.arange(bucket, dtype=torch.int32, device=device) + + +def slot_buffer(bucket: int) -> torch.Tensor: + return _state.slot_buffers[bucket] + + +def workspace(bucket: int, factory: Callable[[], torch.Tensor]) -> torch.Tensor: + """FIA workspace for this bucket, computed outside any capture.""" + ws = _state.workspaces.get(bucket) + if ws is None: + ws = factory() + _state.workspaces[bucket] = ws + return ws + + +def lse_buffer(bucket: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + buf = _state.lse.get(bucket) + if buf is None: + buf = torch.empty(1, dtype=dtype, device=device) + _state.lse[bucket] = buf + return buf + + +def store_graph(bucket: int, graph: Any, hidden_out: torch.Tensor) -> None: + """Register a bucket's captured graph and its baked output tensor. + + These are raw ``NPUGraph`` objects this module owns, the way the codec + step graph owns its own -- deliberately NOT entries in the ACL graph + wrapper. A wrapper-FULL prefill replay poisons the request that follows + it: every codec-step replay afterwards blocks the host for the queued + backbone work (+9 ms per decode step, the bisect4/fixtest arms), and + neither the capture order nor the pre-replay barrier explains it. The + codec graph itself proves raw replays coexist cleanly with everything, + so the prefill graphs live the same way. + """ + _state.graphs[bucket] = (graph, hidden_out) + _state.captured.add(bucket) + + +def mark_captured(bucket: int) -> None: + _state.captured.add(bucket) + + +def bucket_for(n: int) -> int | None: + """Smallest captured bucket that holds an ``n``-token prefill.""" + for bucket in BUCKETS: + if n <= bucket and bucket in _state.captured: + return bucket + return None + + +def maybe_mark_step(attn_metadata: Any, prefill_state: Any, block_size_hint: int = 0) -> None: + """Builder hook: route an eligible prefill step onto its captured graph. + + Cheap, and refuses everything unusual: called once per step from + ``OmniFixedKVMetadataBuilder.build``. When it marks the step, it also + fills the persistent slot buffer -- real slots for the real rows, the + request's own next slots for the padding, one device add with no sync. + """ + if not _state.captured: + return + if getattr(attn_metadata, "attn_state", None) != prefill_state: + return + if not getattr(attn_metadata, "causal", False): + return + lengths = getattr(attn_metadata, "actual_seq_lengths_q", None) + if not lengths or len(lengths) != 1: + return + n = int(lengths[0]) + seq_lens_list = getattr(attn_metadata, "seq_lens_list", None) + if not seq_lens_list or int(seq_lens_list[0]) != n: + # A chunked continuation (offset > 0) needs the earlier KV, which the + # captured PrefillNoCache graph does not read. + return + slot_mapping = getattr(attn_metadata, "slot_mapping", None) + if slot_mapping is None or slot_mapping.shape[0] < n: + return + bucket = bucket_for(n) + if bucket is None: + return + block_size = int(block_size_hint or 0) + if block_size <= 0: + return + if (n - 1) // block_size != (bucket - 1) // block_size: + # The padding rows must land on the request's own allocated blocks. + # Slots are contiguous only within a block, and a request's block + # table covers ceil(n / block_size) blocks -- a bucket that needs one + # more block than the prompt does would write into somebody else's. + # (The first stage-0 round hit exactly this: a 148-token prompt spans + # two blocks, and a slot0+arange fill pushed rows past 127 into the + # physical block after block 0 -- runaway text, 7/32 finished.) + return + block_table = getattr(attn_metadata, "block_tables", None) + if bucket <= block_size: + # One block: slots run contiguously from slot0. + torch.add(_state.aranges[bucket], slot_mapping[0], out=_state.slot_buffers[bucket]) + else: + if block_table is None or block_table.shape[0] < 1: + return + # Per-row: slot(i) = block_table[i // bs] * bs + i % bs. Three device + # ops, no sync; every index stays inside blocks the prompt owns + # because of the same-block-count gate above. + arange = _state.aranges[bucket] + blocks = block_table[0].index_select(0, (arange // block_size).to(torch.int64)) + torch.add(arange % block_size, blocks * block_size, out=_state.slot_buffers[bucket]) + attn_metadata.fixed_kv_prefill_bucket = bucket + + +def replay(attn_metadata: Any) -> torch.Tensor | None: + """Replay the marked bucket's graph and hand back its baked hidden. + + The caller skips the model call entirely: the graph read the runner's + persistent input buffers when it was captured, the builder refreshed the + persistent slot buffer when it marked the step, and the output tensor is + the same storage every time. + """ + if isinstance(attn_metadata, dict): + attn_metadata = next(iter(attn_metadata.values()), None) + bucket = getattr(attn_metadata, "fixed_kv_prefill_bucket", None) + if bucket is None: + return None + entry = _state.graphs.get(bucket) + if entry is None: + return None + graph, hidden_out = entry + graph.replay() + global _LOGGED_REPLAY + if not _LOGGED_REPLAY: + _LOGGED_REPLAY = True + logger.info("[minicpmo] prefill replays its own captured %d-token graph", bucket) + return hidden_out + + diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index c0221f94e64..da5c630a0d2 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -54,6 +54,46 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): _omni_enum = OmniPlatformEnum.NPU dist_backend: str = "hccl" + @classmethod + def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: int | None = None): + """Route short-context decoders onto the fixed-KV decode backend. + + Under FULL_DECODE_ONLY the captured decode step has to re-issue attention on + every layer on every step, because the op takes the KV length as a host + argument that grows each step. On the small Talker decoder that rebind is + ~38% of its busy time. Scoped by ``max_model_len`` (Talker 4096 engages, + Thinker 32768 does not). ``VLLM_OMNI_FIXED_KV_DECODE=0`` restores stock. + + The one-query capture this backend exists for is what makes the 910B + family work at all: the stock builder faults its vector core on + spec-width rows, so without this route there the K-step cannot capture. + """ + resolved = super().get_attn_backend_cls(selected_backend, attn_selector_config, num_heads) + if resolved != "vllm_ascend.attention.attention_v1.AscendAttentionBackend": + return resolved + from vllm.config import get_current_vllm_config + + from vllm_omni.platforms.npu.attention import fixed_kv_decode + + if not fixed_kv_decode.is_enabled(): + return resolved + vllm_config = get_current_vllm_config() + max_model_len = getattr(getattr(vllm_config, "model_config", None), "max_model_len", None) + block_size = getattr(getattr(vllm_config, "cache_config", None), "block_size", None) + capacity = ( + fixed_kv_decode.capacity_for(max_model_len, block_size) + if max_model_len and block_size + else None + ) + fixed_kv_decode.install_into_ascend_aclgraph() + logger.info( + "[minicpmo] fixed-KV decode attention on (max_model_len=%s, kv_capacity=%s); " + "set VLLM_OMNI_FIXED_KV_DECODE=0 to disable", + max_model_len, + capacity, + ) + return "vllm_omni.platforms.npu.attention.fixed_kv_backend.OmniFixedKVAttentionBackend" + # conv2d convolution operator in the code2wav module of Qwen3-TTS not being able to run on Aclnn def __init__(self) -> None: from vllm_ascend.utils import adapt_patch diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 073317150c4..b2b652f1f4e 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -115,7 +115,17 @@ class ExecuteModelState(NamedTuple): class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin): """Autoregressive NPU model runner that returns hidden states per request.""" + # vllm-ascend's ascend-graph wrapper reads this flag from the runner, but + # the vllm-ascend build this image ships never sets it -- the fixed-KV + # capture path would raise AttributeError without a default here. + enable_hamming_sparse: bool = False + def __init__(self, *args, **kwargs): + from vllm_omni.platforms.npu.attention import fixed_kv_decode + + # Before parent init: vLLM's selector uses current_platform + # (vllm-ascend), not NPUOmniPlatform.get_attn_backend_cls. + fixed_kv_decode.install_into_ascend_backend() super().__init__(*args, **kwargs) self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) # each model stage has their own hidden size @@ -169,6 +179,37 @@ def _check_and_update_cudagraph_mode(self, *args, **kwargs): with self._narrow_decode_query_len(): return super()._check_and_update_cudagraph_mode(*args, **kwargs) + def _capture_cudagraphs(self, *args, **kwargs): + """Capture fixed-KV decode graphs per KV capacity bucket. + + Each bucket needs its own pass: the default path leaves the capacity unset + and would file the graphs under a key decode never looks up. Signature is + fully forwarded ("*args/**kwargs") because the upstream parent takes an + extra ``profiler`` argument that this override must not swallow.""" + from vllm.config import CUDAGraphMode + + from vllm_omni.platforms.npu.attention import fixed_kv_decode + + runtime_mode = kwargs.get("cudagraph_runtime_mode") + if runtime_mode is None and len(args) >= 2: + runtime_mode = args[1] + buckets: tuple[int, ...] = () + if runtime_mode == CUDAGraphMode.FULL: + buckets = fixed_kv_decode.buckets_for( + fixed_kv_decode.capacity_for( + self.vllm_config.model_config.max_model_len, + self.vllm_config.cache_config.block_size, + ), + self.vllm_config.cache_config.block_size, + ) + if not buckets: + return super()._capture_cudagraphs(*args, **kwargs) + with self._narrow_decode_query_len(): + for capacity in sorted(buckets, reverse=True): + logger.info("[minicpmo] capturing fixed-KV decode graphs at kv_capacity=%d", capacity) + with fixed_kv_decode.capturing_bucket(capacity): + super()._capture_cudagraphs(*args, **kwargs) + def propose_draft_token_ids(self, valid_sampled_token_ids, *args, **kwargs): """Carry the Talker's frame count to the scheduler, not a prediction. From 70612659816f8e3ce35dc604c09c92823178cde2 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 02:07:31 +0800 Subject: [PATCH 24/89] fix(npu): skip the 910B rejection-sampler warmup unconditionally With the K-step armed, the stock rejection_sampler_triton_warmup faults the 910B vector core (acl 507035) during its dummy-token setup, which is why the guard in ascend_warmup_patch exists. The pre-port tree could tell from the environment; this tree asks the engine config instead, and the stage worker does not have it yet at warmup time -- the engine hands the runner its vllm_config only after load_model, and a deploy-layer record does not survive the spawn that starts these processes. The warmup therefore ran and faulted inside capture_model(), so stage 1 never came up. * _skipped_names(): on Ascend910B return {"rejection_sampler"} without asking _kstep_armed(). The torch-native sampler serves the armed path, and the stock path only pays a first-request JIT of that kernel. * _kstep_armed(): keep the engine-config answer first, fall back to the deploy layer, so callers that do have the runtime config keep their precise answer. * stage_config: record stage 1's declared frame count when a deploy YAML is parsed and expose talker_frames_per_step() for those out-of-band readers. Verified on 910B3: the prof deploy config now reaches "Application startup complete" with stage 2's two replicas, and stage 1 logs the per-capacity fixed-KV captures. Signed-off-by: y-null --- vllm_omni/config/stage_config.py | 35 ++++++++++++++++ .../platforms/npu/ascend_warmup_patch.py | 41 ++++++++++++------- 2 files changed, 61 insertions(+), 15 deletions(-) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 06cb4ac7827..ff11d2b225a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -785,6 +785,40 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: # model-side gate accepts. _MINICPMO_TALKER_FRAMES_MAX = 16 +# Frames the deploy layer declared for stage 1's Talker multi-frame decode, +# recorded when a deploy config is parsed. NPU workers inherit the value across +# the fork, so a reader that runs before the engine hands the runner its +# vllm_config -- the warmup guard is one -- sees the same decision. +_resolved_talker_frames: int = 1 + + +def talker_frames_per_step() -> int: + """Codec frames one stage-1 step produces, per the last deploy config parsed. + + 1 means the deploy layer did not ask for the multi-frame loop. + """ + return _resolved_talker_frames + + +def _record_talker_frames(stages: list[StageDeployConfig]) -> None: + """Remember stage 1's declared frame count for out-of-band readers. + + Stage 1 carries the multi-frame decode as an n-gram speculative_config + whose ``num_speculative_tokens`` is K - 1, so the deploy YAML is the only + place the decision is written down. + """ + global _resolved_talker_frames + for stage in stages: + if stage.stage_id != 1: + continue + spec = (stage.engine_extras or {}).get("speculative_config") or {} + frames = 1 + if isinstance(spec, dict) and spec.get("method") == "ngram": + num_spec = spec.get("num_speculative_tokens", 0) or 0 + if num_spec > 0: + frames = min(int(num_spec) + 1, _MINICPMO_TALKER_FRAMES_MAX) + _resolved_talker_frames = frames + def load_deploy_config(path: str | Path) -> DeployConfig: """Load a deploy YAML (with optional base_config inheritance).""" @@ -796,6 +830,7 @@ def load_deploy_config(path: str | Path) -> DeployConfig: ) stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])] + _record_talker_frames(stages) model_runner = raw_dict.get("model_runner", "v1") if model_runner not in ("v1", "v2"): diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py index 0944ce1a0ab..e965425797c 100644 --- a/vllm_omni/platforms/npu/ascend_warmup_patch.py +++ b/vllm_omni/platforms/npu/ascend_warmup_patch.py @@ -85,22 +85,32 @@ def _probe_soc_name() -> str: def _kstep_armed() -> bool: """Whether this worker runs the Talker multi-frame decode. - Read off the engine's speculative_config, which the deploy YAML's stage-1 - block provides; an unreadable config counts as not armed so the baseline - warmup never changes by accident. + Preferred source is the engine's speculative_config, which the deploy + YAML's stage-1 block provides. The warmup guard runs before the engine + hands the runner that config, so fall back to the deploy layer, which + ``stage_config`` records when the YAML is parsed: both answer the same + question, and answering "not armed" there would run the rejection-sampler + Triton warmup this guard exists to skip, faulting the 910B vector core + (acl 507035) during capture. An unreadable answer still counts as not + armed, so the baseline warmup never changes by accident. """ try: from vllm.config import get_current_vllm_config_or_none cfg = get_current_vllm_config_or_none() except Exception: - return False + cfg = None spec = getattr(cfg, "speculative_config", None) if cfg is not None else None - if spec is None: + if spec is not None: + method = getattr(spec, "method", None) + num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 + return method == "ngram" and num_spec > 0 + try: + from vllm_omni.config.stage_config import talker_frames_per_step + + return talker_frames_per_step() > 1 + except Exception: return False - method = getattr(spec, "method", None) - num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 - return method == "ngram" and num_spec > 0 def _skipped_names() -> set[str]: @@ -111,13 +121,14 @@ def _skipped_names() -> set[str]: return {part.strip() for part in raw.split(",") if part.strip()} soc = _probe_soc_name() if soc.startswith("Ascend910B"): - if _kstep_armed(): - # The torch-native sampler replaces the reject kernels whenever - # the K-step is armed, so this warmup would only compile kernels - # nobody runs. Everything else keeps the stock warmup. - return {"rejection_sampler"} - # Verified clean here; keep the stock warmup untouched. - return set() + # On this family the rejection-sampler Triton warmup faults the vector + # core (acl 507035) whenever the K-step is armed, and the worker cannot + # tell yet at warmup time: the engine hands the runner its vllm_config + # only after ``load_model``, and a deploy-layer record does not survive + # the spawn that starts these stage processes. Skipping is the safe + # side on both paths -- the torch-native sampler serves the armed path, + # and the stock path merely JITs that kernel on its first request. + return {"rejection_sampler"} return set(_DEFAULT_SKIP) From 2bce893d9d24dd4fe527ffa5d5b24bc8069ec031 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 02:49:50 +0800 Subject: [PATCH 25/89] docs(config): correct the talker_frames_per_step comment on stage workers Stage workers are spawned, so they do not inherit the recorded value; the comment claimed a fork. The warmup guard on that side decides without it (see ascend_warmup_patch._skipped_names). Signed-off-by: y-null --- vllm_omni/config/stage_config.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index ff11d2b225a..59c09f27687 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -786,9 +786,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: _MINICPMO_TALKER_FRAMES_MAX = 16 # Frames the deploy layer declared for stage 1's Talker multi-frame decode, -# recorded when a deploy config is parsed. NPU workers inherit the value across -# the fork, so a reader that runs before the engine hands the runner its -# vllm_config -- the warmup guard is one -- sees the same decision. +# recorded when a deploy config is parsed. A reader in the same process that +# runs before the engine hands the runner its vllm_config -- the warmup guard +# is one -- sees the same decision here. Stage workers are spawned, so they do +# NOT inherit it: the guard on that side decides without asking, see +# ascend_warmup_patch._skipped_names. _resolved_talker_frames: int = 1 From 436cb686bcdff16d843c17f1866a7ce0c7a1fda0 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 03:17:18 +0800 Subject: [PATCH 26/89] docs(npu): tighten the comments added by the K-step fixes Same facts, shorter: the deploy-layer fallback, the 910B warmup skip and the fixed-KV route each lose a paragraph of restatement. Signed-off-by: y-null --- vllm_omni/config/stage_config.py | 20 +++++------------ .../platforms/npu/ascend_warmup_patch.py | 22 ++++++------------- vllm_omni/platforms/npu/platform.py | 6 ++--- 3 files changed, 14 insertions(+), 34 deletions(-) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 59c09f27687..1bb6753e9ce 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -785,30 +785,20 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: # model-side gate accepts. _MINICPMO_TALKER_FRAMES_MAX = 16 -# Frames the deploy layer declared for stage 1's Talker multi-frame decode, -# recorded when a deploy config is parsed. A reader in the same process that -# runs before the engine hands the runner its vllm_config -- the warmup guard -# is one -- sees the same decision here. Stage workers are spawned, so they do -# NOT inherit it: the guard on that side decides without asking, see +# Stage 1's declared frame count, recorded when a deploy config is parsed. +# Same-process readers that run before the engine hands the runner its +# vllm_config see it here; spawned stage workers do not, see # ascend_warmup_patch._skipped_names. _resolved_talker_frames: int = 1 def talker_frames_per_step() -> int: - """Codec frames one stage-1 step produces, per the last deploy config parsed. - - 1 means the deploy layer did not ask for the multi-frame loop. - """ + """Codec frames one stage-1 step produces; 1 when the deploy layer asked for none.""" return _resolved_talker_frames def _record_talker_frames(stages: list[StageDeployConfig]) -> None: - """Remember stage 1's declared frame count for out-of-band readers. - - Stage 1 carries the multi-frame decode as an n-gram speculative_config - whose ``num_speculative_tokens`` is K - 1, so the deploy YAML is the only - place the decision is written down. - """ + """Record stage 1's frame count (its n-gram ``num_speculative_tokens`` is K - 1).""" global _resolved_talker_frames for stage in stages: if stage.stage_id != 1: diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py index e965425797c..969391e2983 100644 --- a/vllm_omni/platforms/npu/ascend_warmup_patch.py +++ b/vllm_omni/platforms/npu/ascend_warmup_patch.py @@ -85,14 +85,8 @@ def _probe_soc_name() -> str: def _kstep_armed() -> bool: """Whether this worker runs the Talker multi-frame decode. - Preferred source is the engine's speculative_config, which the deploy - YAML's stage-1 block provides. The warmup guard runs before the engine - hands the runner that config, so fall back to the deploy layer, which - ``stage_config`` records when the YAML is parsed: both answer the same - question, and answering "not armed" there would run the rejection-sampler - Triton warmup this guard exists to skip, faulting the 910B vector core - (acl 507035) during capture. An unreadable answer still counts as not - armed, so the baseline warmup never changes by accident. + Prefers the engine's speculative_config. The warmup guard runs before the + engine hands the runner that config, so fall back to the deploy layer. """ try: from vllm.config import get_current_vllm_config_or_none @@ -121,13 +115,11 @@ def _skipped_names() -> set[str]: return {part.strip() for part in raw.split(",") if part.strip()} soc = _probe_soc_name() if soc.startswith("Ascend910B"): - # On this family the rejection-sampler Triton warmup faults the vector - # core (acl 507035) whenever the K-step is armed, and the worker cannot - # tell yet at warmup time: the engine hands the runner its vllm_config - # only after ``load_model``, and a deploy-layer record does not survive - # the spawn that starts these stage processes. Skipping is the safe - # side on both paths -- the torch-native sampler serves the armed path, - # and the stock path merely JITs that kernel on its first request. + # This warmup faults the vector core (acl 507035) on this family, and the + # worker cannot tell at warmup time whether the K-step is armed (spawned + # process, engine config not handed over yet). Skip either way: the armed + # path uses the torch-native sampler, the stock path only JITs the kernel + # on its first request. return {"rejection_sampler"} return set(_DEFAULT_SKIP) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index da5c630a0d2..7ae2b5ba7aa 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -63,10 +63,8 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: argument that grows each step. On the small Talker decoder that rebind is ~38% of its busy time. Scoped by ``max_model_len`` (Talker 4096 engages, Thinker 32768 does not). ``VLLM_OMNI_FIXED_KV_DECODE=0`` restores stock. - - The one-query capture this backend exists for is what makes the 910B - family work at all: the stock builder faults its vector core on - spec-width rows, so without this route there the K-step cannot capture. + The 910B family also needs it to capture at all: the stock builder + faults there on spec-width rows. """ resolved = super().get_attn_backend_cls(selected_backend, attn_selector_config, num_heads) if resolved != "vllm_ascend.attention.attention_v1.AscendAttentionBackend": From 6e1f366f91903cb03d86e2f7a9b8327b1a5392b7 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 09:15:31 +0800 Subject: [PATCH 27/89] style: satisfy pre-commit across the branch pre-commit failed on ruff-format (7 files) and ruff-check (6 findings). Run ruff format and ruff check --fix over vllm_omni/, and drop the unused prev_code assignment in minicpmo_4_5_omni_tts.py that F841 flagged. Signed-off-by: y-null --- vllm_omni/core/sched/omni_ar_scheduler.py | 2 -- .../models/minicpmo_4_5/minicpmo_4_5_omni_tts.py | 3 --- .../platforms/npu/attention/fixed_kv_backend.py | 13 +++---------- .../platforms/npu/attention/fixed_kv_decode.py | 9 +++------ .../platforms/npu/attention/fixed_kv_prefill.py | 6 +++--- .../platforms/npu/models/minicpmo_4_5_code2wav.py | 10 +++++++++- vllm_omni/platforms/npu/platform.py | 6 +----- 7 files changed, 19 insertions(+), 30 deletions(-) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 9f70ac09d76..b1ce50c0a33 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -3,7 +3,6 @@ from __future__ import annotations -import os import time from collections import defaultdict from collections.abc import Iterable, Iterator @@ -369,7 +368,6 @@ def _talker_kstep_armed(self) -> bool: cache = self._omni_talker_kstep_cache = is_talker and is_ngram and num_spec > 0 return cache - def _drop_talker_drafts_if_prefill_pending(self) -> None: """Keep the Talker's K-frame decode out of steps with uneven spans. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 45d95f0535a..f3e877f611e 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -340,8 +340,6 @@ def resolve_codec_sampling_params( return resolved - - def _codec_int_param(state: Any, key: str, fallback: int) -> int: """Read an integer codec knob from a request state, None meaning unset. @@ -1065,7 +1063,6 @@ def make_omni_output( # stub this method with the sampled tensor itself. sampled = stochastic_result.reshape(()).to(torch.long) sampled_id = int(sampled.item()) - prev_code = state.get("last_code") if isinstance(state, Mapping) else None is_eos = sampled_id == self._codec_eos_id state["step"] = _codec_int_param(state, "step", 0) + 1 reached_limit = int(state["step"]) >= _codec_int_param(state, "max_tokens", self._codec_max_tokens) diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_backend.py b/vllm_omni/platforms/npu/attention/fixed_kv_backend.py index dc2c37d12ed..cbe6069522b 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_backend.py +++ b/vllm_omni/platforms/npu/attention/fixed_kv_backend.py @@ -34,7 +34,6 @@ import torch import torch_npu from vllm.logger import init_logger - from vllm_ascend.ascend_forward_context import _EXTRA_CTX from vllm_ascend.attention.attention_v1 import ( SWA_INT_MAX, @@ -106,9 +105,7 @@ def build(self, *args, **kwargs): if self._omni_buckets and attn_metadata.fixed_kv_decode_step: seq_lens_list = attn_metadata.seq_lens_list or [1] max_seq_len = getattr(common, "max_seq_len", None) or max(seq_lens_list) - fixed_kv_decode.set_runtime_bucket( - fixed_kv_decode.select_bucket(int(max_seq_len), self._omni_buckets) - ) + fixed_kv_decode.set_runtime_bucket(fixed_kv_decode.select_bucket(int(max_seq_len), self._omni_buckets)) else: fixed_kv_decode.set_runtime_bucket(None) # A single fresh prefill inside a captured bucket replays a graph @@ -161,9 +158,7 @@ def __init__(self, *args, **kwargs) -> None: # +2 matches the query_start_loc buffer, which carries the FIA padding # request on a full batch. self._omni_max_rows = int(self.vllm_config.scheduler_config.max_num_seqs) + 2 - self._omni_buckets = fixed_kv_decode.buckets_for( - self._omni_capacity, self.vllm_config.cache_config.block_size - ) + self._omni_buckets = fixed_kv_decode.buckets_for(self._omni_capacity, self.vllm_config.cache_config.block_size) # A step that schedules several tokens per sequence (the Talker's # multi-frame decode declares them through a speculative_config) is # still a uniform decode, and the fixed-capacity mask covers it as @@ -324,9 +319,7 @@ def _fixed_kv_graph_fia( # The bias is per step, not per layer, and only the FIA path reads it, # so it is recorded by the first layer that actually wants it. if not attn_metadata.fixed_kv_pse_recorded: - fixed_kv_decode.emit_mask_refresh( - pse, seq_lens_device, self.num_heads, capacity, q_len - ) + fixed_kv_decode.emit_mask_refresh(pse, seq_lens_device, self.num_heads, capacity, q_len) attn_metadata.fixed_kv_pse_recorded = True softmax_lse = fixed_kv_decode.keep_alive( diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py b/vllm_omni/platforms/npu/attention/fixed_kv_decode.py index fab34a28df4..8366f627d51 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py +++ b/vllm_omni/platforms/npu/attention/fixed_kv_decode.py @@ -179,9 +179,7 @@ def _call(self, *args, **kwargs): wrapper_cls.__call__ = _call wrapper_cls._omni_fixed_kv_call = True - logger.info( - "[minicpmo] fixed-KV: ACLGraphWrapper keyed per KV capacity, pre-replay barrier dropped" - ) + logger.info("[minicpmo] fixed-KV: ACLGraphWrapper keyed per KV capacity, pre-replay barrier dropped") def _is_fixed_graph_here(capacity: int) -> bool: @@ -205,6 +203,7 @@ def install_into_ascend_backend() -> None: if not is_enabled(): return from vllm_ascend.attention.attention_v1 import AscendAttentionBackend + from vllm_omni.platforms.npu.attention.fixed_kv_backend import ( OmniFixedKVAttentionBackendImpl, OmniFixedKVMetadataBuilder, @@ -382,9 +381,7 @@ def emit_mask_refresh( """ ar = _get_arange(num_heads, capacity, pse.device, q_len) neg, zero = _get_scalars(pse.dtype, pse.device) - lens = ( - (seq_lens_device - (q_len - 1)).clamp(min=_MIN_SEQ_LEN).to(torch.int32).view(-1, 1, 1, 1) - ) + lens = (seq_lens_device - (q_len - 1)).clamp(min=_MIN_SEQ_LEN).to(torch.int32).view(-1, 1, 1, 1) torch.where(ar >= lens, neg, zero, out=pse) diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py b/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py index 9dc126dee49..a69315a35c8 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py +++ b/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py @@ -39,7 +39,8 @@ from __future__ import annotations import os -from typing import Any, Callable +from collections.abc import Callable +from typing import Any import torch from vllm.logger import init_logger @@ -58,6 +59,7 @@ def configure(buckets: tuple[int, ...]) -> None: global BUCKETS BUCKETS = tuple(sorted(buckets)) + _LOGGED_REPLAY = False @@ -223,5 +225,3 @@ def replay(attn_metadata: Any) -> torch.Tensor | None: _LOGGED_REPLAY = True logger.info("[minicpmo] prefill replays its own captured %d-token graph", bucket) return hidden_out - - diff --git a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py index df519f3f541..b11c14cab3d 100644 --- a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py +++ b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py @@ -164,6 +164,7 @@ def _step_kwargs(rest): if has_mask: kwargs["attn_mask"] = rest[0] return kwargs + if cnn_cache is None: return graph_runner.run( "cfm_estimator", @@ -187,7 +188,14 @@ def _step_kwargs(rest): "cfm_estimator", (x, mu, time_embedding, speakers, cond, cnn_cache, att_cache, *mask_inputs), ((True, has_mask) if has_mask else (True,)), - lambda step_x, step_mu, step_time, step_speakers, step_cond, step_cnn, step_att, *rest: _graphable_estimator_step( + lambda step_x, + step_mu, + step_time, + step_speakers, + step_cond, + step_cnn, + step_att, + *rest: _graphable_estimator_step( self, estimator, x=step_x, diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 7ae2b5ba7aa..ce3d3249ba1 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -78,11 +78,7 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: vllm_config = get_current_vllm_config() max_model_len = getattr(getattr(vllm_config, "model_config", None), "max_model_len", None) block_size = getattr(getattr(vllm_config, "cache_config", None), "block_size", None) - capacity = ( - fixed_kv_decode.capacity_for(max_model_len, block_size) - if max_model_len and block_size - else None - ) + capacity = fixed_kv_decode.capacity_for(max_model_len, block_size) if max_model_len and block_size else None fixed_kv_decode.install_into_ascend_aclgraph() logger.info( "[minicpmo] fixed-KV decode attention on (max_model_len=%s, kv_capacity=%s); " From 00e9ee06a1d72868d8ac615b36c21ba97ff60c3d Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 11:28:43 +0800 Subject: [PATCH 28/89] docs(minicpmo_4_5): align the comments with the code and drop dead knobs Review pass over this branch's own additions. Comments and docstrings still described removed environment switches, internal item numbers, measurement rounds and out-of-tree design notes. They now describe what the code does, in the wording the rest of the repo uses. Dead code removed: * the unused ``MINICPMO_FAST_CODEC_PENALTY`` read (nothing consumed it); * ``fixed_kv_prefill.is_enabled()`` / ``STAGE0_BUCKETS`` / the ``VLLM_OMNI_FIXED_KV_PREFILL`` gate, which nothing called; * the scheduler's per-request ``states`` list, built and never read; * the unused ``_CODEC_EOS_TOKEN_ID`` left behind by the stop-chain fix; * the frame-count clamp whose only reader compares against 1. Behavior changes are limited to the two that make a gate match its reader: the vocab gate now requires the same arming predicate as the multi-frame path itself (``num_spec_tokens > 0``) rather than the wrapper class alone, and ``get_attn_backend_cls`` logs which of the two fixed-KV outcomes it took. The environment variables this branch introduces are registered in the inventory and the configuration docs. Signed-off-by: y-null --- docs/configuration/environment_variables.md | 4 ++ tests/config/test_environment_variables.py | 6 +-- .../config/environment_variable_inventory.py | 6 +++ vllm_omni/config/stage_config.py | 17 +++---- vllm_omni/core/sched/omni_ar_scheduler.py | 16 +++--- .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 5 +- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 16 +++--- .../models/minicpmo_4_5/pipeline.py | 29 ++++++----- .../platforms/npu/ascend_warmup_patch.py | 41 ++++++++------- .../npu/attention/fixed_kv_backend.py | 21 ++++---- .../npu/attention/fixed_kv_decode.py | 50 +++++++++---------- .../npu/attention/fixed_kv_prefill.py | 39 ++++++--------- .../platforms/npu/native_rejection_sampler.py | 22 ++++---- vllm_omni/platforms/npu/platform.py | 28 +++++++---- .../npu/worker/npu_ar_model_runner.py | 21 +++++--- .../platforms/npu/worker/npu_model_runner.py | 38 +++++++------- .../platforms/npu/worker/talker_multiframe.py | 36 ++++++------- 17 files changed, 197 insertions(+), 198 deletions(-) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index eaefb1015af..e6ecb1092c6 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -122,6 +122,10 @@ settings. | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | | --- | --- | --- | --- | --- | +| `VLLM_OMNI_FIXED_KV_DECODE` | `0` disables; default `1` (on) | Fixed-capacity decode attention backend and its ACL-graph hook; read when the NPU platform resolves the attention backend | Only the exact value `0` disables it. Scoped by `max_model_len`: a long-context stage (e.g. a 32768-token Thinker) declines and keeps the stock path. | Experimental performance control | +| `VLLM_OMNI_FIXED_KV_DECODE_BUCKETS` | Comma-separated KV capacities in tokens; default `512,4096` | Capture buckets of the fixed-KV decode graphs; read when the buckets are computed | The full capacity is always kept as the top bucket, so a longer sequence still has a graph. Empty or malformed entries are skipped. | Experimental performance control | +| `VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN` | Integer tokens; default `8192` | Fixed-KV decode eligibility; read per step when attention decides whether to engage | At or below this length the fixed-KV path is used; above it the stock path runs. | Experimental performance control | +| `VLLM_OMNI_NPU_SKIP_WARMUPS` | Comma-separated warmup-name substrings, or `none`; default `penalties,rejection_sampler` on 910_93 and `rejection_sampler` on 910B | vllm-ascend Triton kernel warmups on the NPU worker; resolved per warmup call | Skips the warmups whose dummy-token setup faults the vector core (acl 507035). `none` restores the stock warmups. | Experimental | | `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control | | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary | | `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 201c9d0e789..4df7e92b5f1 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,10 +164,10 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 30, + EnvironmentVariableCategory.PUBLIC_OMNI: 35, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 59, + EnvironmentVariableCategory.MODEL_SPECIFIC: 61, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 2, } @@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 31, + ModelEnvironmentVariableDisposition.PROMOTE: 33, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 16, diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index c3871bacec7..abbf4769e6b 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -67,10 +67,14 @@ def is_public_omni(self) -> bool: "VLLM_OMNI_AR_DIFFUSION_KV_GATHER", "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT", "VLLM_OMNI_EVENT_DRIVEN_ORCH", + "VLLM_OMNI_FIXED_KV_DECODE", + "VLLM_OMNI_FIXED_KV_DECODE_BUCKETS", + "VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN", "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "VLLM_OMNI_NIXL_LEASE_S", "VLLM_OMNI_NIXL_XFER_TIMEOUT_S", + "VLLM_OMNI_NPU_SKIP_WARMUPS", "VLLM_OMNI_ORCH_MONITOR_PATH", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", "VLLM_OMNI_SERVER_STORAGE__FILE_TTL", @@ -162,6 +166,8 @@ def is_public_omni(self) -> bool: "MINICPMO_TOKEN2WAV_TRT", "MINICPMO_TOKEN2WAV_TRT_CACHE", "MINICPMO_TOKEN2WAV_TRT_DTYPE", + "MINICPMO_TTS_NPU_TOPK_TOPP", + "MINICPMO_TTS_NPU_TOPK_TOPP_MAX_BATCH", "NEMOTRON_VOICECHAT_LLM_PATH", "OMNIVOICE_CUDA_GRAPH", "OMNIVOICE_TF32", diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 1bb6753e9ce..e7bdd6994d8 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -779,16 +779,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: return merged -# Same TND ceiling as stage 0: at most 16 query positions per sequence. The -# Talker's multi-frame decode is configured through stage 1's -# speculative_config in the deploy YAML; this only bounds the frame count the -# model-side gate accepts. -_MINICPMO_TALKER_FRAMES_MAX = 16 - -# Stage 1's declared frame count, recorded when a deploy config is parsed. -# Same-process readers that run before the engine hands the runner its -# vllm_config see it here; spawned stage workers do not, see -# ascend_warmup_patch._skipped_names. +# Stage 1's multi-frame decode is configured through its ``speculative_config`` +# in the deploy YAML. The parse records whether it is armed here, because the +# reader (``ascend_warmup_patch._kstep_armed``) can run in a spawned stage +# worker that has no vllm_config to read it from. It only answers "armed or +# not" -- 1 means not armed. _resolved_talker_frames: int = 1 @@ -808,7 +803,7 @@ def _record_talker_frames(stages: list[StageDeployConfig]) -> None: if isinstance(spec, dict) and spec.get("method") == "ngram": num_spec = spec.get("num_speculative_tokens", 0) or 0 if num_spec > 0: - frames = min(int(num_spec) + 1, _MINICPMO_TALKER_FRAMES_MAX) + frames = int(num_spec) + 1 _resolved_talker_frames = frames diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index b1ce50c0a33..28942c76126 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -400,14 +400,11 @@ def _drop_talker_drafts_if_prefill_pending(self) -> None: max_len = None prefill_pending = bool(self.waiting) widths: set[int] = set() - states: list = [] for req in self.running: computed = int(req.num_computed_tokens) prompt_len = len(req.prompt_token_ids) total = int(getattr(req, "num_tokens", prompt_len)) delta = total - computed - spec_len = len(req.spec_token_ids or []) - states.append((str(getattr(req, "request_id", "?"))[:12], computed, prompt_len, total, spec_len, delta)) if computed < prompt_len: prefill_pending = True continue @@ -560,10 +557,9 @@ def update_from_output( stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): if num_tokens_scheduled <= 0: - # P17 SCHED0 diagnostic: a zero-token schedule entry used to - # die on the bare assert below with no context, which is what - # made the K-step stall so expensive to scope. Report WHICH - # request and in what state first; the assert still fires. + # A zero-token schedule entry used to die on the bare assert + # below with no context. Report which request and in what state + # first; the assert still fires. _r = self.requests.get(req_id) _spec = scheduler_output.scheduled_spec_decode_tokens.get(req_id) logger.error( @@ -672,7 +668,7 @@ def update_from_output( request_id=req_id, ) elif scheduled_spec_token_ids and sampled_token_ids: - # P17 layer C: this request had drafts scheduled but its row + # This request had drafts scheduled but its row # came back with no generated tokens while the step itself # sampled (a logprob-contract failure emptied the row above, # the rejection sampler dropped every token, ...). The draft @@ -686,7 +682,7 @@ def update_from_output( # skips it; the guard above keeps that path intact. # A prefill-chunk row is different: the chunk's prompt tokens # did land in the KV cache, so only the D drafts roll back - # (the case entry_02 measured; the base token stays advanced). + # (the base token stays advanced). # A pure-decode row that came back empty rolls the base token # back too -- schedule() optimistically advanced all scheduled # tokens and not one of them produced output. @@ -1035,7 +1031,7 @@ def _update_request_as_session(self, session: Request, update: StreamingUpdate) session.async_tokens_to_discard = 1 session.num_computed_tokens -= session.num_output_placeholders session.num_output_placeholders = 0 - # P17 layer B: clear stale drafts unconditionally, not only under + # Clear stale drafts unconditionally, not only under # async scheduling. With async off -- which the K-step vehicle # requires -- a draft proposed before the segment boundary survived # into the next segment's prefill; the runner discards prefill rows diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index d7306af05d0..fe3f9565f74 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -649,7 +649,10 @@ def make_omni_output(self, model_outputs, **kwargs): # and the stage-1 engine dies with empty audio. @property def supports_multi_frame_decode(self) -> bool: - # Present on both stage kinds; only the Talker (tts) path uses it. + # "This wrapper implements the multi-frame path", not "the loop is + # running": the Talker (tts) stage has it, the Thinker does not. Whether + # a step actually runs it is decided per deployment by stage 1's + # speculative_config (the runner checks its num_spec_tokens for that). return self.model_stage == "tts" def take_batch_stop_logits(self): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index f3e877f611e..eaf285367c7 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -64,7 +64,6 @@ _CODEC_REPETITION_PENALTY = 1.05 _CODEC_MIN_TOKENS = 50 _CODEC_MAX_TOKENS = 2048 -_FAST_CODEC_PENALTY = os.getenv("MINICPMO_FAST_CODEC_PENALTY", "1") == "1" # YAML key -> (tts_config attribute, hardcoded fallback, type) _CODEC_SAMPLING_SOURCES: tuple[tuple[str, str, Any, Any], ...] = ( ("seed", "seed", _CODEC_SEED, int), @@ -309,15 +308,14 @@ def resolve_codec_sampling_params( ) -> dict[str, Any]: """Resolve Talker codec knobs: deploy YAML, then checkpoint, then defaults. - The deploy YAML that starts a ranked run comes from the organizer's - baseline branch and carries no ``codec_sampling_params`` block, so - requiring one there is a startup failure rather than a configuration - error: the Talker stage would not come up at all under the official - launch. Falling back to the checkpoint's ``tts_config`` is what the - upstream Talker does, and it is what makes this tree drop into the - official deploy config unchanged. + A deploy YAML is not required to carry a ``codec_sampling_params`` block -- + the stock MiniCPM-o configs do not -- so requiring one would turn a + missing block into a startup failure: the Talker stage would not come up + at all. Falling back to the checkpoint's ``tts_config`` is what the + upstream Talker does, and it is what lets an existing deploy config work + unchanged. - A YAML block still wins key by key, so our own configs keep overriding. + A YAML block still wins key by key, so a deployment can override any knob. """ provided = yaml_params if isinstance(yaml_params, Mapping) else {} resolved: dict[str, Any] = {} diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index ae4d400ae19..68c77082e23 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -20,24 +20,23 @@ MINICPMO45_REFERENCE_AUDIO_KEY = "_minicpmo45_reference_audio" -_CODEC_EOS_TOKEN_ID = 6561 # tts_config.num_audio_tokens - 1 - - def _talker_stop_token_ids() -> list[int]: - """Stage 1's stop token, matching the head that will actually run. - - Two shapes exist for the Talker's vLLM-level head: + """Stage 1's default stop id: the one the multi-frame head can emit. - * multi-frame decode on -- ``compute_logits`` collapses to the two-wide - continue/stop row, so the only sampleable ids are 0 and 1 and the codec - EOS can never appear at this level (the model samples codec ids inside - the loop and forwards them to stage 2 itself). + Stage 1's deploy config carries the ``speculative_config`` that arms the + multi-frame decode, and ``compute_logits`` then collapses the Talker's + vLLM-level head to the two-wide continue/stop row. The only sampleable ids + are therefore 0 and 1, and the codec EOS can never appear at this level (the + model samples codec ids inside the loop and forwards them to stage 2 + itself). - Pinning this to [6561] regardless is what made every multi-frame request run - to ``max_tokens``: the model did emit the stop marker, ``check_stop`` - compared it against 6561, and the request kept its slot until the context - ran out. A deployment that turns multi-frame decode off overrides this per - stage with the codec EOS. + Pinning this to the codec EOS regardless is what made every multi-frame + request run to ``max_tokens``: the model did emit the stop marker, + ``check_stop`` compared it against 6561, and the request kept its slot until + the context ran out. A deployment that drops the ``speculative_config`` from + stage 1 falls back to one frame per step, where the head is the full codec + vocabulary again and the codec EOS is the stop id; it overrides this per + stage through the usual ``stop_token_ids`` sampling parameter. """ from vllm_omni.platforms.npu.worker.talker_multiframe import STOP_TOKEN_ID diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py index 969391e2983..d8684028ba6 100644 --- a/vllm_omni/platforms/npu/ascend_warmup_patch.py +++ b/vllm_omni/platforms/npu/ascend_warmup_patch.py @@ -15,36 +15,35 @@ error 507035 -- which kills the stage worker during initialization and takes the whole orchestrator down with it. -The identical line runs clean on 910B3 under the same CANN 9.1.0 and -torch_npu 2.10.0.post4 (verified directly, with the same shapes), so this is -a 910_93 operator fault rather than a shape, dtype or memory problem -- the -warmup input is only ``[max_num_seqs, 257]`` token ids. Rewriting the line -would merely relocate the fault into an operator nobody can verify on the -target box, so the warmup is skipped instead. +The identical line runs clean on 910B3 on the same CANN and torch_npu stack, +so this is a 910_93 operator fault rather than a shape, dtype or memory +problem -- the warmup input is only ``[max_num_seqs, 257]`` token ids. +Rewriting the line would merely relocate the fault into an operator nobody +can verify on the target box, so the warmup is skipped instead. Skipping is safe by construction: the warmup only pre-compiles Triton kernels (JIT), and the sampling path runs those kernels over real token tensors, never through this dummy construction. The cost is a one-off compile on the first penalised sample of the first request. Correctness is unaffected. -Scoped so the 910B baseline keeps the stock warmup byte-for-byte: the guards -delegate to the original warmup there. The skip set is resolved when a warmup -actually runs rather than when the platform is constructed, because the -constructor can run before the device is set -- and a probe that failed there -would silently strip a warmup the baseline relies on. A probe that still -cannot name the SoC skips the faulting warmup: a missing warmup costs latency, -a faulting one costs the whole run. +Scoped per family: 910B3 skips only ``rejection_sampler`` -- its dummy +``(batch, spec_len, vocab)`` sweep faults the vector core before any kernel +runs, the armed path uses the torch-native sampler anyway, and the stock path +merely JITs on its first request. Every other warmup stays stock there. The +skip set is resolved when a warmup actually runs rather than when the platform +is constructed, because the constructor can run before the device is set -- +and a probe that failed there would silently strip a warmup the baseline +relies on. A probe that still cannot name the SoC skips the faulting warmup: a +missing warmup costs latency, a faulting one costs the whole run. ``VLLM_OMNI_NPU_SKIP_WARMUPS`` overrides the set (comma separated substrings of the warmup names, or "none" to restore stock behaviour). -``rejection_sampler_triton_warmup`` joined the default skip set after the A3 -run of 02:38 pinned the acl 507035 aivec fault on it: the dummy -(batch, spec_len, vocab) sweep "completed" on the host half a second before -the pre-capture synchronize raised, and nothing else was enqueued in between -(penalties was already skipped, rms is a no-op there). Skipping is lossless -the same way: on this stack no 910C runtime path launches the reject_sample -kernels -- the Talker's K-step sample() and its SoC guards fall back to -torch argmax, so the warmup only ever compiled kernels nobody runs. +``rejection_sampler_triton_warmup`` is in the default skip set for the same +reason: its dummy (batch, spec_len, vocab) sweep faults the 910_93 vector +core, and the fault surfaces at the pre-capture synchronize rather than +inside the sweep. Skipping is lossless here in the same way as above: on this +stack no 910C runtime path launches the reject_sample kernels, because the +Talker's K-step sample() and its SoC guards fall back to torch argmax. """ from __future__ import annotations diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_backend.py b/vllm_omni/platforms/npu/attention/fixed_kv_backend.py index cbe6069522b..8f9a8d5a0de 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_backend.py +++ b/vllm_omni/platforms/npu/attention/fixed_kv_backend.py @@ -7,10 +7,8 @@ ``SymInt[]``, so its tiling is baked into the captured task. The KV length grows by one every step, so vllm-ascend re-issues the op for every layer on every step (``update_full_graph_params`` -> ``graph_task_update_begin``/``_end``). On the -20-layer, 768-hidden MiniCPM-o Talker that rebind *is* the step: py-spy puts 26% -of the stage-1 main-thread wall -- ~38% of its busy time -- inside it, and -removing it measured RTF 0.1896 -> 0.1599 on A3. -See ``server-env/design/R1_UNDER_FULL_DECODE_20260827.md``. +20-layer, 768-hidden MiniCPM-o Talker that rebind *is* the step: it dominates +the stage-1 main-thread wall, which is what this module removes. The mechanism lives in :mod:`fixed_kv_decode`; this module is only the seam. It is a subclass rather than a patch because every place vllm-ascend needs to be @@ -24,9 +22,9 @@ ACLGraphWrapper.__call__ key + replay barrier -> wrapped by :func:`fixed_kv_decode.install_into_ascend_aclgraph` -That matters: the ranked evaluation installs only this tree and leaves -vllm-ascend as the image ships it, so anything that has to be applied *to* -vllm-ascend cannot reach a scored run. +That matters: vllm-ascend is used exactly as the image ships it, so anything +that would have to be applied *to* vllm-ascend is out of reach and a factory +seam has to be used instead. """ from __future__ import annotations @@ -191,8 +189,8 @@ def update_graph_params( num_tokens, vllm_config, speculative_config=None, - # T2/910B compat: num_dcp_pcp_tokens was dropped from the upstream - # signature (vllm-ascend removed PCP from MRV1, see #12592). + # The vllm-ascend build this targets dropped ``num_dcp_pcp_tokens`` from + # the upstream signature (PCP removed from MRV1, see #12592). draft_attn_metadatas=None, ): """Skip the per-step rebind for a step that captured no updatable tasks.""" @@ -359,8 +357,9 @@ def _fixed_kv_prefill_fia(self, query, key, value, attn_metadata, output): the runtime replays the recorded tasks and never re-enters here. The call mirrors the eager ``PrefillNoCache`` branch exactly -- causal FIA against the singleton mask, KV length = query length, no block table -- - with the workspace and the softmax LSE held outside the graph pool so - the capture strands nothing (the fixed-KV workspace-leak lesson). + with the workspace and the softmax LSE held outside the graph pool, so + the capture strands nothing (an allocation made while capturing comes out + of the graph's private pool). """ bucket = int(attn_metadata.fixed_kv_prefill_bucket) fia_kwargs = dict( diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py b/vllm_omni/platforms/npu/attention/fixed_kv_decode.py index 8366f627d51..90f4a440085 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py +++ b/vllm_omni/platforms/npu/attention/fixed_kv_decode.py @@ -16,13 +16,10 @@ # """Fixed-capacity decode attention: full-graph replay without per-step rebinding. -Vendored from vllm-ascend (Apache-2.0) so that vLLM-Omni can deliver it. The -ranked evaluation installs only this tree -- ``pip install -e .`` over our -source into the official image -- and leaves vllm-ascend exactly as the image -ships it, so a patch to vllm-ascend cannot reach a scored run. Everything this -module needs from vllm-ascend is reachable through factories vLLM-Omni already -overrides, so the rest of the change lives in ``fixed_kv_backend.py`` and in -``NPUOmniPlatform``. +Vendored from vllm-ascend (Apache-2.0) so that vLLM-Omni can deliver it +without patching vllm-ascend: everything this module needs from vllm-ascend is +reachable through factories vLLM-Omni already overrides, so the rest of the +change lives in ``fixed_kv_backend.py`` and in ``NPUOmniPlatform``. ``FULL_DECODE_ONLY`` captures a whole decode step, but ``npu_fused_infer_attention_score`` takes ``actual_seq_lengths_kv`` as a host @@ -130,23 +127,22 @@ def install_into_ascend_aclgraph() -> None: baked into the captured tasks, so one batch size holds one graph *per bucket*. Without that, the second capture pass finds the first bucket's entry, replays it instead of capturing, and hangs on the device - synchronize that closes the pass -- round ``20260827T170843Z``, exactly. + synchronize that closes the pass. * It synchronizes the stream before every FULL replay, to order ``update_attn_params`` against the previous one. This path issues no updates, so the barrier is pure cost. ``enable_enpu`` gates that one line - and nothing else in ``__call__`` (stock ``8092d3f6`` line 264), which is - what lets us borrow it instead of copying ~140 lines. + and nothing else in the stock ``__call__``, which is what lets us borrow it + instead of copying the whole wrapper. This wraps ``__call__`` rather than subclassing ``ACLGraphWrapper``, for the same reason :func:`install_into_ascend_backend` patches the backend class: ``current_platform`` is vllm-ascend's ``NPUPlatform``, so ``NPUOmniPlatform.get_graph_wrapper_cls`` is never consulted. - It deliberately does **not** touch ``vllm_ascend.attention.fixed_kv_decode``. - That module, and the ``graph_key()`` call site in the stock wrapper, exist - only where our own ``vllm-ascend-fixed-kv-decode.patch`` has been applied -- - on the A3 box, not in the ranked image. Reaching for it made this look - delivered while it was still leaning on the patch it exists to replace. + It deliberately does **not** touch ``vllm_ascend.attention.fixed_kv_decode``: + that module only exists in vllm-ascend builds carrying the out-of-tree + fixed-KV patch, and leaning on it would leave this change dependent on the + patch it exists to replace. """ from vllm_ascend.compilation import acl_graph @@ -196,9 +192,9 @@ def install_into_ascend_backend() -> None: """Put the omni impl/builder on the class vLLM's selector actually uses. ``current_platform.get_attn_backend_cls`` is vllm-ascend's ``NPUPlatform``, - not ``NPUOmniPlatform``. Round ``20260827T175303Z`` captured extra buckets - then ran the stock builder, so runtime looked up ``(batch, None)``. Patch - the stock backend class in this process instead. + not ``NPUOmniPlatform``. Without this the extra buckets are captured but the + stock builder still runs for them, so runtime looks up ``(batch, None)``. + Patch the stock backend class in this process instead. """ if not is_enabled(): return @@ -237,13 +233,11 @@ def buckets_for(capacity: int, block_size: int) -> tuple[int, ...]: """ if not capacity: return () - # The default has to be a code default, not an environment one: a ranked run - # sets no environment variables. Attention reads the declared capacity every - # step, so a bucket that just covers the sequence is worth having -- the - # Talker's is ~135 (a ~15-token condition plus ~118 codec frames), and 512 - # against 4096 alone measured RTF 0.1592 vs 0.1721 on A3, -7.5%. The full - # capacity is always kept as the top bucket, so a longer sequence still has - # a graph to land on. + # The default has to be a code default, not an environment one: attention + # reads the declared capacity every step, so a bucket that just covers the + # sequence is worth having -- the Talker's is ~135 (a ~15-token condition + # plus ~118 codec frames). The full capacity is always kept as the top + # bucket, so a longer sequence still has a graph to land on. raw = os.getenv(_BUCKETS_ENV, _DEFAULT_BUCKETS) candidates = set() for part in raw.split(","): @@ -290,7 +284,11 @@ def current_capacity() -> int | None: def graph_key() -> int | None: - """Extra ACL graph dispatch key, so one batch size can hold several buckets.""" + """Extra ACL graph dispatch key, so one batch size can hold several buckets. + + Read by vllm-ascend's ``ACLGraphWrapper`` in builds that carry the fixed-KV + hook; nothing in this package calls it directly. + """ return current_capacity() if is_enabled() else None diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py b/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py index a69315a35c8..bde20c30502 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py +++ b/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py @@ -1,16 +1,15 @@ """Captured graphs for the Talker's prefill, bucketed by token count. One Talker prefill forward is ~18.4 ms of host dispatch wrapped around 1-3 ms -of device work (`DEVICE_BOUND_FRAME_20260830.md` §4a': three back-to-back -forwards never accumulate more than 0.07 ms of drain), and the ranked case -runs 1.26 of them per request -- ~23 ms, all but a couple of milliseconds of it -removable. The decode side of this problem was solved by `fixed_kv_decode`; -this module is the same idea for the prefill shape. +of device work (three back-to-back forwards never accumulate more than 0.07 ms +of drain), and a request pays roughly one of them -- all but a couple of +milliseconds of it removable. The decode side of this problem was solved by +`fixed_kv_decode`; this module is the same idea for the prefill shape. ## Why a padded prefill replays exactly -The ranked Talker prefill is a single fresh request of 9-25 tokens (measured -distribution; `{16, 32}` covers it). A graph captured at bucket ``B`` with +The Talker prefill this targets is a single fresh request of 9-25 tokens, so +`{16, 32}` covers it. A graph captured at bucket ``B`` with ``actual_seq_lengths = [B]`` replays a real ``n <= B`` prefill bit-exactly for the rows that matter: @@ -33,12 +32,10 @@ offset 0, ``n`` inside a captured bucket -- and everything else takes today's eager path unchanged. -Off with ``VLLM_OMNI_FIXED_KV_PREFILL=0``. """ from __future__ import annotations -import os from collections.abc import Callable from typing import Any @@ -47,11 +44,9 @@ logger = init_logger(__name__) -_ENV = "VLLM_OMNI_FIXED_KV_PREFILL" -# Per-stage: the Talker's prefill is 9-25 tokens, stage 0's ranked prefill is -# 144-160. The capture routine picks one set per process by model traits. +# The Talker's prefill is 9-25 tokens; a stage that prefills longer selects its +# own set through ``configure``. TALKER_BUCKETS: tuple[int, ...] = (16, 32) -STAGE0_BUCKETS: tuple[int, ...] = (160, 192) BUCKETS: tuple[int, ...] = TALKER_BUCKETS @@ -63,10 +58,6 @@ def configure(buckets: tuple[int, ...]) -> None: _LOGGED_REPLAY = False -def is_enabled() -> bool: - return os.getenv(_ENV, "1") == "1" - - class _State: """Per-process buffers and the captured-bucket registry.""" @@ -87,8 +78,8 @@ def prewarm(device: torch.device) -> None: """Allocate every runtime buffer before any capture can start. An allocation made while a graph is capturing comes out of that graph's - private pool and is stranded there (the fixed-KV workspace leak), so the - slot buffers and aranges exist before `capture` runs. + private pool and is stranded there, so the slot buffers and aranges exist + before any capture runs. """ for bucket in BUCKETS: if bucket not in _state.slot_buffers: @@ -124,8 +115,8 @@ def store_graph(bucket: int, graph: Any, hidden_out: torch.Tensor) -> None: step graph owns its own -- deliberately NOT entries in the ACL graph wrapper. A wrapper-FULL prefill replay poisons the request that follows it: every codec-step replay afterwards blocks the host for the queued - backbone work (+9 ms per decode step, the bisect4/fixtest arms), and - neither the capture order nor the pre-replay barrier explains it. The + backbone work, and neither the capture order nor the pre-replay barrier + explains it. The codec graph itself proves raw replays coexist cleanly with everything, so the prefill graphs live the same way. """ @@ -182,9 +173,9 @@ def maybe_mark_step(attn_metadata: Any, prefill_state: Any, block_size_hint: int # Slots are contiguous only within a block, and a request's block # table covers ceil(n / block_size) blocks -- a bucket that needs one # more block than the prompt does would write into somebody else's. - # (The first stage-0 round hit exactly this: a 148-token prompt spans - # two blocks, and a slot0+arange fill pushed rows past 127 into the - # physical block after block 0 -- runaway text, 7/32 finished.) + # (A prompt that spans two blocks hits exactly this: a slot0+arange fill + # pushes rows past the first block's 127 slots into the physical block + # after it -- runaway text instead of a captured prefill.) return block_table = getattr(attn_metadata, "block_tables", None) if bucket <= block_size: diff --git a/vllm_omni/platforms/npu/native_rejection_sampler.py b/vllm_omni/platforms/npu/native_rejection_sampler.py index b5b881f281a..040581b4aab 100644 --- a/vllm_omni/platforms/npu/native_rejection_sampler.py +++ b/vllm_omni/platforms/npu/native_rejection_sampler.py @@ -5,23 +5,23 @@ vllm-ascend module-patches vllm's ``rejection_sample`` (and its two helpers) onto Triton kernels (``vllm_ascend/patch/worker/patch_rejection_sampler.py``). -On the 910_93 vector core those kernels fault: the 02:38 run pinned acl -507035 on their warmup, which is why ``rejection_sampler`` sits in the -warmup skip set. The skip moves the fault to the first request -- every -verify then JIT-compiles the kernel family on the fly, and both stages -stall there: stage 0 verifies the n-gram drafts of every decode step, and -stage 1 verifies the always-``continue`` drafts behind the K-step -block-table growth. +On the 910_93 vector core those kernels fault (acl 507035) during their +warmup, which is why ``rejection_sampler`` sits in the warmup skip set. The +skip only moves the fault to the first request -- every verify then +JIT-compiles the kernel family on the fly, and both stages stall there: +stage 0 verifies the n-gram drafts of every decode step, and stage 1 +verifies the always-``continue`` drafts behind the K-step block-table +growth. The torch-native implementation the patch replaces is pure tensor ops and exact for both callers, so this module reloads vllm's own sampler source into a fresh module object and re-points the three names at the pre-patch functions. Scoped to the same SoC that skips the warmup; the 910B baseline keeps the Triton kernels it has been measured with -- except when the Talker -K-step is explicitly armed there. That is the one path whose spec width the -910B kernels cannot take (aivec 507035 on every one of the 8 combos tried), -so the restore on a 910B is tied to exactly that arming, never to the SoC -alone: the stock baseline (no arming env) keeps its sampler byte-for-byte. +multi-frame decode is armed there (stage 1's ``speculative_config``). That is +the one path whose spec width the 910B kernels cannot take, so the restore on +a 910B is tied to that arming, never to the SoC alone: a 910B deployment +without the Talker ``speculative_config`` keeps its sampler byte-for-byte. """ import importlib.util diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index ce3d3249ba1..7a799067de7 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -61,10 +61,10 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: Under FULL_DECODE_ONLY the captured decode step has to re-issue attention on every layer on every step, because the op takes the KV length as a host argument that grows each step. On the small Talker decoder that rebind is - ~38% of its busy time. Scoped by ``max_model_len`` (Talker 4096 engages, - Thinker 32768 does not). ``VLLM_OMNI_FIXED_KV_DECODE=0`` restores stock. - The 910B family also needs it to capture at all: the stock builder - faults there on spec-width rows. + most of its busy time. Scoped by ``max_model_len`` (Talker 4096 engages, + a 32768-token Thinker does not). ``VLLM_OMNI_FIXED_KV_DECODE=0`` restores + stock. The 910B family also needs it to capture at all: with the K-step + armed the stock capture path faults there (acl 507035). """ resolved = super().get_attn_backend_cls(selected_backend, attn_selector_config, num_heads) if resolved != "vllm_ascend.attention.attention_v1.AscendAttentionBackend": @@ -80,12 +80,20 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: block_size = getattr(getattr(vllm_config, "cache_config", None), "block_size", None) capacity = fixed_kv_decode.capacity_for(max_model_len, block_size) if max_model_len and block_size else None fixed_kv_decode.install_into_ascend_aclgraph() - logger.info( - "[minicpmo] fixed-KV decode attention on (max_model_len=%s, kv_capacity=%s); " - "set VLLM_OMNI_FIXED_KV_DECODE=0 to disable", - max_model_len, - capacity, - ) + if capacity: + logger.info( + "[minicpmo] fixed-KV decode attention on (max_model_len=%s, kv_capacity=%s); " + "set VLLM_OMNI_FIXED_KV_DECODE=0 to disable", + max_model_len, + capacity, + ) + else: + # The backend class is still returned; it declines for this + # max_model_len, so no bucket is captured and nothing engages. + logger.info( + "[minicpmo] fixed-KV decode attention not engaged (max_model_len=%s is above the context ceiling)", + max_model_len, + ) return "vllm_omni.platforms.npu.attention.fixed_kv_backend.OmniFixedKVAttentionBackend" # conv2d convolution operator in the code2wav module of Qwen3-TTS not being able to run on Aclnn diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index b2b652f1f4e..508efda86a0 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -1063,12 +1063,21 @@ def sample_tokens( self.input_batch.sampling_metadata.logitsprocs, logits.shape[-1], ) - if getattr(self.model, "supports_multi_frame_decode", False): - # K-step only: this layer censors the stop token itself, and its - # release condition is a spec-decoding bookkeeping counter we do - # not control -- a request whose count never advances can never - # stop. The model already masks the codec EOS by its own frame - # count (talker_codec_sample), so clear this list every step. + # K-step only, and both halves matter: the model class has to + # implement the multi-frame path (``supports_multi_frame_decode``) + # *and* this deployment has to arm it -- stage 1's speculative_config + # is what sets ``num_spec_tokens``. With the loop off, vLLM's + # min_tokens layer behaves as it does on any other stage and must + # keep working, so it is only cleared when both hold. + # + # Why clear it at all: this layer censors the stop token itself, and + # its release condition is a spec-decoding bookkeeping counter we do + # not control -- a request whose count never advances can never + # stop. The model already masks the codec EOS by its own frame count + # (talker_codec_sample), so the vLLM-level layer is redundant here. + if getattr(self.model, "supports_multi_frame_decode", False) and int( + getattr(self, "num_spec_tokens", 0) or 0 + ) > 0: from vllm_omni.platforms.npu.worker import talker_multiframe talker_multiframe.neutralize_kstep_min_tokens( diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 84a0c934d3d..dffbaa05aa6 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -482,12 +482,12 @@ def _model_forward( run_model = partial(self.model, **model_inputs) # K-step (Talker multi-frame decode). A decode step that schedules K - # positions per request belongs to ``talker_multiframe``: it replays the - # captured decode graph K times -- refreshing the graph's parameters - # before every replay, which is what keeps positions and embeddings in - # step with the frame being produced -- and samples one codec frame per - # replay. Everything else (prefill, a mixed step, another stage) falls - # through to the ordinary single-forward path below. + # positions per request belongs to ``talker_multiframe``: it replays a + # captured decode graph once per frame and samples one codec frame per + # replay. The one-query path writes each frame's embedding straight into + # the buffer the graph reads; the wide path refreshes the per-frame + # inputs instead. Everything else (prefill, a mixed step, another stage) + # falls through to the ordinary single-forward path below. from vllm_omni.platforms.npu.worker import talker_multiframe frames = talker_multiframe.applies(self.model, model_kwargs_extra) @@ -526,20 +526,18 @@ def _model_forward( ) finally: talker_multiframe.end_narrow_step(forward_context, narrow) - # Restore the vocab gate the ed8c7db8 refactor dropped along with - # the _model_forward override it lived in. The Talker's vLLM-level - # head is the two-wide continue/stop row, but input_batch reports - # vocab_size=0, and the rejection sampler's parse_output then masks - # every accepted token as out-of-vocab: the request comes back with - # nothing, the scheduler never rolls back the K tokens it advanced, - # and the next steps schedule a negative count -- the engine hangs - # right after the first 'engaged' step. Single-frame steps never - # took the branch that reads vocab_size; only the K-step does. - # The gate needs one thing from this call: that this really is a - # multi-frame decode step. The width it writes is the two-wide stop - # row, never the hidden width this tensor happens to carry. Hand it - # the stop rows when they are at hand so the intent is visible, and - # keep the hidden states only as the "this is a K-step" signal. + # The vocab gate the multi-frame path needs, applied here because + # this is the wrapper the K-step runs through. The Talker's + # vLLM-level head is the two-wide continue/stop row, but input_batch + # reports vocab_size=0, and the rejection sampler's parse_output + # then masks every accepted token as out-of-vocab: the request comes + # back with nothing, the scheduler never rolls back the K tokens it + # advanced, and the next steps schedule a negative count -- the + # engine hangs right after the first 'engaged' step. Single-frame + # steps never took the branch that reads vocab_size; only the K-step + # does. The gate takes only "not None" from this argument: the width + # it writes is the two-wide stop row, never the width of whatever + # tensor it is handed. stop_rows = getattr(self.model, "_batch_stop_logits", None) talker_multiframe.ensure_stop_token_vocab( self, diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py index 8398742d433..c577d3ad1b2 100644 --- a/vllm_omni/platforms/npu/worker/talker_multiframe.py +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -1,10 +1,10 @@ """Several codec frames inside one stage-1 ``execute_model``. -Measured on A3 / 910C (``server-env/design/RTF_CEILING_20260828.md``): one -MiniCPM-o Talker decode step is ~2.9 ms of vLLM host work wrapped around a -0.83 ms device forward, for a 190M-parameter model that emits exactly one codec -frame. Over ~118 frames per request that host work is ~340 ms of the ~540 ms -that follows TTFT -- by a wide margin the largest single term left in the score. +One MiniCPM-o Talker decode step is almost all vLLM host work wrapped around a +small device forward: a 190M-parameter model that emits exactly one codec +frame, but pays scheduler, metadata, sampling and IPC costs per step. Over the +~118 frames of a typical request that host work is the largest single term +after TTFT. None of it is per *frame*. The scheduler step, input preparation, attention metadata, sampling, output assembly and engine-core IPC are per *step*, and a @@ -27,9 +27,8 @@ them, and the next replay overwrites them with the real thing. After the K'th replay every row and every KV entry is what a K-step decode would have written. -A replay of the 20-layer forward is 0.83 ms of device and 8 us of host -(``server-env/tools/dispatch_cost.py``), so K replays cost what K frames cost -and the ~2.9 ms is paid once. At K=4 that is a cadence of ~1.9 ms against ~4.0. +A replay of the 20-layer forward is 0.83 ms of device and microseconds of host, +so K replays cost what K frames cost and the per-step host work is paid once. Nothing in the loop touches the host: the embedding lookup, the codec sample (the captured codec-sampling step), the stop row and the emitted delta are all device @@ -326,7 +325,7 @@ def begin_narrow_step( fixed_kv_decode = None if fixed_kv_decode is None: - return _decline("fixed-KV decode (item 6) is not present in this tree") + return _decline("the fixed-KV decode backend is not present in this build") if not narrow_replay_enabled(runner) or inputs_embeds is None or positions is None: return None rows = len(spans) @@ -617,17 +616,14 @@ def ensure_stop_token_vocab(runner: Any, logits: Any) -> None: head is the two-wide continue/stop row `compute_logits` builds, so two is what every `vocab_size` reader has to see. - The caller used to hand over `text_hidden_states`, so this landed on the - hidden width (768). That passes the `parse_output` filter by accident (0 - and 1 are both below 768) and breaks the one reader that repays checking: - `InputBatch.add_request` stores `top_k = vocab_size` as its "no top-k" - sentinel, which only holds for 0 and 2. At 768 `top_k_reqs` stops being - empty, `sampling_metadata.top_k` stops being None, and the Ascend - `enable_reduce_sample` branch starts treating a two-column distribution as - a 768-way one. The stop row then does not survive into the request's token - list, so the request runs to `max_tokens` instead of stopping on EOS -- and - the frame stream that keeps arriving after the codec sequence ended is what - eventually hands the scheduler a step it cannot merge. + The width has to be the row the model emits, not the width of whatever + tensor the caller happens to be holding: `InputBatch.add_request` stores + `top_k = vocab_size` as its "no top-k" sentinel, which only holds for 0 and + 2. A hidden width (768) instead passes the `parse_output` filter by accident + but leaves `top_k_reqs` non-empty, and the Ascend `enable_reduce_sample` + branch then treats the two-column distribution as a 768-way one -- the stop + row does not survive into the request's token list and the request runs to + `max_tokens`. """ if logits is None or not getattr(runner.model, "supports_multi_frame_decode", False): return From 3de892c23f0ab3738bd855310f627fecb52a2240 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 12:00:21 +0800 Subject: [PATCH 29/89] refactor(npu): drop the fixed-KV decode context-length switch ``VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN`` decided the same thing the deploy YAML already decides: the fixed-KV path engages when a stage's ``max_model_len`` is small enough to have fixed-capacity graphs. A second knob for that rule can only be set to a value the captured shapes do not support, so the ceiling is a code constant now and the stage's own ``max_model_len`` is what selects the outcome (Talker 4096 engages, a 32768-token Thinker keeps the stock path). No behavior change at the default: the constant is the old default 8192. The inventory, the configuration reference and the snapshot counts are updated with it. Signed-off-by: y-null --- docs/configuration/environment_variables.md | 3 +-- tests/config/test_environment_variables.py | 2 +- vllm_omni/config/environment_variable_inventory.py | 1 - .../platforms/npu/attention/fixed_kv_decode.py | 14 +++++++++----- 4 files changed, 11 insertions(+), 9 deletions(-) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index e6ecb1092c6..2338beebb87 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -123,8 +123,7 @@ settings. | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | | --- | --- | --- | --- | --- | | `VLLM_OMNI_FIXED_KV_DECODE` | `0` disables; default `1` (on) | Fixed-capacity decode attention backend and its ACL-graph hook; read when the NPU platform resolves the attention backend | Only the exact value `0` disables it. Scoped by `max_model_len`: a long-context stage (e.g. a 32768-token Thinker) declines and keeps the stock path. | Experimental performance control | -| `VLLM_OMNI_FIXED_KV_DECODE_BUCKETS` | Comma-separated KV capacities in tokens; default `512,4096` | Capture buckets of the fixed-KV decode graphs; read when the buckets are computed | The full capacity is always kept as the top bucket, so a longer sequence still has a graph. Empty or malformed entries are skipped. | Experimental performance control | -| `VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN` | Integer tokens; default `8192` | Fixed-KV decode eligibility; read per step when attention decides whether to engage | At or below this length the fixed-KV path is used; above it the stock path runs. | Experimental performance control | +| `VLLM_OMNI_FIXED_KV_DECODE_BUCKETS` | Comma-separated KV capacities in tokens; default `512` | Capture buckets of the fixed-KV decode graphs; read when the buckets are computed | The full capacity is always kept as the top bucket, so a longer sequence still has a graph. Empty or malformed entries are skipped. | Experimental performance control | | `VLLM_OMNI_NPU_SKIP_WARMUPS` | Comma-separated warmup-name substrings, or `none`; default `penalties,rejection_sampler` on 910_93 and `rejection_sampler` on 910B | vllm-ascend Triton kernel warmups on the NPU worker; resolved per warmup call | Skips the warmups whose dummy-token setup faults the vector core (acl 507035). `none` restores the stock warmups. | Experimental | | `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control | | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 4df7e92b5f1..5507b98ce72 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 35, + EnvironmentVariableCategory.PUBLIC_OMNI: 34, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 61, diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index abbf4769e6b..46aacd9332e 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -69,7 +69,6 @@ def is_public_omni(self) -> bool: "VLLM_OMNI_EVENT_DRIVEN_ORCH", "VLLM_OMNI_FIXED_KV_DECODE", "VLLM_OMNI_FIXED_KV_DECODE_BUCKETS", - "VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN", "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "VLLM_OMNI_NIXL_LEASE_S", diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py b/vllm_omni/platforms/npu/attention/fixed_kv_decode.py index 90f4a440085..87bd03ba57a 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py +++ b/vllm_omni/platforms/npu/attention/fixed_kv_decode.py @@ -43,8 +43,6 @@ Env switches: ``VLLM_OMNI_FIXED_KV_DECODE`` 0 to disable (default 1) - ``VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN`` only engage at or below this - ``max_model_len`` (default 8192) ``VLLM_OMNI_FIXED_KV_DECODE_BUCKETS`` comma-separated candidate capacities, e.g. "1024,4096"; default "512". Each model keeps the candidates below its own @@ -91,12 +89,19 @@ _runtime_bucket: int | None = None +_ENV = "VLLM_OMNI_FIXED_KV_DECODE" _BUCKETS_ENV = "VLLM_OMNI_FIXED_KV_DECODE_BUCKETS" _DEFAULT_BUCKETS = "512" +# Only a small decoder has fixed-capacity graphs worth capturing; a long-context +# stage (a 32768-token Thinker) keeps the stock path. The ceiling is a property of +# the shapes that get captured, not a per-deployment preference, so it is a code +# constant: the stage's own ``max_model_len`` is what selects the outcome. +_MAX_CONTEXT = 8192 + def is_enabled() -> bool: - return os.getenv("VLLM_OMNI_FIXED_KV_DECODE", "1") == "1" + return os.getenv(_ENV, "1") == "1" class _CapacityEntries(dict): @@ -218,8 +223,7 @@ def capacity_for(max_model_len: int, block_size: int) -> int: """Full KV capacity this decoder would run at, or 0 when it should not engage.""" if not is_enabled(): return 0 - limit = int(os.getenv("VLLM_OMNI_FIXED_KV_DECODE_MAX_LEN", "8192")) - if max_model_len > limit: + if max_model_len > _MAX_CONTEXT: return 0 # Round to the block table, not to max_model_len: the op indexes whole blocks. return -(-int(max_model_len) // int(block_size)) * int(block_size) From ab8c2e6079482172b31acc5350f59e285fff1018 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 12:02:25 +0800 Subject: [PATCH 30/89] refactor(npu): name the attention modules after what they do ``fixed_kv_*`` named one decode-side op argument, and the prefill module does not fix a KV length at all -- it buckets the prefill by token count. What all three do is capture a step whose shapes are static so that a replay needs nothing from the host, so they are named for that: fixed_kv_decode.py -> static_shape_decode.py fixed_kv_prefill.py -> static_shape_prefill.py fixed_kv_backend.py -> static_shape_backend.py The symbols follow the files: OmniStaticShapeAttentionBackend and its impl, OmniStaticShapeMetadataBuilder, the ``_static_shape_*`` helpers and the ``static_shape_*`` attention-metadata attributes. The reference to vllm-ascend's own ``vllm_ascend.attention.fixed_kv_decode`` module stays as it is. No behavior change. Signed-off-by: y-null --- ..._kv_backend.py => static_shape_backend.py} | 102 +++++++++--------- ...ed_kv_decode.py => static_shape_decode.py} | 42 ++++---- ..._kv_prefill.py => static_shape_prefill.py} | 8 +- vllm_omni/platforms/npu/platform.py | 16 +-- .../npu/worker/npu_ar_model_runner.py | 18 ++-- .../platforms/npu/worker/talker_multiframe.py | 16 +-- 6 files changed, 101 insertions(+), 101 deletions(-) rename vllm_omni/platforms/npu/attention/{fixed_kv_backend.py => static_shape_backend.py} (82%) rename vllm_omni/platforms/npu/attention/{fixed_kv_decode.py => static_shape_decode.py} (93%) rename vllm_omni/platforms/npu/attention/{fixed_kv_prefill.py => static_shape_prefill.py} (96%) diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_backend.py b/vllm_omni/platforms/npu/attention/static_shape_backend.py similarity index 82% rename from vllm_omni/platforms/npu/attention/fixed_kv_backend.py rename to vllm_omni/platforms/npu/attention/static_shape_backend.py index 8f9a8d5a0de..8190d33f4c3 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_backend.py +++ b/vllm_omni/platforms/npu/attention/static_shape_backend.py @@ -10,7 +10,7 @@ 20-layer, 768-hidden MiniCPM-o Talker that rebind *is* the step: it dominates the stage-1 main-thread wall, which is what this module removes. -The mechanism lives in :mod:`fixed_kv_decode`; this module is only the seam. It +The mechanism lives in :mod:`static_shape_decode`; this module is only the seam. It is a subclass rather than a patch because every place vllm-ascend needs to be told about it is a factory vLLM-Omni already controls: @@ -20,7 +20,7 @@ NPUOmniPlatform.get_attn_backend_cls -> selects this NPUARModelRunner._capture_cudagraphs -> captures buckets ACLGraphWrapper.__call__ key + replay barrier -> wrapped by - :func:`fixed_kv_decode.install_into_ascend_aclgraph` + :func:`static_shape_decode.install_into_ascend_aclgraph` That matters: vllm-ascend is used exactly as the image ships it, so anything that would have to be applied *to* vllm-ascend is out of reach and a factory @@ -41,18 +41,18 @@ AscendAttentionState, ) -from vllm_omni.platforms.npu.attention import fixed_kv_decode, fixed_kv_prefill +from vllm_omni.platforms.npu.attention import static_shape_decode, static_shape_prefill logger = init_logger(__name__) -class OmniFixedKVMetadataBuilder(AscendAttentionMetadataBuilder): +class OmniStaticShapeMetadataBuilder(AscendAttentionMetadataBuilder): """Attach the device-side state a captured step needs, and pick the bucket.""" def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) - self._omni_buckets = fixed_kv_decode.buckets_for( - fixed_kv_decode.capacity_for( + self._omni_buckets = static_shape_decode.buckets_for( + static_shape_decode.capacity_for( self.vllm_config.model_config.max_model_len, self.vllm_config.cache_config.block_size, ), @@ -70,7 +70,7 @@ def build(self, *args, **kwargs): # Positional call: (common_prefix_len, common_attn_metadata, ...). common = next((a for a in args if hasattr(a, "block_table_tensor")), None) if common is None: - fixed_kv_decode.set_runtime_bucket(None) + static_shape_decode.set_runtime_bucket(None) return attn_metadata # `seq_lens` on the metadata is a CPU tensor; a captured graph has to @@ -79,11 +79,11 @@ def build(self, *args, **kwargs): # build() materializes padded copies of seq_lens/block_table for the FIA # batch-size check when the batch is full. Those are fresh tensors every # step, so a graph cannot bake them in. The block table is the one that - # is not caught by the row check in `_fixed_kv_applies`; identity against + # is not caught by the row check in `_static_shape_applies`; identity against # the runtime's own persistent buffer is the reliable test. attn_metadata.fia_inputs_persistent = attn_metadata.block_tables is common.block_table_tensor - attn_metadata.fixed_kv_mask_recorded = False - attn_metadata.fixed_kv_pse_recorded = False + attn_metadata.static_shape_mask_recorded = False + attn_metadata.static_shape_pse_recorded = False # Whether this step is a decode cannot be read off `attn_state`. A step # that schedules several tokens per request comes back labelled @@ -93,23 +93,23 @@ def build(self, *args, **kwargs): # the same number of query positions for every request, none of them # still consuming a prompt -- and it is what the captured graph and the # causal bias actually require. - attn_metadata.fixed_kv_decode_step = ( + attn_metadata.static_shape_decode_step = ( attn_metadata.attn_state == AscendAttentionState.DecodeOnly or self._is_uniform_multi_query_decode(attn_metadata, common) ) # Settle the capacity before the forward runs: it is part of the ACL # graph dispatch key, so the wrapper has to be able to look it up. - if self._omni_buckets and attn_metadata.fixed_kv_decode_step: + if self._omni_buckets and attn_metadata.static_shape_decode_step: seq_lens_list = attn_metadata.seq_lens_list or [1] max_seq_len = getattr(common, "max_seq_len", None) or max(seq_lens_list) - fixed_kv_decode.set_runtime_bucket(fixed_kv_decode.select_bucket(int(max_seq_len), self._omni_buckets)) + static_shape_decode.set_runtime_bucket(static_shape_decode.select_bucket(int(max_seq_len), self._omni_buckets)) else: - fixed_kv_decode.set_runtime_bucket(None) + static_shape_decode.set_runtime_bucket(None) # A single fresh prefill inside a captured bucket replays a graph # instead of paying ~18 ms of eager dispatch; the hook refuses # everything else (chunked, batched, oversized, other stages). - fixed_kv_prefill.maybe_mark_step( + static_shape_prefill.maybe_mark_step( attn_metadata, AscendAttentionState.PrefillNoCache, int(self.vllm_config.cache_config.block_size), @@ -138,7 +138,7 @@ def _is_uniform_multi_query_decode(self, attn_metadata, common) -> bool: return not bool(is_prefilling[:num_reqs].any()) -class OmniFixedKVAttentionBackendImpl(AscendAttentionBackendImpl): +class OmniStaticShapeAttentionBackendImpl(AscendAttentionBackendImpl): """Decode attention whose host-side arguments never change between steps.""" def __init__(self, *args, **kwargs) -> None: @@ -149,14 +149,14 @@ def __init__(self, *args, **kwargs) -> None: for _flag in ("enable_c8_quant", "enable_hamming_sparse"): if not hasattr(self, _flag): setattr(self, _flag, False) - self._omni_capacity = fixed_kv_decode.capacity_for( + self._omni_capacity = static_shape_decode.capacity_for( self.vllm_config.model_config.max_model_len, self.vllm_config.cache_config.block_size, ) # +2 matches the query_start_loc buffer, which carries the FIA padding # request on a full batch. self._omni_max_rows = int(self.vllm_config.scheduler_config.max_num_seqs) + 2 - self._omni_buckets = fixed_kv_decode.buckets_for(self._omni_capacity, self.vllm_config.cache_config.block_size) + self._omni_buckets = static_shape_decode.buckets_for(self._omni_capacity, self.vllm_config.cache_config.block_size) # A step that schedules several tokens per sequence (the Talker's # multi-frame decode declares them through a speculative_config) is # still a uniform decode, and the fixed-capacity mask covers it as @@ -173,7 +173,7 @@ def __init__(self, *args, **kwargs) -> None: if self._omni_capacity: # At model load, not during capture: an allocation made while a graph # is capturing comes out of that graph's private pool. - fixed_kv_decode.prewarm( + static_shape_decode.prewarm( self._omni_max_rows, self.num_heads, self._omni_buckets, @@ -194,10 +194,10 @@ def update_graph_params( draft_attn_metadatas=None, ): """Skip the per-step rebind for a step that captured no updatable tasks.""" - if fixed_kv_decode.is_fixed_graph(num_tokens): + if static_shape_decode.is_static_shape_graph(num_tokens): metadata = forward_context.attn_metadata first = next(iter(metadata.values())) if isinstance(metadata, dict) else metadata - fixed_kv_decode.sync_captured_inputs( + static_shape_decode.sync_captured_inputs( num_tokens, getattr(first, "seq_lens_device", None), getattr(first, "block_tables", None), @@ -212,11 +212,11 @@ def update_graph_params( draft_attn_metadatas, ) - def _fixed_kv_applies(self, attn_metadata, block_table, block_size) -> bool: + def _static_shape_applies(self, attn_metadata, block_table, block_size) -> bool: """Whether this layer can capture a decode step with constant host args.""" if not self._omni_capacity: return False - if not getattr(attn_metadata, "fixed_kv_decode_step", False): + if not getattr(attn_metadata, "static_shape_decode_step", False): return False if not attn_metadata.causal or not getattr(attn_metadata, "fia_inputs_persistent", False): return False @@ -230,7 +230,7 @@ def _fixed_kv_applies(self, attn_metadata, block_table, block_size) -> bool: return False if _EXTRA_CTX.is_draft_model: return False - capacity = fixed_kv_decode.current_capacity() + capacity = static_shape_decode.current_capacity() if capacity is None or capacity not in self._omni_buckets: return False if block_size * block_table.shape[1] < capacity: @@ -267,7 +267,7 @@ def _uniform_query_len(self, attn_metadata) -> int | None: return None return q_len - def _fixed_kv_graph_fia( + def _static_shape_graph_fia( self, query: torch.Tensor, key: torch.Tensor, @@ -285,19 +285,19 @@ def _fixed_kv_graph_fia( sequence length arrives through ``pse_shift``, a device tensor the graph refreshes for itself. """ - capacity = fixed_kv_decode.current_capacity() + capacity = static_shape_decode.current_capacity() rows = len(attn_metadata.actual_seq_lengths_q) q_len = self._uniform_query_len(attn_metadata) or 1 - pse_all = fixed_kv_decode.get_pse_buffer( + pse_all = static_shape_decode.get_pse_buffer( self._omni_max_rows, self.num_heads, capacity, query.dtype, query.device, q_len ) pse = pse_all[:rows] seq_lens_device = attn_metadata.seq_lens_device - if not attn_metadata.fixed_kv_mask_recorded: - attn_metadata.fixed_kv_mask_recorded = True - fixed_kv_decode.remember_captured_inputs( + if not attn_metadata.static_shape_mask_recorded: + attn_metadata.static_shape_mask_recorded = True + static_shape_decode.remember_captured_inputs( num_tokens, capacity, seq_lens_device, @@ -305,7 +305,7 @@ def _fixed_kv_graph_fia( getattr(attn_metadata, "slot_mapping", None), ) logger.info( - "[minicpmo] fixed-KV decode capture: graph_size=%d rows=%d q_len=%d " + "[minicpmo] static-shape decode capture: graph_size=%d rows=%d q_len=%d " "kv_capacity=%d heads=%d attention=FIA", num_tokens, rows, @@ -316,11 +316,11 @@ def _fixed_kv_graph_fia( # The bias is per step, not per layer, and only the FIA path reads it, # so it is recorded by the first layer that actually wants it. - if not attn_metadata.fixed_kv_pse_recorded: - fixed_kv_decode.emit_mask_refresh(pse, seq_lens_device, self.num_heads, capacity, q_len) - attn_metadata.fixed_kv_pse_recorded = True + if not attn_metadata.static_shape_pse_recorded: + static_shape_decode.emit_mask_refresh(pse, seq_lens_device, self.num_heads, capacity, q_len) + attn_metadata.static_shape_pse_recorded = True - softmax_lse = fixed_kv_decode.keep_alive( + softmax_lse = static_shape_decode.keep_alive( num_tokens, capacity, torch.empty(1, dtype=query.dtype, device=query.device) ) fia_kwargs = dict( @@ -341,7 +341,7 @@ def _fixed_kv_graph_fia( pre_tokens=SWA_INT_MAX, next_tokens=0, ) - workspace = fixed_kv_decode.graph_workspace( + workspace = static_shape_decode.graph_workspace( num_tokens, capacity, (rows, self.num_heads, self.num_kv_heads, self.head_size, block_size, capacity, query.dtype), @@ -350,7 +350,7 @@ def _fixed_kv_graph_fia( torch_npu.npu_fused_infer_attention_score.out(workspace=workspace, out=[output, softmax_lse], **fia_kwargs) return output.view(num_tokens, self.num_heads, self.head_size), num_tokens - def _fixed_kv_prefill_fia(self, query, key, value, attn_metadata, output): + def _static_shape_prefill_fia(self, query, key, value, attn_metadata, output): """Capture one prefill attention layer with all-constant host arguments. Runs only while a graph is being captured (and during its warmups): @@ -361,7 +361,7 @@ def _fixed_kv_prefill_fia(self, query, key, value, attn_metadata, output): the capture strands nothing (an allocation made while capturing comes out of the graph's private pool). """ - bucket = int(attn_metadata.fixed_kv_prefill_bucket) + bucket = int(attn_metadata.static_shape_prefill_bucket) fia_kwargs = dict( query=query[:bucket], key=key[:bucket], @@ -376,21 +376,21 @@ def _fixed_kv_prefill_fia(self, query, key, value, attn_metadata, output): scale=self.scale, sparse_mode=3, ) - workspace = fixed_kv_prefill.workspace( + workspace = static_shape_prefill.workspace( bucket, lambda: torch_npu._npu_fused_infer_attention_score_get_max_workspace(**fia_kwargs), ) - lse = fixed_kv_prefill.lse_buffer(bucket, query.dtype, query.device) + lse = static_shape_prefill.lse_buffer(bucket, query.dtype, query.device) torch_npu.npu_fused_infer_attention_score.out(workspace=workspace, out=[output, lse], **fia_kwargs) return output.view(bucket, self.num_heads, self.head_size), bucket def full_graph_fia(self, query, key, value, attn_metadata, output, layer=None): - if getattr(attn_metadata, "fixed_kv_prefill_bucket", None) is not None: - return self._fixed_kv_prefill_fia(query, key, value, attn_metadata, output) + if getattr(attn_metadata, "static_shape_prefill_bucket", None) is not None: + return self._static_shape_prefill_fia(query, key, value, attn_metadata, output) if self._omni_capacity: fixed_key, fixed_value, block_size, block_table, _ = self._get_fia_params(key, value, attn_metadata) - if self._fixed_kv_applies(attn_metadata, block_table, block_size): - return self._fixed_kv_graph_fia( + if self._static_shape_applies(attn_metadata, block_table, block_size): + return self._static_shape_graph_fia( query, fixed_key, fixed_value, @@ -402,17 +402,17 @@ def full_graph_fia(self, query, key, value, attn_metadata, output, layer=None): ) # `_get_fia_params` only reshapes cache views, so recomputing it in the # base implementation costs nothing beyond two `view` calls, and only on - # the paths fixed-KV declines. + # the paths static-shape declines. return super().full_graph_fia(query, key, value, attn_metadata, output, layer) -class OmniFixedKVAttentionBackend(AscendAttentionBackend): - """The stock Ascend backend with the fixed-KV decode impl and builder.""" +class OmniStaticShapeAttentionBackend(AscendAttentionBackend): + """The stock Ascend backend with the static-shape decode impl and builder.""" @staticmethod - def get_impl_cls() -> type[OmniFixedKVAttentionBackendImpl]: - return OmniFixedKVAttentionBackendImpl + def get_impl_cls() -> type[OmniStaticShapeAttentionBackendImpl]: + return OmniStaticShapeAttentionBackendImpl @staticmethod - def get_builder_cls() -> type[OmniFixedKVMetadataBuilder]: - return OmniFixedKVMetadataBuilder + def get_builder_cls() -> type[OmniStaticShapeMetadataBuilder]: + return OmniStaticShapeMetadataBuilder diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py b/vllm_omni/platforms/npu/attention/static_shape_decode.py similarity index 93% rename from vllm_omni/platforms/npu/attention/fixed_kv_decode.py rename to vllm_omni/platforms/npu/attention/static_shape_decode.py index 87bd03ba57a..ca552d2377c 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_decode.py +++ b/vllm_omni/platforms/npu/attention/static_shape_decode.py @@ -14,12 +14,12 @@ # limitations under the License. # This file is a part of the vllm-ascend project. # -"""Fixed-capacity decode attention: full-graph replay without per-step rebinding. +"""Static-shape decode attention: full-graph replay without per-step rebinding. Vendored from vllm-ascend (Apache-2.0) so that vLLM-Omni can deliver it without patching vllm-ascend: everything this module needs from vllm-ascend is reachable through factories vLLM-Omni already overrides, so the rest of the -change lives in ``fixed_kv_backend.py`` and in ``NPUOmniPlatform``. +change lives in ``static_shape_backend.py`` and in ``NPUOmniPlatform``. ``FULL_DECODE_ONLY`` captures a whole decode step, but ``npu_fused_infer_attention_score`` takes ``actual_seq_lengths_kv`` as a host @@ -105,7 +105,7 @@ def is_enabled() -> bool: class _CapacityEntries(dict): - """The wrapper's entries dict, plus one dict per fixed-KV capacity. + """The wrapper's entries dict, plus one dict per static-shape capacity. Sleep mode's ``reset_all_graph_params`` clears whatever dict is installed on the wrapper, so the per-capacity dicts hang off it rather than off the @@ -126,7 +126,7 @@ def clear(self) -> None: def install_into_ascend_aclgraph() -> None: """Give the image's ACLGraphWrapper a per-capacity key and drop a dead barrier. - Two things the stock wrapper does that are wrong for fixed-KV decode: + Two things the stock wrapper does that are wrong for static-shape decode: * It keys captured graphs on ``batch_descriptor`` alone. The KV capacity is baked into the captured tasks, so one batch size holds one graph *per @@ -146,13 +146,13 @@ def install_into_ascend_aclgraph() -> None: It deliberately does **not** touch ``vllm_ascend.attention.fixed_kv_decode``: that module only exists in vllm-ascend builds carrying the out-of-tree - fixed-KV patch, and leaning on it would leave this change dependent on the + static-shape patch, and leaning on it would leave this change dependent on the patch it exists to replace. """ from vllm_ascend.compilation import acl_graph wrapper_cls = acl_graph.ACLGraphWrapper - if getattr(wrapper_cls, "_omni_fixed_kv_call", False): + if getattr(wrapper_cls, "_omni_static_shape_call", False): return original_call = wrapper_cls.__call__ @@ -170,7 +170,7 @@ def _call(self, *args, **kwargs): # Only once the capture has recorded this (size, capacity) is there a # fixed graph to replay with nothing to order against; during the # capture itself the barrier still has to stand. - if not saved_enpu and _is_fixed_graph_here(capacity): + if not saved_enpu and _is_static_shape_graph_here(capacity): self.enable_enpu = True try: return original_call(self, *args, **kwargs) @@ -179,12 +179,12 @@ def _call(self, *args, **kwargs): self.enable_enpu = saved_enpu wrapper_cls.__call__ = _call - wrapper_cls._omni_fixed_kv_call = True - logger.info("[minicpmo] fixed-KV: ACLGraphWrapper keyed per KV capacity, pre-replay barrier dropped") + wrapper_cls._omni_static_shape_call = True + logger.info("[minicpmo] static-shape: ACLGraphWrapper keyed per KV capacity, pre-replay barrier dropped") -def _is_fixed_graph_here(capacity: int) -> bool: - """``is_fixed_graph`` for the batch size the forward context is about to run.""" +def _is_static_shape_graph_here(capacity: int) -> bool: + """``is_static_shape_graph`` for the batch size the forward context is about to run.""" from vllm.forward_context import get_forward_context descriptor = getattr(get_forward_context(), "batch_descriptor", None) @@ -205,18 +205,18 @@ def install_into_ascend_backend() -> None: return from vllm_ascend.attention.attention_v1 import AscendAttentionBackend - from vllm_omni.platforms.npu.attention.fixed_kv_backend import ( - OmniFixedKVAttentionBackendImpl, - OmniFixedKVMetadataBuilder, + from vllm_omni.platforms.npu.attention.static_shape_backend import ( + OmniStaticShapeAttentionBackendImpl, + OmniStaticShapeMetadataBuilder, ) - if getattr(AscendAttentionBackend, "_omni_fixed_kv_installed", False): + if getattr(AscendAttentionBackend, "_omni_static_shape_installed", False): return - AscendAttentionBackend.get_impl_cls = staticmethod(lambda: OmniFixedKVAttentionBackendImpl) - AscendAttentionBackend.get_builder_cls = staticmethod(lambda: OmniFixedKVMetadataBuilder) - AscendAttentionBackend._omni_fixed_kv_installed = True + AscendAttentionBackend.get_impl_cls = staticmethod(lambda: OmniStaticShapeAttentionBackendImpl) + AscendAttentionBackend.get_builder_cls = staticmethod(lambda: OmniStaticShapeMetadataBuilder) + AscendAttentionBackend._omni_static_shape_installed = True install_into_ascend_aclgraph() - logger.info("[minicpmo] installed omni fixed-KV impl/builder onto AscendAttentionBackend") + logger.info("[minicpmo] installed omni static-shape impl/builder onto AscendAttentionBackend") def capacity_for(max_model_len: int, block_size: int) -> int: @@ -290,7 +290,7 @@ def current_capacity() -> int | None: def graph_key() -> int | None: """Extra ACL graph dispatch key, so one batch size can hold several buckets. - Read by vllm-ascend's ``ACLGraphWrapper`` in builds that carry the fixed-KV + Read by vllm-ascend's ``ACLGraphWrapper`` in builds that carry the static-shape hook; nothing in this package calls it directly. """ return current_capacity() if is_enabled() else None @@ -449,7 +449,7 @@ def captured_seq_lens(num_tokens: int) -> torch.Tensor | None: return None if captured is None else captured[0] -def is_fixed_graph(num_tokens: int) -> bool: +def is_static_shape_graph(num_tokens: int) -> bool: capacity = current_capacity() return capacity is not None and (num_tokens, capacity) in _captured_inputs diff --git a/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py b/vllm_omni/platforms/npu/attention/static_shape_prefill.py similarity index 96% rename from vllm_omni/platforms/npu/attention/fixed_kv_prefill.py rename to vllm_omni/platforms/npu/attention/static_shape_prefill.py index bde20c30502..9c1e515fcc2 100644 --- a/vllm_omni/platforms/npu/attention/fixed_kv_prefill.py +++ b/vllm_omni/platforms/npu/attention/static_shape_prefill.py @@ -4,7 +4,7 @@ of device work (three back-to-back forwards never accumulate more than 0.07 ms of drain), and a request pays roughly one of them -- all but a couple of milliseconds of it removable. The decode side of this problem was solved by -`fixed_kv_decode`; this module is the same idea for the prefill shape. +`static_shape_decode`; this module is the same idea for the prefill shape. ## Why a padded prefill replays exactly @@ -140,7 +140,7 @@ def maybe_mark_step(attn_metadata: Any, prefill_state: Any, block_size_hint: int """Builder hook: route an eligible prefill step onto its captured graph. Cheap, and refuses everything unusual: called once per step from - ``OmniFixedKVMetadataBuilder.build``. When it marks the step, it also + ``OmniStaticShapeMetadataBuilder.build``. When it marks the step, it also fills the persistent slot buffer -- real slots for the real rows, the request's own next slots for the padding, one device add with no sync. """ @@ -190,7 +190,7 @@ def maybe_mark_step(attn_metadata: Any, prefill_state: Any, block_size_hint: int arange = _state.aranges[bucket] blocks = block_table[0].index_select(0, (arange // block_size).to(torch.int64)) torch.add(arange % block_size, blocks * block_size, out=_state.slot_buffers[bucket]) - attn_metadata.fixed_kv_prefill_bucket = bucket + attn_metadata.static_shape_prefill_bucket = bucket def replay(attn_metadata: Any) -> torch.Tensor | None: @@ -203,7 +203,7 @@ def replay(attn_metadata: Any) -> torch.Tensor | None: """ if isinstance(attn_metadata, dict): attn_metadata = next(iter(attn_metadata.values()), None) - bucket = getattr(attn_metadata, "fixed_kv_prefill_bucket", None) + bucket = getattr(attn_metadata, "static_shape_prefill_bucket", None) if bucket is None: return None entry = _state.graphs.get(bucket) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 7a799067de7..481170318f4 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -56,7 +56,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): @classmethod def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: int | None = None): - """Route short-context decoders onto the fixed-KV decode backend. + """Route short-context decoders onto the static-shape decode backend. Under FULL_DECODE_ONLY the captured decode step has to re-issue attention on every layer on every step, because the op takes the KV length as a host @@ -71,18 +71,18 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: return resolved from vllm.config import get_current_vllm_config - from vllm_omni.platforms.npu.attention import fixed_kv_decode + from vllm_omni.platforms.npu.attention import static_shape_decode - if not fixed_kv_decode.is_enabled(): + if not static_shape_decode.is_enabled(): return resolved vllm_config = get_current_vllm_config() max_model_len = getattr(getattr(vllm_config, "model_config", None), "max_model_len", None) block_size = getattr(getattr(vllm_config, "cache_config", None), "block_size", None) - capacity = fixed_kv_decode.capacity_for(max_model_len, block_size) if max_model_len and block_size else None - fixed_kv_decode.install_into_ascend_aclgraph() + capacity = static_shape_decode.capacity_for(max_model_len, block_size) if max_model_len and block_size else None + static_shape_decode.install_into_ascend_aclgraph() if capacity: logger.info( - "[minicpmo] fixed-KV decode attention on (max_model_len=%s, kv_capacity=%s); " + "[minicpmo] static-shape decode attention on (max_model_len=%s, kv_capacity=%s); " "set VLLM_OMNI_FIXED_KV_DECODE=0 to disable", max_model_len, capacity, @@ -91,10 +91,10 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: # The backend class is still returned; it declines for this # max_model_len, so no bucket is captured and nothing engages. logger.info( - "[minicpmo] fixed-KV decode attention not engaged (max_model_len=%s is above the context ceiling)", + "[minicpmo] static-shape decode attention not engaged (max_model_len=%s is above the context ceiling)", max_model_len, ) - return "vllm_omni.platforms.npu.attention.fixed_kv_backend.OmniFixedKVAttentionBackend" + return "vllm_omni.platforms.npu.attention.static_shape_backend.OmniStaticShapeAttentionBackend" # conv2d convolution operator in the code2wav module of Qwen3-TTS not being able to run on Aclnn def __init__(self) -> None: diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 508efda86a0..4c2cf339f99 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -116,16 +116,16 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex """Autoregressive NPU model runner that returns hidden states per request.""" # vllm-ascend's ascend-graph wrapper reads this flag from the runner, but - # the vllm-ascend build this image ships never sets it -- the fixed-KV + # the vllm-ascend build this image ships never sets it -- the static-shape # capture path would raise AttributeError without a default here. enable_hamming_sparse: bool = False def __init__(self, *args, **kwargs): - from vllm_omni.platforms.npu.attention import fixed_kv_decode + from vllm_omni.platforms.npu.attention import static_shape_decode # Before parent init: vLLM's selector uses current_platform # (vllm-ascend), not NPUOmniPlatform.get_attn_backend_cls. - fixed_kv_decode.install_into_ascend_backend() + static_shape_decode.install_into_ascend_backend() super().__init__(*args, **kwargs) self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) # each model stage has their own hidden size @@ -180,7 +180,7 @@ def _check_and_update_cudagraph_mode(self, *args, **kwargs): return super()._check_and_update_cudagraph_mode(*args, **kwargs) def _capture_cudagraphs(self, *args, **kwargs): - """Capture fixed-KV decode graphs per KV capacity bucket. + """Capture static-shape decode graphs per KV capacity bucket. Each bucket needs its own pass: the default path leaves the capacity unset and would file the graphs under a key decode never looks up. Signature is @@ -188,15 +188,15 @@ def _capture_cudagraphs(self, *args, **kwargs): extra ``profiler`` argument that this override must not swallow.""" from vllm.config import CUDAGraphMode - from vllm_omni.platforms.npu.attention import fixed_kv_decode + from vllm_omni.platforms.npu.attention import static_shape_decode runtime_mode = kwargs.get("cudagraph_runtime_mode") if runtime_mode is None and len(args) >= 2: runtime_mode = args[1] buckets: tuple[int, ...] = () if runtime_mode == CUDAGraphMode.FULL: - buckets = fixed_kv_decode.buckets_for( - fixed_kv_decode.capacity_for( + buckets = static_shape_decode.buckets_for( + static_shape_decode.capacity_for( self.vllm_config.model_config.max_model_len, self.vllm_config.cache_config.block_size, ), @@ -206,8 +206,8 @@ def _capture_cudagraphs(self, *args, **kwargs): return super()._capture_cudagraphs(*args, **kwargs) with self._narrow_decode_query_len(): for capacity in sorted(buckets, reverse=True): - logger.info("[minicpmo] capturing fixed-KV decode graphs at kv_capacity=%d", capacity) - with fixed_kv_decode.capturing_bucket(capacity): + logger.info("[minicpmo] capturing static-shape decode graphs at kv_capacity=%d", capacity) + with static_shape_decode.capturing_bucket(capacity): super()._capture_cudagraphs(*args, **kwargs) def propose_draft_token_ids(self, valid_sampled_token_ids, *args, **kwargs): diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py index c577d3ad1b2..b798d22ad06 100644 --- a/vllm_omni/platforms/npu/worker/talker_multiframe.py +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -316,22 +316,22 @@ def begin_narrow_step( from vllm.forward_context import BatchDescriptor # The KV geometry required by narrow replay lives in - # platforms/npu/attention/fixed_kv_decode, which may be absent. Import + # platforms/npu/attention/static_shape_decode, which may be absent. Import # defensively: when missing, decline as documented and let the caller fall # back to wide replay; multi-frame decode and the stop chain are unaffected. try: - from vllm_omni.platforms.npu.attention import fixed_kv_decode + from vllm_omni.platforms.npu.attention import static_shape_decode except Exception: - fixed_kv_decode = None + static_shape_decode = None - if fixed_kv_decode is None: - return _decline("the fixed-KV decode backend is not present in this build") + if static_shape_decode is None: + return _decline("the static-shape decode backend is not present in this build") if not narrow_replay_enabled(runner) or inputs_embeds is None or positions is None: return None rows = len(spans) if rows <= 0: return None - seq_lens = fixed_kv_decode.captured_seq_lens(rows) + seq_lens = static_shape_decode.captured_seq_lens(rows) if seq_lens is None: return _decline(f"no one-query graph captured for {rows} rows") metadata = getattr(forward_context, "attn_metadata", None) @@ -339,7 +339,7 @@ def begin_narrow_step( metadata = next(iter(metadata.values()), None) step_slot_mapping = getattr(metadata, "slot_mapping", None) live_seq_lens = getattr(metadata, "seq_lens_device", None) - graph_slot_mapping = fixed_kv_decode.captured_slot_mapping(rows) + graph_slot_mapping = static_shape_decode.captured_slot_mapping(rows) if step_slot_mapping is None or live_seq_lens is None or graph_slot_mapping is None: return _decline("this step's attention metadata carries no slot mapping") if step_slot_mapping.shape[0] < rows * frames: @@ -431,7 +431,7 @@ def run( ``after_forward`` is the runner's post-forward graph bookkeeping. It is called after every replay because that is where it sits in the single-frame - path; under fixed-KV decode it is two pointer comparisons, and the frames of + path; under static-shape decode it is two pointer comparisons, and the frames of one step share their sequence lengths so there is nothing for it to rebind. """ global _LOGGED_ENGAGE From 7e9ba1a31ae29b70fcb01c88ebaaf08645b629ec Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 12:12:33 +0800 Subject: [PATCH 31/89] refactor(npu): make the static-shape decode defaults code constants The 910B family needs this backend to capture the K-step at all, so the module stays on that ground; what goes are its two switches. Their measured-stable values are code constants now: * whether a stage engages is decided only by its own ``max_model_len`` (via the 8192-token ceiling), so a long-context stage keeps the stock path and nothing else changes; * the candidate capacity stays "512", with the model's own capacity always added as the top bucket. ``is_enabled()`` and both names are gone, together with their inventory entries, their rows in the configuration reference and the snapshot counts. No behavior change at the defaults. Signed-off-by: y-null --- docs/configuration/environment_variables.md | 2 - tests/config/test_environment_variables.py | 2 +- .../config/environment_variable_inventory.py | 2 - .../npu/attention/static_shape_decode.py | 38 ++++++------------- vllm_omni/platforms/npu/platform.py | 12 +++--- 5 files changed, 17 insertions(+), 39 deletions(-) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 2338beebb87..714e693a8b5 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -122,8 +122,6 @@ settings. | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | | --- | --- | --- | --- | --- | -| `VLLM_OMNI_FIXED_KV_DECODE` | `0` disables; default `1` (on) | Fixed-capacity decode attention backend and its ACL-graph hook; read when the NPU platform resolves the attention backend | Only the exact value `0` disables it. Scoped by `max_model_len`: a long-context stage (e.g. a 32768-token Thinker) declines and keeps the stock path. | Experimental performance control | -| `VLLM_OMNI_FIXED_KV_DECODE_BUCKETS` | Comma-separated KV capacities in tokens; default `512` | Capture buckets of the fixed-KV decode graphs; read when the buckets are computed | The full capacity is always kept as the top bucket, so a longer sequence still has a graph. Empty or malformed entries are skipped. | Experimental performance control | | `VLLM_OMNI_NPU_SKIP_WARMUPS` | Comma-separated warmup-name substrings, or `none`; default `penalties,rejection_sampler` on 910_93 and `rejection_sampler` on 910B | vllm-ascend Triton kernel warmups on the NPU worker; resolved per warmup call | Skips the warmups whose dummy-token setup faults the vector core (acl 507035). `none` restores the stock warmups. | Experimental | | `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control | | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 5507b98ce72..d1c6093897d 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 34, + EnvironmentVariableCategory.PUBLIC_OMNI: 32, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 61, diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 46aacd9332e..b6b2e3d2549 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -67,8 +67,6 @@ def is_public_omni(self) -> bool: "VLLM_OMNI_AR_DIFFUSION_KV_GATHER", "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT", "VLLM_OMNI_EVENT_DRIVEN_ORCH", - "VLLM_OMNI_FIXED_KV_DECODE", - "VLLM_OMNI_FIXED_KV_DECODE_BUCKETS", "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "VLLM_OMNI_NIXL_LEASE_S", diff --git a/vllm_omni/platforms/npu/attention/static_shape_decode.py b/vllm_omni/platforms/npu/attention/static_shape_decode.py index ca552d2377c..3779633a832 100644 --- a/vllm_omni/platforms/npu/attention/static_shape_decode.py +++ b/vllm_omni/platforms/npu/attention/static_shape_decode.py @@ -41,17 +41,11 @@ bucket that covers it. With one bucket equal to the model's whole block table this degenerates to a fixed capacity, which is only sensible for short contexts. -Env switches: - ``VLLM_OMNI_FIXED_KV_DECODE`` 0 to disable (default 1) - ``VLLM_OMNI_FIXED_KV_DECODE_BUCKETS`` comma-separated candidate capacities, - e.g. "1024,4096"; default "512". Each - model keeps the candidates below its own - block-table capacity and always adds that - capacity as the top bucket. Empty means a - single bucket: the full capacity. +There is no switch: the numbers below are the measured-stable ones, and the only +thing that decides whether a stage engages is its own ``max_model_len`` (see +:func:`capacity_for`), so a long-context stage keeps the stock path untouched. """ -import os from contextlib import contextmanager import torch @@ -89,9 +83,12 @@ _runtime_bucket: int | None = None -_ENV = "VLLM_OMNI_FIXED_KV_DECODE" -_BUCKETS_ENV = "VLLM_OMNI_FIXED_KV_DECODE_BUCKETS" -_DEFAULT_BUCKETS = "512" +# Candidate capacities to capture graphs for, smallest first. A bucket that just +# covers the sequence is worth having: attention reads the declared capacity every +# step, and the Talker's sequence is ~135 (a ~15-token condition plus ~118 codec +# frames) against a 4096-token table. The model's own capacity is always added as +# the top bucket, so a longer sequence still has a graph to land on. +_CANDIDATE_BUCKETS = "512" # Only a small decoder has fixed-capacity graphs worth capturing; a long-context # stage (a 32768-token Thinker) keeps the stock path. The ceiling is a property of @@ -100,10 +97,6 @@ _MAX_CONTEXT = 8192 -def is_enabled() -> bool: - return os.getenv(_ENV, "1") == "1" - - class _CapacityEntries(dict): """The wrapper's entries dict, plus one dict per static-shape capacity. @@ -201,8 +194,6 @@ def install_into_ascend_backend() -> None: stock builder still runs for them, so runtime looks up ``(batch, None)``. Patch the stock backend class in this process instead. """ - if not is_enabled(): - return from vllm_ascend.attention.attention_v1 import AscendAttentionBackend from vllm_omni.platforms.npu.attention.static_shape_backend import ( @@ -221,8 +212,6 @@ def install_into_ascend_backend() -> None: def capacity_for(max_model_len: int, block_size: int) -> int: """Full KV capacity this decoder would run at, or 0 when it should not engage.""" - if not is_enabled(): - return 0 if max_model_len > _MAX_CONTEXT: return 0 # Round to the block table, not to max_model_len: the op indexes whole blocks. @@ -237,12 +226,7 @@ def buckets_for(capacity: int, block_size: int) -> tuple[int, ...]: """ if not capacity: return () - # The default has to be a code default, not an environment one: attention - # reads the declared capacity every step, so a bucket that just covers the - # sequence is worth having -- the Talker's is ~135 (a ~15-token condition - # plus ~118 codec frames). The full capacity is always kept as the top - # bucket, so a longer sequence still has a graph to land on. - raw = os.getenv(_BUCKETS_ENV, _DEFAULT_BUCKETS) + raw = _CANDIDATE_BUCKETS candidates = set() for part in raw.split(","): part = part.strip() @@ -293,7 +277,7 @@ def graph_key() -> int | None: Read by vllm-ascend's ``ACLGraphWrapper`` in builds that carry the static-shape hook; nothing in this package calls it directly. """ - return current_capacity() if is_enabled() else None + return current_capacity() def get_pse_buffer( diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 481170318f4..7b92273eeaa 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -62,9 +62,10 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: every layer on every step, because the op takes the KV length as a host argument that grows each step. On the small Talker decoder that rebind is most of its busy time. Scoped by ``max_model_len`` (Talker 4096 engages, - a 32768-token Thinker does not). ``VLLM_OMNI_FIXED_KV_DECODE=0`` restores - stock. The 910B family also needs it to capture at all: with the K-step - armed the stock capture path faults there (acl 507035). + a 32768-token Thinker does not), which is the only thing that decides: the + backend declines above the ceiling, so nothing else changes. The 910B + family also needs it to capture at all: with the K-step armed the stock + capture path faults there (acl 507035). """ resolved = super().get_attn_backend_cls(selected_backend, attn_selector_config, num_heads) if resolved != "vllm_ascend.attention.attention_v1.AscendAttentionBackend": @@ -73,8 +74,6 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: from vllm_omni.platforms.npu.attention import static_shape_decode - if not static_shape_decode.is_enabled(): - return resolved vllm_config = get_current_vllm_config() max_model_len = getattr(getattr(vllm_config, "model_config", None), "max_model_len", None) block_size = getattr(getattr(vllm_config, "cache_config", None), "block_size", None) @@ -82,8 +81,7 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: static_shape_decode.install_into_ascend_aclgraph() if capacity: logger.info( - "[minicpmo] static-shape decode attention on (max_model_len=%s, kv_capacity=%s); " - "set VLLM_OMNI_FIXED_KV_DECODE=0 to disable", + "[minicpmo] static-shape decode attention on (max_model_len=%s, kv_capacity=%s)", max_model_len, capacity, ) From f3b962d513e421e3b7576f7bdeb052aff515eb49 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 15:50:11 +0800 Subject: [PATCH 32/89] [Style][Ascend] Wrap over-long lines in static_shape_backend to satisfy ruff E501 Signed-off-by: y-null --- vllm_omni/platforms/npu/attention/static_shape_backend.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm_omni/platforms/npu/attention/static_shape_backend.py b/vllm_omni/platforms/npu/attention/static_shape_backend.py index 8190d33f4c3..f1f23f73947 100644 --- a/vllm_omni/platforms/npu/attention/static_shape_backend.py +++ b/vllm_omni/platforms/npu/attention/static_shape_backend.py @@ -103,7 +103,9 @@ def build(self, *args, **kwargs): if self._omni_buckets and attn_metadata.static_shape_decode_step: seq_lens_list = attn_metadata.seq_lens_list or [1] max_seq_len = getattr(common, "max_seq_len", None) or max(seq_lens_list) - static_shape_decode.set_runtime_bucket(static_shape_decode.select_bucket(int(max_seq_len), self._omni_buckets)) + static_shape_decode.set_runtime_bucket( + static_shape_decode.select_bucket(int(max_seq_len), self._omni_buckets) + ) else: static_shape_decode.set_runtime_bucket(None) # A single fresh prefill inside a captured bucket replays a graph @@ -156,7 +158,9 @@ def __init__(self, *args, **kwargs) -> None: # +2 matches the query_start_loc buffer, which carries the FIA padding # request on a full batch. self._omni_max_rows = int(self.vllm_config.scheduler_config.max_num_seqs) + 2 - self._omni_buckets = static_shape_decode.buckets_for(self._omni_capacity, self.vllm_config.cache_config.block_size) + self._omni_buckets = static_shape_decode.buckets_for( + self._omni_capacity, self.vllm_config.cache_config.block_size + ) # A step that schedules several tokens per sequence (the Talker's # multi-frame decode declares them through a speculative_config) is # still a uniform decode, and the fixed-capacity mask covers it as From 522b3ada499b4710770b26ad38c79e2f8a8f87e7 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 15:50:11 +0800 Subject: [PATCH 33/89] [Clean][Ascend][TTS] Drop fused top-k/top-p sampling path and stale env vars Remove the vendor fused torch_npu.npu_top_k_top_p wrapper and its two env knobs (MINICPMO_TTS_NPU_TOPK_TOPP / _MAX_BATCH): the pre-warp path is unused by the sampler and the env vars changed no behavior. Sampling returns to the sampler-native warpers. Also simplify ascend_warmup_patch and sync the env-var doc, inventory snapshot and its test. Signed-off-by: y-null --- docs/configuration/environment_variables.md | 1 - tests/config/test_environment_variables.py | 6 +- .../config/environment_variable_inventory.py | 3 - .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 115 ------------------ .../platforms/npu/ascend_warmup_patch.py | 12 +- 5 files changed, 4 insertions(+), 133 deletions(-) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 714e693a8b5..eaefb1015af 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -122,7 +122,6 @@ settings. | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | | --- | --- | --- | --- | --- | -| `VLLM_OMNI_NPU_SKIP_WARMUPS` | Comma-separated warmup-name substrings, or `none`; default `penalties,rejection_sampler` on 910_93 and `rejection_sampler` on 910B | vllm-ascend Triton kernel warmups on the NPU worker; resolved per warmup call | Skips the warmups whose dummy-token setup faults the vector core (acl 507035). `none` restores the stock warmups. | Experimental | | `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control | | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary | | `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index d1c6093897d..201c9d0e789 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,10 +164,10 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 32, + EnvironmentVariableCategory.PUBLIC_OMNI: 30, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 61, + EnvironmentVariableCategory.MODEL_SPECIFIC: 59, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 2, } @@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 33, + ModelEnvironmentVariableDisposition.PROMOTE: 31, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 16, diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index b6b2e3d2549..c3871bacec7 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -71,7 +71,6 @@ def is_public_omni(self) -> bool: "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "VLLM_OMNI_NIXL_LEASE_S", "VLLM_OMNI_NIXL_XFER_TIMEOUT_S", - "VLLM_OMNI_NPU_SKIP_WARMUPS", "VLLM_OMNI_ORCH_MONITOR_PATH", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", "VLLM_OMNI_SERVER_STORAGE__FILE_TTL", @@ -163,8 +162,6 @@ def is_public_omni(self) -> bool: "MINICPMO_TOKEN2WAV_TRT", "MINICPMO_TOKEN2WAV_TRT_CACHE", "MINICPMO_TOKEN2WAV_TRT_DTYPE", - "MINICPMO_TTS_NPU_TOPK_TOPP", - "MINICPMO_TTS_NPU_TOPK_TOPP_MAX_BATCH", "NEMOTRON_VOICECHAT_LLM_PATH", "OMNIVOICE_CUDA_GRAPH", "OMNIVOICE_TF32", diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index eaf285367c7..534a5a370af 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -12,14 +12,8 @@ 5. Next decode embeds that id with emb_code and emits it to Code2Wav """ -import os from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace - -try: - import torch_npu # noqa: F401 (NPU platform guarantee) -except Exception: # pragma: no cover - torch_npu = None from typing import Any import torch @@ -196,112 +190,6 @@ def _apply_batched_repetition_penalty( return penalized -def _apply_top_k_top_p( - logits: torch.Tensor, - *, - top_k: int | None, - top_p: float | None, - min_tokens_to_keep: int = 3, - inplace: bool = False, -) -> torch.Tensor: - """Reference warper: same candidate floors as the upstream warpers.""" - filtered = logits if inplace else logits.clone() - vocab_size = filtered.shape[-1] - if top_p is not None and 0.0 < top_p < 1.0: - sorted_logits, sorted_indices = torch.sort(filtered, descending=False, dim=-1) - cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) - remove = cumulative_probs <= (1.0 - float(top_p)) - remove[..., -min_tokens_to_keep:] = False - remove = remove.scatter(-1, sorted_indices, remove) - filtered.masked_fill_(remove, float("-inf")) - if top_k is not None and top_k > 0: - keep = min(vocab_size, max(int(top_k), min_tokens_to_keep)) - threshold = torch.topk(filtered, keep, dim=-1).values[..., -1, None] - filtered.masked_fill_(filtered < threshold, float("-inf")) - return filtered - - -_NPU_TOPK_CACHE: dict = {} -_NPU_TOP_K_TOP_P = os.environ.get("MINICPMO_TTS_NPU_TOPK_TOPP", "1") != "0" -# npu_top_k_top_p wins on launches at small batch but loses to the two-op -# torch path at large batch (128/8 A/B: 1.391 gated-off vs on). -_NPU_TOPK_MAX_BATCH = int(os.environ.get("MINICPMO_TTS_NPU_TOPK_TOPP_MAX_BATCH", "64")) - - -def _npu_top_k_top_p_warp(logits, *, top_k, top_p): - """Fused top-k/top-p floor via torch_npu.npu_top_k_top_p. - - Kernel semantics: keep top-k by value, then keep the top-p probability - mass over the retained set. Falls back to the exact PyTorch warper when - the kernel is unavailable or fails. - """ - if top_k is None or top_p is None or not 0.0 < top_p < 1.0: - return logits - npu = getattr(torch_npu, "npu_top_k_top_p", None) - if npu is None: - return logits - key = (str(logits.device), str(logits.dtype), float(top_p), int(top_k)) - cached = _NPU_TOPK_CACHE.get(key) - if cached is None: - dev = torch.full((1,), float(top_p), device=logits.device, dtype=logits.dtype) - dk = torch.full((1,), int(top_k), device=logits.device, dtype=torch.int32) - _NPU_TOPK_CACHE[key] = (dev, dk) - else: - dev, dk = cached - try: - return npu( - logits, - dev.expand(logits.shape[0]).contiguous(), - dk.expand(logits.shape[0]).contiguous(), - ) - except Exception: - return _apply_top_k_top_p(logits, top_k=top_k, top_p=top_p, min_tokens_to_keep=3, inplace=True) - - -def _maybe_prewarp_top_k_top_p(logits, sampling_metadata): - """Pre-apply the fused floor and neutralize the sampler's own warpers. - - Only engages when every row shares one (top_k, top_p) at temperature 1.0, - where the floor on raw logits is exactly the floor the sampler would - compute itself (top-k is rank based, so temperature invariance holds and - top-p mass matches at T=1). Any structural surprise returns None and the - native sampler path runs untouched. - """ - if not _NPU_TOP_K_TOP_P: - return None - if not isinstance(logits, torch.Tensor) or logits.ndim != 2 or logits.shape[0] == 0: - return None - if logits.shape[0] > _NPU_TOPK_MAX_BATCH: - return None - if getattr(torch_npu, "npu_top_k_top_p", None) is None: - return None - try: - params = getattr(sampling_metadata, "sampling_params", None) - if not isinstance(params, (list, tuple)) or len(params) != logits.shape[0]: - return None - top_k = top_p = None - for sp in params: - k = getattr(sp, "top_k", None) - pp = getattr(sp, "top_p", None) - t = getattr(sp, "temperature", None) - if t is None or float(t) != 1.0: - return None - if k is None or int(k) <= 0: - return None - if pp is None or not 0.0 < float(pp) < 1.0: - return None - if top_k is None: - top_k, top_p = int(k), float(pp) - elif int(k) != top_k or float(pp) != top_p: - return None - logits = _npu_top_k_top_p_warp(logits, top_k=top_k, top_p=top_p) - new_params = [replace(sp, top_p=1.0, top_k=-1) for sp in params] - sampling_metadata = replace(sampling_metadata, sampling_params=list(new_params)) - return logits, sampling_metadata - except Exception: - return None - - def resolve_codec_sampling_params( yaml_params: Mapping[str, Any] | None, tts_config: Any | None = None, @@ -1656,9 +1544,6 @@ def sample(self, logits, sampling_metadata): prompt_token_ids=blank_scheduler_prompt_for_penalties(prompt_ids, logits.shape[-1]), ) logits, sampling_metadata = self._apply_codec_repetition_penalty(logits, sampling_metadata) - prewarped = _maybe_prewarp_top_k_top_p(logits, sampling_metadata) - if prewarped is not None: - logits, sampling_metadata = prewarped force_eos = self._pending_force_eos_rows self._pending_force_eos_rows = None output = Sampler()(logits, sampling_metadata) diff --git a/vllm_omni/platforms/npu/ascend_warmup_patch.py b/vllm_omni/platforms/npu/ascend_warmup_patch.py index d8684028ba6..da44597f1ff 100644 --- a/vllm_omni/platforms/npu/ascend_warmup_patch.py +++ b/vllm_omni/platforms/npu/ascend_warmup_patch.py @@ -35,8 +35,6 @@ and a probe that failed there would silently strip a warmup the baseline relies on. A probe that still cannot name the SoC skips the faulting warmup: a missing warmup costs latency, a faulting one costs the whole run. -``VLLM_OMNI_NPU_SKIP_WARMUPS`` overrides the set (comma separated substrings -of the warmup names, or "none" to restore stock behaviour). ``rejection_sampler_triton_warmup`` is in the default skip set for the same reason: its dummy (batch, spec_len, vocab) sweep faults the 910_93 vector @@ -48,7 +46,6 @@ from __future__ import annotations -import os from collections.abc import Callable from typing import Any @@ -57,7 +54,6 @@ logger = init_logger(__name__) _PATCHED = False -_ENV = "VLLM_OMNI_NPU_SKIP_WARMUPS" # Substrings matched against the warmup function names in kernel_warmup.py. _DEFAULT_SKIP = ("penalties", "rejection_sampler") _WARMUP_NAMES = ( @@ -107,11 +103,6 @@ def _kstep_armed() -> bool: def _skipped_names() -> set[str]: - raw = os.environ.get(_ENV, "").strip().lower() - if raw: - if raw in ("none", "off", "0", "false", "no"): - return set() - return {part.strip() for part in raw.split(",") if part.strip()} soc = _probe_soc_name() if soc.startswith("Ascend910B"): # This warmup faults the vector core (acl 507035) on this family, and the @@ -176,7 +167,6 @@ def apply_ascend_warmup_patch() -> None: if guarded: logger.info( "[npu] ascend Triton warmups guarded (%s); the skip set is resolved " - "per call, so the SoC probe sees a live device. Override with %s.", + "per call, so the SoC probe sees a live device.", ", ".join(guarded), - _ENV, ) From c7f7263b9c6cadd0a44e609af2f8b626522e26fa Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 20:59:10 +0800 Subject: [PATCH 34/89] fix(minicpmo_4_5): scope the NPU-only stage overrides to the NPU platform The Talker multi-frame decode and the Thinker prefix cache were declared in the platform-shared stage blocks, so a CUDA deployment inherited both: * stage 1 ran vLLM's real n-gram proposer next to the rejection sampler, which asserts on the Talker's codec ids (device-side assert out of get_token_bin_counts_and_mask); * stage 0 turned prefix caching on, and a turn-2 prefix hit left the Thinker with an empty tts handoff (ValueError in minicpmo_4_5_omni_tts.preprocess) together with a device-side assert out of indexSelectSmallIndex. Either one kills the stage-1 replica, and every later duplex case then fails with "stage 1 has no live replica". Move both blocks into platforms.npu: the NPU behaviour is unchanged, because the NPU worker intercepts the decode step and injects the constant continue drafts itself. CUDA now falls back to the top-level enable_prefix_caching: false and to the codec EOS stop id, which is also what the stop-id helper follows. Verification on L20X (single card), numbers against the PR-free baseline in the PR description: * duplex, 11 cases: 7 passed / 4 failed versus 11 failed before; the 4 remaining are the upstream response timeouts the baseline shows as well, and the engine stayed alive with 0 ValueError and 0 device-side asserts. * offline seed-tts perf (32@1 / 64@4 / 128@8) and the Seed-TTS WER suite are unchanged. * unit tests: 25 config cases and 563 MiniCPM-o model executor cases. Signed-off-by: y-null --- tests/config/test_config_factory.py | 48 +++++++++++++++ .../models/minicpmo_4_5/test_pipeline.py | 21 ++++--- .../test_talker_kstep_feedback.py | 13 +++- vllm_omni/config/stage_config.py | 5 ++ vllm_omni/deploy/minicpmo_4_5.yaml | 60 ++++++++++++------- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 60 ++++++++++++------- .../models/minicpmo_4_5/pipeline.py | 14 +++++ 7 files changed, 169 insertions(+), 52 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index fdba7e90a2a..6e37a87ad20 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1363,6 +1363,54 @@ def test_load_minicpmo_default_deploy_config(self): assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] assert "codec_sampling_params" not in stages[1].yaml_engine_args + @pytest.mark.parametrize("platform", ["cuda", "npu"]) + def test_minicpmo_talker_multi_frame_is_npu_scoped(self, platform: str): + """Stage 1's K-frame decode needs the NPU worker to inject constant + continue drafts, so the deploy block that arms it must not sit in the + platform-shared stage. + + On NPU the ``speculative_config`` is a fingerprint: the worker replays + the decode graph and vLLM's n-gram proposer is never consulted. On CUDA + the same block hands the Talker to the real proposer plus the rejection + sampler, which asserts on its codec ids (device-side assert out of + ``get_token_bin_counts_and_mask``). + + Stage 0's prefix caching is NPU-scoped for the same reason: a CUDA + stage-0 prefix hit rewrites the prefill bookkeeping, so the Thinker + hands an empty tts handoff and the Talker indexes past the codec vocab + (ValueError in ``minicpmo_4_5_omni_tts.preprocess`` plus a + device-side assert out of ``indexSelectSmallIndex``). + """ + deploy = _apply_platform_overrides( + load_deploy_config(Path(get_deploy_config_path("minicpmo_4_5.yaml"))), + platform=platform, + ) + # Assert the deploy-side stop id before the merge: afterwards it is + # unioned with the pipeline's own platform-aware constraint, so an + # exact-list check here would depend on the host running the suite. + deploy_stage1 = next(stage for stage in deploy.stages if stage.stage_id == 1) + stages = merge_pipeline_deploy(resolve_pipeline_config("minicpmo_4_5"), deploy) + + if platform == "npu": + assert deploy_stage1.default_sampling_params["stop_token_ids"] == [1] + assert stages[0].yaml_engine_args["enable_prefix_caching"] is True + assert stages[1].yaml_engine_args["speculative_config"] == { + "method": "ngram", + "num_speculative_tokens": 7, + "prompt_lookup_min": 1, + "prompt_lookup_max": 1, + } + assert stages[1].yaml_engine_args["async_scheduling"] is False + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [8, 8, 8] + else: + assert "stop_token_ids" not in (deploy_stage1.default_sampling_params or {}) + # Either absent (falls back to the top-level false) or explicitly + # false; the regression this guards against is true. + assert stages[0].yaml_engine_args.get("enable_prefix_caching") is not True + assert "speculative_config" not in stages[1].yaml_engine_args + assert stages[1].yaml_engine_args["async_scheduling"] is True + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4] + @pytest.mark.parametrize( ("filename", "stage0_devices", "stage1_devices", "stage2_devices", "stage1_replicas"), [ diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 855af4185bf..15fbb54eaa2 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -29,7 +29,11 @@ load_deploy_config, merge_pipeline_deploy, ) +from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import ( + _CODEC_EOS_TOKEN_ID, +) from vllm_omni.model_executor.models.registry import _OMNI_MODELS +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -106,12 +110,14 @@ def test_talker_stage(self, pipeline: PipelineConfig) -> None: assert talker.engine_output_type == "latent" # scope KV cache / mrope sizing to talker sub-config assert talker.hf_config_name == "tts_config" + # The multi-frame head collapses the vLLM-level head to the two-wide + # continue/stop row, whose only stop id is 1. It is an NPU-worker + # feature (that is where the continue drafts are injected), so on any + # other platform stage 1 keeps its full codec head and the codec EOS. + expected_stop = [1] if (current_omni_platform.device_name or "").lower() == "npu" else [_CODEC_EOS_TOKEN_ID] assert talker.sampling_constraints == { "detokenize": False, - # The multi-frame head: stage 1's deploy config carries the - # speculative_config that collapses the vLLM-level head to the - # two-wide continue/stop row, whose only stop id is 1. - "stop_token_ids": [1], + "stop_token_ids": expected_stop, } assert talker.custom_process_next_stage_input_func == ( "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload" @@ -177,9 +183,10 @@ def test_deploy_resolves_three_stage_pipeline(self, filename: str, devices: list assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor) assert "hf_overrides" not in stages[1].yaml_engine_args if filename == "minicpmo_4_5.yaml": - # All three stages cap out at 8 so the top tier is not gated by a - # single stage. - assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [8, 8, 8] + # The NPU overlay caps all three stages at 8 so the top tier is not + # gated by a single stage; every other platform keeps the base 4. + expected_seqs = 8 if (current_omni_platform.device_name or "").lower() == "npu" else 4 + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [expected_seqs] * 3 memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] assert memory_utilizations == [ 0.55, diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index a0f5ba31e77..ff953aba172 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -247,14 +247,21 @@ def test_talker_stop_token_ids_match_the_multi_frame_head(): is the two-wide continue/stop row, so the only sampleable ids are 0/1, `check_stop` never matched, and every request ran to ``max_tokens`` (142s per request, model finished at frame ~116). The pipeline default therefore - follows the multi-frame head; a deployment that turns the loop off - overrides stage 1's stop token per stage. + follows the head that runs: the multi-frame row next to the NPU worker, + which is the only place that injects the continue drafts, and the codec EOS + everywhere else. """ from vllm_omni.model_executor.models.minicpmo_4_5 import pipeline as mcp_pipeline + from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms.npu.worker import talker_multiframe - assert mcp_pipeline._talker_stop_token_ids() == [talker_multiframe.STOP_TOKEN_ID] assert talker_multiframe.STOP_TOKEN_ID == 1 + if (current_omni_platform.device_name or "").lower() == "npu": + assert mcp_pipeline._talker_stop_token_ids() == [talker_multiframe.STOP_TOKEN_ID] + else: + # No NPU worker means no constant-draft injection, so the K-frame loop + # is off and the two-wide continue/stop row never runs. + assert mcp_pipeline._talker_stop_token_ids() == [mcp_pipeline._CODEC_EOS_TOKEN_ID] def test_multiframe_gate_matrix(): diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index e7bdd6994d8..04a602da4a8 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -779,6 +779,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: return merged +# Same TND ceiling as stage 0: at most 16 query positions per sequence. +# Consumed by minicpmo_4_5_omni_tts._parse_k_step_frames as the hard cap on +# the K-step frame count accepted from the deploy YAML. +_MINICPMO_TALKER_FRAMES_MAX = 16 + # Stage 1's multi-frame decode is configured through its ``speculative_config`` # in the deploy YAML. The parse records whether it is armed here, because the # reader (``ascend_warmup_patch._kstep_armed``) can run in a spawned stage diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 05104f335dd..b73d702e89f 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -41,13 +41,18 @@ stages: # Leave headroom on the shared GPU for Talker KV + concurrent Code2Wav. gpu_memory_utilization: 0.55 enforce_eager: false - # Stage-level prefix caching override; the top level stays false for duplex safety. - # Stage-0 n-gram speculative decoding is intentionally not enabled here: the + # Stage-level prefix caching is NPU-only (see platforms.npu below); + # the top level stays false for duplex safety. Observed with this flag + # in the platform-shared block on CUDA: the first duplex case dies with + # an empty tts handoff (token_ids/hidden_states=None -> ValueError in + # minicpmo_4_5_omni_tts.preprocess) alongside a device-side assert out + # of indexSelectSmallIndex; the stage-1 replica is killed and every later + # case fails with "stage 1 has no live replica". Stage-0 n-gram + # speculative decoding stays off too: the # omni prefix cache refuses the pair (stage_prefix_cache_config -> - # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError as soon - # as a stage has enable_prefix_caching and a speculative_config; vLLM itself - # allows the combination). - enable_prefix_caching: true + # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError + # as soon as a stage has enable_prefix_caching and a speculative_config; + # vLLM itself allows the combination). tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -84,19 +89,6 @@ stages: # Talker context is 4096 (tts_config.max_position_embeddings). max_model_len: 4096 devices: "0" - # Multi-frame decode: the Talker produces 8 codec frames per scheduler step. - # vLLM learns that each request advances 8 tokens through this - # speculative_config (the n-gram proposer is never consulted -- the runner - # supplies constant continue drafts), and async scheduling is rejected next - # to an n-gram proposer. stop_token_ids above must stay the two-wide - # continue/stop id the head emits in this mode. Remove the block (or set it - # to null) to fall back to one frame per step. - speculative_config: - method: ngram - num_speculative_tokens: 7 - prompt_lookup_min: 1 - prompt_lookup_max: 1 - async_scheduling: false output_connectors: to_stage_2: connector_of_shared_memory # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler @@ -112,8 +104,6 @@ stages: top_k: 25 repetition_penalty: 1.05 min_tokens: 50 - # The Talker's stop head emits this id to finish a request. - stop_token_ids: [1] max_tokens: 4096 seed: 42 detokenize: false @@ -152,6 +142,10 @@ platforms: # not gated by a single stage. max_num_seqs: 8 max_num_batched_tokens: 8192 + # Stage-0 prefix caching was moved here from the platform-shared + # block: keep the PR's NPU behaviour unchanged while CUDA falls back + # to the top-level false (duplex safety, see the base stage comment). + enable_prefix_caching: true compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame # expansion; PIECEWISE silently drops both. Measured on 910B3: @@ -162,6 +156,30 @@ platforms: # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 + # Multi-frame decode: the Talker produces 8 codec frames per scheduler + # step. vLLM learns that each request advances 8 tokens through this + # speculative_config; the n-gram proposer itself is never consulted + # because the NPU worker replays the decode graph with constant continue + # drafts, and async scheduling is rejected next to an n-gram proposer. + # stop_token_ids must stay the two-wide continue/stop id the head emits + # in this mode. Remove the block (or set it to null) to fall back to one + # frame per step. + # + # NPU-only on purpose: only + # vllm_omni.platforms.npu.worker.talker_multiframe intercepts the step + # and injects those drafts. Left in the base stage, the same block hands + # stage 1 to vLLM's real n-gram proposer and rejection sampler on CUDA, + # which asserts on the Talker's codec ids (device-side assert out of + # get_token_bin_counts_and_mask). + speculative_config: + method: ngram + num_speculative_tokens: 7 + prompt_lookup_min: 1 + prompt_lookup_max: 1 + async_scheduling: false + default_sampling_params: + # The Talker's stop head emits this id to finish a request. + stop_token_ids: [1] compilation_config: cudagraph_mode: FULL_DECODE_ONLY - stage_id: 2 diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 6d1949af723..52cd8f8e735 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -45,13 +45,18 @@ stages: # replicas below; the NPU block caps stage 0/1 KV explicitly. gpu_memory_utilization: 0.40 enforce_eager: false - # Stage-level prefix caching override; the top level stays false for duplex safety. - # Stage-0 n-gram speculative decoding is intentionally not enabled here: the + # Stage-level prefix caching is NPU-only (see platforms.npu below); + # the top level stays false for duplex safety. Observed with this flag + # in the platform-shared block on CUDA: the first duplex case dies with + # an empty tts handoff (token_ids/hidden_states=None -> ValueError in + # minicpmo_4_5_omni_tts.preprocess) alongside a device-side assert out + # of indexSelectSmallIndex; the stage-1 replica is killed and every later + # case fails with "stage 1 has no live replica". Stage-0 n-gram + # speculative decoding stays off too: the # omni prefix cache refuses the pair (stage_prefix_cache_config -> - # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError as soon - # as a stage has enable_prefix_caching and a speculative_config; vLLM itself - # allows the combination). - enable_prefix_caching: true + # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError + # as soon as a stage has enable_prefix_caching and a speculative_config; + # vLLM itself allows the combination). tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -88,19 +93,6 @@ stages: # Talker context is 4096 (tts_config.max_position_embeddings). max_model_len: 4096 devices: "0" - # Multi-frame decode: the Talker produces 8 codec frames per scheduler step. - # vLLM learns that each request advances 8 tokens through this - # speculative_config (the n-gram proposer is never consulted -- the runner - # supplies constant continue drafts), and async scheduling is rejected next - # to an n-gram proposer. stop_token_ids above must stay the two-wide - # continue/stop id the head emits in this mode. Remove the block (or set it - # to null) to fall back to one frame per step. - speculative_config: - method: ngram - num_speculative_tokens: 7 - prompt_lookup_min: 1 - prompt_lookup_max: 1 - async_scheduling: false output_connectors: to_stage_2: connector_of_shared_memory # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler @@ -116,8 +108,6 @@ stages: top_k: 25 repetition_penalty: 1.05 min_tokens: 50 - # The Talker's stop head emits this id to finish a request. - stop_token_ids: [1] max_tokens: 4096 seed: 42 detokenize: false @@ -159,6 +149,10 @@ platforms: # not gated by a single stage. max_num_seqs: 8 max_num_batched_tokens: 8192 + # Stage-0 prefix caching was moved here from the platform-shared + # block: keep the PR's NPU behaviour unchanged while CUDA falls back + # to the top-level false (duplex safety, see the base stage comment). + enable_prefix_caching: true # Explicit KV cap: the auto-computed value does not leave room for the # two stage-2 replicas on this card. kv_cache_memory_bytes: 6442450944 @@ -172,6 +166,30 @@ platforms: # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 + # Multi-frame decode: the Talker produces 8 codec frames per scheduler + # step. vLLM learns that each request advances 8 tokens through this + # speculative_config; the n-gram proposer itself is never consulted + # because the NPU worker replays the decode graph with constant continue + # drafts, and async scheduling is rejected next to an n-gram proposer. + # stop_token_ids must stay the two-wide continue/stop id the head emits + # in this mode. Remove the block (or set it to null) to fall back to one + # frame per step. + # + # NPU-only on purpose: only + # vllm_omni.platforms.npu.worker.talker_multiframe intercepts the step + # and injects those drafts. Left in the base stage, the same block hands + # stage 1 to vLLM's real n-gram proposer and rejection sampler on CUDA, + # which asserts on the Talker's codec ids (device-side assert out of + # get_token_bin_counts_and_mask). + speculative_config: + method: ngram + num_speculative_tokens: 7 + prompt_lookup_min: 1 + prompt_lookup_max: 1 + async_scheduling: false + default_sampling_params: + # The Talker's stop head emits this id to finish a request. + stop_token_ids: [1] # Explicit KV cap: the auto-computed value does not leave room for the # two stage-2 replicas on this card. kv_cache_memory_bytes: 2684354560 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 68c77082e23..d252bc4712c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -18,6 +18,7 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni" MINICPMO45_REFERENCE_AUDIO_KEY = "_minicpmo45_reference_audio" +_CODEC_EOS_TOKEN_ID = 6561 # tts_config.num_audio_tokens - 1 def _talker_stop_token_ids() -> list[int]: @@ -38,6 +39,19 @@ def _talker_stop_token_ids() -> list[int]: vocabulary again and the codec EOS is the stop id; it overrides this per stage through the usual ``stop_token_ids`` sampling parameter. """ + from vllm_omni.platforms import current_omni_platform + + # ``speculative_config`` only means "constant continue drafts" next to the + # NPU worker that intercepts the step + # (``vllm_omni.platforms.npu.worker.talker_multiframe``); on every other + # platform the same block hands the Talker to vLLM's real n-gram proposer and + # rejection sampler, which asserts on the codec ids. The block is therefore + # confined to ``platforms.npu`` in the deploy config, and the stop id follows: + # NPU keeps the two-wide continue/stop row, everyone else keeps the full codec + # head and stops on the codec EOS. + if (current_omni_platform.device_name or "").lower() != "npu": + return [_CODEC_EOS_TOKEN_ID] + from vllm_omni.platforms.npu.worker.talker_multiframe import STOP_TOKEN_ID return [STOP_TOKEN_ID] From 1dcde09ff286dc1470fc64d9718d9c60ddebde28 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 21:42:09 +0800 Subject: [PATCH 35/89] docs(minicpmo_4_5): trim the NPU-scoping comments to what the code needs Review pass over the comments the NPU-scoping change added: they narrated the CUDA failure that motivated the scoping at a length the code does not need, and the pipeline docstring still described the previous unconditional codec-EOS pinning instead of what the function now does. - deploy configs (both profiles): the stage-0 prefix-caching note and the multi-frame block keep the reason and the two things a reader must not break (the block is NPU-only; its stop id travels with it). The CUDA traces shrink to one clause, and the stage-0 note next to the moved flag becomes a pointer to it. - pipeline.py: the docstring now states that the stop id follows the head this platform runs; the inline comment keeps only why the early return exists. - tests: the two long rationales keep the conclusion, drop the incident log. Comments and docstrings only: no behavior change (verified by comparing normalized token streams for the .py files and the parsed YAML for the two deploy configs). Signed-off-by: y-null --- tests/config/test_config_factory.py | 11 ++--- .../test_talker_kstep_feedback.py | 10 +---- vllm_omni/deploy/minicpmo_4_5.yaml | 41 +++++-------------- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 41 +++++-------------- .../models/minicpmo_4_5/pipeline.py | 31 ++++---------- 5 files changed, 36 insertions(+), 98 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 6e37a87ad20..299a7d54957 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1376,18 +1376,15 @@ def test_minicpmo_talker_multi_frame_is_npu_scoped(self, platform: str): ``get_token_bin_counts_and_mask``). Stage 0's prefix caching is NPU-scoped for the same reason: a CUDA - stage-0 prefix hit rewrites the prefill bookkeeping, so the Thinker - hands an empty tts handoff and the Talker indexes past the codec vocab - (ValueError in ``minicpmo_4_5_omni_tts.preprocess`` plus a - device-side assert out of ``indexSelectSmallIndex``). + stage-0 prefix hit hands the Thinker an empty tts handoff, and the + Talker then indexes past the codec vocab. """ deploy = _apply_platform_overrides( load_deploy_config(Path(get_deploy_config_path("minicpmo_4_5.yaml"))), platform=platform, ) - # Assert the deploy-side stop id before the merge: afterwards it is - # unioned with the pipeline's own platform-aware constraint, so an - # exact-list check here would depend on the host running the suite. + # Check the deploy-side stop id before the merge: afterwards it is + # unioned with the pipeline's platform-aware constraint. deploy_stage1 = next(stage for stage in deploy.stages if stage.stage_id == 1) stages = merge_pipeline_deploy(resolve_pipeline_config("minicpmo_4_5"), deploy) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index ff953aba172..7ae1509cbbc 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -242,14 +242,8 @@ def test_guard_keeps_drafts_when_all_decodes_share_the_padded_width(): def test_talker_stop_token_ids_match_the_multi_frame_head(): """Stage 1's stop id must be one the head that actually runs can emit. - The pipeline constraint used to pin ``stop_token_ids`` to the codec EOS - (6561) unconditionally. With the multi-frame loop on, the vLLM-level head - is the two-wide continue/stop row, so the only sampleable ids are 0/1, - `check_stop` never matched, and every request ran to ``max_tokens`` (142s - per request, model finished at frame ~116). The pipeline default therefore - follows the head that runs: the multi-frame row next to the NPU worker, - which is the only place that injects the continue drafts, and the codec EOS - everywhere else. + The default follows the head: the two-wide continue/stop row next to the + NPU worker, and the codec EOS on every platform without that worker. """ from vllm_omni.model_executor.models.minicpmo_4_5 import pipeline as mcp_pipeline from vllm_omni.platforms import current_omni_platform diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index b73d702e89f..162ef0b9f70 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -41,18 +41,11 @@ stages: # Leave headroom on the shared GPU for Talker KV + concurrent Code2Wav. gpu_memory_utilization: 0.55 enforce_eager: false - # Stage-level prefix caching is NPU-only (see platforms.npu below); - # the top level stays false for duplex safety. Observed with this flag - # in the platform-shared block on CUDA: the first duplex case dies with - # an empty tts handoff (token_ids/hidden_states=None -> ValueError in - # minicpmo_4_5_omni_tts.preprocess) alongside a device-side assert out - # of indexSelectSmallIndex; the stage-1 replica is killed and every later - # case fails with "stage 1 has no live replica". Stage-0 n-gram - # speculative decoding stays off too: the - # omni prefix cache refuses the pair (stage_prefix_cache_config -> - # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError - # as soon as a stage has enable_prefix_caching and a speculative_config; - # vLLM itself allows the combination). + # Stage-level prefix caching is NPU-only (see platforms.npu below): in the + # platform-shared block it killed the stage-1 replica on CUDA (empty tts + # handoff, then an indexSelectSmallIndex device assert). Stage-0 n-gram + # speculation stays off either way -- the omni prefix cache refuses it next + # to enable_prefix_caching. tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -142,9 +135,7 @@ platforms: # not gated by a single stage. max_num_seqs: 8 max_num_batched_tokens: 8192 - # Stage-0 prefix caching was moved here from the platform-shared - # block: keep the PR's NPU behaviour unchanged while CUDA falls back - # to the top-level false (duplex safety, see the base stage comment). + # Moved here from the platform-shared block; see the base stage note. enable_prefix_caching: true compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame @@ -156,21 +147,11 @@ platforms: # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 - # Multi-frame decode: the Talker produces 8 codec frames per scheduler - # step. vLLM learns that each request advances 8 tokens through this - # speculative_config; the n-gram proposer itself is never consulted - # because the NPU worker replays the decode graph with constant continue - # drafts, and async scheduling is rejected next to an n-gram proposer. - # stop_token_ids must stay the two-wide continue/stop id the head emits - # in this mode. Remove the block (or set it to null) to fall back to one - # frame per step. - # - # NPU-only on purpose: only - # vllm_omni.platforms.npu.worker.talker_multiframe intercepts the step - # and injects those drafts. Left in the base stage, the same block hands - # stage 1 to vLLM's real n-gram proposer and rejection sampler on CUDA, - # which asserts on the Talker's codec ids (device-side assert out of - # get_token_bin_counts_and_mask). + # Multi-frame decode: 8 codec frames per step. NPU-only -- only the NPU + # worker intercepts the step and injects the constant continue drafts, + # so the block stays in this section. stop_token_ids below is the + # two-wide continue/stop id that head emits; drop the block (or set it + # to null) for one frame per step, and drop that id with it. speculative_config: method: ngram num_speculative_tokens: 7 diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 52cd8f8e735..f3415068ad2 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -45,18 +45,11 @@ stages: # replicas below; the NPU block caps stage 0/1 KV explicitly. gpu_memory_utilization: 0.40 enforce_eager: false - # Stage-level prefix caching is NPU-only (see platforms.npu below); - # the top level stays false for duplex safety. Observed with this flag - # in the platform-shared block on CUDA: the first duplex case dies with - # an empty tts handoff (token_ids/hidden_states=None -> ValueError in - # minicpmo_4_5_omni_tts.preprocess) alongside a device-side assert out - # of indexSelectSmallIndex; the stage-1 replica is killed and every later - # case fails with "stage 1 has no live replica". Stage-0 n-gram - # speculative decoding stays off too: the - # omni prefix cache refuses the pair (stage_prefix_cache_config -> - # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError - # as soon as a stage has enable_prefix_caching and a speculative_config; - # vLLM itself allows the combination). + # Stage-level prefix caching is NPU-only (see platforms.npu below): in the + # platform-shared block it killed the stage-1 replica on CUDA (empty tts + # handoff, then an indexSelectSmallIndex device assert). Stage-0 n-gram + # speculation stays off either way -- the omni prefix cache refuses it next + # to enable_prefix_caching. tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -149,9 +142,7 @@ platforms: # not gated by a single stage. max_num_seqs: 8 max_num_batched_tokens: 8192 - # Stage-0 prefix caching was moved here from the platform-shared - # block: keep the PR's NPU behaviour unchanged while CUDA falls back - # to the top-level false (duplex safety, see the base stage comment). + # Moved here from the platform-shared block; see the base stage note. enable_prefix_caching: true # Explicit KV cap: the auto-computed value does not leave room for the # two stage-2 replicas on this card. @@ -166,21 +157,11 @@ platforms: # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 - # Multi-frame decode: the Talker produces 8 codec frames per scheduler - # step. vLLM learns that each request advances 8 tokens through this - # speculative_config; the n-gram proposer itself is never consulted - # because the NPU worker replays the decode graph with constant continue - # drafts, and async scheduling is rejected next to an n-gram proposer. - # stop_token_ids must stay the two-wide continue/stop id the head emits - # in this mode. Remove the block (or set it to null) to fall back to one - # frame per step. - # - # NPU-only on purpose: only - # vllm_omni.platforms.npu.worker.talker_multiframe intercepts the step - # and injects those drafts. Left in the base stage, the same block hands - # stage 1 to vLLM's real n-gram proposer and rejection sampler on CUDA, - # which asserts on the Talker's codec ids (device-side assert out of - # get_token_bin_counts_and_mask). + # Multi-frame decode: 8 codec frames per step. NPU-only -- only the NPU + # worker intercepts the step and injects the constant continue drafts, + # so the block stays in this section. stop_token_ids below is the + # two-wide continue/stop id that head emits; drop the block (or set it + # to null) for one frame per step, and drop that id with it. speculative_config: method: ngram num_speculative_tokens: 7 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index d252bc4712c..7eb9878720d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -22,33 +22,18 @@ def _talker_stop_token_ids() -> list[int]: - """Stage 1's default stop id: the one the multi-frame head can emit. + """Stage 1's default stop id, derived from the head this platform runs. - Stage 1's deploy config carries the ``speculative_config`` that arms the - multi-frame decode, and ``compute_logits`` then collapses the Talker's - vLLM-level head to the two-wide continue/stop row. The only sampleable ids - are therefore 0 and 1, and the codec EOS can never appear at this level (the - model samples codec ids inside the loop and forwards them to stage 2 - itself). - - Pinning this to the codec EOS regardless is what made every multi-frame - request run to ``max_tokens``: the model did emit the stop marker, - ``check_stop`` compared it against 6561, and the request kept its slot until - the context ran out. A deployment that drops the ``speculative_config`` from - stage 1 falls back to one frame per step, where the head is the full codec - vocabulary again and the codec EOS is the stop id; it overrides this per - stage through the usual ``stop_token_ids`` sampling parameter. + On NPU, stage 1's ``speculative_config`` collapses the vLLM-level head to + the two-wide continue/stop row: only 0/1 are sampleable, so the codec EOS + can never appear. Everywhere else that block stays out of the config (it is + confined to ``platforms.npu``), the head is the full codec vocabulary, and + the codec EOS is the only valid stop. """ from vllm_omni.platforms import current_omni_platform - # ``speculative_config`` only means "constant continue drafts" next to the - # NPU worker that intercepts the step - # (``vllm_omni.platforms.npu.worker.talker_multiframe``); on every other - # platform the same block hands the Talker to vLLM's real n-gram proposer and - # rejection sampler, which asserts on the codec ids. The block is therefore - # confined to ``platforms.npu`` in the deploy config, and the stop id follows: - # NPU keeps the two-wide continue/stop row, everyone else keeps the full codec - # head and stops on the codec EOS. + # The deploy config confines stage 1's speculative_config to platforms.npu, + # so a non-NPU stage 1 keeps the full codec head. if (current_omni_platform.device_name or "").lower() != "npu": return [_CODEC_EOS_TOKEN_ID] From 79f17121af83e0cf08cef6e23fb3348ae60d0d48 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 21:50:12 +0800 Subject: [PATCH 36/89] fix(config): record the K step through the NPU overlay Stage 1's ``speculative_config`` now lives under ``platforms.npu``, but ``_record_talker_frames`` still read the top-level stage list, so ``talker_frames_per_step()`` reported 1 for both deploy profiles. Its only reader is ``ascend_warmup_patch._kstep_armed``, which -- when the engine config is not readable yet -- decides from that answer whether a 910B restores the native rejection sampler, i.e. whether the armed multi-frame path runs the sampler its spec width needs. Merge the NPU overlay the same way ``_apply_platform_overrides`` does before recording. Verified: both deploy profiles report 8 frames per step again. Signed-off-by: y-null --- vllm_omni/config/stage_config.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 04a602da4a8..a061f86a9f2 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -797,13 +797,30 @@ def talker_frames_per_step() -> int: return _resolved_talker_frames -def _record_talker_frames(stages: list[StageDeployConfig]) -> None: - """Record stage 1's frame count (its n-gram ``num_speculative_tokens`` is K - 1).""" +def _record_talker_frames( + stages: list[StageDeployConfig], + platforms: dict[str, Any] | None = None, +) -> None: + """Record stage 1's frame count (its n-gram ``num_speculative_tokens`` is K - 1). + + The K block sits under ``platforms.npu`` -- the loop is NPU-only -- so the + NPU overlay is merged here. This runs at load time, before the per-platform + merge, and its reader (``ascend_warmup_patch._kstep_armed``) runs in a + spawned stage worker that may have no vllm_config to read it from. + """ global _resolved_talker_frames + npu_overlay: dict[int, dict[str, Any]] = {} + for stage_override in ((platforms or {}).get("npu") or {}).get("stages") or []: + if isinstance(stage_override, dict) and "stage_id" in stage_override: + npu_overlay[stage_override["stage_id"]] = stage_override for stage in stages: if stage.stage_id != 1: continue - spec = (stage.engine_extras or {}).get("speculative_config") or {} + spec = (stage.engine_extras or {}).get("speculative_config") + if spec is None: + # Platform overrides follow _extract_platform_overrides: every key + # but stage_id/devices/env is an engine or stage override. + spec = (npu_overlay.get(1) or {}).get("speculative_config") frames = 1 if isinstance(spec, dict) and spec.get("method") == "ngram": num_spec = spec.get("num_speculative_tokens", 0) or 0 @@ -822,7 +839,7 @@ def load_deploy_config(path: str | Path) -> DeployConfig: ) stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])] - _record_talker_frames(stages) + _record_talker_frames(stages, raw_dict.get("platforms")) model_runner = raw_dict.get("model_runner", "v1") if model_runner not in ("v1", "v2"): From b804b9e1b1561d4df6bf32f5b9147560c47a735c Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 17:24:30 +0800 Subject: [PATCH 37/89] fix(minicpmo_4_5): rebuild the codec sampler's device state per duplex condition _sample_audio_code builds a request's device state (step counter, EOS floor, max_tokens, history window) only when the request_id is absent, and _request_codec_device_inputs caches the min_tokens tensor the same way. A duplex resume reuses the same request id, so the next segment inherited the previous segment's step counter and its chunk budget: the EOS floor and the cadence boundaries were computed against the old segment. The preprocess path already rebuilds the host-side state per condition; drop the device-side caches there too (duplex only), so the next decode re-creates them from the fresh budget. The repetition-penalty window is carried through the state's recent codes, matching main's clean-slate-except-penalty semantics. Signed-off-by: y-null (cherry picked from commit d490ff703821654cc7c4f94e82c5ad37ca2bdf2a) --- .../models/minicpmo_4_5/minicpmo_4_5_omni_tts.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 534a5a370af..68149124865 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -745,6 +745,20 @@ def preprocess( request_states = {} self._request_audio_states = request_states request_states[request_id] = state + if native_duplex: + # A new duplex condition starts a fresh codec segment: drop the + # sampler's device-side caches for this request so the next + # decode rebuilds them from this condition's state (step, EOS + # floor, limits) instead of resuming the previous segment's. + # The repetition-penalty window is carried through the state + # and ``codes`` above, so only the counters and the cached + # input tensors reset. + codec_states = getattr(self, "_request_codec_device_states", None) + if isinstance(codec_states, dict): + codec_states.pop(request_id, None) + codec_inputs = getattr(self, "_request_codec_device_inputs", None) + if isinstance(codec_inputs, dict): + codec_inputs.pop(request_id, None) empty_codes = torch.empty(0, dtype=torch.long, device=embeds.device) return ( input_ids, From ee12b1adc80f34385c1ae10e4c6c8a68dac0173a Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 17:27:11 +0800 Subject: [PATCH 38/89] fix(minicpmo_4_5): wire the duplex turn-end EOS window into the K-step sampler The single-frame path masks codec EOS within 5 steps after each 25-frame cadence boundary while a turn-end chunk drains (mask_eos_rows), so the Talker finishes the text it still owes instead of stopping on the periodic cadence EOS. The K-step path samples through talker_codec_sample.prepare_codec_logits / greedy_codec_sample, whose only EOS criterion was step < min_tokens, so that window was left unwired and a turn-end chunk could end a few frames early. Thread an eos_window_masked flag through both sampling cores and set it from the same turn_end_drain + _turn_end_boundary_eos_masked(step) criterion the single-frame path uses. Signed-off-by: y-null (cherry picked from commit c63b276a9d5d4af5e41cc5a76ad08c87d0a58f66) --- .../test_talker_kstep_feedback.py | 2 +- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 21 ++++++---- .../minicpmo_4_5/talker_codec_sample.py | 41 +++++++++++++------ 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 7ae1509cbbc..13be179f6c4 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -57,7 +57,7 @@ def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): model.emb_code = nn.ModuleList([emb]) queue = iter(scripted_samples) - def _greedy(_hidden, _codes, _request_id, _step, _min_tokens, _max_tokens): + def _greedy(_hidden, _codes, _request_id, _step, _min_tokens, _max_tokens, _eos_window_masked=False): return torch.tensor(next(queue), dtype=torch.long) model._sample_audio_code_greedy = _greedy diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 68149124865..0358f478726 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -936,15 +936,13 @@ def make_omni_output( # _turn_end_boundary_eos_masked and state["turn_end_drain"]; # MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS is 4 chunks, so the window # lands on steps 25-29 / 50-54 / 75-79 / 100-103). The EOS masking - # moved into the sampling core (talker_codec_sample. - # prepare_codec_logits / greedy_codec_sample), whose criterion is - # state.step < min_tokens with min_tokens a per-request cached - # tensor, so that periodic window cannot be expressed per frame here - # and is left unwired. It only affects duplex turn-end; the - # offline/TTS path is unaffected. The constants, budget and - # turn_end_drain flag from main are kept in this file. + # lives in the sampling core (talker_codec_sample. + # prepare_codec_logits / greedy_codec_sample), so the periodic + # window rides along as the eos_window_masked flag below -- the same + # criterion the single-frame path's mask_eos_rows applies. min_tokens = _codec_int_param(state, "min_tokens", self._codec_min_tokens) max_tokens = _codec_int_param(state, "max_tokens", self._codec_max_tokens) + eos_window_masked = bool(state.get("turn_end_drain")) and _turn_end_boundary_eos_masked(step) if self._codec_temperature == 0.0: sampled = self._sample_audio_code_greedy( hidden[end - 1 : end], @@ -953,9 +951,12 @@ def make_omni_output( step, min_tokens, max_tokens, + eos_window_masked, ) else: - stochastic_result = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step) + stochastic_result = self._sample_audio_code( + hidden[end - 1 : end], codes, request_id, step, eos_window_masked + ) if isinstance(stochastic_result, TalkerCodecSampleResult): sampled = stochastic_result.sampled_token.reshape(()).to(torch.long) else: @@ -1088,6 +1089,7 @@ def _sample_audio_code( history: torch.Tensor, request_id: str, step: int, + eos_window_masked: bool = False, ) -> TalkerCodecSampleResult: device_states = getattr(self, "_request_codec_device_states", None) if device_states is None: @@ -1138,6 +1140,7 @@ def _sample_audio_code( top_k=self._codec_top_k, top_p=self._codec_top_p, min_tokens_to_keep=3, + eos_window_masked=eos_window_masked, ) probabilities = torch.softmax(logits, dim=-1) sampled = torch.multinomial( @@ -1161,6 +1164,7 @@ def _sample_audio_code_greedy( step: int, min_tokens: int, max_tokens: int, + eos_window_masked: bool = False, ) -> torch.Tensor: """Run the greedy codec boundary while keeping sampler state on device. @@ -1205,6 +1209,7 @@ def _sample_audio_code_greedy( penalty_tensor, top_k=self._codec_top_k, eos_token_id=self._codec_eos_id, + eos_window_masked=eos_window_masked, ) device_states[request_id] = result.state # The kernel ABI is int32, while the existing connector/audio-code diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index 79abb49f822..78ee40f0c47 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -65,8 +65,14 @@ def prepare_codec_logits( top_k: int, top_p: float, min_tokens_to_keep: int = MIN_TOKENS_TO_KEEP, + eos_window_masked: bool = False, ) -> torch.Tensor: - """Prepare and filter logits without changing native NPU RNG semantics.""" + """Prepare and filter logits without changing native NPU RNG semantics. + + ``eos_window_masked`` masks the codec EOS for this frame regardless of the + step counter -- the duplex turn-end drain window (a few frames after each + cadence boundary), which the ``step < min_tokens`` criterion cannot express. + """ logits = raw_logits.float() / temperature.reshape(1, 1) positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) valid = positions < state.history_len.reshape(1, 1) @@ -76,11 +82,14 @@ def prepare_codec_logits( alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) logits = torch.where(logits < 0, logits * alpha, logits / alpha) eos = logits[..., eos_token_id : eos_token_id + 1] - eos = torch.where( - (state.step < min_tokens).reshape(1, 1), - torch.full_like(eos, float("-inf")), - eos, - ) + if eos_window_masked: + eos = torch.full_like(eos, float("-inf")) + else: + eos = torch.where( + (state.step < min_tokens).reshape(1, 1), + torch.full_like(eos, float("-inf")), + eos, + ) logits = torch.cat([logits[..., :eos_token_id], eos, logits[..., eos_token_id + 1 :]], dim=-1) if 0.0 < float(top_p) < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=False, dim=-1) @@ -142,8 +151,13 @@ def greedy_codec_sample( *, top_k: int, eos_token_id: int, + eos_window_masked: bool = False, ) -> TalkerCodecSampleResult: - """Greedy codec sample in graph-capturable tensor ops.""" + """Greedy codec sample in graph-capturable tensor ops. + + ``eos_window_masked`` masks the codec EOS for this frame regardless of the + step counter -- see prepare_codec_logits. + """ positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) valid = positions < state.history_len.reshape(1, 1) safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) @@ -153,11 +167,14 @@ def greedy_codec_sample( penalized = torch.where(raw_logits < 0, raw_logits * alpha, raw_logits / alpha) eos = penalized[..., eos_token_id : eos_token_id + 1] - eos = torch.where( - (state.step < min_tokens).reshape(1, 1), - torch.full_like(eos, float("-inf")), - eos, - ) + if eos_window_masked: + eos = torch.full_like(eos, float("-inf")) + else: + eos = torch.where( + (state.step < min_tokens).reshape(1, 1), + torch.full_like(eos, float("-inf")), + eos, + ) penalized = torch.cat([penalized[..., :eos_token_id], eos, penalized[..., eos_token_id + 1 :]], dim=-1) if top_k > 0: keep = min(VOCAB_SIZE, max(int(top_k), MIN_TOKENS_TO_KEEP)) From 3f23cffc43caa3e0ac9257c7b7fcd383ccc7212d Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 17:31:12 +0800 Subject: [PATCH 39/89] fix(minicpmo_4_5): feed the real codec history to streaming prompt recompute Under K-step the scheduler's token stream records the 0/1 control rows the two-wide head emits, not the codec ids the model sampled. The streaming prompt recompute path reads its confirmed-token slice from that stream and embeds it with emb_code, so after a rollover the rebuilt Talker condition was made of continue/stop marker embeddings: ids 0 and 1 have valid embeddings and pass the < EOS range check, so nothing failed loudly. Record the confirmed frame per forwarded position in a per-request codec history and swap the values in at recompute time -- same length, one row per confirmed frame, and the slice always covers the tail of that history. The history is dropped with the rest of the request's caches in _flush_deferred_cleanup. Signed-off-by: y-null (cherry picked from commit ea3eee3323dfc0a39041d14e1bfda54be1f07d8e) --- .../test_talker_kstep_feedback.py | 5 ++ .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 55 +++++++++++++++++-- 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 13be179f6c4..446da5079ec 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -49,6 +49,7 @@ def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): model._num_audio_tokens = _NUM_AUDIO_TOKENS model._codec_eos_id = _EOS_ID model._request_audio_states = {} + model._request_codec_history = {} emb = nn.Embedding(_NUM_AUDIO_TOKENS, 4) with torch.no_grad(): emb.weight.zero_() @@ -94,6 +95,10 @@ def test_kstep_frame_sample_is_recorded_for_next_frame(): assert state["last_code"] == 43 assert state["finished"] is True assert out2.multimodal_outputs["codes"]["audio"][0].numel() == 0 + # Only confirmed frames are recorded: the terminating frame carries no + # codec id for the next step to embed, so streaming prompt recompute must + # not see it in the history either. + assert model._request_codec_history["r1"] == [42, 43] def test_legacy_single_frame_path_does_not_touch_last_code(): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 0358f478726..a37352a20a8 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -339,6 +339,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self._request_audio_states: dict[str, dict[str, Any]] = {} self._request_codec_device_states: dict[str, TalkerCodecDeviceState] = {} self._request_codec_device_inputs: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + # Confirmed codec frames this model sampled, one entry per forwarded + # frame. Under K-step the scheduler's token stream carries the 0/1 + # control rows the two-wide head emits instead, so streaming prompt + # recompute swaps its confirmed-token slice for these real ids. + self._request_codec_history: dict[str, list[int]] = {} self._init_native_talker(prefix) @@ -568,6 +573,23 @@ def _build_streaming_recompute_embeddings( ids = info_dict.get("ids") previous_codes = ids.get("streaming_prompt_previous_codes") if isinstance(ids, Mapping) else None + if self._k_step_frames > 0 and previous_codes is not None: + # K-step: the scheduler's confirmed slice carries the 0/1 control + # rows the two-wide head emits, not codec ids -- feeding them to + # emb_code would silently rebuild the prompt from continue/stop + # markers (both have valid embeddings and pass the range check + # below). The model recorded the real frame it confirmed at each + # of those positions, one row per forwarded frame, so swap the + # values in. The slice only ever covers the frames confirmed + # since the current prompt, i.e. the tail of that history. + codec_history = getattr(self, "_request_codec_history", {}).get(request_id) + if codec_history is None or len(codec_history) < len(previous_codes): + raise ValueError( + "K-step streaming prompt recompute is missing confirmed codec frames: " + f"history={len(codec_history) if codec_history is not None else 0}, " + f"confirmed={len(previous_codes)}" + ) + previous_codes = codec_history[-len(previous_codes) :] if isinstance(previous_codes, torch.Tensor): code_ids = previous_codes.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1) elif isinstance(previous_codes, (list, tuple)): @@ -979,8 +1001,24 @@ def make_omni_output( # vLLM schedules for the next step carry placeholder # continue ids. Record this frame's sample so the decode # preprocess embeds the real previous frame's codec id - # (the k_last branch) instead of the placeholder. + # (the k_last branch) instead of the placeholder, and keep + # the confirmed-frame history that streaming prompt + # recompute rebuilds its window from -- the scheduler's + # token slice holds the placeholder ids. state["last_code"] = sampled_id + histories = getattr(self, "_request_codec_history", None) + if not isinstance(histories, dict): + # Instances built without __init__ (CPU tests, + # subclasses) still need a place for the record -- the + # readers at the recompute and cleanup sites already + # tolerate a missing dict the same way. + histories = {} + self._request_codec_history = histories + codec_history = histories.get(request_id) + if codec_history is None: + codec_history = [] + histories[request_id] = codec_history + codec_history.append(sampled_id) else: delta = empty_delta state["codes"] = codes @@ -1278,11 +1316,13 @@ def _flush_deferred_cleanup(self) -> None: ``_request_codec_device_states`` carries the step counter and the penalty window the EOS mask is derived from, and ``_request_codec_device_inputs`` carries the ``min_tokens`` copied at - first use. Leaving them behind means a recycled request id resumes with - another request's counters -- the EOS mask then holds or lifts at the - wrong frame, and ``_request_generators`` leaks a generator per request - on top of that. All three key off the same request id, so they are - dropped together with the rest. + first use. ``_request_codec_history`` carries the confirmed K-step + frames streaming prompt recompute reads. Leaving them behind means a + recycled request id resumes with another request's counters -- the EOS + mask then holds or lifts at the wrong frame, and + ``_request_generators`` leaks a generator per request on top of that. + All of them key off the same request id, so they are dropped together + with the rest. """ request_audio_states = getattr(self, "_request_audio_states", {}) request_condition_states = getattr(self, "_request_condition_states", {}) @@ -1290,6 +1330,7 @@ def _flush_deferred_cleanup(self) -> None: penalty_freqs = getattr(self, "_penalty_freqs_dev", None) codec_device_states = getattr(self, "_request_codec_device_states", None) codec_device_inputs = getattr(self, "_request_codec_device_inputs", None) + codec_histories = getattr(self, "_request_codec_history", None) request_generators = getattr(self, "_request_generators", None) for request_id in self._deferred_cleanup_ids: request_audio_states.pop(request_id, None) @@ -1302,6 +1343,8 @@ def _flush_deferred_cleanup(self) -> None: codec_device_states.pop(request_id, None) if isinstance(codec_device_inputs, dict): codec_device_inputs.pop(request_id, None) + if isinstance(codec_histories, dict): + codec_histories.pop(request_id, None) if isinstance(request_generators, dict): request_generators.pop(request_id, None) self._deferred_cleanup_ids.clear() From 8d61b99e2c993bac6f3446fc891cd7119ac9f15f Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 22:07:39 +0800 Subject: [PATCH 40/89] fix(minicpmo_4_5): derive the Talker stop id from the armed head, not the platform The deploy config that arms the multi-frame head already pins that head's stop marker (1) through stage 1's ``default_sampling_params``, and ``merge_sampling_constraints`` unions it with whatever the pipeline supplies. So the pipeline's only job is to supply the stop id of the *other* head -- the codec EOS of a stage that samples the full codec vocabulary. Keyed on the platform instead, the NPU branch returned the marker even for an NPU deployment that drops stage 1's ``speculative_config`` (one frame per step again): its head is the full codec vocabulary there, so 1 is an ordinary audio id and the request ends mid-utterance -- the same defect this default was introduced to remove, just narrower, and the one case the review called out. No shipped configuration changes: armed NPU configs still merge to [1, 6561] (1 ends the request; 6561 is unsampleable in the collapsed head) and non-NPU configs still stop on 6561 alone. Only NPU-without-the-block moves from a wrong 1 to the correct 6561. Dropping the platform probe also drops the import of the NPU worker module from the pipeline, so resolving the pipeline no longer drags in a platform worker. Tests: the topology test and the stop-id test now pin the pipeline default to the codec EOS on every platform, and the deploy-scoping test pins both merged lists ([1, 6561] armed, [6561] without the block). Verified: 69 passed across test_pipeline.py, test_talker_kstep_feedback.py and test_config_factory.py. Signed-off-by: y-null --- tests/config/test_config_factory.py | 9 +++++- .../models/minicpmo_4_5/test_pipeline.py | 12 ++++---- .../test_talker_kstep_feedback.py | 11 +++----- .../models/minicpmo_4_5/pipeline.py | 28 ++++++++----------- 4 files changed, 30 insertions(+), 30 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 299a7d54957..84af75fedc6 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1384,12 +1384,15 @@ def test_minicpmo_talker_multi_frame_is_npu_scoped(self, platform: str): platform=platform, ) # Check the deploy-side stop id before the merge: afterwards it is - # unioned with the pipeline's platform-aware constraint. + # unioned with the pipeline's codec-EOS default. deploy_stage1 = next(stage for stage in deploy.stages if stage.stage_id == 1) stages = merge_pipeline_deploy(resolve_pipeline_config("minicpmo_4_5"), deploy) if platform == "npu": assert deploy_stage1.default_sampling_params["stop_token_ids"] == [1] + # The armed head ends on 1; the union keeps the codec EOS inert + # (the collapsed head cannot sample it). + assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1, 6561] assert stages[0].yaml_engine_args["enable_prefix_caching"] is True assert stages[1].yaml_engine_args["speculative_config"] == { "method": "ngram", @@ -1401,6 +1404,10 @@ def test_minicpmo_talker_multi_frame_is_npu_scoped(self, platform: str): assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [8, 8, 8] else: assert "stop_token_ids" not in (deploy_stage1.default_sampling_params or {}) + # Without the block the head is the full codec vocabulary, so the + # codec EOS is the only stop id -- the marker would end a request + # on an ordinary codec id. + assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] # Either absent (falls back to the top-level false) or explicitly # false; the regression this guards against is true. assert stages[0].yaml_engine_args.get("enable_prefix_caching") is not True diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 15fbb54eaa2..d8084fb2f0d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -110,14 +110,14 @@ def test_talker_stage(self, pipeline: PipelineConfig) -> None: assert talker.engine_output_type == "latent" # scope KV cache / mrope sizing to talker sub-config assert talker.hf_config_name == "tts_config" - # The multi-frame head collapses the vLLM-level head to the two-wide - # continue/stop row, whose only stop id is 1. It is an NPU-worker - # feature (that is where the continue drafts are injected), so on any - # other platform stage 1 keeps its full codec head and the codec EOS. - expected_stop = [1] if (current_omni_platform.device_name or "").lower() == "npu" else [_CODEC_EOS_TOKEN_ID] + # The pipeline default is the codec EOS of the one-frame head. The + # deploy config that arms the multi-frame row (stage 1's + # speculative_config, under platforms.npu) adds that row's stop marker + # (1) through its default_sampling_params; see + # test_minicpmo_talker_multi_frame_is_npu_scoped for the merged list. assert talker.sampling_constraints == { "detokenize": False, - "stop_token_ids": expected_stop, + "stop_token_ids": [_CODEC_EOS_TOKEN_ID], } assert talker.custom_process_next_stage_input_func == ( "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload" diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 446da5079ec..93a9f2a306e 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -251,16 +251,13 @@ def test_talker_stop_token_ids_match_the_multi_frame_head(): NPU worker, and the codec EOS on every platform without that worker. """ from vllm_omni.model_executor.models.minicpmo_4_5 import pipeline as mcp_pipeline - from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms.npu.worker import talker_multiframe assert talker_multiframe.STOP_TOKEN_ID == 1 - if (current_omni_platform.device_name or "").lower() == "npu": - assert mcp_pipeline._talker_stop_token_ids() == [talker_multiframe.STOP_TOKEN_ID] - else: - # No NPU worker means no constant-draft injection, so the K-frame loop - # is off and the two-wide continue/stop row never runs. - assert mcp_pipeline._talker_stop_token_ids() == [mcp_pipeline._CODEC_EOS_TOKEN_ID] + # The pipeline default cannot depend on the platform: the block that + # collapses the head is a deploy-config decision, so the marker is added + # there (test_minicpmo_talker_multi_frame_is_npu_scoped pins the merge). + assert mcp_pipeline._talker_stop_token_ids() == [mcp_pipeline._CODEC_EOS_TOKEN_ID] def test_multiframe_gate_matrix(): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 7eb9878720d..c91ce39f530 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -22,24 +22,20 @@ def _talker_stop_token_ids() -> list[int]: - """Stage 1's default stop id, derived from the head this platform runs. + """Stage 1's default stop id: the codec EOS of the one-frame head. - On NPU, stage 1's ``speculative_config`` collapses the vLLM-level head to - the two-wide continue/stop row: only 0/1 are sampleable, so the codec EOS - can never appear. Everywhere else that block stays out of the config (it is - confined to ``platforms.npu``), the head is the full codec vocabulary, and - the codec EOS is the only valid stop. - """ - from vllm_omni.platforms import current_omni_platform - - # The deploy config confines stage 1's speculative_config to platforms.npu, - # so a non-NPU stage 1 keeps the full codec head. - if (current_omni_platform.device_name or "").lower() != "npu": - return [_CODEC_EOS_TOKEN_ID] + The deploy config that arms the multi-frame decode (stage 1's + ``speculative_config``, under ``platforms.npu``) collapses the vLLM-level + head to the two-wide continue/stop row, where the codec EOS can never + appear; that same config adds the stop marker (1) through its + ``default_sampling_params``, and ``merge_sampling_constraints`` unions both + lists. The marker therefore travels with the block that needs it, and this + default covers every deployment without one. - from vllm_omni.platforms.npu.worker.talker_multiframe import STOP_TOKEN_ID - - return [STOP_TOKEN_ID] + A deployment that keeps the marker here instead ends every one-frame + request on the first ordinary codec id 1. + """ + return [_CODEC_EOS_TOKEN_ID] MINICPMO_4_5_PIPELINE = PipelineConfig( From dfda666c127715d0b4d30062e8a6921c62c61199 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 20:59:10 +0800 Subject: [PATCH 41/89] fix(minicpmo_4_5): scope the NPU-only stage overrides to the NPU platform The Talker multi-frame decode and the Thinker prefix cache were declared in the platform-shared stage blocks, so a CUDA deployment inherited both: * stage 1 ran vLLM's real n-gram proposer next to the rejection sampler, which asserts on the Talker's codec ids (device-side assert out of get_token_bin_counts_and_mask); * stage 0 turned prefix caching on, and a turn-2 prefix hit left the Thinker with an empty tts handoff (ValueError in minicpmo_4_5_omni_tts.preprocess) together with a device-side assert out of indexSelectSmallIndex. Either one kills the stage-1 replica, and every later duplex case then fails with "stage 1 has no live replica". Move both blocks into platforms.npu: the NPU behaviour is unchanged, because the NPU worker intercepts the decode step and injects the constant continue drafts itself. CUDA now falls back to the top-level enable_prefix_caching: false and to the codec EOS stop id, which is also what the stop-id helper follows. Verification on L20X (single card), numbers against the PR-free baseline in the PR description: * duplex, 11 cases: 7 passed / 4 failed versus 11 failed before; the 4 remaining are the upstream response timeouts the baseline shows as well, and the engine stayed alive with 0 ValueError and 0 device-side asserts. * offline seed-tts perf (32@1 / 64@4 / 128@8) and the Seed-TTS WER suite are unchanged. * unit tests: 25 config cases and 563 MiniCPM-o model executor cases. Signed-off-by: y-null --- tests/config/test_config_factory.py | 48 +++++++++++++++ .../models/minicpmo_4_5/test_pipeline.py | 21 ++++--- .../test_talker_kstep_feedback.py | 13 +++- vllm_omni/config/stage_config.py | 5 ++ vllm_omni/deploy/minicpmo_4_5.yaml | 60 ++++++++++++------- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 60 ++++++++++++------- .../models/minicpmo_4_5/pipeline.py | 14 +++++ 7 files changed, 169 insertions(+), 52 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index fdba7e90a2a..6e37a87ad20 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1363,6 +1363,54 @@ def test_load_minicpmo_default_deploy_config(self): assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] assert "codec_sampling_params" not in stages[1].yaml_engine_args + @pytest.mark.parametrize("platform", ["cuda", "npu"]) + def test_minicpmo_talker_multi_frame_is_npu_scoped(self, platform: str): + """Stage 1's K-frame decode needs the NPU worker to inject constant + continue drafts, so the deploy block that arms it must not sit in the + platform-shared stage. + + On NPU the ``speculative_config`` is a fingerprint: the worker replays + the decode graph and vLLM's n-gram proposer is never consulted. On CUDA + the same block hands the Talker to the real proposer plus the rejection + sampler, which asserts on its codec ids (device-side assert out of + ``get_token_bin_counts_and_mask``). + + Stage 0's prefix caching is NPU-scoped for the same reason: a CUDA + stage-0 prefix hit rewrites the prefill bookkeeping, so the Thinker + hands an empty tts handoff and the Talker indexes past the codec vocab + (ValueError in ``minicpmo_4_5_omni_tts.preprocess`` plus a + device-side assert out of ``indexSelectSmallIndex``). + """ + deploy = _apply_platform_overrides( + load_deploy_config(Path(get_deploy_config_path("minicpmo_4_5.yaml"))), + platform=platform, + ) + # Assert the deploy-side stop id before the merge: afterwards it is + # unioned with the pipeline's own platform-aware constraint, so an + # exact-list check here would depend on the host running the suite. + deploy_stage1 = next(stage for stage in deploy.stages if stage.stage_id == 1) + stages = merge_pipeline_deploy(resolve_pipeline_config("minicpmo_4_5"), deploy) + + if platform == "npu": + assert deploy_stage1.default_sampling_params["stop_token_ids"] == [1] + assert stages[0].yaml_engine_args["enable_prefix_caching"] is True + assert stages[1].yaml_engine_args["speculative_config"] == { + "method": "ngram", + "num_speculative_tokens": 7, + "prompt_lookup_min": 1, + "prompt_lookup_max": 1, + } + assert stages[1].yaml_engine_args["async_scheduling"] is False + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [8, 8, 8] + else: + assert "stop_token_ids" not in (deploy_stage1.default_sampling_params or {}) + # Either absent (falls back to the top-level false) or explicitly + # false; the regression this guards against is true. + assert stages[0].yaml_engine_args.get("enable_prefix_caching") is not True + assert "speculative_config" not in stages[1].yaml_engine_args + assert stages[1].yaml_engine_args["async_scheduling"] is True + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4] + @pytest.mark.parametrize( ("filename", "stage0_devices", "stage1_devices", "stage2_devices", "stage1_replicas"), [ diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 855af4185bf..15fbb54eaa2 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -29,7 +29,11 @@ load_deploy_config, merge_pipeline_deploy, ) +from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import ( + _CODEC_EOS_TOKEN_ID, +) from vllm_omni.model_executor.models.registry import _OMNI_MODELS +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -106,12 +110,14 @@ def test_talker_stage(self, pipeline: PipelineConfig) -> None: assert talker.engine_output_type == "latent" # scope KV cache / mrope sizing to talker sub-config assert talker.hf_config_name == "tts_config" + # The multi-frame head collapses the vLLM-level head to the two-wide + # continue/stop row, whose only stop id is 1. It is an NPU-worker + # feature (that is where the continue drafts are injected), so on any + # other platform stage 1 keeps its full codec head and the codec EOS. + expected_stop = [1] if (current_omni_platform.device_name or "").lower() == "npu" else [_CODEC_EOS_TOKEN_ID] assert talker.sampling_constraints == { "detokenize": False, - # The multi-frame head: stage 1's deploy config carries the - # speculative_config that collapses the vLLM-level head to the - # two-wide continue/stop row, whose only stop id is 1. - "stop_token_ids": [1], + "stop_token_ids": expected_stop, } assert talker.custom_process_next_stage_input_func == ( "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload" @@ -177,9 +183,10 @@ def test_deploy_resolves_three_stage_pipeline(self, filename: str, devices: list assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor) assert "hf_overrides" not in stages[1].yaml_engine_args if filename == "minicpmo_4_5.yaml": - # All three stages cap out at 8 so the top tier is not gated by a - # single stage. - assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [8, 8, 8] + # The NPU overlay caps all three stages at 8 so the top tier is not + # gated by a single stage; every other platform keeps the base 4. + expected_seqs = 8 if (current_omni_platform.device_name or "").lower() == "npu" else 4 + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [expected_seqs] * 3 memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] assert memory_utilizations == [ 0.55, diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index a0f5ba31e77..ff953aba172 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -247,14 +247,21 @@ def test_talker_stop_token_ids_match_the_multi_frame_head(): is the two-wide continue/stop row, so the only sampleable ids are 0/1, `check_stop` never matched, and every request ran to ``max_tokens`` (142s per request, model finished at frame ~116). The pipeline default therefore - follows the multi-frame head; a deployment that turns the loop off - overrides stage 1's stop token per stage. + follows the head that runs: the multi-frame row next to the NPU worker, + which is the only place that injects the continue drafts, and the codec EOS + everywhere else. """ from vllm_omni.model_executor.models.minicpmo_4_5 import pipeline as mcp_pipeline + from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms.npu.worker import talker_multiframe - assert mcp_pipeline._talker_stop_token_ids() == [talker_multiframe.STOP_TOKEN_ID] assert talker_multiframe.STOP_TOKEN_ID == 1 + if (current_omni_platform.device_name or "").lower() == "npu": + assert mcp_pipeline._talker_stop_token_ids() == [talker_multiframe.STOP_TOKEN_ID] + else: + # No NPU worker means no constant-draft injection, so the K-frame loop + # is off and the two-wide continue/stop row never runs. + assert mcp_pipeline._talker_stop_token_ids() == [mcp_pipeline._CODEC_EOS_TOKEN_ID] def test_multiframe_gate_matrix(): diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index e7bdd6994d8..04a602da4a8 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -779,6 +779,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: return merged +# Same TND ceiling as stage 0: at most 16 query positions per sequence. +# Consumed by minicpmo_4_5_omni_tts._parse_k_step_frames as the hard cap on +# the K-step frame count accepted from the deploy YAML. +_MINICPMO_TALKER_FRAMES_MAX = 16 + # Stage 1's multi-frame decode is configured through its ``speculative_config`` # in the deploy YAML. The parse records whether it is armed here, because the # reader (``ascend_warmup_patch._kstep_armed``) can run in a spawned stage diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 05104f335dd..b73d702e89f 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -41,13 +41,18 @@ stages: # Leave headroom on the shared GPU for Talker KV + concurrent Code2Wav. gpu_memory_utilization: 0.55 enforce_eager: false - # Stage-level prefix caching override; the top level stays false for duplex safety. - # Stage-0 n-gram speculative decoding is intentionally not enabled here: the + # Stage-level prefix caching is NPU-only (see platforms.npu below); + # the top level stays false for duplex safety. Observed with this flag + # in the platform-shared block on CUDA: the first duplex case dies with + # an empty tts handoff (token_ids/hidden_states=None -> ValueError in + # minicpmo_4_5_omni_tts.preprocess) alongside a device-side assert out + # of indexSelectSmallIndex; the stage-1 replica is killed and every later + # case fails with "stage 1 has no live replica". Stage-0 n-gram + # speculative decoding stays off too: the # omni prefix cache refuses the pair (stage_prefix_cache_config -> - # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError as soon - # as a stage has enable_prefix_caching and a speculative_config; vLLM itself - # allows the combination). - enable_prefix_caching: true + # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError + # as soon as a stage has enable_prefix_caching and a speculative_config; + # vLLM itself allows the combination). tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -84,19 +89,6 @@ stages: # Talker context is 4096 (tts_config.max_position_embeddings). max_model_len: 4096 devices: "0" - # Multi-frame decode: the Talker produces 8 codec frames per scheduler step. - # vLLM learns that each request advances 8 tokens through this - # speculative_config (the n-gram proposer is never consulted -- the runner - # supplies constant continue drafts), and async scheduling is rejected next - # to an n-gram proposer. stop_token_ids above must stay the two-wide - # continue/stop id the head emits in this mode. Remove the block (or set it - # to null) to fall back to one frame per step. - speculative_config: - method: ngram - num_speculative_tokens: 7 - prompt_lookup_min: 1 - prompt_lookup_max: 1 - async_scheduling: false output_connectors: to_stage_2: connector_of_shared_memory # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler @@ -112,8 +104,6 @@ stages: top_k: 25 repetition_penalty: 1.05 min_tokens: 50 - # The Talker's stop head emits this id to finish a request. - stop_token_ids: [1] max_tokens: 4096 seed: 42 detokenize: false @@ -152,6 +142,10 @@ platforms: # not gated by a single stage. max_num_seqs: 8 max_num_batched_tokens: 8192 + # Stage-0 prefix caching was moved here from the platform-shared + # block: keep the PR's NPU behaviour unchanged while CUDA falls back + # to the top-level false (duplex safety, see the base stage comment). + enable_prefix_caching: true compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame # expansion; PIECEWISE silently drops both. Measured on 910B3: @@ -162,6 +156,30 @@ platforms: # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 + # Multi-frame decode: the Talker produces 8 codec frames per scheduler + # step. vLLM learns that each request advances 8 tokens through this + # speculative_config; the n-gram proposer itself is never consulted + # because the NPU worker replays the decode graph with constant continue + # drafts, and async scheduling is rejected next to an n-gram proposer. + # stop_token_ids must stay the two-wide continue/stop id the head emits + # in this mode. Remove the block (or set it to null) to fall back to one + # frame per step. + # + # NPU-only on purpose: only + # vllm_omni.platforms.npu.worker.talker_multiframe intercepts the step + # and injects those drafts. Left in the base stage, the same block hands + # stage 1 to vLLM's real n-gram proposer and rejection sampler on CUDA, + # which asserts on the Talker's codec ids (device-side assert out of + # get_token_bin_counts_and_mask). + speculative_config: + method: ngram + num_speculative_tokens: 7 + prompt_lookup_min: 1 + prompt_lookup_max: 1 + async_scheduling: false + default_sampling_params: + # The Talker's stop head emits this id to finish a request. + stop_token_ids: [1] compilation_config: cudagraph_mode: FULL_DECODE_ONLY - stage_id: 2 diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 6d1949af723..52cd8f8e735 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -45,13 +45,18 @@ stages: # replicas below; the NPU block caps stage 0/1 KV explicitly. gpu_memory_utilization: 0.40 enforce_eager: false - # Stage-level prefix caching override; the top level stays false for duplex safety. - # Stage-0 n-gram speculative decoding is intentionally not enabled here: the + # Stage-level prefix caching is NPU-only (see platforms.npu below); + # the top level stays false for duplex safety. Observed with this flag + # in the platform-shared block on CUDA: the first duplex case dies with + # an empty tts handoff (token_ids/hidden_states=None -> ValueError in + # minicpmo_4_5_omni_tts.preprocess) alongside a device-side assert out + # of indexSelectSmallIndex; the stage-1 replica is killed and every later + # case fails with "stage 1 has no live replica". Stage-0 n-gram + # speculative decoding stays off too: the # omni prefix cache refuses the pair (stage_prefix_cache_config -> - # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError as soon - # as a stage has enable_prefix_caching and a speculative_config; vLLM itself - # allows the combination). - enable_prefix_caching: true + # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError + # as soon as a stage has enable_prefix_caching and a speculative_config; + # vLLM itself allows the combination). tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -88,19 +93,6 @@ stages: # Talker context is 4096 (tts_config.max_position_embeddings). max_model_len: 4096 devices: "0" - # Multi-frame decode: the Talker produces 8 codec frames per scheduler step. - # vLLM learns that each request advances 8 tokens through this - # speculative_config (the n-gram proposer is never consulted -- the runner - # supplies constant continue drafts), and async scheduling is rejected next - # to an n-gram proposer. stop_token_ids above must stay the two-wide - # continue/stop id the head emits in this mode. Remove the block (or set it - # to null) to fall back to one frame per step. - speculative_config: - method: ngram - num_speculative_tokens: 7 - prompt_lookup_min: 1 - prompt_lookup_max: 1 - async_scheduling: false output_connectors: to_stage_2: connector_of_shared_memory # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler @@ -116,8 +108,6 @@ stages: top_k: 25 repetition_penalty: 1.05 min_tokens: 50 - # The Talker's stop head emits this id to finish a request. - stop_token_ids: [1] max_tokens: 4096 seed: 42 detokenize: false @@ -159,6 +149,10 @@ platforms: # not gated by a single stage. max_num_seqs: 8 max_num_batched_tokens: 8192 + # Stage-0 prefix caching was moved here from the platform-shared + # block: keep the PR's NPU behaviour unchanged while CUDA falls back + # to the top-level false (duplex safety, see the base stage comment). + enable_prefix_caching: true # Explicit KV cap: the auto-computed value does not leave room for the # two stage-2 replicas on this card. kv_cache_memory_bytes: 6442450944 @@ -172,6 +166,30 @@ platforms: # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 + # Multi-frame decode: the Talker produces 8 codec frames per scheduler + # step. vLLM learns that each request advances 8 tokens through this + # speculative_config; the n-gram proposer itself is never consulted + # because the NPU worker replays the decode graph with constant continue + # drafts, and async scheduling is rejected next to an n-gram proposer. + # stop_token_ids must stay the two-wide continue/stop id the head emits + # in this mode. Remove the block (or set it to null) to fall back to one + # frame per step. + # + # NPU-only on purpose: only + # vllm_omni.platforms.npu.worker.talker_multiframe intercepts the step + # and injects those drafts. Left in the base stage, the same block hands + # stage 1 to vLLM's real n-gram proposer and rejection sampler on CUDA, + # which asserts on the Talker's codec ids (device-side assert out of + # get_token_bin_counts_and_mask). + speculative_config: + method: ngram + num_speculative_tokens: 7 + prompt_lookup_min: 1 + prompt_lookup_max: 1 + async_scheduling: false + default_sampling_params: + # The Talker's stop head emits this id to finish a request. + stop_token_ids: [1] # Explicit KV cap: the auto-computed value does not leave room for the # two stage-2 replicas on this card. kv_cache_memory_bytes: 2684354560 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 68c77082e23..d252bc4712c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -18,6 +18,7 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni" MINICPMO45_REFERENCE_AUDIO_KEY = "_minicpmo45_reference_audio" +_CODEC_EOS_TOKEN_ID = 6561 # tts_config.num_audio_tokens - 1 def _talker_stop_token_ids() -> list[int]: @@ -38,6 +39,19 @@ def _talker_stop_token_ids() -> list[int]: vocabulary again and the codec EOS is the stop id; it overrides this per stage through the usual ``stop_token_ids`` sampling parameter. """ + from vllm_omni.platforms import current_omni_platform + + # ``speculative_config`` only means "constant continue drafts" next to the + # NPU worker that intercepts the step + # (``vllm_omni.platforms.npu.worker.talker_multiframe``); on every other + # platform the same block hands the Talker to vLLM's real n-gram proposer and + # rejection sampler, which asserts on the codec ids. The block is therefore + # confined to ``platforms.npu`` in the deploy config, and the stop id follows: + # NPU keeps the two-wide continue/stop row, everyone else keeps the full codec + # head and stops on the codec EOS. + if (current_omni_platform.device_name or "").lower() != "npu": + return [_CODEC_EOS_TOKEN_ID] + from vllm_omni.platforms.npu.worker.talker_multiframe import STOP_TOKEN_ID return [STOP_TOKEN_ID] From 2ff445b9c54b2d07edd9fca3f6367cea0e5b626a Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 21:42:09 +0800 Subject: [PATCH 42/89] docs(minicpmo_4_5): trim the NPU-scoping comments to what the code needs Review pass over the comments the NPU-scoping change added: they narrated the CUDA failure that motivated the scoping at a length the code does not need, and the pipeline docstring still described the previous unconditional codec-EOS pinning instead of what the function now does. - deploy configs (both profiles): the stage-0 prefix-caching note and the multi-frame block keep the reason and the two things a reader must not break (the block is NPU-only; its stop id travels with it). The CUDA traces shrink to one clause, and the stage-0 note next to the moved flag becomes a pointer to it. - pipeline.py: the docstring now states that the stop id follows the head this platform runs; the inline comment keeps only why the early return exists. - tests: the two long rationales keep the conclusion, drop the incident log. Comments and docstrings only: no behavior change (verified by comparing normalized token streams for the .py files and the parsed YAML for the two deploy configs). Signed-off-by: y-null --- tests/config/test_config_factory.py | 11 ++--- .../test_talker_kstep_feedback.py | 10 +---- vllm_omni/deploy/minicpmo_4_5.yaml | 41 +++++-------------- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 41 +++++-------------- .../models/minicpmo_4_5/pipeline.py | 31 ++++---------- 5 files changed, 36 insertions(+), 98 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 6e37a87ad20..299a7d54957 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1376,18 +1376,15 @@ def test_minicpmo_talker_multi_frame_is_npu_scoped(self, platform: str): ``get_token_bin_counts_and_mask``). Stage 0's prefix caching is NPU-scoped for the same reason: a CUDA - stage-0 prefix hit rewrites the prefill bookkeeping, so the Thinker - hands an empty tts handoff and the Talker indexes past the codec vocab - (ValueError in ``minicpmo_4_5_omni_tts.preprocess`` plus a - device-side assert out of ``indexSelectSmallIndex``). + stage-0 prefix hit hands the Thinker an empty tts handoff, and the + Talker then indexes past the codec vocab. """ deploy = _apply_platform_overrides( load_deploy_config(Path(get_deploy_config_path("minicpmo_4_5.yaml"))), platform=platform, ) - # Assert the deploy-side stop id before the merge: afterwards it is - # unioned with the pipeline's own platform-aware constraint, so an - # exact-list check here would depend on the host running the suite. + # Check the deploy-side stop id before the merge: afterwards it is + # unioned with the pipeline's platform-aware constraint. deploy_stage1 = next(stage for stage in deploy.stages if stage.stage_id == 1) stages = merge_pipeline_deploy(resolve_pipeline_config("minicpmo_4_5"), deploy) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index ff953aba172..7ae1509cbbc 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -242,14 +242,8 @@ def test_guard_keeps_drafts_when_all_decodes_share_the_padded_width(): def test_talker_stop_token_ids_match_the_multi_frame_head(): """Stage 1's stop id must be one the head that actually runs can emit. - The pipeline constraint used to pin ``stop_token_ids`` to the codec EOS - (6561) unconditionally. With the multi-frame loop on, the vLLM-level head - is the two-wide continue/stop row, so the only sampleable ids are 0/1, - `check_stop` never matched, and every request ran to ``max_tokens`` (142s - per request, model finished at frame ~116). The pipeline default therefore - follows the head that runs: the multi-frame row next to the NPU worker, - which is the only place that injects the continue drafts, and the codec EOS - everywhere else. + The default follows the head: the two-wide continue/stop row next to the + NPU worker, and the codec EOS on every platform without that worker. """ from vllm_omni.model_executor.models.minicpmo_4_5 import pipeline as mcp_pipeline from vllm_omni.platforms import current_omni_platform diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index b73d702e89f..162ef0b9f70 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -41,18 +41,11 @@ stages: # Leave headroom on the shared GPU for Talker KV + concurrent Code2Wav. gpu_memory_utilization: 0.55 enforce_eager: false - # Stage-level prefix caching is NPU-only (see platforms.npu below); - # the top level stays false for duplex safety. Observed with this flag - # in the platform-shared block on CUDA: the first duplex case dies with - # an empty tts handoff (token_ids/hidden_states=None -> ValueError in - # minicpmo_4_5_omni_tts.preprocess) alongside a device-side assert out - # of indexSelectSmallIndex; the stage-1 replica is killed and every later - # case fails with "stage 1 has no live replica". Stage-0 n-gram - # speculative decoding stays off too: the - # omni prefix cache refuses the pair (stage_prefix_cache_config -> - # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError - # as soon as a stage has enable_prefix_caching and a speculative_config; - # vLLM itself allows the combination). + # Stage-level prefix caching is NPU-only (see platforms.npu below): in the + # platform-shared block it killed the stage-1 replica on CUDA (empty tts + # handoff, then an indexSelectSmallIndex device assert). Stage-0 n-gram + # speculation stays off either way -- the omni prefix cache refuses it next + # to enable_prefix_caching. tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -142,9 +135,7 @@ platforms: # not gated by a single stage. max_num_seqs: 8 max_num_batched_tokens: 8192 - # Stage-0 prefix caching was moved here from the platform-shared - # block: keep the PR's NPU behaviour unchanged while CUDA falls back - # to the top-level false (duplex safety, see the base stage comment). + # Moved here from the platform-shared block; see the base stage note. enable_prefix_caching: true compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame @@ -156,21 +147,11 @@ platforms: # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 - # Multi-frame decode: the Talker produces 8 codec frames per scheduler - # step. vLLM learns that each request advances 8 tokens through this - # speculative_config; the n-gram proposer itself is never consulted - # because the NPU worker replays the decode graph with constant continue - # drafts, and async scheduling is rejected next to an n-gram proposer. - # stop_token_ids must stay the two-wide continue/stop id the head emits - # in this mode. Remove the block (or set it to null) to fall back to one - # frame per step. - # - # NPU-only on purpose: only - # vllm_omni.platforms.npu.worker.talker_multiframe intercepts the step - # and injects those drafts. Left in the base stage, the same block hands - # stage 1 to vLLM's real n-gram proposer and rejection sampler on CUDA, - # which asserts on the Talker's codec ids (device-side assert out of - # get_token_bin_counts_and_mask). + # Multi-frame decode: 8 codec frames per step. NPU-only -- only the NPU + # worker intercepts the step and injects the constant continue drafts, + # so the block stays in this section. stop_token_ids below is the + # two-wide continue/stop id that head emits; drop the block (or set it + # to null) for one frame per step, and drop that id with it. speculative_config: method: ngram num_speculative_tokens: 7 diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 52cd8f8e735..f3415068ad2 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -45,18 +45,11 @@ stages: # replicas below; the NPU block caps stage 0/1 KV explicitly. gpu_memory_utilization: 0.40 enforce_eager: false - # Stage-level prefix caching is NPU-only (see platforms.npu below); - # the top level stays false for duplex safety. Observed with this flag - # in the platform-shared block on CUDA: the first duplex case dies with - # an empty tts handoff (token_ids/hidden_states=None -> ValueError in - # minicpmo_4_5_omni_tts.preprocess) alongside a device-side assert out - # of indexSelectSmallIndex; the stage-1 replica is killed and every later - # case fails with "stage 1 has no live replica". Stage-0 n-gram - # speculative decoding stays off too: the - # omni prefix cache refuses the pair (stage_prefix_cache_config -> - # check_prefix_cache_token_accounting raises OmniPrefixCacheUnmatchError - # as soon as a stage has enable_prefix_caching and a speculative_config; - # vLLM itself allows the combination). + # Stage-level prefix caching is NPU-only (see platforms.npu below): in the + # platform-shared block it killed the stage-1 replica on CUDA (empty tts + # handoff, then an indexSelectSmallIndex device assert). Stage-0 n-gram + # speculation stays off either way -- the omni prefix cache refuses it next + # to enable_prefix_caching. tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -149,9 +142,7 @@ platforms: # not gated by a single stage. max_num_seqs: 8 max_num_batched_tokens: 8192 - # Stage-0 prefix caching was moved here from the platform-shared - # block: keep the PR's NPU behaviour unchanged while CUDA falls back - # to the top-level false (duplex safety, see the base stage comment). + # Moved here from the platform-shared block; see the base stage note. enable_prefix_caching: true # Explicit KV cap: the auto-computed value does not leave room for the # two stage-2 replicas on this card. @@ -166,21 +157,11 @@ platforms: # stages 0 and 2, so stage 1 kept the base value of 4. max_num_seqs: 8 max_num_batched_tokens: 8192 - # Multi-frame decode: the Talker produces 8 codec frames per scheduler - # step. vLLM learns that each request advances 8 tokens through this - # speculative_config; the n-gram proposer itself is never consulted - # because the NPU worker replays the decode graph with constant continue - # drafts, and async scheduling is rejected next to an n-gram proposer. - # stop_token_ids must stay the two-wide continue/stop id the head emits - # in this mode. Remove the block (or set it to null) to fall back to one - # frame per step. - # - # NPU-only on purpose: only - # vllm_omni.platforms.npu.worker.talker_multiframe intercepts the step - # and injects those drafts. Left in the base stage, the same block hands - # stage 1 to vLLM's real n-gram proposer and rejection sampler on CUDA, - # which asserts on the Talker's codec ids (device-side assert out of - # get_token_bin_counts_and_mask). + # Multi-frame decode: 8 codec frames per step. NPU-only -- only the NPU + # worker intercepts the step and injects the constant continue drafts, + # so the block stays in this section. stop_token_ids below is the + # two-wide continue/stop id that head emits; drop the block (or set it + # to null) for one frame per step, and drop that id with it. speculative_config: method: ngram num_speculative_tokens: 7 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index d252bc4712c..7eb9878720d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -22,33 +22,18 @@ def _talker_stop_token_ids() -> list[int]: - """Stage 1's default stop id: the one the multi-frame head can emit. + """Stage 1's default stop id, derived from the head this platform runs. - Stage 1's deploy config carries the ``speculative_config`` that arms the - multi-frame decode, and ``compute_logits`` then collapses the Talker's - vLLM-level head to the two-wide continue/stop row. The only sampleable ids - are therefore 0 and 1, and the codec EOS can never appear at this level (the - model samples codec ids inside the loop and forwards them to stage 2 - itself). - - Pinning this to the codec EOS regardless is what made every multi-frame - request run to ``max_tokens``: the model did emit the stop marker, - ``check_stop`` compared it against 6561, and the request kept its slot until - the context ran out. A deployment that drops the ``speculative_config`` from - stage 1 falls back to one frame per step, where the head is the full codec - vocabulary again and the codec EOS is the stop id; it overrides this per - stage through the usual ``stop_token_ids`` sampling parameter. + On NPU, stage 1's ``speculative_config`` collapses the vLLM-level head to + the two-wide continue/stop row: only 0/1 are sampleable, so the codec EOS + can never appear. Everywhere else that block stays out of the config (it is + confined to ``platforms.npu``), the head is the full codec vocabulary, and + the codec EOS is the only valid stop. """ from vllm_omni.platforms import current_omni_platform - # ``speculative_config`` only means "constant continue drafts" next to the - # NPU worker that intercepts the step - # (``vllm_omni.platforms.npu.worker.talker_multiframe``); on every other - # platform the same block hands the Talker to vLLM's real n-gram proposer and - # rejection sampler, which asserts on the codec ids. The block is therefore - # confined to ``platforms.npu`` in the deploy config, and the stop id follows: - # NPU keeps the two-wide continue/stop row, everyone else keeps the full codec - # head and stops on the codec EOS. + # The deploy config confines stage 1's speculative_config to platforms.npu, + # so a non-NPU stage 1 keeps the full codec head. if (current_omni_platform.device_name or "").lower() != "npu": return [_CODEC_EOS_TOKEN_ID] From 492659a3ce51968cb1d6daea0f3eb2c64205be2f Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 21:50:12 +0800 Subject: [PATCH 43/89] fix(config): record the K step through the NPU overlay Stage 1's ``speculative_config`` now lives under ``platforms.npu``, but ``_record_talker_frames`` still read the top-level stage list, so ``talker_frames_per_step()`` reported 1 for both deploy profiles. Its only reader is ``ascend_warmup_patch._kstep_armed``, which -- when the engine config is not readable yet -- decides from that answer whether a 910B restores the native rejection sampler, i.e. whether the armed multi-frame path runs the sampler its spec width needs. Merge the NPU overlay the same way ``_apply_platform_overrides`` does before recording. Verified: both deploy profiles report 8 frames per step again. Signed-off-by: y-null --- vllm_omni/config/stage_config.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 04a602da4a8..a061f86a9f2 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -797,13 +797,30 @@ def talker_frames_per_step() -> int: return _resolved_talker_frames -def _record_talker_frames(stages: list[StageDeployConfig]) -> None: - """Record stage 1's frame count (its n-gram ``num_speculative_tokens`` is K - 1).""" +def _record_talker_frames( + stages: list[StageDeployConfig], + platforms: dict[str, Any] | None = None, +) -> None: + """Record stage 1's frame count (its n-gram ``num_speculative_tokens`` is K - 1). + + The K block sits under ``platforms.npu`` -- the loop is NPU-only -- so the + NPU overlay is merged here. This runs at load time, before the per-platform + merge, and its reader (``ascend_warmup_patch._kstep_armed``) runs in a + spawned stage worker that may have no vllm_config to read it from. + """ global _resolved_talker_frames + npu_overlay: dict[int, dict[str, Any]] = {} + for stage_override in ((platforms or {}).get("npu") or {}).get("stages") or []: + if isinstance(stage_override, dict) and "stage_id" in stage_override: + npu_overlay[stage_override["stage_id"]] = stage_override for stage in stages: if stage.stage_id != 1: continue - spec = (stage.engine_extras or {}).get("speculative_config") or {} + spec = (stage.engine_extras or {}).get("speculative_config") + if spec is None: + # Platform overrides follow _extract_platform_overrides: every key + # but stage_id/devices/env is an engine or stage override. + spec = (npu_overlay.get(1) or {}).get("speculative_config") frames = 1 if isinstance(spec, dict) and spec.get("method") == "ngram": num_spec = spec.get("num_speculative_tokens", 0) or 0 @@ -822,7 +839,7 @@ def load_deploy_config(path: str | Path) -> DeployConfig: ) stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])] - _record_talker_frames(stages) + _record_talker_frames(stages, raw_dict.get("platforms")) model_runner = raw_dict.get("model_runner", "v1") if model_runner not in ("v1", "v2"): From 8207eb0aa80ccb7cd099e402637dd6b117fac8fa Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 17:24:30 +0800 Subject: [PATCH 44/89] fix(minicpmo_4_5): rebuild the codec sampler's device state per duplex condition _sample_audio_code builds a request's device state (step counter, EOS floor, max_tokens, history window) only when the request_id is absent, and _request_codec_device_inputs caches the min_tokens tensor the same way. A duplex resume reuses the same request id, so the next segment inherited the previous segment's step counter and its chunk budget: the EOS floor and the cadence boundaries were computed against the old segment. The preprocess path already rebuilds the host-side state per condition; drop the device-side caches there too (duplex only), so the next decode re-creates them from the fresh budget. The repetition-penalty window is carried through the state's recent codes, matching main's clean-slate-except-penalty semantics. Signed-off-by: y-null (cherry picked from commit d490ff703821654cc7c4f94e82c5ad37ca2bdf2a) --- .../models/minicpmo_4_5/minicpmo_4_5_omni_tts.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 534a5a370af..68149124865 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -745,6 +745,20 @@ def preprocess( request_states = {} self._request_audio_states = request_states request_states[request_id] = state + if native_duplex: + # A new duplex condition starts a fresh codec segment: drop the + # sampler's device-side caches for this request so the next + # decode rebuilds them from this condition's state (step, EOS + # floor, limits) instead of resuming the previous segment's. + # The repetition-penalty window is carried through the state + # and ``codes`` above, so only the counters and the cached + # input tensors reset. + codec_states = getattr(self, "_request_codec_device_states", None) + if isinstance(codec_states, dict): + codec_states.pop(request_id, None) + codec_inputs = getattr(self, "_request_codec_device_inputs", None) + if isinstance(codec_inputs, dict): + codec_inputs.pop(request_id, None) empty_codes = torch.empty(0, dtype=torch.long, device=embeds.device) return ( input_ids, From 2449c9dde24bb71a845ce0e5255bf4c28d833ab6 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 17:27:11 +0800 Subject: [PATCH 45/89] fix(minicpmo_4_5): wire the duplex turn-end EOS window into the K-step sampler The single-frame path masks codec EOS within 5 steps after each 25-frame cadence boundary while a turn-end chunk drains (mask_eos_rows), so the Talker finishes the text it still owes instead of stopping on the periodic cadence EOS. The K-step path samples through talker_codec_sample.prepare_codec_logits / greedy_codec_sample, whose only EOS criterion was step < min_tokens, so that window was left unwired and a turn-end chunk could end a few frames early. Thread an eos_window_masked flag through both sampling cores and set it from the same turn_end_drain + _turn_end_boundary_eos_masked(step) criterion the single-frame path uses. Signed-off-by: y-null (cherry picked from commit c63b276a9d5d4af5e41cc5a76ad08c87d0a58f66) --- .../test_talker_kstep_feedback.py | 2 +- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 21 ++++++---- .../minicpmo_4_5/talker_codec_sample.py | 41 +++++++++++++------ 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 7ae1509cbbc..13be179f6c4 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -57,7 +57,7 @@ def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): model.emb_code = nn.ModuleList([emb]) queue = iter(scripted_samples) - def _greedy(_hidden, _codes, _request_id, _step, _min_tokens, _max_tokens): + def _greedy(_hidden, _codes, _request_id, _step, _min_tokens, _max_tokens, _eos_window_masked=False): return torch.tensor(next(queue), dtype=torch.long) model._sample_audio_code_greedy = _greedy diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 68149124865..0358f478726 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -936,15 +936,13 @@ def make_omni_output( # _turn_end_boundary_eos_masked and state["turn_end_drain"]; # MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS is 4 chunks, so the window # lands on steps 25-29 / 50-54 / 75-79 / 100-103). The EOS masking - # moved into the sampling core (talker_codec_sample. - # prepare_codec_logits / greedy_codec_sample), whose criterion is - # state.step < min_tokens with min_tokens a per-request cached - # tensor, so that periodic window cannot be expressed per frame here - # and is left unwired. It only affects duplex turn-end; the - # offline/TTS path is unaffected. The constants, budget and - # turn_end_drain flag from main are kept in this file. + # lives in the sampling core (talker_codec_sample. + # prepare_codec_logits / greedy_codec_sample), so the periodic + # window rides along as the eos_window_masked flag below -- the same + # criterion the single-frame path's mask_eos_rows applies. min_tokens = _codec_int_param(state, "min_tokens", self._codec_min_tokens) max_tokens = _codec_int_param(state, "max_tokens", self._codec_max_tokens) + eos_window_masked = bool(state.get("turn_end_drain")) and _turn_end_boundary_eos_masked(step) if self._codec_temperature == 0.0: sampled = self._sample_audio_code_greedy( hidden[end - 1 : end], @@ -953,9 +951,12 @@ def make_omni_output( step, min_tokens, max_tokens, + eos_window_masked, ) else: - stochastic_result = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step) + stochastic_result = self._sample_audio_code( + hidden[end - 1 : end], codes, request_id, step, eos_window_masked + ) if isinstance(stochastic_result, TalkerCodecSampleResult): sampled = stochastic_result.sampled_token.reshape(()).to(torch.long) else: @@ -1088,6 +1089,7 @@ def _sample_audio_code( history: torch.Tensor, request_id: str, step: int, + eos_window_masked: bool = False, ) -> TalkerCodecSampleResult: device_states = getattr(self, "_request_codec_device_states", None) if device_states is None: @@ -1138,6 +1140,7 @@ def _sample_audio_code( top_k=self._codec_top_k, top_p=self._codec_top_p, min_tokens_to_keep=3, + eos_window_masked=eos_window_masked, ) probabilities = torch.softmax(logits, dim=-1) sampled = torch.multinomial( @@ -1161,6 +1164,7 @@ def _sample_audio_code_greedy( step: int, min_tokens: int, max_tokens: int, + eos_window_masked: bool = False, ) -> torch.Tensor: """Run the greedy codec boundary while keeping sampler state on device. @@ -1205,6 +1209,7 @@ def _sample_audio_code_greedy( penalty_tensor, top_k=self._codec_top_k, eos_token_id=self._codec_eos_id, + eos_window_masked=eos_window_masked, ) device_states[request_id] = result.state # The kernel ABI is int32, while the existing connector/audio-code diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index 79abb49f822..78ee40f0c47 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -65,8 +65,14 @@ def prepare_codec_logits( top_k: int, top_p: float, min_tokens_to_keep: int = MIN_TOKENS_TO_KEEP, + eos_window_masked: bool = False, ) -> torch.Tensor: - """Prepare and filter logits without changing native NPU RNG semantics.""" + """Prepare and filter logits without changing native NPU RNG semantics. + + ``eos_window_masked`` masks the codec EOS for this frame regardless of the + step counter -- the duplex turn-end drain window (a few frames after each + cadence boundary), which the ``step < min_tokens`` criterion cannot express. + """ logits = raw_logits.float() / temperature.reshape(1, 1) positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) valid = positions < state.history_len.reshape(1, 1) @@ -76,11 +82,14 @@ def prepare_codec_logits( alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) logits = torch.where(logits < 0, logits * alpha, logits / alpha) eos = logits[..., eos_token_id : eos_token_id + 1] - eos = torch.where( - (state.step < min_tokens).reshape(1, 1), - torch.full_like(eos, float("-inf")), - eos, - ) + if eos_window_masked: + eos = torch.full_like(eos, float("-inf")) + else: + eos = torch.where( + (state.step < min_tokens).reshape(1, 1), + torch.full_like(eos, float("-inf")), + eos, + ) logits = torch.cat([logits[..., :eos_token_id], eos, logits[..., eos_token_id + 1 :]], dim=-1) if 0.0 < float(top_p) < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=False, dim=-1) @@ -142,8 +151,13 @@ def greedy_codec_sample( *, top_k: int, eos_token_id: int, + eos_window_masked: bool = False, ) -> TalkerCodecSampleResult: - """Greedy codec sample in graph-capturable tensor ops.""" + """Greedy codec sample in graph-capturable tensor ops. + + ``eos_window_masked`` masks the codec EOS for this frame regardless of the + step counter -- see prepare_codec_logits. + """ positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) valid = positions < state.history_len.reshape(1, 1) safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) @@ -153,11 +167,14 @@ def greedy_codec_sample( penalized = torch.where(raw_logits < 0, raw_logits * alpha, raw_logits / alpha) eos = penalized[..., eos_token_id : eos_token_id + 1] - eos = torch.where( - (state.step < min_tokens).reshape(1, 1), - torch.full_like(eos, float("-inf")), - eos, - ) + if eos_window_masked: + eos = torch.full_like(eos, float("-inf")) + else: + eos = torch.where( + (state.step < min_tokens).reshape(1, 1), + torch.full_like(eos, float("-inf")), + eos, + ) penalized = torch.cat([penalized[..., :eos_token_id], eos, penalized[..., eos_token_id + 1 :]], dim=-1) if top_k > 0: keep = min(VOCAB_SIZE, max(int(top_k), MIN_TOKENS_TO_KEEP)) From 209558f57e3fd401c1df653601e69102ab9a5f4a Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 17:31:12 +0800 Subject: [PATCH 46/89] fix(minicpmo_4_5): feed the real codec history to streaming prompt recompute Under K-step the scheduler's token stream records the 0/1 control rows the two-wide head emits, not the codec ids the model sampled. The streaming prompt recompute path reads its confirmed-token slice from that stream and embeds it with emb_code, so after a rollover the rebuilt Talker condition was made of continue/stop marker embeddings: ids 0 and 1 have valid embeddings and pass the < EOS range check, so nothing failed loudly. Record the confirmed frame per forwarded position in a per-request codec history and swap the values in at recompute time -- same length, one row per confirmed frame, and the slice always covers the tail of that history. The history is dropped with the rest of the request's caches in _flush_deferred_cleanup. Signed-off-by: y-null (cherry picked from commit ea3eee3323dfc0a39041d14e1bfda54be1f07d8e) --- .../test_talker_kstep_feedback.py | 5 ++ .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 55 +++++++++++++++++-- 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 13be179f6c4..446da5079ec 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -49,6 +49,7 @@ def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): model._num_audio_tokens = _NUM_AUDIO_TOKENS model._codec_eos_id = _EOS_ID model._request_audio_states = {} + model._request_codec_history = {} emb = nn.Embedding(_NUM_AUDIO_TOKENS, 4) with torch.no_grad(): emb.weight.zero_() @@ -94,6 +95,10 @@ def test_kstep_frame_sample_is_recorded_for_next_frame(): assert state["last_code"] == 43 assert state["finished"] is True assert out2.multimodal_outputs["codes"]["audio"][0].numel() == 0 + # Only confirmed frames are recorded: the terminating frame carries no + # codec id for the next step to embed, so streaming prompt recompute must + # not see it in the history either. + assert model._request_codec_history["r1"] == [42, 43] def test_legacy_single_frame_path_does_not_touch_last_code(): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 0358f478726..a37352a20a8 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -339,6 +339,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self._request_audio_states: dict[str, dict[str, Any]] = {} self._request_codec_device_states: dict[str, TalkerCodecDeviceState] = {} self._request_codec_device_inputs: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + # Confirmed codec frames this model sampled, one entry per forwarded + # frame. Under K-step the scheduler's token stream carries the 0/1 + # control rows the two-wide head emits instead, so streaming prompt + # recompute swaps its confirmed-token slice for these real ids. + self._request_codec_history: dict[str, list[int]] = {} self._init_native_talker(prefix) @@ -568,6 +573,23 @@ def _build_streaming_recompute_embeddings( ids = info_dict.get("ids") previous_codes = ids.get("streaming_prompt_previous_codes") if isinstance(ids, Mapping) else None + if self._k_step_frames > 0 and previous_codes is not None: + # K-step: the scheduler's confirmed slice carries the 0/1 control + # rows the two-wide head emits, not codec ids -- feeding them to + # emb_code would silently rebuild the prompt from continue/stop + # markers (both have valid embeddings and pass the range check + # below). The model recorded the real frame it confirmed at each + # of those positions, one row per forwarded frame, so swap the + # values in. The slice only ever covers the frames confirmed + # since the current prompt, i.e. the tail of that history. + codec_history = getattr(self, "_request_codec_history", {}).get(request_id) + if codec_history is None or len(codec_history) < len(previous_codes): + raise ValueError( + "K-step streaming prompt recompute is missing confirmed codec frames: " + f"history={len(codec_history) if codec_history is not None else 0}, " + f"confirmed={len(previous_codes)}" + ) + previous_codes = codec_history[-len(previous_codes) :] if isinstance(previous_codes, torch.Tensor): code_ids = previous_codes.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1) elif isinstance(previous_codes, (list, tuple)): @@ -979,8 +1001,24 @@ def make_omni_output( # vLLM schedules for the next step carry placeholder # continue ids. Record this frame's sample so the decode # preprocess embeds the real previous frame's codec id - # (the k_last branch) instead of the placeholder. + # (the k_last branch) instead of the placeholder, and keep + # the confirmed-frame history that streaming prompt + # recompute rebuilds its window from -- the scheduler's + # token slice holds the placeholder ids. state["last_code"] = sampled_id + histories = getattr(self, "_request_codec_history", None) + if not isinstance(histories, dict): + # Instances built without __init__ (CPU tests, + # subclasses) still need a place for the record -- the + # readers at the recompute and cleanup sites already + # tolerate a missing dict the same way. + histories = {} + self._request_codec_history = histories + codec_history = histories.get(request_id) + if codec_history is None: + codec_history = [] + histories[request_id] = codec_history + codec_history.append(sampled_id) else: delta = empty_delta state["codes"] = codes @@ -1278,11 +1316,13 @@ def _flush_deferred_cleanup(self) -> None: ``_request_codec_device_states`` carries the step counter and the penalty window the EOS mask is derived from, and ``_request_codec_device_inputs`` carries the ``min_tokens`` copied at - first use. Leaving them behind means a recycled request id resumes with - another request's counters -- the EOS mask then holds or lifts at the - wrong frame, and ``_request_generators`` leaks a generator per request - on top of that. All three key off the same request id, so they are - dropped together with the rest. + first use. ``_request_codec_history`` carries the confirmed K-step + frames streaming prompt recompute reads. Leaving them behind means a + recycled request id resumes with another request's counters -- the EOS + mask then holds or lifts at the wrong frame, and + ``_request_generators`` leaks a generator per request on top of that. + All of them key off the same request id, so they are dropped together + with the rest. """ request_audio_states = getattr(self, "_request_audio_states", {}) request_condition_states = getattr(self, "_request_condition_states", {}) @@ -1290,6 +1330,7 @@ def _flush_deferred_cleanup(self) -> None: penalty_freqs = getattr(self, "_penalty_freqs_dev", None) codec_device_states = getattr(self, "_request_codec_device_states", None) codec_device_inputs = getattr(self, "_request_codec_device_inputs", None) + codec_histories = getattr(self, "_request_codec_history", None) request_generators = getattr(self, "_request_generators", None) for request_id in self._deferred_cleanup_ids: request_audio_states.pop(request_id, None) @@ -1302,6 +1343,8 @@ def _flush_deferred_cleanup(self) -> None: codec_device_states.pop(request_id, None) if isinstance(codec_device_inputs, dict): codec_device_inputs.pop(request_id, None) + if isinstance(codec_histories, dict): + codec_histories.pop(request_id, None) if isinstance(request_generators, dict): request_generators.pop(request_id, None) self._deferred_cleanup_ids.clear() From d95250a62f61371af3a1325b7a5d94726117d457 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 22:07:39 +0800 Subject: [PATCH 47/89] fix(minicpmo_4_5): derive the Talker stop id from the armed head, not the platform The deploy config that arms the multi-frame head already pins that head's stop marker (1) through stage 1's ``default_sampling_params``, and ``merge_sampling_constraints`` unions it with whatever the pipeline supplies. So the pipeline's only job is to supply the stop id of the *other* head -- the codec EOS of a stage that samples the full codec vocabulary. Keyed on the platform instead, the NPU branch returned the marker even for an NPU deployment that drops stage 1's ``speculative_config`` (one frame per step again): its head is the full codec vocabulary there, so 1 is an ordinary audio id and the request ends mid-utterance -- the same defect this default was introduced to remove, just narrower, and the one case the review called out. No shipped configuration changes: armed NPU configs still merge to [1, 6561] (1 ends the request; 6561 is unsampleable in the collapsed head) and non-NPU configs still stop on 6561 alone. Only NPU-without-the-block moves from a wrong 1 to the correct 6561. Dropping the platform probe also drops the import of the NPU worker module from the pipeline, so resolving the pipeline no longer drags in a platform worker. Tests: the topology test and the stop-id test now pin the pipeline default to the codec EOS on every platform, and the deploy-scoping test pins both merged lists ([1, 6561] armed, [6561] without the block). Verified: 69 passed across test_pipeline.py, test_talker_kstep_feedback.py and test_config_factory.py. Signed-off-by: y-null --- tests/config/test_config_factory.py | 9 +++++- .../models/minicpmo_4_5/test_pipeline.py | 12 ++++---- .../test_talker_kstep_feedback.py | 11 +++----- .../models/minicpmo_4_5/pipeline.py | 28 ++++++++----------- 4 files changed, 30 insertions(+), 30 deletions(-) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 299a7d54957..84af75fedc6 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1384,12 +1384,15 @@ def test_minicpmo_talker_multi_frame_is_npu_scoped(self, platform: str): platform=platform, ) # Check the deploy-side stop id before the merge: afterwards it is - # unioned with the pipeline's platform-aware constraint. + # unioned with the pipeline's codec-EOS default. deploy_stage1 = next(stage for stage in deploy.stages if stage.stage_id == 1) stages = merge_pipeline_deploy(resolve_pipeline_config("minicpmo_4_5"), deploy) if platform == "npu": assert deploy_stage1.default_sampling_params["stop_token_ids"] == [1] + # The armed head ends on 1; the union keeps the codec EOS inert + # (the collapsed head cannot sample it). + assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1, 6561] assert stages[0].yaml_engine_args["enable_prefix_caching"] is True assert stages[1].yaml_engine_args["speculative_config"] == { "method": "ngram", @@ -1401,6 +1404,10 @@ def test_minicpmo_talker_multi_frame_is_npu_scoped(self, platform: str): assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [8, 8, 8] else: assert "stop_token_ids" not in (deploy_stage1.default_sampling_params or {}) + # Without the block the head is the full codec vocabulary, so the + # codec EOS is the only stop id -- the marker would end a request + # on an ordinary codec id. + assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] # Either absent (falls back to the top-level false) or explicitly # false; the regression this guards against is true. assert stages[0].yaml_engine_args.get("enable_prefix_caching") is not True diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 15fbb54eaa2..d8084fb2f0d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -110,14 +110,14 @@ def test_talker_stage(self, pipeline: PipelineConfig) -> None: assert talker.engine_output_type == "latent" # scope KV cache / mrope sizing to talker sub-config assert talker.hf_config_name == "tts_config" - # The multi-frame head collapses the vLLM-level head to the two-wide - # continue/stop row, whose only stop id is 1. It is an NPU-worker - # feature (that is where the continue drafts are injected), so on any - # other platform stage 1 keeps its full codec head and the codec EOS. - expected_stop = [1] if (current_omni_platform.device_name or "").lower() == "npu" else [_CODEC_EOS_TOKEN_ID] + # The pipeline default is the codec EOS of the one-frame head. The + # deploy config that arms the multi-frame row (stage 1's + # speculative_config, under platforms.npu) adds that row's stop marker + # (1) through its default_sampling_params; see + # test_minicpmo_talker_multi_frame_is_npu_scoped for the merged list. assert talker.sampling_constraints == { "detokenize": False, - "stop_token_ids": expected_stop, + "stop_token_ids": [_CODEC_EOS_TOKEN_ID], } assert talker.custom_process_next_stage_input_func == ( "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload" diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 446da5079ec..93a9f2a306e 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -251,16 +251,13 @@ def test_talker_stop_token_ids_match_the_multi_frame_head(): NPU worker, and the codec EOS on every platform without that worker. """ from vllm_omni.model_executor.models.minicpmo_4_5 import pipeline as mcp_pipeline - from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms.npu.worker import talker_multiframe assert talker_multiframe.STOP_TOKEN_ID == 1 - if (current_omni_platform.device_name or "").lower() == "npu": - assert mcp_pipeline._talker_stop_token_ids() == [talker_multiframe.STOP_TOKEN_ID] - else: - # No NPU worker means no constant-draft injection, so the K-frame loop - # is off and the two-wide continue/stop row never runs. - assert mcp_pipeline._talker_stop_token_ids() == [mcp_pipeline._CODEC_EOS_TOKEN_ID] + # The pipeline default cannot depend on the platform: the block that + # collapses the head is a deploy-config decision, so the marker is added + # there (test_minicpmo_talker_multi_frame_is_npu_scoped pins the merge). + assert mcp_pipeline._talker_stop_token_ids() == [mcp_pipeline._CODEC_EOS_TOKEN_ID] def test_multiframe_gate_matrix(): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 7eb9878720d..c91ce39f530 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -22,24 +22,20 @@ def _talker_stop_token_ids() -> list[int]: - """Stage 1's default stop id, derived from the head this platform runs. + """Stage 1's default stop id: the codec EOS of the one-frame head. - On NPU, stage 1's ``speculative_config`` collapses the vLLM-level head to - the two-wide continue/stop row: only 0/1 are sampleable, so the codec EOS - can never appear. Everywhere else that block stays out of the config (it is - confined to ``platforms.npu``), the head is the full codec vocabulary, and - the codec EOS is the only valid stop. - """ - from vllm_omni.platforms import current_omni_platform - - # The deploy config confines stage 1's speculative_config to platforms.npu, - # so a non-NPU stage 1 keeps the full codec head. - if (current_omni_platform.device_name or "").lower() != "npu": - return [_CODEC_EOS_TOKEN_ID] + The deploy config that arms the multi-frame decode (stage 1's + ``speculative_config``, under ``platforms.npu``) collapses the vLLM-level + head to the two-wide continue/stop row, where the codec EOS can never + appear; that same config adds the stop marker (1) through its + ``default_sampling_params``, and ``merge_sampling_constraints`` unions both + lists. The marker therefore travels with the block that needs it, and this + default covers every deployment without one. - from vllm_omni.platforms.npu.worker.talker_multiframe import STOP_TOKEN_ID - - return [STOP_TOKEN_ID] + A deployment that keeps the marker here instead ends every one-frame + request on the first ordinary codec id 1. + """ + return [_CODEC_EOS_TOKEN_ID] MINICPMO_4_5_PIPELINE = PipelineConfig( From d96289ed656be70c91303be4c2d073b84c05f5b7 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 22 Sep 2026 23:17:58 +0800 Subject: [PATCH 48/89] docs(deploy): drop the unverified PIECEWISE number from the cudagraph note The -32.6% figure came from the source line's A2 runs and was never measured here. Our own PIECEWISE A/B cannot produce a number in this deployment either: with the K step armed, stage 1 dies on the first request (RuntimeError: CUDA graph capturing detected at an inappropriate time, out of talker_multiframe.run), so the note keeps the mechanism only. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 5 +++-- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 162ef0b9f70..12d7e722e76 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -139,8 +139,9 @@ platforms: enable_prefix_caching: true compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame - # expansion; PIECEWISE silently drops both. Measured on 910B3: - # -32.6% RTF versus PIECEWISE. + # expansion: with PIECEWISE the K-step replay trips vLLM's capture-time + # guard (CUDA graph capturing detected at an inappropriate time) and + # stage 1 dies. cudagraph_mode: FULL_DECODE_ONLY - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index f3415068ad2..9ca381d5cf6 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -149,8 +149,9 @@ platforms: kv_cache_memory_bytes: 6442450944 compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame - # expansion; PIECEWISE silently drops both. Measured on 910B3: - # -32.6% RTF versus PIECEWISE. + # expansion: with PIECEWISE the K-step replay trips vLLM's capture-time + # guard (CUDA graph capturing detected at an inappropriate time) and + # stage 1 dies. cudagraph_mode: FULL_DECODE_ONLY - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered From 83601d133d2332fa03cd65429f71122c60db6828 Mon Sep 17 00:00:00 2001 From: y-null Date: Wed, 23 Sep 2026 11:05:26 +0800 Subject: [PATCH 49/89] fix(minicpmo_4_5): keep static-shape decode buckets FULL-decode-only The runtime bucket is part of the builder's dispatch key and is only answerable where decode runs as one FULL graph. Under PIECEWISE/eager setting it sent the wrapper after never-captured per-capacity entries and the runtime capture attempt died in vLLM's capturing-window check. Empty buckets now leave every consumer on its stock path, so PIECEWISE deploys run again (K included, verified on 910B1); FULL_DECODE_ONLY stays the default because K's RTF win is FULL-only. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 7 +++++ vllm_omni/deploy/minicpmo_4_5_prof.yaml | 7 +++++ .../npu/attention/static_shape_backend.py | 26 ++++++++++++++++++- 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 12d7e722e76..180fb8711d5 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -139,9 +139,16 @@ platforms: enable_prefix_caching: true compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame +<<<<<<< HEAD # expansion: with PIECEWISE the K-step replay trips vLLM's capture-time # guard (CUDA graph capturing detected at an inappropriate time) and # stage 1 dies. +======= + # expansion. The static-shape builder gates its runtime buckets on + # the decode graph mode, so PIECEWISE also runs (K included, + # verified on 910B1) but K's RTF win is FULL-only there. Measured + # on 910B3: -32.6% RTF versus PIECEWISE. +>>>>>>> a125d09d1 (fix(minicpmo_4_5): keep static-shape decode buckets FULL-decode-only) cudagraph_mode: FULL_DECODE_ONLY - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 9ca381d5cf6..1ae390599d9 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -149,9 +149,16 @@ platforms: kv_cache_memory_bytes: 6442450944 compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame +<<<<<<< HEAD # expansion: with PIECEWISE the K-step replay trips vLLM's capture-time # guard (CUDA graph capturing detected at an inappropriate time) and # stage 1 dies. +======= + # expansion. The static-shape builder gates its runtime buckets on + # the decode graph mode, so PIECEWISE also runs (K included, + # verified on 910B1) but K's RTF win is FULL-only there. Measured + # on 910B3: -32.6% RTF versus PIECEWISE. +>>>>>>> a125d09d1 (fix(minicpmo_4_5): keep static-shape decode buckets FULL-decode-only) cudagraph_mode: FULL_DECODE_ONLY - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered diff --git a/vllm_omni/platforms/npu/attention/static_shape_backend.py b/vllm_omni/platforms/npu/attention/static_shape_backend.py index f1f23f73947..e8ccb50356d 100644 --- a/vllm_omni/platforms/npu/attention/static_shape_backend.py +++ b/vllm_omni/platforms/npu/attention/static_shape_backend.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 """Ascend attention backend that captures a decode step with constant host args. ``FULL_DECODE_ONLY`` captures a whole decode step, but @@ -31,6 +31,7 @@ import torch import torch_npu +from vllm.config import CUDAGraphMode from vllm.logger import init_logger from vllm_ascend.ascend_forward_context import _EXTRA_CTX from vllm_ascend.attention.attention_v1 import ( @@ -62,6 +63,29 @@ def __init__(self, *args, **kwargs) -> None: self._omni_max_q_len = 1 if spec is not None and getattr(spec, "num_speculative_tokens", 0): self._omni_max_q_len = 1 + int(spec.num_speculative_tokens) + # The runtime bucket is part of the dispatch key: it is only answerable + # where decode steps run as one FULL graph (FULL / FULL_DECODE_ONLY / + # FULL_AND_PIECEWISE), because those are the only modes that capture + # per-capacity decode graphs to look up. Under PIECEWISE or eager, + # setting it sends the wrapper after a per-capacity entry dict that was + # never captured, and the runtime capture attempt crashes vLLM's + # capturing-window check ("Cannot capture CUDA graphs in an + # inappropriate time"). With empty buckets build() leaves the bucket + # unset and every consumer (wrapper __call__, static-shape attention, + # prefill hook) takes its stock path -- the behaviour this stage would + # have had without the static-shape seam at all. + mode = self.vllm_config.compilation_config.cudagraph_mode + if self._omni_buckets and mode not in ( + CUDAGraphMode.FULL, + CUDAGraphMode.FULL_DECODE_ONLY, + CUDAGraphMode.FULL_AND_PIECEWISE, + ): + logger.info( + "[minicpmo] static-shape decode graphs need FULL decode graphs " + "(cudagraph_mode=%s); keeping the stock attention path", + mode, + ) + self._omni_buckets = () def build(self, *args, **kwargs): attn_metadata = super().build(*args, **kwargs) From 620e44b885f8b97cd794d39841a2600cba346b99 Mon Sep 17 00:00:00 2001 From: y-null Date: Wed, 23 Sep 2026 03:15:20 +0000 Subject: [PATCH 50/89] fix(deploy): drop the leftover conflict markers from the prof yaml note Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 1ae390599d9..20aca4988b2 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -149,16 +149,11 @@ platforms: kv_cache_memory_bytes: 6442450944 compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame -<<<<<<< HEAD - # expansion: with PIECEWISE the K-step replay trips vLLM's capture-time - # guard (CUDA graph capturing detected at an inappropriate time) and - # stage 1 dies. -======= # expansion. The static-shape builder gates its runtime buckets on # the decode graph mode, so PIECEWISE also runs (K included, # verified on 910B1) but K's RTF win is FULL-only there. Measured - # on 910B3: -32.6% RTF versus PIECEWISE. ->>>>>>> a125d09d1 (fix(minicpmo_4_5): keep static-shape decode buckets FULL-decode-only) + # on 910B1: PIECEWISE without K is +75% RTF at 32/1 versus + # FULL_DECODE_ONLY with K=8. cudagraph_mode: FULL_DECODE_ONLY - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered From 6a5bd1c2b6c134faf158d9285d7766acf0024905 Mon Sep 17 00:00:00 2001 From: y-null Date: Wed, 23 Sep 2026 03:19:09 +0000 Subject: [PATCH 51/89] fix(deploy): drop the leftover conflict markers from the main yaml note Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 180fb8711d5..d5bde22b660 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -139,16 +139,11 @@ platforms: enable_prefix_caching: true compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame -<<<<<<< HEAD - # expansion: with PIECEWISE the K-step replay trips vLLM's capture-time - # guard (CUDA graph capturing detected at an inappropriate time) and - # stage 1 dies. -======= # expansion. The static-shape builder gates its runtime buckets on # the decode graph mode, so PIECEWISE also runs (K included, # verified on 910B1) but K's RTF win is FULL-only there. Measured - # on 910B3: -32.6% RTF versus PIECEWISE. ->>>>>>> a125d09d1 (fix(minicpmo_4_5): keep static-shape decode buckets FULL-decode-only) + # on 910B1: PIECEWISE without K is +75% RTF at 32/1 versus + # FULL_DECODE_ONLY with K=8. cudagraph_mode: FULL_DECODE_ONLY - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered From 1883b4fb3ef15ae2ae5ce9880c1fc83a1d7a6e6b Mon Sep 17 00:00:00 2001 From: y-null Date: Thu, 24 Sep 2026 03:53:07 +0000 Subject: [PATCH 52/89] fix(minicpmo_4_5): forward sample eligibility and honor deploy/request codec params Two review fixes for PR #7929: - Forward requires_request_sample_eligibility from the outer MiniCPMO45OmniForConditionalGeneration wrapper to the inner Talker. The runner-side getattr silently read False, so request_sample_eligible was never passed and mid-prefill chunks sampled codec tokens, making audio output depend on prefill chunking. - Source the K-step codec sampler from deploy config instead of the checkpoint: resolve_codec_sampling_params now inserts default_sampling_params (the stage defaults the single-frame path uses) between the yaml codec block and tts_config, and the runner passes each request resolved SamplingParams so temperature/top_k/top_p/repetition_penalty/seed are pinned per request. Before, tts_config top_k=100/top_p=0.8/max_tokens=2048 won while the deploy yaml asked for 25/0.85/4096. Tests: tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py 14 passed, incl. 4 new locking the two behaviors. Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 155 ++++++++++++++++++ .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 15 ++ .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 92 +++++++++-- vllm_omni/worker/gpu_model_runner.py | 9 + 4 files changed, 259 insertions(+), 12 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 93a9f2a306e..f986a1ac268 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """No-device regression tests for the MiniCPM-o Talker K-step feedback chain. Layers pinned down here: @@ -379,3 +382,155 @@ def test_kstep_min_tokens_neutralization_clears_the_censor_list(): # An empty list and odd inputs must not raise either. talker_multiframe.neutralize_kstep_min_tokens(SimpleNamespace(non_argmax_invariant=[])) talker_multiframe.neutralize_kstep_min_tokens(None) + + +def test_incomplete_prefill_chunk_skips_sampling(): + """The runner's eligibility flag must gate the K-step sampling branch. + + A request whose prompt spans several prefill chunks reports + ``request_sample_eligible=False`` for the incomplete chunks; sampling + there would advance codec history and RNG state, making the generated + audio depend on how the prompt happened to be chunked (review on PR + #7929). The flag only reaches the Talker while the outer wrapper + forwards ``requires_request_sample_eligibility`` -- without that forward + the runner never sends the flag and this branch degrades to the + ``[True] * len(infos)`` fallback. + """ + model = _make_talker(k_step_frames=8, scripted_samples=[42]) + state = {"step": 0, "codes": torch.tensor([10, 11])} + model._request_audio_states["r1"] = state + + out = model.make_omni_output( + torch.randn(1, 8), + model_intermediate_buffer=[{"request_id": "r1"}], + request_token_spans=[(0, 1)], + request_sample_eligible=[False], + ) + # No codec frame, no state advance, no history touch. + assert out.multimodal_outputs["codes"]["audio"][0].numel() == 0 + assert state["step"] == 0 + assert "last_code" not in state + assert model._request_codec_history.get("r1", []) == [] + + # Eligible again (the chunking completed): sampling resumes. + out2 = _frame_call(model, torch.randn(1, 8)) + assert state["last_code"] == 42 + assert out2.multimodal_outputs["codes"]["audio"][0].reshape(-1).tolist() == [42] + + +def test_outer_wrapper_forwards_sampling_eligibility_flag(): + """The runner only sees the wrapper, so the flag must resolve through it. + + gpu/npu runners arm the ``request_sample_eligible`` transmission with + ``getattr(self.model, "requires_request_sample_eligibility", False)``, + and ``self.model`` is the stage's registered architecture -- the outer + ``MiniCPMO45OmniForConditionalGeneration`` wrapper, never the inner + Talker that declares the flag. + """ + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + wrapper = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + nn.Module.__init__(wrapper) + # A wrapper without a Talker (stage 0 LLM) must not arm the transmission. + assert wrapper.requires_request_sample_eligibility is False + + wrapper.talker = _make_talker(k_step_frames=8, scripted_samples=[]) + assert wrapper.requires_request_sample_eligibility is True + + +def test_request_sampling_params_pin_codec_knobs(): + """Per-request SamplingParams override the statically resolved knobs. + + Single-frame contract: temperature/seed/top-k/top-p/penalty overrides + steer the codec stream exactly like they steer the vLLM sampler in the + single-frame path. A request pinning temperature to 0 must also take + the deterministic boundary sampler. + """ + model = _make_talker(k_step_frames=8, scripted_samples=[42]) + # Deployment resolved a warm stochastic profile; the request pins its own. + model._codec_temperature = 0.8 + model._codec_top_k = 100 + model._codec_top_p = 0.8 + model._codec_repetition_penalty = 1.05 + model._codec_seed = 42 + state = {"step": 0, "codes": torch.tensor([10, 11])} + model._request_audio_states["r1"] = state + + model.make_omni_output( + torch.randn(1, 8), + model_intermediate_buffer=[{"request_id": "r1"}], + request_token_spans=[(0, 1)], + request_sample_eligible=[True], + request_sampling_params=[ + SimpleNamespace(temperature=0.0, top_k=25, top_p=0.85, repetition_penalty=1.0, seed=7) + ], + ) + # temperature 0 -> the deterministic boundary sampler was taken. + assert state["last_code"] == 42 + # The knobs are pinned for this request's samplers. + assert state["codec_temperature"] == 0.0 + assert state["codec_top_k"] == 25 + assert state["codec_top_p"] == 0.85 + assert state["codec_repetition_penalty"] == 1.0 + gen = model._request_generator("r1", torch.device("cpu")) + assert gen.initial_seed() == 7 + + # A request without SamplingParams (dummy runs, CPU tests) keeps the + # statically resolved deployment profile. + model2 = _make_talker(k_step_frames=8, scripted_samples=[43]) + model2._codec_seed = 42 + model2._request_audio_states["r2"] = {"step": 0, "codes": torch.tensor([10, 11])} + model2.make_omni_output( + torch.randn(1, 8), + model_intermediate_buffer=[{"request_id": "r2"}], + request_token_spans=[(0, 1)], + request_sample_eligible=[True], + ) + gen2 = model2._request_generator("r2", torch.device("cpu")) + assert gen2.initial_seed() == 42 + + +def test_default_sampling_params_feed_codec_resolution(): + """``default_sampling_params`` sits between the YAML block and tts_config. + + It is the source the single-frame path's SamplingParams are built from, + so a deployment tuning stage 1 through it must control the K-step codec + sampler too (review on PR #7929). The explicit YAML block still wins key + by key, and untouched keys fall through to the checkpoint config. + """ + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + resolve_codec_sampling_params, + ) + + tts_config = SimpleNamespace( + seed=42, + temperature=0.8, + top_k=100, + top_p=0.8, + repetition_penalty=1.05, + min_new_tokens=50, + max_new_tokens=2048, + ) + + # No YAML block: the stage defaults drive the knobs the deployment set, + # untouched keys keep the checkpoint values. + resolved = resolve_codec_sampling_params( + None, tts_config, deploy_defaults={"top_k": 25, "top_p": 0.85, "max_tokens": 4096} + ) + assert resolved["top_k"] == 25 + assert resolved["top_p"] == 0.85 + assert resolved["max_tokens"] == 4096 + assert resolved["temperature"] == 0.8 + assert resolved["min_tokens"] == 50 + + # The explicit YAML block still wins key by key. + resolved2 = resolve_codec_sampling_params({"top_k": 9}, tts_config, deploy_defaults={"top_k": 25}) + assert resolved2["top_k"] == 9 + + # No stage defaults either: back to the old chain (tts_config, then + # module fallbacks). + resolved3 = resolve_codec_sampling_params(None, tts_config) + assert resolved3["top_k"] == 100 + assert resolved3["seed"] == 42 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index fe3f9565f74..fb7d50275d9 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -634,6 +634,21 @@ def make_omni_output(self, model_outputs, **kwargs): return model_outputs return self.talker.make_omni_output(model_outputs, **kwargs) + @property + def requires_request_sample_eligibility(self) -> bool: + """Forward the Talker's sampling-eligibility contract to the runner. + + The runner only supplies ``request_sample_eligible`` when the model it + sees declares the flag, and it always sees this wrapper -- never the + inner Talker (the same resolution rule as talker_multiframe.applies + below). Without the forward, the Talker's K-step branch falls back to + treating every request as eligible, so an incomplete prefill chunk + advances codec history and RNG state and the generated audio starts + depending on prefill chunking. + """ + talker = getattr(self, "talker", None) + return talker is not None and bool(getattr(talker, "requires_request_sample_eligibility", False)) + # Runner-side forwards for multi-frame decode. The baseline declares these # four methods on this outer wrapper class; the runner only ever sees the # wrapper (talker_multiframe.applies / is_multi_token_decode / diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index a37352a20a8..c32067f297a 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -193,8 +193,11 @@ def _apply_batched_repetition_penalty( def resolve_codec_sampling_params( yaml_params: Mapping[str, Any] | None, tts_config: Any | None = None, + *, + deploy_defaults: Mapping[str, Any] | None = None, ) -> dict[str, Any]: - """Resolve Talker codec knobs: deploy YAML, then checkpoint, then defaults. + """Resolve Talker codec knobs: deploy YAML, then stage defaults, then + checkpoint, then defaults. A deploy YAML is not required to carry a ``codec_sampling_params`` block -- the stock MiniCPM-o configs do not -- so requiring one would turn a @@ -204,6 +207,10 @@ def resolve_codec_sampling_params( unchanged. A YAML block still wins key by key, so a deployment can override any knob. + ``deploy_defaults`` is the stage's ``default_sampling_params`` -- the same + source the single-frame path's SamplingParams are built from -- so a + deployment tuning stage 1 through it keeps controlling the K-step codec + sampler instead of silently drifting to checkpoint/module values. """ provided = yaml_params if isinstance(yaml_params, Mapping) else {} resolved: dict[str, Any] = {} @@ -211,6 +218,9 @@ def resolve_codec_sampling_params( for key, attribute, fallback, cast in _CODEC_SAMPLING_SOURCES: value = provided.get(key) source = "yaml" + if value is None and isinstance(deploy_defaults, Mapping): + value = deploy_defaults.get(key) + source = "default_sampling_params" if value is None: value = getattr(tts_config, attribute, None) if tts_config is not None else None source = "tts_config" @@ -226,6 +236,12 @@ def resolve_codec_sampling_params( return resolved +def _codec_float_param(state: Any, key: str, fallback: float) -> float: + """Float counterpart of _codec_int_param; same None-means-unset rule.""" + value = state.get(key) if isinstance(state, Mapping) else None + return float(fallback if value is None else value) + + def _codec_int_param(state: Any, key: str, fallback: int) -> int: """Read an integer codec knob from a request state, None meaning unset. @@ -295,7 +311,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # module fallbacks. The multi-frame path reads the resolved values # straight off the model. yaml_codec = getattr(getattr(vllm_config, "model_config", None), "codec_sampling_params", None) - resolved = resolve_codec_sampling_params(yaml_codec, tts_config) + # The stage's default_sampling_params sits between the YAML block + # and the checkpoint: it is the source the single-frame path's + # SamplingParams are built from, so tuning stage 1 through it must + # control the K-step codec sampler too. + deploy_defaults = getattr(getattr(vllm_config, "model_config", None), "default_sampling_params", None) + resolved = resolve_codec_sampling_params(yaml_codec, tts_config, deploy_defaults=deploy_defaults) self._codec_seed = resolved["seed"] self._codec_temperature = resolved["temperature"] self._codec_top_k = resolved["top_k"] @@ -856,6 +877,12 @@ def make_omni_output( raise RuntimeError( f"MiniCPM-o continuous Talker received {len(sample_eligible)} sampling flags for {len(infos)} requests" ) + request_sampling_params = kwargs.get("request_sampling_params") + if request_sampling_params is not None and len(request_sampling_params) != len(infos): + raise RuntimeError( + f"MiniCPM-o continuous Talker received {len(request_sampling_params)} " + f"sampling params for {len(infos)} requests" + ) emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos) stop_rows: list[torch.Tensor] = [] @@ -932,6 +959,10 @@ def make_omni_output( if not isinstance(state, dict): state = dict(info.get("audio_state", {}) or {}) request_states[request_id] = state + self._merge_request_codec_params( + state, + request_sampling_params[index] if request_sampling_params is not None else None, + ) if state.get("finished"): stop_rows.append(row_stop) codec_deltas.append(empty_delta) @@ -965,7 +996,9 @@ def make_omni_output( min_tokens = _codec_int_param(state, "min_tokens", self._codec_min_tokens) max_tokens = _codec_int_param(state, "max_tokens", self._codec_max_tokens) eos_window_masked = bool(state.get("turn_end_drain")) and _turn_end_boundary_eos_masked(step) - if self._codec_temperature == 0.0: + # A request pinning temperature to 0 (override or stage default) + # takes the deterministic boundary sampler, same as single-frame. + if _codec_float_param(state, "codec_temperature", self._codec_temperature) == 0.0: sampled = self._sample_audio_code_greedy( hidden[end - 1 : end], codes, @@ -1113,11 +1146,38 @@ def merge_frame_outputs( multimodal_outputs={"codes": {"audio": codec_deltas}, "meta": meta_outputs}, ) + def _merge_request_codec_params(self, state: dict[str, Any], sampling_params: Any) -> None: + """Pin one request's effective codec sampling knobs into its state. + + The request's SamplingParams -- the engine merges the stage's + default_sampling_params with per-request overrides into it -- wins + field by field, so K-step in-model sampling keeps the single-frame + configuration contract: a request that overrides temperature or seed + steers the codec stream the same way it would under single-frame + decoding. Values already pinned (offline duplex states) and missing + sampling params (CPU tests, dummy runs without the runner hook) fall + back to the statically resolved deployment values. + """ + if sampling_params is None or not isinstance(state, dict): + return + for key, attr in ( + ("codec_temperature", "temperature"), + ("codec_top_k", "top_k"), + ("codec_top_p", "top_p"), + ("codec_repetition_penalty", "repetition_penalty"), + ("codec_seed", "seed"), + ): + if state.get(key) is None: + value = getattr(sampling_params, attr, None) + if value is not None: + state[key] = value + def _request_generator(self, request_id: str, device: torch.device) -> torch.Generator: generator = self._request_generators.get(request_id) if generator is None: + request_state = getattr(self, "_request_audio_states", {}).get(request_id, {}) generator = torch.Generator(device=device) - generator.manual_seed(self._codec_seed) + generator.manual_seed(_codec_int_param(request_state, "codec_seed", self._codec_seed)) self._request_generators[request_id] = generator return generator @@ -1149,18 +1209,22 @@ def _sample_audio_code( device_inputs_by_request = {} self._request_codec_device_inputs = device_inputs_by_request device_inputs = device_inputs_by_request.get(request_id) + request_states = getattr(self, "_request_audio_states", {}) + request_state = request_states.get(request_id, {}) if device_inputs is None: - request_states = getattr(self, "_request_audio_states", {}) - request_state = request_states.get(request_id, {}) device_inputs = ( torch.tensor( [_codec_int_param(request_state, "min_tokens", self._codec_min_tokens)], dtype=torch.int32, device=hidden_state.device, ), - torch.tensor([self._codec_temperature], dtype=torch.float32, device=hidden_state.device), torch.tensor( - [self._codec_repetition_penalty], + [_codec_float_param(request_state, "codec_temperature", self._codec_temperature)], + dtype=torch.float32, + device=hidden_state.device, + ), + torch.tensor( + [_codec_float_param(request_state, "codec_repetition_penalty", self._codec_repetition_penalty)], dtype=torch.float32, device=hidden_state.device, ), @@ -1175,8 +1239,11 @@ def _sample_audio_code( temperature_tensor, penalty_tensor, eos_token_id=eos_id, - top_k=self._codec_top_k, - top_p=self._codec_top_p, + # top-k/top-p stay host-side constants: they branch the sampler at + # capture time, and the per-request sampling graph bakes in this + # request's pinned values from its own state. + top_k=_codec_int_param(request_state, "codec_top_k", self._codec_top_k), + top_p=_codec_float_param(request_state, "codec_top_p", self._codec_top_p), min_tokens_to_keep=3, eos_window_masked=eos_window_masked, ) @@ -1228,12 +1295,13 @@ def _sample_audio_code_greedy( device_inputs_by_request = {} self._request_codec_device_inputs = device_inputs_by_request device_inputs = device_inputs_by_request.get(request_id) + request_state = getattr(self, "_request_audio_states", {}).get(request_id, {}) if device_inputs is None: device_inputs = ( torch.tensor([min_tokens], dtype=torch.int32, device=hidden_state.device), torch.tensor([1.0], dtype=torch.float32, device=hidden_state.device), torch.tensor( - [self._codec_repetition_penalty], + [_codec_float_param(request_state, "codec_repetition_penalty", self._codec_repetition_penalty)], dtype=torch.float32, device=hidden_state.device, ), @@ -1245,7 +1313,7 @@ def _sample_audio_code_greedy( device_state, min_tokens_tensor, penalty_tensor, - top_k=self._codec_top_k, + top_k=_codec_int_param(request_state, "codec_top_k", self._codec_top_k), eos_token_id=self._codec_eos_id, eos_window_masked=eos_window_masked, ) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 80eaf3e6793..55cae816fe2 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1432,6 +1432,15 @@ def _build_model_kwargs_extra(self) -> dict: ) for req_index, req_id in enumerate(self.input_batch.req_ids) ] + # Effective per-request sampling params (stage defaults + # merged with request overrides) for in-model sampling + # paths, so the K-step codec sampler keeps the + # single-frame configuration contract. Same order as + # request_sample_eligible / request_token_spans. + model_kwargs_extra["request_sampling_params"] = [ + req.sampling_params if (req := self.requests.get(req_id)) is not None else None + for req_id in self.input_batch.req_ids + ] except Exception as e: # Visible on purpose: the fallback is the equal rows-per-request # split, which can re-introduce the cross-request corruption this From 1b4cd23454e82333da281f907e4fb518def54f87 Mon Sep 17 00:00:00 2001 From: y-null Date: Thu, 24 Sep 2026 07:46:48 +0000 Subject: [PATCH 53/89] fix(minicpmo_4_5): drop dead K-step RNG draft; forward _batch_stop_logits probe on wrapper - remove _k_step_rngs/_k_step_base_seed: zero references repo-wide (incl. tests), superseded by _request_generators seeded from per-request codec_seed - add wrapper _batch_stop_logits property: npu_model_runner probes the attribute on the registered architecture; without the forward it always read None and the gate fell back to text_hidden_states (behavior unchanged, probe now resolves) Signed-off-by: y-null --- .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 11 +++++++++++ .../models/minicpmo_4_5/minicpmo_4_5_omni_tts.py | 4 ---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index fb7d50275d9..31d9d6ad90c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -657,6 +657,8 @@ def requires_request_sample_eligibility(self) -> bool: # * missing supports_multi_frame_decode -> the check silently reads False, # the multi-frame loop never runs and no reason is logged; # * missing take/set_batch_stop_logits -> per-frame stop rows unavailable; + # * missing _batch_stop_logits -> the runner's gate probe reads None and + # always falls back to text_hidden_states; # * missing merge_frame_outputs -> multi-frame results cannot be merged. # The speculative config alone is not enough: the runner-side loop and its # vocab-width correction (ensure_stop_token_vocab) stay dormant, the @@ -670,6 +672,15 @@ def supports_multi_frame_decode(self) -> bool: # speculative_config (the runner checks its num_spec_tokens for that). return self.model_stage == "tts" + @property + def _batch_stop_logits(self): + # The runner probes this attribute on the registered architecture + # (getattr in npu_model_runner); without the forward it reads None and + # the gate always falls back to text_hidden_states. + if self.model_stage != "tts": + return None + return self.talker._batch_stop_logits + def take_batch_stop_logits(self): if self.model_stage != "tts": return None diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index c32067f297a..784fa3da1cd 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -338,10 +338,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # read back from the injected stage-1 speculative_config. self._k_step_frames = self._parse_k_step_frames(vllm_config) self.supports_multi_frame_decode = self._k_step_frames > 0 - # Per-request device RNG streams for in-model codec sampling (the - # multi-frame loop cannot run the vLLM host sampler per frame). - self._k_step_rngs: dict[str, torch.Generator] = {} - self._k_step_base_seed = 0x5EEDC0DE # Last frame's stop/continue rows and the flattened per-token stop # sequence the merged step presents to compute_logits. self._k_last_stop_rows: list[bool] | None = None From cf447c12da2906c9670b417e768a8b26a4b56ecb Mon Sep 17 00:00:00 2001 From: y-null Date: Thu, 24 Sep 2026 09:01:19 +0000 Subject: [PATCH 54/89] fix(minicpmo_4_5): propagate request min_tokens into K-step state; keep top_k exact - pin a positive request min_tokens/min_new_tokens floor into the K-step codec sampling state: the NPU runner neutralizes vLLM MinTokensLogitsProcessor, so the in-model sampler is the only min-length guard left and offline requests had none; engine default 0 never overwrites the stage-resolved codec minimum - keep an explicit top_k exact in prepare_codec_logits: min_tokens_to_keep is a top_p-only floor in vLLM semantics, not a top_k widening, so top_k=1 stays greedy under the single-frame-per-step K-frame contract - initialize _request_generators in the test fixture (it bypasses __init__) and add regressions for both fixes Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 93 +++++++++++++++++++ .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 11 +++ .../minicpmo_4_5/talker_codec_sample.py | 7 +- 3 files changed, 110 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index f986a1ac268..f07200f566e 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -53,6 +53,7 @@ def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): model._codec_eos_id = _EOS_ID model._request_audio_states = {} model._request_codec_history = {} + model._request_generators = {} emb = nn.Embedding(_NUM_AUDIO_TOKENS, 4) with torch.no_grad(): emb.weight.zero_() @@ -534,3 +535,95 @@ def test_default_sampling_params_feed_codec_resolution(): resolved3 = resolve_codec_sampling_params(None, tts_config) assert resolved3["top_k"] == 100 assert resolved3["seed"] == 42 + + +def test_request_min_tokens_reaches_codec_state(): + """A positive request ``min_tokens`` floor reaches the K-step state. + + The NPU runner neutralizes vLLM's MinTokensLogitsProcessor, so the + in-model codec sampler is the only min-length guard left (PR #7929 + review). The engine default 0 means "no extra floor" and must not + overwrite the stage-resolved codec minimum. + """ + model = _make_talker(k_step_frames=8, scripted_samples=[42]) + model._codec_min_tokens = 50 + + base = dict(temperature=0.8, top_k=25, top_p=0.85, repetition_penalty=1.05, seed=None) + state = {"step": 0} + model._merge_request_codec_params(state, SimpleNamespace(min_tokens=100, **base)) + assert state["min_tokens"] == 100 + + # 0 (engine default) is not a floor: the key stays unset so the K-step + # loop keeps falling back to the stage-resolved minimum. + state_zero = {"step": 0} + model._merge_request_codec_params(state_zero, SimpleNamespace(min_tokens=0, **base)) + assert "min_tokens" not in state_zero + + # The min_new_tokens alias is honored too. + state_alias = {"step": 0} + model._merge_request_codec_params(state_alias, SimpleNamespace(min_new_tokens=7, **base)) + assert state_alias["min_tokens"] == 7 + + +def test_top_k_filter_honors_explicit_top_k(): + """top_k is kept exact; min_tokens_to_keep is a top_p-only floor. + + Passing min_tokens_to_keep=3 must not widen an explicit top_k of 1 or 2 + -- otherwise K-step sampling changes the single-frame contract (top_k=1 + would no longer be greedy) (PR #7929 review). + """ + from vllm_omni.model_executor.models.minicpmo_4_5.talker_codec_sample import ( + make_device_state, + prepare_codec_logits, + ) + + logits = torch.full((1, _NUM_AUDIO_TOKENS), -1.0) + logits[0, 10] = 5.0 + logits[0, 11] = 4.0 + logits[0, 12] = 3.0 + device_state = make_device_state(torch.zeros(0, dtype=torch.int32), step=0, max_tokens=100, finished=False) + kwargs = dict( + state=device_state, + min_tokens=torch.tensor([0]), + temperature=torch.tensor([0.8]), + repetition_penalty=torch.tensor([1.0]), + eos_token_id=_EOS_ID, + top_p=1.0, + min_tokens_to_keep=3, + ) + for top_k, expected in ((1, 1), (2, 2)): + filtered = prepare_codec_logits(logits.clone(), top_k=top_k, **kwargs) + assert int(torch.isfinite(filtered).sum()) == expected + + +def test_eos_window_mask_hides_codec_eos(): + """``eos_window_masked=True`` masks codec EOS regardless of the step. + + This is the turn-end drain path: duplex meta ``turn_end`` pins + ``state["turn_end_drain"]`` (tts preprocess), the K-step loop forwards it + as ``eos_window_masked``, and the sampler must hide EOS so the chunk can + drain its remaining cadence frames instead of ending early. + """ + from vllm_omni.model_executor.models.minicpmo_4_5.talker_codec_sample import ( + make_device_state, + prepare_codec_logits, + ) + + logits = torch.full((1, _NUM_AUDIO_TOKENS), 1.0) + logits[0, _EOS_ID] = 10.0 + device_state = make_device_state(torch.zeros(0, dtype=torch.int32), step=5, max_tokens=100, finished=False) + kwargs = dict( + state=device_state, + min_tokens=torch.tensor([0]), + temperature=torch.tensor([0.8]), + repetition_penalty=torch.tensor([1.0]), + eos_token_id=_EOS_ID, + top_p=1.0, + top_k=0, + ) + # step >= min_tokens: EOS is eligible and keeps its dominant logit. + unmasked = prepare_codec_logits(logits.clone(), eos_window_masked=False, **kwargs) + assert unmasked[0, _EOS_ID] != float("-inf") + # Drain window: EOS is forced out even though the step allows it. + masked = prepare_codec_logits(logits.clone(), eos_window_masked=True, **kwargs) + assert masked[0, _EOS_ID] == float("-inf") diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 784fa3da1cd..b559a78e7f8 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -1167,6 +1167,17 @@ def _merge_request_codec_params(self, state: dict[str, Any], sampling_params: An value = getattr(sampling_params, attr, None) if value is not None: state[key] = value + # The K-step codec sampler is the only min-length guard left (the NPU + # runner neutralizes vLLM's MinTokensLogitsProcessor), so the request + # floor must reach it. Pin only a positive floor: the engine default 0 + # means "no extra floor" and would otherwise overwrite the stage's + # resolved codec minimum. + if state.get("min_tokens") is None: + requested = getattr(sampling_params, "min_tokens", None) + if requested is None: + requested = getattr(sampling_params, "min_new_tokens", None) + if requested is not None and int(requested) > 0: + state["min_tokens"] = int(requested) def _request_generator(self, request_id: str, device: torch.device) -> torch.Generator: generator = self._request_generators.get(request_id) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index 78ee40f0c47..8efa775b4dd 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Framework boundaries for Talker codec sampling. Production stochastic sampling uses a logits-filter boundary implemented in @@ -99,7 +102,9 @@ def prepare_codec_logits( remove = remove.scatter(-1, sorted_indices, remove) logits.masked_fill_(remove, float("-inf")) if top_k > 0: - keep = min(VOCAB_SIZE, max(int(top_k), int(min_tokens_to_keep))) + # Honor the request's explicit top_k exactly: min_tokens_to_keep is a + # top_p-only floor (vLLM semantics), not a top_k widening. + keep = min(VOCAB_SIZE, int(top_k)) threshold = torch.topk(logits, keep, dim=-1).values[..., -1, None] logits.masked_fill_(logits < threshold, float("-inf")) return logits From cdbe8d6a9b47c1cf9e38c0b677e8b42a664dacb4 Mon Sep 17 00:00:00 2001 From: y-null Date: Fri, 25 Sep 2026 12:05:47 +0800 Subject: [PATCH 55/89] fix(minicpmo_4_5): align the K-step codec filtering with the single-frame sampler Mirror the engine's own filters on the in-model K-step codec logits: run top-k before top-p with one candidate kept, cap the segment frame budget by the request's remaining max_tokens, and read the request min_tokens (plus the min_new_tokens alias). Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 146 ++++++++++++++++-- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 66 +++++++- .../minicpmo_4_5/talker_codec_sample.py | 71 +++++---- .../platforms/npu/worker/talker_multiframe.py | 8 +- vllm_omni/worker/gpu_model_runner.py | 18 +++ 5 files changed, 255 insertions(+), 54 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index f07200f566e..3a186becaac 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -493,6 +493,110 @@ def test_request_sampling_params_pin_codec_knobs(): assert gen2.initial_seed() == 42 +def test_request_max_tokens_remaining_caps_kstep_frames(): + """The request's remaining output budget caps the K-step frame ceiling. + + The scheduler truncates the sampled ids at the request limit while the + connector concatenates every emitted codec frame, so a request whose limit + falls inside a K-frame step must stop emitting frames at the limit instead + of running on to the stage's codec budget (review on PR #7929). + """ + model = _make_talker(k_step_frames=8, scripted_samples=[42]) + model._codec_max_tokens = 4032 + state = {"step": 0, "max_tokens": 4032, "codes": torch.tensor([10, 11])} + model._request_audio_states["r1"] = state + + model.make_omni_output( + torch.randn(1, 8), + model_intermediate_buffer=[{"request_id": "r1"}], + request_token_spans=[(0, 1)], + request_sample_eligible=[True], + request_sampling_params=[ + SimpleNamespace(temperature=0.0, top_k=25, top_p=0.85, repetition_penalty=1.0, seed=7) + ], + request_max_tokens_remaining=[3], + ) + # K=8 frames were available, but the request had only 3 tokens left. + assert state["max_tokens"] == 3 + + # A request without a limit keeps the stage-resolved codec budget. + model2 = _make_talker(k_step_frames=8, scripted_samples=[43]) + model2._codec_max_tokens = 4032 + state2 = {"step": 0, "max_tokens": 4032, "codes": torch.tensor([10, 11])} + model2._request_audio_states["r2"] = state2 + model2.make_omni_output( + torch.randn(1, 8), + model_intermediate_buffer=[{"request_id": "r2"}], + request_token_spans=[(0, 1)], + request_sample_eligible=[True], + request_sampling_params=[ + SimpleNamespace(temperature=0.0, top_k=25, top_p=0.85, repetition_penalty=1.0, seed=7) + ], + request_max_tokens_remaining=[None], + ) + assert state2["max_tokens"] == 4032 + + +def test_reused_device_state_follows_tightened_max_tokens(): + """A reused codec device state picks up the caller's tightened ceiling. + + ``_merge_request_codec_params`` clamps ``state["max_tokens"]`` to the + request's remaining output budget on later steps, but the device state is + built once and reused. If it kept its creation-time budget, the codec loop + would emit past the frame the engine stops accepting ids at -- the + alignment the review asked to close (PR #7929). + """ + model = _make_talker(k_step_frames=8, scripted_samples=[42]) + model._codec_max_tokens = 4032 + model._codec_repetition_penalty = 1.05 + model.head_code = nn.ModuleList([nn.Linear(8, _NUM_AUDIO_TOKENS, bias=False)]) + # _make_talker installs a scripted stub; this test needs the real boundary. + real_greedy = type(model)._sample_audio_code_greedy + model._sample_audio_code_greedy = real_greedy.__get__(model, type(model)) + state = {"step": 0, "max_tokens": 4032, "codes": torch.tensor([10, 11])} + model._request_audio_states["r1"] = state + hidden = torch.randn(1, 8) + + # Step 1: no request limit, so the device state starts at the stage budget. + model._sample_audio_code_greedy(hidden, state["codes"], "r1", 0, 0, state["max_tokens"]) + assert int(model._request_codec_device_states["r1"].max_tokens.item()) == 4032 + + # Step 2: the host ceiling was tightened to 3; the reused device state has + # to follow it instead of staying at 4032. + state["max_tokens"] = 3 + model._sample_audio_code_greedy(hidden, state["codes"], "r1", 1, 0, state["max_tokens"]) + assert int(model._request_codec_device_states["r1"].max_tokens.item()) == 3 + + +def test_reused_device_state_follows_tightened_max_tokens_stochastic(): + """The stochastic boundary refreshes the reused device state the same way. + + ``_sample_audio_code`` reads the ceiling from the request state rather than + from an argument, so it needs its own check that a reused device state does + not keep the creation-time budget (PR #7929). + """ + model = _make_talker(k_step_frames=8, scripted_samples=[42]) + model._codec_max_tokens = 4032 + model._codec_temperature = 0.8 + model._codec_repetition_penalty = 1.05 + model._codec_top_k = 0 + model._codec_top_p = 1.0 + model._codec_seed = 42 + model.head_code = nn.ModuleList([nn.Linear(8, _NUM_AUDIO_TOKENS, bias=False)]) + real_sample = type(model)._sample_audio_code + model._sample_audio_code = real_sample.__get__(model, type(model)) + state = {"step": 0, "max_tokens": 4032, "codes": torch.tensor([10, 11])} + model._request_audio_states["r1"] = state + hidden = torch.randn(1, 8) + + model._sample_audio_code(hidden, state["codes"], "r1", 0) + assert int(model._request_codec_device_states["r1"].max_tokens.item()) == 4032 + + state["max_tokens"] = 5 + model._sample_audio_code(hidden, state["codes"], "r1", 1) + assert int(model._request_codec_device_states["r1"].max_tokens.item()) == 5 + + def test_default_sampling_params_feed_codec_resolution(): """``default_sampling_params`` sits between the YAML block and tts_config. @@ -565,35 +669,47 @@ def test_request_min_tokens_reaches_codec_state(): assert state_alias["min_tokens"] == 7 -def test_top_k_filter_honors_explicit_top_k(): - """top_k is kept exact; min_tokens_to_keep is a top_p-only floor. +def test_top_k_then_top_p_matches_single_frame_order(): + """K-step filtering follows the engine's order: top-k, then top-p. - Passing min_tokens_to_keep=3 must not widen an explicit top_k of 1 or 2 - -- otherwise K-step sampling changes the single-frame contract (top_k=1 - would no longer be greedy) (PR #7929 review). + The single-frame path applies top-k first and then computes top-p over the + top-k-filtered distribution, keeping at least one candidate + (vllm/v1/sample/ops/topk_topp_sampler.py:392/:404/:415). Filtering top-p + first -- or holding a floor of three candidates -- changes the sampling + distribution even for identical logits and sampling parameters, so enabling + K-step decoding would silently alter the single-frame contract + (PR #7929 review). """ from vllm_omni.model_executor.models.minicpmo_4_5.talker_codec_sample import ( make_device_state, prepare_codec_logits, ) - logits = torch.full((1, _NUM_AUDIO_TOKENS), -1.0) - logits[0, 10] = 5.0 - logits[0, 11] = 4.0 - logits[0, 12] = 3.0 device_state = make_device_state(torch.zeros(0, dtype=torch.int32), step=0, max_tokens=100, finished=False) kwargs = dict( state=device_state, min_tokens=torch.tensor([0]), - temperature=torch.tensor([0.8]), + temperature=torch.tensor([1.0]), repetition_penalty=torch.tensor([1.0]), eos_token_id=_EOS_ID, - top_p=1.0, - min_tokens_to_keep=3, ) - for top_k, expected in ((1, 1), (2, 2)): - filtered = prepare_codec_logits(logits.clone(), top_k=top_k, **kwargs) - assert int(torch.isfinite(filtered).sum()) == expected + # The review's example: candidate weights [40, 30, 20, 10] with top_k=3 and + # top_p=0.5. The engine keeps two candidates -- top-p runs over the + # renormalized top-3 distribution [0.444, 0.333, 0.222], whose ascending + # cumsum first exceeds 0.5 at the second candidate. Keeping three + # candidates was the reported bug. + weights = torch.tensor([40.0, 30.0, 20.0, 10.0]) + logits = torch.full((1, _NUM_AUDIO_TOKENS), -1e4) + logits[0, 10:14] = weights.log() + filtered = prepare_codec_logits(logits.clone(), top_k=3, top_p=0.5, **kwargs) + assert int(torch.isfinite(filtered).sum()) == 2 + # An explicit top_k is honored exactly: top_k=1 keeps a single candidate. + filtered_top_k_one = prepare_codec_logits(logits.clone(), top_k=1, top_p=1.0, **kwargs) + assert int(torch.isfinite(filtered_top_k_one).sum()) == 1 + # top-p alone keeps at least one candidate, like the engine's `at least one` + # row (topk_topp_sampler.py:415). + filtered_top_p_only = prepare_codec_logits(logits.clone(), top_k=0, top_p=0.01, **kwargs) + assert int(torch.isfinite(filtered_top_p_only).sum()) >= 1 def test_eos_window_mask_hides_codec_eos(): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 0286c9712ff..1202f13c9b8 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -883,6 +883,17 @@ def make_omni_output( f"MiniCPM-o continuous Talker received {len(request_sampling_params)} " f"sampling params for {len(infos)} requests" ) + # Remaining request output budget per request (runner-side + # ``max_tokens - num_output_tokens``). The engine truncates the sampled + # ids at the request limit while the connector concatenates every + # emitted codec frame, so the codec loop has to stop emitting at the + # same place (PR #7929 review). + request_max_tokens_remaining = kwargs.get("request_max_tokens_remaining") + if request_max_tokens_remaining is not None and len(request_max_tokens_remaining) != len(infos): + raise RuntimeError( + f"MiniCPM-o continuous Talker received {len(request_max_tokens_remaining)} " + f"remaining-token budgets for {len(infos)} requests" + ) emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos) stop_rows: list[torch.Tensor] = [] @@ -962,6 +973,9 @@ def make_omni_output( self._merge_request_codec_params( state, request_sampling_params[index] if request_sampling_params is not None else None, + max_tokens_remaining=( + request_max_tokens_remaining[index] if request_max_tokens_remaining is not None else None + ), ) if state.get("finished"): stop_rows.append(row_stop) @@ -988,7 +1002,7 @@ def make_omni_output( # steps after each 25-frame cadence boundary (see # _turn_end_boundary_eos_masked and state["turn_end_drain"]; # MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS is 4 chunks, so the window - # lands on steps 25-29 / 50-54 / 75-79 / 100-103). The EOS masking + # lands on steps 25-29 / 50-54 / 75-79 / 100-104). The EOS masking # lives in the sampling core (talker_codec_sample. # prepare_codec_logits / greedy_codec_sample), so the periodic # window rides along as the eos_window_masked flag below -- the same @@ -1146,7 +1160,13 @@ def merge_frame_outputs( multimodal_outputs={"codes": {"audio": codec_deltas}, "meta": meta_outputs}, ) - def _merge_request_codec_params(self, state: dict[str, Any], sampling_params: Any) -> None: + def _merge_request_codec_params( + self, + state: dict[str, Any], + sampling_params: Any, + *, + max_tokens_remaining: int | None = None, + ) -> None: """Pin one request's effective codec sampling knobs into its state. The request's SamplingParams -- the engine merges the stage's @@ -1157,6 +1177,12 @@ def _merge_request_codec_params(self, state: dict[str, Any], sampling_params: An decoding. Values already pinned (offline duplex states) and missing sampling params (CPU tests, dummy runs without the runner hook) fall back to the statically resolved deployment values. + + ``max_tokens_remaining`` is the runner-computed request output budget + (``SamplingParams.max_tokens - num_output_tokens``). It tightens this + segment's frame ceiling so the emitted codec frames stay aligned with + the ids the engine accepts; ``None`` leaves the stage-resolved codec + budget in charge. """ if sampling_params is None or not isinstance(state, dict): return @@ -1182,6 +1208,18 @@ def _merge_request_codec_params(self, state: dict[str, Any], sampling_params: An requested = getattr(sampling_params, "min_new_tokens", None) if requested is not None and int(requested) > 0: state["min_tokens"] = int(requested) + # The request's remaining output budget caps this segment's frame + # ceiling: the scheduler truncates the sampled ids at the request limit + # while the connector concatenates every emitted codec frame, so the + # codec loop must stop emitting where the engine stops accepting. The + # stage-resolved codec budget is the ceiling only when the request + # carries no limit of its own; a smaller request limit tightens it, + # never the other way round (PR #7929 review). + if max_tokens_remaining is not None: + current = _codec_int_param(state, "max_tokens", self._codec_max_tokens) + ceiling = int(state.get("step") or 0) + max(int(max_tokens_remaining), 0) + if ceiling < current: + state["max_tokens"] = ceiling def _request_generator(self, request_id: str, device: torch.device) -> torch.Generator: generator = self._request_generators.get(request_id) @@ -1204,10 +1242,10 @@ def _sample_audio_code( if device_states is None: device_states = {} self._request_codec_device_states = device_states + request_states = getattr(self, "_request_audio_states", {}) + request_state = request_states.get(request_id, {}) device_state = device_states.get(request_id) if device_state is None: - request_states = getattr(self, "_request_audio_states", {}) - request_state = request_states.get(request_id, {}) device_state = make_device_state( history, step=step, @@ -1215,6 +1253,18 @@ def _sample_audio_code( finished=False, ) device_states[request_id] = device_state + else: + # The host state can tighten the frame ceiling after the device + # state was built: _merge_request_codec_params clamps + # state["max_tokens"] to the request's remaining output budget, + # while a reused device state still holds its creation-time value. + # Without this refresh the codec loop and the engine can stop at + # different frames, which is what the emitted codes must not do. + # Refresh in place: swapping the tensor would re-capture the + # per-request sampling graph. + budget = _codec_int_param(request_state, "max_tokens", self._codec_max_tokens) + if int(device_state.max_tokens.item()) != budget: + device_state.max_tokens.fill_(budget) device_inputs_by_request = getattr(self, "_request_codec_device_inputs", None) if device_inputs_by_request is None: device_inputs_by_request = {} @@ -1255,7 +1305,6 @@ def _sample_audio_code( # request's pinned values from its own state. top_k=_codec_int_param(request_state, "codec_top_k", self._codec_top_k), top_p=_codec_float_param(request_state, "codec_top_p", self._codec_top_p), - min_tokens_to_keep=3, eos_window_masked=eos_window_masked, ) probabilities = torch.softmax(logits, dim=-1) @@ -1300,6 +1349,12 @@ def _sample_audio_code_greedy( finished=False, ) device_states[request_id] = device_state + elif int(device_state.max_tokens.item()) != int(max_tokens): + # Same refresh as _sample_audio_code: the caller passes the host + # state's current frame ceiling, which the request's remaining + # output budget may have tightened after this device state was + # built. Refreshing in place keeps the per-request sampling graph. + device_state.max_tokens.fill_(int(max_tokens)) device_inputs_by_request = getattr(self, "_request_codec_device_inputs", None) if device_inputs_by_request is None: @@ -1324,7 +1379,6 @@ def _sample_audio_code_greedy( device_state, min_tokens_tensor, penalty_tensor, - top_k=_codec_int_param(request_state, "codec_top_k", self._codec_top_k), eos_token_id=self._codec_eos_id, eos_window_masked=eos_window_masked, ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index 8efa775b4dd..5e135df8b3b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -5,8 +5,9 @@ Production stochastic sampling uses a logits-filter boundary implemented in graph-capturable tensor ops: history counting, temperature, repetition -penalty, EOS mask, top-p and top-k, while native NPU operators keep softmax -and ``torch.multinomial`` semantics (including Generator state) unchanged. +penalty, EOS mask, top-k, then top-p over the top-k-filtered distribution +keeping at least one candidate, while native NPU operators keep softmax and +``torch.multinomial`` semantics (including Generator state) unchanged. """ from __future__ import annotations @@ -17,7 +18,6 @@ VOCAB_SIZE = 6562 HISTORY_WINDOW = 16 -MIN_TOKENS_TO_KEEP = 3 @dataclass(frozen=True) @@ -67,7 +67,6 @@ def prepare_codec_logits( eos_token_id: int, top_k: int, top_p: float, - min_tokens_to_keep: int = MIN_TOKENS_TO_KEEP, eos_window_masked: bool = False, ) -> torch.Tensor: """Prepare and filter logits without changing native NPU RNG semantics. @@ -76,14 +75,8 @@ def prepare_codec_logits( step counter -- the duplex turn-end drain window (a few frames after each cadence boundary), which the ``step < min_tokens`` criterion cannot express. """ - logits = raw_logits.float() / temperature.reshape(1, 1) - positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) - valid = positions < state.history_len.reshape(1, 1) - safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) - counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=raw_logits.device) - counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) - alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) - logits = torch.where(logits < 0, logits * alpha, logits / alpha) + logits = raw_logits.float() + # Engine step 5a: the stop/EOS mask is applied before the penalties. eos = logits[..., eos_token_id : eos_token_id + 1] if eos_window_masked: eos = torch.full_like(eos, float("-inf")) @@ -94,19 +87,31 @@ def prepare_codec_logits( eos, ) logits = torch.cat([logits[..., :eos_token_id], eos, logits[..., eos_token_id + 1 :]], dim=-1) + # Engine step 6a: repetition penalty over the codec history window. + positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) + valid = positions < state.history_len.reshape(1, 1) + safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) + counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=raw_logits.device) + counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) + alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) + logits = torch.where(logits < 0, logits * alpha, logits / alpha) + # Engine step 7b: temperature is applied after the penalties. The division + # also gives the in-place masks below a tensor of their own. + logits = logits / temperature.reshape(1, 1) + # Engine step 7d: top-k first (topk_topp_sampler.py:392). + if top_k > 0: + keep = min(VOCAB_SIZE, int(top_k)) + threshold = torch.topk(logits, keep, dim=-1).values[..., -1, None] + logits.masked_fill_(logits < threshold, float("-inf")) + # Engine step 7d: top-p runs over the top-k-filtered distribution and keeps + # at least one candidate (topk_topp_sampler.py:404/:415). if 0.0 < float(top_p) < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=False, dim=-1) cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) remove = cumulative_probs <= (1.0 - float(top_p)) - remove[..., -min_tokens_to_keep:] = False + remove[..., -1:] = False remove = remove.scatter(-1, sorted_indices, remove) logits.masked_fill_(remove, float("-inf")) - if top_k > 0: - # Honor the request's explicit top_k exactly: min_tokens_to_keep is a - # top_p-only floor (vLLM semantics), not a top_k widening. - keep = min(VOCAB_SIZE, int(top_k)) - threshold = torch.topk(logits, keep, dim=-1).values[..., -1, None] - logits.masked_fill_(logits < threshold, float("-inf")) return logits @@ -154,23 +159,21 @@ def greedy_codec_sample( min_tokens: torch.Tensor, repetition_penalty: torch.Tensor, *, - top_k: int, eos_token_id: int, eos_window_masked: bool = False, ) -> TalkerCodecSampleResult: """Greedy codec sample in graph-capturable tensor ops. + The engine's greedy branch returns the argmax after the stop/EOS mask and + the penalties and skips temperature, top-k and top-p entirely + (``vllm/v1/sample/sampler.py:30-60``), so no top-k mask belongs on this + path -- masking candidates below the argmax cannot change the argmax itself. + ``eos_window_masked`` masks the codec EOS for this frame regardless of the step counter -- see prepare_codec_logits. """ - positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) - valid = positions < state.history_len.reshape(1, 1) - safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) - counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=raw_logits.device) - counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) - alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) - penalized = torch.where(raw_logits < 0, raw_logits * alpha, raw_logits / alpha) - + penalized = raw_logits.float() + # Engine step 5a: the stop/EOS mask is applied before the penalties. eos = penalized[..., eos_token_id : eos_token_id + 1] if eos_window_masked: eos = torch.full_like(eos, float("-inf")) @@ -181,10 +184,14 @@ def greedy_codec_sample( eos, ) penalized = torch.cat([penalized[..., :eos_token_id], eos, penalized[..., eos_token_id + 1 :]], dim=-1) - if top_k > 0: - keep = min(VOCAB_SIZE, max(int(top_k), MIN_TOKENS_TO_KEEP)) - threshold = torch.topk(penalized, keep, dim=-1).values[..., -1, None] - penalized = penalized.masked_fill(penalized < threshold, float("-inf")) + # Engine step 6a: repetition penalty over the codec history window. + positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) + valid = positions < state.history_len.reshape(1, 1) + safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) + counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=raw_logits.device) + counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) + alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) + penalized = torch.where(penalized < 0, penalized * alpha, penalized / alpha) sampled = torch.argmax(penalized, dim=-1).to(torch.int32) active = ~state.finished diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py index b798d22ad06..bb8882fe5d6 100644 --- a/vllm_omni/platforms/npu/worker/talker_multiframe.py +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -38,11 +38,17 @@ ## What makes it correct * **Stop is still per frame.** ``make_omni_output`` evaluates the codec EOS and - the token limit on every frame exactly as it does today. A request that ends + the frame ceiling on every frame exactly as it does today. A request that ends at frame j reports its stop row there; frames j+1..K-1 take the "already finished" early return, emit an empty delta and report a stop row too. vLLM's rejection sampler sees the first stop and truncates the request to j+1 tokens, which is the same sequence a one-frame-per-step run would produce. + The ceiling is not only the stage-resolved codec budget: the runner also + passes each request's remaining output budget + (``SamplingParams.max_tokens - num_output_tokens``) and the codec merge clamps + it into ``state["max_tokens"]``, so a request limit that falls inside a + K-frame step ends the loop where the engine stops accepting ids instead of + emitting frames past it (PR #7929 review). * **The emitted codes are unchanged.** The extra frames after a stop contribute nothing: their deltas are empty or the ``-1`` sentinel the connector already drops (``_codec_scalars``). So the audio is expected to be bit-identical, not diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index f8cf0f064ed..d047b479678 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1441,6 +1441,24 @@ def _build_model_kwargs_extra(self) -> dict: req.sampling_params if (req := self.requests.get(req_id)) is not None else None for req_id in self.input_batch.req_ids ] + # Remaining request output budget per request, so the + # in-model K-step codec loop stops emitting frames where the + # engine stops accepting tokens: the scheduler truncates the + # sampled ids at the request limit, but the connector still + # concatenates every emitted codec frame (PR #7929 review). + # None means the request carries no max_tokens, which leaves + # the stage-resolved codec budget in charge. Same order as + # request_sampling_params. + model_kwargs_extra["request_max_tokens_remaining"] = [ + ( + max(int(req.sampling_params.max_tokens) - int(req.num_output_tokens), 0) + if (req := self.requests.get(req_id)) is not None + and req.sampling_params is not None + and req.sampling_params.max_tokens is not None + else None + ) + for req_id in self.input_batch.req_ids + ] except Exception as e: # Visible on purpose: the fallback is the equal rows-per-request # split, which can re-introduce the cross-request corruption this From 61cf0d964a828dcf597a96c9987547be15e400c4 Mon Sep 17 00:00:00 2001 From: y-null Date: Fri, 25 Sep 2026 14:21:19 +0800 Subject: [PATCH 56/89] fix(minicpmo_4_5): mirror the remaining engine sampling filters in the K-step codec path Propagate the request-level knobs the engine would still apply to codec logits or stop conditions: min_p, ignore_eos, stop_token_ids, allowed_token_ids, logit_bias, frequency/presence penalties, the min_tokens stop-id set, max_model_len, and repetition detection. Clamp min_p <= top_p and keep bad_words unimplemented (a no-op on codec ids in both paths). Signed-off-by: y-null --- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 222 ++++++++++- .../minicpmo_4_5/talker_codec_sample.py | 351 +++++++++++++++--- vllm_omni/worker/gpu_model_runner.py | 28 +- 3 files changed, 539 insertions(+), 62 deletions(-) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 1202f13c9b8..5c9ddc09a34 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -34,6 +34,7 @@ MINICPMO45_DUPLEX_TURN_END_CODEC_TOKENS, ) from vllm_omni.model_executor.models.minicpmo_4_5.talker_codec_sample import ( + VOCAB_SIZE as _CODEC_VOCAB_SIZE, TalkerCodecDeviceState, TalkerCodecSampleResult, codec_sample_result, @@ -234,6 +235,14 @@ def resolve_codec_sampling_params( raise ValueError("codec_sampling_params.min_tokens must be >= 0") if resolved["max_tokens"] <= 0: raise ValueError("codec_sampling_params.max_tokens must be > 0") + # Same constraint the engine enforces on SamplingParams + # (vllm/sampling_params.py:610-613): a floor above the ceiling can never be + # satisfied, so reject it here instead of silently accepting it. + if resolved["min_tokens"] > resolved["max_tokens"]: + raise ValueError( + "codec_sampling_params.min_tokens must be <= max_tokens, got " + f"{resolved['min_tokens']} > {resolved['max_tokens']}" + ) logger.info("MiniCPM-o Talker codec sampling %s (from %s)", resolved, sources) return resolved @@ -244,6 +253,159 @@ def _codec_float_param(state: Any, key: str, fallback: float) -> float: return float(fallback if value is None else value) +def _codec_bool_param(state: Any, key: str, fallback: bool) -> bool: + """Bool counterpart of _codec_int_param; same None-means-unset rule.""" + value = state.get(key) if isinstance(state, Mapping) else None + return bool(fallback if value is None else value) + + +def _codec_stop_ids(state: Any) -> tuple[int, ...]: + """The request's extra stop ids, deduplicated in a stable order. + + Populated by ``_merge_request_codec_params`` from + ``SamplingParams.stop_token_ids``. The engine validates these against + ``model_config.get_vocab_size()`` before the request is accepted + (``vllm/sampling_params.py:874-895``), so surviving ids are in-vocabulary + for the model -- but the codec head is only ``_CODEC_VOCAB_SIZE`` wide, so + anything at or beyond it is dropped here, the same way + ``_codec_allowed_ids_mask`` and ``_codec_logit_bias_items`` drop theirs, + which keeps the index tensors in range. The engine censors and stops on + ``all_stop_token_ids`` = eos + these ids + extra eos ids + (``builtin.py:196-207``, ``v1/core/sched/utils.py:105``); K-step has no + engine layer left, so it censors and stops on them itself. Empty means the + request set none, i.e. the codec-EOS-only behaviour. + """ + value = state.get("codec_stop_token_ids") if isinstance(state, Mapping) else None + if not value: + return () + return tuple( + sorted({int(token_id) for token_id in value if 0 <= int(token_id) < _CODEC_VOCAB_SIZE}) + ) + + +def _codec_allowed_ids_mask(state: Any, device: torch.device) -> torch.Tensor | None: + """Censor mask for the request's codec whitelist, or None when unset. + + The engine builds ``allowed_token_ids_mask`` with True on the ids to + censor (``v1/worker/gpu_input_batch.py:463-466``) and applies it as + ``masked_fill_(-inf)`` (``sampler.py:393-394``). The request's list is + keyed by the text vocabulary, so ids outside the codec vocabulary are + dropped; a list with nothing left in range is treated as unset rather + than producing an all ``-inf`` row. + """ + value = state.get("codec_allowed_token_ids") if isinstance(state, Mapping) else None + if not value: + return None + keep = sorted({int(token_id) for token_id in value if 0 <= int(token_id) < _CODEC_VOCAB_SIZE}) + if not keep: + return None + mask = torch.ones((1, _CODEC_VOCAB_SIZE), dtype=torch.bool, device=device) + mask[0, torch.tensor(keep, dtype=torch.long, device=device)] = False + return mask + + +def _codec_logit_bias_items(state: Any) -> tuple[tuple[int, ...], tuple[float, ...]]: + """The request's codec logit biases as parallel tuples, in a stable order. + + ``LogitBiasLogitsProcessor`` adds ``bias`` to ``logits[req, tok]`` + (``builtin.py:161-164``). The request's dict is keyed by the text + vocabulary, so keys outside the codec vocabulary are dropped. + """ + value = state.get("codec_logit_bias") if isinstance(state, Mapping) else None + if not isinstance(value, Mapping): + return (), () + pairs = sorted( + (int(token_id), float(bias)) + for token_id, bias in value.items() + if 0 <= int(token_id) < _CODEC_VOCAB_SIZE + ) + if not pairs: + return (), () + return tuple(token_id for token_id, _ in pairs), tuple(bias for _, bias in pairs) + + +def _codec_filter_inputs(state: Any, device: torch.device) -> dict[str, Any]: + """Whitelist mask and logit-bias tuples for the codec filters. + + Returns an empty mapping when the request set neither, so the call site + keeps the sampler's no-op defaults. + """ + inputs: dict[str, Any] = {} + mask = _codec_allowed_ids_mask(state, device) + if mask is not None: + inputs["allowed_token_ids_mask"] = mask + bias_ids, bias_values = _codec_logit_bias_items(state) + if bias_ids: + inputs["logit_bias_ids"] = bias_ids + inputs["logit_bias_values"] = bias_values + return inputs + + +def _codec_full_bins_input(state: Any, device: torch.device) -> torch.Tensor | None: + """A zeroed whole-stream counter when the request uses the stream penalties. + + Returns None when the request sets neither ``frequency_penalty`` nor + ``presence_penalty``, so the sampler skips both the tensor and the per-frame + scatter. The tensor lives in the request's device state, which survives + across chunks, so the counts stay whole-stream; it is advanced in place. + """ + if not ( + _codec_float_param(state, "codec_frequency_penalty", 0.0) + or _codec_float_param(state, "codec_presence_penalty", 0.0) + ): + return None + return torch.zeros(_CODEC_VOCAB_SIZE, dtype=torch.float32, device=device) + + +def _codec_has_repeating_pattern( + token_ids: Sequence[int], + pattern_len: int, + repetition_min_count: int, +) -> bool: + """Tail-repetition test, copied from ``v1/core/sched/utils.py:10-25``.""" + for n in range(1, pattern_len + 1): + target_token = token_ids[-n] + for m in range(1, repetition_min_count): + if token_ids[-(pattern_len * m + n)] != target_token: + return False + return True + + +def _codec_repetition_detected(state: Any) -> bool: + """Whether the codec stream hit the request's repetition pattern. + + Mirrors ``check_sequence_repetition`` (``v1/core/sched/utils.py:28-59``): + every pattern length in ``[min_pattern_size, max_pattern_size]`` is tried and + the tail must repeat at least ``min_count`` times. The engine runs this in + ``check_stop`` over the request's whole output; the K-step path keeps the same + list host-side, because the model samples the codec ids itself. An unset + ``repetition_detection`` returns immediately, so the default pays nothing. + """ + params = state.get("codec_repetition_detection") if isinstance(state, Mapping) else None + if params is None: + return False + token_ids = state.get("codec_full_ids") + if not token_ids: + return False + max_pattern_size = int(getattr(params, "max_pattern_size", 0) or 0) + min_pattern_size = int(getattr(params, "min_pattern_size", 0) or 0) + min_count = int(getattr(params, "min_count", 0) or 0) + if min_pattern_size <= 0: + min_pattern_size = 1 + if max_pattern_size <= 0 or min_count < 2 or min_pattern_size > max_pattern_size: + return False + for pattern_len in range(min_pattern_size, max_pattern_size + 1): + if pattern_len * min_count > len(token_ids): + return False + if _codec_has_repeating_pattern(token_ids, pattern_len, min_count): + return True + return False + + + + + + def _codec_int_param(state: Any, key: str, fallback: int) -> int: """Read an integer codec knob from a request state, None meaning unset. @@ -1033,10 +1195,28 @@ def make_omni_output( # stub this method with the sampled tensor itself. sampled = stochastic_result.reshape(()).to(torch.long) sampled_id = int(sampled.item()) - is_eos = sampled_id == self._codec_eos_id + full_ids = state.get("codec_full_ids") + if isinstance(full_ids, list): + full_ids.append(sampled_id) + if _codec_bool_param(state, "codec_ignore_eos", False): + # ignore_eos blanks the engine's _eos_token_id + # (vllm/sampling_params.py:670-671), so a sampled codec EOS is + # just another frame and the request runs to its length budget. + is_eos = False + else: + is_eos = sampled_id == self._codec_eos_id + # stop_token_ids: the engine ends the request on any of these ids + # (v1/core/sched/utils.py:105); the codec stream ends it on the codec + # EOS, so both stop conditions apply here. + if not is_eos and sampled_id in _codec_stop_ids(state): + is_eos = True state["step"] = _codec_int_param(state, "step", 0) + 1 reached_limit = int(state["step"]) >= _codec_int_param(state, "max_tokens", self._codec_max_tokens) - finished = is_eos or reached_limit + # The engine stops the request when check_stop sees the pattern + # (v1/core/sched/utils.py:125-133); the K-step path owns the + # codec ids, so it has to make that call itself. + repetition_stop = _codec_repetition_detected(state) + finished = is_eos or reached_limit or repetition_stop state["finished"] = finished # MiniCPMTTS.generate_chunk consumes the boundary sample but # returns only codes that were fed into the retained KV state. @@ -1192,11 +1372,34 @@ def _merge_request_codec_params( ("codec_top_p", "top_p"), ("codec_repetition_penalty", "repetition_penalty"), ("codec_seed", "seed"), + ("codec_min_p", "min_p"), + ("codec_ignore_eos", "ignore_eos"), + ("codec_allowed_token_ids", "allowed_token_ids"), + ("codec_logit_bias", "logit_bias"), + ("codec_frequency_penalty", "frequency_penalty"), + ("codec_presence_penalty", "presence_penalty"), ): if state.get(key) is None: value = getattr(sampling_params, attr, None) if value is not None: state[key] = value + # stop_token_ids joins the codec censor set: while a request sits below + # its min_tokens floor the engine masks every id of all_stop_token_ids + # (builtin.py:196-207) and stops on them afterwards + # (v1/core/sched/utils.py:105). Only a non-empty list is pinned, so the + # default keeps the codec-EOS-only behaviour. + if state.get("codec_stop_token_ids") is None: + requested_stop_ids = getattr(sampling_params, "stop_token_ids", None) + if requested_stop_ids: + state["codec_stop_token_ids"] = [int(token_id) for token_id in requested_stop_ids] + # repetition_detection is a whole-stream pattern test + # (v1/core/sched/utils.py:28-59), so its record has to keep every + # emitted codec id. Started here, appended per frame below. + if state.get("codec_repetition_detection") is None: + detection = getattr(sampling_params, "repetition_detection", None) + if detection is not None: + state["codec_repetition_detection"] = detection + state.setdefault("codec_full_ids", []) # The K-step codec sampler is the only min-length guard left (the NPU # runner neutralizes vLLM's MinTokensLogitsProcessor), so the request # floor must reach it. Pin only a positive floor: the engine default 0 @@ -1251,6 +1454,7 @@ def _sample_audio_code( step=step, max_tokens=_codec_int_param(request_state, "max_tokens", self._codec_max_tokens), finished=False, + full_bins=_codec_full_bins_input(request_state, hidden_state.device), ) device_states[request_id] = device_state else: @@ -1305,6 +1509,11 @@ def _sample_audio_code( # request's pinned values from its own state. top_k=_codec_int_param(request_state, "codec_top_k", self._codec_top_k), top_p=_codec_float_param(request_state, "codec_top_p", self._codec_top_p), + min_p=_codec_float_param(request_state, "codec_min_p", 0.0), + min_tokens_stop_ids=_codec_stop_ids(request_state), + **_codec_filter_inputs(request_state, hidden_state.device), + frequency_penalty=_codec_float_param(request_state, "codec_frequency_penalty", 0.0), + presence_penalty=_codec_float_param(request_state, "codec_presence_penalty", 0.0), eos_window_masked=eos_window_masked, ) probabilities = torch.softmax(logits, dim=-1) @@ -1317,6 +1526,7 @@ def _sample_audio_code( device_state, sampled, eos_token_id=eos_id, + ignore_eos=_codec_bool_param(request_state, "codec_ignore_eos", False), ) device_states[request_id] = result.state return result @@ -1340,6 +1550,7 @@ def _sample_audio_code_greedy( if device_states is None: device_states = {} self._request_codec_device_states = device_states + request_state = getattr(self, "_request_audio_states", {}).get(request_id, {}) device_state = device_states.get(request_id) if device_state is None: device_state = make_device_state( @@ -1347,6 +1558,7 @@ def _sample_audio_code_greedy( step=step, max_tokens=max_tokens, finished=False, + full_bins=_codec_full_bins_input(request_state, hidden_state.device), ) device_states[request_id] = device_state elif int(device_state.max_tokens.item()) != int(max_tokens): @@ -1361,7 +1573,6 @@ def _sample_audio_code_greedy( device_inputs_by_request = {} self._request_codec_device_inputs = device_inputs_by_request device_inputs = device_inputs_by_request.get(request_id) - request_state = getattr(self, "_request_audio_states", {}).get(request_id, {}) if device_inputs is None: device_inputs = ( torch.tensor([min_tokens], dtype=torch.int32, device=hidden_state.device), @@ -1380,6 +1591,11 @@ def _sample_audio_code_greedy( min_tokens_tensor, penalty_tensor, eos_token_id=self._codec_eos_id, + min_tokens_stop_ids=_codec_stop_ids(request_state), + **_codec_filter_inputs(request_state, hidden_state.device), + frequency_penalty=_codec_float_param(request_state, "codec_frequency_penalty", 0.0), + presence_penalty=_codec_float_param(request_state, "codec_presence_penalty", 0.0), + ignore_eos=_codec_bool_param(request_state, "codec_ignore_eos", False), eos_window_masked=eos_window_masked, ) device_states[request_id] = result.state diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index 5e135df8b3b..a48d2459514 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -4,10 +4,18 @@ """Framework boundaries for Talker codec sampling. Production stochastic sampling uses a logits-filter boundary implemented in -graph-capturable tensor ops: history counting, temperature, repetition -penalty, EOS mask, top-k, then top-p over the top-k-filtered distribution -keeping at least one candidate, while native NPU operators keep softmax and -``torch.multinomial`` semantics (including Generator state) unchanged. +graph-capturable tensor ops: stop-token censoring, whitelist censoring, +per-id bias, repetition penalty, frequency/presence penalties, temperature, +min-p, top-k, then top-p over the top-k-filtered distribution keeping at least +one candidate, while native NPU operators keep softmax and ``torch.multinomial`` +semantics (including Generator state) unchanged. + +Order and defaults follow the engine's single-frame path so that enabling the +multi-frame (K-step) decode does not change the sampling semantics: +``vllm/v1/sample/sampler.py`` for the temperature/min-p/top-k/top-p order, +``logits_processor/builtin.py`` for the censoring, bias and min-p processors and +``model_executor/layers/utils.py`` for the penalties. Every added knob is a +no-op at its default value. """ from __future__ import annotations @@ -27,6 +35,15 @@ class TalkerCodecDeviceState: step: torch.Tensor max_tokens: torch.Tensor finished: torch.Tensor + # Whole-stream occurrence counts, one entry per codec id, or ``None`` when + # the request set neither frequency_penalty nor presence_penalty. The engine + # derives both penalties from the request's entire prompt+output sequence + # (``model_executor/layers/utils.py:73-78``), which the 16-frame ``history`` + # window cannot express; this tensor is rebuilt from the host-side record at + # every chunk boundary and advanced in place per frame. It is mutated rather + # than re-allocated because a fresh VOCAB_SIZE-wide tensor per frame would + # be an allocation per decode step. + full_bins: torch.Tensor | None = None @dataclass(frozen=True) @@ -42,6 +59,7 @@ def make_device_state( step: int, max_tokens: int, finished: bool, + full_bins: torch.Tensor | None = None, ) -> TalkerCodecDeviceState: """Create one persistent B=1 state at request/segment setup time.""" recent = codes.reshape(-1)[-HISTORY_WINDOW:].to(dtype=torch.int32) @@ -54,7 +72,163 @@ def make_device_state( step=torch.tensor([step], dtype=torch.int32, device=codes.device), max_tokens=torch.tensor([max_tokens], dtype=torch.int32, device=codes.device), finished=torch.tensor([finished], dtype=torch.bool, device=codes.device), + full_bins=full_bins, + ) + + +def _stop_id_tuple(eos_token_id: int, min_tokens_stop_ids: tuple[int, ...]) -> tuple[int, ...]: + """The censored stop set, codec EOS first and de-duplicated in a stable order. + + Taken from the engine's ``all_stop_token_ids`` (``sampling_params.py:537``, + ``:676``, ``:688``). An empty ``min_tokens_stop_ids`` keeps the + codec-EOS-only behaviour, which is what the K-step path used before this + knob existed. + """ + return tuple(dict.fromkeys((int(eos_token_id), *min_tokens_stop_ids))) + + +def _censor_stop_logits( + logits: torch.Tensor, + state: TalkerCodecDeviceState, + min_tokens: torch.Tensor, + *, + stop_ids: tuple[int, ...], + eos_window_masked: bool, +) -> torch.Tensor: + """Censor every stop id while the request is below its ``min_tokens`` floor. + + The engine's ``MinTokensLogitsProcessor`` masks ``all_stop_token_ids`` + (eos + ``stop_token_ids`` + extra eos ids) and lifts the mask by a spec-aware + counter (``builtin.py:196-207`` / ``:296-373``); here the counter is the + state's frame step, which the single-frame path uses as well + (``minicpmo_4_5_omni_tts.py``'s ``mask_eos_rows``). ``eos_window_masked`` + overrides the floor for the duplex turn-end drain window, which the + ``step < min_tokens`` criterion cannot express. + """ + stop_index = torch.tensor(stop_ids, dtype=torch.long, device=logits.device) + stop_logits = logits.index_select(-1, stop_index) + if eos_window_masked: + censored = torch.full_like(stop_logits, float("-inf")) + else: + censored = torch.where( + (state.step < min_tokens).reshape(1, 1), + torch.full_like(stop_logits, float("-inf")), + stop_logits, + ) + return logits.index_copy(-1, stop_index, censored) + + +def _censor_disallowed_logits( + logits: torch.Tensor, + allowed_token_ids_mask: torch.Tensor | None, +) -> torch.Tensor: + """Blank every id the request's whitelist excludes. + + The engine precomputes a boolean mask whose True entries are the ids to + censor (``v1/worker/gpu_input_batch.py:463-466``) and applies it as + ``logits.masked_fill_(mask, -inf)`` (``sampler.py:393-394``). ``None`` means + the request set no whitelist, which is the existing behaviour. + """ + if allowed_token_ids_mask is None: + return logits + return logits.masked_fill(allowed_token_ids_mask, float("-inf")) + + +def _apply_logit_bias( + logits: torch.Tensor, + logit_bias_ids: tuple[int, ...], + logit_bias_values: tuple[float, ...], +) -> torch.Tensor: + """Add the request's per-id biases. + + Mirrors ``LogitBiasLogitsProcessor.apply`` (``builtin.py:161-164``), which + the engine runs after the censoring masks and before the penalties, so a + censored id stays censored. Empty ids means the request set no bias. + """ + if not logit_bias_ids: + return logits + bias_index = torch.tensor(logit_bias_ids, dtype=torch.long, device=logits.device) + bias_values = torch.tensor(logit_bias_values, dtype=logits.dtype, device=logits.device) + current = logits.index_select(-1, bias_index) + return logits.index_copy(-1, bias_index, current + bias_values) + + +def _filter_codec_logits( + logits: torch.Tensor, + state: TalkerCodecDeviceState, + min_tokens: torch.Tensor, + *, + stop_ids: tuple[int, ...], + allowed_token_ids_mask: torch.Tensor | None, + logit_bias_ids: tuple[int, ...], + logit_bias_values: tuple[float, ...], + eos_window_masked: bool, +) -> torch.Tensor: + """The masks both sampling paths share, in the engine's order. + + ``allowed ids -> stop censoring -> logit bias`` are all applied by the engine + before its penalties (``sampler.py:393-405``); the bias is last because it is + additive while the other two write ``-inf``, so a censored id must not be + revived by a bias. + """ + logits = _censor_disallowed_logits(logits, allowed_token_ids_mask) + logits = _censor_stop_logits( + logits, + state, + min_tokens, + stop_ids=stop_ids, + eos_window_masked=eos_window_masked, ) + return _apply_logit_bias(logits, logit_bias_ids, logit_bias_values) + + +def _repetition_penalty_scaling( + logits: torch.Tensor, + state: TalkerCodecDeviceState, + repetition_penalty: torch.Tensor, +) -> torch.Tensor: + """Repetition penalty over the codec history window (engine step 6a). + + ``alpha = penalty ** frequency`` with the sign-dependent branch, matching the + model's single-frame ``_apply_batched_repetition_penalty``. + """ + positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=logits.device).reshape(1, -1) + valid = positions < state.history_len.reshape(1, 1) + safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) + counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=logits.device) + counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) + alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) + return torch.where(logits < 0, logits * alpha, logits / alpha) + + +def _apply_stream_penalties( + logits: torch.Tensor, + state: TalkerCodecDeviceState, + frequency_penalty: float, + presence_penalty: float, +) -> torch.Tensor: + """Whole-stream frequency and presence penalties (engine steps 6b/6c). + + Straight from ``model_executor/layers/utils.py:86-87``: + + logits -= frequency_penalties * output_bin_counts + logits -= presence_penalties * output_mask + + ``output_bin_counts`` is the occurrence count over the whole prompt+output + sequence and ``output_mask`` is its ``> 0`` image; both come from the + request's stream record. Skipped entirely when the request set neither, or + when the record is absent (nothing to subtract from). + """ + if state.full_bins is None: + return logits + if not frequency_penalty and not presence_penalty: + return logits + bins = state.full_bins.reshape(1, -1) + if frequency_penalty: + logits = logits - bins * frequency_penalty + if presence_penalty: + logits = logits - (bins > 0).to(logits.dtype) * presence_penalty + return logits def prepare_codec_logits( @@ -67,37 +241,53 @@ def prepare_codec_logits( eos_token_id: int, top_k: int, top_p: float, + min_p: float = 0.0, + min_tokens_stop_ids: tuple[int, ...] = (), + allowed_token_ids_mask: torch.Tensor | None = None, + logit_bias_ids: tuple[int, ...] = (), + logit_bias_values: tuple[float, ...] = (), + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, eos_window_masked: bool = False, ) -> torch.Tensor: """Prepare and filter logits without changing native NPU RNG semantics. - ``eos_window_masked`` masks the codec EOS for this frame regardless of the - step counter -- the duplex turn-end drain window (a few frames after each - cadence boundary), which the ``step < min_tokens`` criterion cannot express. + ``min_p`` mirrors the engine's argmax-invariant ``MinPLogitsProcessor``, + which runs after temperature and before top-k/top-p (``builtin.py:102-113``). + ``min_tokens_stop_ids`` widens the censored stop set beyond the codec EOS; + ``allowed_token_ids_mask`` applies the request's whitelist; + ``logit_bias_ids``/``logit_bias_values`` add the request's per-id biases; + ``frequency_penalty``/``presence_penalty`` subtract the whole-stream counts. + Every one of them is a no-op at its default value. """ + stop_ids = _stop_id_tuple(eos_token_id, min_tokens_stop_ids) logits = raw_logits.float() - # Engine step 5a: the stop/EOS mask is applied before the penalties. - eos = logits[..., eos_token_id : eos_token_id + 1] - if eos_window_masked: - eos = torch.full_like(eos, float("-inf")) - else: - eos = torch.where( - (state.step < min_tokens).reshape(1, 1), - torch.full_like(eos, float("-inf")), - eos, - ) - logits = torch.cat([logits[..., :eos_token_id], eos, logits[..., eos_token_id + 1 :]], dim=-1) + # Engine steps 5a / 3-4 / 5b: whitelist, stop censoring and bias, in the + # order the engine applies them, before any penalty. + logits = _filter_codec_logits( + logits, + state, + min_tokens, + stop_ids=stop_ids, + allowed_token_ids_mask=allowed_token_ids_mask, + logit_bias_ids=logit_bias_ids, + logit_bias_values=logit_bias_values, + eos_window_masked=eos_window_masked, + ) # Engine step 6a: repetition penalty over the codec history window. - positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) - valid = positions < state.history_len.reshape(1, 1) - safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) - counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=raw_logits.device) - counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) - alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) - logits = torch.where(logits < 0, logits * alpha, logits / alpha) + logits = _repetition_penalty_scaling(logits, state, repetition_penalty) + # Engine steps 6b/6c: frequency and presence penalties over the whole stream. + logits = _apply_stream_penalties(logits, state, frequency_penalty, presence_penalty) # Engine step 7b: temperature is applied after the penalties. The division # also gives the in-place masks below a tensor of their own. logits = logits / temperature.reshape(1, 1) + # Engine step 7c: min_p. The engine runs it as an argmax-invariant + # processor, i.e. after temperature and before top-k/top-p, and the greedy + # branch never sees it -- which is why greedy_codec_sample has no min_p. + if min_p > 0.0: + probabilities = torch.softmax(logits, dim=-1) + min_p_threshold = probabilities.amax(dim=-1, keepdim=True) * min_p + logits.masked_fill_(probabilities < min_p_threshold, float("-inf")) # Engine step 7d: top-k first (topk_topp_sampler.py:392). if top_k > 0: keep = min(VOCAB_SIZE, int(top_k)) @@ -115,17 +305,45 @@ def prepare_codec_logits( return logits +def _advance_full_bins( + state: TalkerCodecDeviceState, + sampled: torch.Tensor, + emit: torch.Tensor, +) -> None: + """Add one occurrence for the emitted id, in place. + + In-place on purpose: re-allocating a VOCAB_SIZE-wide tensor every frame would + be an allocation per decode step. The tensor belongs to the request's state + record, which is replaced every frame, so mutating it cannot be observed by + an earlier state. + """ + if state.full_bins is None: + return + state.full_bins.scatter_add_( + 0, + sampled.reshape(-1).to(torch.long), + emit.reshape(-1).to(state.full_bins.dtype), + ) + + def advance_codec_device_state( state: TalkerCodecDeviceState, sampled_token: torch.Tensor, *, eos_token_id: int, + ignore_eos: bool = False, ) -> TalkerCodecDeviceState: """Advance persistent stochastic state with graph-capturable tensor ops.""" sampled = sampled_token.reshape(1).to(torch.int32) active = ~state.finished next_step = torch.where(active, state.step + 1, state.step) - is_eos = sampled == int(eos_token_id) + if ignore_eos: + # ``ignore_eos`` blanks the engine's ``_eos_token_id`` + # (vllm/sampling_params.py:670-671), so a sampled codec EOS is just + # another frame and must not finish the request. + is_eos = torch.zeros_like(state.finished) + else: + is_eos = sampled == int(eos_token_id) reached_limit = active & (next_step >= state.max_tokens) finished = state.finished | is_eos | reached_limit emit = active & (~is_eos) & (~reached_limit) @@ -135,7 +353,10 @@ def advance_codec_device_state( candidate = torch.where((state.history_len >= HISTORY_WINDOW).reshape(1, 1), shifted, appended) history = torch.where(emit.reshape(1, 1), candidate, state.history) history_len = torch.clamp(state.history_len + emit.to(torch.int32), max=HISTORY_WINDOW) - return TalkerCodecDeviceState(history, history_len, next_step, state.max_tokens, finished) + _advance_full_bins(state, sampled, emit) + return TalkerCodecDeviceState( + history, history_len, next_step, state.max_tokens, finished, state.full_bins + ) def codec_sample_result( @@ -143,13 +364,23 @@ def codec_sample_result( sampled_token: torch.Tensor, *, eos_token_id: int, + ignore_eos: bool = False, ) -> TalkerCodecSampleResult: """Advance device state and retain the device-side emit decision.""" sampled = sampled_token.reshape(1).to(torch.int32) active = ~state.finished next_step = torch.where(active, state.step + 1, state.step) - emit = active & (sampled != int(eos_token_id)) & (next_step < state.max_tokens) - next_state = advance_codec_device_state(state, sampled, eos_token_id=eos_token_id) + if ignore_eos: + is_eos = torch.zeros_like(state.finished) + else: + is_eos = sampled == int(eos_token_id) + emit = active & (~is_eos) & (next_step < state.max_tokens) + next_state = advance_codec_device_state( + state, + sampled, + eos_token_id=eos_token_id, + ignore_eos=ignore_eos, + ) return TalkerCodecSampleResult(sampled_token=sampled, state=next_state, emit=emit) @@ -160,43 +391,52 @@ def greedy_codec_sample( repetition_penalty: torch.Tensor, *, eos_token_id: int, + min_tokens_stop_ids: tuple[int, ...] = (), + allowed_token_ids_mask: torch.Tensor | None = None, + logit_bias_ids: tuple[int, ...] = (), + logit_bias_values: tuple[float, ...] = (), + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + ignore_eos: bool = False, eos_window_masked: bool = False, ) -> TalkerCodecSampleResult: """Greedy codec sample in graph-capturable tensor ops. - The engine's greedy branch returns the argmax after the stop/EOS mask and - the penalties and skips temperature, top-k and top-p entirely - (``vllm/v1/sample/sampler.py:30-60``), so no top-k mask belongs on this - path -- masking candidates below the argmax cannot change the argmax itself. + The engine's greedy branch returns the argmax after the censoring masks, the + logit bias and the penalties, and skips temperature, min-p, top-k and top-p + entirely (``vllm/v1/sample/sampler.py:30-60``), so no top-k mask belongs on + this path -- masking candidates below the argmax cannot change the argmax + itself. Every other filter is the same one the random path uses. ``eos_window_masked`` masks the codec EOS for this frame regardless of the step counter -- see prepare_codec_logits. """ + stop_ids = _stop_id_tuple(eos_token_id, min_tokens_stop_ids) penalized = raw_logits.float() - # Engine step 5a: the stop/EOS mask is applied before the penalties. - eos = penalized[..., eos_token_id : eos_token_id + 1] - if eos_window_masked: - eos = torch.full_like(eos, float("-inf")) - else: - eos = torch.where( - (state.step < min_tokens).reshape(1, 1), - torch.full_like(eos, float("-inf")), - eos, - ) - penalized = torch.cat([penalized[..., :eos_token_id], eos, penalized[..., eos_token_id + 1 :]], dim=-1) - # Engine step 6a: repetition penalty over the codec history window. - positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=raw_logits.device).reshape(1, -1) - valid = positions < state.history_len.reshape(1, 1) - safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) - counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=raw_logits.device) - counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) - alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) - penalized = torch.where(penalized < 0, penalized * alpha, penalized / alpha) + # Engine steps 5a / 3-4 / 5b: same masks as the random path. + penalized = _filter_codec_logits( + penalized, + state, + min_tokens, + stop_ids=stop_ids, + allowed_token_ids_mask=allowed_token_ids_mask, + logit_bias_ids=logit_bias_ids, + logit_bias_values=logit_bias_values, + eos_window_masked=eos_window_masked, + ) + # Engine steps 6a/6b/6c: same penalties as the random path. + penalized = _repetition_penalty_scaling(penalized, state, repetition_penalty) + penalized = _apply_stream_penalties(penalized, state, frequency_penalty, presence_penalty) sampled = torch.argmax(penalized, dim=-1).to(torch.int32) active = ~state.finished next_step = torch.where(active, state.step + 1, state.step) - is_eos = sampled.to(torch.int64) == int(eos_token_id) + if ignore_eos: + # Same rule as advance_codec_device_state: ``ignore_eos`` disables the + # engine's EOS stop, so it disables this one too. + is_eos = torch.zeros_like(state.finished) + else: + is_eos = sampled.to(torch.int64) == int(eos_token_id) reached_limit = active & (next_step >= state.max_tokens) finished = state.finished | is_eos | reached_limit emit = active & (~is_eos) & (~reached_limit) @@ -207,8 +447,11 @@ def greedy_codec_sample( candidate = torch.where((state.history_len >= HISTORY_WINDOW).unsqueeze(1), shifted, appended) history = torch.where(emit.unsqueeze(1), candidate, state.history) history_len = torch.clamp(state.history_len + emit.to(torch.int32), max=HISTORY_WINDOW) + _advance_full_bins(state, sampled, emit) return TalkerCodecSampleResult( sampled_token=sampled, - state=TalkerCodecDeviceState(history, history_len, next_step, state.max_tokens, finished), + state=TalkerCodecDeviceState( + history, history_len, next_step, state.max_tokens, finished, state.full_bins + ), emit=emit, ) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index d047b479678..fe15d5f9915 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1446,18 +1446,36 @@ def _build_model_kwargs_extra(self) -> dict: # engine stops accepting tokens: the scheduler truncates the # sampled ids at the request limit, but the connector still # concatenates every emitted codec frame (PR #7929 review). - # None means the request carries no max_tokens, which leaves - # the stage-resolved codec budget in charge. Same order as - # request_sampling_params. + # Both of the engine's length stops are folded in + # (v1/core/sched/utils.py:111-117): the request's + # max_tokens - num_output_tokens, and the context's + # max_model_len - num_tokens. None means the request carries + # no max_tokens, which leaves the stage-resolved codec budget + # in charge. Same order as request_sampling_params. + context_budget = [ + ( + self.max_model_len - int(req.num_tokens) + if (req := self.requests.get(req_id)) is not None + else None + ) + for req_id in self.input_batch.req_ids + ] model_kwargs_extra["request_max_tokens_remaining"] = [ ( - max(int(req.sampling_params.max_tokens) - int(req.num_output_tokens), 0) + max( + min( + int(req.sampling_params.max_tokens) - int(req.num_output_tokens), + context_budget[index], + ), + 0, + ) if (req := self.requests.get(req_id)) is not None and req.sampling_params is not None and req.sampling_params.max_tokens is not None + and context_budget[index] is not None else None ) - for req_id in self.input_batch.req_ids + for index, req_id in enumerate(self.input_batch.req_ids) ] except Exception as e: # Visible on purpose: the fallback is the equal rows-per-request From 2f74c093cbe93b97004a7fef7115906888160774 Mon Sep 17 00:00:00 2001 From: y-null Date: Fri, 25 Sep 2026 06:50:34 +0000 Subject: [PATCH 57/89] style: satisfy ruff in the MiniCPM-o Talker codec path Signed-off-by: y-null --- .../models/minicpmo_4_5/minicpmo_4_5_omni_tts.py | 14 ++++---------- .../models/minicpmo_4_5/talker_codec_sample.py | 8 ++------ 2 files changed, 6 insertions(+), 16 deletions(-) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 5c9ddc09a34..3b7a8381388 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -35,6 +35,8 @@ ) from vllm_omni.model_executor.models.minicpmo_4_5.talker_codec_sample import ( VOCAB_SIZE as _CODEC_VOCAB_SIZE, +) +from vllm_omni.model_executor.models.minicpmo_4_5.talker_codec_sample import ( TalkerCodecDeviceState, TalkerCodecSampleResult, codec_sample_result, @@ -278,9 +280,7 @@ def _codec_stop_ids(state: Any) -> tuple[int, ...]: value = state.get("codec_stop_token_ids") if isinstance(state, Mapping) else None if not value: return () - return tuple( - sorted({int(token_id) for token_id in value if 0 <= int(token_id) < _CODEC_VOCAB_SIZE}) - ) + return tuple(sorted({int(token_id) for token_id in value if 0 <= int(token_id) < _CODEC_VOCAB_SIZE})) def _codec_allowed_ids_mask(state: Any, device: torch.device) -> torch.Tensor | None: @@ -315,9 +315,7 @@ def _codec_logit_bias_items(state: Any) -> tuple[tuple[int, ...], tuple[float, . if not isinstance(value, Mapping): return (), () pairs = sorted( - (int(token_id), float(bias)) - for token_id, bias in value.items() - if 0 <= int(token_id) < _CODEC_VOCAB_SIZE + (int(token_id), float(bias)) for token_id, bias in value.items() if 0 <= int(token_id) < _CODEC_VOCAB_SIZE ) if not pairs: return (), () @@ -402,10 +400,6 @@ def _codec_repetition_detected(state: Any) -> bool: return False - - - - def _codec_int_param(state: Any, key: str, fallback: int) -> int: """Read an integer codec knob from a request state, None meaning unset. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index a48d2459514..3bef0c9a6f5 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -354,9 +354,7 @@ def advance_codec_device_state( history = torch.where(emit.reshape(1, 1), candidate, state.history) history_len = torch.clamp(state.history_len + emit.to(torch.int32), max=HISTORY_WINDOW) _advance_full_bins(state, sampled, emit) - return TalkerCodecDeviceState( - history, history_len, next_step, state.max_tokens, finished, state.full_bins - ) + return TalkerCodecDeviceState(history, history_len, next_step, state.max_tokens, finished, state.full_bins) def codec_sample_result( @@ -450,8 +448,6 @@ def greedy_codec_sample( _advance_full_bins(state, sampled, emit) return TalkerCodecSampleResult( sampled_token=sampled, - state=TalkerCodecDeviceState( - history, history_len, next_step, state.max_tokens, finished, state.full_bins - ), + state=TalkerCodecDeviceState(history, history_len, next_step, state.max_tokens, finished, state.full_bins), emit=emit, ) From 86b312b32693ba262c25a118d4def24cf2a40993 Mon Sep 17 00:00:00 2001 From: y-null Date: Fri, 25 Sep 2026 15:35:12 +0800 Subject: [PATCH 58/89] fix(minicpmo_4_5): drop the K-step allowed_token_ids / logit_bias filters Both parameters are keyed by the text vocabulary but the engines index their own narrow tensors with those ids: InputBatch sizes the allowed-ids mask by model_config.get_vocab_size() (0 for the Talker stage, raised to the 2-wide stop row on the K-step path) and LogitBiasLogitsProcessor indexes the logits columns directly, so any codec-space id raises IndexError inside the engine on both the single-frame and the multi-frame paths. Mirroring them in the K-step sampler cannot make them usable, so drop the filters and keep the K-step path equal to the engine by not implementing them. Signed-off-by: y-null --- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 65 +---------------- .../minicpmo_4_5/talker_codec_sample.py | 69 ++----------------- 2 files changed, 8 insertions(+), 126 deletions(-) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 3b7a8381388..f6fea189d85 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -269,9 +269,8 @@ def _codec_stop_ids(state: Any) -> tuple[int, ...]: ``model_config.get_vocab_size()`` before the request is accepted (``vllm/sampling_params.py:874-895``), so surviving ids are in-vocabulary for the model -- but the codec head is only ``_CODEC_VOCAB_SIZE`` wide, so - anything at or beyond it is dropped here, the same way - ``_codec_allowed_ids_mask`` and ``_codec_logit_bias_items`` drop theirs, - which keeps the index tensors in range. The engine censors and stops on + anything at or beyond it is dropped here, which keeps the index tensor in + range. The engine censors and stops on ``all_stop_token_ids`` = eos + these ids + extra eos ids (``builtin.py:196-207``, ``v1/core/sched/utils.py:105``); K-step has no engine layer left, so it censors and stops on them itself. Empty means the @@ -283,62 +282,6 @@ def _codec_stop_ids(state: Any) -> tuple[int, ...]: return tuple(sorted({int(token_id) for token_id in value if 0 <= int(token_id) < _CODEC_VOCAB_SIZE})) -def _codec_allowed_ids_mask(state: Any, device: torch.device) -> torch.Tensor | None: - """Censor mask for the request's codec whitelist, or None when unset. - - The engine builds ``allowed_token_ids_mask`` with True on the ids to - censor (``v1/worker/gpu_input_batch.py:463-466``) and applies it as - ``masked_fill_(-inf)`` (``sampler.py:393-394``). The request's list is - keyed by the text vocabulary, so ids outside the codec vocabulary are - dropped; a list with nothing left in range is treated as unset rather - than producing an all ``-inf`` row. - """ - value = state.get("codec_allowed_token_ids") if isinstance(state, Mapping) else None - if not value: - return None - keep = sorted({int(token_id) for token_id in value if 0 <= int(token_id) < _CODEC_VOCAB_SIZE}) - if not keep: - return None - mask = torch.ones((1, _CODEC_VOCAB_SIZE), dtype=torch.bool, device=device) - mask[0, torch.tensor(keep, dtype=torch.long, device=device)] = False - return mask - - -def _codec_logit_bias_items(state: Any) -> tuple[tuple[int, ...], tuple[float, ...]]: - """The request's codec logit biases as parallel tuples, in a stable order. - - ``LogitBiasLogitsProcessor`` adds ``bias`` to ``logits[req, tok]`` - (``builtin.py:161-164``). The request's dict is keyed by the text - vocabulary, so keys outside the codec vocabulary are dropped. - """ - value = state.get("codec_logit_bias") if isinstance(state, Mapping) else None - if not isinstance(value, Mapping): - return (), () - pairs = sorted( - (int(token_id), float(bias)) for token_id, bias in value.items() if 0 <= int(token_id) < _CODEC_VOCAB_SIZE - ) - if not pairs: - return (), () - return tuple(token_id for token_id, _ in pairs), tuple(bias for _, bias in pairs) - - -def _codec_filter_inputs(state: Any, device: torch.device) -> dict[str, Any]: - """Whitelist mask and logit-bias tuples for the codec filters. - - Returns an empty mapping when the request set neither, so the call site - keeps the sampler's no-op defaults. - """ - inputs: dict[str, Any] = {} - mask = _codec_allowed_ids_mask(state, device) - if mask is not None: - inputs["allowed_token_ids_mask"] = mask - bias_ids, bias_values = _codec_logit_bias_items(state) - if bias_ids: - inputs["logit_bias_ids"] = bias_ids - inputs["logit_bias_values"] = bias_values - return inputs - - def _codec_full_bins_input(state: Any, device: torch.device) -> torch.Tensor | None: """A zeroed whole-stream counter when the request uses the stream penalties. @@ -1368,8 +1311,6 @@ def _merge_request_codec_params( ("codec_seed", "seed"), ("codec_min_p", "min_p"), ("codec_ignore_eos", "ignore_eos"), - ("codec_allowed_token_ids", "allowed_token_ids"), - ("codec_logit_bias", "logit_bias"), ("codec_frequency_penalty", "frequency_penalty"), ("codec_presence_penalty", "presence_penalty"), ): @@ -1505,7 +1446,6 @@ def _sample_audio_code( top_p=_codec_float_param(request_state, "codec_top_p", self._codec_top_p), min_p=_codec_float_param(request_state, "codec_min_p", 0.0), min_tokens_stop_ids=_codec_stop_ids(request_state), - **_codec_filter_inputs(request_state, hidden_state.device), frequency_penalty=_codec_float_param(request_state, "codec_frequency_penalty", 0.0), presence_penalty=_codec_float_param(request_state, "codec_presence_penalty", 0.0), eos_window_masked=eos_window_masked, @@ -1586,7 +1526,6 @@ def _sample_audio_code_greedy( penalty_tensor, eos_token_id=self._codec_eos_id, min_tokens_stop_ids=_codec_stop_ids(request_state), - **_codec_filter_inputs(request_state, hidden_state.device), frequency_penalty=_codec_float_param(request_state, "codec_frequency_penalty", 0.0), presence_penalty=_codec_float_param(request_state, "codec_presence_penalty", 0.0), ignore_eos=_codec_bool_param(request_state, "codec_ignore_eos", False), diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index 3bef0c9a6f5..c1a75165019 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -118,68 +118,26 @@ def _censor_stop_logits( return logits.index_copy(-1, stop_index, censored) -def _censor_disallowed_logits( - logits: torch.Tensor, - allowed_token_ids_mask: torch.Tensor | None, -) -> torch.Tensor: - """Blank every id the request's whitelist excludes. - - The engine precomputes a boolean mask whose True entries are the ids to - censor (``v1/worker/gpu_input_batch.py:463-466``) and applies it as - ``logits.masked_fill_(mask, -inf)`` (``sampler.py:393-394``). ``None`` means - the request set no whitelist, which is the existing behaviour. - """ - if allowed_token_ids_mask is None: - return logits - return logits.masked_fill(allowed_token_ids_mask, float("-inf")) - - -def _apply_logit_bias( - logits: torch.Tensor, - logit_bias_ids: tuple[int, ...], - logit_bias_values: tuple[float, ...], -) -> torch.Tensor: - """Add the request's per-id biases. - - Mirrors ``LogitBiasLogitsProcessor.apply`` (``builtin.py:161-164``), which - the engine runs after the censoring masks and before the penalties, so a - censored id stays censored. Empty ids means the request set no bias. - """ - if not logit_bias_ids: - return logits - bias_index = torch.tensor(logit_bias_ids, dtype=torch.long, device=logits.device) - bias_values = torch.tensor(logit_bias_values, dtype=logits.dtype, device=logits.device) - current = logits.index_select(-1, bias_index) - return logits.index_copy(-1, bias_index, current + bias_values) - - def _filter_codec_logits( logits: torch.Tensor, state: TalkerCodecDeviceState, min_tokens: torch.Tensor, *, stop_ids: tuple[int, ...], - allowed_token_ids_mask: torch.Tensor | None, - logit_bias_ids: tuple[int, ...], - logit_bias_values: tuple[float, ...], eos_window_masked: bool, ) -> torch.Tensor: - """The masks both sampling paths share, in the engine's order. + """The stop censoring both sampling paths share. - ``allowed ids -> stop censoring -> logit bias`` are all applied by the engine - before its penalties (``sampler.py:393-405``); the bias is last because it is - additive while the other two write ``-inf``, so a censored id must not be - revived by a bias. + The engine blanks every ``all_stop_token_ids`` entry before its penalties + (``sampler.py:393-405``), so the K-step path does the same here. """ - logits = _censor_disallowed_logits(logits, allowed_token_ids_mask) - logits = _censor_stop_logits( + return _censor_stop_logits( logits, state, min_tokens, stop_ids=stop_ids, eos_window_masked=eos_window_masked, ) - return _apply_logit_bias(logits, logit_bias_ids, logit_bias_values) def _repetition_penalty_scaling( @@ -243,9 +201,6 @@ def prepare_codec_logits( top_p: float, min_p: float = 0.0, min_tokens_stop_ids: tuple[int, ...] = (), - allowed_token_ids_mask: torch.Tensor | None = None, - logit_bias_ids: tuple[int, ...] = (), - logit_bias_values: tuple[float, ...] = (), frequency_penalty: float = 0.0, presence_penalty: float = 0.0, eos_window_masked: bool = False, @@ -255,23 +210,17 @@ def prepare_codec_logits( ``min_p`` mirrors the engine's argmax-invariant ``MinPLogitsProcessor``, which runs after temperature and before top-k/top-p (``builtin.py:102-113``). ``min_tokens_stop_ids`` widens the censored stop set beyond the codec EOS; - ``allowed_token_ids_mask`` applies the request's whitelist; - ``logit_bias_ids``/``logit_bias_values`` add the request's per-id biases; ``frequency_penalty``/``presence_penalty`` subtract the whole-stream counts. Every one of them is a no-op at its default value. """ stop_ids = _stop_id_tuple(eos_token_id, min_tokens_stop_ids) logits = raw_logits.float() - # Engine steps 5a / 3-4 / 5b: whitelist, stop censoring and bias, in the - # order the engine applies them, before any penalty. + # Engine step 3-4: stop censoring, before any penalty. logits = _filter_codec_logits( logits, state, min_tokens, stop_ids=stop_ids, - allowed_token_ids_mask=allowed_token_ids_mask, - logit_bias_ids=logit_bias_ids, - logit_bias_values=logit_bias_values, eos_window_masked=eos_window_masked, ) # Engine step 6a: repetition penalty over the codec history window. @@ -390,9 +339,6 @@ def greedy_codec_sample( *, eos_token_id: int, min_tokens_stop_ids: tuple[int, ...] = (), - allowed_token_ids_mask: torch.Tensor | None = None, - logit_bias_ids: tuple[int, ...] = (), - logit_bias_values: tuple[float, ...] = (), frequency_penalty: float = 0.0, presence_penalty: float = 0.0, ignore_eos: bool = False, @@ -411,15 +357,12 @@ def greedy_codec_sample( """ stop_ids = _stop_id_tuple(eos_token_id, min_tokens_stop_ids) penalized = raw_logits.float() - # Engine steps 5a / 3-4 / 5b: same masks as the random path. + # Engine step 3-4: same stop censoring as the random path. penalized = _filter_codec_logits( penalized, state, min_tokens, stop_ids=stop_ids, - allowed_token_ids_mask=allowed_token_ids_mask, - logit_bias_ids=logit_bias_ids, - logit_bias_values=logit_bias_values, eos_window_masked=eos_window_masked, ) # Engine steps 6a/6b/6c: same penalties as the random path. From 7c28f57e9462cc490bb800fe99476c7ac457d83d Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 28 Sep 2026 04:18:01 +0000 Subject: [PATCH 59/89] Fix two review findings in the K-step request param merge 1. Control-head stop ids no longer join the codec censor set. The default NPU K=8 profile merges request stop_token_ids to [1, 6561]: 1 finishes the request on the two-wide continue/stop control row, 6561 is the codec EOS. Copying the list verbatim into codec_stop_token_ids made the sampler censor codec id 1 below the min_tokens floor and treat a sampled 1 as EOS above it -- both wrong in the 6562-wide codec vocabulary where 1 is an ordinary audio code. The merge now drops _SCHEDULER_CONTROL_STOP_IDS (0, 1) when the multi-frame path is armed; the engine-side list stays untouched. 2. An explicit request min_tokens=0 now disables the K-step floor. The merge pinned only positive values, so a caller override of 0 (turn the floor off, as the single-frame path honours) fell through to the stage-resolved minimum. The absent vs explicitly-set-to-0 cases are told apart via pydantic model_fields_set, which survives the engine clone() deepcopy; stub objects without the bookkeeping keep the old positive-only rule. Tests: test_request_min_tokens_reaches_codec_state now pins an explicit 0 and asserts the resolved floor reads 0; new test_control_stop_ids_stay_out_of_codec_censor_set covers the merge, the non-armed legacy copy, and an end-to-end stream where a sampled codec 1 is emitted without terminating while a later codec EOS still does. 25 passed. pre-commit env is uninstallable offline; ruff check+format run manually, both green. Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 83 ++++++++++++++++++- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 45 ++++++++-- 2 files changed, 119 insertions(+), 9 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 3a186becaac..9e394b3cac4 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -27,6 +27,7 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( MiniCPMO45OmniTTSForConditionalGeneration, + _codec_int_param, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -35,6 +36,17 @@ _EOS_ID = _NUM_AUDIO_TOKENS - 1 +class _FakeStageParams: + """Stands in for a real pydantic ``SamplingParams``: like pydantic v2 it + records the explicitly-set fields in ``model_fields_set``, which the + request merge reads to tell "absent" apart from "explicitly set to 0".""" + + def __init__(self, **kwargs): + self.model_fields_set = set(kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): """Bare talker instance: no config, no weights, deterministic sampler. @@ -647,7 +659,9 @@ def test_request_min_tokens_reaches_codec_state(): The NPU runner neutralizes vLLM's MinTokensLogitsProcessor, so the in-model codec sampler is the only min-length guard left (PR #7929 review). The engine default 0 means "no extra floor" and must not - overwrite the stage-resolved codec minimum. + overwrite the stage-resolved codec minimum -- but an explicit + request-level 0 disables the floor, matching the single-frame path + where MinTokensLogitsProcessor.add_request skips ``min_tokens=0``. """ model = _make_talker(k_step_frames=8, scripted_samples=[42]) model._codec_min_tokens = 50 @@ -657,18 +671,81 @@ def test_request_min_tokens_reaches_codec_state(): model._merge_request_codec_params(state, SimpleNamespace(min_tokens=100, **base)) assert state["min_tokens"] == 100 - # 0 (engine default) is not a floor: the key stays unset so the K-step - # loop keeps falling back to the stage-resolved minimum. + # 0 without pydantic bookkeeping (the engine default flowing through a + # stub) is not a floor: the key stays unset so the K-step loop keeps + # falling back to the stage-resolved minimum. state_zero = {"step": 0} model._merge_request_codec_params(state_zero, SimpleNamespace(min_tokens=0, **base)) assert "min_tokens" not in state_zero + # An explicitly set 0 (recorded in model_fields_set, as a real pydantic + # SamplingParams would carry after a caller override) pins 0 and turns + # the floor off, mirroring the single-frame semantics. + state_explicit = {"step": 0} + model._merge_request_codec_params(state_explicit, _FakeStageParams(min_tokens=0, **base)) + assert state_explicit["min_tokens"] == 0 + assert _codec_int_param(state_explicit, "min_tokens", 50) == 0 + # The min_new_tokens alias is honored too. state_alias = {"step": 0} model._merge_request_codec_params(state_alias, SimpleNamespace(min_new_tokens=7, **base)) assert state_alias["min_tokens"] == 7 +def test_control_stop_ids_stay_out_of_codec_censor_set(): + """The scheduler's control-head ids never join the codec censor set. + + Under multi-frame decode the default NPU K=8 profile merges the request + list to [1, 6561]: 1 finishes the request on the two-wide continue/stop + control row, 6561 is the codec EOS. Inside the 6562-wide codec + vocabulary 1 is an ordinary audio code -- censoring it would skew the + below-floor distribution and truncate real speech once sampled. The + single-frame path (no stop_token_ids: [1] block) never carries it. + """ + model = _make_talker(k_step_frames=8, scripted_samples=[1, 42, _EOS_ID]) + model._codec_min_tokens = 50 + + # The merge drops the control ids and keeps the codec-side stop ids. + state = {"step": 0} + model._merge_request_codec_params(state, SimpleNamespace(stop_token_ids=[1, 6561], min_tokens=50)) + assert state["codec_stop_token_ids"] == [6561] + + # The drop is gated on the multi-frame arm: a non-armed deployment keeps + # the legacy verbatim copy (the control-row contract does not apply). + model_legacy = _make_talker(k_step_frames=1, scripted_samples=[1]) + state_legacy = {"step": 0} + model_legacy._merge_request_codec_params(state_legacy, SimpleNamespace(stop_token_ids=[1, 6561], min_tokens=50)) + assert state_legacy["codec_stop_token_ids"] == [1, 6561] + + # End to end across the sampling boundary: a codec 1 sampled past the + # floor is emitted as ordinary audio and does not finish the request; + # a later codec EOS still terminates it normally. + state_stream = {"step": 0} + model._request_audio_states["r1"] = state_stream + params = SimpleNamespace(stop_token_ids=[1, 6561], min_tokens=50) + + def _frame(): + return model.make_omni_output( + torch.randn(1, 8), + model_intermediate_buffer=[{"request_id": "r1"}], + request_token_spans=[(0, 1)], + request_sample_eligible=[True], + request_sampling_params=[params], + ) + + out0 = _frame() + assert out0.multimodal_outputs["codes"]["audio"][0].reshape(-1).tolist() == [1] + assert not state_stream.get("finished") + + out1 = _frame() + assert out1.multimodal_outputs["codes"]["audio"][0].reshape(-1).tolist() == [42] + assert not state_stream.get("finished") + + out2 = _frame() + assert out2.multimodal_outputs["codes"]["audio"][0].numel() == 0 + assert state_stream["finished"] is True + + def test_top_k_then_top_p_matches_single_frame_order(): """K-step filtering follows the engine's order: top-k, then top-p. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index f6fea189d85..830552e9e0f 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -61,6 +61,13 @@ _CODEC_REPETITION_PENALTY = 1.05 _CODEC_MIN_TOKENS = 50 _CODEC_MAX_TOKENS = 2048 +# The Talker's two-wide stop head emits continue/stop control rows under +# multi-frame decode: 0=continue, 1=stop. Those ids belong to the +# scheduler's control channel, not the 6562-wide codec stream, where both +# values are ordinary audio codes. The engine needs the stop id inside the +# request's stop_token_ids to finish the request on a control row; the +# codec censor set must not inherit it. +_SCHEDULER_CONTROL_STOP_IDS = (0, 1) # YAML key -> (tts_config attribute, hardcoded fallback, type) _CODEC_SAMPLING_SOURCES: tuple[tuple[str, str, Any, Any], ...] = ( ("seed", "seed", _CODEC_SEED, int), @@ -1326,7 +1333,22 @@ def _merge_request_codec_params( if state.get("codec_stop_token_ids") is None: requested_stop_ids = getattr(sampling_params, "stop_token_ids", None) if requested_stop_ids: - state["codec_stop_token_ids"] = [int(token_id) for token_id in requested_stop_ids] + stop_ids = [int(token_id) for token_id in requested_stop_ids] + if self._k_step_frames >= 2: + # Under multi-frame decode the merged request list mixes + # two channels: the two-wide control head's stop id (the + # yaml platforms.npu stop_token_ids block, which the + # scheduler needs to finish the request on a control row) + # and ids acting on the codec stream itself (the merged + # codec EOS 6561). In the 6562-wide codec vocabulary the + # control ids are ordinary audio codes, so censoring them + # would skew the below-floor distribution and truncate + # real speech once sampled; single-frame deployments + # never carry them. Drop them from the codec set only -- + # the engine-side list stays untouched. + stop_ids = [token_id for token_id in stop_ids if token_id not in _SCHEDULER_CONTROL_STOP_IDS] + if stop_ids: + state["codec_stop_token_ids"] = stop_ids # repetition_detection is a whole-stream pattern test # (v1/core/sched/utils.py:28-59), so its record has to keep every # emitted codec id. Started here, appended per frame below. @@ -1337,15 +1359,26 @@ def _merge_request_codec_params( state.setdefault("codec_full_ids", []) # The K-step codec sampler is the only min-length guard left (the NPU # runner neutralizes vLLM's MinTokensLogitsProcessor), so the request - # floor must reach it. Pin only a positive floor: the engine default 0 - # means "no extra floor" and would otherwise overwrite the stage's - # resolved codec minimum. + # floor must reach it. The merged stage params always carry the stage + # defaults, and the engine's own registration treats a zero min_tokens + # as "no floor" (MinTokensLogitsProcessor.add_request: ``if not + # min_tokens``), so a 0 in a merged param object can only be an + # explicit caller override -- which the single-frame path honours by + # disabling the floor. Distinguish "absent" from "explicitly set to + # 0" through pydantic's model_fields_set (kept across clone(), a + # deepcopy); a positive value needs no marker, and stub objects + # without the bookkeeping keep the old positive-only rule. if state.get("min_tokens") is None: + attr = "min_tokens" requested = getattr(sampling_params, "min_tokens", None) if requested is None: + attr = "min_new_tokens" requested = getattr(sampling_params, "min_new_tokens", None) - if requested is not None and int(requested) > 0: - state["min_tokens"] = int(requested) + if requested is not None: + value = int(requested) + fields_set = getattr(sampling_params, "model_fields_set", None) + if value > 0 or (fields_set is not None and attr in fields_set): + state["min_tokens"] = value # The request's remaining output budget caps this segment's frame # ceiling: the scheduler truncates the sampled ids at the request limit # while the connector concatenates every emitted codec frame, so the From 7ad07435db160e063389403ff8d9e8b929ec298d Mon Sep 17 00:00:00 2001 From: y-null Date: Mon, 28 Sep 2026 08:18:30 +0000 Subject: [PATCH 60/89] Compat: drive vllm 0.29 and 0.30 from one omni build - omni_scheduler_mixin: vllm 0.30 moved the structured-output accept flow onto the manager (accept_tokens); on the 0.29 line it still lives inside the scheduler update_from_output, which this class overrides. Probe the manager and replay the 0.29 flow when accept_tokens is missing so the grammar advance is not silently lost. - npu_model_runner: vllm-ascend only grew the kv_cache_allocation_context parameter on the 0.30 line. Probe NPUModelRunner.initialize_kv_cache signature and forward the context only when accepted, so one omni build drives both vllm-ascend versions. Signed-off-by: y-null --- vllm_omni/core/sched/omni_scheduler_mixin.py | 22 ++++++++++++++++++- .../platforms/npu/worker/npu_model_runner.py | 10 ++++++++- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 0f17fb54da5..f803dc1f83e 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -763,7 +763,27 @@ def _attach_finished_request_sets( def _reject_invalid_grammar_tokens(self, request: Request, new_token_ids: list[int]) -> bool: """Mark rejected tokens terminal before callers capture the finish reason.""" - if not new_token_ids or self.structured_output_manager.accept_tokens(request, new_token_ids): + if not new_token_ids: + return False + # vllm 0.30 put the whole accept flow on the manager (``accept_tokens``); + # on the 0.29 line it still lives inside the scheduler's own + # ``update_from_output``, which this class overrides -- so replay that + # flow here instead of losing the grammar advance entirely. + manager = self.structured_output_manager + accept_tokens = getattr(manager, "accept_tokens", None) + if accept_tokens is not None: + accepted = accept_tokens(request, new_token_ids) + else: + accepted = True + if manager.should_advance(request, new_token_ids=new_token_ids): + structured_req = getattr(request, "structured_output_request", None) + grammar = getattr(structured_req, "grammar", None) + if grammar is not None: + advance_token_ids = manager.trim_reasoning_for_advance(request, new_token_ids) + accepted = not advance_token_ids or bool( + grammar.accept_tokens(request.request_id, advance_token_ids) + ) + if accepted: return False logger.error( "Unexpected: grammar rejected tokens %s for request %s. Terminating request.", diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 836272f35f6..7779dcdbbda 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import inspect from contextlib import AbstractContextManager from functools import partial from typing import Any @@ -41,8 +42,15 @@ def initialize_kv_cache( self, kv_cache_config, kv_cache_allocation_context: AbstractContextManager | None = None ) -> None: """Stage the omni prefix-cache config (hidden / mm tensors reused on hits).""" + # vllm-ascend's override only grew the ``kv_cache_allocation_context`` + # parameter on the 0.30 line; probe the signature so one omni build + # drives both the 0.29 and the 0.30 vllm-ascend. + params = inspect.signature(NPUModelRunner.initialize_kv_cache).parameters + accepts_ctx = "kv_cache_allocation_context" in params NPUModelRunner.initialize_kv_cache( - self, kv_cache_config, kv_cache_allocation_context=kv_cache_allocation_context + self, + kv_cache_config, + **({"kv_cache_allocation_context": kv_cache_allocation_context} if accepts_ctx else {}), ) if getattr(self, "_omni_prefix_cache_cfg", None) is None: # Same gate as the GPU runner (pooling stage, kv_consumer / From fd78f82247dbbfbb657afc760ae2039e7b7f499a Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 29 Sep 2026 01:07:31 +0000 Subject: [PATCH 61/89] fix(minicpmo): address review comments on codec sampling and duplex metadata 1. deploy/minicpmo_4_5.yaml: enable_prefix_caching must be false on the duplex stage - the duplex KV window code (window_kv.py) requires prefix caching disabled. 2. talker_codec_sample.py: greedy_codec_sample now delegates the state advance (step / finished / emit / history) to codec_sample_result instead of duplicating it, so the greedy and random paths share one state machine and cannot drift apart. 3. minicpmo_4_5_omni_tts.py: forward and make_omni_output extracted native-duplex row metadata independently. Both now call the shared _native_duplex_row_metadata(), which also validates that epoch and turn_id are non-negative ints instead of silently passing garbage. Validation: tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py + duplex/test_window_wiring.py + duplex/test_policy.py - 51 passed. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 3 +- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 179 ++++++------------ .../minicpmo_4_5/talker_codec_sample.py | 33 +--- 3 files changed, 70 insertions(+), 145 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index d5bde22b660..7ca529f69af 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -136,7 +136,8 @@ platforms: max_num_seqs: 8 max_num_batched_tokens: 8192 # Moved here from the platform-shared block; see the base stage note. - enable_prefix_caching: true + # Duplex KV window (window_kv.py) requires prefix caching disabled. + enable_prefix_caching: false compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame # expansion. The static-shape builder gates its runtime buckets on diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 830552e9e0f..e259f7bb757 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -104,6 +104,44 @@ _DUPLEX_TURN_END_BOUNDARY_MASK_STEPS = 5 +def _native_duplex_row_metadata(info_dict: dict) -> tuple[bool, int, int, str, bool]: + """Extract one row's native-duplex metadata. + + Shared by forward and make_omni_output so the two extraction sites cannot + drift apart. Returns ``(native_duplex, epoch, turn_id, segment_text, + turn_end)``; ``turn_end`` is already gated on ``native_duplex``. + """ + native_duplex = info_dict.get("native_duplex") is True + duplex_info = info_dict.get("duplex") + if not isinstance(duplex_info, dict): + duplex_info = {} + epoch = duplex_info.get("epoch", -1) + turn_id = duplex_info.get("turn_id", -1) + if native_duplex and not all( + isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in (epoch, turn_id) + ): + raise RuntimeError( + "MiniCPM-o native duplex Talker requires non-negative integer " + f"epoch and turn_id, got epoch={epoch!r}, turn_id={turn_id!r}" + ) + meta_info = info_dict.get("meta") + if not isinstance(meta_info, dict): + meta_info = {} + segment_text = meta_info.get("native_duplex_segment_text", "") if native_duplex else "" + if not isinstance(segment_text, str): + segment_text = "" + turn_eos_id = meta_info.get("turn_eos_token_id") + ids_info = info_dict.get("ids") + tts_ids = ids_info.get("tts") if native_duplex and isinstance(ids_info, dict) else None + if isinstance(tts_ids, torch.Tensor): + contains_turn_eos = isinstance(turn_eos_id, int) and bool(torch.any(tts_ids.reshape(-1) == turn_eos_id).item()) + elif isinstance(tts_ids, (list, tuple)): + contains_turn_eos = isinstance(turn_eos_id, int) and turn_eos_id in tts_ids + else: + contains_turn_eos = False + return native_duplex, epoch, turn_id, segment_text, (native_duplex and contains_turn_eos) + + def _native_duplex_chunk_budget(meta: Mapping[str, Any] | None) -> tuple[int, int]: """Return ``(max_tokens, min_tokens)`` for one native-duplex Talker request.""" turn_start = isinstance(meta, Mapping) and bool(meta.get("turn_start")) @@ -861,32 +899,6 @@ def preprocess( } if retained_codes: state["recent_codes"] = retained_codes - # Rebuild the device-side penalty window once per chunk from the - # carried-over ids; the per-frame path below then appends on - # device without touching the host again. - retained = [int(code_id) for code_id in retained_codes][-_CODEC_PENALTY_WINDOW:] - if retained: - penalty_windows = getattr(self, "_penalty_windows_dev", None) - if not isinstance(penalty_windows, dict): - penalty_windows = {} - self._penalty_windows_dev = penalty_windows - penalty_windows[request_id] = torch.tensor(retained, dtype=torch.long, device=embeds.device) - pf_freq = torch.zeros( - int(getattr(self, "_num_audio_tokens", 0) or 0), - device=embeds.device, - dtype=torch.float32, - ) - if pf_freq.numel() > 0: - pf_freq.index_add_( - 0, - penalty_windows[request_id], - torch.ones(len(retained), device=embeds.device, dtype=torch.float32), - ) - pf_dict = getattr(self, "_penalty_freqs_dev", None) - if not isinstance(pf_dict, dict): - pf_dict = {} - self._penalty_freqs_dev = pf_dict - pf_dict[request_id] = pf_freq request_states = getattr(self, "_request_audio_states", None) if request_states is None: request_states = {} @@ -942,10 +954,13 @@ def preprocess( dtype=torch.long, device=self.emb_code[0].weight.device, ) + # k_last is already the host int this tensor was built from; + # .item() would only sync it right back. + code_id = k_last else: code = input_ids.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1)[-1:] + code_id = int(code.item()) embeds = self.emb_code[0](code) - code_id = int(code.item()) if code_id == int(self._codec_eos_id): if isinstance(state, dict): state["finished"] = True @@ -1013,37 +1028,8 @@ def make_omni_output( row_continue, row_stop, flag_false, flag_true, empty_delta = self._step_constants(hidden) for index, info in enumerate(infos): info_dict = info if isinstance(info, dict) else {} - native_duplex = info_dict.get("native_duplex") is True if emit_duplex_metadata: - duplex_info = info_dict.get("duplex") - if not isinstance(duplex_info, dict): - duplex_info = {} - epoch = duplex_info.get("epoch", -1) - turn_id = duplex_info.get("turn_id", -1) - if native_duplex and not all( - isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in (epoch, turn_id) - ): - raise RuntimeError( - "MiniCPM-o native duplex Talker requires non-negative integer " - f"epoch and turn_id, got epoch={epoch!r}, turn_id={turn_id!r}" - ) - meta_info = info_dict.get("meta") - if not isinstance(meta_info, dict): - meta_info = {} - segment_text = meta_info.get("native_duplex_segment_text", "") if native_duplex else "" - if not isinstance(segment_text, str): - segment_text = "" - turn_eos_id = meta_info.get("turn_eos_token_id") - ids_info = info_dict.get("ids") - tts_ids = ids_info.get("tts") if native_duplex and isinstance(ids_info, dict) else None - if isinstance(tts_ids, torch.Tensor): - contains_turn_eos = isinstance(turn_eos_id, int) and bool( - torch.any(tts_ids.reshape(-1) == turn_eos_id).item() - ) - elif isinstance(tts_ids, (list, tuple)): - contains_turn_eos = isinstance(turn_eos_id, int) and turn_eos_id in tts_ids - else: - contains_turn_eos = False + native_duplex, epoch, turn_id, segment_text, turn_end = _native_duplex_row_metadata(info_dict) native_duplex_flags.append(torch.tensor(native_duplex, dtype=torch.bool)) duplex_epochs.append(torch.tensor(epoch if isinstance(epoch, int) else -1, dtype=torch.long)) duplex_turn_ids.append(torch.tensor(turn_id if isinstance(turn_id, int) else -1, dtype=torch.long)) @@ -1053,7 +1039,7 @@ def make_omni_output( dtype=torch.uint8, ) ) - turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool)) + turn_end_flags.append(torch.tensor(turn_end, dtype=torch.bool)) if not isinstance(info, dict): stop_rows.append(row_continue) @@ -1435,8 +1421,12 @@ def _sample_audio_code( # Refresh in place: swapping the tensor would re-capture the # per-request sampling graph. budget = _codec_int_param(request_state, "max_tokens", self._codec_max_tokens) - if int(device_state.max_tokens.item()) != budget: - device_state.max_tokens.fill_(budget) + # Refresh in place, unconditionally: fill_ is a no-op when the + # budget did not move, and the .item() compare would sync every + # request every frame just to decide whether to write a value we + # can write regardless. Swapping the tensor would re-capture the + # per-request sampling graph. + device_state.max_tokens.fill_(budget) device_inputs_by_request = getattr(self, "_request_codec_device_inputs", None) if device_inputs_by_request is None: device_inputs_by_request = {} @@ -1528,11 +1518,12 @@ def _sample_audio_code_greedy( full_bins=_codec_full_bins_input(request_state, hidden_state.device), ) device_states[request_id] = device_state - elif int(device_state.max_tokens.item()) != int(max_tokens): + else: # Same refresh as _sample_audio_code: the caller passes the host # state's current frame ceiling, which the request's remaining # output budget may have tightened after this device state was - # built. Refreshing in place keeps the per-request sampling graph. + # built. Refreshing in place keeps the per-request sampling graph, + # and fill_ is a no-op when the value already matches. device_state.max_tokens.fill_(int(max_tokens)) device_inputs_by_request = getattr(self, "_request_codec_device_inputs", None) @@ -1641,8 +1632,6 @@ def _flush_deferred_cleanup(self) -> None: """ request_audio_states = getattr(self, "_request_audio_states", {}) request_condition_states = getattr(self, "_request_condition_states", {}) - penalty_windows = getattr(self, "_penalty_windows_dev", None) - penalty_freqs = getattr(self, "_penalty_freqs_dev", None) codec_device_states = getattr(self, "_request_codec_device_states", None) codec_device_inputs = getattr(self, "_request_codec_device_inputs", None) codec_histories = getattr(self, "_request_codec_history", None) @@ -1650,10 +1639,6 @@ def _flush_deferred_cleanup(self) -> None: for request_id in self._deferred_cleanup_ids: request_audio_states.pop(request_id, None) request_condition_states.pop(request_id, None) - if isinstance(penalty_windows, dict): - penalty_windows.pop(request_id, None) - if isinstance(penalty_freqs, dict): - penalty_freqs.pop(request_id, None) if isinstance(codec_device_states, dict): codec_device_states.pop(request_id, None) if isinstance(codec_device_inputs, dict): @@ -1734,37 +1719,8 @@ def _make_omni_output_single_frame( penalty_histories = [empty_history for _ in infos] for index, info in enumerate(infos): info_dict = info if isinstance(info, dict) else {} - native_duplex = info_dict.get("native_duplex") is True if emit_duplex_metadata: - duplex_info = info_dict.get("duplex") - if not isinstance(duplex_info, dict): - duplex_info = {} - epoch = duplex_info.get("epoch", -1) - turn_id = duplex_info.get("turn_id", -1) - if native_duplex and not all( - isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in (epoch, turn_id) - ): - raise RuntimeError( - "MiniCPM-o native duplex Talker requires non-negative integer " - f"epoch and turn_id, got epoch={epoch!r}, turn_id={turn_id!r}" - ) - meta_info = info_dict.get("meta") - if not isinstance(meta_info, dict): - meta_info = {} - segment_text = meta_info.get("native_duplex_segment_text", "") if native_duplex else "" - if not isinstance(segment_text, str): - segment_text = "" - turn_eos_id = meta_info.get("turn_eos_token_id") - ids_info = info_dict.get("ids") - tts_ids = ids_info.get("tts") if native_duplex and isinstance(ids_info, dict) else None - if isinstance(tts_ids, torch.Tensor): - contains_turn_eos = isinstance(turn_eos_id, int) and bool( - torch.any(tts_ids.reshape(-1) == turn_eos_id).item() - ) - elif isinstance(tts_ids, (list, tuple)): - contains_turn_eos = isinstance(turn_eos_id, int) and turn_eos_id in tts_ids - else: - contains_turn_eos = False + native_duplex, epoch, turn_id, segment_text, turn_end = _native_duplex_row_metadata(info_dict) native_duplex_flags.append(torch.tensor(native_duplex, dtype=torch.bool)) duplex_epochs.append(torch.tensor(epoch if isinstance(epoch, int) else -1, dtype=torch.long)) duplex_turn_ids.append(torch.tensor(turn_id if isinstance(turn_id, int) else -1, dtype=torch.long)) @@ -1774,7 +1730,7 @@ def _make_omni_output_single_frame( dtype=torch.uint8, ) ) - turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool)) + turn_end_flags.append(torch.tensor(turn_end, dtype=torch.bool)) if not isinstance(info, dict): continue @@ -1969,28 +1925,11 @@ def _apply_codec_repetition_penalty(self, logits, sampling_metadata): or not isinstance(penalties, torch.Tensor) or getattr(sampling_metadata, "no_penalties", False) ): - self._penalty_freq_rows = None return logits, sampling_metadata - freq_rows = getattr(self, "_penalty_freq_rows", None) - self._penalty_freq_rows = None - if ( - isinstance(freq_rows, list) - and len(freq_rows) == logits.shape[0] - and all( - isinstance(f, torch.Tensor) and f.device == logits.device and f.numel() == logits.shape[-1] - for f in freq_rows - ) - and bool((penalties.reshape(-1) != 1.0).any()) - ): - # Incremental histogram: freq rows are maintained on device by - # make_omni_output (+1 append / -1 evict), identical to a full - # scatter_add rebuild of the sliding window. - freqs = torch.stack([f.to(dtype=torch.float32) for f in freq_rows], dim=0) - alpha = torch.pow(penalties.to(device=logits.device, dtype=torch.float32).reshape(-1, 1), freqs).to( - dtype=logits.dtype - ) - logits = torch.where(logits < 0, logits * alpha, logits / alpha) - return logits, replace(sampling_metadata, repetition_penalties=torch.ones_like(penalties)) + # The incremental-histogram branch (_penalty_freq_rows) was removed: + # that attribute was only ever assigned None, so the branch was + # unreachable -- frequency/presence penalties use full_bins and the + # sliding repetition window lives on device_state.history. logits = _apply_batched_repetition_penalty( logits, histories, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index c1a75165019..86c2b1206d7 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -369,28 +369,13 @@ def greedy_codec_sample( penalized = _repetition_penalty_scaling(penalized, state, repetition_penalty) penalized = _apply_stream_penalties(penalized, state, frequency_penalty, presence_penalty) sampled = torch.argmax(penalized, dim=-1).to(torch.int32) - - active = ~state.finished - next_step = torch.where(active, state.step + 1, state.step) - if ignore_eos: - # Same rule as advance_codec_device_state: ``ignore_eos`` disables the - # engine's EOS stop, so it disables this one too. - is_eos = torch.zeros_like(state.finished) - else: - is_eos = sampled.to(torch.int64) == int(eos_token_id) - reached_limit = active & (next_step >= state.max_tokens) - finished = state.finished | is_eos | reached_limit - emit = active & (~is_eos) & (~reached_limit) - - shifted = torch.cat([state.history[:, 1:], sampled.to(torch.int32).unsqueeze(1)], dim=1) - append_at = state.history_len.clamp(min=0, max=HISTORY_WINDOW - 1).to(torch.long).unsqueeze(1) - appended = state.history.scatter(1, append_at, sampled.to(torch.int32).unsqueeze(1)) - candidate = torch.where((state.history_len >= HISTORY_WINDOW).unsqueeze(1), shifted, appended) - history = torch.where(emit.unsqueeze(1), candidate, state.history) - history_len = torch.clamp(state.history_len + emit.to(torch.int32), max=HISTORY_WINDOW) - _advance_full_bins(state, sampled, emit) - return TalkerCodecSampleResult( - sampled_token=sampled, - state=TalkerCodecDeviceState(history, history_len, next_step, state.max_tokens, finished, state.full_bins), - emit=emit, + # The state advance / emit / finished transition lives in exactly one + # place: codec_sample_result. The sampled id is already in [0, vocab), so + # the random and greedy paths share it verbatim (RNG untouched: the + # multinomial call lives in the caller, not here). + return codec_sample_result( + state, + sampled, + eos_token_id=eos_token_id, + ignore_eos=ignore_eos, ) From 1cca045845a0b815019a565354769df6deb2fe05 Mon Sep 17 00:00:00 2001 From: y_null Date: Tue, 29 Sep 2026 09:46:10 +0800 Subject: [PATCH 62/89] fix(minicpmo_4_5): restore the CPU codec transport contract in make_omni_output 6142cbfa3 rewrote the codec/penalty bookkeeping for the multi-frame-decode-off path and moved the payloads onto the hidden_states device. Main keeps these tensors on the CPU on purpose: the cross-stage transport serializes codes.audio to the next stage and host-side bookkeeping reads the penalty history back, so the GPU placement regresses the contract -- 3 parametrized unit tests fail on H100 while the same file passes on main (63/63). - make_omni_output: return codes.audio deltas moved to CPU - build the sliding-window history with .cpu() readbacks - keep penalty_histories on CPU (as on main); the batched penalty step packs them host-side and uploads one buffer per step Validation: tests/model_executor/models/minicpmo_4_5/ 745 passed 5 skipped on H100 (fd78f8224 + this fix). Signed-off-by: y_null --- .../models/minicpmo_4_5/minicpmo_4_5_omni_tts.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index e259f7bb757..07285ad9e34 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -1748,11 +1748,11 @@ def _make_omni_output_single_frame( # ``audio`` is the id sampled last step, i.e. exactly the # history the logits computed below are scored against. recent = state.get("recent_codes") - recent = (recent if isinstance(recent, list) else []) + codec_deltas[index].reshape(-1).tolist() + recent = (recent if isinstance(recent, list) else []) + codec_deltas[index].cpu().reshape(-1).tolist() state["recent_codes"] = recent[-_CODEC_PENALTY_WINDOW:] recent_codes = state.get("recent_codes") if recent_codes: - penalty_histories[index] = torch.tensor(recent_codes, dtype=torch.long, device=hidden.device) + penalty_histories[index] = torch.tensor(recent_codes, dtype=torch.long, device="cpu") max_tokens = state.get("max_tokens") min_tokens = state.get("min_tokens") step = int(state.get("step", 0)) @@ -1792,10 +1792,14 @@ def _make_omni_output_single_frame( "turn_end": turn_end_flags, } ) + # Cross-stage codec transport must stay on CPU: these deltas are + # serialized to the next stage and read back by host-side + # bookkeeping; penalty histories stay on CPU as well -- the batched + # penalty step packs them host-side and uploads one buffer per step. return OmniOutput( text_hidden_states=hidden, multimodal_outputs={ - "codes": {"audio": codec_deltas}, + "codes": {"audio": [delta.to("cpu") for delta in codec_deltas]}, "meta": meta_outputs, }, ) From e77b5a3cce309bf1f1f820f1fc59268635f09654 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 29 Sep 2026 02:42:19 +0000 Subject: [PATCH 63/89] fix(minicpmo): route codec stop ids through the device state make_omni_output recomputed the eos/stop-id/limit transition host-side while codec_sample_result ran the same transition on the device -- the same class of drift the P2 review comment flagged. The device state now owns all three conditions (stop_token_ids included) and the host stop routing consumes result.state.finished / result.emit directly. - advance_codec_device_state / codec_sample_result / greedy_codec_sample take the request's codec stop ids and fold them into finished/emit; stop ids are request-level constants, so the comparison loop unrolls at capture time like the pinned top-k/top-p values - _sample_audio_code_greedy returns the full TalkerCodecSampleResult like its stochastic twin - make_omni_output consumes the device decision; the host-side recomputation only remains as the scripted-stub fallback for CPU unit tests - repetition detection stays host-side by design: the scheduler never sees codec ids, so the engine's check_stop cannot police the codec stream; make_omni_output mirrors it over codec_full_ids exactly as the engine does over its own tokens Four tests cover stop-id routing on both boundaries, the normal-then-stop framing sequence, and the ignore_eos/stop-id split. Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 125 ++++++++++++++++++ .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 84 ++++++++---- .../minicpmo_4_5/talker_codec_sample.py | 26 +++- 3 files changed, 203 insertions(+), 32 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 9e394b3cac4..9c2cc7ee751 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -820,3 +820,128 @@ def test_eos_window_mask_hides_codec_eos(): # Drain window: EOS is forced out even though the step allows it. masked = prepare_codec_logits(logits.clone(), eos_window_masked=True, **kwargs) assert masked[0, _EOS_ID] == float("-inf") + + +_STOP_ID = 6561 # in-vocabulary for the codec head, not the codec EOS + + +def _real_boundary_talker(temperature: float): + """A talker on the real (non-stub) greedy or stochastic boundary.""" + model = _make_talker(k_step_frames=8, scripted_samples=[]) + model._codec_temperature = temperature + model._codec_repetition_penalty = 1.05 + model.head_code = nn.ModuleList([nn.Linear(8, _NUM_AUDIO_TOKENS, bias=False)]) + if temperature == 0.0: + real_greedy = type(model)._sample_audio_code_greedy + model._sample_audio_code_greedy = real_greedy.__get__(model, type(model)) + else: + model._codec_top_k = 0 + model._codec_top_p = 1.0 + model._codec_seed = 42 + real_sample = type(model)._sample_audio_code + model._sample_audio_code = real_sample.__get__(model, type(model)) + return model + + +def _bias_head_to(model, token_id: int): + """Make the head's argmax land exactly on ``token_id``.""" + with torch.no_grad(): + model.head_code[0].weight.zero_() + model.head_code[0].weight[token_id] = 1.0 + + +def test_stop_id_finishes_request_through_device_state_greedy(): + """A codec stop id finishes the request via the device state alone. + + Review follow-up (PR #7929): make_omni_output used to recompute the + eos/stop-id/limit transition host-side while codec_sample_result ran the + same transition on the device. The device state is now the single source, + so a stop-id frame finishes the request and drops out of the audio stream + without the host ever recomputing the stop set. + """ + model = _real_boundary_talker(0.0) + state = { + "step": 0, + "max_tokens": 4032, + "codes": torch.tensor([10, 11]), + "codec_stop_token_ids": [_STOP_ID], + } + model._request_audio_states["r1"] = state + _bias_head_to(model, _STOP_ID) + + out = _frame_call(model, torch.ones(1, 8)) + + assert state["finished"] is True + # The stop frame itself must not join the emitted audio stream. + assert out.multimodal_outputs["codes"]["audio"][0].numel() == 0 + # K-step bookkeeping: a terminating frame records no confirmed codec id. + assert "last_code" not in state + assert model._request_codec_history.get("r1", []) == [] + # The device state is what decided this. + assert bool(model._request_codec_device_states["r1"].finished.item()) is True + + +def test_normal_frame_then_stop_id_frame_device_routing(): + """A normal frame emits and records; the following stop frame terminates.""" + model = _real_boundary_talker(0.0) + state = { + "step": 0, + "max_tokens": 4032, + "codes": torch.tensor([10, 11]), + "codec_stop_token_ids": [_STOP_ID], + } + model._request_audio_states["r1"] = state + _bias_head_to(model, 42) + out0 = _frame_call(model, torch.ones(1, 8)) + assert state["last_code"] == 42 + assert state["finished"] is False + assert out0.multimodal_outputs["codes"]["audio"][0].reshape(-1).tolist() == [42] + + _bias_head_to(model, _STOP_ID) + out1 = _frame_call(model, torch.ones(1, 8)) + assert state["finished"] is True + assert out1.multimodal_outputs["codes"]["audio"][0].numel() == 0 + # Only confirmed frames reach the history: the stop frame does not. + assert model._request_codec_history["r1"] == [42] + + +def test_stop_id_finishes_request_stochastic_path(): + """The stochastic boundary routes stop ids through the same device state.""" + model = _real_boundary_talker(0.8) + state = { + "step": 0, + "max_tokens": 4032, + "codes": torch.tensor([10, 11]), + "codec_stop_token_ids": [_STOP_ID], + } + model._request_audio_states["r1"] = state + _bias_head_to(model, _STOP_ID) + + _frame_call(model, torch.ones(1, 8)) + + assert state["finished"] is True + assert bool(model._request_codec_device_states["r1"].finished.item()) is True + + +def test_ignore_eos_does_not_ignore_stop_ids(): + """``ignore_eos`` blanks the EOS id, not the request's stop ids. + + The engine treats the two independently (a sampled stop id finishes the + request even when ``ignore_eos`` is set), and the device-side transition + keeps that split: only the EOS comparison is blanked. + """ + model = _real_boundary_talker(0.0) + state = { + "step": 0, + "max_tokens": 4032, + "codes": torch.tensor([10, 11]), + "codec_stop_token_ids": [_STOP_ID], + "codec_ignore_eos": True, + } + model._request_audio_states["r1"] = state + _bias_head_to(model, _STOP_ID) + + _frame_call(model, torch.ones(1, 8)) + + assert state["finished"] is True + assert bool(model._request_codec_device_states["r1"].finished.item()) is True diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 07285ad9e34..9e350416dad 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -1104,8 +1104,14 @@ def make_omni_output( eos_window_masked = bool(state.get("turn_end_drain")) and _turn_end_boundary_eos_masked(step) # A request pinning temperature to 0 (override or stage default) # takes the deterministic boundary sampler, same as single-frame. + # Both boundary samplers return the full sample result; the device + # state inside it owns the eos/stop-id/limit transition (single + # source). CPU unit tests historically stub these methods with the + # sampled tensor itself -- a bare tensor means no device state, and + # the stop routing falls back to the host recomputation below. + sample_result = None if _codec_float_param(state, "codec_temperature", self._codec_temperature) == 0.0: - sampled = self._sample_audio_code_greedy( + greedy_result = self._sample_audio_code_greedy( hidden[end - 1 : end], codes, request_id, @@ -1114,43 +1120,63 @@ def make_omni_output( max_tokens, eos_window_masked, ) + if isinstance(greedy_result, TalkerCodecSampleResult): + sample_result = greedy_result + else: + sampled = greedy_result.reshape(()).to(torch.long) else: stochastic_result = self._sample_audio_code( hidden[end - 1 : end], codes, request_id, step, eos_window_masked ) if isinstance(stochastic_result, TalkerCodecSampleResult): - sampled = stochastic_result.sampled_token.reshape(()).to(torch.long) + sample_result = stochastic_result else: - # CPU unit tests and downstream subclasses historically - # stub this method with the sampled tensor itself. sampled = stochastic_result.reshape(()).to(torch.long) + if sample_result is not None: + # The kernel ABI is int32, while the existing connector/audio-code + # payload is int64. Keep the compatibility cast outside the fused op. + sampled = sample_result.sampled_token.to(torch.long).reshape(()) sampled_id = int(sampled.item()) full_ids = state.get("codec_full_ids") if isinstance(full_ids, list): full_ids.append(sampled_id) - if _codec_bool_param(state, "codec_ignore_eos", False): - # ignore_eos blanks the engine's _eos_token_id - # (vllm/sampling_params.py:670-671), so a sampled codec EOS is - # just another frame and the request runs to its length budget. - is_eos = False - else: - is_eos = sampled_id == self._codec_eos_id - # stop_token_ids: the engine ends the request on any of these ids - # (v1/core/sched/utils.py:105); the codec stream ends it on the codec - # EOS, so both stop conditions apply here. - if not is_eos and sampled_id in _codec_stop_ids(state): - is_eos = True state["step"] = _codec_int_param(state, "step", 0) + 1 - reached_limit = int(state["step"]) >= _codec_int_param(state, "max_tokens", self._codec_max_tokens) - # The engine stops the request when check_stop sees the pattern - # (v1/core/sched/utils.py:125-133); the K-step path owns the - # codec ids, so it has to make that call itself. - repetition_stop = _codec_repetition_detected(state) - finished = is_eos or reached_limit or repetition_stop + if sample_result is not None: + # Single-source stop routing: the device state owns the + # eos/stop-id/limit transition (codec_sample_result), so the + # host only consumes its finished/emit decision. Repetition + # detection stays host-side -- the scheduler never sees codec + # ids, so the engine's check_stop cannot police the codec + # stream; this mirrors it over codec_full_ids exactly as the + # engine does over its tokens (v1/core/sched/utils.py:125-133). + repetition_stop = _codec_repetition_detected(state) + finished = bool(sample_result.state.finished.item()) or repetition_stop + emit_frame = bool(sample_result.emit.item()) + else: + # Stub fallback: no device state behind the scripted sample. + if _codec_bool_param(state, "codec_ignore_eos", False): + # ignore_eos blanks the engine's _eos_token_id + # (vllm/sampling_params.py:670-671), so a sampled codec EOS is + # just another frame and the request runs to its length budget. + is_eos = False + else: + is_eos = sampled_id == self._codec_eos_id + # stop_token_ids: the engine ends the request on any of these ids + # (v1/core/sched/utils.py:105); the codec stream ends it on the codec + # EOS, so both stop conditions apply here. + if not is_eos and sampled_id in _codec_stop_ids(state): + is_eos = True + reached_limit = int(state["step"]) >= _codec_int_param(state, "max_tokens", self._codec_max_tokens) + # The engine stops the request when check_stop sees the pattern + # (v1/core/sched/utils.py:125-133); the K-step path owns the + # codec ids, so it has to make that call itself. + repetition_stop = _codec_repetition_detected(state) + finished = is_eos or reached_limit or repetition_stop + emit_frame = not is_eos and not reached_limit state["finished"] = finished # MiniCPMTTS.generate_chunk consumes the boundary sample but # returns only codes that were fed into the retained KV state. - if not is_eos and not reached_limit: + if emit_frame: codes = torch.cat([codes[-(_REPETITION_WINDOW - 1) :], sampled.reshape(1)]) delta = sampled.reshape(1, 1) if self._k_step_frames > 0: @@ -1484,6 +1510,7 @@ def _sample_audio_code( sampled, eos_token_id=eos_id, ignore_eos=_codec_bool_param(request_state, "codec_ignore_eos", False), + stop_token_ids=_codec_stop_ids(request_state), ) device_states[request_id] = result.state return result @@ -1497,10 +1524,12 @@ def _sample_audio_code_greedy( min_tokens: int, max_tokens: int, eos_window_masked: bool = False, - ) -> torch.Tensor: + ) -> TalkerCodecSampleResult: """Run the greedy codec boundary while keeping sampler state on device. - Host stop routing still consumes ``sampled.item()`` below; removing that + Returns the full sample result so the host stop routing consumes the + device-side finished/emit decision instead of recomputing it. The host + still reads ``sampled.item()`` for the connector payload; removing that sync requires a runner/connector state refactor and is not hidden here. """ device_states = getattr(self, "_request_codec_device_states", None) @@ -1554,11 +1583,10 @@ def _sample_audio_code_greedy( presence_penalty=_codec_float_param(request_state, "codec_presence_penalty", 0.0), ignore_eos=_codec_bool_param(request_state, "codec_ignore_eos", False), eos_window_masked=eos_window_masked, + stop_token_ids=_codec_stop_ids(request_state), ) device_states[request_id] = result.state - # The kernel ABI is int32, while the existing connector/audio-code - # payload is int64. Keep the compatibility cast outside the fused op. - return result.sampled_token.to(torch.long).reshape(()) + return result def _step_constants(self, hidden: torch.Tensor): """Constant per-step tensors, cached per (device, dtype). diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index 86c2b1206d7..b4054cc1e2b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -281,8 +281,16 @@ def advance_codec_device_state( *, eos_token_id: int, ignore_eos: bool = False, + stop_token_ids: tuple[int, ...] = (), ) -> TalkerCodecDeviceState: - """Advance persistent stochastic state with graph-capturable tensor ops.""" + """Advance persistent stochastic state with graph-capturable tensor ops. + + ``stop_token_ids`` are the request's extra codec stop ids -- the + ``stop_token_ids`` analogue for the self-sampled codec stream. They are + request-level constants, so the Python loop unrolls into a fixed set of + comparisons at capture time, the same way the pinned top-k/top-p values + are baked into the per-request sampling graph. + """ sampled = sampled_token.reshape(1).to(torch.int32) active = ~state.finished next_step = torch.where(active, state.step + 1, state.step) @@ -293,9 +301,12 @@ def advance_codec_device_state( is_eos = torch.zeros_like(state.finished) else: is_eos = sampled == int(eos_token_id) + is_stop = torch.zeros_like(state.finished) + for token_id in stop_token_ids: + is_stop = is_stop | (sampled == int(token_id)) reached_limit = active & (next_step >= state.max_tokens) - finished = state.finished | is_eos | reached_limit - emit = active & (~is_eos) & (~reached_limit) + finished = state.finished | is_eos | is_stop | reached_limit + emit = active & (~is_eos) & (~is_stop) & (~reached_limit) shifted = torch.cat([state.history[:, 1:], sampled.reshape(1, 1)], dim=1) append_at = state.history_len.clamp(min=0, max=HISTORY_WINDOW - 1).to(torch.long).reshape(1, 1) appended = state.history.scatter(1, append_at, sampled.reshape(1, 1)) @@ -312,6 +323,7 @@ def codec_sample_result( *, eos_token_id: int, ignore_eos: bool = False, + stop_token_ids: tuple[int, ...] = (), ) -> TalkerCodecSampleResult: """Advance device state and retain the device-side emit decision.""" sampled = sampled_token.reshape(1).to(torch.int32) @@ -321,12 +333,16 @@ def codec_sample_result( is_eos = torch.zeros_like(state.finished) else: is_eos = sampled == int(eos_token_id) - emit = active & (~is_eos) & (next_step < state.max_tokens) + is_stop = torch.zeros_like(state.finished) + for token_id in stop_token_ids: + is_stop = is_stop | (sampled == int(token_id)) + emit = active & (~is_eos) & (~is_stop) & (next_step < state.max_tokens) next_state = advance_codec_device_state( state, sampled, eos_token_id=eos_token_id, ignore_eos=ignore_eos, + stop_token_ids=stop_token_ids, ) return TalkerCodecSampleResult(sampled_token=sampled, state=next_state, emit=emit) @@ -343,6 +359,7 @@ def greedy_codec_sample( presence_penalty: float = 0.0, ignore_eos: bool = False, eos_window_masked: bool = False, + stop_token_ids: tuple[int, ...] = (), ) -> TalkerCodecSampleResult: """Greedy codec sample in graph-capturable tensor ops. @@ -378,4 +395,5 @@ def greedy_codec_sample( sampled, eos_token_id=eos_token_id, ignore_eos=ignore_eos, + stop_token_ids=stop_token_ids, ) From 93168b5c3f1c151de38a83a9ab9793cf59a74ae0 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 29 Sep 2026 15:59:42 +0800 Subject: [PATCH 64/89] fix(minicpmo_4_5): use PIECEWISE cudagraphs to keep vllm 0.30.0 runnable FULL_DECODE_ONLY deadlocks on vllm-ascend 0.30.0 + 910B3: the first request after capture hangs in the cross-stage D2H (npu_ar_model_runner.py:1222). 0.30.0 replays FULL decode graphs through the updatable-graph path (acl_graph.py:291) and updates graph params on a separate stream afterward; the replay never completes on this SoC. Unworkable alternatives, all verified: pre-sync before the D2H (still deadlocks inside stream synchronize), ENPU_ENABLE=true (+26 GiB device memory, KV budget goes negative), use_updatable_graph=False (FIA update_graph_params raises NotImplementedError in 0.30.0). PIECEWISE keeps the stage-0 piecewise capture (4 graphs, mixed prefill-decode) and completes the three bench tiers with 0 failures (32/1, 64/4, 128/8; ynx_omni_030, 2026-09-29). Stage 0 and stage 1 in this yaml now match minicpmo_4_5_2gpu.yaml and minicpmo_4_5_3gpu.yaml. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index a8133e246cc..1a6ffe2fa35 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -151,7 +151,7 @@ platforms: # verified on 910B1) but K's RTF win is FULL-only there. Measured # on 910B1: PIECEWISE without K is +75% RTF at 32/1 versus # FULL_DECODE_ONLY with K=8. - cudagraph_mode: FULL_DECODE_ONLY + cudagraph_mode: PIECEWISE - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered # stages 0 and 2, so stage 1 kept the base value of 4. @@ -172,7 +172,7 @@ platforms: # The Talker's stop head emits this id to finish a request. stop_token_ids: [1] compilation_config: - cudagraph_mode: FULL_DECODE_ONLY + cudagraph_mode: PIECEWISE - stage_id: 2 max_num_seqs: 8 # Stage 2 stays outer-eager and captures only exact-shape CFM DiT From bfd36cc1257ad08fac4ff9ad6e519c447d8f1cff Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 29 Sep 2026 16:07:17 +0800 Subject: [PATCH 65/89] fix(minicpmo_4_5_prof): use PIECEWISE cudagraphs to keep vllm 0.30.0 runnable Same fix as minicpmo_4_5.yaml (commit 93168b5c3): FULL_DECODE_ONLY deadlocks on vllm-ascend 0.30.0 + 910B3 in the first-request cross-stage D2H (updatable-graph dual-stream path, acl_graph.py:291). Both stage 0 (:157) and stage 1 (:181) switch to PIECEWISE, matching minicpmo_4_5{,_2gpu,_3gpu}.yaml. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 20aca4988b2..8eaecfc80cd 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -154,7 +154,7 @@ platforms: # verified on 910B1) but K's RTF win is FULL-only there. Measured # on 910B1: PIECEWISE without K is +75% RTF at 32/1 versus # FULL_DECODE_ONLY with K=8. - cudagraph_mode: FULL_DECODE_ONLY + cudagraph_mode: PIECEWISE - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered # stages 0 and 2, so stage 1 kept the base value of 4. @@ -178,7 +178,7 @@ platforms: # two stage-2 replicas on this card. kv_cache_memory_bytes: 2684354560 compilation_config: - cudagraph_mode: FULL_DECODE_ONLY + cudagraph_mode: PIECEWISE - stage_id: 2 max_num_seqs: 8 # Stage 2 stays outer-eager and captures only exact-shape CFM DiT From a45528e88138772ee4152592a864014a9933c56a Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 29 Sep 2026 16:33:32 +0800 Subject: [PATCH 66/89] fix(minicpmo_4_5): stage0 PIECEWISE + stage1 FULL recovers c1 and beats 0.29 baseline The 0.30.0 FULL deadlock is confined to stage 0: its FULL decode replay goes through the updatable-graph dual-stream path (acl_graph.py:291) and never completes on 910B3, hanging the cross-stage D2H. Stage 1 under FULL_DECODE_ONLY is safe (verified: smoke request completed, both stages' replays active, 0 engine tracebacks) and re-enables the talker static-shape graphs (16 captured per capacity) that PIECEWISE gate disables (static_shape_backend.py:77-88). Those graphs are what kept stage 1 at 35 ms/step on 0.29; without them K=8 runs 8 eager forwards per step (149-434 ms/step on pure PIECEWISE). Three tiers (mean_audio_rtf, 0 failures, 910B3, 0.30.0): 32/1 0.4334 vs 0.29 baseline 0.4560 (-5.0%) 64/4 1.0808 vs 1.1289 (-4.3%) 128/8 2.1264 vs 2.2071 (-3.7%) Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 1a6ffe2fa35..edc23de23ba 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -172,7 +172,7 @@ platforms: # The Talker's stop head emits this id to finish a request. stop_token_ids: [1] compilation_config: - cudagraph_mode: PIECEWISE + cudagraph_mode: FULL_DECODE_ONLY - stage_id: 2 max_num_seqs: 8 # Stage 2 stays outer-eager and captures only exact-shape CFM DiT From 19105e82d36844f1be490aed271e436e77c18327 Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 29 Sep 2026 19:54:55 +0800 Subject: [PATCH 67/89] fix(minicpmo_4_5_prof): align stage 1 cudagraph_mode with minicpmo_4_5.yaml minicpmo_4_5_prof.yaml stage 1 was PIECEWISE while minicpmo_4_5.yaml uses FULL_DECODE_ONLY for stage 1 (stage 0 PIECEWISE on both). Align both files to the same mixed mode. Signed-off-by: y-null --- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml index 8eaecfc80cd..02870401c08 100644 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_prof.yaml @@ -178,7 +178,7 @@ platforms: # two stage-2 replicas on this card. kv_cache_memory_bytes: 2684354560 compilation_config: - cudagraph_mode: PIECEWISE + cudagraph_mode: FULL_DECODE_ONLY - stage_id: 2 max_num_seqs: 8 # Stage 2 stays outer-eager and captures only exact-shape CFM DiT From 538347b9f873d27513ea371c62dbccf15111f928 Mon Sep 17 00:00:00 2001 From: y_null Date: Tue, 29 Sep 2026 21:10:24 +0800 Subject: [PATCH 68/89] [test][Ascend] Expect PIECEWISE stage-0 graphs under vLLM 0.30.0 FULL_DECODE_ONLY cannot run on the NPU stage-0 engine under vLLM 0.30.0, so the single-card deploy now pins stage 0 to PIECEWISE. The Talker (stage 1) keeps its full decode graph on the single-card deploy (in-graph sampling + multi-frame K expansion); multi-card variants stay PIECEWISE. Align the stage-scoped topology test with the per-stage configuration. Signed-off-by: y_null --- .../models/minicpmo_4_5/test_pipeline.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 6d9a4c199b2..423ee60056f 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -223,11 +223,13 @@ def test_npu_graph_configuration_is_stage_scoped(self, filename: str) -> None: deploy = _apply_platform_overrides(deploy, platform="npu") stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy) - # The single-card deploy captures a full decode graph, which keeps - # in-graph sampling and multi-frame expansion; multi-card variants keep - # the base mode. + # vLLM 0.30.0 cannot run FULL_DECODE_ONLY on the NPU stage-0 engine, + # so every variant pins PIECEWISE there. The single-card deploy keeps + # the Talker (stage 1) on a full decode graph, which preserves + # in-graph sampling and multi-frame expansion; multi-card variants + # keep the base mode. + assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE" expected_graph_mode = "FULL_DECODE_ONLY" if filename == "minicpmo_4_5.yaml" else "PIECEWISE" - assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == expected_graph_mode assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == expected_graph_mode assert stages[2].yaml_engine_args["enforce_eager"] is True # Only the single-card deploy raises the Code2Wav graph pool; the From a9cee1eb0dcd40295187fde360974bba8be51dad Mon Sep 17 00:00:00 2001 From: y-null Date: Tue, 29 Sep 2026 22:03:56 +0800 Subject: [PATCH 69/89] fix(review amy-why-3459): P2-1/P2-2/P2-3/P3-5/P3-6 + claims-2 code fixes - P2-1: reject logit_bias/allowed_token_ids at admission when the K-step codec stream is armed (OmniSchedulerMixin.add_request -> finish_requests FINISHED_ERROR); the in-model filter has no reader for them, and the talker_codec_sample module docstring no longer claims support - P2-2: gate the empty-row draft rollback on _talker_kstep_armed() (omni_ar_scheduler.py) -- rollback semantics verified for the Talker K-step loop only - P2-3: pack sampled_id/finished/emit into one D2H per frame in the minicpmo_4_5_omni_tts.py K loop (3 device syncs per frame -> 1) - P3-5: native_rejection_sampler._RESTORED latches only on a confirmed non-target SoC; an empty probe name retries instead of locking out the restore for the whole process lifetime - P3-6: single transition source _advance_classified shared by advance_codec_device_state and codec_sample_result (emit/finished computed once per frame); rep_counts histogram preallocated in the device state and zeroed in place per frame - claims-2: minicpmo_4_5.yaml stage0 comment attributes the K-step FULL-only RTF win to stage 1 Validation: py_compile all touched files; pytest tests/core/sched + tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py = 255 passed (3 logprobs tests fixed by adding _talker_kstep_armed to the SimpleNamespace stub fixture) Signed-off-by: y-null --- .../sched/test_omni_ar_scheduler_logprobs.py | 5 + vllm_omni/core/sched/omni_ar_scheduler.py | 7 +- vllm_omni/core/sched/omni_scheduler_mixin.py | 28 ++++++ vllm_omni/deploy/minicpmo_4_5.yaml | 10 +- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 22 ++++- .../minicpmo_4_5/talker_codec_sample.py | 94 ++++++++++++++----- .../platforms/npu/native_rejection_sampler.py | 10 +- 7 files changed, 145 insertions(+), 31 deletions(-) diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 747be2ed1a9..fb473ac405c 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py @@ -129,7 +129,12 @@ def remove_requests(self, requests) -> None: def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace: """Build a minimal stub for OmniARScheduler.update_from_output().""" + # update_from_output gates its K-step empty-row rollback on + # _talker_kstep_armed (omni_ar_scheduler.py); the stub models an + # un-armed scheduler so the rollback branch stays out of these + # generic logprob-behavior tests. scheduler = SimpleNamespace( + _talker_kstep_armed=lambda: False, perf_metrics=None, connector=None, chunk_transfer_adapter=None, diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 46db45c9f6e..96afd6e2872 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -680,7 +680,7 @@ def update_from_output( num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, request_id=req_id, ) - elif scheduled_spec_token_ids and sampled_token_ids: + elif self._talker_kstep_armed() and scheduled_spec_token_ids and sampled_token_ids: # This request had drafts scheduled but its row # came back with no generated tokens while the step itself # sampled (a logprob-contract failure emptied the row above, @@ -692,7 +692,10 @@ def update_from_output( # it from `.get(req_id)` and a request with no drafts this # step gets None -- an empty row without drafts is normal # (a non-final prefill chunk produces no tokens) and upstream - # skips it; the guard above keeps that path intact. + # skips it; the guard above keeps that path intact. The + # rollback semantics are verified for the Talker K-step loop + # only, hence the armed gate: a non-K-step spec-decode stage + # keeps upstream's own handling of the empty-row case. # A prefill-chunk row is different: the chunk's prompt tokens # did land in the KV cache, so only the D drafts roll back # (the base token stays advanced). diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index f803dc1f83e..a3cd0122f2d 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -921,6 +921,34 @@ def add_request(self, request: Request) -> None: the end of the last segment. Emit the terminal chunk here so the downstream stage terminates on the payload path either way. """ + # Review P2-1 reject gate: the K-step in-model codec sampler has no + # reader for logit_bias / allowed_token_ids (the talker_codec_sample + # filter implements censoring/penalties/temperature/min-p/top-k/top-p + # only). The engine applies both fields to the 6562-wide codec logits + # on a single-frame stage, so a request setting either would decode + # differently than the same request on a single-frame deployment, + # silently. Honoring them needs a LogitBiasState-style per-request + # machine the codec path does not have; finish the request as + # FINISHED_ERROR instead of letting it diverge unnoticed. + kstep_armed = getattr(self, "_talker_kstep_armed", None) + if kstep_armed is not None and kstep_armed(): + sp = getattr(request, "sampling_params", None) + unsupported = [ + name + for name in ("logit_bias", "allowed_token_ids") + if sp is not None and getattr(sp, name, None) + ] + if unsupported: + logger.warning( + "Rejecting request %s: %s not supported by the K-step codec " + "sampler (no reader in the in-model filter); the request " + "would decode differently than on a single-frame stage.", + request.request_id, + ", ".join(unsupported), + ) + super().add_request(request) + self.finish_requests([request.request_id], RequestStatus.FINISHED_ERROR) + return existing = self.requests.get(request.request_id) if existing is not None and existing.status == RequestStatus.WAITING_FOR_STREAMING_REQ: adapter = None if getattr(request, "resumable", False) else self._adapter_owing_terminal(existing) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index edc23de23ba..7f4fed0ce24 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -147,10 +147,12 @@ platforms: compilation_config: # A full decode-step graph keeps in-graph sampling and multi-frame # expansion. The static-shape builder gates its runtime buckets on - # the decode graph mode, so PIECEWISE also runs (K included, - # verified on 910B1) but K's RTF win is FULL-only there. Measured - # on 910B1: PIECEWISE without K is +75% RTF at 32/1 versus - # FULL_DECODE_ONLY with K=8. + # the decode graph mode, so PIECEWISE also runs here (K included, + # verified on 910B1). This is stage 0 and it ships PIECEWISE; the + # K-step's measured RTF win is FULL-only on stage 1, which ships + # FULL_DECODE_ONLY below. Stage-0 data point for scale (910B1): + # PIECEWISE without K is +75% RTF at 32/1 versus FULL_DECODE_ONLY + # with K=8. cudagraph_mode: PIECEWISE - stage_id: 1 # Concurrency cap raised from 4 to 8; the override block only covered diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 9e350416dad..a1493210ca5 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -1136,7 +1136,23 @@ def make_omni_output( # The kernel ABI is int32, while the existing connector/audio-code # payload is int64. Keep the compatibility cast outside the fused op. sampled = sample_result.sampled_token.to(torch.long).reshape(()) - sampled_id = int(sampled.item()) + # Review P2-3: the host needs sampled_id plus the device + # finished/emit decision every frame; pack the three scalars + # into one D2H transfer so a K-frame step costs one device + # sync per frame instead of three. + packed = torch.stack( + ( + sampled, + sample_result.state.finished.reshape(()).to(torch.long), + sample_result.emit.reshape(()).to(torch.long), + ) + ) + packed_host = packed.cpu().tolist() + sampled_id = int(packed_host[0]) + device_finished = bool(packed_host[1]) + device_emit = bool(packed_host[2]) + else: + sampled_id = int(sampled.item()) full_ids = state.get("codec_full_ids") if isinstance(full_ids, list): full_ids.append(sampled_id) @@ -1150,8 +1166,8 @@ def make_omni_output( # stream; this mirrors it over codec_full_ids exactly as the # engine does over its tokens (v1/core/sched/utils.py:125-133). repetition_stop = _codec_repetition_detected(state) - finished = bool(sample_result.state.finished.item()) or repetition_stop - emit_frame = bool(sample_result.emit.item()) + finished = device_finished or repetition_stop + emit_frame = device_emit else: # Stub fallback: no device state behind the scripted sample. if _codec_bool_param(state, "codec_ignore_eos", False): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py index b4054cc1e2b..f3e12ed6d9f 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_codec_sample.py @@ -4,11 +4,15 @@ """Framework boundaries for Talker codec sampling. Production stochastic sampling uses a logits-filter boundary implemented in -graph-capturable tensor ops: stop-token censoring, whitelist censoring, -per-id bias, repetition penalty, frequency/presence penalties, temperature, -min-p, top-k, then top-p over the top-k-filtered distribution keeping at least -one candidate, while native NPU operators keep softmax and ``torch.multinomial`` -semantics (including Generator state) unchanged. +graph-capturable tensor ops: stop-token censoring, repetition penalty, +frequency/presence penalties, temperature, min-p, top-k, then top-p over the +top-k-filtered distribution keeping at least one candidate, while native NPU +operators keep softmax and ``torch.multinomial`` semantics (including +Generator state) unchanged. ``logit_bias`` and ``allowed_token_ids`` are NOT +implemented here -- this filter has no reader for them -- so when the K-step +stream is armed the scheduler rejects requests that set either field +(finished as FINISHED_ERROR at admission) instead of silently sampling a +different codec distribution than a single-frame deployment would. Order and defaults follow the engine's single-frame path so that enabling the multi-frame (K-step) decode does not change the sampling semantics: @@ -44,6 +48,11 @@ class TalkerCodecDeviceState: # than re-allocated because a fresh VOCAB_SIZE-wide tensor per frame would # be an allocation per decode step. full_bins: torch.Tensor | None = None + # Per-frame scratch for the repetition-penalty histogram over the history + # window, (1, VOCAB_SIZE) float32, allocated once per request/segment and + # zeroed in place each frame -- same rationale as full_bins: a fresh + # VOCAB_SIZE-wide tensor per frame would be an allocation per decode step. + rep_counts: torch.Tensor | None = None @dataclass(frozen=True) @@ -73,6 +82,7 @@ def make_device_state( max_tokens=torch.tensor([max_tokens], dtype=torch.int32, device=codes.device), finished=torch.tensor([finished], dtype=torch.bool, device=codes.device), full_bins=full_bins, + rep_counts=torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=codes.device), ) @@ -153,7 +163,13 @@ def _repetition_penalty_scaling( positions = torch.arange(HISTORY_WINDOW, dtype=torch.int32, device=logits.device).reshape(1, -1) valid = positions < state.history_len.reshape(1, 1) safe_tokens = torch.where(valid, state.history, torch.zeros_like(state.history)).to(torch.long) - counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=logits.device) + if state.rep_counts is not None: + # Review P3-6: reuse the per-request histogram buffer instead of + # allocating a VOCAB_SIZE-wide tensor every frame. zero_() re-runs + # under graph replay, so no stale count survives a replay. + counts = state.rep_counts.zero_() + else: + counts = torch.zeros((1, VOCAB_SIZE), dtype=torch.float32, device=logits.device) counts.scatter_add_(1, safe_tokens, valid.to(torch.float32)) alpha = torch.pow(repetition_penalty.reshape(1, 1), counts) return torch.where(logits < 0, logits * alpha, logits / alpha) @@ -292,8 +308,26 @@ def advance_codec_device_state( are baked into the per-request sampling graph. """ sampled = sampled_token.reshape(1).to(torch.int32) - active = ~state.finished - next_step = torch.where(active, state.step + 1, state.step) + is_eos, is_stop = _classify_stop( + sampled, + state, + eos_token_id=eos_token_id, + ignore_eos=ignore_eos, + stop_token_ids=stop_token_ids, + ) + next_state, _emit, _finished = _advance_classified(state, sampled, is_eos, is_stop) + return next_state + + +def _classify_stop( + sampled: torch.Tensor, + state: TalkerCodecDeviceState, + *, + eos_token_id: int, + ignore_eos: bool, + stop_token_ids: tuple[int, ...], +) -> tuple[torch.Tensor, torch.Tensor]: + """The eos / stop-id masks shared by every advance entry point.""" if ignore_eos: # ``ignore_eos`` blanks the engine's ``_eos_token_id`` # (vllm/sampling_params.py:670-671), so a sampled codec EOS is just @@ -304,6 +338,25 @@ def advance_codec_device_state( is_stop = torch.zeros_like(state.finished) for token_id in stop_token_ids: is_stop = is_stop | (sampled == int(token_id)) + return is_eos, is_stop + + +def _advance_classified( + state: TalkerCodecDeviceState, + sampled: torch.Tensor, + is_eos: torch.Tensor, + is_stop: torch.Tensor, +) -> tuple[TalkerCodecDeviceState, torch.Tensor, torch.Tensor]: + """Single transition source: next state plus the emit/finished decision. + + Both entry points (``advance_codec_device_state`` and + ``codec_sample_result``) route through here, so the eos/stop/limit + transition and its emit mask are computed exactly once per frame + (review P3-6: they used to be computed twice -- once in + ``codec_sample_result`` and again inside ``advance_codec_device_state``). + """ + active = ~state.finished + next_step = torch.where(active, state.step + 1, state.step) reached_limit = active & (next_step >= state.max_tokens) finished = state.finished | is_eos | is_stop | reached_limit emit = active & (~is_eos) & (~is_stop) & (~reached_limit) @@ -314,7 +367,10 @@ def advance_codec_device_state( history = torch.where(emit.reshape(1, 1), candidate, state.history) history_len = torch.clamp(state.history_len + emit.to(torch.int32), max=HISTORY_WINDOW) _advance_full_bins(state, sampled, emit) - return TalkerCodecDeviceState(history, history_len, next_step, state.max_tokens, finished, state.full_bins) + next_state = TalkerCodecDeviceState( + history, history_len, next_step, state.max_tokens, finished, state.full_bins, state.rep_counts + ) + return next_state, emit, finished def codec_sample_result( @@ -327,23 +383,19 @@ def codec_sample_result( ) -> TalkerCodecSampleResult: """Advance device state and retain the device-side emit decision.""" sampled = sampled_token.reshape(1).to(torch.int32) - active = ~state.finished - next_step = torch.where(active, state.step + 1, state.step) - if ignore_eos: - is_eos = torch.zeros_like(state.finished) - else: - is_eos = sampled == int(eos_token_id) - is_stop = torch.zeros_like(state.finished) - for token_id in stop_token_ids: - is_stop = is_stop | (sampled == int(token_id)) - emit = active & (~is_eos) & (~is_stop) & (next_step < state.max_tokens) - next_state = advance_codec_device_state( - state, + is_eos, is_stop = _classify_stop( sampled, + state, eos_token_id=eos_token_id, ignore_eos=ignore_eos, stop_token_ids=stop_token_ids, ) + # One transition for state + emit (review P3-6): the emit mask here used + # to be recomputed independently of the one inside + # ``advance_codec_device_state``. Semantics unchanged: for an active + # request ``~reached_limit`` == ``next_step < max_tokens``, and for an + # already-finished one both forms emit False. + next_state, emit, _finished = _advance_classified(state, sampled, is_eos, is_stop) return TalkerCodecSampleResult(sampled_token=sampled, state=next_state, emit=emit) diff --git a/vllm_omni/platforms/npu/native_rejection_sampler.py b/vllm_omni/platforms/npu/native_rejection_sampler.py index 040581b4aab..494450fa4a9 100644 --- a/vllm_omni/platforms/npu/native_rejection_sampler.py +++ b/vllm_omni/platforms/npu/native_rejection_sampler.py @@ -51,9 +51,17 @@ def restore_native_rejection_sampler() -> None: global _RESTORED if _RESTORED: return - _RESTORED = True # probe once; a non-target SoC never retries if not _target_soc(): + # Latch only on a confirmed non-target SoC: "probe once" describes a + # real part. An empty probe result is "unknown", not "non-target" -- + # leaving _RESTORED False here lets the next call retry instead of + # skipping the restore for the whole process lifetime. The probe is + # an env/name lookup and this function runs once per worker init, so + # the retry costs nothing. + if _probe_soc_name().strip(): + _RESTORED = True return + _RESTORED = True try: import vllm.v1.sample.rejection_sampler as rs except ImportError: # pragma: no cover - vllm core moved the module From 9ad237e405e06eaf31bdbbacc4cb0b65c86f04a8 Mon Sep 17 00:00:00 2001 From: y_null Date: Tue, 29 Sep 2026 23:19:45 +0800 Subject: [PATCH 70/89] [Deploy][MiniCPM-o 4.5] Turn the stage-2 replicas profile into a base_config overlay The profiling profile was a 193-line copy of minicpmo_4_5.yaml and had already drifted from it (prefix caching, code2wav_max_npu_graphs, captured-shape counts). It also armed enable_prefix_caching on the NPU stage-0 engine of a duplex deploy, which validate_duplex_window_install rejects at startup. Replace it with a thin minicpmo_4_5_stage2_replicas.yaml overlay (base_config: minicpmo_4_5.yaml) that spells out only the two-replica topology and its sizing, inheriting the duplex contracts (including enable_prefix_caching=false) from the base config. The rename follows the *_stage1_replicas naming and the file is marked as a non-production profiling entrypoint. Signed-off-by: y_null --- vllm_omni/deploy/minicpmo_4_5_prof.yaml | 193 ------------------ .../deploy/minicpmo_4_5_stage2_replicas.yaml | 74 +++++++ 2 files changed, 74 insertions(+), 193 deletions(-) delete mode 100644 vllm_omni/deploy/minicpmo_4_5_prof.yaml create mode 100644 vllm_omni/deploy/minicpmo_4_5_stage2_replicas.yaml diff --git a/vllm_omni/deploy/minicpmo_4_5_prof.yaml b/vllm_omni/deploy/minicpmo_4_5_prof.yaml deleted file mode 100644 index 02870401c08..00000000000 --- a/vllm_omni/deploy/minicpmo_4_5_prof.yaml +++ /dev/null @@ -1,193 +0,0 @@ -# MiniCPM-o 4.5 three-stage deploy profile: all three stages on one device, -# with stage 2 running two colocated replicas (the base deploy config runs a -# single stage-2 replica). Only that difference is spelled out here. -# MiniCPM-o 4.5 three-stage deploy on one large-memory GPU. -pipeline: minicpmo_4_5 -async_chunk: true -# Use minicpmo_4_5_turn.yaml to select the ordinary turn-based serving engine. -session_mode: duplex -# Keep the active window aligned with the session and stage capacities below. -# A smaller window serializes first-packet generation for concurrent chat turns. -duplex_session: - # Matches Stage 0/1 max_num_seqs below. - max_sessions: 4 - -trust_remote_code: true -enable_prefix_caching: false - -connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - connector_get_sleep_s: 0.01 - connector_get_max_wait_first_chunk: 3000 - connector_get_max_wait: 300 - # Steady-state payload length in codec frames. - codec_chunk_frames: 75 - # Smaller first payload so the first audio packet is emitted earlier. - # Takes effect only when initial < codec_chunk_frames; values far below - # the steady-state chunk length add a seam between chunks and hurt WER. - initial_codec_chunk_frames: 25 - codec_left_context_frames: 3 - # HiFT capture buckets are derived from codec_chunk_frames and - # codec_left_context_frames. (#5869) - # A future initial_codec_chunk_frames shape will use lazy capture unless - # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. - enable_hift_graph: true - enable_cfm_graph: true - cfm_max_graphs: 32 - cfm_graph_bucket_frames: 16 - -stages: - - stage_id: 0 - max_num_seqs: 4 - # Leave headroom on the shared GPU for Talker KV + the two stage-2 - # replicas below; the NPU block caps stage 0/1 KV explicitly. - gpu_memory_utilization: 0.40 - enforce_eager: false - # Stage-level prefix caching is NPU-only (see platforms.npu below): in the - # platform-shared block it killed the stage-1 replica on CUDA (empty tts - # handoff, then an indexSelectSmallIndex device assert). Stage-0 n-gram - # speculation stays off either way -- the omni prefix cache refuses it next - # to enable_prefix_caching. - tensor_parallel_size: 1 - max_num_batched_tokens: 16384 - devices: "0" - # Required for MiniCPM interleaved image/audio Daily-Omni packs (~78% recipe). - # Keep these in YAML so they are applied through --deploy-config. - interleave_mm_strings: true - media_io_kwargs: - video: - fps: 1 - num_frames: 128 - limit_mm_per_prompt: - # Daily-Omni: up to 64 interleaved image/audio pairs (1fps). - image: 64 - audio: 64 - video: 1 - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 2048 - seed: 42 - detokenize: true - # 1.0 for MCQ accuracy; 1.1 can bias short A–D answers. - repetition_penalty: 1.0 - - - stage_id: 1 - max_num_seqs: 4 - # Preserve the checkpoint's full-attention Talker prompt while it fits; - # native duplex rolls over to a one-chunk window at context capacity. - # Leave device-level headroom for the colocated HiFi-GAN cuDNN workspace. - gpu_memory_utilization: 0.08 - enforce_eager: false - max_num_batched_tokens: 8192 - # Talker context is 4096 (tts_config.max_position_embeddings). - max_model_len: 4096 - devices: "0" - output_connectors: - to_stage_2: connector_of_shared_memory - # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler - # (head_code logits). EOS is tts_config.eos_token_id (num_audio_tokens-1). - # Values track upstream utils.TTSSamplingParams (min_p/win_size/tau_r are - # declared there but never read, so they stay out). repetition_penalty is - # consumed by the Talker, which rescores it over upstream's 16-frame window - # instead of the Sampler's whole-stream presence penalty. - # min_tokens is upstream's min_new_token. - default_sampling_params: - temperature: 0.8 - top_p: 0.85 - top_k: 25 - repetition_penalty: 1.05 - min_tokens: 50 - max_tokens: 4096 - seed: 42 - detokenize: false - - - stage_id: 2 - max_num_seqs: 4 - gpu_memory_utilization: 0.18 - enforce_eager: true - enable_chunked_prefill: false - max_num_batched_tokens: 65536 - max_model_len: 65536 - # Two replicas on the same device: `devices` repeats the card once per - # replica (replica r lands on card r), so this needs no extra hardware. - num_replicas: 2 - devices: "0,0" - input_connectors: - from_stage_1: connector_of_shared_memory - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 65536 - detokenize: true - -platforms: - cuda: - stages: - - stage_id: 1 - # Full-attention occupancy is bounded by max_num_seqs * max_model_len - # (4 * 4096 = 16384 tokens). Keep the reported CUDA KV token capacity - # above that bound because physical-preemption replay is unsupported. - # The explicit cap leaves device headroom for concurrent Stage 2. - kv_cache_memory_bytes: 2147483648 # 2 GiB - - npu: - stages: - - stage_id: 0 - # Keep the three stages' concurrency caps aligned so the top tier is - # not gated by a single stage. - max_num_seqs: 8 - max_num_batched_tokens: 8192 - # Moved here from the platform-shared block; see the base stage note. - enable_prefix_caching: true - # Explicit KV cap: the auto-computed value does not leave room for the - # two stage-2 replicas on this card. - kv_cache_memory_bytes: 6442450944 - compilation_config: - # A full decode-step graph keeps in-graph sampling and multi-frame - # expansion. The static-shape builder gates its runtime buckets on - # the decode graph mode, so PIECEWISE also runs (K included, - # verified on 910B1) but K's RTF win is FULL-only there. Measured - # on 910B1: PIECEWISE without K is +75% RTF at 32/1 versus - # FULL_DECODE_ONLY with K=8. - cudagraph_mode: PIECEWISE - - stage_id: 1 - # Concurrency cap raised from 4 to 8; the override block only covered - # stages 0 and 2, so stage 1 kept the base value of 4. - max_num_seqs: 8 - max_num_batched_tokens: 8192 - # Multi-frame decode: 8 codec frames per step. NPU-only -- only the NPU - # worker intercepts the step and injects the constant continue drafts, - # so the block stays in this section. stop_token_ids below is the - # two-wide continue/stop id that head emits; drop the block (or set it - # to null) for one frame per step, and drop that id with it. - speculative_config: - method: ngram - num_speculative_tokens: 7 - prompt_lookup_min: 1 - prompt_lookup_max: 1 - async_scheduling: false - default_sampling_params: - # The Talker's stop head emits this id to finish a request. - stop_token_ids: [1] - # Explicit KV cap: the auto-computed value does not leave room for the - # two stage-2 replicas on this card. - kv_cache_memory_bytes: 2684354560 - compilation_config: - cudagraph_mode: FULL_DECODE_ONLY - - stage_id: 2 - max_num_seqs: 8 - # Stage 2 stays outer-eager and captures only exact-shape CFM DiT - # estimator kernels through the NPU platform patch. - additional_config: - code2wav_enable_npu_graph: true - # One pool entry per captured shape. With two stage-2 replicas the pool - # is per replica, and the CFM mel bucketing already present in main - # (cfm_graph_bucket_frames: 16) keeps the captured shape count near 31, - # so 32 covers it; 64 per replica would only add graph memory that is - # never replayed. - code2wav_max_npu_graphs: 32 diff --git a/vllm_omni/deploy/minicpmo_4_5_stage2_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_stage2_replicas.yaml new file mode 100644 index 00000000000..318ba18e137 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_stage2_replicas.yaml @@ -0,0 +1,74 @@ +# MiniCPM-o 4.5 profiling overlay: all three stages on one device, with two +# colocated stage-2 (Code2Wav) replicas. Experimental profiling entrypoint -- +# not a recommended production profile. +# +# Everything not spelled out here is inherited from minicpmo_4_5.yaml. +base_config: minicpmo_4_5.yaml +pipeline: minicpmo_4_5 + +# Stage 0/1 run 4 concurrent sessions instead of 16 to leave device headroom +# for the second stage-2 replica. +duplex_session: + max_sessions: 4 + +connectors: + connector_of_shared_memory: + extra: + # Cap whole-Euler CFM graphs at the stage-2 batch (4); the base 16-batch + # graph arena would cost ~10 GiB shared with the second replica. + max_graph_batch: 4 + # Pin the CFM mel-frame bucket step so the capture pool stays small on + # the shared device. + cfm_graph_bucket_frames: 16 + +stages: + - stage_id: 0 + max_num_seqs: 4 + # Leave headroom on the shared GPU for Talker KV + the two stage-2 + # replicas below; the NPU block caps stage 0/1 KV explicitly. + gpu_memory_utilization: 0.40 + + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.08 + + - stage_id: 2 + max_num_seqs: 4 + # Two replicas on the same device: `devices` repeats the card once per + # replica (replica r lands on card r), so this needs no extra hardware. + num_replicas: 2 + devices: "0,0" + +platforms: + cuda: + stages: + - stage_id: 1 + # Full-attention occupancy is bounded by max_num_seqs * max_model_len + # (4 * 4096 = 16384 tokens). Keep the reported CUDA KV token capacity + # above that bound because physical-preemption replay is unsupported. + # The explicit cap leaves device headroom for concurrent Stage 2. + kv_cache_memory_bytes: 2147483648 # 2 GiB + + npu: + stages: + - stage_id: 0 + # Explicit KV cap: the auto-computed value does not leave room for the + # two stage-2 replicas on this card. + kv_cache_memory_bytes: 6442450944 + + - stage_id: 1 + # Explicit KV cap: the auto-computed value does not leave room for the + # two stage-2 replicas on this card. + kv_cache_memory_bytes: 2684354560 + + - stage_id: 2 + # One pool entry per captured shape. With two stage-2 replicas the pool + # is per replica, and the CFM mel bucketing above + # (cfm_graph_bucket_frames: 16) keeps the captured shape count near 31, + # so 32 covers it; 64 per replica would only add graph memory that is + # never replayed. Both keys are spelled out: stage fields merge by + # replacement, so a partial additional_config would drop the base + # enable flag. + additional_config: + code2wav_enable_npu_graph: true + code2wav_max_npu_graphs: 32 From 1cb0958cb67c88fa85cab1cda79f9606797effe6 Mon Sep 17 00:00:00 2001 From: y-null Date: Wed, 30 Sep 2026 01:08:46 +0800 Subject: [PATCH 71/89] style: fix pre-commit violations (ruff format) CI's pre-commit run fails because ruff-format would reformat the kstep unsupported-params comprehension in omni_scheduler_mixin and reformat talker_multiframe and test_talker_kstep_feedback. Apply the formatter output verbatim; no functional changes. Signed-off-by: y-null --- .../test_talker_kstep_feedback.py | 40 +++++++++++++++++++ vllm_omni/core/sched/omni_scheduler_mixin.py | 4 +- .../platforms/npu/worker/talker_multiframe.py | 9 +++++ 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 9c2cc7ee751..0bf566a1139 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -17,6 +17,9 @@ - The multi-frame gate matrix: non-uniform decode spans stay blocked at ``applies()``; the runner raise remains the assertion of last resort for a combination the guard makes unschedulable. +- The draft rule: a stop-truncated row (a request that accepted fewer + tokens than the step ran) folds the whole batch to a draft-free step, + the same way an empty row already did. """ from types import SimpleNamespace @@ -316,6 +319,43 @@ def test_multiframe_gate_matrix(): assert talker_multiframe.is_multi_token_decode(model, mixed_prefill) is True +def test_constant_drafts_fold_when_a_stop_truncates_a_row(): + """A request whose codec stop row fired mid-step accepted fewer tokens + than the step ran; its next-step schedule is short the same way, so the + batch folds and the whole step takes the single-frame path instead of + scheduling a non-uniform span set that ``applies`` would refuse (and + ``_model_forward`` would turn into a fatal error).""" + from vllm_omni.platforms.npu.worker import talker_multiframe + + # Five staggered codecs, K=6: one stopped at frame 3 (4 tokens), three + # at frame 4 (5 tokens), one ran the step out (6 tokens). + sampled = [ + [1, 2, 3, 4], + [1, 2, 3, 4, 5], + [1, 2, 3, 4, 5], + [1, 2, 3, 4, 5], + [1, 2, 3, 4, 5, 6], + ] + assert talker_multiframe.constant_drafts(sampled, frames=6, num_reqs=5) == [[] for _ in range(5)] + + +def test_constant_drafts_fold_when_a_row_samples_nothing(): + from vllm_omni.platforms.npu.worker import talker_multiframe + + sampled = [[1, 2, 3, 4, 5, 6], [], [1, 2, 3, 4, 5, 6]] + assert talker_multiframe.constant_drafts(sampled, frames=6, num_reqs=3) == [[] for _ in range(3)] + + +def test_constant_drafts_keep_drafts_when_every_row_is_full(): + """No stop anywhere: every request accepted the step's K tokens and the + drafts arm the next K-frame step exactly as before the fold rule.""" + from vllm_omni.platforms.npu.worker import talker_multiframe + + sampled = [[1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5, 6]] + drafts = talker_multiframe.constant_drafts(sampled, frames=6, num_reqs=2) + assert drafts == [[talker_multiframe.CONTINUE_TOKEN_ID] * 5 for _ in range(2)] + + def _vocab_runner(*, supports_multi_frame: bool, vocab_size: int): return SimpleNamespace( model=SimpleNamespace(supports_multi_frame_decode=supports_multi_frame), diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index f69a95c9ff1..ceb75f6282f 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -945,9 +945,7 @@ def add_request(self, request: Request) -> None: if kstep_armed is not None and kstep_armed(): sp = getattr(request, "sampling_params", None) unsupported = [ - name - for name in ("logit_bias", "allowed_token_ids") - if sp is not None and getattr(sp, name, None) + name for name in ("logit_bias", "allowed_token_ids") if sp is not None and getattr(sp, name, None) ] if unsupported: logger.warning( diff --git a/vllm_omni/platforms/npu/worker/talker_multiframe.py b/vllm_omni/platforms/npu/worker/talker_multiframe.py index bb8882fe5d6..0d1b49d6cb4 100644 --- a/vllm_omni/platforms/npu/worker/talker_multiframe.py +++ b/vllm_omni/platforms/npu/worker/talker_multiframe.py @@ -586,6 +586,12 @@ def constant_drafts( into a fatal error rather than let the single-forward path sample one frame from the last row of a K-row span. + A request whose codec stop row fired mid-step is the same refusal by + another route: the rejection sampler truncates it to the frames it actually + accepted, so its next-step schedule is j+1 wide against the neighbours' K + -- non-uniform all the same. The empty check misses it, so the row length + is checked too: a short row folds the whole batch. + So: every request drafts, or nobody does. A step with no drafts is an ordinary one-frame-per-request decode, and the frames resume on the step after it. @@ -597,6 +603,9 @@ def constant_drafts( if not sampled: _log_block_once("a request in this batch sampled nothing; no drafts this step") return [[] for _ in range(num_reqs)] + if isinstance(sampled, list) and len(sampled) < frames: + _log_block_once("a request stopped mid-step, short of its drafts; no drafts this step") + return [[] for _ in range(num_reqs)] return [[CONTINUE_TOKEN_ID] * (frames - 1) for _ in range(num_reqs)] From 5bc44edc4b7fe9cd7d5c9f7d149aa6a2a8735b8e Mon Sep 17 00:00:00 2001 From: y-null Date: Wed, 30 Sep 2026 10:24:52 +0800 Subject: [PATCH 72/89] Fix ruff-format blank lines in minicpmo_4_5_omni_tts.py pre-commit (ruff format) flagged two double blank lines introduced by the merge conflict resolution; collapse both to single blank lines. Signed-off-by: y-null --- .../model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 003b7c2014c..4577b7c5c66 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -693,7 +693,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # recompute swaps its confirmed-token slice for these real ids. self._request_codec_history: dict[str, list[int]] = {} - self._init_native_talker(prefix) # Model Runner V2 keeps the sampled id, codec history and EOS state on # the GPU (see make_omni_output_mrv2 and mrv2_custom_sampler). @@ -2376,7 +2375,6 @@ def _force_eos_on_sampled_ids(self, output: Any, force_eos: list[bool] | None) - sampled.masked_fill_(rows.view(-1, *([1] * (sampled.ndim - 1))), int(self._codec_eos_id)) return output - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): return self._load_native_weights(weights) From d03dca05cfa4ef94f2b5f8212e2e17b24dc4c7af Mon Sep 17 00:00:00 2001 From: y-null Date: Wed, 30 Sep 2026 12:17:42 +0800 Subject: [PATCH 73/89] Fix num_output_tokens AttributeError in request_max_tokens_remaining CachedRequestState in vLLM 0.30.0 has no num_output_tokens attribute (only num_tokens as prompt+output); this made every decode step fall back to uniform budget splitting. Use len(req.output_token_ids), the equivalent generated-token count, instead. Signed-off-by: y-null --- vllm_omni/worker/gpu_model_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 5faba1ab135..042c560d013 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1491,7 +1491,7 @@ def _build_model_kwargs_extra(self) -> dict: ( max( min( - int(req.sampling_params.max_tokens) - int(req.num_output_tokens), + int(req.sampling_params.max_tokens) - len(req.output_token_ids), context_budget[index], ), 0, From 618ba5df200f60b20bce365d6ae9e4cf0b1a51ec Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:29 +0800 Subject: [PATCH 74/89] [Perf][Core] Opt-in in-place growth for CUMULATIVE stage outputs Outputs: grow CUMULATIVE dim-0 streams of matrices in place A CUMULATIVE stream re-consolidates the stream from the start on every step. Reusing capacity behind the previous view avoids O(T^2) recopying over long duplex epochs while returning value-identical views. Signed-off-by: BeatSeat --- tests/engine/test_output_cumulative_growth.py | 147 ++++++++++++++++++ vllm_omni/outputs/mm_outputs.py | 90 ++++++++++- vllm_omni/outputs/output_processor.py | 8 +- 3 files changed, 242 insertions(+), 3 deletions(-) create mode 100644 tests/engine/test_output_cumulative_growth.py diff --git a/tests/engine/test_output_cumulative_growth.py b/tests/engine/test_output_cumulative_growth.py new file mode 100644 index 00000000000..81ac587a0fb --- /dev/null +++ b/tests/engine/test_output_cumulative_growth.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUMULATIVE outputs re-consolidate the growing stream every step. + +The growth buffer must return exactly what ``torch.cat`` returned and keep +already-emitted snapshots intact. +""" + +import pytest +import torch +from vllm.outputs import PoolingRequestOutput +from vllm.sampling_params import RequestOutputKind +from vllm.v1.engine import FinishReason + +from vllm_omni.outputs.mm_outputs import MultimodalPayload +from vllm_omni.outputs.output_modality import OutputModality +from vllm_omni.outputs.output_processor import OmniRequestState + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _state(kind: RequestOutputKind) -> OmniRequestState: + return OmniRequestState( + request_id="audio", + external_req_id="audio", + parent_req=None, + request_index=0, + lora_request=None, + prompt=None, + prompt_token_ids=[0], + prompt_embeds=None, + logprobs_processor=None, + detokenizer=None, + max_tokens_param=None, + arrival_time=0.0, + queue=None, + log_stats=False, + stream_interval=1, + output_kind=kind, + ) + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + return tensor.contiguous().view(torch.int32) + + +def _chunk(index: int, length: int = 24) -> torch.Tensor: + generator = torch.Generator().manual_seed(index) + return torch.randn(length, generator=generator) + + +def test_cumulative_stream_matches_torch_cat_and_keeps_old_snapshots(): + state = _state(RequestOutputKind.CUMULATIVE) + chunks: list[torch.Tensor] = [] + emitted: list[torch.Tensor] = [] + for index in range(200): + chunk = _chunk(index, length=17 + index % 5) + chunks.append(chunk) + state.add_multimodal_tensor({"model_outputs": chunk, "sr": torch.tensor(24000)}, "audio") + output = state.make_request_output([], None, None, None) + assert output is not None and not isinstance(output, PoolingRequestOutput) + audio = output.outputs[0].multimodal_output["audio"] + assert audio.is_contiguous() and audio.dtype == torch.float32 + assert torch.equal(_bits(audio), _bits(torch.cat(chunks))) + emitted.append(audio) + assert state._mm_growth + + # Snapshots handed out earlier still hold exactly their prefix. + for index, audio in enumerate(emitted): + assert torch.equal(_bits(audio), _bits(torch.cat(chunks[: index + 1]))) + + final = state.make_request_output([], None, FinishReason.STOP, None) + assert final is not None and final.finished + assert torch.equal(_bits(final.outputs[0].multimodal_output["audio"]), _bits(torch.cat(chunks))) + + +@pytest.mark.parametrize("kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY]) +def test_non_cumulative_kinds_keep_plain_concatenation(kind): + state = _state(kind) + chunks = [_chunk(index) for index in range(5)] + for chunk in chunks: + state.add_multimodal_tensor({"model_outputs": chunk}, "audio") + state.make_request_output([], None, None, None) + final = state.make_request_output([], None, FinishReason.STOP, None) + assert state._mm_growth == {} + if kind == RequestOutputKind.FINAL_ONLY: + assert torch.equal(_bits(final.outputs[0].multimodal_output["audio"]), _bits(torch.cat(chunks))) + + +def _consolidated( + payload: MultimodalPayload, growth: dict | None, key: str = "audio", modality=OutputModality.AUDIO +) -> torch.Tensor: + payload.consolidate_tensors(modality, growth=growth) + return payload.tensors[key] + + +def test_growth_falls_back_for_mixed_dtype_and_multi_dim_chunks(): + growth: dict = {} + mixed = MultimodalPayload(tensors={"audio": [torch.ones(3), torch.ones(2, dtype=torch.float64)]}) + assert torch.equal(_consolidated(mixed, growth), torch.cat([torch.ones(3), torch.ones(2, dtype=torch.float64)])) + assert _consolidated(mixed, growth).dtype == torch.float64 + assert "audio" not in growth + + frames = [torch.arange(6.0).reshape(2, 3), torch.arange(6.0, 9.0).reshape(1, 3)] + stacked = MultimodalPayload(tensors={"audio": list(frames)}) + # CONCAT_LAST on 2-D chunks keeps its existing shape-mismatch fallback. + expected = MultimodalPayload(tensors={"audio": list(frames)}) + assert torch.equal(_consolidated(stacked, growth), _consolidated(expected, None)) + assert "audio" not in growth + + +def test_growth_restarts_when_the_head_is_not_the_last_view(): + growth: dict = {} + first = MultimodalPayload(tensors={"audio": [torch.ones(4), torch.full((2,), 2.0)]}) + view = _consolidated(first, growth) + # A foreign head (e.g. a replaced stream) must not append after `view`. + other = MultimodalPayload(tensors={"audio": [torch.zeros(3), torch.full((1,), 5.0)]}) + assert torch.equal(_consolidated(other, growth), torch.tensor([0.0, 0.0, 0.0, 5.0])) + assert torch.equal(view, torch.tensor([1.0, 1.0, 1.0, 1.0, 2.0, 2.0])) + + +def test_growth_extends_dim0_streams_of_matrices(): + """A LATENT stream (e.g. the Thinker's [T, hidden] rows) grows in place along dim 0.""" + growth: dict = {} + chunks: list[torch.Tensor] = [] + emitted: list[tuple[torch.Tensor, torch.Tensor]] = [] + head = None + for index in range(300): + chunk = torch.randn(1 + index % 3, 8, generator=torch.Generator().manual_seed(index)).to(torch.bfloat16) + chunks.append(chunk) + payload = MultimodalPayload(tensors={"latent": [chunk] if head is None else [head, chunk]}) + head = _consolidated(payload, growth, "latent", OutputModality.LATENT) + assert torch.equal(head, torch.cat(chunks, dim=0)) + emitted.append((head, torch.cat(chunks, dim=0))) + assert "latent" in growth + # Snapshots handed out earlier are never written. + for view, expected in emitted[::37]: + assert torch.equal(view, expected) + + +def test_growth_falls_back_when_dim0_chunks_change_trailing_shape(): + growth: dict = {} + chunks = [torch.ones(2, 4), torch.ones(1, 3)] + payload = MultimodalPayload(tensors={"latent": list(chunks)}) + with pytest.raises(RuntimeError): + _consolidated(payload, growth, "latent", OutputModality.LATENT) + assert "latent" not in growth diff --git a/vllm_omni/outputs/mm_outputs.py b/vllm_omni/outputs/mm_outputs.py index 45d1fc6fcdd..080b37ffe4c 100644 --- a/vllm_omni/outputs/mm_outputs.py +++ b/vllm_omni/outputs/mm_outputs.py @@ -70,6 +70,75 @@ def _consolidate_tensor_list( return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0) +# Spare capacity reserved when a growth buffer is (re)allocated: each sample is +# copied about 1 / (factor - 1) times amortized instead of once per +# consolidation, and at most 25% of the stream is slack. +_GROWTH_FACTOR = 1.25 + + +def _append_into_growth_buffer( + key: str, + tensor_list: list[torch.Tensor], + growth: dict[str, tuple[torch.Tensor, torch.Tensor]], + *, + along_dim0: bool = False, +) -> torch.Tensor | None: + """Concatenate chunks in place, or return ``None`` to use ``torch.cat``. + + A CUMULATIVE stream consolidates ``[everything so far, new chunk]`` on + every step, so ``torch.cat`` recopies the whole stream per chunk (O(T^2) + over a long duplex epoch). When the head is the view this function + returned last time, the new chunks are written into the storage behind it + and a longer prefix view is returned. Earlier views are never written, so + outputs already handed out keep their values. + + Chunks must share dtype and device. They must be 1-D, or, when + ``along_dim0`` (the strategy concatenates along dim 0), any rank with + the same trailing shape -- e.g. the Thinker's ``[T, hidden]`` latent rows. + The view then equals ``torch.cat(tensor_list, dim=0)`` element for element. + """ + head = tensor_list[0] + dtype = head.dtype + device = head.device + trailing = tuple(head.shape[1:]) if isinstance(head, torch.Tensor) else () + for tensor in tensor_list: + if ( + not isinstance(tensor, torch.Tensor) + or tensor.ndim == 0 + or (tensor.ndim != 1 and not along_dim0) + or tuple(tensor.shape[1:]) != trailing + or tensor.dtype != dtype + or tensor.device != device + or tensor.requires_grad + ): + growth.pop(key, None) + return None + state = growth.get(key) + if state is not None and state[1] is head: + storage, _ = state + start = head.shape[0] + chunks = tensor_list[1:] + else: + storage = None + start = 0 + chunks = tensor_list + needed = start + sum(chunk.shape[0] for chunk in chunks) + if storage is None or needed > storage.shape[0]: + rows = max(needed, int(needed * _GROWTH_FACTOR)) + grown = torch.empty((rows, *trailing), dtype=dtype, device=device) + if start: + grown[:start].copy_(storage[:start]) + storage = grown + offset = start + for chunk in chunks: + length = chunk.shape[0] + storage[offset : offset + length].copy_(chunk) + offset += length + view = storage[:needed] + growth[key] = (storage, view) + return view + + def _append_entries(store: dict[str, Any], incoming: dict[str, Any]) -> None: """Append *incoming* entries into *store*. @@ -174,7 +243,11 @@ def merged_with(self, incoming: MultimodalPayload) -> MultimodalPayload: self.metadata[key] = value return self - def consolidate_tensors(self, modality: OutputModality) -> None: + def consolidate_tensors( + self, + modality: OutputModality, + growth: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None, + ) -> None: """Concatenate deferred tensor lists into single tensors. Tensors are generated content accumulated as chunks, so each key's @@ -185,6 +258,9 @@ def consolidate_tensors(self, modality: OutputModality) -> None: key can be registered for a different strategy -- e.g. a codec-frame matrix that grows along dim 0 rather than the waveform-tuned default for its modality -- via ``register_key_accumulation_strategy``. + + *growth* is caller-owned state for a stream consolidated repeatedly + (see ``_append_into_growth_buffer``); results are value-identical. """ # Relocate scalar metadata tensors (e.g. sample rate) that from_dict # routed into .tensors, so they take the REPLACE path via .metadata @@ -198,7 +274,17 @@ def consolidate_tensors(self, modality: OutputModality) -> None: for key, value in list(self.tensors.items()): if _is_tensor_list(value): strategy = get_accumulation_strategy(modality, key) - self.tensors[key] = _consolidate_tensor_list(key, value, strategy) + # 1-D chunks concatenate identically along dim 0 and -1, so any + # concatenating strategy may reuse the caller's growth buffers; + # higher-rank chunks only when the strategy grows dim 0. + grown = ( + _append_into_growth_buffer( + key, value, growth, along_dim0=strategy is not TensorAccumulationStrategy.CONCAT_LAST + ) + if growth is not None and strategy is not TensorAccumulationStrategy.REPLACE + else None + ) + self.tensors[key] = grown if grown is not None else _consolidate_tensor_list(key, value, strategy) def consolidate_metadata(self) -> None: """Resolve deferred tensor lists in metadata by keeping the latest value. diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index cac8d999009..f486992b68d 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -109,6 +109,9 @@ def __init__( # types (e.g. dict[str, str]) for future multi-output models. self.mm_type: str | None = None self.mm_accumulated: MultimodalPayload = MultimodalPayload() + # Spare-capacity buffers behind keys that CUMULATIVE outputs + # re-consolidate on every step (see MultimodalPayload.consolidate_tensors). + self._mm_growth: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} def apply_streaming_update(self, update) -> None: super().apply_streaming_update(update) @@ -151,8 +154,11 @@ def _consolidate_multimodal_tensors(self) -> None: except (ValueError, KeyError): modality = OutputModality.TEXT + # Only a CUMULATIVE stream consolidates the same growing keys again on + # the next step; DELTA drains them and FINAL_ONLY consolidates once. + growth = self._mm_growth if self.output_kind == RequestOutputKind.CUMULATIVE else None try: - self.mm_accumulated.consolidate_tensors(modality) + self.mm_accumulated.consolidate_tensors(modality, growth=growth) self.mm_accumulated.consolidate_metadata() except (RuntimeError, TypeError, KeyError): logger.exception("Error consolidating multimodal tensors") From ba1e1959f7df1c9fa3ab74c930ac9ec355901c43 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:29 +0800 Subject: [PATCH 75/89] [Perf][Duplex] Framework hooks for pacing, unit continuation, and stage idle wait Framework hooks for pacing, unit continuation, and stage idle wait Add minimal opt-in hooks for duplex pacing, decoupled unit continuation, event-driven orchestrator loop, and stage idle waiting. Default behaviors for existing non-duplex and turn-based pipelines are preserved. Signed-off-by: BeatSeat --- .../test_generation_scheduler_deadline.py | 267 ++++++++++ .../test_chunk_transfer_adapter.py | 103 ++++ .../omni_connectors/test_shm_connector.py | 48 ++ tests/engine/duplex/test_duplex_barge_cut.py | 452 +++++++++++++++++ .../engine/duplex/test_duplex_intermediate.py | 118 ++++- tests/engine/duplex/test_duplex_onset_lead.py | 355 +++++++++++++ tests/engine/duplex/test_duplex_pacing.py | 304 ++++++++++++ .../duplex/test_duplex_unit_decouple.py | 426 ++++++++++++++++ tests/engine/duplex/test_session_runner.py | 170 ++++++- .../duplex/test_session_state_contract.py | 125 +++++ .../engine/test_orchestrator_event_driven.py | 7 +- tests/engine/test_stage_engine_core_proc.py | 185 +++++++ vllm_omni/config/stage_config.py | 105 ++++ .../core/sched/omni_generation_scheduler.py | 75 +++ .../connectors/shm_connector.py | 27 + .../omni_connectors/transfer_adapter/base.py | 70 ++- .../chunk_transfer_adapter.py | 19 +- vllm_omni/engine/duplex/config.py | 5 + vllm_omni/engine/duplex/intermediate.py | 57 ++- vllm_omni/engine/duplex/plugin.py | 32 ++ .../engine/duplex/session/append_task.py | 2 + vllm_omni/engine/duplex/session/context.py | 4 + vllm_omni/engine/duplex/session/helpers.py | 8 +- .../engine/duplex/session/model_channel.py | 117 ++++- vllm_omni/engine/duplex/session/pacing.py | 467 ++++++++++++++++++ vllm_omni/engine/duplex/session/runner.py | 251 +++++++++- vllm_omni/engine/duplex_orchestrator.py | 1 + vllm_omni/engine/orchestrator.py | 2 +- vllm_omni/engine/stage_engine_core_proc.py | 99 ++++ 29 files changed, 3865 insertions(+), 36 deletions(-) create mode 100644 tests/core/sched/test_generation_scheduler_deadline.py create mode 100644 tests/engine/duplex/test_duplex_barge_cut.py create mode 100644 tests/engine/duplex/test_duplex_onset_lead.py create mode 100644 tests/engine/duplex/test_duplex_pacing.py create mode 100644 tests/engine/duplex/test_duplex_unit_decouple.py create mode 100644 tests/engine/duplex/test_session_state_contract.py create mode 100644 vllm_omni/engine/duplex/session/pacing.py diff --git a/tests/core/sched/test_generation_scheduler_deadline.py b/tests/core/sched/test_generation_scheduler_deadline.py new file mode 100644 index 00000000000..f9fb6ff83fc --- /dev/null +++ b/tests/core/sched/test_generation_scheduler_deadline.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``OmniGenerationScheduler`` with deadline batching (``additional_config.codec_deadline``).""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from vllm.v1.core.sched.request_queue import create_request_queue + +import vllm_omni.core.sched.omni_generation_scheduler as scheduler_module +from tests.core.sched.test_generation_scheduler_restore import ( + FakeAdapter, + _chunk_request, + _make_generation_scheduler, +) +from vllm_omni.core.sched.codec_deadline import CodecDeadlineConfig, CodecDeadlinePolicy, StreamLedger +from vllm_omni.core.sched.omni_generation_scheduler import ( + OmniGenerationScheduler, + _build_codec_deadline, + _codec_deadline_plan, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _ReadyAdapter(FakeAdapter): + """Chunk adapter whose queued requests all hold a loaded chunk.""" + + def __init__(self) -> None: + super().__init__() + self.requests_with_ready_chunks: set[str] = set() + + def process_pending_chunks(self, waiting, running, scheduler_requests=None): + pass + + def restore_queues(self, waiting, running, scheduler_requests=None): + self.restore_called = True + + +def _meta(chunk_seq: int) -> dict: + return { + "meta": { + "chunk_seq": chunk_seq, + "cache_epoch": 0, + "last_chunk": False, + "turn_end": False, + "codec_chunk_frames": 25, + "duplex_epoch": 0, + "tts_is_last_chunk": True, + } + } + + +def _deadline_scheduler(monkeypatch, requests, *, policy: CodecDeadlinePolicy | None): + monkeypatch.setattr(scheduler_module, "create_request_queue", create_request_queue) + scheduler = _make_generation_scheduler(requests[0]) + scheduler.max_num_scheduled_tokens = 64 + scheduler.max_num_running_reqs = 8 + scheduler.chunk_transfer_adapter = _ReadyAdapter() + scheduler.waiting = create_request_queue(scheduler.policy) + scheduler.requests = {} + for request in requests: + scheduler.waiting.add_request(request) + scheduler.requests[request.request_id] = request + scheduler.chunk_transfer_adapter.requests_with_ready_chunks.add(request.request_id) + scheduler._codec_deadline = policy + return scheduler + + +def _requests(): + return [ + _chunk_request("cont", prompt_token_ids=[1, 2], additional_information=_meta(3)), + _chunk_request("first", prompt_token_ids=[1, 2], additional_information=_meta(0)), + ] + + +def test_without_a_policy_both_ready_chunks_are_scheduled(monkeypatch) -> None: + scheduler = _deadline_scheduler(monkeypatch, _requests(), policy=None) + output = OmniGenerationScheduler.schedule(scheduler) + assert set(output.num_scheduled_tokens) == {"cont", "first"} + + +def test_a_continuation_with_slack_is_held_and_released_by_its_hold_bound(monkeypatch) -> None: + policy = CodecDeadlinePolicy(CodecDeadlineConfig(enabled=True)) + now = scheduler_module.time.monotonic() + policy.ledgers["cont"] = StreamLedger(key=(0, 0), anchor=now - 5, play_end=now + 30, audio_s=30.0, d1=1.0) + requests = _requests() + scheduler = _deadline_scheduler(monkeypatch, requests, policy=policy) + + output = OmniGenerationScheduler.schedule(scheduler) + # The first chunk goes alone; the continuation stays queued, ready mark kept. + assert set(output.num_scheduled_tokens) == {"first"} + assert [request.request_id for request in scheduler.waiting] == ["cont"] + assert "cont" in scheduler.chunk_transfer_adapter.requests_with_ready_chunks + assert 0.0 < scheduler.next_release_in() <= 0.35 + assert policy.dispatched["first"][0].chunk_seq == 0 + + # The first chunk ran (the adapter clears its ready mark); past its max + # hold the continuation is released. + requests[1].num_computed_tokens = 2 + scheduler.chunk_transfer_adapter.requests_with_ready_chunks.discard("first") + policy.ready_since["cont"] -= 1.0 + output = OmniGenerationScheduler.schedule(scheduler) + assert set(output.num_scheduled_tokens) == {"cont"} + assert policy.dispatched["cont"][1] is True # it was held + assert scheduler.next_release_in() is None + + +def test_held_running_requests_are_skipped_in_place(monkeypatch) -> None: + policy = CodecDeadlinePolicy(CodecDeadlineConfig(enabled=True)) + now = scheduler_module.time.monotonic() + policy.ledgers["cont"] = StreamLedger(key=(0, 0), anchor=now - 5, play_end=now + 30, audio_s=30.0, d1=1.0) + requests = _requests() + scheduler = _deadline_scheduler(monkeypatch, requests[1:], policy=policy) + scheduler.running = [requests[0]] + scheduler.requests["cont"] = requests[0] + scheduler.chunk_transfer_adapter.requests_with_ready_chunks.add("cont") + output = OmniGenerationScheduler.schedule(scheduler) + assert set(output.num_scheduled_tokens) == {"first"} + assert requests[0] in scheduler.running + + +def test_plan_is_none_for_mock_or_missing_policies() -> None: + assert _codec_deadline_plan(MagicMock(), 0.0) is None + assert _codec_deadline_plan(SimpleNamespace(), 0.0) is None + assert OmniGenerationScheduler.next_release_in(MagicMock()) is None + assert OmniGenerationScheduler.next_release_in(SimpleNamespace()) is None + + +def test_free_request_forgets_the_stream(monkeypatch) -> None: + policy = CodecDeadlinePolicy(CodecDeadlineConfig(enabled=True)) + policy.ledgers["cont"] = StreamLedger(key=(0, 0), anchor=0.0, play_end=1.0, audio_s=1.0, d1=1.0) + policy.buffered.add("cont") + scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + scheduler._codec_deadline = policy + scheduler.input_coordinator = None + monkeypatch.setattr( + scheduler_module.VLLMScheduler, "_free_request", lambda self, request, delay=False: (None, None) + ) + scheduler._free_request(SimpleNamespace(request_id="cont")) + assert "cont" not in policy.ledgers and "cont" not in policy.buffered + + +def _build_target(*, additional_config, async_scheduling=False, native=False): + return SimpleNamespace( + vllm_config=SimpleNamespace(additional_config=additional_config), + scheduler_config=SimpleNamespace(async_scheduling=async_scheduling), + _native_data_plane=native, + _batch_window=False, + _first_chunk_express=False, + ) + + +@pytest.mark.parametrize( + ("idle_wait", "async_scheduling", "native", "enabled"), + [ + ("0.05", False, False, True), + ("0", False, False, False), + ("0.05", True, False, False), + ("0.05", False, True, False), + ], +) +def test_build_reads_additional_config_and_its_preconditions(monkeypatch, idle_wait, async_scheduling, native, enabled): + monkeypatch.setenv("VLLM_OMNI_STAGE_IDLE_WAIT_S", idle_wait) + target = _build_target( + additional_config={"codec_deadline": {"enabled": True}}, async_scheduling=async_scheduling, native=native + ) + assert isinstance(_build_codec_deadline(target), CodecDeadlinePolicy) is enabled + assert _build_codec_deadline(_build_target(additional_config={})) is None + + +def test_a_stream_without_its_first_chunk_marks_the_onset_pending(monkeypatch) -> None: + policy = CodecDeadlinePolicy(CodecDeadlineConfig(enabled=True)) + seen: list[bool] = [] + real_plan = policy.plan + + def plan(ready, now, *, onset_pending=False): + seen.append(onset_pending) + return real_plan(ready, now, onset_pending=onset_pending) + + monkeypatch.setattr(policy, "plan", plan) + now = scheduler_module.time.monotonic() + prewarmed = _chunk_request("prewarmed", prompt_token_ids=[], additional_information=_meta(0)) + _codec_deadline_plan(_deadline_scheduler(monkeypatch, [*_requests(), prewarmed], policy=policy), now) + _codec_deadline_plan(_deadline_scheduler(monkeypatch, _requests(), policy=policy), now) + assert seen == [True, False] + + +def test_update_from_output_feeds_the_ledger_and_the_step_time() -> None: + """The real ``update_from_output``: dispatch snapshot -> ledger, and the step's row count -> T(1).""" + import torch + from vllm.sampling_params import SamplingParams + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.outputs import ModelRunnerOutput + from vllm.v1.request import Request, RequestStatus + + from vllm_omni.core.sched.codec_deadline import seed_step_s + + request = Request( + request_id="first", + prompt_token_ids=[1, 2], + sampling_params=SamplingParams(max_tokens=8), + pooling_params=None, + arrival_time=0.0, + block_hasher=None, + ) + request.status = RequestStatus.RUNNING + request.num_computed_tokens = 2 + request.num_in_flight_tokens = 1 + + policy = CodecDeadlinePolicy(CodecDeadlineConfig(enabled=True)) + policy.on_scheduled("first", scheduler_module.ChunkMeta(0, 0, False, False, 25, 0), 1.0) + policy.last_schedule_t = scheduler_module.time.monotonic() - 0.5 + + sched = MagicMock() + sched._codec_deadline = policy + sched._audio_seconds = OmniGenerationScheduler._audio_seconds + sched._first_chunk_express = False + sched._express_min_slack_s = 0.0 + sched._batch_window = False + sched.requests = {"first": request} + sched.perf_metrics = None + sched.chunk_transfer_adapter = SimpleNamespace( + is_done_receiving_chunks=lambda _request_id: False, + segment_finished_requests=set(), + ) + sched.running = [request] + sched.waiting = MagicMock() + sched.skipped_waiting = MagicMock() + sched.structured_output_manager.accept_tokens.return_value = True + sched._pending_finish_reqs = [] + sched.recompute_kv_load_failures = False + sched.connector = None + sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 + sched.finished_req_ids_dict = {} + sched.make_stats.return_value = None + sched._handle_stopped_request.return_value = False + sched._free_request.return_value = (None, None) + + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {"first": 1} + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [[]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {"first": 0} + model_runner_output.routed_experts = None + model_runner_output.multimodal_outputs = [{"model_outputs": torch.zeros(2400), "sr": 24000}] + + OmniGenerationScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + ledger = policy.ledgers["first"] + assert ledger.audio_s == pytest.approx(0.1) + assert ledger.play_end - ledger.anchor == pytest.approx(0.5 + 0.1) # first segment: client prebuffer + assert "first" not in policy.dispatched + assert policy.step_time.mean[0] > seed_step_s(1) # one-row step of ~0.5 s observed diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 8f85366e3b2..2318675a401 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -1195,6 +1195,27 @@ def recording_processor(**kwargs): ] +@pytest.mark.parametrize("prompt_token_ids", [list(range(4096)), (1, 2, 3), None]) +def test_snapshot_freezes_prompt_ids_exactly_like_the_elementwise_walk(prompt_token_ids): + request = SimpleNamespace( + prompt_token_ids=prompt_token_ids, + additional_information={"ids": {"prompt": [5, 6]}, "meta": {"segment": "old"}}, + _all_token_ids=[1, 2], + _output_token_ids=[2], + ) + + snapshot = OmniChunkTransferAdapter._snapshot_processor_request(request) + + assert snapshot.prompt_token_ids == prompt_token_ids + assert type(snapshot.prompt_token_ids) is type(prompt_token_ids) + if isinstance(prompt_token_ids, list): + assert snapshot.prompt_token_ids is not prompt_token_ids + prompt_token_ids.append(-1) + assert snapshot.prompt_token_ids[-1] == 4095 + assert snapshot.additional_information == request.additional_information + assert snapshot.additional_information["ids"]["prompt"] is not request.additional_information["ids"]["prompt"] + + def test_save_without_custom_processor_does_not_snapshot_request(build_adapter, mocker): adapter, _ = build_adapter(stage_id=1) request = _req("req-direct", RequestStatus.RUNNING, external_req_id="ext-direct") @@ -3685,3 +3706,85 @@ def test_save_async_boundary_holds_generation_without_request_counter(build_adap assert len(adapter._pending_save_reqs) == queued_before + 1 assert adapter._pending_save_reqs[-1]["request"] is follow_up + + +# --- event-driven receive (VLLM_OMNI_STAGE_IDLE_WAIT_S) --------------------------------------- + + +class _WakeConnector: + """Connector double with arrival wakeups: ``get`` misses until ``deliver``.""" + + def __init__(self, stage_id: int, payload) -> None: + self.stage_id = stage_id + self.config = {"extra": {}} + self.payload = payload + self.ready = False + self.generation = 0 + self.waits: list[tuple[int, float]] = [] + self.wakes = 0 + + def get(self, *args, **kwargs): + return (self.payload, 16) if self.ready else None + + def get_wakeup_generation(self) -> int: + return self.generation + + def wait_for_change(self, generation: int, *, timeout: float) -> bool: + self.waits.append((generation, timeout)) + self.ready = True # the producer's put lands while the receiver waits + self.generation += 1 + return True + + def wake_self(self) -> None: + self.wakes += 1 + + +def _generation_payload(): + return { + "codes": {"audio": [1, 2, 3]}, + "meta": {"finished": torch.tensor(False, dtype=torch.bool)}, + } + + +def test_recv_loop_waits_on_connector_wakeup_and_signals_chunk_ready(build_adapter): + adapter, _ = build_adapter(stage_id=2, model_mode="generation") + connector = _WakeConnector(2, _generation_payload()) + adapter.connector = connector + adapter._recv_fifo_wait_s = 0.25 + ready = [] + + def on_ready(): + # The commit is visible before the wakeup: a woken schedule() sees it. + ready.append(set(adapter._finished_load_reqs)) + adapter.stop_event.set() + + adapter.set_chunk_ready_callback(on_ready) + request = _req("gen-1", RequestStatus.WAITING_FOR_CHUNK) + adapter.load_async(request) + assert connector.wakes == 1 # registration interrupts a blocked wait + adapter.recv_loop() + # One miss, one bounded wait taken with the pre-pass generation, then the hit. + assert connector.waits == [(0, 0.25)] + assert ready == [{"gen-1"}] + + +def test_recv_loop_keeps_the_1ms_poll_without_the_switch(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="generation") + adapter.load_async(_req("gen-2", RequestStatus.WAITING_FOR_CHUNK)) + timeouts = [] + + def wait(timeout): + timeouts.append(timeout) + adapter.stop_event.set() + + adapter._recv_cond.wait = lambda timeout: wait(timeout) + adapter.recv_loop() + assert timeouts == [0.001] + connector.wait_for_change.assert_not_called() + connector.wake_self.assert_not_called() + + +def test_chunk_ready_callback_failure_is_swallowed(build_adapter): + adapter, _ = build_adapter(stage_id=2, model_mode="generation") + adapter.set_chunk_ready_callback(Mock(side_effect=RuntimeError("queue closed"))) + adapter._notify_chunk_ready() # a lost hint only costs the fallback bound diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index 39a0e8327b4..b9bc7c168ed 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/tests/distributed/omni_connectors/test_shm_connector.py @@ -296,6 +296,38 @@ def fail_deserialize(_data): assert key not in connector._pending_keys +def test_poll_miss_does_not_attach_shared_memory(connector, monkeypatch): + """A miss is answered by the existence probe, without a raising shm_open.""" + from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module + + if module._POSIX_SHM_DIR is None: + pytest.skip("no /dev/shm on this host") + + def unexpected_attach(*args, **kwargs): + raise AssertionError("missing segment must not be attached") + + monkeypatch.setattr(module.shm_pkg, "SharedMemory", unexpected_attach) + assert connector.get("0", "1", f"missing_{uuid.uuid4().hex}") is None + + +@pytest.mark.parametrize("probe", [True, False]) +def test_key_read_is_unchanged_with_or_without_the_probe(connector, monkeypatch, probe): + from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module + + if not probe: + monkeypatch.setattr(module, "_POSIX_SHM_DIR", None) + key = f"probe_{probe}_{uuid.uuid4().hex}" + assert connector.get("0", "1", key) is None + payload = {"codes": torch.arange(28, dtype=torch.long), "meta": {"chunk_seq": 3}} + assert connector.put("0", "1", key, payload)[0] + + obj, _ = connector.get("0", "1", key) + + assert torch.equal(obj["codes"], payload["codes"]) and obj["meta"] == {"chunk_seq": 3} + assert connector.get("0", "1", key) is None + assert not os.path.exists(f"/dev/shm/{key}") + + # ── Arrival wakeups ─────────────────────────────────────────────────── @@ -405,6 +437,22 @@ def test_wakeups_can_be_disabled_and_close_unlinks_fifo(monkeypatch): receiver.close() +def test_wake_self_interrupts_own_wait_and_drains(): + receiver = _stage_connector(74) + try: + generation = receiver.get_wakeup_generation() + assert generation is not None + receiver.wake_self() + start = time.monotonic() + assert receiver.wait_for_change(generation, timeout=5) is True + assert time.monotonic() - start < 0.5 + # Drained: the next wait blocks for its timeout again. + assert receiver.wait_for_change(receiver.get_wakeup_generation(), timeout=0.02) is False + finally: + receiver.close() + receiver.wake_self() # closed: a no-op, never raises + + def _sibling_receiver(stage_id, ready, result): receiver = _stage_connector(stage_id) generation = receiver.get_wakeup_generation() diff --git a/tests/engine/duplex/test_duplex_barge_cut.py b/tests/engine/duplex/test_duplex_barge_cut.py new file mode 100644 index 00000000000..fd842717214 --- /dev/null +++ b/tests/engine/duplex/test_duplex_barge_cut.py @@ -0,0 +1,452 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``barge_cut_on_model_yield``: cut the reply's tail when the model yields to a new user utterance.""" + +from __future__ import annotations + +import asyncio +import base64 +from types import SimpleNamespace + +import numpy as np +import pytest + +from tests.engine.duplex.test_session_runner import ( + LISTEN_TOKEN_ID, + append_audio, + close_harness, + find, + listen_output, + open_harness, + tts_output, + types, +) +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex import commands +from vllm_omni.engine.duplex.config import DuplexSessionState +from vllm_omni.engine.duplex.session.pacing import SessionPacing + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _pace(**runtime: object) -> SessionPacing: + config: dict[str, object] = {"barge_cut_on_model_yield": True} + config.update(runtime) + pace = SessionPacing(DuplexSessionRuntimeConfig(**config), SimpleNamespace(active_response_id="resp_1")) + for index in range(10): + pace.on_audio_emit(0.1 * index, 1.0) # ~10 s of reply audio queued at the client + return pace + + +def _feed(pace: SessionPacing, t: float, *, quiet_s: float, speech_s: float, step: float = 0.2) -> float: + for _ in range(round(quiet_s / step)): + pace.on_client_append(t, is_speech=False, duration_s=step) + t += step + for _ in range(round(speech_s / step)): + pace.on_client_append(t, is_speech=True, duration_s=step) + t += step + return t + + +def test_new_utterance_over_playing_audio_arms_after_the_speech_floor() -> None: + pace = _pace() + t = _feed(pace, 1.2, quiet_s=0.8, speech_s=0.2) + assert pace.armed is False # a 0.2 s backchannel stays below the 0.4 s floor + t = _feed(pace, t, quiet_s=0.0, speech_s=0.2) + assert pace.armed is True + assert pace.cut_ready(t) is True + pace.session.active_response_id = None # the reply ended meanwhile + assert pace.cut_ready(t) is False + + +def test_a_short_pause_inside_the_interruption_does_not_reset_it() -> None: + pace = _pace() + t = _feed(pace, 1.2, quiet_s=0.8, speech_s=0.2) + t = _feed(pace, t, quiet_s=0.2, speech_s=0.2) # one low-energy 200 ms chunk mid-sentence + assert pace.armed is True + + +def test_a_long_pause_ends_the_utterance() -> None: + pace = _pace() + t = _feed(pace, 1.2, quiet_s=0.8, speech_s=0.2) + # 0.4 s pause: the run ends; the next 0.2 s has too little quiet before it to start a new one. + _feed(pace, t, quiet_s=0.4, speech_s=0.2) + assert pace.armed is False + + +@pytest.mark.parametrize( + ("start", "quiet_s", "speech_s", "runtime"), + [ + (1.2, 0.4, 1.0, {}), # not enough quiet before the speech + (0.2, 0.8, 1.0, {"barge_arm_min_playback_s": 2.0}), # reply not audible long enough + (1.2, 0.8, 0.2, {}), # backchannel + (1.2, 0.8, 1.0, {"barge_cut_on_model_yield": False}), + ], +) +def test_arming_conditions(start: float, quiet_s: float, speech_s: float, runtime: dict) -> None: + pace = _pace(**runtime) + _feed(pace, start, quiet_s=quiet_s, speech_s=speech_s) + assert pace.armed is False + + +def test_continuous_speech_from_before_the_reply_never_arms() -> None: + session = SimpleNamespace(active_response_id=None) + pace = SessionPacing(DuplexSessionRuntimeConfig(barge_cut_on_model_yield=True), session) + # The user is still talking when the reply's first audio arrives. + t = _feed(pace, 0.0, quiet_s=0.0, speech_s=1.0) + session.active_response_id = "resp_1" + pace.on_audio_emit(t, 1.0) + _feed(pace, t, quiet_s=0.0, speech_s=2.0) + assert pace.armed is False + + +def test_the_replay_keeps_the_latest_units_of_the_arming_utterance() -> None: + pace = _pace() + pace.on_input_unit({"seq": "before"}) # no utterance running: not kept + t = _feed(pace, 1.2, quiet_s=0.8, speech_s=0.2) + for index in range(8): + pace.on_input_unit({"seq": index}) + assert pace.take_replay() == [] # not armed: nothing to replay + for index in range(8): + pace.on_input_unit({"seq": index}) + _feed(pace, t, quiet_s=0.0, speech_s=0.4) + assert pace.armed is True + assert [unit["seq"] for unit in pace.take_replay()] == [2, 3, 4, 5, 6, 7] + assert pace.take_replay() == [] + + +def test_silence_after_the_armed_utterance_never_pushes_its_start_out_of_the_replay() -> None: + pace = _pace() + t = _feed(pace, 1.2, quiet_s=0.8, speech_s=0.6) + assert pace.armed is True + for index in range(4): + pace.on_input_unit({"seq": f"speech{index}", "is_speech": True}) + _feed(pace, t, quiet_s=1.0, speech_s=0.0) # the utterance ended; the model has not yielded yet + for index in range(4): + pace.on_input_unit({"seq": f"quiet{index}", "is_speech": False}) + # Trailing silence fills the free slots only (the model yields on it); speech still evicts the oldest. + pace.on_input_unit({"seq": "speech4", "is_speech": True}) + assert [unit["seq"] for unit in pace.take_replay()] == [ + "speech1", + "speech2", + "speech3", + "quiet0", + "quiet1", + "speech4", + ] + + +def test_take_replay_forgets_the_arming_of_a_reply_that_already_ended() -> None: + """The route-M listen abort takes the replay without asking ``cut_ready``: a stale arming is not replayed.""" + pace = _pace() + _feed(pace, 1.2, quiet_s=0.8, speech_s=0.6) + pace.on_input_unit({"seq": 0}) + assert pace.armed is True + pace.session.active_response_id = "resp_2" # a new reply opened; nothing synced the pacing since + assert pace.take_replay() == [] + assert pace.armed is False + + +def test_cut_needs_unplayed_audio_and_disarms_on_epoch_and_new_reply() -> None: + pace = _pace() + t = _feed(pace, 1.2, quiet_s=0.8, speech_s=0.6) + assert pace.armed is True + assert pace.cut_ready(100.0) is False # everything played out + pace.on_epoch_advance(t) + assert pace.armed is False + pace = _pace() + _feed(pace, 1.2, quiet_s=0.8, speech_s=0.6) + pace.session.active_response_id = "resp_2" + assert pace.cut_ready(2.0) is False + assert pace.armed is False + + +async def _armed_harness(**runtime: object): + config: dict[str, object] = {"barge_cut_on_model_yield": True, "barge_arm_min_playback_s": 0.0} + config.update(runtime) + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(**config)) + await h.run(append_audio()) + request_id = h.stage0_request_id() + spoken = await h.deliver_and_settle(tts_output(request_id, samples=240000, text="hello", tts_is_last_chunk=True)) + response_id = find(spoken, "response.created").response_id + await h.run(append_audio(12800, value=0.0, is_speech=False)) + await h.run(append_audio(8000, value=0.05, is_speech=True)) + return h, request_id, response_id + + +def _unit_audio(submission) -> str: + return submission.prompt["model_intermediate_buffer"]["duplex"]["payload"]["audio"] + + +def _epoch_units(h, epoch: int) -> list[str]: + """The audio of every Stage-0 append submitted in ``epoch``, in submission order.""" + return [_unit_audio(sub) for sub in h.port.submissions if f".e.{epoch}." in sub.context.request_id] + + +def _samples(audio: str) -> np.ndarray: + return np.frombuffer(base64.b64decode(audio), dtype=" None: + h, request_id, response_id = await _armed_harness() + try: + assert h.runner.run.pace.armed is True + events = await _model_listens(h, request_id) + cancelled = find(events, "response.done") + assert cancelled.status == "cancelled" + assert cancelled.response_id == response_id + assert h.session.epoch == 1 + assert h.port.aborts and h.port.aborts[0] == [request_id] + assert h.runner.run.pace.cuts == 1 + assert h.runner.run.pace.armed is False + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_the_cut_replays_the_interrupting_utterance_into_the_new_epoch() -> None: + """Stage 0 drops the session context with the aborted request: the utterance's units go again.""" + h, request_id, _ = await _armed_harness() + try: + # The 0.5 s of speech completed a unit (0.8 s quiet + 0.2 s speech) Stage 0 took in epoch 0. + assert len(h.runner.run.pace._utterance_units) == 1 + await _model_listens(h, request_id) + assert h.session.epoch == 1 + assert h.runner.tasks.append_tail is not None + assert await h.runner.tasks.append_tail + new_epoch = [sub for sub in h.port.submissions if sub.context.request_id != request_id] + assert len(new_epoch) == 1 + assert ".e.1." in new_epoch[0].context.request_id + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_natural_turn_end_cuts_an_armed_reply() -> None: + h, request_id, response_id = await _armed_harness() + try: + natural = SimpleNamespace( + request_id=request_id, + finished=True, + outputs=[SimpleNamespace(text="hello", token_ids=[], multimodal_output={})], + multimodal_output={"meta.turn_eos_token_id": 9}, + ) + h.deliver( + natural, + stage_id=0, + segment_finished=True, + segment_token_ids=[9, 10], + segment_output_metadata={"meta.turn_eos_token_id": 9}, + ) + events = await h.settle() + cancelled = find(events, "response.done") + assert cancelled.status == "cancelled" + assert cancelled.response_id == response_id + assert h.session.epoch == 1 + tail = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="tail"), epoch=0) + assert "response.output_audio.delta" not in types(tail) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_unarmed_reply_drains_as_before() -> None: + # Speech straight after the reply opened (no quiet run): not a new utterance. + h = await open_harness( + runtime_config=DuplexSessionRuntimeConfig(barge_cut_on_model_yield=True, barge_arm_min_playback_s=0.0) + ) + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=240000, text="hello", tts_is_last_chunk=True)) + response_id = h.session.active_response_id + await h.run(append_audio(8000, value=0.05, is_speech=True)) + assert h.runner.run.pace.armed is False + await _model_listens(h, request_id) + assert h.port.aborts == [] + assert h.session.epoch == 0 + assert h.session.active_response_id == response_id + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_switch_off_drains_the_same_utterance_as_before() -> None: + """The armed sequence with the switch off: no pacing state, the listen drains the reply.""" + h, request_id, response_id = await _armed_harness(barge_cut_on_model_yield=False) + try: + assert h.runner.run.pace is None + events = await _model_listens(h, request_id) + assert h.port.aborts == [] + assert h.session.epoch == 0 + assert "audio.cancelled" not in types(events) + assert h.session.active_response_id == response_id + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_the_replay_goes_before_the_buffered_tail_and_the_tail_goes_once() -> None: + """Audio still buffered at the cut is no unit yet: not replayed, it opens the new epoch after the replay.""" + h, request_id, _ = await _armed_harness() + try: + replay = [unit["audio"] for unit in h.runner.run.pace._utterance_units] + assert h.runner.model_state.audio_buffer.pending_byte_count == 4800 * 4 # 0.3 s of the interruption + await _model_listens(h, request_id) + assert h.session.epoch == 1 + await h.run(append_audio(11200, value=0.07)) + units = _epoch_units(h, 1) + assert units[: len(replay)] == replay + assert len(units) == len(replay) + 1 + tail = _samples(units[-1]) + assert tail.size == 16000 + assert np.allclose(tail[:4800], 0.05) and np.allclose(tail[4800:], 0.07) + assert h.session.pending_input_bytes == 0 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_the_replay_includes_the_unit_a_commit_flushed_before_the_cut() -> None: + """An explicit commit flushes the utterance's last, zero-padded unit into the old epoch: the cut replays it.""" + h, request_id, _ = await _armed_harness() + try: + events = await h.run(commands.Commit()) + assert "error" not in types(events) + flushed = _epoch_units(h, 0)[-1] + samples = _samples(flushed) + assert samples.size == 16000 + assert np.allclose(samples[:4800], 0.05) and not samples[4800:].any() + replay = [unit["audio"] for unit in h.runner.run.pace._utterance_units] + assert len(replay) == 2 and replay[-1] == flushed + await _model_listens(h, request_id) + assert h.session.epoch == 1 + assert h.runner.tasks.append_tail is not None and await h.runner.tasks.append_tail + new_epoch = [sub for sub in h.port.submissions if ".e.1." in sub.context.request_id] + assert [_unit_audio(sub) for sub in new_epoch] == replay + last = new_epoch[-1].prompt["model_intermediate_buffer"]["duplex"] + assert last["final"] is False and "final" not in last["payload"] + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_units_still_in_the_append_chain_at_the_cut_reach_the_new_epoch_once() -> None: + """Units planned in the old epoch but not accepted there are dropped by the epoch check: the replay delivers them.""" + h, request_id, _ = await _armed_harness() + try: + replay, _ = await _cut_with_an_append_in_flight(h, request_id) + h.port.submit_gate.set() + await h.settle() + assert _epoch_units(h, 1) == replay + assert replay[-1] not in _epoch_units(h, 0) # the queued unit never reached the old epoch's Stage 0 + assert h.session.pending_input_bytes == 0 + assert h.runner.model_state.audio_buffer.pending_byte_count == 0 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_an_append_in_flight_at_the_cut_does_not_fail_the_replay_behind_it() -> None: + """The parked append fails its epoch check after submit: dropped like a queued one, not a runtime failure. + + Failing it rolled its PCM back and stopped the append chain, so the whole + replay behind it was abandoned silently, the client got a non-retryable + ``runtime_append_failed`` and its next commit ``commit_aborted``. + """ + h, request_id, _ = await _armed_harness() + try: + replay, events = await _cut_with_an_append_in_flight(h, request_id) + h.port.submit_gate.set() + events += await h.settle() + assert "error" not in types(events) + tail = h.runner.tasks.append_tail + assert tail is not None and tail.done() and tail.result() is True + assert _epoch_units(h, 1) == replay + events = await h.run(commands.Commit()) + assert "error" not in types(events) + assert "input_audio_buffer.committed" in types(events) + assert h.session.state == DuplexSessionState.OPEN + finally: + await close_harness(h) + + +async def _cut_after_a_commit(h, request_id: str) -> list[str]: + """Armed, the client commits (flushing its buffer) and stops sending; the model listens: the cut.""" + await h.run(commands.Commit()) + replay = [unit["audio"] for unit in h.runner.run.pace._utterance_units] + await _model_listens(h, request_id) + assert h.session.epoch == 1 + assert _epoch_units(h, 1) == replay + return replay + + +@pytest.mark.asyncio +async def test_a_quiet_client_still_gets_silence_units_after_the_cut() -> None: + """The cut reply's continuations kept Stage 0 of a quiet client moving; the replayed epoch has none. + + Without them nothing would submit another unit after the replay's + listen: no answer until the client sends audio again. + """ + h, request_id, _ = await _armed_harness() + try: + replay = await _cut_after_a_commit(h, request_id) + events = await _model_listens(h, h.stage0_request_id()) # Stage 0 listens on the replay + assert "response.listen" in types(events) + units = _epoch_units(h, 1) + assert len(units) == len(replay) + 1 + assert not _samples(units[-1]).any() + assert h.runner.model_state.continuation_units == 1 + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_a_client_that_keeps_streaming_after_the_cut_gets_no_silence_units() -> None: + h, request_id, _ = await _armed_harness() + try: + replay = await _cut_after_a_commit(h, request_id) + await h.run(append_audio(16000, value=0.0, is_speech=False)) # the microphone stays on + await _model_listens(h, h.stage0_request_id()) + assert len(_epoch_units(h, 1)) == len(replay) + 1 # the client's own unit, no silence after it + assert h.runner.model_state.continuation_units == 0 + finally: + await close_harness(h) diff --git a/tests/engine/duplex/test_duplex_intermediate.py b/tests/engine/duplex/test_duplex_intermediate.py index d542f3cabfb..713bd45be19 100644 --- a/tests/engine/duplex/test_duplex_intermediate.py +++ b/tests/engine/duplex/test_duplex_intermediate.py @@ -2,8 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest +import torch -from vllm_omni.engine.duplex.intermediate import get_stream_request_key +from vllm_omni.engine.duplex.intermediate import ( + drop_static_append_configs, + get_stream_request_key, + get_tts_handoff, + pack_transport_tensor, + set_ref_audio, + set_tts_handoff, + unpack_transport_tensor, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -15,3 +24,110 @@ def test_get_stream_request_key_requires_stable_identifier(): def test_get_stream_request_key_accepts_global_request_id(): assert get_stream_request_key({"global_request_id": ["duplex-sid-stage0"]}) == "duplex-sid-stage0" + + +# ── Tensor transport through model_intermediate_buffer ─────────────── + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + """Reinterpret as integers so equality is bit-exact (NaN, -0.0).""" + integer = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}[tensor.element_size()] + return tensor.contiguous().view(integer) + + +@pytest.mark.parametrize( + "tensor", + [ + torch.randn(7, 16), + torch.randn(5, 8).to(torch.bfloat16), + torch.arange(12, dtype=torch.int64), + torch.tensor(3.5), + torch.empty(0, 16), + torch.randn(8, 6).t(), # non-contiguous + torch.tensor([float("nan"), -0.0, float("inf"), 1e-45]), + ], +) +def test_packed_tensor_round_trip_is_bit_exact(tensor): + packed = pack_transport_tensor(tensor) + assert isinstance(packed["data"], bytes) + restored = unpack_transport_tensor(packed) + assert restored.dtype == tensor.dtype + assert restored.shape == tensor.shape + assert torch.equal(_bits(restored), _bits(tensor)) + # Writable, independent storage: the Talker may mutate it freely. + if restored.numel(): + restored.reshape(-1)[0] = 0 + assert torch.equal(_bits(unpack_transport_tensor(packed)), _bits(tensor)) + + +def test_non_tensor_values_pass_through_unchanged(): + rows = [[0.25, -0.5]] + assert pack_transport_tensor(rows) is rows + assert unpack_transport_tensor(rows) is rows + assert unpack_transport_tensor(None) is None + plain = {"dtype": "float32", "shape": [1], "data": b"\0\0\0\0"} + assert unpack_transport_tensor(plain) is plain + + +def test_drop_static_append_configs_keeps_runtime_fields_and_payload(): + prompt = { + "prompt_token_ids": [0, 0], + "model_intermediate_buffer": { + "duplex": { + "payload": {"audio": "unit"}, + "scheduler_token_budget": 2, + "session_config": {"conversation": ["old"]}, + "runtime_config": {"ref_audio_data": "large"}, + } + }, + } + + assert drop_static_append_configs(prompt) is prompt + duplex = prompt["model_intermediate_buffer"]["duplex"] + assert duplex == {"payload": {"audio": "unit"}, "scheduler_token_budget": 2} + + +def test_packed_handoff_matches_the_legacy_list_transport(): + """The old path shipped ``tolist()`` and rebuilt with ``as_tensor``.""" + hidden = torch.randn(9, 32) + legacy = torch.as_tensor(hidden.tolist(), dtype=torch.float32) + buffer: dict = {} + set_tts_handoff(buffer, [1, 2, 3], hidden) + token_ids, restored = get_tts_handoff(buffer) + assert token_ids == [1, 2, 3] + assert torch.equal(_bits(restored), _bits(legacy)) + + +def test_legacy_list_handoff_is_still_readable(): + buffer = {"ids": {"tts": [5]}, "hidden_states": {"tts": [[0.5, 1.5]]}} + assert get_tts_handoff(buffer) == ([5], [[0.5, 1.5]]) + assert get_tts_handoff({"tts_token_ids": [4], "tts_hidden_states": [[2.0]]}) == ([4], [[2.0]]) + + +def test_packed_handoff_survives_engine_core_request_msgpack(): + """Round-trip the exact orchestrator -> Stage-1 wire format.""" + from vllm import SamplingParams + from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder + + from vllm_omni.engine import OmniEngineCoreRequest + from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens + + hidden = torch.randn(33, 64) + ref = torch.randn(4000) + buffer: dict = {"global_request_id": ["req"], "meta": {"next_stage_prompt_len": 35}} + set_tts_handoff(buffer, list(range(33)), hidden) + set_ref_audio(buffer, ref, 16000) + request = build_engine_core_request_from_tokens( + request_id="req", + prompt={"prompt_token_ids": [0] * 35, "model_intermediate_buffer": buffer}, + params=SamplingParams(max_tokens=4), + ) + + decoded = MsgpackDecoder(OmniEngineCoreRequest).decode(MsgpackEncoder().encode(request)) + + info = decoded.model_intermediate_buffer + token_ids, restored = get_tts_handoff(info) + assert token_ids == list(range(33)) + assert torch.equal(_bits(restored), _bits(hidden)) + assert torch.equal(_bits(unpack_transport_tensor(info["codes"]["ref"])), _bits(ref)) + assert info["meta"] == {"next_stage_prompt_len": 35, "ref_audio_sr": 16000} diff --git a/tests/engine/duplex/test_duplex_onset_lead.py b/tests/engine/duplex/test_duplex_onset_lead.py new file mode 100644 index 00000000000..5c8cf0ee597 --- /dev/null +++ b/tests/engine/duplex/test_duplex_onset_lead.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Paced onset (``duplex_session.pacing``) through the production continuation scheduler. + +Same fake clock as ``test_deadline_pacing_scheduler``: ``asyncio.sleep`` +advances it, so the submission time of each continuation is exact. +""" + +from __future__ import annotations + +import asyncio +import time + +import pytest + +import vllm_omni.engine.duplex.session.runner as runner_module +from tests.engine.duplex.test_deadline_pacing_scheduler import ( + FakeClock, + _active_response_harness, + _continuation_kwargs, + _install_fake_clock, +) +from tests.engine.duplex.test_session_runner import ( + LISTEN_TOKEN_ID, + append_audio, + close_harness, + find, + listen_output, + open_harness, + tts_output, + types, +) +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex.session.pacing import PaceState, SessionPacing + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _paced(h, **pacing: object) -> SessionPacing: + """Install a pacing state on an open harness (the runtime config of ``open_harness`` has none).""" + config: dict[str, object] = {"enabled": True, "onset_lead_max_s": 0.9} + config.update(pacing) + pace = SessionPacing(DuplexSessionRuntimeConfig(pacing=config), h.session) + h.runner.run.pace = pace + return pace + + +async def _continue(h, **kwargs: object) -> None: + scheduled = await h.runner._schedule_silence_continuation( + h.runner.model.silence_unit_payload(), + **_continuation_kwargs(h), + **kwargs, + ) + assert scheduled is True + task = h.runner.tasks.append_tail + assert task is not None + assert await task + + +@pytest.mark.asyncio +async def test_pacing_is_off_by_default() -> None: + h = await open_harness() + try: + assert h.runner.run.pace is None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_runner_creates_pacing_when_enabled() -> None: + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(pacing={"enabled": True})) + try: + assert isinstance(h.runner.run.pace, SessionPacing) + assert h.runner.run.pace.state == PaceState.LISTEN + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_quiet_client_submits_ahead_of_the_nominal_cadence(monkeypatch: pytest.MonkeyPatch) -> None: + clock = FakeClock(start=100.0) + _install_fake_clock(monkeypatch, clock=clock, real_sleep=asyncio.sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + pace = _paced(h) + pace.on_commit(100.0) # the harness set-up advanced the fake clock + assert pace.state == PaceState.ONSET + # Unit 1 submitted at 100.0; its segment ends at 100.2, inside the + # 0.25 s quiet guard: unit 2 goes once the guard passes, not at 101.0, + # and the nominal chain is untouched (next deadline 102.0). + state.last_native_submit_monotonic = 100.0 + state.silence_deadline_monotonic = None + clock.value = 100.2 + await _continue(h) + assert state.last_native_submit_monotonic == pytest.approx(100.25) + assert state.silence_deadline_monotonic == pytest.approx(102.0) + + # The first chunk was 0.84 s: lead 0.26, unit 3 goes at 102 - 0.26. + pace.on_audio_emit(100.5, 0.84) + clock.value = 100.6 + await _continue(h) + assert state.last_native_submit_monotonic == pytest.approx(101.74) + assert state.silence_deadline_monotonic == pytest.approx(103.0) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_lead_rechecked_while_waiting_only_delays(monkeypatch: pytest.MonkeyPatch) -> None: + """d1 arriving mid-wait shrinks the lead: the unit fires later, never earlier.""" + clock = FakeClock(start=200.0) + _install_fake_clock(monkeypatch, clock=clock, real_sleep=asyncio.sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + pace = _paced(h) + pace.on_commit(200.0) + state.last_native_submit_monotonic = 200.0 + state.silence_deadline_monotonic = 202.0 + clock.value = 200.5 + real_speech_fire_at = pace.speech_fire_at + calls = {"n": 0} + + def speech_fire_at(nominal: float, now: float) -> float: + calls["n"] += 1 + if calls["n"] == 3: + # d1 lands while the unit waits for 202.0 - 0.9. + pace.on_audio_emit(now, 1.0) + return real_speech_fire_at(nominal, now) + + monkeypatch.setattr(pace, "speech_fire_at", speech_fire_at) + await _continue(h) + assert state.last_native_submit_monotonic == pytest.approx(201.9) + assert state.silence_deadline_monotonic == pytest.approx(203.0) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_an_append_inside_the_quiet_guard_keeps_unit_2_on_the_wall_clock( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A mic-open client (200 ms chunks) whose next append lands after unit 1's segment ended.""" + clock = FakeClock(start=99.0) + _install_fake_clock(monkeypatch, clock=clock, real_sleep=asyncio.sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + pace = _paced(h) + for t in (99.4, 99.6, 99.8): + pace.on_client_append(t, is_speech=True, duration_s=0.2) + pace.on_commit(100.0) # guard 1.5 x 0.2 = 0.3 s + state.last_native_submit_monotonic = 100.0 + state.silence_deadline_monotonic = None + clock.value = 100.15 + real_speech_fire_at = pace.speech_fire_at + + def speech_fire_at(nominal: float, now: float) -> float: + if now >= 100.2 and pace.quiet_since_commit: + pace.on_client_append(now, is_speech=False, duration_s=0.2) + return real_speech_fire_at(nominal, now) + + monkeypatch.setattr(pace, "speech_fire_at", speech_fire_at) + await _continue(h) + assert pace.quiet_since_commit is False + assert state.last_native_submit_monotonic == pytest.approx(101.0) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_input_after_commit_restores_the_wall_clock_cadence(monkeypatch: pytest.MonkeyPatch) -> None: + clock = FakeClock(start=300.0) + _install_fake_clock(monkeypatch, clock=clock, real_sleep=asyncio.sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + pace = _paced(h) + pace.on_commit(clock.value) + pace.on_client_append(clock.value, is_speech=False, duration_s=0.2) + assert pace.lead_s(clock.value + 1.0) == 0.0 + state.last_native_submit_monotonic = 300.0 + state.silence_deadline_monotonic = None + clock.value = 300.2 + await _continue(h) + assert state.last_native_submit_monotonic == pytest.approx(301.0) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_runner_reports_client_appends_to_the_pacing_state() -> None: + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(pacing={"enabled": True})) + try: + pace = h.runner.run.pace + assert isinstance(pace, SessionPacing) + pace.on_commit(0.0) + assert pace.quiet_since_commit is True + await h.run(append_audio(3200, value=0.0, is_speech=False)) + assert pace.quiet_since_commit is False + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_idle_continuations_keep_the_nominal_deadline(monkeypatch: pytest.MonkeyPatch) -> None: + clock = FakeClock(start=400.0) + _install_fake_clock(monkeypatch, clock=clock, real_sleep=asyncio.sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + pace = _paced(h) + pace.on_commit(clock.value) + state.last_native_submit_monotonic = 400.0 + state.silence_deadline_monotonic = None + clock.value = 400.2 + await _continue(h, pace_kind="idle") + assert state.last_native_submit_monotonic == pytest.approx(401.0) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_same_turn_response_does_not_stale_a_paced_model_turn_unit() -> None: + h = await _active_response_harness() + try: + session = h.session + # The response that opened while unit 2 waited is the unit's own turn's. + session.bind_response_turn(session.turn_id) + turn_id = session.turn_id + kwargs = { + "request_id": session.active_request_id, + "response_id": None, + "response_owned": False, + "expected_epoch": session.epoch, + "expected_model_turn_id": turn_id, + } + # Historical rule: any active response makes a model-turn continuation stale. + assert h.runner.model.silence_continuation_is_stale(**kwargs) is True + _paced(h) # skip switch off: unchanged + assert h.runner.model.silence_continuation_is_stale(**kwargs) is True + _paced(h, onset_skip_response_wait=True) + assert h.runner.model.silence_continuation_is_stale(**kwargs) is False + # A response of another turn still makes it stale. + assert h.runner.model.silence_continuation_is_stale(**{**kwargs, "expected_model_turn_id": turn_id + 1}) is True + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stop_token_trigger_skips_the_response_wait_on_a_paced_onset(monkeypatch: pytest.MonkeyPatch) -> None: + h = await open_harness() + try: + session = h.session + session.bind_request(h.stage0_request_id()) + assert session.active_response_id is None + calls: list[dict] = [] + + async def fake_continue(**kwargs) -> None: + calls.append(kwargs) + + monkeypatch.setattr(h.runner.model, "maybe_continue_response", fake_continue) + sleeps: list[float] = [] + real_sleep = asyncio.sleep + + async def counting_sleep(delay: float) -> None: + sleeps.append(delay) + await real_sleep(0) + + monkeypatch.setattr(runner_module.asyncio, "sleep", counting_sleep) + pace = _paced(h, onset_skip_response_wait=True) + pace.on_commit(time.monotonic() - 1.0) # past the quiet guard + await h.runner._continue_response_on_stop_token_task(session.epoch) + assert sleeps == [] + assert calls == [{"expected_epoch": session.epoch, "expected_model_turn_id": session.turn_id}] + finally: + monkeypatch.undo() + await close_harness(h) + + +@pytest.mark.asyncio +async def test_paced_unit_snaps_to_the_fire_grid(monkeypatch: pytest.MonkeyPatch) -> None: + clock = FakeClock(start=100.0) + _install_fake_clock(monkeypatch, clock=clock, real_sleep=asyncio.sleep) + h = await _active_response_harness() + state = h.runner.model_state + try: + pace = _paced(h, fire_grid_ms=250) + pace.on_commit(100.0) + pace.on_audio_emit(100.5, 0.84) + assert pace.state == PaceState.PACED + state.last_native_submit_monotonic = 101.0 + state.silence_deadline_monotonic = 102.0 + clock.value = 100.6 + await _continue(h) + # Lead 0.26 gives 101.74; the grid moves it to 101.5 (lead 0.5 <= 0.9). + assert state.last_native_submit_monotonic == pytest.approx(101.5) + assert state.silence_deadline_monotonic == pytest.approx(103.0) + finally: + await close_harness(h) + + +async def _speaking_quiet_client(pacing: dict[str, object]): + """A committed client 1 whose reply is speaking (TTS in flight), with ``abort_on_model_listen``.""" + h = await open_harness( + runtime_config=DuplexSessionRuntimeConfig(abort_on_model_listen=True, pacing=pacing), + ) + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + assert h.session.active_response_id is not None + return h, request_id + + +async def _model_listens(h, request_id: str): + return await h.deliver_and_settle( + listen_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + +@pytest.mark.asyncio +async def test_a_quiet_clients_model_listen_does_not_abort_the_paced_reply() -> None: + """Unit 2 decides listen while unit 1 is still synthesized: no barge-in, the reply drains.""" + h, request_id = await _speaking_quiet_client({"enabled": True, "onset_lead_max_s": 0.9}) + try: + h.runner.run.pace.on_commit(time.monotonic()) + response_id = h.session.active_response_id + events = await _model_listens(h, request_id) + assert h.port.aborts == [] + assert h.session.epoch == 0 + assert "audio.cancelled" not in types(events) + assert all(event.status != "cancelled" for event in events if event.type == "response.done") + assert h.session.active_response_id == response_id + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_a_streaming_clients_model_listen_still_aborts() -> None: + h, request_id = await _speaking_quiet_client({"enabled": True, "onset_lead_max_s": 0.9}) + try: + pace = h.runner.run.pace + pace.on_commit(time.monotonic()) + pace.on_client_append(time.monotonic(), is_speech=True, duration_s=0.2) # input after the commit + events = await _model_listens(h, request_id) + assert find(events, "response.done").status == "cancelled" + assert h.session.epoch == 1 + finally: + await close_harness(h) diff --git a/tests/engine/duplex/test_duplex_pacing.py b/tests/engine/duplex/test_duplex_pacing.py new file mode 100644 index 00000000000..562f2b73bc1 --- /dev/null +++ b/tests/engine/duplex/test_duplex_pacing.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``duplex_session.pacing``: the playback mirror, the session pacing state machine and its lead.""" + +from __future__ import annotations + +import random +from types import SimpleNamespace + +import pytest + +from vllm_omni.config.stage_config import DuplexPacingConfig, DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex.session.pacing import ( + QUIET_GUARD_MIN_S, + PaceState, + PlaybackMirror, + SessionPacing, + same_turn_response_keeps_continuation, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +#: Past the quiet guard of a commit at 0.0 with no appends before it. +QUIET = QUIET_GUARD_MIN_S + + +def _reference_playback(chunks: list[tuple[float, float]], buffer_s: float) -> tuple[float, float]: + """duplex_rt / rtf_load ``playback``: (stall, play end) from the first arrival plus ``buffer_s``.""" + play_end = chunks[0][0] + buffer_s + chunks[0][1] + stall = 0.0 + for arrival, duration in chunks[1:]: + if arrival > play_end: + stall += arrival - play_end + play_end = arrival + play_end += duration + return stall, play_end + + +def _pacing(**pacing: object) -> SessionPacing: + defaults: dict[str, object] = {"enabled": True, "onset_lead_max_s": 0.9} + defaults.update(pacing) + session = SimpleNamespace(active_response_id=None) + return SessionPacing(DuplexSessionRuntimeConfig(pacing=defaults), session) + + +def _emit(pace: SessionPacing, now: float, seconds: float, response_id: str = "resp_1") -> None: + """One audio delta of ``response_id`` (the session's active response while it plays).""" + pace.session.active_response_id = response_id + pace.on_audio_emit(now, seconds) + + +@pytest.mark.parametrize("seed", range(20)) +def test_mirror_matches_the_duplex_rt_playback_rule(seed: int) -> None: + rng = random.Random(seed) + mirror = PlaybackMirror(prebuffer_s=0.5) + t = 10.0 + chunks: list[tuple[float, float]] = [] + stall = 0.0 + for _ in range(rng.randint(1, 30)): + t += rng.uniform(0.0, 1.6) + duration = rng.choice([0.12, 0.5, 1.0, 1.12]) + chunks.append((t, duration)) + stall += mirror.on_emit(t, duration) + want_stall, want_end = _reference_playback(chunks, 0.5) + assert stall == pytest.approx(want_stall) + assert mirror.play_end == pytest.approx(want_end) + + +def test_mirror_buffers_once_per_session_and_cuts() -> None: + mirror = PlaybackMirror() + assert mirror.prebuffer_s == 0.5 + assert mirror.unplayed(0.0) == 0.0 + mirror.on_emit(1.0, 0.12) + assert mirror.play_end == pytest.approx(1.62) + assert mirror.unplayed(1.0) == pytest.approx(0.62) + mirror.cut(1.2) + assert mirror.play_end == pytest.approx(1.2) + assert mirror.response_first_audio_t is None + # A later response's first chunk re-anchors without a prebuffer and is not a stall. + assert mirror.on_emit(5.0, 1.0) == 0.0 + assert mirror.play_end == pytest.approx(6.0) + # Inside a response a late chunk is a stall. + assert mirror.on_emit(6.5, 1.0) == pytest.approx(0.5) + assert mirror.response_stall_s == pytest.approx(0.5) + + +@pytest.mark.parametrize(("d1", "lead"), [(0.04, 0.9), (0.12, 0.9), (0.84, 0.26), (1.0, 0.1), (1.12, 0.0)]) +def test_lead_follows_the_first_chunk(d1: float, lead: float) -> None: + pace = _pacing() + pace.on_commit(0.0) + assert pace.state == PaceState.ONSET + # d1 unknown: the full lead, so unit 2 fires as soon as unit 1's segment ends. + assert pace.lead_s(QUIET) == pytest.approx(0.9) + _emit(pace, 0.4, d1) + assert pace.state == PaceState.PACED + assert pace.lead_s(0.5) == pytest.approx(lead) + + +def test_state_machine_critical_and_recovery() -> None: + pace = _pacing() + pace.on_commit(0.0) + _emit(pace, 0.3, 0.84) # play end 0.3 + 0.5 + 0.84 = 1.64 + assert pace.lead_s(0.5) == pytest.approx(0.26) + _emit(pace, 1.75, 1.0) # 0.11 s stall + assert pace.state == PaceState.CRITICAL + assert pace.lead_s(1.8) == pytest.approx(0.9) + _emit(pace, 2.0, 1.0) + assert pace.state == PaceState.CRITICAL + _emit(pace, 2.5, 1.0) + assert pace.state == PaceState.PACED + + +def test_the_quiet_guard_holds_the_lead_back_right_after_the_commit() -> None: + pace = _pacing(onset_skip_response_wait=True) + pace.on_commit(10.0) + assert pace.lead_s(10.1) == 0.0 + assert pace.skip_response_wait(10.1) is False + assert pace.recheck_in(10.2) == pytest.approx(QUIET_GUARD_MIN_S - 0.2) # wakes when the guard passes + assert pace.lead_s(10.0 + QUIET) == pytest.approx(0.9) + assert pace.skip_response_wait(10.0 + QUIET) is True + + +def test_a_long_chunk_client_with_its_mic_open_is_not_taken_for_quiet() -> None: + """Client 2 pushing 600 ms chunks: its next append after the commit lands inside the guard.""" + pace = _pacing(onset_skip_response_wait=True) + for index in range(5): + pace.on_client_append(10.0 + 0.6 * index, is_speech=True, duration_s=0.6) + pace.on_commit(12.5) + assert pace.lead_s(12.8) == 0.0 # guard 1.5 x 0.6 = 0.9 s + assert pace.skip_response_wait(12.8) is False + assert pace.lead_s(13.45) == pytest.approx(0.9) # a client that really went quiet still leads + # The pause before the next utterance is no chunk interval: that commit uses the floor. + pace.on_client_append(20.0, is_speech=True, duration_s=0.2) + pace.on_client_append(20.2, is_speech=True, duration_s=0.2) + pace.on_commit(20.4) + assert pace.lead_s(20.4 + 0.35) == pytest.approx(0.9) # guard 1.5 x 0.2 = 0.3 s + + +def test_any_input_after_commit_returns_the_reply_to_the_wall_clock() -> None: + pace = _pacing(onset_skip_response_wait=True) + pace.on_commit(0.0) + assert pace.skip_response_wait(QUIET) is True + pace.on_client_append(0.3, is_speech=False, duration_s=0.2) + assert pace.lead_s(0.4) == 0.0 + assert pace.skip_response_wait(0.4) is False + _emit(pace, 0.4, 0.12) + assert pace.lead_s(0.5) == 0.0 + # The next commit makes the client quiet again. + pace.session.active_response_id = None + pace.on_commit(5.0) + assert pace.lead_s(5.0 + QUIET) == pytest.approx(0.9) + + +def test_a_reply_ended_by_any_path_resets_the_next_turn_to_onset() -> None: + """Turn 1 ends with a listen (no end_of_turn hook): turn 2's commit still starts at ONSET.""" + pace = _pacing(onset_skip_response_wait=True) + pace.on_commit(0.0) + _emit(pace, 0.4, 0.5) + assert pace.state == PaceState.PACED + pace.session.active_response_id = None # ended: completed, listen, drained ... + assert pace.lead_s(3.0) == 0.0 + assert pace.state == PaceState.LISTEN + pace.on_commit(5.0) + assert pace.state == PaceState.ONSET and pace.d1 is None + assert pace.skip_response_wait(5.0 + QUIET) is True + assert pace.lead_s(5.0 + QUIET) == pytest.approx(0.9) + + +def test_a_draining_replys_tail_plays_but_is_not_the_new_replys_first_chunk() -> None: + pace = _pacing() + pace.on_commit(0.0) + pace.session.active_response_id = "resp_2" + pace.on_audio_emit(0.3, 1.0, draining=True) # resp_1's tail, still draining + assert pace.d1 is None and pace.state == PaceState.ONSET + assert pace.mirror.play_end == pytest.approx(0.3 + 0.5 + 1.0) + pace.on_audio_emit(0.4, 0.12) + assert pace.d1 == pytest.approx(0.12) and pace.state == PaceState.PACED + assert pace.mirror.play_end == pytest.approx(1.92) + + +def test_a_reply_that_spoke_before_the_commit_stays_paced() -> None: + pace = _pacing() + _emit(pace, 0.4, 0.84) # Stage-0 handoff before the client's commit + pace.on_commit(1.0) + assert pace.state == PaceState.PACED + assert pace.lead_s(1.0 + QUIET) == pytest.approx(0.26) + + +def test_skip_response_wait_needs_onset_and_its_switch() -> None: + assert _pacing().skip_response_wait(QUIET) is False + pace = _pacing(onset_skip_response_wait=True) + assert pace.skip_response_wait(QUIET) is False # no commit yet + pace.on_commit(0.0) + assert pace.skip_response_wait(QUIET) is True + _emit(pace, 0.3, 0.5) + assert pace.skip_response_wait(0.4) is False # PACED + + +def test_a_quiet_client_keeps_its_reply_on_a_model_listen() -> None: + pace = _pacing() + assert pace.keeps_reply_on_listen() is False # no commit: a listen may be a barge-in + pace.on_commit(0.0) + assert pace.keeps_reply_on_listen() is True # also inside the guard: no input since the commit + pace.on_client_append(0.5, is_speech=True, duration_s=0.2) + assert pace.keeps_reply_on_listen() is False + off = _pacing(enabled=False) + off.on_commit(0.0) + assert off.keeps_reply_on_listen() is False + + +def test_epoch_advance_resets_to_listen_and_cuts_the_mirror() -> None: + pace = _pacing() + pace.on_commit(0.0) + _emit(pace, 1.0, 1.0) + pace.session.active_response_id = None + pace.on_epoch_advance(1.2) + assert pace.state == PaceState.LISTEN + assert pace.lead_s(1.3) == 0.0 + assert pace.mirror.unplayed(1.2) == 0.0 + + +def test_lead_is_zero_when_disabled_or_unset() -> None: + off = _pacing(enabled=False) + off.on_commit(0.0) + assert off.lead_s(QUIET) == 0.0 + no_lead = _pacing(onset_lead_max_s=0.0) + no_lead.on_commit(0.0) + assert no_lead.lead_s(QUIET) == 0.0 + assert no_lead.speech_fire_at(10.0, 9.0) == 10.0 + assert no_lead.idle_fire_at(10.0, 9.0) == 10.0 + + +def test_speech_fire_at_subtracts_the_lead() -> None: + pace = _pacing() + pace.on_commit(0.0) + assert pace.speech_fire_at(1.0, QUIET) == pytest.approx(0.1) + + +def test_same_turn_keep_needs_the_skip_switch() -> None: + assert same_turn_response_keeps_continuation(None) is False + assert same_turn_response_keeps_continuation(_pacing()) is False + assert same_turn_response_keeps_continuation(_pacing(onset_skip_response_wait=True)) is True + + +def test_pacing_config_defaults_are_off() -> None: + config = DuplexPacingConfig() + assert config.enabled is False + assert config.onset_lead_max_s == 0.0 + assert config.onset_skip_response_wait is False + assert config.fire_grid_ms == 0 + assert config.idle_grid is False + assert DuplexSessionRuntimeConfig().pacing == config + + +def _paced_grid(*, lead_max: float = 0.9, d1: float = 0.84, **pacing: object) -> SessionPacing: + pace = _pacing(onset_lead_max_s=lead_max, fire_grid_ms=250, **pacing) + pace.on_commit(100.0) + _emit(pace, 100.3, d1) + assert pace.state == PaceState.PACED + return pace + + +def test_snap_is_the_identity_without_a_grid() -> None: + pace = _pacing() + pace.on_commit(100.0) + _emit(pace, 100.3, 0.84) + assert pace.snap(101.74, 102.0, 100.6) == 101.74 + assert pace.snap_holds(100.6) is False + + +def test_snap_only_moves_earlier_within_the_lead_budget() -> None: + pace = _paced_grid() + assert pace.snap(101.74, 102.0, 100.6) == pytest.approx(101.5) + assert pace.snap_holds(100.6) is True + # The earlier grid point is already past: unchanged, never rounded up. + assert pace.snap(101.74, 102.0, 101.6) == pytest.approx(101.74) + # 102.03 -> 102.0 would lead 0.13 s > 0.1: unchanged. + tight = _paced_grid(lead_max=0.1, d1=1.0) + assert tight.snap(102.03, 102.13, 100.6) == pytest.approx(102.03) + + +def test_eager_states_never_snap() -> None: + pace = _pacing(fire_grid_ms=250) + pace.on_commit(100.0) + assert pace.state == PaceState.ONSET + assert pace.snap(101.74, 102.0, 100.6) == 101.74 + critical = _paced_grid() + _emit(critical, 102.5, 1.0) # a stall + assert critical.state == PaceState.CRITICAL + assert critical.snap(101.74, 102.0, 100.6) == 101.74 + streaming = _paced_grid() + streaming.on_client_append(100.7, is_speech=False, duration_s=0.2) + assert streaming.snap(101.74, 102.0, 100.8) == 101.74 + assert streaming.snap_holds(100.8) is False + + +def test_idle_grid_rounds_up_by_at_most_one_grid_step() -> None: + pace = _paced_grid(idle_grid=True) + assert pace.idle_fire_at(101.1, 100.6) == pytest.approx(101.25) + assert pace.idle_fire_at(101.25, 100.6) == pytest.approx(101.25) + assert _paced_grid().idle_fire_at(101.1, 100.6) == 101.1 # idle_grid off + pace.on_client_append(101.0, is_speech=True, duration_s=0.2) + assert pace.idle_fire_at(101.1, 101.0) == 101.1 # not a quiet client any more diff --git a/tests/engine/duplex/test_duplex_unit_decouple.py b/tests/engine/duplex/test_duplex_unit_decouple.py new file mode 100644 index 00000000000..62063cdf804 --- /dev/null +++ b/tests/engine/duplex/test_duplex_unit_decouple.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Duplex unit continuation decoupling (VLLM_OMNI_DUPLEX_CONTINUE_ON / _UNIT_DEPTH). + +The historical chain plans a spoken unit's next append only after the unit's +audio deltas have been emitted, which holds each session to one unit in +flight. These tests pin the opt-in behaviours and the invariants that must +not change when the trigger moves earlier: + +* stop_token mode schedules the next unit's append while the current unit's + audio has not even been produced yet, without double-scheduling when the + audio-side trigger later fires for the same unit; +* client audio emission order is untouched (it follows stage-output arrival, + not append submission time) even when a later unit's append is submitted + before an earlier unit's audio is emitted; +* default (audio / depth 1) behaviour matches the historical trigger exactly; +* depth >= 2 lets a second continuation be planned while the first is still + submitting, where depth 1 keeps waiting. +""" + +from __future__ import annotations + +import asyncio +import time +from types import SimpleNamespace +from typing import Any + +import pytest + +from tests.engine.duplex.test_session_runner import ( + LISTEN_TOKEN_ID, + Harness, + append_audio, + close_harness, + open_harness, + tts_output, + types, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +CONTINUE_ON_ENV = "VLLM_OMNI_DUPLEX_CONTINUE_ON" +DEPTH_ENV = "VLLM_OMNI_DUPLEX_UNIT_DEPTH" + + +def spoken_stage0_output(request_id: str) -> SimpleNamespace: + """A finished spoken Stage-0 segment: stopped on the chunk terminator. + + The last token is not the listen token, so the plugin's ``decide_output`` + returns None (forwarded to the TTS stage) -- exactly the output whose stop + token is the decoupled continuation trigger. + """ + return SimpleNamespace( + request_id=request_id, + finished=True, + outputs=[SimpleNamespace(text="", token_ids=[21, 22, 23], stop_reason=23, multimodal_output={})], + multimodal_output={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + +async def _active_response_harness() -> Harness: + """Open a harness and drive one TTS segment so a response is active.""" + h = await open_harness() + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + assert h.session.active_response_id is not None + assert len(h.port.submissions) == 1 + return h + + +def _make_continuation_due_now(h: Harness) -> None: + """Point the silence-continuation anchor far enough back that no sleep is owed.""" + model_state = h.runner.model_state + model_state.last_native_submit_monotonic = time.monotonic() - 1.2 + model_state.silence_deadline_monotonic = None + + +def _continuation_kwargs(h: Harness) -> dict[str, Any]: + return { + "request_id": h.session.active_request_id, + "owner_id": f"response:{h.session.active_response_id}", + "response_id": h.session.active_response_id, + "response_owned": True, + "expected_epoch": h.session.epoch, + "expected_model_turn_id": h.session.turn_id, + } + + +# --------------------------------------------------------------------------- # +# Trigger switch # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_stop_token_mode_schedules_continuation_before_audio_emission( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The Stage-0 segment finish alone must plan the next unit's append. + + At the moment the assertion runs, the spoken unit's audio has not been + produced at all (no Stage-1 output was delivered), yet the next silence + unit is already submitted -- the decoupling this switch exists for. + """ + monkeypatch.setenv(CONTINUE_ON_ENV, "stop_token") + h = await _active_response_harness() + try: + request_id = h.stage0_request_id() + _make_continuation_due_now(h) + + h.deliver( + spoken_stage0_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[21, 22, 23], + ) + events = await h.settle() + + assert not [event for event in events if event.type == "error"], types(events) + assert len(h.port.submissions) == 2, "the stop token did not schedule the next unit" + silence = h.port.submissions[1] + assert silence.already_submitted is True + assert h.runner.model_state.continuation_units == 1 + + # The audio-side trigger for the same unit must not schedule a second + # silence unit when the audio finally arrives. + await h.deliver_and_settle( + tts_output(request_id, samples=48000, text="hello", tts_is_last_chunk=True, finished=True) + ) + assert len(h.port.submissions) == 2, "the audio-emit trigger double-scheduled a unit" + assert h.runner.model_state.continuation_units == 1 + assert h.session.active_response_id is not None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_audio_mode_parity_keeps_the_audio_side_trigger( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Default mode: the Stage-0 segment finish alone schedules nothing. + + Bit-parity with the historical behaviour: only the TTS segment end (audio + emit) plans the next unit. + """ + monkeypatch.delenv(CONTINUE_ON_ENV, raising=False) + h = await _active_response_harness() + try: + request_id = h.stage0_request_id() + _make_continuation_due_now(h) + + h.deliver( + spoken_stage0_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[21, 22, 23], + ) + await h.settle() + assert len(h.port.submissions) == 1, "stop token scheduled a unit in audio mode" + + await h.deliver_and_settle( + tts_output(request_id, samples=48000, text="hello", tts_is_last_chunk=True, finished=True) + ) + assert len(h.port.submissions) == 2, "the audio-emit trigger did not schedule the unit" + assert h.runner.model_state.continuation_units == 1 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stop_token_mode_preserves_client_audio_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Unit k+1's append may be submitted before unit k's audio; order holds. + + The correctness red line of the decoupling: submission order and emission + order are different queues. The test drives two spoken units whose appends + are submitted at their stop tokens, with each unit's audio delivered only + afterwards, and asserts the client still hears unit k's full audio before + unit k+1's first delta. + """ + monkeypatch.setenv(CONTINUE_ON_ENV, "stop_token") + h = await _active_response_harness() + try: + request_id = h.stage0_request_id() + + # Unit 2's stop token: its continuation (unit 3's append) is submitted + # before unit 2's audio exists. + _make_continuation_due_now(h) + h.deliver( + spoken_stage0_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[21, 22, 23], + ) + await h.settle() + assert len(h.port.submissions) == 2 + + # Unit 2's audio arrives late; the client sees its delta now. (The + # data plane slices cumulative audio, so each unit's snapshot must + # grow past the previous unit's offset.) + events_unit2 = await h.deliver_and_settle(tts_output(request_id, samples=48000, text="second")) + delta_unit2 = [event for event in events_unit2 if event.type == "response.output_audio.delta"] + assert len(delta_unit2) == 1 + + # Unit 3's stop token schedules unit 4's append while unit 3's audio + # is also still pending. + _make_continuation_due_now(h) + h.deliver( + spoken_stage0_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[31, 32, 33], + ) + await h.settle() + assert len(h.port.submissions) == 3 + + events_unit3 = await h.deliver_and_settle(tts_output(request_id, samples=72000, text="third")) + delta_unit3 = [event for event in events_unit3 if event.type == "response.output_audio.delta"] + assert len(delta_unit3) == 1 + + # Emission order follows delivery order, which follows per-session + # stage processing order -- never the append submission order. + assert h.events.index(delta_unit2[0]) < h.events.index(delta_unit3[0]) + assert not [event for event in events_unit3 if event.type == "error"], types(events_unit3) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stop_token_mode_listen_segment_keeps_listen_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A listen decision segment is not a stop-token trigger. + + The listen path already continued on Stage 0 output before this switch + existed; the hook must stay out of its way (no second unit for one + segment). + """ + from tests.engine.duplex.test_session_runner import listen_output + + monkeypatch.setenv(CONTINUE_ON_ENV, "stop_token") + h = await _active_response_harness() + try: + request_id = h.stage0_request_id() + submissions_before = len(h.port.submissions) + + listen = listen_output(request_id) + listen.finished = False + await h.deliver_and_settle( + listen, + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + # Exactly the listen path's one continuation, not two. + assert len(h.port.submissions) == submissions_before + 1 + assert h.runner.model_state.continuation_units == 1 + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Unit depth # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_depth_two_plans_second_continuation_while_first_submits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A parked first continuation append must not block planning at depth 2.""" + monkeypatch.setenv(DEPTH_ENV, "2") + h = await _active_response_harness() + try: + _make_continuation_due_now(h) + gate = asyncio.Event() + h.port.submit_gate = gate + h.port.submit_started.clear() + kwargs = _continuation_kwargs(h) + payload = h.runner.model.silence_unit_payload() + + first = asyncio.ensure_future(h.runner._schedule_silence_continuation(payload, **kwargs)) + await asyncio.wait_for(h.port.submit_started.wait(), timeout=2.0) + assert len(h.runner._silence_in_flight()) == 1 + + second = asyncio.ensure_future(h.runner._schedule_silence_continuation(payload, **kwargs)) + scheduled = await asyncio.wait_for(asyncio.shield(second), timeout=1.0) + assert scheduled is True, "depth 2 did not plan the second unit while the first submits" + assert len(h.runner._silence_in_flight()) == 2 + + gate.set() + assert await asyncio.wait_for(first, timeout=2.0) is True + assert await asyncio.wait_for(second, timeout=2.0) is True + await h.settle() + assert len(h.port.submissions) == 3 # user append + two silence units + # The scheduler entry was called directly (not through the model + # channel), so the continuation-unit counter stays untouched; the + # in-flight window is the observable and it drains to empty. + assert h.runner._silence_in_flight() == [] + finally: + h.port.submit_gate = None + await close_harness(h) + + +@pytest.mark.asyncio +async def test_depth_one_still_waits_for_the_parked_continuation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Default depth keeps the historical single-flight wait.""" + monkeypatch.delenv(DEPTH_ENV, raising=False) + h = await _active_response_harness() + try: + _make_continuation_due_now(h) + gate = asyncio.Event() + h.port.submit_gate = gate + h.port.submit_started.clear() + kwargs = _continuation_kwargs(h) + payload = h.runner.model.silence_unit_payload() + + first = asyncio.ensure_future(h.runner._schedule_silence_continuation(payload, **kwargs)) + await asyncio.wait_for(h.port.submit_started.wait(), timeout=2.0) + + second = asyncio.ensure_future(h.runner._schedule_silence_continuation(payload, **kwargs)) + with pytest.raises((TimeoutError, asyncio.TimeoutError)): + await asyncio.wait_for(asyncio.shield(second), timeout=0.3) + assert len(h.runner._silence_in_flight()) == 1 + + gate.set() + assert await asyncio.wait_for(first, timeout=2.0) is True + assert await asyncio.wait_for(second, timeout=2.0) is True + await h.settle() + assert len(h.port.submissions) == 3 + finally: + h.port.submit_gate = None + await close_harness(h) + + +@pytest.mark.asyncio +async def test_real_append_invalidates_a_planned_continuation_at_any_depth( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A real (user) append accepted while a continuation sleeps supersedes it. + + The generation check replaced the raw anchor-timestamp equality when the + depth window was introduced; this pins the invariant it protected: the + stale unit is called off before submission, at every depth. + """ + monkeypatch.setenv(DEPTH_ENV, "2") + h = await _active_response_harness() + try: + model_state = h.runner.model_state + # A continuation planned 0.4 s into the current anchor's period. + model_state.last_native_submit_monotonic = time.monotonic() - 0.4 + model_state.silence_deadline_monotonic = None + kwargs = _continuation_kwargs(h) + + planned = asyncio.ensure_future( + h.runner._schedule_silence_continuation(h.runner.model.silence_unit_payload(), **kwargs) + ) + await asyncio.sleep(0) + + # A real append is accepted while the continuation still sleeps: it + # re-anchors the chain and bumps the generation. + model_state.last_native_submit_monotonic = time.monotonic() + model_state.silence_deadline_monotonic = None + model_state.native_input_generation += 1 + + await asyncio.wait_for(planned, timeout=3.0) + await h.settle() + # Called off before submission: no silence unit reached the stage. + assert len(h.port.submissions) == 1 + assert h.runner._silence_in_flight() == [] + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Flag parsing # +# --------------------------------------------------------------------------- # + + +def test_continue_on_flag_parsing(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.engine.duplex.config import duplex_continue_on_stop_token + + monkeypatch.delenv(CONTINUE_ON_ENV, raising=False) + assert duplex_continue_on_stop_token() is False + + monkeypatch.setenv(CONTINUE_ON_ENV, "audio") + assert duplex_continue_on_stop_token() is False + + monkeypatch.setenv(CONTINUE_ON_ENV, "stop_token") + assert duplex_continue_on_stop_token() is True + + monkeypatch.setenv(CONTINUE_ON_ENV, " Stop-Token ") + assert duplex_continue_on_stop_token() is True + + monkeypatch.setenv(CONTINUE_ON_ENV, "bogus") + assert duplex_continue_on_stop_token() is False + + +def test_unit_depth_flag_parsing(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.engine.duplex.config import duplex_unit_depth + + monkeypatch.delenv(DEPTH_ENV, raising=False) + assert duplex_unit_depth() == 1 + + monkeypatch.setenv(DEPTH_ENV, "2") + assert duplex_unit_depth() == 2 + + monkeypatch.setenv(DEPTH_ENV, "0") + assert duplex_unit_depth() == 1 + + monkeypatch.setenv(DEPTH_ENV, "-3") + assert duplex_unit_depth() == 1 + + monkeypatch.setenv(DEPTH_ENV, "64") + assert duplex_unit_depth() == 8 + + monkeypatch.setenv(DEPTH_ENV, "not-a-number") + assert duplex_unit_depth() == 1 + + monkeypatch.setenv(DEPTH_ENV, "") + assert duplex_unit_depth() == 1 diff --git a/tests/engine/duplex/test_session_runner.py b/tests/engine/duplex/test_session_runner.py index c51fc243cd9..d12bbd433e2 100644 --- a/tests/engine/duplex/test_session_runner.py +++ b/tests/engine/duplex/test_session_runner.py @@ -436,7 +436,8 @@ async def test_commands_for_a_closed_session_are_answered_with_unknown_session() @pytest.mark.asyncio -async def test_append_plans_and_submits_one_model_unit_in_wire_order() -> None: +async def test_append_plans_and_submits_one_model_unit_in_wire_order(monkeypatch) -> None: + monkeypatch.setenv("VLLM_OMNI_DUPLEX_STATIC_CONFIG", "1") h = await open_harness() try: events = await h.run(append_audio()) @@ -459,7 +460,10 @@ async def test_append_plans_and_submits_one_model_unit_in_wire_order() -> None: await h.run(append_audio()) second = h.port.submissions[1] assert second.already_submitted is True - assert second.prompt["model_intermediate_buffer"]["duplex"]["seq"] == 2 + second_duplex = second.prompt["model_intermediate_buffer"]["duplex"] + assert second_duplex["seq"] == 2 + assert "session_config" not in second_duplex + assert "runtime_config" not in second_duplex assert h.session.input_seq == 2 finally: await close_harness(h) @@ -844,6 +848,168 @@ async def test_listen_decision_is_consumed_and_never_forwarded_to_tts() -> None: assert events[0].details["model_listen"] is True assert events[0].to_realtime()["response"]["status"] == "listening" assert h.session.active_response_id is None + assert h.port.aborts == [] + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_model_listen_abort_is_opt_in_and_advances_epoch() -> None: + """Route-M listen can cut queued TTS without changing the default path.""" + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(abort_on_model_listen=True)) + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + spoken = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + response_created = find(spoken, "response.created") + assert response_created.response_id is not None + assert h.session.active_response_id == response_created.response_id + + events = await h.deliver_and_settle( + listen_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + cancelled = find(events, "response.done") + assert cancelled.status == "cancelled" + assert cancelled.response["status_details"]["reason"] == "model_listen" + assert find(events, "response.listen").epoch == 1 + assert h.port.aborts == [[request_id]] + assert h.port.cleanups == [([request_id], True)] + assert h.session.epoch == 1 + assert h.session.active_response_id is None + + # A late Stage2 delta from the cancelled epoch must not reach the + # client after the cancellation event. + late = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="stale"), epoch=0) + assert "response.output_audio.delta" not in types(late) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_model_listen_after_tts_segment_end_does_not_abort() -> None: + """A resident Stage-0 listen after a completed TTS segment is harmless.""" + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(abort_on_model_listen=True)) + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + spoken = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello", tts_is_last_chunk=True)) + response_created = find(spoken, "response.created") + response_id = response_created.response_id + + await h.deliver_and_settle( + listen_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + assert h.port.aborts == [] + assert h.session.epoch == 0 + assert h.session.active_response_id == response_id + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_model_listen_without_response_does_not_abort_resident_stage0() -> None: + """Pure listening has a resident request but no assistant response to cut.""" + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(abort_on_model_listen=True)) + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle( + listen_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + assert h.port.aborts == [] + assert h.session.epoch == 0 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_natural_stage0_end_does_not_abort_tts_tail() -> None: + """S0 turn-eos protects S2 tail audio that is still draining.""" + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(abort_on_model_listen=True)) + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + spoken = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + response_id = find(spoken, "response.created").response_id + + natural = SimpleNamespace( + request_id=request_id, + finished=True, + outputs=[SimpleNamespace(text="hello", token_ids=[], multimodal_output={})], + multimodal_output={"meta.turn_eos_token_id": 9}, + ) + h.deliver( + natural, + stage_id=0, + segment_finished=True, + segment_token_ids=[9, 10], + segment_output_metadata={"meta.turn_eos_token_id": 9}, + ) + await h.settle() + + await h.deliver_and_settle( + listen_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + assert h.port.aborts == [] + assert h.session.epoch == 0 + assert h.session.active_response_id == response_id + + tail = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="tail")) + assert "response.output_audio.delta" in types(tail) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_model_listen_before_first_audio_aborts_handoff() -> None: + """A spoken S0 handoff can be aborted before Talker emits its first audio.""" + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(abort_on_model_listen=True)) + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + handoff = SimpleNamespace( + request_id=request_id, + finished=True, + outputs=[SimpleNamespace(text="hello", token_ids=[8], multimodal_output={})], + multimodal_output={"meta.speak_token_id": 8}, + ) + h.deliver( + handoff, + stage_id=0, + segment_finished=True, + segment_token_ids=[8], + segment_output_metadata={"meta.speak_token_id": 8}, + ) + await h.settle() + assert h.session.active_response_id is None + + await h.deliver_and_settle( + listen_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + assert h.port.aborts == [[request_id]] + assert h.session.epoch == 1 finally: await close_harness(h) diff --git a/tests/engine/duplex/test_session_state_contract.py b/tests/engine/duplex/test_session_state_contract.py new file mode 100644 index 00000000000..dc846fe9e65 --- /dev/null +++ b/tests/engine/duplex/test_session_state_contract.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Every duplex model session state carries the runner's full flag contract. + +``DuplexModelSessionState`` only annotates its fields; a plugin state that +subclasses the ABC directly (Qwen3-Omni, AURA) must declare each one itself. +The session runner's real-append acceptance callback bumps +``native_input_generation`` on every accepted non-silence append, so a state +without that field fails the append (AttributeError) and closes the session. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import inspect +from types import SimpleNamespace + +import pytest + +from vllm_omni.engine.duplex.plugin import DuplexModelSessionState +from vllm_omni.engine.duplex.session import helpers +from vllm_omni.engine.duplex.session import runner as runner_mod +from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner +from vllm_omni.model_executor.models.aura_omni.duplex.session import AuraServingSessionState +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import MiniCPMO45ServingSessionState +from vllm_omni.model_executor.models.qwen3_omni.duplex.session import QwenDuplexSessionState + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +STATE_CLASSES = [QwenDuplexSessionState, AuraServingSessionState, MiniCPMO45ServingSessionState] +CONTRACT_FIELDS = sorted(inspect.get_annotations(DuplexModelSessionState)) + + +@pytest.mark.parametrize("state_cls", STATE_CLASSES, ids=lambda cls: cls.__name__) +def test_state_declares_every_contract_field(state_cls: type[DuplexModelSessionState]) -> None: + state = state_cls() + declared = {f.name for f in dataclasses.fields(state_cls)} + assert [name for name in CONTRACT_FIELDS if name not in declared] == [] + for name in CONTRACT_FIELDS: + getattr(state, name) + # Mutable defaults are per instance. + assert state.pending_silence_tasks is not state_cls().pending_silence_tasks + + +@pytest.mark.parametrize("state_cls", STATE_CLASSES, ids=lambda cls: cls.__name__) +def test_clear_continuation_resets_silence_bookkeeping(state_cls: type[DuplexModelSessionState]) -> None: + state = state_cls() + state.last_native_submit_monotonic = 1.0 + state.silence_deadline_monotonic = 2.0 + state.pending_silence_tasks.append(object()) # type: ignore[arg-type] + state.native_input_generation = 3 + state.silence_append_seq = 4 + state.clear_continuation() + assert state.last_native_submit_monotonic is None + assert state.silence_deadline_monotonic is None + assert state.pending_silence_tasks == [] + assert state.native_input_generation == 0 + assert state.silence_append_seq == 0 + + +class _AcceptingAttempt: + """Stands in for ``AppendAttempt``: the stage accepts the append at t=123.0. + + Mirrors ``AppendAttempt._submit``: an exception from the acceptance + callback fails the append and the session. + """ + + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + + async def run_in_wire_order(self, predecessor: object) -> bool: + del predecessor + try: + self.kwargs["on_append_accepted"](123.0) # type: ignore[operator] + except Exception as exc: + self.kwargs["fail_session"](f"runtime_append_task_failed: {exc!r}") # type: ignore[operator] + return False + return True + + def release_on_failure(self, task: object) -> None: + del task + + def clear_pending_silence(self, task: object) -> None: + del task + + +def _fake_runner(model_state: DuplexModelSessionState, failures: list[str]) -> SimpleNamespace: + session = SimpleNamespace( + epoch=0, + active_response_id=None, + capabilities=SimpleNamespace(supports_core_resumable_request=True), + ) + return SimpleNamespace( + session=session, + model_state=model_state, + ctx=SimpleNamespace(manager=SimpleNamespace(stage_request_id=lambda fence, **kw: "sess.r.stage0")), + out=None, + model=None, + tasks=SimpleNamespace(append_tail=None, track_append_task=lambda task, **kw: None), + _mark_pending_silence_superseded=lambda: None, + _fail_session_from_append=failures.append, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state_cls", STATE_CLASSES, ids=lambda cls: cls.__name__) +async def test_accepted_real_append_reanchors_chain( + state_cls: type[DuplexModelSessionState], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(runner_mod, "AppendAttempt", _AcceptingAttempt) + monkeypatch.setattr(helpers, "append_fence", lambda session, payload, *, epoch=None: SimpleNamespace(turn_id=0)) + state = state_cls() + failures: list[str] = [] + fake = _fake_runner(state, failures) + + task = await DuplexSessionRunner._start_append(fake, {"type": "audio"}, final=False) # type: ignore[arg-type] + ok = await asyncio.wait_for(task, timeout=2.0) + + assert failures == [] + assert ok is True + assert state.native_input_generation == 1 + assert state.last_native_submit_monotonic == 123.0 + assert state.silence_deadline_monotonic is None diff --git a/tests/engine/test_orchestrator_event_driven.py b/tests/engine/test_orchestrator_event_driven.py index 10972144688..e0ac1659f0d 100644 --- a/tests/engine/test_orchestrator_event_driven.py +++ b/tests/engine/test_orchestrator_event_driven.py @@ -128,9 +128,10 @@ def test_flag_parsing(monkeypatch) -> None: assert _event_driven_orch_enabled() is False -def test_only_qwen3_tts_has_a_pipeline_default() -> None: - assert _event_driven_orch_default_for_pipeline("qwen3_tts") is True - for model_type in (None, "qwen3_omni_moe", "minicpmo_4_5", "moss_tts_delay"): +def test_event_driven_pipeline_default() -> None: + for model_type in ("qwen3_tts", "minicpmo_4_5"): + assert _event_driven_orch_default_for_pipeline(model_type) is True + for model_type in (None, "qwen3_omni_moe", "moss_tts_delay"): assert _event_driven_orch_default_for_pipeline(model_type) is False diff --git a/tests/engine/test_stage_engine_core_proc.py b/tests/engine/test_stage_engine_core_proc.py index 217f13e53a4..0aca951413f 100644 --- a/tests/engine/test_stage_engine_core_proc.py +++ b/tests/engine/test_stage_engine_core_proc.py @@ -54,3 +54,188 @@ def test_preprocess_add_request_preserves_omni_fields(): assert current_wave == 3 assert result.external_req_id == "external" assert result.additional_information == {"conditioning": "payload"} + + +# --- event-driven idle wait (VLLM_OMNI_STAGE_IDLE_WAIT_S) ------------------------------------- + + +import queue # noqa: E402 +import threading # noqa: E402 +import time # noqa: E402 +from unittest.mock import MagicMock # noqa: E402 + +import pytest # noqa: E402 +from vllm.v1.engine import EngineCoreRequestType # noqa: E402 + + +def _idle_engine(*, idle_wait_s: float, step_result=({}, False), has_requests=True): + """A StageEngineCoreProc shell with just the loop state the idle wait touches.""" + engine = StageEngineCoreProc.__new__(StageEngineCoreProc) + engine.input_queue = queue.Queue() + engine.output_queue = queue.Queue() + engine.batch_queue = None + engine._idle_wait_s = idle_wait_s + engine._last_step_idle = False + engine.scheduler = MagicMock() + engine.scheduler.has_requests.return_value = has_requests + engine.step_fn = MagicMock(return_value=step_result) + engine.post_step = MagicMock() + engine.is_running = MagicMock(return_value=True) + engine.handled = [] + engine._handle_client_request = lambda request_type, request: engine.handled.append(request_type) + return engine + + +def test_idle_step_marks_the_loop_parked_without_sleeping(): + engine = _idle_engine(idle_wait_s=5.0) + with patch("vllm_omni.engine.stage_engine_core_proc.time.sleep") as sleep: + assert engine._process_engine_step() is False + assert engine._last_step_idle is True + sleep.assert_not_called() + + +@pytest.mark.parametrize( + ("step_result", "has_requests"), + [ + (({}, True), True), # executed: not idle + (({0: object()}, False), True), # produced output: not idle + (({}, False), False), # no requests: vLLM's own blocking get applies + (({0: SimpleNamespace(outputs=[], finished_requests={"r"})}, False), True), # finished a request + ], +) +def test_busy_or_empty_steps_do_not_park(step_result, has_requests): + engine = _idle_engine(idle_wait_s=5.0, step_result=step_result, has_requests=has_requests) + engine._process_engine_step() + assert engine._last_step_idle is False + + +def test_stats_only_step_output_still_parks(): + # log_stats attaches scheduler stats to an otherwise empty output every step. + stats_only = SimpleNamespace(outputs=[], scheduler_stats=object(), finished_requests=None, utility_output=None) + engine = _idle_engine(idle_wait_s=5.0, step_result=({0: stats_only}, False)) + engine._process_engine_step() + assert engine._last_step_idle is True + assert engine.output_queue.get_nowait() == (0, stats_only) + + +def test_parked_loop_wakes_on_chunk_ready_before_the_fallback(): + engine = _idle_engine(idle_wait_s=5.0) + engine._last_step_idle = True + timer = threading.Timer(0.05, engine._post_idle_wakeup) + timer.start() + start = time.monotonic() + with patch.object(EngineCoreProc, "_process_input_queue") as upstream: + engine._process_input_queue() + elapsed = time.monotonic() - start + timer.join() + assert elapsed < 2.0 + assert engine.handled == [EngineCoreRequestType.WAKEUP] + assert engine._last_step_idle is False + upstream.assert_called_once() + + +def test_parked_loop_falls_back_after_the_bound_when_no_wakeup_arrives(): + engine = _idle_engine(idle_wait_s=0.05) + engine._last_step_idle = True + start = time.monotonic() + with patch.object(EngineCoreProc, "_process_input_queue"): + engine._process_input_queue() + elapsed = time.monotonic() - start + assert 0.04 <= elapsed < 2.0 + assert engine.handled == [] + # The next loop pass steps once (deadlines, missed hints) instead of parking again. + assert engine._last_step_idle is False + + +def test_loop_does_not_park_with_pending_input_or_batches(): + engine = _idle_engine(idle_wait_s=5.0) + engine._last_step_idle = True + engine.batch_queue = [object()] + start = time.monotonic() + with patch.object(EngineCoreProc, "_process_input_queue"): + engine._process_input_queue() + assert time.monotonic() - start < 1.0 + + +def test_idle_wait_off_keeps_the_upstream_step(): + engine = _idle_engine(idle_wait_s=0.0) + with patch.object(EngineCoreProc, "_process_engine_step", return_value=False) as upstream: + assert engine._process_engine_step() is False + upstream.assert_called_once() + engine.step_fn.assert_not_called() + + +def test_init_idle_wait_binds_the_chunk_ready_wakeup(monkeypatch): + monkeypatch.setenv("VLLM_OMNI_STAGE_IDLE_WAIT_S", "0.2") + engine = _idle_engine(idle_wait_s=0.0) + adapter = MagicMock() + engine.scheduler._batch_window_s = 0.05 + engine.scheduler.chunk_transfer_adapter = adapter + engine._init_idle_wait() + assert engine._idle_wait_s == pytest.approx(0.05) # clamped to the admission window + adapter.set_chunk_ready_callback.assert_called_once() + callback = adapter.set_chunk_ready_callback.call_args.args[0] + callback() + assert engine.input_queue.get_nowait() == (EngineCoreRequestType.WAKEUP, None) + + +def test_init_idle_wait_default_off(monkeypatch): + monkeypatch.delenv("VLLM_OMNI_STAGE_IDLE_WAIT_S", raising=False) + engine = _idle_engine(idle_wait_s=1.0) + adapter = MagicMock() + engine.scheduler.chunk_transfer_adapter = adapter + engine._init_idle_wait() + assert engine._idle_wait_s == 0.0 + adapter.set_chunk_ready_callback.assert_not_called() + + +def test_init_idle_wait_refuses_a_stage_without_chunk_wakeup(monkeypatch): + # e.g. the native MRv2 data plane: every chunk would wait for the fallback bound. + monkeypatch.setenv("VLLM_OMNI_STAGE_IDLE_WAIT_S", "0.05") + engine = _idle_engine(idle_wait_s=0.0) + engine.scheduler.chunk_transfer_adapter = None + engine._init_idle_wait() + assert engine._idle_wait_s == 0.0 + + +# --- deadline batching: park bounded by the next release ---------------------------------------- + + +def test_park_is_bounded_by_the_schedulers_next_release(): + engine = _idle_engine(idle_wait_s=5.0) + engine.scheduler.next_release_in = MagicMock(return_value=0.05) + engine._last_step_idle = True + start = time.monotonic() + with patch.object(EngineCoreProc, "_process_input_queue"): + engine._process_input_queue() + elapsed = time.monotonic() - start + assert 0.04 <= elapsed < 2.0 + assert engine.handled == [] + + +@pytest.mark.parametrize("release_in", [None, MagicMock(), 1, "0.01"]) +def test_park_ignores_non_float_release_hints(release_in): + from vllm_omni.engine.stage_engine_core_proc import _scheduler_next_release_in + + scheduler = MagicMock() + scheduler.next_release_in = MagicMock(return_value=release_in) + assert _scheduler_next_release_in(scheduler) is None + assert _scheduler_next_release_in(object()) is None + + +def test_release_hint_has_a_floor(): + import vllm_omni.engine.stage_engine_core_proc as proc_module + + engine = _idle_engine(idle_wait_s=5.0) + engine.scheduler.next_release_in = MagicMock(return_value=0.0) + engine._last_step_idle = True + timeouts = [] + + def fake_get(timeout): + timeouts.append(timeout) + raise queue.Empty + + engine.input_queue.get = fake_get + with patch.object(EngineCoreProc, "_process_input_queue"): + engine._process_input_queue() + assert timeouts == [proc_module._MIN_RELEASE_WAIT_S] diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index a0ea5530f11..81c2907b5b8 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -525,6 +525,68 @@ class StageDeployConfig: engine_extras: dict[str, Any] = field(default_factory=dict) +@dataclass(frozen=True) +class DuplexPacingConfig: + """Opt-in duplex output pacing (``duplex_session.pacing``); every field off by default. + + A client that went quiet after its commit may have the silence + continuations of its reply submitted up to ``onset_lead_max_s`` ahead of + the nominal 1 s cadence (``onset_lead_target_s`` minus the reply's first + chunk length once that is known), so a short first audio chunk no longer + turns into a structural stall at the second unit. + ``onset_skip_response_wait`` lets that second unit go before its response + opens; a stall above ``critical_stall_s`` raises the lead to its maximum + until two clean chunks. ``fire_grid_ms`` snaps paced continuations onto a + shared grid so units of different sessions reach Stage 0 together; + ``idle_grid`` rounds the continuations after a non-terminal listen up to it. + """ + + enabled: bool = False + onset_lead_max_s: float = 0.0 + onset_lead_target_s: float = 1.1 + onset_skip_response_wait: bool = False + critical_stall_s: float = 0.05 + fire_grid_ms: int = 0 + idle_grid: bool = False + + def __post_init__(self) -> None: + for name in ("enabled", "onset_skip_response_wait", "idle_grid"): + if not isinstance(getattr(self, name), bool): + raise ValueError(f"duplex_session.pacing.{name} must be a boolean") + for name in ("onset_lead_max_s", "onset_lead_target_s", "critical_stall_s"): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int | float) or value < 0: + raise ValueError(f"duplex_session.pacing.{name} must be a non-negative number") + # The lead must stay below one model unit: unit_continuation.depth 2 then + # bounds the synthesized silence ahead of the wall clock to a single unit. + if self.onset_lead_max_s >= 1.0: + raise ValueError("duplex_session.pacing.onset_lead_max_s must be < 1.0 (one model unit)") + grid = self.fire_grid_ms + if isinstance(grid, bool) or not isinstance(grid, int) or grid < 0 or (grid and 1000 % grid): + raise ValueError("duplex_session.pacing.fire_grid_ms must be 0 or a positive divisor of 1000") + + +@dataclass(frozen=True) +class DuplexUnitContinuationConfig: + """When a duplex session plans its next silence unit (``duplex_session.unit_continuation``). + + ``continue_on: audio`` (default) plans it after the previous unit's audio is + emitted; ``stop_token`` plans it as soon as Stage 0 finishes the segment, so + Stage 0's next unit overlaps Stage 1/2 of the previous one. ``depth`` is how + many continuations of one session may be scheduled or submitting at once + (1 = single flight, the default). Cadence and emission order are unchanged. + """ + + continue_on: str = "audio" + depth: int = 1 + + def __post_init__(self) -> None: + if self.continue_on not in ("audio", "stop_token"): + raise ValueError("duplex_session.unit_continuation.continue_on must be 'audio' or 'stop_token'") + if isinstance(self.depth, bool) or not isinstance(self.depth, int) or not 1 <= self.depth <= 8: + raise ValueError("duplex_session.unit_continuation.depth must be an integer in [1, 8]") + + @dataclass(frozen=True) class DuplexSessionRuntimeConfig: """Server-owned lifecycle and per-session buffering limits.""" @@ -552,8 +614,51 @@ class DuplexSessionRuntimeConfig: # compile their real shapes. The empty per-stage JIT registry does not. # A negative value disables every startup warmup. warmup_frames: int = 0 + # Opt-in playback cut: a new user utterance over audible assistant audio + # (``barge_arm_min_silence_s`` of non-speech appends, then at least + # ``barge_arm_min_speech_s`` of speech, starting at least + # ``barge_arm_min_playback_s`` after the reply's first audio) cancels the + # reply once the model yields (listens or ends its turn) while the client + # still holds ``barge_cut_min_unplayed_s`` of unplayed audio. The speech + # floor keeps short backchannels ("mm") from cutting the reply. + barge_cut_on_model_yield: bool = False + barge_arm_min_silence_s: float = 0.6 + barge_arm_min_speech_s: float = 0.4 + barge_arm_min_playback_s: float = 1.0 + barge_cut_min_unplayed_s: float = 0.15 + # Mapping in YAML; coerced to ``DuplexPacingConfig`` (all off by default). + pacing: DuplexPacingConfig = field(default_factory=DuplexPacingConfig) + # Mapping in YAML; coerced to ``DuplexUnitContinuationConfig``. + unit_continuation: DuplexUnitContinuationConfig = field(default_factory=DuplexUnitContinuationConfig) def __post_init__(self) -> None: + pacing = self.pacing + if pacing is None: + pacing = DuplexPacingConfig() + elif isinstance(pacing, Mapping): + pacing = DuplexPacingConfig(**pacing) + elif not isinstance(pacing, DuplexPacingConfig): + raise ValueError("duplex_session.pacing must be a mapping") + object.__setattr__(self, "pacing", pacing) + continuation = self.unit_continuation + if continuation is None: + continuation = DuplexUnitContinuationConfig() + elif isinstance(continuation, Mapping): + continuation = DuplexUnitContinuationConfig(**continuation) + elif not isinstance(continuation, DuplexUnitContinuationConfig): + raise ValueError("duplex_session.unit_continuation must be a mapping") + object.__setattr__(self, "unit_continuation", continuation) + if not isinstance(self.barge_cut_on_model_yield, bool): + raise ValueError("duplex_session.barge_cut_on_model_yield must be a boolean") + for name in ( + "barge_arm_min_silence_s", + "barge_arm_min_speech_s", + "barge_arm_min_playback_s", + "barge_cut_min_unplayed_s", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int | float) or value < 0: + raise ValueError(f"duplex_session.{name} must be a non-negative number") positive = { "disconnect_grace_s": self.disconnect_grace_s, "reaper_interval_s": self.reaper_interval_s, diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index cefa2a37a81..f880ddf96c0 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -6,11 +6,13 @@ import os import time from collections import defaultdict +from collections.abc import Container, Iterable from typing import TYPE_CHECKING, Any import torch from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import KVConnectorBlockState, SchedulerOutput @@ -49,6 +51,51 @@ def _has_async_chunk_payload_to_run(request: Request) -> bool: return bool(audio) +class StepReleasePolicy: + """Opt-in hold/release of ready chunks across steps (``additional_config.step_release_policy``). + + The key is a ``module.path.factory`` name; the factory is called once with + the scheduler and returns a policy, or None to stay off. A held request + stays queued with its chunk ready and is offered again next step. The base + class holds nothing. + """ + + def held_requests(self, scheduler: OmniGenerationScheduler, now: float) -> Container[str]: + return () + + def on_scheduled(self, request_ids: Iterable[str], now: float) -> None: + pass + + def on_output(self, request_id: str, mm_output: Any, now: float) -> None: + pass + + def observe_step(self, rows: int, now: float) -> None: + pass + + def forget(self, request_id: str) -> None: + pass + + def next_release_in(self, now: float) -> float | None: + """Seconds until a held chunk is due (None: nothing held).""" + return None + + +def _release_policy(scheduler: object) -> StepReleasePolicy | None: + policy = getattr(scheduler, "_step_release_policy", None) + return policy if isinstance(policy, StepReleasePolicy) else None + + +def _build_release_policy(scheduler: OmniGenerationScheduler) -> StepReleasePolicy | None: + additional_config = getattr(scheduler.vllm_config, "additional_config", None) + name = additional_config.get("step_release_policy") if isinstance(additional_config, dict) else None + if not name: + return None + policy = resolve_obj_by_qualname(name)(scheduler) + if policy is not None and not isinstance(policy, StepReleasePolicy): + raise TypeError(f"additional_config.step_release_policy {name!r} must return a StepReleasePolicy or None") + return policy + + class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): waiting: RequestQueue @@ -75,6 +122,7 @@ def __init__(self, *args, **kwargs): "Generation scheduler decodes first chunks in express steps (min continuation slack %.2f s)", self._express_min_slack_s, ) + self._step_release_policy = _build_release_policy(self) def _build_generation_scheduler_output( self, @@ -258,6 +306,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self._process_pending_omni_inputs(model_mode="generation") self._drop_aborted_queued_requests() self._resync_streaming_input_counter() + release_policy = _release_policy(self) + held = release_policy.held_requests(self, scheduled_timestamp) if release_policy is not None else () async_chunk_transport = self._async_chunk_transport_enabled() native_chunks = bool(getattr(self, "_native_data_plane", False) and async_chunk_transport) # Parking releases an execution slot, but stateful codecs retain their @@ -303,6 +353,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: req_index += 1 continue + if held and request.request_id in held: + # The release policy lets this stream wait for a fuller step. + req_index += 1 + continue + if native_chunks and request.num_in_flight_tokens > 0: # Overlap different streams, never two stateful chunks from # the same stream. Readiness remains pending for the next step. @@ -374,6 +429,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.waiting.pop_request() skipped_waiting_requests.add_request(request) continue + if held and request.request_id in held: + self.waiting.pop_request() + skipped_waiting_requests.add_request(request) + continue if native_chunks and request.num_in_flight_tokens > 0: # A restored waiting entry can still own an outstanding batch. # Do not let it block other ready streams or execute it twice. @@ -447,6 +506,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: if first_chunk_express: self._chunk_started.update(num_scheduled_tokens) + if release_policy is not None: + release_policy.on_scheduled(num_scheduled_tokens, scheduled_timestamp) # Return skipped waiting requests if skipped_waiting_requests: @@ -543,6 +604,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: return self._wrap_omni_scheduler_output(scheduler_output) + def next_release_in(self) -> float | None: + """Seconds until the release policy's earliest held chunk is due (None: nothing held).""" + release_policy = _release_policy(self) + return None if release_policy is None else release_policy.next_release_in(time.monotonic()) + def _continuations_have_slack(self) -> bool: """Whether every started stream with a chunk ready to decode can wait one express step.""" now = time.monotonic() @@ -589,6 +655,9 @@ def _free_request( if getattr(self, "_first_chunk_express", False): self._chunk_started.discard(request.request_id) getattr(self, "_stream_audio", {}).pop(request.request_id, None) + release_policy = _release_policy(self) + if release_policy is not None: + release_policy.forget(request.request_id) if self.input_coordinator is None: return super()._free_request(request, delay_free_blocks) @@ -614,6 +683,7 @@ def update_from_output( ec_connector_output = getattr(model_runner_output, "ec_connector_output", None) cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats + release_policy = _release_policy(self) # Every GPU write enqueued by this and earlier steps has completed, so # it is safe to return deferred-free blocks to the pool. This is the @@ -736,6 +806,8 @@ def update_from_output( mm_output = mm_outputs[req_index] if mm_outputs else None if getattr(self, "_express_min_slack_s", 0) > 0 and self._first_chunk_express: self._record_stream_audio(req_id, mm_output) + if release_policy is not None: + release_policy.on_output(req_id, mm_output, time.monotonic()) status_before_stop = request.status finish_reason = None is_segment_finished = False @@ -930,6 +1002,9 @@ def update_from_output( self._capture_omni_connector_output(model_runner_output) + if release_policy is not None: + release_policy.observe_step(len(num_scheduled_tokens), time.monotonic()) + return engine_core_outputs def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 09854d22d27..4d73414a9d1 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -21,6 +21,10 @@ logger = get_connector_logger(__name__) +# Where glibc places POSIX shared-memory objects (the lock files live there +# too). None disables the existence probe on hosts without it. +_POSIX_SHM_DIR = "/dev/shm" if os.path.isdir("/dev/shm") else None + def _wakeup_enabled() -> bool: return os.environ.get("VLLM_OMNI_SHM_WAKEUP", "1") == "1" @@ -147,6 +151,22 @@ def wait_for_change(self, generation: int, *, timeout: float) -> bool: self._wake_generation += 1 return self._wake_generation != generation + def wake_self(self) -> None: + """Interrupt this receiver's own ``wait_for_change`` (e.g. a new key to poll). + + Writes one byte through the receiver's own write end, so the next (or a + blocked) ``wait_for_change`` returns as if a ``put`` had arrived. A full + FIFO already holds a pending wakeup, so ``EAGAIN`` is ignored. + """ + with self._wake_lock: + fd = self._wake_hold_fd + if self._wake_closed or fd is None: + return + try: + os.write(fd, b"\0") + except OSError: + pass + def _wake_receiver(self, to_stage: Any) -> None: if not _wakeup_enabled(): return @@ -250,6 +270,13 @@ def _get_data_with_lock(self, lock_file: str, shm_handle: dict[str, Any]) -> tup def _get_by_key(self, get_key: str) -> tuple[Any, int] | None: """Read a SHM segment addressed purely by *get_key*.""" + # Receivers poll every pending request about once per millisecond and + # nearly every poll misses. os.access reports a missing segment without + # building a FileNotFoundError, which is most of the miss cost under + # the GIL. A segment that appears after the probe is read next poll, + # exactly as a miss inside shm_open is. + if _POSIX_SHM_DIR is not None and not os.access(f"{_POSIX_SHM_DIR}/{get_key}", os.F_OK): + return None shm = None try: shm = shm_pkg.SharedMemory(name=get_key) diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py index 8b64e3f246e..f158c8fc8c3 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py @@ -3,6 +3,7 @@ import threading from collections import deque +from collections.abc import Callable, Mapping from typing import Any from ..utils.logging import get_connector_logger @@ -10,6 +11,21 @@ logger = get_connector_logger(__name__) +def stage_idle_wait_s(additional_config: object) -> float: + """Fallback bound (s) of a chunk-fed stage's event-driven wakeups; ``0`` keeps polling. + + ``additional_config.chunk_idle_wait_s`` > 0 lets the chunk receive thread + block on the connector's arrival wakeup instead of re-polling every + millisecond, and lets the stage engine loop block between steps while every + request only waits for an upstream chunk. Wakeups are hints: the value + bounds how long a missed one can delay a chunk. + """ + value = additional_config.get("chunk_idle_wait_s") if isinstance(additional_config, Mapping) else None + if isinstance(value, bool) or not isinstance(value, int | float): + return 0.0 + return float(value) if value > 0 else 0.0 + + class OmniTransferAdapterBase: """Base class for managing data transfer via OmniConnector. @@ -38,6 +54,14 @@ def __init__(self, config: Any): self._send_failure_lock = threading.Lock() self._receive_failures: dict[str, str] = {} self._receive_failure_lock = threading.Lock() + # Event-driven receive (off unless ``chunk_idle_wait_s`` > 0 and the + # connector exposes arrival wakeups): decided before the thread starts. + self._recv_fifo_wait_s = ( + stage_idle_wait_s(getattr(config, "additional_config", None)) + if callable(getattr(self.connector, "wait_for_change", None)) + else 0.0 + ) + self._chunk_ready_callback: Callable[[], None] | None = None self.recv_thread = threading.Thread(target=self.recv_loop, daemon=True) self.recv_thread.start() @@ -49,14 +73,50 @@ def __init__(self, config: Any): def create_connector(cls, model_config: Any): raise NotImplementedError + def set_chunk_ready_callback(self, callback: Callable[[], None] | None) -> None: + """Call ``callback`` from the receive thread after a pass that committed data. + + The stage engine uses it to wake its loop; it runs after the commit is + visible to the scheduler (``_finished_load_reqs``), so a woken + ``schedule()`` always sees the chunk. + """ + self._chunk_ready_callback = callback + + def _notify_chunk_ready(self) -> None: + callback = getattr(self, "_chunk_ready_callback", None) + if callback is None: + return + try: + callback() + except Exception: # noqa: BLE001 - a lost hint only costs the fallback bound + logger.debug("chunk-ready callback failed", exc_info=True) + + def _wake_recv_wait(self) -> None: + """Interrupt a receive thread blocked on the connector wakeup (new work to poll).""" + if getattr(self, "_recv_fifo_wait_s", 0.0) <= 0: + return + wake = getattr(self.connector, "wake_self", None) + if callable(wake): + wake() + def recv_loop(self): """Loop to poll for incoming data. Process each pending request exactly once per pass. When no request made progress, back off 1 ms instead of tight-spinning on failed shm_open syscalls (which can burn a full CPU core). + + With ``additional_config.chunk_idle_wait_s`` > 0 and a connector that exposes + arrival wakeups, a pass without progress blocks until the next ``put`` + to this stage (or a local ``wake_self``) instead, bounded by that value. + The generation is taken before the pass, so a ``put`` that lands while + the pass runs makes the wait return at once. """ + fifo_wait_s = getattr(self, "_recv_fifo_wait_s", 0.0) + snapshot = getattr(self.connector, "get_wakeup_generation", None) if fifo_wait_s > 0 else None + wait_for_change = getattr(self.connector, "wait_for_change", None) if callable(snapshot) else None while not self.stop_event.is_set(): + generation = snapshot() if wait_for_change is not None else None n = len(self._pending_load_reqs) any_success = False for _ in range(n): @@ -74,12 +134,20 @@ def recv_loop(self): self._pending_load_reqs.append(request) logger.warning(f"Error receiving data for {request_id}: {e}") + if any_success: + self._notify_chunk_ready() + # Timeout is the fallback for lock-free append/notify races. with self._recv_cond: if not self._pending_load_reqs and not self.stop_event.is_set(): self._recv_cond.wait(timeout=0.1) - elif not any_success and not self.stop_event.is_set(): + continue + if any_success or self.stop_event.is_set(): + continue + if generation is None: self._recv_cond.wait(timeout=0.001) + continue + wait_for_change(generation, timeout=fifo_wait_s) def record_send_failure(self, request_id: str | None, reason: str) -> None: """Note that a chunk for *request_id* will never be delivered. 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 bd6d7abf15b..1df680a9d4c 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 @@ -260,12 +260,16 @@ def snapshot_container(value: Any) -> Any: return value snapshot = copy.copy(request) - for name in ( - "additional_information", - "prompt_token_ids", - ): - if hasattr(request, name): - setattr(snapshot, name, snapshot_container(getattr(request, name))) + if hasattr(request, "additional_information"): + snapshot.additional_information = snapshot_container(request.additional_information) + if hasattr(request, "prompt_token_ids"): + prompt_token_ids = request.prompt_token_ids + # Token ids are immutable ints, so a C-level shallow copy freezes the + # list exactly like the element-wise walk, which costs O(prompt) + # Python work per request on every sender step. + snapshot.prompt_token_ids = ( + prompt_token_ids.copy() if type(prompt_token_ids) is list else snapshot_container(prompt_token_ids) + ) for private_name, public_name in ( ("_all_token_ids", "all_token_ids"), @@ -339,6 +343,9 @@ def load_async(self, request: Request): self._pending_load_reqs.append(entry) with self._recv_cond: self._recv_cond.notify() + # A receive thread blocked on the connector wakeup must poll the new + # key too: its chunk may have been put before this registration. + self._wake_recv_wait() def _accepts_new_token_ids(self, processor: Callable[..., Any]) -> bool: cached = self._processor_accepts_step_tokens.get(processor) diff --git a/vllm_omni/engine/duplex/config.py b/vllm_omni/engine/duplex/config.py index a576220045c..f4849dbc797 100644 --- a/vllm_omni/engine/duplex/config.py +++ b/vllm_omni/engine/duplex/config.py @@ -200,6 +200,11 @@ def validate_append_modalities(self, *, has_audio: bool, has_video: bool) -> str return None +#: Audio a duplex client buffers before it plays the first chunk of a session, +#: then playing each chunk on arrival. Duplex pacing models that client with it. +DUPLEX_CLIENT_PREBUFFER_S = 0.5 + + @dataclass class DuplexPlaybackCursor: generated_ms: int = 0 diff --git a/vllm_omni/engine/duplex/intermediate.py b/vllm_omni/engine/duplex/intermediate.py index 175259dd990..fd77cae4281 100644 --- a/vllm_omni/engine/duplex/intermediate.py +++ b/vllm_omni/engine/duplex/intermediate.py @@ -6,9 +6,13 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, TypedDict +import torch + if TYPE_CHECKING: from vllm_omni.engine.duplex.contracts import DuplexFence +_TRANSPORT_TENSOR_MARKER = "__tensor__" + class DuplexIntermediateBuffer(TypedDict, total=False): """Structured keys carried through ``model_intermediate_buffer``. @@ -79,6 +83,22 @@ def build_duplex_append_prompt( } +def drop_static_append_configs(prompt: dict[str, object]) -> dict[str, object]: + """Remove the session-sized config snapshots from a later append. + + The model-intermediate buffer is deliberately mutable while a submission + is being prepared. Keeping this helper at the transport seam makes the + optimization explicit and leaves the prompt/token calculation (which may + need the full config) in the model plugin unchanged. + """ + intermediate = prompt.get("model_intermediate_buffer") + duplex = intermediate.get("duplex") if isinstance(intermediate, dict) else None + if isinstance(duplex, dict): + duplex.pop("session_config", None) + duplex.pop("runtime_config", None) + return prompt + + def build_duplex_intermediate_buffer( *, request_id: str, @@ -109,8 +129,39 @@ def build_duplex_intermediate_buffer( return buffer +def pack_transport_tensor(value: object) -> object: + """Pack a tensor as raw bytes so it survives ``model_intermediate_buffer`` transport. + + The buffer is ``dict[str, Any]``, so a raw tensor would decode as vLLM's + ``(dtype, shape, aux-index)`` tuple, and a ``tolist()`` round trip costs + ~100x a memcpy on the orchestrator thread. Non-tensors pass through. + """ + if not isinstance(value, torch.Tensor): + return value + tensor = value.detach().to("cpu").contiguous() + return { + _TRANSPORT_TENSOR_MARKER: True, + "dtype": str(tensor.dtype).removeprefix("torch."), + "shape": list(tensor.shape), + "data": tensor.reshape(-1).view(torch.uint8).numpy().tobytes(), + } + + +def unpack_transport_tensor(value: object) -> object: + """Inverse of :func:`pack_transport_tensor`; other values pass through.""" + if not isinstance(value, Mapping) or value.get(_TRANSPORT_TENSOR_MARKER) is not True: + return value + dtype = getattr(torch, str(value["dtype"])) + shape = [int(dim) for dim in value["shape"]] + data = value["data"] + if not data: + return torch.empty(shape, dtype=dtype) + # Copy into a writable buffer: msgpack hands back immutable bytes. + return torch.frombuffer(bytearray(data), dtype=dtype).reshape(shape) + + def set_ref_audio(buffer: DuplexIntermediateBuffer, waveform: object, sample_rate_hz: int) -> None: - buffer.setdefault("codes", {})["ref"] = waveform + buffer.setdefault("codes", {})["ref"] = pack_transport_tensor(waveform) buffer.setdefault("meta", {})["ref_audio_sr"] = int(sample_rate_hz) @@ -119,7 +170,7 @@ def set_tts_handoff(buffer: DuplexIntermediateBuffer, token_ids: object | None, if token_ids is not None: buffer.setdefault("ids", {})["tts"] = token_ids if hidden_states is not None: - buffer.setdefault("hidden_states", {})["tts"] = hidden_states + buffer.setdefault("hidden_states", {})["tts"] = pack_transport_tensor(hidden_states) def get_tts_handoff(info: dict[str, object]) -> tuple[object | None, object | None]: @@ -130,7 +181,7 @@ def get_tts_handoff(info: dict[str, object]) -> tuple[object | None, object | No hidden_states = hidden_info.get("tts") if isinstance(hidden_info, dict) else None return ( info.get("tts_token_ids") if token_ids is None else token_ids, - info.get("tts_hidden_states") if hidden_states is None else hidden_states, + unpack_transport_tensor(info.get("tts_hidden_states") if hidden_states is None else hidden_states), ) diff --git a/vllm_omni/engine/duplex/plugin.py b/vllm_omni/engine/duplex/plugin.py index 26402f108ba..91a9e4300d1 100644 --- a/vllm_omni/engine/duplex/plugin.py +++ b/vllm_omni/engine/duplex/plugin.py @@ -147,6 +147,14 @@ class DuplexModelSessionState(ABC): # continuation deadline. A real (non-silence) input resets the chain. last_native_submit_monotonic: float | None silence_deadline_monotonic: float | None + # Multi-unit continuation chain, read and written only when + # ``duplex_session.unit_continuation.depth`` > 1. The class defaults serve + # every state at depth 1. In-flight continuations, oldest first: + pending_silence_tasks: tuple[asyncio.Task[bool], ...] = () + # Bumped per accepted real append and per clear: older plans are stale. + native_input_generation: int = 0 + # Accepted silence appends: bounds the sibling submissions a plan may see. + silence_append_seq: int = 0 @abstractmethod def retain_committed_audio( @@ -220,6 +228,9 @@ def clear_continuation(self) -> None: self.pending_silence_owner_id = None self.last_native_submit_monotonic = None self.silence_deadline_monotonic = None + # Read only at unit_continuation.depth > 1: plans made before the clear are stale. + self.pending_silence_tasks = () + self.native_input_generation += 1 @dataclass(frozen=True, slots=True) @@ -292,6 +303,11 @@ class DuplexModelPlugin(ABC): projects_intermediate_outputs: bool = False plugin_id: str = "" private_runtime_config_keys: frozenset[str] = frozenset() + #: If true, the model consumes session/runtime append configuration on the + #: first append of a request and keeps it in its own session state. The + #: engine still sends the snapshots when the config generation changes. + #: Plugins that inspect configuration on every append leave this false. + append_configs_are_session_static: bool = False #: Samples per silence unit the runner appends to keep a model turn going. silence_continuation_samples: int = 16000 #: Sample rate of that unit: the runner submits it through ``plan_append`` @@ -364,6 +380,22 @@ def decide_output( output: object, ) -> DuplexOutputDecision | None: ... + def stage0_naturally_ended( + self, + *, + segment_finished: bool, + segment_token_ids: tuple[int, ...], + segment_output_metadata: Mapping[str, object], + output: object, + ) -> bool: + """Whether Stage 0 emitted the model's natural response-end marker. + + Read only while the session paces its output or may cut playback + (``duplex_session.pacing`` / ``barge_cut_on_model_yield``). + """ + del segment_finished, segment_token_ids, segment_output_metadata, output + return False + def project_intermediate_output( self, *, diff --git a/vllm_omni/engine/duplex/session/append_task.py b/vllm_omni/engine/duplex/session/append_task.py index 7c07234714b..37d1be84fb9 100644 --- a/vllm_omni/engine/duplex/session/append_task.py +++ b/vllm_omni/engine/duplex/session/append_task.py @@ -130,6 +130,8 @@ def clear_pending_silence(self, done: asyncio.Task[bool]) -> None: if model_state.pending_silence_task is done: model_state.pending_silence_task = None model_state.pending_silence_owner_id = None + if done in model_state.pending_silence_tasks: + model_state.pending_silence_tasks = tuple(t for t in model_state.pending_silence_tasks if t is not done) # ------------------------------------------------------------------ # # Running # diff --git a/vllm_omni/engine/duplex/session/context.py b/vllm_omni/engine/duplex/session/context.py index fde9425be84..588b09ebebb 100644 --- a/vllm_omni/engine/duplex/session/context.py +++ b/vllm_omni/engine/duplex/session/context.py @@ -33,6 +33,7 @@ from vllm_omni.engine.duplex.plugin import DuplexModelPlugin, DuplexModelSessionState from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession from vllm_omni.engine.duplex.session.manager import DuplexSessionManager + from vllm_omni.engine.duplex.session.pacing import SessionPacing from vllm_omni.metrics.stats import StageRequestStats _OffloadT = TypeVar("_OffloadT") @@ -62,6 +63,9 @@ class DuplexRunState: #: assistant audio is still draining. Cleared on barge-in/cancel, or when #: the next ephemeral turn begins. Orthogonal to barge-in (which aborts). concurrent_turn_requests_released: bool = False + #: Output pacing / playback-cut state (``duplex_session.pacing`` or + #: ``barge_cut_on_model_yield``); None keeps every historical path. + pace: SessionPacing | None = None class RunnerServices(Protocol): diff --git a/vllm_omni/engine/duplex/session/helpers.py b/vllm_omni/engine/duplex/session/helpers.py index 5a74d3c7831..7fe691070cd 100644 --- a/vllm_omni/engine/duplex/session/helpers.py +++ b/vllm_omni/engine/duplex/session/helpers.py @@ -15,6 +15,7 @@ import asyncio import binascii +import time from typing import TYPE_CHECKING import pybase64 as base64 @@ -29,6 +30,7 @@ from vllm_omni.engine.duplex.session.context import DuplexSessionTasks from vllm_omni.engine.duplex.session.engine_session import DuplexCommittedInput, DuplexEngineSession + from vllm_omni.engine.duplex.session.pacing import SessionPacing def task_is_cancelling(task: asyncio.Task[object] | None) -> bool: @@ -130,11 +132,15 @@ def audio_payload_size_bytes(payload: Mapping[str, object]) -> int: # --------------------------------------------------------------------------- # -def advance_barge_in_epoch(session: DuplexEngineSession) -> tuple[int, dict[str, int]]: +def advance_barge_in_epoch( + session: DuplexEngineSession, *, pace: SessionPacing | None = None +) -> tuple[int, dict[str, int]]: """Start a new epoch for a barge-in, returning it with the playback it cut off.""" old_playback = session.playback.as_dict() new_epoch = session.barge_in() session.clear_playback_cursor() + if pace is not None: + pace.on_epoch_advance(time.monotonic()) return new_epoch, old_playback diff --git a/vllm_omni/engine/duplex/session/model_channel.py b/vllm_omni/engine/duplex/session/model_channel.py index e2bccf9b795..9e8db81a50f 100644 --- a/vllm_omni/engine/duplex/session/model_channel.py +++ b/vllm_omni/engine/duplex/session/model_channel.py @@ -42,6 +42,7 @@ duplex_same_turn_request_ids, duplex_session_id_from_request_id, ) +from vllm_omni.engine.duplex.intermediate import drop_static_append_configs from vllm_omni.engine.duplex.plugin import ( DuplexRuntimeConfigError, coerce_int, @@ -52,6 +53,7 @@ from vllm_omni.engine.duplex.session.emitter import SessionEmitter from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.engine.duplex.session.pacing import continue_on_stop_token, same_turn_response_keeps_continuation from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats from vllm_omni.outputs import OmniRequestOutput from vllm_omni.outputs.duplex import attach_duplex_output_decision @@ -59,6 +61,8 @@ if TYPE_CHECKING: from vllm.outputs import RequestOutput + from vllm_omni.engine.duplex.session.pacing import SessionPacing + logger = init_logger(__name__) @@ -75,6 +79,7 @@ async def __call__( response_owned: bool, expected_epoch: int | None, expected_model_turn_id: int | None, + pace_kind: str = "speech", ) -> bool: ... @@ -93,12 +98,18 @@ def __init__( close_from_runtime: Callable[[str], Awaitable[None]], schedule_silence_continuation: SilenceContinuationScheduler, abort_request: Callable[..., Awaitable[None]], + cancel_model_response: Callable[[str], Awaitable[bool]] | None = None, ) -> None: self._ctx = ctx self._out = out self._close_from_runtime = close_from_runtime self._schedule_silence_continuation = schedule_silence_continuation self._abort_request = abort_request + self._cancel_model_response = cancel_model_response + # Config snapshots are needed by a session-static plugin only when a + # request starts or the session config generation changes. Keep the + # last successful submission so a failed append retries the snapshot. + self._last_append_config_key: tuple[str, int] | None = None def _draining_stage_ids(self) -> frozenset[int]: """Stages the plugin keeps across a concurrent turn. Empty if undeclared.""" @@ -108,6 +119,34 @@ def _draining_stage_ids(self) -> frozenset[int]: stage_count = int(getattr(self._ctx.stage_port, "stage_count", 0) or 0) return frozenset(int(stage_id) for stage_id in declare(stage_count=stage_count)) + def _pace(self) -> SessionPacing | None: + """The session's pacing state; None when pacing and the playback cut are off.""" + return getattr(getattr(self._ctx, "run", None), "pace", None) + + def _barge_cut_ready(self) -> bool: + pace = self._pace() + return pace is not None and pace.cut_ready(time.monotonic()) + + def _barge_cut_enabled(self) -> bool: + pace = self._pace() + return pace is not None and pace.barge_enabled + + async def cut_on_model_yield(self) -> bool: + """``barge_cut_on_model_yield``: Stage 0 ended its turn after a new user utterance; cut the tail. + + Same cancellation as an explicit barge-in (the old epoch's stage + requests aborted, ``audio.cancelled``, the next append opens epoch+1); + a no-op when the reply already finished or the client has nearly + played it out. + """ + pace = self._pace() + if not self._out.auto_responds() or self._cancel_model_response is None or not self._barge_cut_ready(): + return False + pace.note_cut(time.monotonic(), "model_yield") + await self._cancel_model_response("model_yield") + self._ctx.model_state.clear_continuation() + return True + @staticmethod def should_commit_response_to_history(session: DuplexEngineSession, response_id: str | None) -> bool: if response_id is not None and response_id != session.active_response_id: @@ -184,6 +223,18 @@ async def append_runtime_input( except asyncio.CancelledError: raise except Exception as exc: + if ( + isinstance(exc, DuplexFenceMismatchError) + and expected_epoch is not None + and session.epoch != expected_epoch + and self._barge_cut_enabled() + ): + # The epoch moved while this append was in flight (a barge cut + # or another cancel): drop it like the check above drops a + # queued one. Failing it would roll its PCM back and stop the + # append chain, abandoning every append behind it -- the cut's + # replay, which also carries this unit. + return True, False logger.exception("Failed to append duplex runtime input: %s", exc) self.send_runtime_error("runtime_append_failed", exc) return False, False @@ -319,6 +370,13 @@ async def _append_via_data_plane( raise DuplexFenceMismatchError(session.fence, fence) if not isinstance(append_plan, DuplexAppendPlan): raise TypeError("duplex plugin plan_append() must return DuplexAppendPlan") + append_config_key = (request_id, request_context.config_generation) + plugin = self._ctx.plugin + if ( + getattr(plugin, "append_configs_are_session_static", False) is True + and self._last_append_config_key == append_config_key + ): + drop_static_append_configs(append_plan.prompt) submission = DuplexStageSubmission( context=request_context, prompt=append_plan.prompt, @@ -333,6 +391,7 @@ async def _append_via_data_plane( raise DuplexFenceMismatchError(session.fence, fence) update = session.commit_append(reservation) session.bind_stage_request(stage_id, request_id, fence=fence) + self._last_append_config_key = append_config_key except BaseException: try: await self._ctx.stage_port.cleanup([request_id]) @@ -733,6 +792,16 @@ async def _on_model_listen( auto_response = self._out.auto_responds() close_reason: str | None = None emitted_response = False + pace = self._pace() + # barge_cut_on_model_yield: the model listens to a new user utterance + # while the client still plays the reply (its TTS may be done already). + if auto_response and self._cancel_model_response is not None and self._barge_cut_ready(): + pace.note_cut(time.monotonic(), "model_listen") + # The old Stage-1/2 requests may still be synthesizing queued + # audio: cancel like an explicit barge-in, so they are aborted, + # their unplayed output is discarded and the next append gets epoch+1. + await self._cancel_model_response("model_listen") + self._ctx.model_state.clear_continuation() self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id) if ( session.active_response_id is not None @@ -753,9 +822,32 @@ async def _on_model_listen( ) if non_terminal_auto_listen: self._ctx.services.spawn( - self.maybe_continue_response(expected_epoch=expected_epoch), name="duplex-continue" + self.maybe_continue_response(expected_epoch=expected_epoch, pace_kind="idle"), + name="duplex-continue", ) return close_reason, emitted_response + if ( + auto_response + and active_response_id is None + and session.active_request_id is not None + and model_result.get("end_of_turn") is not True + and expected_epoch is not None + and pace is not None + and pace.drives_replay(expected_epoch) + ): + # barge_cut_on_model_yield: a cut replayed the interrupted user's + # utterance into this epoch, where no reply is open. The cut reply's + # continuations would have kept Stage 0 of a client that went quiet + # moving; offer the same silence under the model-turn owner. Real + # input supersedes it, and any client append or a new reply ends it. + self._ctx.services.spawn( + self.maybe_continue_response( + expected_epoch=expected_epoch, + expected_model_turn_id=session.turn_id, + pace_kind="idle", + ), + name="duplex-continue", + ) if not auto_response and data_plane_request_id == session.active_request_id: session.clear_request() model_listen = model_result.get("model_listen") @@ -897,7 +989,9 @@ async def _send_one_model_output_event( tts_segment_ended and auto_response and (model_turn_id is not None or session.active_response_id is not None) + and not continue_on_stop_token(self._ctx.manager) ): + # Stop-token mode: the runner's Stage-0 hook already planned it. self._ctx.services.spawn( self.maybe_continue_response(expected_epoch=expected_epoch, expected_model_turn_id=model_turn_id), name="duplex-continue", @@ -1042,6 +1136,13 @@ async def _send_one_model_output_event( {"text_chars": max(0, int(mark_text_chars)), "audio_end_ms": max(0, int(mark_duration_ms))} ] payload["playback"] = session.playback_for_response(target_id).as_dict() + pace = self._pace() + if pace is not None and has_audio and mark_duration_ms is not None: + pace.on_audio_emit( + time.monotonic(), + max(0, mark_duration_ms - previous_sent_ms) / 1000.0, + draining=target_id is not None, + ) sample_rate_hz = model_result.get("sample_rate_hz") or model_result.get("audio_sample_rate_hz") if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) > 0: payload["sample_rate_hz"] = int(sample_rate_hz) @@ -1057,7 +1158,10 @@ async def _send_one_model_output_event( and model_result.get("stage_role") == "tts" and model_result.get("abort_data_plane_request") is True and auto_response + and not continue_on_stop_token(self._ctx.manager) ): + # Audio-emit trigger (default); in stop-token mode the runner plans + # the next unit when Stage 0 finishes the segment instead. self._ctx.services.spawn( self.maybe_continue_response(expected_epoch=expected_epoch), name="duplex-continue" ) @@ -1285,6 +1389,15 @@ def silence_continuation_is_stale( stale_model_turn_owner = not response_owned and ( session.active_response_id is not None or session.turn_id != expected_model_turn_id ) + if ( + stale_model_turn_owner + and session.turn_id == expected_model_turn_id + and session.active_response_turn_id == expected_model_turn_id + and same_turn_response_keeps_continuation(self._pace()) + ): + # Paced onset: unit 2 was planned before its own turn's response + # opened; that response opening is not a newer owner. + stale_model_turn_owner = False return stale_common_owner or stale_response_owner or stale_model_turn_owner async def maybe_continue_response( @@ -1292,6 +1405,7 @@ async def maybe_continue_response( *, expected_epoch: int | None, expected_model_turn_id: int | None = None, + pace_kind: str = "speech", ) -> None: session = self._ctx.session model_state = self._ctx.model_state @@ -1366,6 +1480,7 @@ async def maybe_continue_response( response_owned=response_owned, expected_epoch=expected_epoch, expected_model_turn_id=expected_model_turn_id, + pace_kind=pace_kind, ) except asyncio.CancelledError: raise diff --git a/vllm_omni/engine/duplex/session/pacing.py b/vllm_omni/engine/duplex/session/pacing.py new file mode 100644 index 00000000000..2077d0d5996 --- /dev/null +++ b/vllm_omni/engine/duplex/session/pacing.py @@ -0,0 +1,467 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Per-session output pacing for full-duplex auto responses (``duplex_session.pacing``). + +A duplex reply is paced by its input: one model unit per chunk period, each +continuation unit of synthesized silence submitted on the session's nominal +cadence (``compute_silence_continuation_deadline``). Played back with a jitter +buffer of ``B`` (``DUPLEX_CLIENT_PREBUFFER_S``, counted once per session), the +second unit's audio must arrive within ``B + d1`` of the first chunk's, where +``d1`` is that first chunk's length. With a short first chunk (0.12 s is common +at 16 concurrent sessions) the gap at unit 2 is about one chunk period plus the +latency difference of the two units -- a structural stall no Stage-2 batching +removes. + +``SessionPacing`` lets a client that went quiet after its commit ("client 1") +submit its speech continuations up to ``onset_lead_max_s`` (< one unit) ahead of +the nominal cadence, so the model timeline runs that much ahead of the wall +clock. Quiet means no input append since the commit for at least +``max(QUIET_GUARD_MIN_S, 1.5 x the longest append gap of the committed +utterance)``: a client that keeps its microphone open and streams in long +chunks is never mistaken for a quiet one. Any input append after the commit +makes the rest of the reply wall-clock paced again (lead 0); unsubmitted +silence is then dropped by the runner's existing validity checks. The nominal +deadline chain is never moved: the lead only shortens each sleep, so it does +not accumulate. + +The per-reply state (first chunk, playback mirror) follows the session's +active response lazily: whichever path ends a reply, the next decision sees a +different response id and resets, and a quiet client's pacing ends with the +reply its commit started. + +``PlaybackMirror`` is the API's model of that client's playback (the same rule +as ``duplex_rt``): ``B`` before the first chunk of the session, then play on +arrival, re-anchoring after a late chunk, truncated by a cut. + +``barge_cut_on_model_yield`` reuses the mirror for the playback cut: a new +user utterance over audible assistant audio (``barge_arm_min_silence_s`` of +non-speech input, then ``barge_arm_min_speech_s`` of speech, where pauses of +up to ``BARGE_SPEECH_HANGOVER_S`` do not end the utterance, starting at least +``barge_arm_min_playback_s`` after the reply's first audio) arms the session; +when the model then yields -- listens, or ends its turn -- while the client +still holds ``barge_cut_min_unplayed_s`` of unplayed audio, the reply is +cancelled like an explicit barge-in instead of draining its tail. Stage 0 +drops a session's context with its aborted request, so the real input units of +the interrupting utterance (at most ``BARGE_REPLAY_MAX_UNITS``, the latest) are +replayed as the new epoch's first appends. Until the client appends again or a +reply opens, a model listen there gets a silence unit, as the cut reply's +continuations would have given a client that went quiet. + +``fire_grid_ms`` (PersonaPlex-style lockstep) snaps the paced speech +continuations of quiet clients onto one grid of the shared monotonic clock, so +units of different sessions reach Stage 0 together: only earlier, within the +lead budget (the phase is paid by the lead). ``idle_grid`` rounds the +continuations after a non-terminal listen (no audio deadline) up to the grid. +Eager units -- the commit, the onset, CRITICAL, any real input -- never snap. + +Everything here is pure Python on the orchestrator loop; with both switches +off no ``SessionPacing`` exists and the runner keeps its historical path. +""" + +from __future__ import annotations + +import math +from collections import deque +from enum import Enum +from typing import TYPE_CHECKING + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.config import DUPLEX_CLIENT_PREBUFFER_S + +if TYPE_CHECKING: + from vllm_omni.config.stage_config import DuplexPacingConfig, DuplexSessionRuntimeConfig + from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession + +logger = init_logger(__name__) + +#: How often a paced continuation re-reads its lead while it waits. +PACE_RECHECK_S = 0.1 +#: Shortest quiet after a commit before the onset lead applies. +QUIET_GUARD_MIN_S = 0.25 +#: Non-speech input inside a user utterance that does not end it (playback-cut arming). +BARGE_SPEECH_HANGOVER_S = 0.3 +#: Input units of the interrupting utterance replayed into the new epoch after a cut. +BARGE_REPLAY_MAX_UNITS = 6 + + +class PlaybackMirror: + """The client's playback timeline as the API expects it (``duplex_rt`` rule), in API monotonic time.""" + + def __init__(self, *, prebuffer_s: float = DUPLEX_CLIENT_PREBUFFER_S) -> None: + self.prebuffer_s = prebuffer_s + #: Time the last scheduled audio ends (None before any audio). + self.play_end: float | None = None + #: Time of the current response's first audio (None until it emits). + self.response_first_audio_t: float | None = None + + def new_response(self) -> None: + self.response_first_audio_t = None + + def on_emit(self, now: float, seconds: float, *, current: bool = True) -> float: + """Account one emitted audio delta; return the stall it caused (0 for a response's first chunk). + + ``current=False``: the tail of an older, draining reply. It still plays + before anything newer, but is no part of the current response. + """ + stall = 0.0 + if self.play_end is None: + self.play_end = now + self.prebuffer_s + seconds + else: + if now > self.play_end: + stall = now - self.play_end + self.play_end = now + self.play_end += seconds + if not current: + return 0.0 + if self.response_first_audio_t is None: + # A gap before a response's first chunk is the model listening, not a stall. + self.response_first_audio_t = now + stall = 0.0 + return stall + + def unplayed(self, now: float) -> float: + """Seconds of scheduled audio the client has not played yet.""" + if self.play_end is None: + return 0.0 + return max(0.0, self.play_end - now) + + def cut(self, now: float) -> None: + """The client stopped playing at ``now`` (a cancel or barge-in).""" + if self.play_end is not None: + self.play_end = min(self.play_end, now) + self.new_response() + + +class PaceState(str, Enum): + LISTEN = "listen" + #: Committed; the reply's first audio has not been emitted yet. + ONSET = "onset" + PACED = "paced" + #: The mirror saw a stall above ``critical_stall_s``: full lead until it recovers. + CRITICAL = "critical" + + +class SessionPacing: + """Client class, pacing state and lead of one duplex session (see module docstring).""" + + def __init__(self, runtime_config: DuplexSessionRuntimeConfig, session: DuplexEngineSession) -> None: + self.runtime_config = runtime_config + self.config: DuplexPacingConfig = runtime_config.pacing + self.session = session + self.mirror = PlaybackMirror() + self.state = PaceState.LISTEN + #: A commit happened and no input append arrived since (client 1). + self.quiet_since_commit = False + #: Length (s) of the current response's first audio chunk. + self.d1: float | None = None + self._clean_deltas = 0 + #: The response the per-reply state belongs to. + self._response_id: str | None = None + self._commit_t: float | None = None + self._quiet_guard_s = QUIET_GUARD_MIN_S + self._last_append_t: float | None = None + self._utterance_gap_s = 0.0 + # Playback cut (``barge_cut_on_model_yield``). + self.barge_enabled = runtime_config.barge_cut_on_model_yield is True + self._quiet_input_s = 0.0 + self._speech_run_s = 0.0 + self._speech_run_over_playback = False + #: When the current reply was armed for a cut (None: not armed). + self.armed_at: float | None = None + #: Real input units Stage 0 took during the current user utterance (replayed after a cut). + self._utterance_units: deque[dict[str, object]] = deque(maxlen=BARGE_REPLAY_MAX_UNITS) + #: The epoch a cut's replay opened, until the client appends again or a reply opens there. + self._replay_epoch: int | None = None + + @property + def paces_continuations(self) -> bool: + return self.config.enabled is True + + def _sync(self) -> None: + """Reset the per-reply state once the session's active response changed (any end path).""" + response_id = self.session.active_response_id + if response_id == self._response_id: + return + if self._response_id is not None: + # The reply this state followed is gone: wall clock until the next commit. + self.quiet_since_commit = False + self._set_state(PaceState.LISTEN) + self._response_id = response_id + self._new_response() + + def _new_response(self) -> None: + self.mirror.new_response() + self.d1 = None + self._clean_deltas = 0 + self._disarm() + self._replay_epoch = None + + def quiet(self, now: float) -> bool: + """Quiet since the commit, for at least the guard (see module docstring).""" + return self.quiet_since_commit and self._commit_t is not None and now - self._commit_t >= self._quiet_guard_s + + # ---- inputs -------------------------------------------------------- # + + def on_commit(self, now: float) -> None: + self._sync() + self.quiet_since_commit = True + self._commit_t = now + # A client streaming in long chunks would look quiet between two of them. + self._quiet_guard_s = min(1.0, max(QUIET_GUARD_MIN_S, 1.5 * self._utterance_gap_s)) + self._utterance_gap_s = 0.0 + # A reply that already spoke before the commit (Stage-0 handoff) is paced from its known d1. + self._set_state(PaceState.ONSET if self.d1 is None else PaceState.PACED) + + def on_client_append(self, now: float, *, is_speech: bool, duration_s: float) -> None: + """Any input append after the commit hands the rest of the reply back to the wall clock.""" + self._sync() + self.quiet_since_commit = False + last = self._last_append_t + if last is not None and (self._commit_t is None or last > self._commit_t): + # Gaps inside one utterance only: the first append after a commit opens a new one. + self._utterance_gap_s = max(self._utterance_gap_s, now - last) + self._last_append_t = now + if self.barge_enabled: + self._replay_epoch = None + self._track_utterance(now, is_speech=is_speech, duration_s=duration_s) + + def _track_utterance(self, now: float, *, is_speech: bool, duration_s: float) -> None: + """Arm the playback cut once a new user utterance runs over audible reply audio.""" + if not is_speech: + self._quiet_input_s += duration_s + if self._quiet_input_s > BARGE_SPEECH_HANGOVER_S: + # A pause, not a dip inside the utterance: the speech run ends. + self._speech_run_s = 0.0 + self._speech_run_over_playback = False + if self.armed_at is None: + self._utterance_units.clear() + return + if self._speech_run_s == 0.0: + # A speech run starts: it is a new utterance over the reply only + # after enough quiet input and once the reply is audibly playing. + self._speech_run_over_playback = ( + self._quiet_input_s >= self.runtime_config.barge_arm_min_silence_s and self._playing(now) + ) + self._quiet_input_s = 0.0 + self._speech_run_s += duration_s + if ( + self.armed_at is None + and self._speech_run_over_playback + and self._speech_run_s >= self.runtime_config.barge_arm_min_speech_s + ): + self.armed_at = now + + def _playing(self, now: float) -> bool: + first_audio_t = self.mirror.response_first_audio_t + return ( + first_audio_t is not None + and now - first_audio_t >= self.runtime_config.barge_arm_min_playback_s + and self.mirror.unplayed(now) > 0.0 + ) + + @property + def armed(self) -> bool: + return self.armed_at is not None + + def cut_ready(self, now: float) -> bool: + """Whether a model yield now should cut the reply the client is still playing.""" + self._sync() + return ( + self.barge_enabled + and self.armed_at is not None + and self._response_id is not None + and self.mirror.unplayed(now) >= self.runtime_config.barge_cut_min_unplayed_s + ) + + def note_cut(self, now: float, reason: str) -> None: + logger.info( + "Duplex barge cut (%s) response=%s armed %.2f s ago, %.2f s unplayed", + reason, + self._response_id, + now - (self.armed_at if self.armed_at is not None else now), + self.mirror.unplayed(now), + ) + + def on_input_unit(self, payload: dict[str, object]) -> None: + """A real input unit was planned for Stage 0 while a user utterance runs: kept for the cut replay. + + Recorded when the append is planned, not when Stage 0 accepts it: an + append still queued when the cut moves the epoch is dropped by the + stale-epoch check, so the replay is its only delivery. + """ + if not self.barge_enabled or (self._speech_run_s == 0.0 and self.armed_at is None): + return + if ( + self.armed_at is not None + and payload.get("is_speech") is not True + and len(self._utterance_units) == BARGE_REPLAY_MAX_UNITS + ): + # Non-speech input after arming is what the model hears before it + # yields: kept while there is room, but it never pushes the start + # of the utterance out of the replay. + return + self._utterance_units.append(dict(payload)) + + def take_replay(self) -> list[dict[str, object]]: + """The units of the utterance that armed the cut (empty when not armed); call before the cancel. + + The arming of a reply that already ended (the route-M listen abort does + not ask ``cut_ready`` first) is reset here, not replayed. + """ + if not self.barge_enabled: + return [] + self._sync() + units = list(self._utterance_units) if self.armed_at is not None else [] + self._utterance_units.clear() + return units + + def on_replay(self, epoch: int) -> None: + """A cut replayed the interrupting utterance as the first appends of ``epoch``.""" + self._replay_epoch = epoch + + def drives_replay(self, epoch: int) -> bool: + """Whether a model listen in ``epoch`` should get a silence unit. + + Only in the epoch a cut's replay opened, while no reply opened there and + the client appended nothing since the cut: the cut reply's + continuations would have kept Stage 0 of a quiet client moving (it + answers from the silence after the utterance), the new epoch has none. + """ + if self._replay_epoch != epoch: + return False + self._sync() # a reply that opened meanwhile ends it + return self._replay_epoch == epoch + + def _disarm(self) -> None: + self.armed_at = None + self._speech_run_over_playback = False + self._utterance_units.clear() + + def on_epoch_advance(self, now: float) -> None: + """A barge-in or cancel: the client stops the old reply now.""" + self.mirror.cut(now) + self.quiet_since_commit = False + self._response_id = None + self._new_response() + self._set_state(PaceState.LISTEN) + + # ---- outputs ------------------------------------------------------- # + + def on_audio_emit(self, now: float, seconds: float, *, draining: bool = False) -> None: + if seconds <= 0: + return + self._sync() + if draining: + # An older reply's tail: playback only, never the current reply's d1. + self.mirror.on_emit(now, seconds, current=False) + return + first = self.mirror.response_first_audio_t is None + stall = self.mirror.on_emit(now, seconds) + if first: + self.d1 = seconds + if self.state == PaceState.ONSET: + self._set_state(PaceState.PACED) + return + if stall > self.config.critical_stall_s: + self._clean_deltas = 0 + if self.state == PaceState.PACED: + self._set_state(PaceState.CRITICAL) + elif self.state == PaceState.CRITICAL: + self._clean_deltas += 1 + if self._clean_deltas >= 2: + self._set_state(PaceState.PACED) + + # ---- decisions ----------------------------------------------------- # + + def lead_s(self, now: float) -> float: + """How far ahead of its nominal deadline the next speech continuation may submit.""" + self._sync() + lead_max = float(self.config.onset_lead_max_s) + if not self.config.enabled or lead_max <= 0: + return 0.0 + if not self.quiet(now) or self.state == PaceState.LISTEN: + return 0.0 + if self.state == PaceState.CRITICAL or self.d1 is None: + return lead_max + return min(lead_max, max(0.0, float(self.config.onset_lead_target_s) - self.d1)) + + def recheck_in(self, now: float) -> float: + """How long a waiting continuation may sleep before its lead can change.""" + if self.quiet_since_commit and self._commit_t is not None: + guard_left = self._commit_t + self._quiet_guard_s - now + if guard_left > 0: + return min(PACE_RECHECK_S, guard_left) + return PACE_RECHECK_S + + def skip_response_wait(self, now: float) -> bool: + """Whether the stop-token trigger may continue unit 2 before its response opens.""" + self._sync() + return ( + self.config.enabled is True + and self.config.onset_skip_response_wait is True + and self.quiet(now) + and self.state == PaceState.ONSET + ) + + def speech_fire_at(self, nominal: float, now: float) -> float: + """Submission time of a speech continuation whose nominal deadline is ``nominal``.""" + return nominal - self.lead_s(now) + + def idle_fire_at(self, nominal: float, now: float) -> float: + """Submission time of a continuation after a non-terminal listen (no audio deadline).""" + self._sync() + grid_s = self.config.fire_grid_ms / 1000.0 + if grid_s <= 0 or self.config.idle_grid is not True or not self.quiet(now): + return nominal + return grid_s * math.ceil(nominal / grid_s - 1e-9) + + def snap(self, fire_at: float, nominal: float, now: float) -> float: + """Move a paced speech continuation earlier onto the fire grid, within the lead budget. + + Identity without ``fire_grid_ms``, outside PACED, or when the earlier + grid point is already past or would lead more than ``onset_lead_max_s``. + """ + if not self.snap_holds(now): + return fire_at + grid_s = self.config.fire_grid_ms / 1000.0 + lower = grid_s * math.floor(fire_at / grid_s + 1e-9) + if lower >= now and nominal - lower <= self.config.onset_lead_max_s + 1e-9: + return lower + return fire_at + + def snap_holds(self, now: float) -> bool: + """Whether a grid point applies (a grid, a paced reply and a quiet client).""" + self._sync() + return self.config.fire_grid_ms > 0 and self.state == PaceState.PACED and self.quiet(now) + + def _set_state(self, state: PaceState) -> None: + self.state = state + + +def _unit_continuation(manager: object) -> object: + return getattr(getattr(manager, "runtime_config", None), "unit_continuation", None) + + +def unit_continuation_depth(manager: object) -> int: + """``duplex_session.unit_continuation.depth``: continuations of one session in flight (default 1).""" + depth = getattr(_unit_continuation(manager), "depth", 1) + return depth if isinstance(depth, int) and not isinstance(depth, bool) and depth >= 1 else 1 + + +def continue_on_stop_token(manager: object) -> bool: + """``unit_continuation.continue_on: stop_token``: the runner plans the next unit itself.""" + return getattr(_unit_continuation(manager), "continue_on", None) == "stop_token" + + +def same_turn_response_keeps_continuation(pace: object) -> bool: + """Whether a model-turn continuation stays valid once its own turn's response opens. + + Only with ``onset_skip_response_wait``: the stop-token trigger then submits + unit 2 under the model-turn owner before the response exists, and the + response opening while it waits must not drop it. + """ + if not isinstance(pace, SessionPacing): + return False + return pace.config.enabled is True and pace.config.onset_skip_response_wait is True diff --git a/vllm_omni/engine/duplex/session/runner.py b/vllm_omni/engine/duplex/session/runner.py index ca979393864..7c818a1bd2f 100644 --- a/vllm_omni/engine/duplex/session/runner.py +++ b/vllm_omni/engine/duplex/session/runner.py @@ -98,6 +98,7 @@ from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity from vllm_omni.engine.duplex.session.model_channel import ModelChannel +from vllm_omni.engine.duplex.session.pacing import SessionPacing, continue_on_stop_token, unit_continuation_depth from vllm_omni.engine.duplex.turn_detection import ( TurnDetectionResult, ) @@ -125,6 +126,10 @@ class _Internal: _CANCEL_EVENTS = frozenset({"input.cancel", "response.cancel", "barge_in", "output_audio_buffer.clear"}) +#: How long the stop-token continuation trigger waits for its unit's response +#: to open before continuing under the model-turn owner instead. +_STOP_TOKEN_RESPONSE_WAIT_S = 2.0 + def compute_silence_continuation_deadline( *, @@ -203,6 +208,7 @@ def __init__( close_from_runtime=self._close_from_runtime, schedule_silence_continuation=self._schedule_silence_continuation, abort_request=self._abort_request_background, + cancel_model_response=self._cancel_model_response, ) self.control = SessionControl( self.ctx, @@ -229,6 +235,12 @@ def start(self) -> None: # Every client speaks the Realtime protocol now; the old runner forced # the ACK-only playback ledger for that path. session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value + runtime_config = getattr(self.manager, "runtime_config", None) + if ( + getattr(getattr(runtime_config, "pacing", None), "enabled", False) is True + or getattr(runtime_config, "barge_cut_on_model_yield", False) is True + ): + self.run.pace = SessionPacing(runtime_config, session) self.control.init_turn_detection() self._worker = self._loop.create_task(self._run(), name=f"duplex-session-{session.session_id}") self.emit({"type": "session.created", "session": session.as_public_dict()}) @@ -248,6 +260,22 @@ def on_stage_output( """Accept one stage output (orchestrator loop); return True when it must not be forwarded.""" decision: DuplexOutputDecision | None = None project = False + pace = self.run.pace + if ( + stage_id == 0 + and pace is not None + and pace.armed + and self.plugin.stage0_naturally_ended( + segment_finished=context.segment_finished, + segment_token_ids=context.segment_token_ids, + segment_output_metadata=context.segment_output_metadata, + output=output, + ) + ): + # The model ended its turn after the user started a new utterance + # over its audio: cut the tail (on the mailbox, ordered with this + # segment's own outputs). + self._mailbox.put_nowait(_Internal("model_yield_cut")) if stage_id < context.final_stage_id: decision = self.model.decide_output(stage_id, output, context) # Optional mid-pipeline projection: client sees this stage; TTS still runs. @@ -259,6 +287,10 @@ def on_stage_output( self.run.concurrent_turn_requests_released = True consume = decision is not None or stage_id >= context.final_stage_id project_intermediate = self.plugin.projects_intermediate_outputs and stage_id == 0 + if stage_id == 0 and decision is None and context.segment_finished and continue_on_stop_token(self.manager): + # unit_continuation.continue_on: stop_token; before the early return + # below so forwarded (not consumed) segments are covered too. + self._continue_response_on_stop_token(context) if not consume and not project and not project_intermediate: # Stage0 text without a direct decision feeds the TTS stage as before. # Its metrics still have to reach the client: before sessions moved @@ -520,6 +552,10 @@ async def _on_internal(self, item: _Internal) -> None: if item.kind == "commit": await self._on_commit(dict(item.payload)) return + if item.kind == "model_yield_cut": + if not self.run.closing and self.session.state == DuplexSessionState.OPEN: + await self.model.cut_on_model_yield() + return if item.kind == "run_payload": await self._run_internal_payload(dict(item.payload)) return @@ -755,7 +791,7 @@ async def _barge_in_for_overlap(self, event: dict[str, object], decision: dict[s old_response_id = session.active_response_id committed_ms = session.playback.committed_ms helpers.commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session, pace=self.run.pace) self.emit( { "type": "audio.cancelled", @@ -773,7 +809,7 @@ async def _barge_in_for_overlap(self, event: dict[str, object], decision: dict[s old_epoch = session.epoch committed_ms = session.playback.committed_ms helpers.commit_played_response_history(session, session.last_response_id, committed_ms) - new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session, pace=self.run.pace) self.emit( { "type": "audio.cancelled", @@ -893,6 +929,12 @@ async def _on_append_audio(self, event: dict[str, object]) -> None: payload["is_speech"] = ( False if video_only else overlap_policy.input_looks_like_speech(self.session, event, payload) ) + if self.run.pace is not None: + self.run.pace.on_client_append( + time.monotonic(), + is_speech=payload["is_speech"] is True, + duration_s=overlap_policy.input_audio_duration_ms(event, payload) / 1000.0, + ) auto_responds = self._session_auto_responds() defer_append = False buffer_overlap_audio = True @@ -1001,6 +1043,8 @@ async def _on_append_audio(self, event: dict[str, object]) -> None: if append_payload is None: self._maybe_schedule_vad_commit(vad_result) return + if self.run.pace is not None: + self.run.pace.on_input_unit(append_payload) await self._start_append(append_payload, final=False, pcm_reservation=pcm_reservation) self._maybe_schedule_vad_commit(vad_result) @@ -1030,6 +1074,72 @@ def _clear_completed_pending_silence(self) -> None: if task is not None and task.done(): self.model_state.pending_silence_task = None self.model_state.pending_silence_owner_id = None + pending = self.model_state.pending_silence_tasks + if pending: + self.model_state.pending_silence_tasks = tuple(t for t in pending if not t.done()) + + def _silence_in_flight(self) -> tuple[asyncio.Task[bool], ...]: + """Prune finished continuations; return those still scheduled or submitting (depth > 1).""" + self._clear_completed_pending_silence() + return self.model_state.pending_silence_tasks + + def _continue_response_on_stop_token(self, context: DuplexOutputContext) -> None: + """Plan the next unit when Stage 0 finishes a spoken segment (stop-token mode). + + The audio-emit trigger holds each session to one unit in flight; here + only the planning moves earlier. The anchor cadence still paces + submissions and client emission still follows stage-output arrival. + """ + session = self.session + if ( + self.run.closing + or session.state != DuplexSessionState.OPEN + or session.lease.terminal_reason is not None + or not session.capabilities.supports_core_resumable_request + or not self.out.auto_responds() + ): + return + expected_epoch = context.identity.fence.epoch + self.spawn( + self._continue_response_on_stop_token_task(expected_epoch), + name="duplex-continue-stop-token", + ) + + async def _continue_response_on_stop_token_task(self, expected_epoch: int) -> None: + """Schedule the next silence unit for a segment that just finished. + + The unit's response opens with its first Stage 1/2 output; wait briefly + for it so the continuation is response-owned like the audio-emit + trigger's, else continue under the model-turn owner. + """ + session = self.session + model_state = self.model_state + waited_since: float | None = None + pace = self.run.pace + while ( + session.active_response_id is None + and session.state == DuplexSessionState.OPEN + and not self.run.closing + and session.epoch == expected_epoch + and session.lease.terminal_reason is None + and session.active_request_id is not None + ): + if pace is not None and pace.skip_response_wait(time.monotonic()): + # Paced onset (client quiet since its commit): continue unit 2 + # under the model-turn owner now instead of waiting for its response. + break + if waited_since is None: + waited_since = time.monotonic() + elif time.monotonic() - waited_since > _STOP_TOKEN_RESPONSE_WAIT_S: + break + await asyncio.sleep(0.02) + if session.state == DuplexSessionState.CLOSED or self.run.closing: + model_state.clear_continuation() + return + await self.model.maybe_continue_response( + expected_epoch=expected_epoch, + expected_model_turn_id=None if session.active_response_id is not None else session.turn_id, + ) def _mark_pending_silence_superseded(self) -> None: task = self.model_state.pending_silence_task @@ -1066,9 +1176,13 @@ async def _start_append( # A real (non-silence) input re-anchors the silence pacing # chain: its submission becomes the anchor and the stored # deadline is cleared until the next continuation sets one. + multi_unit = unit_continuation_depth(self.manager) > 1 + def _reanchor_chain(submit_time: float) -> None: model_state.last_native_submit_monotonic = submit_time model_state.silence_deadline_monotonic = None + if multi_unit: + model_state.native_input_generation += 1 on_append_accepted = _reanchor_chain append_epoch = session.epoch @@ -1133,6 +1247,8 @@ def _reanchor_chain(submit_time: float) -> None: ) if silence_continuation: model_state.pending_silence_task = task + if unit_continuation_depth(self.manager) > 1: + model_state.pending_silence_tasks = (*model_state.pending_silence_tasks, task) task.add_done_callback(attempt.clear_pending_silence) # Let this wire-order effect start before the next mailbox item can cancel it. await asyncio.sleep(0) @@ -1166,12 +1282,22 @@ async def _schedule_silence_continuation( response_owned: bool, expected_epoch: int | None, expected_model_turn_id: int | None, + pace_kind: str = "speech", ) -> bool: session = self.session model_state = self.model_state - self._clear_completed_pending_silence() - pending_silence = model_state.pending_silence_task - if pending_silence is not None and not pending_silence.done(): + # unit_continuation.depth > 1: submissions beyond the first wait in the + # engine's per-session streaming-update queue instead of the runner. + depth = unit_continuation_depth(self.manager) + if depth > 1: + pending_window = self._silence_in_flight() + pending_silence = pending_window[0] if len(pending_window) >= depth else None + else: + self._clear_completed_pending_silence() + pending_silence = model_state.pending_silence_task + if pending_silence is not None and pending_silence.done(): + pending_silence = None + if pending_silence is not None: if pending_silence is asyncio.current_task(): return False try: @@ -1183,10 +1309,14 @@ async def _schedule_silence_continuation( return False except Exception: return False - self._clear_completed_pending_silence() - pending_silence = model_state.pending_silence_task - if pending_silence is not None and not pending_silence.done(): - return False + if depth > 1: + if len(self._silence_in_flight()) >= depth: + return False + else: + self._clear_completed_pending_silence() + pending_silence = model_state.pending_silence_task + if pending_silence is not None and not pending_silence.done(): + return False append_tail = self.tasks.append_tail if (append_tail is None or append_tail.done()) and self._real_input_waiting(): return False @@ -1195,6 +1325,11 @@ async def _schedule_silence_continuation( # accepted between the snapshot and this continuation's submission # re-anchors the chain; _still_valid() then skips the stale unit. anchor = model_state.last_native_submit_monotonic + # At depth > 1 the anchor also moves with sibling silence submissions; + # generation and silence sequence tell the two moves apart: a real + # append invalidates the plan, siblings only fill the depth window. + generation = model_state.native_input_generation + planned_silence_seq = model_state.silence_append_seq # Align the next silence unit to submission_time_N + chunk_period and # sleep only the remaining budget. The deadline is stored by the # acceptance callback when the append actually submits, so skipped or @@ -1205,7 +1340,11 @@ async def _schedule_silence_continuation( last_submit=anchor, current_deadline=model_state.silence_deadline_monotonic, ) - if delay_s > 0: + pace = self.run.pace + if pace is not None and pace.paces_continuations and pace.quiet_since_commit and self.out.auto_responds(): + # Only a client quiet since its commit can lead; everyone else keeps the plain sleep. + await self._paced_continuation_wait(pace, delay_s, pace_kind) + elif delay_s > 0: await asyncio.sleep(delay_s) if ( self.tasks.append_tail is not append_tail @@ -1221,10 +1360,19 @@ async def _schedule_silence_continuation( return False def _still_valid() -> bool: - return ( + if depth > 1: + # A real append since planning (generation), or a full window of + # sibling silence submissions, makes this unit outdated. + chain_current = ( + model_state.native_input_generation == generation + and model_state.silence_append_seq - planned_silence_seq < depth + ) + else: # The anchor changed (a real append was accepted) after this # continuation was planned; the unit is outdated. - model_state.last_native_submit_monotonic == anchor + chain_current = model_state.last_native_submit_monotonic == anchor + return ( + chain_current and not self._real_input_waiting() and not self.model.silence_continuation_is_stale( request_id=request_id, @@ -1242,6 +1390,8 @@ def _on_append_accepted(submit_time: float) -> None: # submission. Small delays keep the planned cadence so ordinary # jitter does not accumulate as drift. model_state.last_native_submit_monotonic = submit_time + if depth > 1: + model_state.silence_append_seq += 1 model_state.silence_deadline_monotonic = ( submit_time + chunk_period_s if submit_time > next_silence_deadline else next_silence_deadline ) @@ -1256,6 +1406,37 @@ def _on_append_accepted(submit_time: float) -> None: ) return task is not None + @staticmethod + async def _paced_continuation_wait(pace: SessionPacing, delay_s: float, pace_kind: str) -> None: + """Sleep until a paced continuation of a quiet client may submit (``duplex_session.pacing``). + + ``nominal`` is the unchanged cadence deadline. A speech continuation + fires up to the lead earlier; the lead is re-read while waiting (d1 + arrives, input resumes, the quiet guard passes, CRITICAL), so a unit + never fires earlier than the lead allowed when it fired. A fire-grid + point is chosen once and kept while the session stays paced and quiet. + """ + nominal = time.monotonic() + delay_s + grid_point: float | None = None + first = True + while True: + now = time.monotonic() + if pace_kind == "idle": + target = pace.idle_fire_at(nominal, now) + else: + target = pace.speech_fire_at(nominal, now) + if first: + snapped = pace.snap(target, nominal, now) + grid_point = snapped if snapped != target else None + if grid_point is not None and pace.snap_holds(now): + target = grid_point + else: + grid_point = None + first = False + if now >= target: + return + await asyncio.sleep(min(target - now, pace.recheck_in(now))) + # ------------------------------------------------------------------ # # Runtime-initiated close # # ------------------------------------------------------------------ # @@ -1344,7 +1525,7 @@ async def _on_cancel(self, event: dict[str, object]) -> None: old_response_id = session.active_response_id committed_ms = session.playback.committed_ms helpers.commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session, pace=self.run.pace) self.emit( { "type": "audio.cancelled", @@ -1362,7 +1543,7 @@ async def _on_cancel(self, event: dict[str, object]) -> None: old_epoch = session.epoch committed_ms = session.playback.committed_ms helpers.commit_played_response_history(session, session.last_response_id, committed_ms) - new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session, pace=self.run.pace) self.emit( { "type": "audio.cancelled", @@ -1381,7 +1562,7 @@ async def _on_cancel(self, event: dict[str, object]) -> None: old_response_id = session.active_response_id committed_ms = session.playback.committed_ms helpers.commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session, pace=self.run.pace) self.emit( { "type": "audio.cancelled", @@ -1425,6 +1606,7 @@ async def _cancel_active_response( *, reason: str, notify: bool = True, + finalize_fence: bool = False, ) -> bool: session = self.session has_running_task = active_task is not None and not active_task.done() @@ -1435,6 +1617,7 @@ async def _cancel_active_response( if not has_running_task and old_request_id is None and session.active_response_id is None and not draining_ids: return False + cancelled_fence = session.fence old_epoch = session.epoch old_response_id = session.active_response_id committed_ms = session.playback.committed_ms @@ -1473,7 +1656,7 @@ async def _cancel_active_response( older_abort_ids = [request_id for request_id in abort_ids if request_id not in cancelled_ids] if older_abort_ids: session.release_resources_for_request_ids(older_abort_ids) - new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session, pace=self.run.pace) if old_request_id is not None: # Release projector/parser cursors so cancelled epochs do not # accumulate until the whole session closes. @@ -1490,6 +1673,12 @@ async def _cancel_active_response( # asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11. except (TimeoutError, asyncio.TimeoutError): pass + if finalize_fence: + # Explicit cancel commands signal the fence in ``_on_cancel`` + # after this method returns. A model-owned listen has no outer + # command, so finalize it here or the old request state would stay + # registered after the stage abort. + await self.model.signal_cancel_fence(cancelled_fence) if notify: self.emit( { @@ -1518,6 +1707,27 @@ async def _cancel_active_response( ) return True + async def _cancel_model_response(self, reason: str) -> bool: + """A model-owned cancel (route-M listen abort, playback cut) of the active response. + + Stage 0 drops a session's context with its aborted request, so the + user utterance it heard in the old epoch is gone. With + ``barge_cut_on_model_yield`` the units of that utterance are replayed, + in order, as the new epoch's first appends. + """ + pace = self.run.pace + replay = pace.take_replay() if pace is not None else [] + cancelled = await self._cancel_active_response(None, reason=reason, finalize_fence=True) + if cancelled and self.session.state == DuplexSessionState.OPEN and not self.run.closing: + if replay: + # The new epoch has no reply whose continuations keep Stage 0 + # moving if the client went quiet; its listens get silence + # until the client appends again (``ModelChannel._on_model_listen``). + pace.on_replay(self.session.epoch) + for payload in replay: + await self._start_append(payload, final=False) + return cancelled + async def _abort_request_background(self, request_ids: list[str], *, notify: bool) -> None: try: await self.stage_port.abort_requests(request_ids) @@ -1531,7 +1741,7 @@ async def _abort_request_background(self, request_ids: list[str], *, notify: boo def _cancel_pending_input(self, *, reason: str) -> None: session = self.session cancelled = session.cancel_pending_input() - helpers.advance_barge_in_epoch(session) + helpers.advance_barge_in_epoch(session, pace=self.run.pace) self.emit( { "type": "input.cancelled", @@ -1820,6 +2030,11 @@ async def _commit_and_start_auto_response(self, event: dict[str, object], *, rea chunk_period_ms=session.capabilities.chunk_period_ms or 1000, ) final_payload = commit_reservation.payload + if self.run.pace is not None and final_payload is not None: + # The commit flushes the utterance's last (zero-padded) unit; a + # playback cut after it aborts the epoch Stage 0 heard it in, so the + # cut replays it too (as a plain unit: ``final`` only padded it). + self.run.pace.on_input_unit({key: value for key, value in final_payload.items() if key != "final"}) if model_state.committed_audio_payload is not None: if final_payload is not None: final_payload = overlap_policy.merge_audio_payloads(model_state.committed_audio_payload, final_payload) @@ -1835,6 +2050,8 @@ async def _commit_and_start_auto_response(self, event: dict[str, object], *, rea model_state.deferred_response_create = False model_state.input_since_commit = False model_state.speech_since_commit = False + if self.run.pace is not None: + self.run.pace.on_commit(time.monotonic()) data_plane_turn_id = session.turn_id committed = helpers.commit_audio_input( session, diff --git a/vllm_omni/engine/duplex_orchestrator.py b/vllm_omni/engine/duplex_orchestrator.py index 39719453aec..9b099280c57 100644 --- a/vllm_omni/engine/duplex_orchestrator.py +++ b/vllm_omni/engine/duplex_orchestrator.py @@ -310,6 +310,7 @@ def _sync_bridge_state( "fence": context.fence, "epoch": context.fence.epoch, "turn_id": context.fence.turn_id, + "config_generation": context.config_generation, "session_config": dict(context.session_config), "runtime_config": dict(context.runtime_config), } diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 8f55280cd23..0fb4bfc5534 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -101,7 +101,7 @@ def _event_driven_orch_enabled(*, default: bool = False) -> bool: def _event_driven_orch_default_for_pipeline(pipeline_model_type: str | None) -> bool: """Return whether a pipeline has a validated event-driven default.""" - return pipeline_model_type == "qwen3_tts" + return pipeline_model_type in ("qwen3_tts", "minicpmo_4_5") def _build_terminal_empty_output( diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index c19a5cf7511..9c4c7b66938 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -12,7 +12,9 @@ import contextlib import os +import queue import signal +import time from typing import Any import vllm.v1.engine.core as _vllm_engine_core_module @@ -44,6 +46,30 @@ _SIGNAL_EXIT_BASE = 128 +#: Floor of a park bounded by a held chunk's release time (deadline batching). +_MIN_RELEASE_WAIT_S = 5e-4 + + +def _scheduler_next_release_in(scheduler: object) -> float | None: + """Seconds until the scheduler releases a held request, if it holds any (else None).""" + next_release_in = getattr(scheduler, "next_release_in", None) + if not callable(next_release_in): + return None + value = next_release_in() + return value if type(value) is float else None + + +def _has_step_output(outputs: dict[int, Any] | None) -> bool: + """Whether a step returned anything but the scheduler stats ``log_stats`` attaches to every step.""" + for output in (outputs or {}).values(): + if ( + getattr(output, "outputs", True) + or getattr(output, "finished_requests", None) + or getattr(output, "utility_output", None) is not None + ): + return True + return False + def _install_phase_locks(kwargs: dict[str, Any], local_dp_rank: int) -> None: """Wrap ``kwargs["executor_class"]`` with the SH/EX phase-lock guard. @@ -130,6 +156,79 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: _bind_first_audio_sink(self.model_executor, self.output_queue, self.scheduler) if _bind_native_data_plane_ready_sink(self.model_executor, self.scheduler): logger.info("Bound native MRv2 connector readiness directly to the scheduler inbox.") + self._init_idle_wait() + + # ------------------------------------------------------------------ # + # Event-driven idle wait (additional_config.chunk_idle_wait_s > 0) # + # ------------------------------------------------------------------ # + # vLLM's loop never blocks while ``scheduler.has_requests()``; a stage whose + # resident requests only wait for upstream chunks (Code2Wav in duplex: + # WAITING_FOR_CHUNK still counts as unfinished) then spins schedule -> + # zero-token execute -> sleep(1 ms) ~300 times per second. With the switch + # on, a step that executed nothing and produced no output parks the loop on + # the input queue until a client request, a chunk-ready WAKEUP from the + # receive thread, or the fallback bound -- only the waiting changes, never + # what a step does. + + def _init_idle_wait(self) -> None: + from vllm_omni.distributed.omni_connectors.transfer_adapter.base import stage_idle_wait_s + + self._idle_wait_s = stage_idle_wait_s( + getattr(getattr(self.scheduler, "vllm_config", None), "additional_config", None) + ) + self._last_step_idle = False + if self._idle_wait_s <= 0: + return + adapter = getattr(self.scheduler, "chunk_transfer_adapter", None) + set_callback = getattr(adapter, "set_chunk_ready_callback", None) + if not callable(set_callback): + # Without the receive thread's WAKEUP every chunk would wait for the + # fallback bound (e.g. the native MRv2 data plane): keep vLLM's loop. + logger.warning("[StageEngineCoreProc] event-driven idle wait off: this stage has no chunk-ready wakeup") + self._idle_wait_s = 0.0 + return + set_callback(self._post_idle_wakeup) + logger.info("[StageEngineCoreProc] event-driven idle wait: fallback %.3f s", self._idle_wait_s) + + def _post_idle_wakeup(self) -> None: + """Wake a parked loop (called from the chunk receive thread).""" + self.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None)) + + def _process_input_queue(self) -> None: + if ( + getattr(self, "_last_step_idle", False) + and not self.batch_queue + and self.input_queue.empty() + and self.is_running() + ): + self._last_step_idle = False + timeout = self._idle_wait_s + release_in = _scheduler_next_release_in(self.scheduler) + if release_in is not None: + # A held chunk must be re-planned by its release time, not the fallback bound. + timeout = min(timeout, max(_MIN_RELEASE_WAIT_S, release_in)) + try: + req = self.input_queue.get(timeout=timeout) + except queue.Empty: + pass # fallback bound: step once (deadlines, missed hints), then park again + else: + self._handle_client_request(*req) + super()._process_input_queue() + + def _process_engine_step(self) -> bool: + if getattr(self, "_idle_wait_s", 0.0) <= 0: + return super()._process_engine_step() + # Same as vLLM's EngineCoreProc._process_engine_step, except that an + # idle step parks in _process_input_queue instead of sleeping 1 ms. + outputs, model_executed = self.step_fn() + for output in outputs.items() if outputs else (): + self.output_queue.put_nowait(output) + self.post_step(model_executed) + has_requests = self.scheduler.has_requests() + self._last_step_idle = not model_executed and not _has_step_output(outputs) and has_requests + if not model_executed and not self._last_step_idle and has_requests: + time.sleep(0.001) + return model_executed def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]: """Preserve omni payloads when vLLM builds its scheduler request.""" From 91097aba238ba70cf7f9de95071c83afd9e36c79 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:29 +0800 Subject: [PATCH 76/89] [Perf][Worker] Multi-frame decode hook for Talker K-step Worker multi-frame decode hook for Talker K-step Expose a model-level multi_frame_decode_hook on GPU runners and waiting-safe dispatch in the AR scheduler. Default single-frame execution paths are unchanged. Signed-off-by: BeatSeat --- tests/worker/test_native_duplex_hooks.py | 13 ++++-- vllm_omni/core/sched/omni_ar_scheduler.py | 35 +++++++++++++-- vllm_omni/model_executor/duplex_sampling.py | 28 ++++++++++-- vllm_omni/worker/gpu_ar_model_runner.py | 21 +++++++++ .../worker/gpu_generation_model_runner.py | 1 + vllm_omni/worker/gpu_model_runner.py | 45 ++++++++++++++++++- 6 files changed, 132 insertions(+), 11 deletions(-) diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index 65fb52ad017..4d37a1ab56c 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -392,7 +392,7 @@ def test_plain_model_hook_resolution_does_not_allocate_duplex_tracking(): assert not hasattr(runner, "_duplex_sampling_helper") -def test_minicpmo_non_duplex_sample_skips_duplex_row_scan(): +def test_minicpmo_non_duplex_step_publishes_empty_rows_without_scanning(): from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner calls = [] @@ -415,7 +415,9 @@ def test_minicpmo_non_duplex_sample_skips_duplex_row_scan(): ) assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" - assert calls == [] + # The empty row set is published every step so the model's sample() can + # take its no-duplex fast path without scanning prompt tokens. + assert calls == ["prepare"] def test_minicpmo_duplex_sample_clears_stale_rows_once_without_scanning(): @@ -442,7 +444,7 @@ def test_minicpmo_duplex_sample_clears_stale_rows_once_without_scanning(): assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" - assert calls == [()] + assert calls == [(), ()] def test_generic_ar_runner_has_no_minicpmo_sampler_state_or_typeerror_probe(): @@ -2890,6 +2892,9 @@ def convert_tokens_to_ids(self, token): ) sampled = model.sample(logits, sampling_metadata) + # The sampler decides on the device; the runner's next step commits the + # session latches once the token is on the host. + model._commit_minicpmo45_duplex_pending_samples() assert sampled is not None assert sampled.sampled_token_ids.tolist() == [[151705]] @@ -3155,6 +3160,7 @@ def test_minicpmo_stage0_turn_eos_is_forwarded_not_pending(): logits[0, 151717] = 30.0 sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201])) + model._commit_minicpmo45_duplex_pending_samples() assert sampled.sampled_token_ids.tolist() == [[151717]] assert state.pending_terminator_token is None @@ -3166,6 +3172,7 @@ def test_minicpmo_stage0_turn_eos_is_forwarded_not_pending(): logits = torch.full((1, 151723), -100.0) logits[0, 151718] = 30.0 sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201, 151717])) + model._commit_minicpmo45_duplex_pending_samples() assert sampled.sampled_token_ids.tolist() == [[151718]] assert state.pending_terminator_token == 151718 diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 5044abbc38d..53e747198f2 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -158,6 +158,13 @@ def __init__(self, *args, **kwargs): # Drained into an explicit FinishReason.ERROR output on the next # update_from_output so the client learns why the session ended. self._streaming_context_overflow: dict[str, tuple[int, str]] = {} + # Opt-in (stage ``additional_config.talker_kstep_waiting_safe``): keep the + # Talker K-step drafts while a prefill waits but cannot be admitted + # because the running batch is full, instead of a single-frame step. + additional_config = getattr(getattr(self, "vllm_config", None), "additional_config", None) + self._talker_kstep_waiting_safe = isinstance(additional_config, dict) and ( + additional_config.get("talker_kstep_waiting_safe") is True + ) def _get_confirmed_num_computed_tokens(self, request: Request) -> int: """num_computed_tokens minus async placeholders (KV actually on GPU).""" @@ -419,7 +426,8 @@ def _drop_talker_drafts_if_prefill_pending(self) -> None: max_len = int(max_len) if max_len is not None else None except (TypeError, ValueError): max_len = None - prefill_pending = bool(self.waiting) + waiting = bool(self.waiting) + prefill_pending = False widths: set[int] = set() for req in self.running: computed = int(req.num_computed_tokens) @@ -438,12 +446,25 @@ def _drop_talker_drafts_if_prefill_pending(self) -> None: else: widths.add(delta) uneven = len(widths) > 1 - if not prefill_pending and not uneven: + if not prefill_pending and not uneven and (not waiting or self._talker_kstep_waiting_is_blocked()): return for req in self.running: if req.spec_token_ids: req.spec_token_ids = [] + def _talker_kstep_waiting_is_blocked(self) -> bool: + """Whether no waiting request can be admitted without a free running slot. + + vLLM only moves a waiting request into a free running slot, so with a + full batch the waiting prefill cannot join this step and the uniform + decode spans stay uniform. The full-batch test is deliberately stronger + than a KV allocation probe, which is stateful. + """ + if getattr(self, "_talker_kstep_waiting_safe", False) is not True: + return False + capacity = getattr(self, "max_num_running_reqs", None) + return isinstance(capacity, int) and capacity > 0 and len(self.running) >= capacity + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Remove FINISHED_ABORTED requests before the upstream scheduler sees # them. Upstream vllm raises RuntimeError on this status; omni allows @@ -1271,7 +1292,15 @@ def _prepare_minicpmo45_stage0_window( if not isinstance(duplex, dict) or duplex.get("data_plane") is not True: return False runtime_config = duplex.get("runtime_config") - runtime_config = runtime_config if isinstance(runtime_config, dict) else {} + if isinstance(runtime_config, dict): + # A session-static prompt carries this snapshot only on its first + # append (or after a config generation change); keep a copy for the + # window planner, which runs on every unit boundary. + session._minicpmo45_runtime_config = runtime_config + else: + runtime_config = getattr(session, "_minicpmo45_runtime_config", {}) + if not isinstance(runtime_config, dict): + runtime_config = {} window = runtime_config.get("duplex_window_config") if not isinstance(window, dict): return False diff --git a/vllm_omni/model_executor/duplex_sampling.py b/vllm_omni/model_executor/duplex_sampling.py index dd48431f6d8..96d622779f3 100644 --- a/vllm_omni/model_executor/duplex_sampling.py +++ b/vllm_omni/model_executor/duplex_sampling.py @@ -22,6 +22,20 @@ class DuplexSamplingRow: seq: int | None payload: dict[str, object] | None max_tokens: int | None + # Host copies of the row's sampling parameters (None when the runner has none). + temperature: float | None = None + top_k: int | None = None + top_p: float | None = None + + +def _host_value(values: Any, row_idx: int, cast: type) -> Any: + """``values[row_idx]`` of a host array as ``cast``, or None when there is none.""" + if values is None: + return None + try: + return cast(values[row_idx]) + except (IndexError, TypeError, ValueError): + return None class DuplexSamplingHelper: @@ -60,8 +74,12 @@ def update_states(self, runner: object, scheduler_output: object) -> None: def rows(self, runner: object) -> tuple[DuplexSamplingRow, ...]: rows: list[DuplexSamplingRow] = [] - req_ids = [str(req_id) for req_id in getattr(runner.input_batch, "req_ids", [])] + input_batch = runner.input_batch + req_ids = [str(req_id) for req_id in getattr(input_batch, "req_ids", [])] requests = getattr(runner, "requests", {}) + temperature_cpu = getattr(input_batch, "temperature_cpu", None) + top_k_cpu = getattr(input_batch, "top_k_cpu", None) + top_p_cpu = getattr(input_batch, "top_p_cpu", None) for row_idx, req_id in enumerate(req_ids): if req_id not in self.active_request_ids: continue @@ -93,6 +111,9 @@ def rows(self, runner: object) -> tuple[DuplexSamplingRow, ...]: seq=seq, payload=payload, max_tokens=max_tokens if max_tokens > 0 else None, + temperature=_host_value(temperature_cpu, row_idx, float), + top_k=_host_value(top_k_cpu, row_idx, int), + top_p=_host_value(top_p_cpu, row_idx, float), ) ) return tuple(rows) @@ -142,8 +163,9 @@ def _apply_duplex_sampling(self, logits: torch.Tensor, prepared_sampling_metadat return helper = getattr(self, "_duplex_sampling_helper", None) rows = helper.rows(self) if helper is not None and helper.active_request_ids else () - if rows or (helper is not None and helper.hook_active): - prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) + # Publish the row set even when empty, so the model skips its per-row + # prompt-token scan (a host sync per row) in turn-based serving. + prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) if helper is not None: helper.hook_active = bool(rows) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 330355c640a..aebbc8a94da 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -402,6 +402,11 @@ def __init__(self, *args, **kwargs): def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) self._resolve_duplex_sampling_hook(force=True) + multi_frame = getattr(self, "_multi_frame_decode", None) + if multi_frame is not None: + # Before the serving InputBatch is built (e.g. the MiniCPM-o Talker's + # stage config has no vocab_size, so vLLM would report 0). + multi_frame.ensure_codec_vocab(self) def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size @@ -1392,8 +1397,24 @@ def _sample( sampling_metadata=sampling_metadata, ) + multi_frame = getattr(self, "_multi_frame_decode", None) + if multi_frame is not None: + # A multi-frame step sampled its frames in _model_forward; hand over + # that rejection-sampler-shaped output instead of verifying. + frames_output = multi_frame.take_sampler_output(self) + if frames_output is not None: + return frames_output return super()._sample(logits, spec_decode_metadata) + def propose_draft_token_ids(self, scheduler_output, sampled_token_ids, *args, **kwargs): + multi_frame = getattr(self, "_multi_frame_decode", None) + if multi_frame is not None: + # The multi-frame loop's drafts (e.g. the Talker repeats its last codec id). + drafts = multi_frame.propose_drafts(self, sampled_token_ids) + if drafts is not None: + return drafts + return super().propose_draft_token_ids(scheduler_output, sampled_token_ids, *args, **kwargs) + @staticmethod def _resolve_req_hidden_states( hidden_states_cpu: torch.Tensor | None, diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index a05e1270814..1373064d7ce 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -965,6 +965,7 @@ def _dummy_run( return hidden_states, None def profile_run(self) -> None: + self._run_omni_post_load() # Profile with multimodal encoder & encoder cache. if self.supports_mm_inputs: mm_config = self.model_config.multimodal_config diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 5b08a88abce..7718d220d9a 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -95,6 +95,10 @@ def __init__(self, *args, **kwargs): self.omni_prefix_cache = None self._omni_prefix_cache_cfg = None self._sampled_token_ids_cpu_override = None + # Model hooks resolved in load_model: a setup run once at the start of + # profile_run, and a multi-frame decode loop driven per step. + self._omni_post_load = None + self._multi_frame_decode = None self._omni_query_start_loc_model_kwarg = False # Output-payload constant snapshotted once in load_model; the cache # policy counterpart lives on PrefixCacheRunnerMixin. @@ -212,6 +216,29 @@ def load_model(self, *args, **kwargs) -> None: self._prewarm_attention_capture_workspaces() self._report_model_local_kv() self._warn_unexposed_stage_hooks(model) + # Model-owned setup that needs loaded weights in eval mode (e.g. the + # MiniCPM-o duplex streaming audio encoder's CUDA graphs). Deferred to + # profile_run: vLLM's Worker.load_model scopes max_split_size_mb=20 + # around this method, and a CUDA graph pool captured under it cannot + # reuse its blocks across captures (MiniCPM-o S0 audio graphs: ~5 GiB + # of private pool instead of ~0.4 GiB). In profile_run the profiler + # also counts the memory in the stage's budget. + post_load = getattr(model, "omni_post_load", None) + self._omni_post_load = post_load if callable(post_load) else None + # A model that runs several decode frames per step (the MiniCPM-o + # Talker K-step on CUDA) hands the runner its loop; None otherwise. + self._multi_frame_decode = getattr(model, "multi_frame_decode_hook", None) + + def _run_omni_post_load(self) -> None: + """Run the model's ``omni_post_load`` once (see ``load_model``).""" + post_load = getattr(self, "_omni_post_load", None) + self._omni_post_load = None + if post_load is not None: + post_load() + + def profile_run(self) -> None: + self._run_omni_post_load() + super().profile_run() # Read on the model this runner holds. A multi-stage wrapper that builds its # stage module as a child must re-export them, or the runner silently takes @@ -2213,8 +2240,8 @@ def _model_forward( req_ids=self.input_batch.req_ids, ) - try: - model_output = super()._model_forward( + def run_model(): + return super(OmniGPUModelRunner, self)._model_forward( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, @@ -2222,6 +2249,9 @@ def _model_forward( **model_kwargs, **model_kwargs_extra, ) + + try: + model_output = run_model() finally: finish_decode_step = getattr(self.model, "finish_decode_step_forward", None) if callable(finish_decode_step): @@ -2238,6 +2268,17 @@ def _model_forward( model_output = OmniOutput(*model_output) if not isinstance(model_output, (OmniOutput, IntermediateTensors)) and hasattr(self.model, "make_omni_output"): model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra) + multi_frame = getattr(self, "_multi_frame_decode", None) + if multi_frame is not None: + # Model-owned multi-frame step (a no-op for every other step): + # replays the forward once per extra frame and stashes the sampled ids. + model_output = multi_frame.maybe_run( + self, + model_output, + run_model=run_model, + inputs_embeds=inputs_embeds, + model_kwargs_extra=model_kwargs_extra, + ) # Cache model output so later sample_tokens can consume multimodal results. self._omni_last_model_output = model_output return model_output From f602f4bfe8a385fcbc9556da694d8df0945a8a33 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:30 +0800 Subject: [PATCH 77/89] [Perf][MiniCPM-o] Talker K-step multi-frame decode on CUDA Talker K-step multi-frame decode on CUDA Implement CUDA Talker multi-frame decoding using the #7929 hook interface, with native sampling controls, repetition penalty cache, and pinned copies. Signed-off-by: BeatSeat --- .../minicpmo_4_5/test_talker_batching.py | 10 - .../models/minicpmo_4_5/test_talker_kstep.py | 580 +++++++++++ .../test_talker_kstep_feedback.py | 59 +- .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 972 ++++++++++++++++-- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 23 +- .../models/minicpmo_4_5/pipeline.py | 25 +- .../models/minicpmo_4_5/talker_frame_plan.py | 92 ++ .../models/minicpmo_4_5/talker_kstep.py | 670 ++++++++++++ 8 files changed, 2271 insertions(+), 160 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_talker_kstep.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/talker_frame_plan.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/talker_kstep.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py index 10a91e94a7c..c68a62c50fe 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py @@ -109,17 +109,7 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: talker._request_condition_states = {} talker._deferred_cleanup_ids = set() talker._tts_config = ConditionalChatTTSConfig() - # Codec sampling contract added by the multi-frame decode work: the fixture - # runs the deterministic greedy boundary (argmax over head_code) and never - # arms the multi-frame loop. - talker._codec_temperature = 0.0 - talker._codec_top_k = 0 - talker._codec_top_p = 1.0 - talker._codec_repetition_penalty = 1.0 - talker._codec_min_tokens = 0 - talker._codec_max_tokens = 4032 talker._k_step_frames = 0 - talker.supports_multi_frame_decode = False talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) with torch.no_grad(): talker.head_code[0].weight.copy_(torch.eye(8, 2)) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep.py new file mode 100644 index 00000000000..6832bb6b094 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep.py @@ -0,0 +1,580 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA multi-frame Talker decode vs. the single-frame path it replaces. + +A tiny fake Talker -- the real MiniCPM-o Talker class with toy codec +embedding/head tables and a causal recurrent "backbone" -- is decoded twice: +once a frame per step through the model's own preprocess / make_omni_output / +compute_logits / sample and a check_stop emulation, once K frames per step +through ``talker_kstep`` (frame 0 on the single-frame path, frames +1..K-1 replayed over a K-row buffer whose later rows are stale). The codec +streams, scheduler tokens and final Talker state must be identical. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +from vllm.v1.sample.logits_processor import LogitsProcessors +from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor +from vllm.v1.sample.metadata import SamplingMetadata + +from vllm_omni.model_executor.models.minicpmo_4_5 import talker_frame_plan +from vllm_omni.model_executor.models.minicpmo_4_5 import talker_kstep as mf +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +VOCAB, HIDDEN, EOS = 16, 8, 15 +BIG = 10**6 + +# (vLLM max_tokens, vLLM min_tokens, Talker state max_tokens, state min_tokens) +REQUESTS = [ + (BIG, 6, None, None), # EOS is favoured: stops on a natural EOS at output 7 + (BIG, 0, None, 4), # the Talker's own min_tokens holds EOS back until step 4 + (BIG, 8, 4, 10), # chunk limit forces EOS at step 3, past both min_tokens masks + (6, 0, None, 100), # vLLM max_tokens truncates mid-window, never an EOS +] + + +def _talker(seed: int) -> MiniCPMO45OmniTTSForConditionalGeneration: + torch.manual_seed(seed) + talker = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + nn.Module.__init__(talker) + talker.emb_code = nn.ModuleList([nn.Embedding(VOCAB, HIDDEN)]) + head = nn.Linear(HIDDEN, VOCAB, bias=True) + with torch.no_grad(): + head.bias.zero_() + head.bias[EOS] = 4.0 # EOS wins whenever nothing masks it + talker.head_code = nn.ModuleList([head]) + talker._codec_eos_id = EOS + talker._num_audio_tokens = VOCAB + talker._k_step_frames = 0 + talker._request_audio_states = {} + talker._force_eos_rows = talker._mask_eos_rows = talker._pending_force_eos_rows = None + talker._penalty_histories = None + return talker.requires_grad_(False) + + +class _Backbone: + """Causal toy LM: row state = 0.7 * previous row state + input embedding.""" + + def __init__(self, seed: int): + gen = torch.Generator().manual_seed(seed) + self.proj = torch.randn(HIDDEN, HIDDEN, generator=gen) + + def row(self, prev_state: torch.Tensor, embeds: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + state = 0.7 * prev_state + embeds + return state, torch.tanh(state @ self.proj) * 3.0 + + +def _min_tokens_proc(outputs: list[list[int]]) -> MinTokensLogitsProcessor: + proc = MinTokensLogitsProcessor(None, torch.device("cpu"), False) + for index, (_, min_tokens, _, _) in enumerate(REQUESTS): + if min_tokens: + proc.min_toks[index] = (min_tokens, outputs[index], {EOS}, False) + proc.min_toks[-1] = (1, [0], set(), False) # already reached: forces the first rebuild + proc.update_state(None) + return proc + + +def _metadata(outputs, proc, greedy: bool, generators) -> SamplingMetadata: + b = len(outputs) + return SamplingMetadata( + temperature=None if greedy else torch.full((b,), 0.8), + all_greedy=greedy, + all_random=not greedy, + top_p=None if greedy else torch.full((b,), 0.9), + top_k=None if greedy else torch.full((b,), 5, dtype=torch.int32), + generators=generators, + max_num_logprobs=None, + no_penalties=False, + prompt_token_ids=torch.zeros((b, 2), dtype=torch.long), + frequency_penalties=torch.zeros(b), + presence_penalties=torch.zeros(b), + repetition_penalties=torch.full((b,), 1.3), + output_token_ids=[list(o) for o in outputs], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors([proc]), + ) + + +class _Run: + """Request state shared by both drivers: a fresh talker, same weights.""" + + def __init__(self, greedy: bool): + self.talker = _talker(0) + self.backbone = _Backbone(1) + self.rids = [f"r{i}" for i in range(len(REQUESTS))] + gen = torch.Generator().manual_seed(2) + self.prefix = torch.randn(len(REQUESTS), HIDDEN, generator=gen) + self.last = [3, 5, 7, 9] # codec ids the prefill sampled + self.outputs = [[t] for t in self.last] + self.alive = [True] * len(REQUESTS) + self.streams: list[list[int]] = [[] for _ in REQUESTS] + self.flags = [False] * len(REQUESTS) + self.final: list[tuple | None] = [None] * len(REQUESTS) + for rid, (_, _, max_tokens, min_tokens) in zip(self.rids, REQUESTS): + self.talker._request_audio_states[rid] = { + "finished": False, + "step": 0, + "max_tokens": max_tokens, + "min_tokens": min_tokens, + "turn_end_drain": False, + } + self.proc = _min_tokens_proc(self.outputs) + self.greedy = greedy + self.generators = {} if greedy else {i: torch.Generator().manual_seed(100 + i) for i in range(len(REQUESTS))} + + def preprocess(self, span: int): + embeds, infos = [], [] + for rid, last in zip(self.rids, self.last): + state = self.talker._request_audio_states[rid] + _, emb, upd = self.talker.preprocess( + torch.tensor([last] * span), None, request_id=rid, audio_state=state, _omni_is_prefill=False + ) + embeds.append(emb[0]) + infos.append({"request_id": rid, "audio_state": state, "codes": upd["codes"]}) + return torch.stack(embeds), infos + + def record(self, index: int, output, frame_flag: bool): + audio = output.multimodal_outputs["codes"]["audio"][index] + self.streams[index].extend(int(t) for t in audio.reshape(-1).tolist()) + self.flags[index] |= frame_flag + + def accept(self, index: int, tokens: list[int]) -> None: + """check_stop, token by token; the loop must never need truncating.""" + max_tokens = REQUESTS[index][0] + for position, token in enumerate(tokens): + self.outputs[index].append(token) + if token == EOS or len(self.outputs[index]) >= max_tokens: + assert position == len(tokens) - 1, "the loop emitted a frame past check_stop" + self.alive[index] = False + self.final[index] = self.state()[index] + self.last[index] = tokens[-1] + + def state(self): + states = self.talker._request_audio_states + return [ + (states[r]["step"], list(states[r].get("recent_codes", [])), bool(states[r]["finished"])) for r in self.rids + ] + + +def _single_frame(greedy: bool) -> _Run: + run = _Run(greedy) + spans = [(i, i + 1) for i in range(len(REQUESTS))] + for _ in range(64): + if not any(run.alive): + break + run.proc.update_state(None) + embeds, infos = run.preprocess(1) + state, hidden = run.backbone.row(run.prefix, embeds) + output = run.talker._make_omni_output_single_frame(hidden, infos, spans) + logits = run.talker.compute_logits(hidden) + md = _metadata(run.outputs, run.proc, greedy, run.generators) + sampled = run.talker.sample(logits, md).sampled_token_ids[:, 0].tolist() + for i in range(len(REQUESTS)): + if run.alive[i]: + run.record(i, output, bool(output.multimodal_outputs["meta"]["finished"][i])) + run.prefix[i] = state[i] + run.accept(i, [int(sampled[i])]) + assert not any(run.alive) + return run + + +def _multi_frame(greedy: bool, frames: int) -> _Run: + run = _Run(greedy) + num = len(REQUESTS) + spans = [(i * frames, (i + 1) * frames) for i in range(num)] + for _ in range(64): + if not any(run.alive): + break + run.proc.update_state(None) + embeds0, infos = run.preprocess(frames) + buf = torch.randn(frames, num, HIDDEN) # stale rows past the current frame + buf[0] = embeds0 + row_states: list[torch.Tensor] = [] + + def forward(k: int, embeds: torch.Tensor, buf=buf, row_states=row_states) -> torch.Tensor: + buf[k] = embeds + state, hiddens = run.prefix, [] + row_states.clear() + for j in range(frames): + state, hidden = run.backbone.row(state, buf[j]) + row_states.append(state) + hiddens.append(hidden) + return hiddens[k] + + hidden0 = forward(0, embeds0) + output = run.talker._make_omni_output_single_frame(hidden0, infos, spans) + md = _metadata(run.outputs, run.proc, greedy, run.generators) + first = run.talker.sample(run.talker.compute_logits(hidden0), md).sampled_token_ids[:, 0].long() + plan = talker_frame_plan.plan_codec_frames(run.talker, infos, frames) + controls, _ = mf.build_controls( + plan, + frames=frames, + vocab_size=VOCAB, + budgets=[REQUESTS[i][0] - len(run.outputs[i]) for i in range(num)], + stop_ids=[{EOS}] * num, + min_tokens_state=run.proc.min_toks, + penalties=md.repetition_penalties, + device=torch.device("cpu"), + ) + frame_md = mf.frame_sampling_metadata(md) + sampled, emitted = mf.decode_frames( + first, + frames, + controls, + embed=run.talker.emb_code[0], + forward=forward, + codec_logits=run.talker.compute_logits, + sample=lambda logits: mf.sample_frame(logits, frame_md), + ) + ids = torch.where(emitted, sampled, -1).tolist() + tokens = [[t for t in row if t >= 0] for row in ids] + assert all(row[: len(t)] == t for row, t in zip(ids, tokens)), "emitted frames must be a prefix" + forwarded = [t[:-1] for t in tokens] + flags = talker_frame_plan.commit_codec_frames(run.talker, plan, forwarded) + for i in range(num): + if run.alive[i]: + run.record(i, output, bool(output.multimodal_outputs["meta"]["finished"][i]) or flags[i]) + run.streams[i].extend(forwarded[i]) + run.prefix[i] = row_states[len(tokens[i]) - 1][i] + run.accept(i, tokens[i]) + assert not any(run.alive) + return run + + +# --------------------------------------------------------------------------- +# End to end through the runner hook: maybe_run, vLLM's split of the sampled +# ids, the scheduler's accept/rollback and the next step's drafts. +# --------------------------------------------------------------------------- + + +class _Model: + """What maybe_run / propose_drafts read off the stage-1 model wrapper.""" + + supports_multi_frame_decode = True + requires_request_sample_eligibility = True + codec_eos_token_id = EOS + codec_vocab_size = VOCAB + + def __init__(self, talker: MiniCPMO45OmniTTSForConditionalGeneration): + self.talker = talker + + def plan_codec_frames(self, infos, frames): + return talker_frame_plan.plan_codec_frames(self.talker, infos, frames) + + def commit_codec_frames(self, plan, forwarded): + return talker_frame_plan.commit_codec_frames(self.talker, plan, forwarded) + + def compute_logits(self, hidden): + return self.talker.compute_logits(hidden) + + def embed_input_ids(self, ids): + return self.talker.emb_code[0](ids) + + +class _Runner: + """The GPUARModelRunner surface the hook touches. + + ``batch_vocab`` is ``input_batch.vocab_size``. The real stage 1 reports 0: + MiniCPMTTSConfig has no ``vocab_size``, so vLLM's ``get_vocab_size()`` + falls back to 0 for it. + """ + + def __init__(self, run: _Run, frames: int, batch_vocab: int): + self.model = _Model(run.talker) + self.num_spec_tokens = frames - 1 + self.use_async_scheduling = False + self.max_model_len = BIG + self.input_batch = SimpleNamespace( + num_reqs=len(run.rids), req_ids=list(run.rids), vocab_size=batch_vocab, sampling_metadata=None + ) + self.requests = {rid: SimpleNamespace(sampling_params=_params()) for rid in run.rids} + arch = SimpleNamespace(vocab_size=batch_vocab) + self.model_config = SimpleNamespace(model_arch_config=arch, get_vocab_size=lambda: arch.vocab_size) + + def _sample(self, logits, spec_decode_metadata): + # GPUARModelRunner._sample without drafts: the model's own sampler. + assert spec_decode_metadata is None + return self.model.talker.sample(logits, self.input_batch.sampling_metadata) + + +def _bookkeep(ids: torch.Tensor, vocab_size: int) -> list[list[int]]: + """GPUModelRunner._bookkeeping_sync's split of the step's sampled ids.""" + from vllm.v1.sample.rejection_sampler import RejectionSampler + + if ids.shape[-1] == 1: + return ids.tolist() + return RejectionSampler.parse_output(ids, vocab_size)[0] + + +def _engine(greedy: bool, frames: int, *, batch_vocab: int, force_drafts: bool = False, prepare=None): + """Drive the real maybe_run/propose_drafts the way stage 1 does. + + Each step schedules ``[last] + drafts`` per request; the runner's + preprocess embeds ``last`` at the span's first row (later rows are stale), + make_omni_output does frame 0's bookkeeping and maybe_run the rest. The + ids are split as vLLM splits them, and the scheduler emulation accepts what + comes back: a request's KV is committed up to its last accepted id, and an + empty row rolls the whole span back (OmniARScheduler's empty-row path), so + the next step forwards the same codec id again. + ``force_drafts`` schedules drafts every step, whatever propose_drafts said. + Returns the run and the scheduled span width of every step. + """ + run = _Run(greedy) + runner = _Runner(run, frames, batch_vocab) + if prepare is not None: + prepare(runner) + num = len(REQUESTS) + drafts: list[list[int]] = [[] for _ in range(num)] + stale = torch.Generator().manual_seed(3) + widths: list[int] = [] + for _ in range(64): + if not any(run.alive): + break + run.proc.update_state(None) + width = 1 + len(drafts[0]) + widths.append(width) + embeds0, infos = run.preprocess(width) + starts = [i * width for i in range(num)] + spans = [(s, s + width) for s in starts] + buf = torch.randn(num * width, HIDDEN, generator=stale) + buf[starts] = embeds0 + row_states: list[list[torch.Tensor]] = [[run.prefix[i]] * width for i in range(num)] + + def run_model(buf=buf, width=width, starts=starts, row_states=row_states) -> torch.Tensor: + hidden = torch.empty(num * width, HIDDEN) + for i, start in enumerate(starts): + state = run.prefix[i] + for j in range(width): + state, hidden[start + j] = run.backbone.row(state, buf[start + j]) + row_states[i][j] = state + return hidden + + hidden = run_model() + output = run.talker._make_omni_output_single_frame(hidden, infos, spans) + md = _metadata(run.outputs, run.proc, greedy, run.generators) + runner.input_batch.sampling_metadata = md + extra = { + "request_token_spans": spans, + "model_intermediate_buffer": infos, + "request_max_tokens_remaining": [REQUESTS[i][0] - len(run.outputs[i]) for i in range(num)], + } + output = mf.maybe_run(runner, output, run_model=run_model, inputs_embeds=buf, model_kwargs_extra=extra) + stash = mf.take_sampler_output(runner) + if stash is None: + assert width == 1, "a multi-row Talker step must come back through the stash" + ids = run.talker.sample(run.talker.compute_logits(hidden), md).sampled_token_ids + else: + ids = stash.sampled_token_ids + valid = _bookkeep(ids, runner.input_batch.vocab_size) + for i in range(num): + if not run.alive[i]: + continue + run.record(i, output, bool(output.multimodal_outputs["meta"]["finished"][i])) + if valid[i]: + run.prefix[i] = row_states[i][len(valid[i]) - 1] + run.accept(i, valid[i]) + proposed = mf.propose_drafts(runner, valid) + if force_drafts: + proposed = [[run.last[i]] * (frames - 1) for i in range(num)] + drafts = proposed or [[] for _ in range(num)] + assert not any(run.alive) + return run, widths + + +def _assert_same(got: _Run, ref: _Run) -> None: + assert got.streams == ref.streams, "codec frames sent to Code2Wav differ" + assert got.outputs == ref.outputs + assert got.flags == ref.flags + assert got.final == ref.final + + +@pytest.mark.parametrize("greedy", [True, False], ids=["greedy", "seeded"]) +def test_engaged_through_maybe_run_matches_single_frame(greedy: bool): + ref = _single_frame(greedy) + got, widths = _engine(greedy, 4, batch_vocab=VOCAB) + _assert_same(got, ref) + assert widths[0] == 1 and set(widths[1:]) == {4}, widths + + +@pytest.mark.parametrize("greedy", [True, False], ids=["greedy", "seeded"]) +def test_declined_multi_row_step_is_one_single_frame_step(greedy: bool): + """The real stage 1 (input_batch.vocab_size 0) with drafts scheduled anyway. + + maybe_run must decline, and the step must be exactly one single-frame + step: the same codes, one frame per request to Code2Wav, and the same + Talker state. A -1-padded (B, K) output went through parse_output, which + drops every id >= vocab_size; the scheduler rolled the whole span back and + the next step re-sent the same codec id, so every frame reached Code2Wav + twice. + """ + ref = _single_frame(greedy) + got, widths = _engine(greedy, 4, batch_vocab=0, force_drafts=True) + _assert_same(got, ref) + assert set(widths[1:]) == {4}, widths + + +def test_no_drafts_when_the_loop_would_decline(): + """With vocab_size 0 the next step is not drafted in the first place.""" + ref = _single_frame(greedy=True) + got, widths = _engine(True, 4, batch_vocab=0) + _assert_same(got, ref) + assert set(widths) == {1}, widths + + +@pytest.mark.parametrize("greedy", [True, False], ids=["greedy", "seeded"]) +def test_codec_vocab_lets_a_vocab0_stage_engage(greedy: bool): + """ensure_codec_vocab on a stage whose config reports vocab_size 0.""" + ref = _single_frame(greedy) + got, widths = _engine(greedy, 4, batch_vocab=0, prepare=mf.ensure_codec_vocab) + _assert_same(got, ref) + assert set(widths[1:]) == {4}, widths + + +@pytest.mark.parametrize("frames", [2, 3, 4, 8]) +@pytest.mark.parametrize("greedy", [True, False], ids=["greedy", "seeded"]) +def test_multi_frame_matches_single_frame(greedy: bool, frames: int): + ref = _single_frame(greedy) + _assert_same(_multi_frame(greedy, frames), ref) + + +def _runner(req_params: list[SimpleNamespace], num_spec: int = 3, batch_vocab: int = VOCAB): + md = SimpleNamespace( + max_num_logprobs=None, + logprob_token_ids=None, + bad_words_token_ids={}, + allowed_token_ids_mask=None, + thinking_budget_state_holder=None, + logitsprocs=LogitsProcessors(), + ) + req_ids = [f"r{i}" for i in range(len(req_params))] + return SimpleNamespace( + model=SimpleNamespace( + supports_multi_frame_decode=True, + requires_request_sample_eligibility=True, + plan_codec_frames=lambda *a: None, + codec_eos_token_id=EOS, + codec_vocab_size=VOCAB, + ), + num_spec_tokens=num_spec, + use_async_scheduling=False, + input_batch=SimpleNamespace( + num_reqs=len(req_ids), req_ids=req_ids, sampling_metadata=md, vocab_size=batch_vocab + ), + requests={r: SimpleNamespace(sampling_params=p) for r, p in zip(req_ids, req_params)}, + ) + + +def _params(**overrides): + base = dict( + frequency_penalty=0.0, + presence_penalty=0.0, + structured_outputs=None, + repetition_detection=None, + stop_token_ids=[EOS], + eos_token_id=None, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +def test_drafts_repeat_last_codec_for_the_whole_batch_or_none(): + runner = _runner([_params(), _params()]) + assert mf.propose_drafts(runner, [[4], [7, 9]]) == [[4, 4, 4], [9, 9, 9]] + # One request sampled nothing: nobody drafts, the next step is one frame. + assert mf.propose_drafts(runner, [[4], []]) == [[], []] + # A request whose stop set lacks the codec EOS cannot take the loop. + assert mf.propose_drafts(_runner([_params(), _params(stop_token_ids=[1])]), [[4], [7]]) == [[], []] + assert mf.propose_drafts(_runner([_params(frequency_penalty=0.5)]), [[4]]) == [[]] + # Other stages (and K=1) fall through to vLLM's drafter. + assert mf.propose_drafts(_runner([_params()], num_spec=0), [[4]]) is None + other = _runner([_params()]) + other.model = SimpleNamespace(supports_multi_frame_decode=False) + assert mf.propose_drafts(other, [[4]]) is None + + +def test_no_drafts_for_a_step_maybe_run_would_decline(monkeypatch): + """Everything maybe_run declines on that is known at draft time stops the + drafts, and each distinct reason is logged once, with its numbers.""" + log = SimpleNamespace(calls=[]) + monkeypatch.setattr(mf, "logger", SimpleNamespace(info=lambda *a: log.calls.append(a[0] % a[1:]))) + monkeypatch.setattr(mf, "_LOGGED", set()) + # The real stage 1 without ensure_codec_vocab: input_batch.vocab_size 0. + narrow = _runner([_params(), _params()], batch_vocab=0) + assert mf.propose_drafts(narrow, [[4], [7]]) == [[], []] + assert mf.propose_drafts(narrow, [[5], [8]]) == [[], []] + assert log.calls == [ + "[minicpmo] CUDA multi-frame Talker decode: no drafts: input batch vocab is narrower than the codec " + f"head (input batch vocab 0, codec head {VOCAB})" + ] + asynchronous = _runner([_params()]) + asynchronous.use_async_scheduling = True + assert mf.propose_drafts(asynchronous, [[4]]) == [[]] + no_budget = _runner([_params()]) + no_budget.model.requires_request_sample_eligibility = False + assert mf.propose_drafts(no_budget, [[4]]) == [[]] + penalty = _runner([_params(), _params(presence_penalty=0.5)]) + assert mf.propose_drafts(penalty, [[4], [7]]) == [[], []] + assert len(log.calls) == 4 + assert "(stage 1 needs async_scheduling: false)" in log.calls[1] + assert "requires_request_sample_eligibility" in log.calls[2] + assert "request r1: frequency 0.0, presence 0.5" in log.calls[3] + + +def test_ensure_codec_vocab_raises_only_a_narrow_talker_vocab(): + runner = _runner([_params()], batch_vocab=0) + arch = SimpleNamespace(vocab_size=0) + runner.model_config = SimpleNamespace(model_arch_config=arch, get_vocab_size=lambda: arch.vocab_size) + mf.ensure_codec_vocab(runner) + # model_config too: initialize_kv_cache rebuilds the InputBatch from it. + assert (runner.input_batch.vocab_size, arch.vocab_size) == (VOCAB, VOCAB) + wide = _runner([_params()], batch_vocab=32000) + wide_arch = SimpleNamespace(vocab_size=32000) + wide.model_config = SimpleNamespace(model_arch_config=wide_arch, get_vocab_size=lambda: wide_arch.vocab_size) + mf.ensure_codec_vocab(wide) + assert (wide.input_batch.vocab_size, wide_arch.vocab_size) == (32000, 32000) + other = _runner([_params()], batch_vocab=0) + other.model = SimpleNamespace(supports_multi_frame_decode=False, codec_vocab_size=VOCAB) + mf.ensure_codec_vocab(other) + assert other.input_batch.vocab_size == 0 + + +def test_frame_budgets_bound_a_request_without_max_tokens_by_the_context(): + runner = _runner([_params(), _params()]) + runner.max_model_len = 100 + runner.requests["r1"].num_tokens = 97 + assert mf.frame_budgets(runner, {"request_max_tokens_remaining": [5, None]}, 2) == ([5, 3], None) + budgets, decline = mf.frame_budgets(runner, {"request_max_tokens_remaining": [5]}, 2) + assert budgets == [] and str(decline) == "no per-request token budget (1 budgets for 2 requests)" + + +@pytest.mark.parametrize("platform", ["cuda", "npu"]) +def test_kstep_overlay_arms_cuda_stage1_only(platform: str): + """The opt-in overlay arms K=8 on CUDA stage 1 and keeps the codec EOS as + the only stop id; the base config (and NPU through the overlay) is as before.""" + from pathlib import Path + + from tests.helpers.stage_config import get_deploy_config_path + from vllm_omni.config.pipeline_registry import resolve_pipeline_config + from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config, merge_pipeline_deploy + + def stage1(name: str): + deploy = _apply_platform_overrides(load_deploy_config(Path(get_deploy_config_path(name))), platform=platform) + return merge_pipeline_deploy(resolve_pipeline_config("minicpmo_4_5"), deploy)[1] + + overlay, base = stage1("minicpmo_4_5_kstep.yaml"), stage1("minicpmo_4_5.yaml") + if platform == "cuda": + assert "speculative_config" not in base.yaml_engine_args + assert overlay.yaml_engine_args["speculative_config"]["num_speculative_tokens"] == 7 + assert overlay.yaml_engine_args["async_scheduling"] is False + assert overlay.yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] + else: + assert overlay.yaml_engine_args == base.yaml_engine_args + assert overlay.yaml_extras["default_sampling_params"] == base.yaml_extras["default_sampling_params"] diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py index 0bf566a1139..d83ee39deb3 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_feedback.py @@ -69,6 +69,8 @@ def _make_talker(*, k_step_frames: int, scripted_samples: list[int]): model._request_audio_states = {} model._request_codec_history = {} model._request_generators = {} + model._request_codec_device_states = {} + model._request_codec_device_inputs = {} emb = nn.Embedding(_NUM_AUDIO_TOKENS, 4) with torch.no_grad(): emb.weight.zero_() @@ -190,6 +192,53 @@ def test_guard_drops_drafts_when_waiting_request_pending(): assert decode_req.spec_token_ids == [] +def test_guard_keeps_uniform_drafts_when_waiting_is_at_capacity(): + """A queued prefill cannot enter a full running batch, so K-step is safe.""" + decode_req = _req(computed=100, prompt=100, spec=[0] * 7, total=101) + sched = _make_scheduler(num_spec=7, waiting=[object()], running=[decode_req]) + sched.max_num_running_reqs = 1 + sched._talker_kstep_waiting_safe = True + + sched._drop_talker_drafts_if_prefill_pending() + assert decode_req.spec_token_ids == [0] * 7 + + +def test_guard_keeps_default_downgrade_when_capacity_is_available(): + """The opt-in fast path never delays an actually admissible prefill.""" + decode_req = _req(computed=100, prompt=100, spec=[0] * 7, total=101) + sched = _make_scheduler(num_spec=7, waiting=[object()], running=[decode_req]) + sched.max_num_running_reqs = 2 + sched._talker_kstep_waiting_safe = True + + sched._drop_talker_drafts_if_prefill_pending() + assert decode_req.spec_token_ids == [] + + +def test_kstep_diagnostics_explain_prefill_width_and_max_len(monkeypatch): + from vllm_omni.core.sched import omni_ar_scheduler as module + + decode_req = _req(computed=100, prompt=100, spec=[0] * 7, total=101) + prefill_req = _req(computed=96, prompt=100, spec=[], total=100) + near_max = _req(computed=4090, prompt=4090, spec=[0] * 7, total=4091) + sched = _make_scheduler(num_spec=7, waiting=[object()], running=[decode_req, prefill_req, near_max]) + sched._talker_kstep_stats = True + sched.max_model_len = 4096 + messages: list[tuple[object, ...]] = [] + monkeypatch.setattr(module.logger, "info", lambda *args: messages.append(args)) + + sched._drop_talker_drafts_if_prefill_pending() + + assert decode_req.spec_token_ids == [] + downgrade = [args for args in messages if "downgrade" in str(args[0])] + assert downgrade + detail = str(downgrade[-1]) + assert "waiting_nonempty" in detail + assert "mid_prefill" in detail + assert "rows_width_mismatch" in detail + assert "near_max_model_len" in detail + assert "num_computed_tokens" in detail + + def test_guard_keeps_drafts_when_no_prefill_pending(): decode_req = _req(computed=100, prompt=100, spec=[0] * 7, total=101) sched = _make_scheduler(num_spec=7, waiting=[], running=[decode_req]) @@ -276,7 +325,8 @@ def test_talker_stop_token_ids_match_the_multi_frame_head(): # The pipeline default cannot depend on the platform: the block that # collapses the head is a deploy-config decision, so the marker is added # there (test_minicpmo_talker_multi_frame_is_npu_scoped pins the merge). - assert mcp_pipeline._talker_stop_token_ids() == [mcp_pipeline._CODEC_EOS_TOKEN_ID] + talker_constraints = mcp_pipeline.MINICPMO_4_5_PIPELINE.get_stage(1).sampling_constraints + assert talker_constraints["stop_token_ids"] == [mcp_pipeline._CODEC_EOS_TOKEN_ID] def test_multiframe_gate_matrix(): @@ -750,13 +800,6 @@ def test_control_stop_ids_stay_out_of_codec_censor_set(): model._merge_request_codec_params(state, SimpleNamespace(stop_token_ids=[1, 6561], min_tokens=50)) assert state["codec_stop_token_ids"] == [6561] - # The drop is gated on the multi-frame arm: a non-armed deployment keeps - # the legacy verbatim copy (the control-row contract does not apply). - model_legacy = _make_talker(k_step_frames=1, scripted_samples=[1]) - state_legacy = {"step": 0} - model_legacy._merge_request_codec_params(state_legacy, SimpleNamespace(stop_token_ids=[1, 6561], min_tokens=50)) - assert state_legacy["codec_stop_token_ids"] == [1, 6561] - # End to end across the sampling boundary: a codec 1 sampled past the # floor is emitted as ordinary audio and does not finish the request; # a later codec EOS still terminates it normally. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index 439b46554d7..6b8b806f071 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -19,6 +19,7 @@ from collections.abc import Iterable from contextlib import suppress +from dataclasses import dataclass from functools import cached_property from typing import Any @@ -45,10 +46,48 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights from vllm_omni.platforms import current_omni_platform +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking logger = init_logger(__name__) +@dataclass(slots=True) +class _MiniCPMO45PendingSamples: + """A deferred Stage-0 duplex step awaiting its host commit. + + ``host`` receives (final token, raw stage-2 sample, boundary hit) per + row once ``event`` completes; ``rewinds`` maps a row to the generator + offset before its stage-2 draw, restored if the boundary resolved it. + The row maps are the step's own, since the next step replaces them. + """ + + host: torch.Tensor + event: Any + row_idxs: list[int] + stage2: list[bool] + rewinds: dict[int, tuple[torch.Generator, Any]] + token_ids: dict[str, int] + row_sessions: dict[int, str] | None + row_payloads: dict[int, dict[str, Any]] | None + + +def _generator_mark(generator: torch.Generator) -> Any: + """Where ``generator`` stands: its Philox offset (a host value), else its full state.""" + try: + return ("offset", generator.get_offset()) + except RuntimeError: # not a Philox generator (CPU) + return ("state", generator.get_state()) + + +def _generator_rewind(generator: torch.Generator, mark: Any) -> None: + """Put ``generator`` back where ``_generator_mark`` found it, undoing later draws.""" + kind, value = mark + if kind == "offset": + generator.set_offset(value) + else: + generator.set_state(value) + + @MULTIMODAL_REGISTRY.register_processor( MiniCPMO45OmniLLMMultiModalProcessor, info=MiniCPMO45OmniLLMProcessingInfo, @@ -184,6 +223,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): getattr(self.talker, "mrv2_decode_preprocess_is_identity", False) ) self.logits_vocab_size = int(self.talker.logits_vocab_size) + # Opt-in (hf_overrides talker_kstep_graph_sampling): the CUDA K-step + # loop replays captured per-frame sampling tails (talker_kstep_graph). + if getattr(config, "talker_kstep_graph_sampling", False) is True: + from vllm_omni.model_executor.models.minicpmo_4_5.talker_kstep_graph import ( + TalkerKStepFrameGraphs, + ) + + self._kstep_frame_graphs = TalkerKStepFrameGraphs() + logger.info("[minicpmo] Talker K-step frame graphs on (talker_kstep_graph_sampling)") else: raise ValueError(f"Invalid model stage: {self.model_stage}. Must be one of: 'llm', 'tts'") @@ -244,7 +292,14 @@ def prepare_duplex_sampling( ) -> None: """Apply MiniCPM duplex policy before the standard model sampler.""" del sampling_metadata + # The previous step's deferred decisions update the session state read below. + self._commit_minicpmo45_duplex_pending_samples() self._minicpmo45_active_duplex_rows = [row.row_idx for row in rows] + self._minicpmo45_duplex_row_sampling_host = { + row.row_idx: (row.temperature, row.top_k, row.top_p) + for row in rows + if row.temperature is not None and row.top_k is not None and row.top_p is not None + } self._minicpmo45_duplex_row_sessions = { row.row_idx: row.session_id for row in rows if row.session_id is not None } @@ -306,8 +361,10 @@ def prepare_duplex_sampling( if force_listen and segment_key in force_listen_segments: continue if force_listen: - logits[row_idx, :] = float("-inf") - logits[row_idx, listen_id] = 0.0 + # fill_ keeps the scalars on the host; assigning one element + # copies a host scalar tensor and waits for the device. + logits[row_idx].fill_(float("-inf")) + logits[row_idx].narrow(0, listen_id, 1).fill_(0.0) force_listen_segments.add(segment_key) # -------------------- Device utilities -------------------- @@ -418,53 +475,20 @@ def preprocess( embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": "bad_duplex_payload"}} - session_key = session_id - state = helper.sessions.get(session_key) - if state is None: - from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( - _MiniCPMO45Stage0SessionState, - ) - - state = _MiniCPMO45Stage0SessionState(session_id=session_id) - helper.sessions[session_key] = state - session_config = duplex.get("session_config") - session_config = dict(session_config) if isinstance(session_config, dict) else {} - runtime_config = duplex.get("runtime_config") - runtime_config = dict(runtime_config) if isinstance(runtime_config, dict) else {} - if hasattr(helper.thinker, "audio_past_key_values"): - helper.thinker.audio_past_key_values = None - helper._configure_streaming_processor(state) - helper._prepare_session_context(state, session_config, runtime_config=runtime_config) - - audio_waveform = helper._decode_audio_payload(payload) - try: - video_frames = helper._decode_video_frames_payload(payload) - except ValueError as exc: - embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) - return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": str(exc)}} - seq = duplex.get("seq") - try: - seq = int(seq) if seq is not None else None - except (TypeError, ValueError): - seq = None - epoch = duplex.get("epoch") - try: - epoch = int(epoch) if epoch is not None else None - except (TypeError, ValueError): - epoch = None - result = helper._stage_prefill_embeddings_only( - state, - audio_waveform, - video_frames=video_frames, - epoch=epoch, - seq=seq, - is_speech=bool(payload.get("is_speech", False)), - final=bool(duplex.get("final")), - stage0_window=(duplex.get("stage0_window") if isinstance(duplex.get("stage0_window"), dict) else None), - stage0_reanchor=( - duplex.get("stage0_reanchor") if isinstance(duplex.get("stage0_reanchor"), dict) else None - ), - ) + # A deferred sample of this session updates the latches its append reads. + self._commit_minicpmo45_duplex_pending_samples(session_ids={session_id}) + state = self._minicpmo45_duplex_session_state(helper, session_id, duplex) + prefill_kwargs = self._minicpmo45_duplex_prefill_kwargs(duplex, payload) + seq = prefill_kwargs["seq"] + result = helper.take_staged_prefill(state, prefill_kwargs["epoch"], seq) + if result is None: # not built by this step's preprocess_batch + audio_waveform = helper._decode_audio_payload(payload) + try: + prefill_kwargs.update(self._minicpmo45_duplex_append_frames(helper, duplex, payload)) + except ValueError as exc: + embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) + return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": str(exc)}} + result = helper._stage_prefill_embeddings_only(state, audio_waveform, **prefill_kwargs) update_result = dict(result) if result.get("stage0_window_replaced") is True: window = duplex.get("stage0_window", {}) @@ -542,12 +566,156 @@ def preprocess( if len(input_token_ids) < span_len: input_token_ids.extend([pad_token_id] * (span_len - len(input_token_ids))) if input_token_ids: - req_input_ids = torch.tensor(input_token_ids, dtype=input_ids.dtype, device=input_ids.device) + # Pinned + non-blocking: a pageable copy would wait for all queued GPU work. + req_input_ids = index_to_device(input_token_ids, input_ids.device, input_ids.dtype) update_result["duplex_prompt_token_ids"] = full_input_token_ids else: req_input_ids = torch.full_like(input_ids, helper._required_token_id("unit_token_id")) return req_input_ids, req_embeds, {"duplex": update_result} + def preprocess_batch( + self, + *, + req_ids: list[str], + model_intermediate_buffer: dict[str, dict[str, Any]], + device: torch.device, + ) -> None: + """Runner hook: build this step's new duplex appends across sessions at once. + + Called before the per-request ``preprocess`` loop; both batches only + pay off with two or more appends. The camera frames of every append go + through one vision-tower call (``prefetch_vision``), then every append + not built yet is prefilled with its streaming-encoder units batched + across sessions (``stage_prefill_batch``). ``preprocess`` takes the + staged results; anything skipped here is built or reported there. + """ + del device + if self.model_stage != "llm": + return + appends: list[dict[str, Any]] = [] + for req_id in req_ids: + info = model_intermediate_buffer.get(req_id) + duplex = info.get("duplex") if isinstance(info, dict) else None + if isinstance(duplex, dict) and duplex.get("data_plane") is True: + appends.append(duplex) + frame_appends = [ + duplex + for duplex in appends + if isinstance(duplex.get("payload"), dict) and duplex["payload"].get("video_frames") + ] + helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) + if len(frame_appends) < 2: + if helper is not None: + helper.prefetch_vision([]) # drop entries the previous step did not consume + else: + try: + helper = self._duplex_data_plane_helper() + helper.prefetch_vision(frame_appends) + except Exception: # noqa: BLE001 - preprocess encodes each append itself instead + logger.warning("MiniCPM-o duplex frame prefetch failed; encoding per request", exc_info=True) + if helper is None or len(appends) < 2 or not helper.batches_audio_encoder(): + return + candidates: list[tuple[str, dict[str, Any], dict[str, Any], dict[str, Any]]] = [] + seen: set[str] = set() + for duplex in appends: + session_id = str(duplex.get("session_id") or "") + payload = duplex.get("payload") + if not session_id or not isinstance(payload, dict) or session_id in seen: + # A second append of one session keeps its order: preprocess. + seen.add(session_id) + continue + seen.add(session_id) + prefill_kwargs = self._minicpmo45_duplex_prefill_kwargs(duplex, payload) + seq = prefill_kwargs["seq"] + if seq is not None and helper.needs_prefill(helper.sessions.get(session_id), prefill_kwargs["epoch"], seq): + candidates.append((session_id, duplex, payload, prefill_kwargs)) + if len(candidates) < 2: + return + self._commit_minicpmo45_duplex_pending_samples(session_ids={candidate[0] for candidate in candidates}) + batch = [] + for session_id, duplex, payload, prefill_kwargs in candidates: + try: + audio_waveform = helper._decode_audio_payload(payload) + prefill_kwargs.update(self._minicpmo45_duplex_append_frames(helper, duplex, payload)) + except ValueError: + continue + state = self._minicpmo45_duplex_session_state(helper, session_id, duplex) + batch.append((state, audio_waveform, prefill_kwargs)) + if len(batch) >= 2: + helper.stage_prefill_batch(batch) + + @staticmethod + def _minicpmo45_duplex_append_frames(helper, duplex: dict[str, Any], payload: dict[str, Any]) -> dict[str, Any]: + """The append's prefetched frame embeddings, else its decoded frames (may raise ValueError).""" + encoded_frames = helper.take_prefetched_vision(duplex) + if encoded_frames is not None: + return {"encoded_frames": encoded_frames} + return {"video_frames": helper._decode_video_frames_payload(payload)} + + def _minicpmo45_duplex_session_state(self, helper, session_id: str, duplex: dict[str, Any]): + """The Stage-0 state of ``session_id``, created with its session context on first use.""" + state = helper.sessions.get(session_id) + if state is None: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + state = _MiniCPMO45Stage0SessionState(session_id=session_id) + helper.sessions[session_id] = state + session_config = duplex.get("session_config") + session_config = dict(session_config) if isinstance(session_config, dict) else {} + runtime_config = duplex.get("runtime_config") + runtime_config = dict(runtime_config) if isinstance(runtime_config, dict) else {} + if hasattr(helper.thinker, "audio_past_key_values"): + helper.thinker.audio_past_key_values = None + helper._configure_streaming_processor(state) + helper._prepare_session_context(state, session_config, runtime_config=runtime_config) + return state + + @staticmethod + def _minicpmo45_duplex_prefill_kwargs(duplex: dict[str, Any], payload: dict[str, Any]) -> dict[str, Any]: + """Keyword arguments of ``_stage_prefill_embeddings_only`` for one append, frames aside.""" + seq = duplex.get("seq") + try: + seq = int(seq) if seq is not None else None + except (TypeError, ValueError): + seq = None + epoch = duplex.get("epoch") + try: + epoch = int(epoch) if epoch is not None else None + except (TypeError, ValueError): + epoch = None + return { + "epoch": epoch, + "seq": seq, + "is_speech": bool(payload.get("is_speech", False)), + "final": bool(duplex.get("final")), + "stage0_window": duplex.get("stage0_window") if isinstance(duplex.get("stage0_window"), dict) else None, + "stage0_reanchor": ( + duplex.get("stage0_reanchor") if isinstance(duplex.get("stage0_reanchor"), dict) else None + ), + } + + def omni_post_load(self) -> None: + """Runner hook once the weights are loaded and the model is in eval mode. + + The runner calls it once, at the start of its profile run (outside the + worker's load-time allocator scope). Captures the duplex Stage-0 streaming audio encoder's CUDA graphs on the + thinker's device (``duplex_audio_encoder_cuda_graph``, default on). The + runtime that owns them is built in ``__init__``, where the encoder is not + loaded yet. + """ + helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) + build = getattr(helper, "build_audio_cuda_graph", None) + if not callable(build): + return + device = self._module_device(self.thinker if self.thinker is not None else self) + if device.type == "cuda": + with torch.cuda.device(device): + build() + else: + build() + def _duplex_data_plane_helper(self): helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) if helper is not None: @@ -694,13 +862,12 @@ def forward( if isinstance(key, str) and isinstance(value, int) and value >= 0 } if special_keys: + # Host tensors: every consumer reads these back on the host, and a + # pageable host->device copy per row and key would wait for the + # whole forward, serializing the step. multimodal_outputs["meta"] = { key: [ - torch.tensor( - [int(value)], - dtype=torch.long, - device=text_hidden_states.device, - ) + torch.tensor([int(value)], dtype=torch.long) if isinstance(value, int) and value >= 0 else None for duplex_info in duplex_rows @@ -760,47 +927,26 @@ def make_omni_output(self, model_outputs, **kwargs): @property def requires_request_sample_eligibility(self) -> bool: - """Forward the Talker's sampling-eligibility contract to the runner. - - The runner only supplies ``request_sample_eligible`` when the model it - sees declares the flag, and it always sees this wrapper -- never the - inner Talker (the same resolution rule as talker_multiframe.applies - below). Without the forward, the Talker's K-step branch falls back to - treating every request as eligible, so an incomplete prefill chunk - advances codec history and RNG state and the generated audio starts - depending on prefill chunking. + """Forward the Talker's flag: the runner only sees this wrapper. + + Without ``request_sample_eligible`` an incomplete prefill chunk would + advance the Talker's codec history and RNG state. """ talker = getattr(self, "talker", None) return talker is not None and bool(getattr(talker, "requires_request_sample_eligibility", False)) - # Runner-side forwards for multi-frame decode. The baseline declares these - # four methods on this outer wrapper class; the runner only ever sees the - # wrapper (talker_multiframe.applies / is_multi_token_decode / - # ensure_stop_token_vocab all resolve through getattr(model, ...)), so - # without them: - # * missing supports_multi_frame_decode -> the check silently reads False, - # the multi-frame loop never runs and no reason is logged; - # * missing take/set_batch_stop_logits -> per-frame stop rows unavailable; - # * missing _batch_stop_logits -> the runner's gate probe reads None and - # always falls back to text_hidden_states; - # * missing merge_frame_outputs -> multi-frame results cannot be merged. - # The speculative config alone is not enough: the runner-side loop and its - # vocab-width correction (ensure_stop_token_vocab) stay dormant, the - # rejection sampler consumes width-0 rows, AIV goes out of bounds (507035) - # and the stage-1 engine dies with empty audio. + # Multi-frame Talker decode: the runners resolve these through getattr on + # this wrapper, never the inner Talker. Without them the multi-frame loop + # silently stays off, and on NPU the rejection sampler consumes width-0 + # rows and stage 1 dies with empty audio. @property def supports_multi_frame_decode(self) -> bool: - # "This wrapper implements the multi-frame path", not "the loop is - # running": the Talker (tts) stage has it, the Thinker does not. Whether - # a step actually runs it is decided per deployment by stage 1's - # speculative_config (the runner checks its num_spec_tokens for that). + # Implemented by the Talker stage; stage 1's speculative_config decides + # whether a step runs it. return self.model_stage == "tts" @property def _batch_stop_logits(self): - # The runner probes this attribute on the registered architecture - # (getattr in npu_model_runner); without the forward it reads None and - # the gate always falls back to text_hidden_states. if self.model_stage != "tts": return None return self.talker._batch_stop_logits @@ -820,6 +966,47 @@ def merge_frame_outputs(self, frame_outputs, frame_stop_logits): return frame_outputs return self.talker.merge_frame_outputs(frame_outputs, frame_stop_logits) + # CUDA multi-frame decode (talker_kstep.py, talker_frame_plan.py). + @property + def multi_frame_decode_hook(self) -> Any: + """The K-step loop the GPU runners drive for the Talker; None for the + other stages and on NPU (whose runner owns its own K-step).""" + if self.model_stage != "tts" or current_omni_platform.is_npu(): + return None + from vllm_omni.model_executor.models.minicpmo_4_5 import talker_kstep + + return talker_kstep + + @property + def codec_eos_token_id(self) -> int: + return int(self.talker._codec_eos_id) + + @property + def codec_vocab_size(self) -> int | None: + # Stage 1's config has no vocab_size; the CUDA runner reads the codec + # head width from here (talker_kstep.ensure_codec_vocab). + if self.model_stage != "tts" or self.talker is None: + return None + width = getattr(self.talker, "_num_audio_tokens", None) + return int(width) if width is not None else None + + def plan_codec_frames(self, infos, frames: int): + from vllm_omni.model_executor.models.minicpmo_4_5.talker_frame_plan import plan_codec_frames + + return plan_codec_frames(self.talker, infos, frames) + + def commit_codec_frames(self, plan, forwarded: list[list[int]]) -> list[bool]: + from vllm_omni.model_executor.models.minicpmo_4_5.talker_frame_plan import commit_codec_frames + + return commit_codec_frames(self.talker, plan, forwarded) + + def kstep_decode_frames(self, **kwargs): + """``gpu_talker_multiframe.maybe_run``'s frames-1..K-1 hook: ``(sampled, emitted)`` + from captured per-frame tails, or None for the eager loop (always, unless + ``talker_kstep_graph_sampling`` is on).""" + graphs = getattr(self, "_kstep_frame_graphs", None) + return None if graphs is None else graphs(**kwargs) + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None: # Handle OmniOutput type if isinstance(hidden_states, OmniOutput): @@ -870,6 +1057,7 @@ def _sample_minicpmo45_native_duplex_stage0( ) -> SamplerOutput | None: if self.model_stage != "llm" or logits.ndim != 2 or logits.shape[0] == 0: return None + self._commit_minicpmo45_duplex_pending_samples() token_ids = self._minicpmo45_native_duplex_token_ids() unit_id = token_ids.get("unit_token_id", -1) if unit_id < 0: @@ -885,8 +1073,11 @@ def _sample_minicpmo45_native_duplex_stage0( chunk_terminators = self._minicpmo45_chunk_terminator_token_ids(token_ids) output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] - sampled_ids: list[int] = [] - for row_idx in range(logits.shape[0]): + num_rows = logits.shape[0] + row_params = self._minicpmo45_duplex_row_params(sampling_metadata, num_rows) + sampled_ids: list[int] = [-1] * num_rows + pending_rows: list[int] = [] + for row_idx in range(num_rows): accepted = output_token_ids[row_idx] if row_idx < len(output_token_ids) else [] last_accepted = next((int(t) for t in reversed(accepted) if isinstance(t, int) and t >= 0), None) if last_accepted in chunk_terminators: @@ -896,22 +1087,482 @@ def _sample_minicpmo45_native_duplex_stage0( # token, so decide nothing here: re-emit the terminator and # leave the model-owned policy state exactly as the accepted # history left it. The next append re-injects that terminator. - sampled_ids.append(last_accepted) - continue - row_logits = logits[row_idx : row_idx + 1].clone() - sampled = self._sample_minicpmo45_native_duplex_row( - row_logits, + sampled_ids[row_idx] = last_accepted + else: + pending_rows.append(row_idx) + + if pending_rows: + deferred = self._sample_minicpmo45_native_duplex_rows_deferred( + logits, + sampling_metadata, + row_idxs=pending_rows, + token_ids=token_ids, + row_params=row_params, + ) + if deferred is not None: + # Lookahead rows re-emit their host-known terminator; the rest + # take their device-decided token. + sampled = index_to_device(sampled_ids, logits.device, torch.int32) + sampled[index_to_device(pending_rows, logits.device)] = deferred.to(torch.int32) + return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None) + batched = self._sample_minicpmo45_native_duplex_rows( + logits, sampling_metadata, - row_idx=row_idx, + row_idxs=pending_rows, token_ids=token_ids, + row_params=row_params, ) - self._record_minicpmo45_duplex_terminator(row_idx, sampled, token_ids) - sampled_ids.append(sampled) + for row_idx, sampled in zip(pending_rows, batched, strict=True): + sampled_ids[row_idx] = sampled + self._record_minicpmo45_duplex_terminator(row_idx, sampled, token_ids) + return SamplerOutput( - sampled_token_ids=torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32).unsqueeze(-1), + sampled_token_ids=index_to_device(sampled_ids, logits.device, torch.int32).unsqueeze(-1), logprobs_tensors=None, ) + def _sample_minicpmo45_native_duplex_rows( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + *, + row_idxs: list[int], + token_ids: dict[str, int], + row_params: list[tuple[float, float, float]], + ) -> list[int]: + """``_sample_minicpmo45_native_duplex_row`` for many rows with at most two host reads. + + A row with its own generator draws exactly as in the per-row loop. Rows + sharing the process-wide generator still draw once each, but all + boundary draws come before all stage-2 draws. + """ + chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) + all_greedy = bool(getattr(sampling_metadata, "all_greedy", False)) + generators = getattr(sampling_metadata, "generators", {}) or {} + recent_tokens_by_row = self._minicpmo45_duplex_recent_tokens(sampling_metadata, row_idxs) + + resolved: dict[int, int] = {} + if 0 <= chunk_eos_id < logits.shape[-1]: + capped = self._minicpmo45_duplex_capped_rows(row_idxs, recent_tokens_by_row) + resolved = dict.fromkeys(capped, int(chunk_eos_id)) + boundary_rows = [row_idx for row_idx in row_idxs if row_idx not in capped] + if boundary_rows: + hits = self._duplex_boundary_hits( + logits[index_to_device(boundary_rows, logits.device)], + boundary_rows, + generators, + all_greedy, + chunk_eos_id, + ) + for row_idx, hit in zip(boundary_rows, hits.tolist(), strict=True): + if hit: + resolved[row_idx] = int(chunk_eos_id) + + stage2_rows = [row_idx for row_idx in row_idxs if row_idx not in resolved] + if stage2_rows: + stage2_logits = self._minicpmo45_duplex_stage2_logits( + logits, + index_to_device(stage2_rows, logits.device), + stage2_rows, + token_ids, + recent_tokens_by_row, + ) + samples = self._minicpmo45_duplex_stage2_draws( + stage2_logits, stage2_rows, row_params, all_greedy, generators + ) + for row_idx, sampled in zip(stage2_rows, samples.tolist(), strict=True): + sampled = int(sampled) + self._record_minicpmo45_duplex_generation_token(row_idx, sampled) + sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk( + sampled, + recent_tokens_by_row[row_idx], + token_ids, + ) + resolved[row_idx] = self._finalize_minicpmo45_native_duplex_sample(row_idx, sampled, token_ids) + + return [resolved[row_idx] for row_idx in row_idxs] + + @staticmethod + def _minicpmo45_duplex_recent_tokens( + sampling_metadata: SamplingMetadata, + row_idxs: list[int], + ) -> dict[int, list[int]]: + """Each row's accepted output tokens.""" + output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] + recent_tokens_by_row: dict[int, list[int]] = {} + for row_idx in row_idxs: + raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else [] + recent_tokens_by_row[row_idx] = [ + int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0 + ] + return recent_tokens_by_row + + def _minicpmo45_policy_limit(self, name: str, default: int) -> int: + return int(getattr(self, name, default) or default) + + def _minicpmo45_duplex_capped_rows( + self, + row_idxs: list[int], + recent_tokens_by_row: dict[int, list[int]], + ) -> set[int]: + """Rows at their per-chunk speak-length cap: forced to chunk_eos without sampling.""" + max_speak_tokens = self._minicpmo45_policy_limit( + "max_new_speak_tokens_per_chunk", + MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK, + ) + capped: set[int] = set() + for row_idx in row_idxs: + effective_max_speak_tokens = max_speak_tokens + request_max_tokens = self._minicpmo45_duplex_row_request_max_tokens(row_idx) + if request_max_tokens is not None: + effective_max_speak_tokens = min(effective_max_speak_tokens, request_max_tokens) + if len(recent_tokens_by_row[row_idx]) >= max(1, effective_max_speak_tokens - 1): + capped.add(row_idx) + return capped + + def _minicpmo45_duplex_stage2_logits( + self, + logits: torch.Tensor, + index: torch.Tensor, + row_idxs: list[int], + token_ids: dict[str, int], + recent_tokens_by_row: dict[int, list[int]], + ) -> torch.Tensor: + """Rows ``index`` of ``logits`` with the forbidden ids masked and the repetition penalty applied.""" + vocab = logits.shape[-1] + stage2_logits = logits[index] # a gather: already a copy + forbidden_index = self._minicpmo45_duplex_forbidden_index(token_ids, vocab, logits.device) + if forbidden_index is not None: + stage2_logits.index_fill_(1, forbidden_index, float("-inf")) + repetition_penalty = 1.05 + history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE + penalty_rows: list[int] = [] + penalty_cols: list[int] = [] + for local_idx, row_idx in enumerate(row_idxs): + generated_tokens = getattr(self._minicpmo45_duplex_state_for_row(row_idx), "generated_tokens", None) + repetition_tokens = generated_tokens if generated_tokens else recent_tokens_by_row[row_idx] + penalized = [token_id for token_id in set(repetition_tokens[-history_size:]) if 0 <= token_id < vocab] + penalty_rows.extend([local_idx] * len(penalized)) + penalty_cols.extend(penalized) + if penalty_rows: + # Distinct (row, col) pairs, so one indexed division equals dividing each once. + rows_t = index_to_device(penalty_rows, logits.device) + cols_t = index_to_device(penalty_cols, logits.device) + stage2_logits[rows_t, cols_t] = stage2_logits[rows_t, cols_t] / repetition_penalty + return stage2_logits + + def _minicpmo45_duplex_stage2_draws( + self, + stage2_logits: torch.Tensor, + row_idxs: list[int], + row_params: list[tuple[float, float, float]], + all_greedy: bool, + generators: dict[int, torch.Generator], + marks: dict[int, tuple[torch.Generator, Any]] | None = None, + ) -> torch.Tensor: + """Each row's stage-2 token, drawn in row order with the row's own generator. + + Greedy rows (``all_greedy`` or temperature <= 0) take the argmax; the + rest draw in top-k/top-p candidate space. ``marks`` receives, per local + row, the generator position just before its draw. + """ + device = stage2_logits.device + samples: list[torch.Tensor | None] = [None] * len(row_idxs) + greedy_local: list[int] = [] + sampled_local: list[int] = [] + for local_idx, row_idx in enumerate(row_idxs): + if all_greedy or float(row_params[row_idx][0]) <= 0: + greedy_local.append(local_idx) + else: + sampled_local.append(local_idx) + if greedy_local: + greedy_sample = torch.argmax(stage2_logits[index_to_device(greedy_local, device)], dim=-1) + for pos, local_idx in enumerate(greedy_local): + samples[local_idx] = greedy_sample[pos : pos + 1] + if sampled_local: + cand_probs, cand_indices = self._duplex_stage2_candidate_probs( + stage2_logits[index_to_device(sampled_local, device)], + [row_params[row_idxs[local_idx]] for local_idx in sampled_local], + ) + for pos, local_idx in enumerate(sampled_local): + generator = generators.get(row_idxs[local_idx]) + if marks is not None: + marks[local_idx] = (generator, _generator_mark(generator)) + draw = torch.multinomial(cand_probs[pos], num_samples=1, generator=generator) + samples[local_idx] = cand_indices[pos].gather(0, draw) + return torch.cat(samples, dim=0) # type: ignore[arg-type] + + def _minicpmo45_duplex_row_params( + self, + sampling_metadata: SamplingMetadata, + num_rows: int, + ) -> list[tuple[float, float, float]]: + """Every row's (temperature, top_k, top_p) as ``_sampling_metadata_rows`` reads them. + + The runner's host copies (``DuplexSamplingRow``) hold the values the + device tensors were copied from, so reading them skips a device read. + A parameter the metadata carries no tensor for keeps its default, as in + the device read; without host copies for every row that read is used. + """ + host = getattr(self, "_minicpmo45_duplex_row_sampling_host", None) or {} + columns = [] + for column, (name, default) in enumerate((("temperature", 0.7), ("top_k", 100), ("top_p", 0.8))): + value = getattr(sampling_metadata, name, None) + if ( + isinstance(value, torch.Tensor) + and value.numel() == num_rows + and all(row in host for row in range(num_rows)) + ): + columns.append([float(host[row][column]) for row in range(num_rows)]) + else: + columns.append(self._sampling_metadata_rows(sampling_metadata, name, num_rows, default)) + return list(zip(*columns, strict=True)) + + def _minicpmo45_duplex_forbidden_index(self, token_ids: dict[str, int], vocab: int, device) -> torch.Tensor | None: + """``_minicpmo45_native_forbidden_token_ids`` inside the vocab, as a cached device index.""" + key = (vocab, str(device)) + cache = getattr(self, "_minicpmo45_duplex_forbidden_index_cache", None) + if cache is None or cache[0] != key: + forbidden = self._minicpmo45_native_forbidden_token_ids(token_ids) + valid = sorted({token_id for token_id in forbidden if 0 <= token_id < vocab}) + cache = (key, torch.tensor(valid, dtype=torch.long, device=device) if valid else None) + self._minicpmo45_duplex_forbidden_index_cache = cache + return cache[1] + + def _minicpmo45_duplex_cut_tables(self, token_ids: dict[str, int], vocab: int, device): + """Per-token decoded length and "never cut" mask for ``_maybe_cut_minicpmo45_native_duplex_text_chunk``. + + Byte-level decoding is additive after a prefix that ends on a + character boundary, so ``len(decode(chunk + [t]))`` is + ``len(decode(chunk)) + length[t]``. Returns None (no device cut) when + the tokenizer can not guarantee that: no batch decode, or the + whitespace clean-up that rewrites across token boundaries. + """ + key = (vocab, str(device)) + cache = getattr(self, "_minicpmo45_duplex_cut_tables_cache", None) + if cache is not None and cache[0] == key: + return cache[1] + tables = None + tokenizer = self._minicpmo45_tokenizer() + batch_decode = getattr(tokenizer, "batch_decode", None) + if callable(batch_decode) and not getattr(tokenizer, "clean_up_tokenization_spaces", False): + try: + num_tokens = min(vocab, len(tokenizer)) + texts = batch_decode([[token_id] for token_id in range(num_tokens)], skip_special_tokens=True) + except Exception: + texts = None + if texts is not None and len(texts) == num_tokens: + lengths = torch.zeros(vocab, dtype=torch.int32) + lengths[:num_tokens] = torch.tensor([len(text) for text in texts], dtype=torch.int32) + never_cut = torch.ones(vocab, dtype=torch.bool) + never_cut[:num_tokens] = False + special = [t for t in self._minicpmo45_native_special_token_ids(token_ids) if 0 <= t < vocab] + if special: + never_cut[torch.tensor(special, dtype=torch.long)] = True + tables = (lengths.to(device), never_cut.to(device)) + self._minicpmo45_duplex_cut_tables_cache = (key, tables) + return tables + + def _minicpmo45_duplex_cut_bases( + self, + row_idxs: list[int], + recent_tokens_by_row: dict[int, list[int]], + token_ids: dict[str, int], + ) -> list[int] | None: + """Each row's decoded current-chunk length (-1: never cut), or None when a row needs the host decode. + + Mirrors the early returns of ``_maybe_cut_minicpmo45_native_duplex_text_chunk``. + A chunk whose text ends in U+FFFD may end inside a character, where + the next token's bytes would merge with it, so it is not additive. + """ + max_chars = self._minicpmo45_policy_limit( + "max_speak_chars_per_chunk", + MiniCPMO45DuplexPolicy.DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK, + ) + tokenizer = self._minicpmo45_tokenizer() + decode = getattr(tokenizer, "decode", None) + if max_chars <= 0 or not callable(decode): + return [-1] * len(row_idxs) + bases = [] + for row_idx in row_idxs: + chunk = self._minicpmo45_current_chunk_tokens(recent_tokens_by_row[row_idx], token_ids) + try: + text = decode(chunk, skip_special_tokens=True) + except Exception: + return None + if not isinstance(text, str) or text.endswith("\ufffd"): + return None + bases.append(len(text)) + return bases + + def _sample_minicpmo45_native_duplex_rows_deferred( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + *, + row_idxs: list[int], + token_ids: dict[str, int], + row_params: list[tuple[float, float, float]], + ) -> torch.Tensor | None: + """``_sample_minicpmo45_native_duplex_rows`` decided on the device, without a host sync. + + Returns each row's token as a device tensor, or None when this step + needs the synchronous path. The cap, boundary decision, stage-2 draw, + chunk cut and listen->tts_bos rewrite are combined with ``torch.where``. + The host state they update is committed by + ``_commit_minicpmo45_duplex_pending_samples`` one step later, before + anything reads it. + + Stage 2 is drawn for every uncapped row, because the boundary outcome + is only known on the device; a row whose boundary hits has its + generator rewound at commit, so every generator advances exactly as on + the synchronous path. Rows on the process-wide generator can not be + rewound one at a time and keep the step synchronous. + """ + device = logits.device + vocab = logits.shape[-1] + chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) + has_chunk_eos = 0 <= chunk_eos_id < vocab + all_greedy = bool(getattr(sampling_metadata, "all_greedy", False)) + generators = getattr(sampling_metadata, "generators", {}) or {} + recent_tokens_by_row = self._minicpmo45_duplex_recent_tokens(sampling_metadata, row_idxs) + capped = self._minicpmo45_duplex_capped_rows(row_idxs, recent_tokens_by_row) if has_chunk_eos else set() + stage2_rows = [row_idx for row_idx in row_idxs if row_idx not in capped] + if any( + generators.get(row_idx) is None + for row_idx in stage2_rows + if not (all_greedy or float(row_params[row_idx][0]) <= 0) + ): + return None + cut = None + if has_chunk_eos and stage2_rows: + bases = self._minicpmo45_duplex_cut_bases(stage2_rows, recent_tokens_by_row, token_ids) + if bases is None: + return None + if any(base >= 0 for base in bases): + tables = self._minicpmo45_duplex_cut_tables(token_ids, vocab, device) + if tables is None: + return None + cut = (tables, bases) + + listen_id = token_ids.get("listen_token_id", -1) + tts_bos_id = token_ids.get("tts_bos_token_id", -1) + rewrite_listen = [] + for row_idx in stage2_rows: + state = self._minicpmo45_duplex_state_for_row(row_idx) + payload = self._minicpmo45_duplex_payload_for_row(row_idx) + rewrite_listen.append( + 0 <= tts_bos_id + and state is not None + and not getattr(state, "current_turn_ended", True) + and not (isinstance(payload, dict) and payload.get("force_listen") is True) + ) + + final = torch.full((len(row_idxs),), max(chunk_eos_id, 0), dtype=torch.long, device=device) + raw = torch.full((len(stage2_rows),), -1, dtype=torch.long, device=device) + hit = torch.zeros(len(stage2_rows), dtype=torch.bool, device=device) + rewinds: dict[int, tuple[torch.Generator, Any]] = {} + position = {row_idx: pos for pos, row_idx in enumerate(row_idxs)} + if stage2_rows: + stage2_index = index_to_device(stage2_rows, device) + positions_t = index_to_device([position[row_idx] for row_idx in stage2_rows], device) + if has_chunk_eos: + hit = self._duplex_boundary_hits( + logits[stage2_index], stage2_rows, generators, all_greedy, chunk_eos_id + ) + stage2_logits = self._minicpmo45_duplex_stage2_logits( + logits, stage2_index, stage2_rows, token_ids, recent_tokens_by_row + ) + raw = self._minicpmo45_duplex_stage2_draws( + stage2_logits, + stage2_rows, + row_params, + all_greedy, + generators, + marks=rewinds if has_chunk_eos else None, + ) + + token = raw + if cut is not None: + (lengths, never_cut), bases = cut + max_chars = self._minicpmo45_policy_limit( + "max_speak_chars_per_chunk", + MiniCPMO45DuplexPolicy.DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK, + ) + base_t = index_to_device(bases, device) + cut_t = (base_t >= 0) & ~never_cut[raw] & (base_t + lengths[raw].long() >= max_chars) + token = torch.where(cut_t, torch.full_like(raw, chunk_eos_id), raw) + if 0 <= listen_id and any(rewrite_listen): + rewrite_t = index_to_device([int(flag) for flag in rewrite_listen], device).bool() + token = torch.where((token == listen_id) & rewrite_t, torch.full_like(token, tts_bos_id), token) + if has_chunk_eos: + token = torch.where(hit, torch.full_like(token, chunk_eos_id), token) + final[positions_t] = token + + # One pinned copy of (final, raw stage-2 sample, boundary hit) per row + # for the commit; the host reads it once the runner has waited for it. + packed = torch.stack([final, torch.full_like(final, -1), torch.zeros_like(final)], dim=1) + if stage2_rows: + packed[positions_t, 1] = raw + packed[positions_t, 2] = hit.long() + host = torch.empty(packed.shape, dtype=packed.dtype, pin_memory=device.type == "cuda") + host.copy_(packed, non_blocking=True) + event = None + if device.type == "cuda": + event = torch.cuda.Event() + event.record() + self._minicpmo45_duplex_pending_samples = _MiniCPMO45PendingSamples( + host=host, + event=event, + row_idxs=list(row_idxs), + stage2=[row_idx not in capped for row_idx in row_idxs], + rewinds={position[stage2_rows[local_idx]]: rewind for local_idx, rewind in rewinds.items()}, + token_ids=token_ids, + row_sessions=getattr(self, "_minicpmo45_duplex_row_sessions", None), + row_payloads=getattr(self, "_minicpmo45_duplex_row_payloads", None), + ) + return final + + def _commit_minicpmo45_duplex_pending_samples(self, *, session_ids: set[str] | None = None) -> None: + """Apply the host state updates of the last deferred step, in the synchronous path's order. + + With ``session_ids``, only when that step sampled one of them: a + caller building those sessions' next appends needs their latches, and + must not wait for an unrelated step otherwise. + """ + pending = getattr(self, "_minicpmo45_duplex_pending_samples", None) + if pending is None: + return + if session_ids is not None: + sessions = pending.row_sessions or {} + if not any(sessions.get(row_idx) in session_ids for row_idx in pending.row_idxs): + return + self._minicpmo45_duplex_pending_samples = None + if pending.event is not None: + pending.event.synchronize() + values = pending.host.tolist() + saved = ( + getattr(self, "_minicpmo45_duplex_row_sessions", None), + getattr(self, "_minicpmo45_duplex_row_payloads", None), + ) + self._minicpmo45_duplex_row_sessions = pending.row_sessions + self._minicpmo45_duplex_row_payloads = pending.row_payloads + try: + for pos, row_idx in enumerate(pending.row_idxs): + _final, raw, hit = values[pos] + if not pending.stage2[pos]: + continue + if hit: + rewind = pending.rewinds.get(pos) + if rewind is not None: + _generator_rewind(*rewind) + else: + self._record_minicpmo45_duplex_generation_token(row_idx, int(raw)) + for pos, row_idx in enumerate(pending.row_idxs): + self._record_minicpmo45_duplex_terminator(row_idx, int(values[pos][0]), pending.token_ids) + finally: + self._minicpmo45_duplex_row_sessions, self._minicpmo45_duplex_row_payloads = saved + def _sample_minicpmo45_native_duplex_row( self, logits: torch.Tensor, @@ -919,15 +1570,21 @@ def _sample_minicpmo45_native_duplex_row( *, row_idx: int, token_ids: dict[str, int], + params: tuple[float, float, float] | None = None, ) -> int: + """Sample one duplex row; ``params`` is the row's (temperature, top_k, top_p) if already read.""" chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) generator = getattr(sampling_metadata, "generators", {}).get(row_idx) output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else [] recent_tokens = [int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0] - temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7)) - top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100)) - top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8)) + if params is None: + params = ( + self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7), + self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100), + self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8), + ) + temperature, top_k, top_p = float(params[0]), int(params[1]), float(params[2]) state = self._minicpmo45_duplex_state_for_row(row_idx) if chunk_eos_id >= 0 and chunk_eos_id < logits.shape[-1]: max_speak_tokens = int( @@ -949,12 +1606,8 @@ def _sample_minicpmo45_native_duplex_row( # distribution only to preserve the model's own chunk boundary. # If it does not choose chunk_eos, mask that token before the # normal text/listen/turn sampling pass below. - if getattr(sampling_metadata, "all_greedy", False): - boundary_sample = int(torch.argmax(logits, dim=-1).item()) - else: - boundary_probs = F.softmax(logits, dim=-1) - boundary_sample = int(torch.multinomial(boundary_probs, num_samples=1, generator=generator).item()) - if boundary_sample == chunk_eos_id: + all_greedy = bool(getattr(sampling_metadata, "all_greedy", False)) + if self._duplex_boundary_hits(logits, [row_idx], {row_idx: generator}, all_greedy, chunk_eos_id).item(): return int(chunk_eos_id) logits = logits.clone() @@ -969,10 +1622,11 @@ def _sample_minicpmo45_native_duplex_row( repetition_penalty = 1.05 if repetition_penalty != 1.0 and repetition_tokens: history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE - for token_id in set(repetition_tokens[-history_size:]): - if token_id < 0 or token_id >= logits.shape[-1]: - continue - logits[0, token_id] /= repetition_penalty + penalized = [t for t in set(repetition_tokens[-history_size:]) if 0 <= t < logits.shape[-1]] + if penalized: + # Distinct ids, so one indexed division equals dividing each once. + columns = index_to_device(penalized, logits.device) + logits[0, columns] = logits[0, columns] / repetition_penalty if getattr(sampling_metadata, "all_greedy", False) or temperature <= 0: sampled = int(torch.argmax(logits, dim=-1).item()) @@ -984,10 +1638,9 @@ def _sample_minicpmo45_native_duplex_row( ) return self._finalize_minicpmo45_native_duplex_sample(row_idx, sampled, token_ids) - logits = logits / temperature - logits = self._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p) - probs = F.softmax(logits, dim=-1) - sampled = int(torch.multinomial(probs, num_samples=1, generator=generator).item()) + cand_probs, cand_indices = self._duplex_top_k_top_p_candidates(logits / temperature, top_k=top_k, top_p=top_p) + draw = torch.multinomial(cand_probs[0], num_samples=1, generator=generator) + sampled = int(cand_indices[0].gather(0, draw).item()) self._record_minicpmo45_duplex_generation_token(row_idx, sampled) sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk( sampled, @@ -1261,6 +1914,23 @@ def _minicpmo45_native_special_token_ids(self, token_ids: dict[str, int]) -> set tokenizer_special_ids=getattr(tokenizer, "all_special_ids", []) if tokenizer is not None else [], ) + @classmethod + def _sampling_metadata_rows( + cls, + sampling_metadata: SamplingMetadata, + name: str, + num_rows: int, + default: float, + ) -> list[float]: + """``_sampling_metadata_value`` for rows ``0..num_rows-1`` with a single device read.""" + value = getattr(sampling_metadata, name, None) + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return [default] * num_rows + flat = [float(v) for v in value.reshape(-1).tolist()] + return [flat[min(row, len(flat) - 1)] for row in range(num_rows)] + return [cls._sampling_metadata_value(sampling_metadata, name, 0, default)] * num_rows + @staticmethod def _sampling_metadata_value( sampling_metadata: SamplingMetadata, @@ -1299,6 +1969,80 @@ def _top_k_top_p_filter(logits: torch.Tensor, *, top_k: int, top_p: float) -> to logits = logits.masked_fill(remove, float("-inf")) return logits + @staticmethod + def _duplex_top_k_top_p_candidates( + logits: torch.Tensor, *, top_k: int, top_p: float + ) -> tuple[torch.Tensor, torch.Tensor]: + """``softmax(_top_k_top_p_filter(logits, ...))`` restricted to its support, as ``(probs, indices)``. + + The nucleus is decided on the k sorted candidates, O(V + k log k) + instead of sorting the vocabulary. The only difference from the filter + is an exact tie at the k-th logit: exactly k candidates are kept here. + """ + vocab = logits.shape[-1] + keep = vocab if not top_k or top_k <= 0 else min(int(top_k), vocab) + values, indices = torch.topk(logits, keep, dim=-1) + probs = torch.softmax(values, dim=-1) + if 0.0 < float(top_p) < 1.0: + cumulative = probs.cumsum(dim=-1) + remove = cumulative > float(top_p) + remove[..., 1:] = remove[..., :-1].clone() + remove[..., 0] = False + probs = probs.masked_fill(remove, 0.0) + probs = probs / probs.sum(dim=-1, keepdim=True) + return probs, indices + + def _duplex_stage2_candidate_probs( + self, + rows_logits: torch.Tensor, + row_params: list[tuple[float, float, float]], + ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + """Per-row ``(candidate probs, candidate ids)``; rows sharing ``(top_k, top_p)`` share one top-k.""" + groups: dict[tuple[int, float], list[int]] = {} + for pos, (_temperature, top_k, top_p) in enumerate(row_params): + groups.setdefault((int(top_k), float(top_p)), []).append(pos) + temps = to_device_nonblocking( + torch.tensor([float(p[0]) for p in row_params], dtype=rows_logits.dtype), + rows_logits.device, + ).view(-1, 1) + probs: list[torch.Tensor] = [torch.empty(0)] * len(row_params) + indices: list[torch.Tensor] = [torch.empty(0, dtype=torch.long)] * len(row_params) + for (top_k, top_p), members in groups.items(): + members_t = index_to_device(members, rows_logits.device) + rows = rows_logits[members_t] / temps[members_t] + cand_probs, cand_indices = self._duplex_top_k_top_p_candidates(rows, top_k=top_k, top_p=top_p) + for row, pos in enumerate(members): + probs[pos] = cand_probs[row] + indices[pos] = cand_indices[row] + return probs, indices + + @staticmethod + def _duplex_boundary_chunk_eos_probs(logits: torch.Tensor, chunk_eos_id: int) -> torch.Tensor: + """Per-row ``softmax(logits)[chunk_eos_id] = exp(l_eos - logsumexp(l))``.""" + eos_logits = logits[:, chunk_eos_id] + return torch.exp(eos_logits - torch.logsumexp(logits, dim=-1)) + + @classmethod + def _duplex_boundary_hits( + cls, + rows_logits: torch.Tensor, + row_idxs: list[int], + generators: dict[int, torch.Generator | None], + all_greedy: bool, + chunk_eos_id: int, + ) -> torch.Tensor: + """Whether each row's boundary sample is chunk_eos, drawn in row order with the row's own generator. + + The boundary sample is discarded unless it is chunk_eos, so one + Bernoulli at P(sample == chunk_eos) is the same decision as a + full-vocabulary softmax + multinomial. + """ + if all_greedy: + return torch.argmax(rows_logits, dim=-1) == chunk_eos_id + boundary_p = cls._duplex_boundary_chunk_eos_probs(rows_logits, chunk_eos_id) + draws = [torch.rand((), generator=generators.get(row_idx), device=rows_logits.device) for row_idx in row_idxs] + return torch.stack(draws) < boundary_p + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights for all components of the omni model.""" loaded_weights = set() diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 4577b7c5c66..3892bcb1f54 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -744,6 +744,12 @@ def _parse_k_step_frames(cls, vllm_config: Any = None) -> int: num_spec = getattr(spec, "num_speculative_tokens", 0) or 0 if method != "ngram": return 0 + if not current_omni_platform.is_npu(): + # The in-model sampler and two-wide stop/continue head are the NPU + # runner's contract. On CUDA the same speculative_config drives + # talker_kstep.py, which keeps this single-frame + # codec head and vLLM's sampler. + return 0 frames = int(num_spec) + 1 if frames < 2: # K=1 degenerates to the ordinary one-frame path; treat it as off @@ -1141,11 +1147,7 @@ def preprocess( # which defers this ``.item()`` EOS check to make_omni_output. k_last = state.get("last_code") if isinstance(state, dict) else None if isinstance(k_last, int) and k_last >= 0: - code = torch.tensor( - [k_last], - dtype=torch.long, - device=self.emb_code[0].weight.device, - ) + code = index_to_device([k_last], self.emb_code[0].weight.device) # k_last is already the host int this tensor was built from; # .item() would only sync it right back. code_id = k_last @@ -2172,12 +2174,19 @@ def _make_omni_output_single_frame( audio = torch.tensor([[code_id]], dtype=torch.long, device="cpu") empty_speech = bool(state.get("finished")) if isinstance(audio, torch.Tensor) and audio.numel() > 0: - codec_deltas[index] = audio.to(device=hidden.device, dtype=torch.long).reshape(-1, 1) + # The ids are read on the host from ``audio`` itself when it is + # a host tensor; the device copy and the history never wait. + delta_ids = audio.reshape(-1).tolist() if audio.device.type == "cpu" else None + if delta_ids is not None: + codec_deltas[index] = index_to_device(delta_ids, hidden.device).reshape(-1, 1) + else: + codec_deltas[index] = audio.to(device=hidden.device, dtype=torch.long).reshape(-1, 1) + delta_ids = codec_deltas[index].cpu().reshape(-1).tolist() state["step"] = int(state.get("step", 0)) + 1 # ``audio`` is the id sampled last step, i.e. exactly the # history the logits computed below are scored against. recent = state.get("recent_codes") - recent = (recent if isinstance(recent, list) else []) + codec_deltas[index].cpu().reshape(-1).tolist() + recent = (recent if isinstance(recent, list) else []) + delta_ids state["recent_codes"] = recent[-_CODEC_PENALTY_WINDOW:] recent_codes = state.get("recent_codes") if recent_codes: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index ba2e8a19942..c30a326192a 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -21,23 +21,6 @@ _CODEC_EOS_TOKEN_ID = 6561 # tts_config.num_audio_tokens - 1 -def _talker_stop_token_ids() -> list[int]: - """Stage 1's default stop id: the codec EOS of the one-frame head. - - The deploy config that arms the multi-frame decode (stage 1's - ``speculative_config``, under ``platforms.npu``) collapses the vLLM-level - head to the two-wide continue/stop row, where the codec EOS can never - appear; that same config adds the stop marker (1) through its - ``default_sampling_params``, and ``merge_sampling_constraints`` unions both - lists. The marker therefore travels with the block that needs it, and this - default covers every deployment without one. - - A deployment that keeps the marker here instead ends every one-frame - request on the first ordinary codec id 1. - """ - return [_CODEC_EOS_TOKEN_ID] - - MINICPMO_4_5_PIPELINE = PipelineConfig( model_type="minicpmo_4_5", default_deploy_config_name="minicpmo_4_5.yaml", @@ -89,10 +72,10 @@ def _talker_stop_token_ids() -> list[int]: supports_native_mrv2_data_plane=True, sampling_constraints={ "detokenize": False, - # The stop id has to be one the vLLM-level head can actually - # emit, and that head changes shape when the K-frame loop is - # armed -- see _talker_stop_token_ids. - "stop_token_ids": _talker_stop_token_ids(), + # The one-frame head's codec EOS. A multi-frame deploy config + # adds its stop marker (1) via default_sampling_params; keeping + # it here would end every one-frame request on codec id 1. + "stop_token_ids": [_CODEC_EOS_TOKEN_ID], }, ), StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_frame_plan.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_frame_plan.py new file mode 100644 index 00000000000..4c1c5e689b3 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_frame_plan.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Talker codec-EOS plan for the CUDA multi-frame decode. + +The single-frame EOS force/mask routing (``_make_omni_output_single_frame``) +depends only on the request's frame counter, which advances by one per +forwarded frame, so frames 1..K-1 can be planned on the host up front. +``commit_codec_frames`` then leaves the host state where that many +single-frame steps would have. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + _CODEC_PENALTY_WINDOW, + _turn_end_boundary_eos_masked, +) + +CODEC_PENALTY_WINDOW = _CODEC_PENALTY_WINDOW + + +@dataclass +class CodecFramePlan: + """Per-frame codec-EOS routing (``[frame][request]``; row 0 unused) and the + repetition window after frame 0's bookkeeping.""" + + request_ids: list[str] + eos_token_id: int + force_eos: list[list[bool]] + mask_eos: list[list[bool]] + recent_codes: list[list[int]] + + +def plan_codec_frames(talker: Any, infos: list[Any], frames: int) -> CodecFramePlan: + """Plan frames 1..frames-1 from the state frame 0's bookkeeping left.""" + num_reqs = len(infos) + force = [[False] * num_reqs for _ in range(frames)] + mask = [[False] * num_reqs for _ in range(frames)] + request_ids: list[str] = [] + recent_codes: list[list[int]] = [] + states = talker._request_audio_states + for index, info in enumerate(infos): + request_id = str(info.get("request_id", index)) if isinstance(info, dict) else str(index) + request_ids.append(request_id) + state = states.get(request_id) + state = state if isinstance(state, dict) else {} + finished = bool(state.get("finished")) + step0 = int(state.get("step", 0)) + max_tokens = state.get("max_tokens") + min_tokens = state.get("min_tokens") + drain = bool(state.get("turn_end_drain")) + for k in range(1, frames): + step = step0 + k + chunk_done = finished or (max_tokens is not None and step >= int(max_tokens) - 1) + force[k][index] = chunk_done + mask[k][index] = not chunk_done and ( + (min_tokens is not None and step < int(min_tokens)) or (drain and _turn_end_boundary_eos_masked(step)) + ) + codes = state.get("recent_codes") + recent_codes.append([int(c) for c in codes][-CODEC_PENALTY_WINDOW:] if isinstance(codes, list) else []) + return CodecFramePlan( + request_ids=request_ids, + eos_token_id=int(talker._codec_eos_id), + force_eos=force, + mask_eos=mask, + recent_codes=recent_codes, + ) + + +def commit_codec_frames(talker: Any, plan: CodecFramePlan, forwarded: list[list[int]]) -> list[bool]: + """Apply each request's forwarded ids (sampled on frames 0..n-2) as that many + single-frame steps; return their terminal flags.""" + flags: list[bool] = [] + for index, (request_id, codes) in enumerate(zip(plan.request_ids, forwarded)): + state = talker._request_audio_states.get(request_id) + if not codes or not isinstance(state, dict): + flags.append(False) + continue + state["step"] = int(state.get("step", 0)) + len(codes) + recent = state.get("recent_codes") + state["recent_codes"] = ((recent if isinstance(recent, list) else []) + [int(c) for c in codes])[ + -CODEC_PENALTY_WINDOW: + ] + done = plan.force_eos[len(codes)][index] + if done: + state["finished"] = True + flags.append(done) + return flags diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_kstep.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_kstep.py new file mode 100644 index 00000000000..4f028ea94b5 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_kstep.py @@ -0,0 +1,670 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CUDA multi-frame (K-step) decode for the MiniCPM-o Talker. + +Stage 1's n-gram ``speculative_config`` (``num_speculative_tokens == K - 1``) +schedules K positions per request, and one ``execute_model`` produces up to K +codec frames, so the per-step host work is paid once per K frames. Unlike the +NPU runner, the codec head, the scheduler's codec ids and vLLM's ``Sampler`` +stay as in the single-frame path: + +* frame 0 is the ordinary single-frame step (the drafts repeat the last id, so + preprocess embeds it at the span's first row); +* frames 1..K-1 replay the same forward after writing the previous frame's + embedding into its row (attention is causal, so later stale rows are + harmless), with the stop decision kept on device; the step reads its K ids + back once and returns a rejection-sampler-shaped output (``-1`` padded). + +A declined step is exactly one single-frame step: it returns frame 0's id as a +``(B, 1)`` output and the scheduler rejects the drafts. ``input_batch.vocab_size`` +must cover the codec head (``ensure_codec_vocab``), since vLLM drops ids at or +above it from multi-token outputs. + +The GPU runners reach this module through the model's +``multi_frame_decode_hook`` (``None`` on every other model and on NPU, whose +in-model K-step lives in ``platforms/npu/worker/talker_multiframe.py``). +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Any + +import torch +from vllm.logger import init_logger +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.logits_processor import LogitsProcessors +from vllm.v1.sample.logits_processor.builtin import ( + LogitBiasLogitsProcessor, + MinPLogitsProcessor, + MinTokensLogitsProcessor, +) +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids + +logger = init_logger(__name__) + +_LOGGED_BLOCKS: set[str] = set() + + +def _log_block_once(reason: str, detail: str = "") -> None: + """Log once per distinct reason (with its first numbers) why the loop did not engage.""" + if reason in _LOGGED_BLOCKS: + return + _LOGGED_BLOCKS.add(reason) + suffix = f" ({detail})" if detail else "" + logger.info("[minicpmo] multi-frame Talker decode not engaged: %s%s", reason, suffix) + + +def _block(reason: str, detail: str = "") -> int: + _log_block_once(reason, detail) + return 0 + + +def applies( + model: Any, + model_kwargs_extra: dict[str, Any], + *, + require_sampled_frame: bool = True, +) -> int: + """The uniform per-request query length of a pure Talker decode step, or 0. + + A prefill, a mixed step or non-uniform spans return 0 (the single-forward + path). ``require_sampled_frame`` is the NPU contract (frame 0 embeds the + model's recorded previous code); the CUDA runner passes ``False``. A + multi-token step that is refused must not take the single-forward path; + the runners raise on it (``is_multi_token_decode``). + """ + if not getattr(model, "supports_multi_frame_decode", False): + return 0 + spans = model_kwargs_extra.get("request_token_spans") + infos = model_kwargs_extra.get("model_intermediate_buffer") + if not spans or not infos or len(spans) != len(infos): + return _block( + "no request_token_spans/model_intermediate_buffer for this step", + f"{len(spans or ())} spans, {len(infos or ())} buffers", + ) + frames = int(spans[0][1]) - int(spans[0][0]) + if frames <= 1: + return 0 + for start, end in spans: + if int(end) - int(start) != frames: + return _block( + "request token spans are not uniform", + f"frames {frames}, spans {[(int(s), int(e)) for s, e in spans]}", + ) + for index, info in enumerate(infos): + if not isinstance(info, dict): + return _block("a request carries no intermediate buffer", f"request {index} of {len(infos)}") + if bool(info.get("_omni_is_prefill", False)): + return 0 + state = info.get("audio_state") + if not isinstance(state, dict): + return _block("a request carries no audio_state", f"request {info.get('request_id', index)}") + if require_sampled_frame and int(state.get("step", 0)) <= 0: + return _block("a request has not sampled a codec token yet", f"request {info.get('request_id', index)}") + return frames + + +def is_multi_token_decode(model: Any, model_kwargs_extra: dict[str, Any]) -> bool: + """True when this Talker step schedules several tokens for a decoding request.""" + if not getattr(model, "supports_multi_frame_decode", False): + return False + spans = model_kwargs_extra.get("request_token_spans") + infos = model_kwargs_extra.get("model_intermediate_buffer") + if not spans or not infos or len(spans) != len(infos): + return False + for (start, end), info in zip(spans, infos): + if int(end) - int(start) <= 1: + continue + if isinstance(info, dict) and not bool(info.get("_omni_is_prefill", False)): + return True + return False + + +def drafts_this_step(runner: Any) -> int: + """Frames the *next* step is scheduled for (``num_spec_tokens + 1``), or 0. + + Speculative tokens are vLLM V1's only way to reserve KV for several + positions; the frames are generated sequentially, not verified. + """ + if not getattr(getattr(runner, "model", None), "supports_multi_frame_decode", False): + return 0 + num_spec = int(getattr(runner, "num_spec_tokens", 0) or 0) + if num_spec <= 0: + return 0 + return num_spec + 1 + + +def constant_drafts( + valid_sampled_token_ids: Any, + frames: int, + num_reqs: int, + *, + draft_token: int | None, +) -> list[list[int]]: + """``frames - 1`` drafts per request -- for the whole batch, or for none of it. + + The loop needs uniform spans, so one request that sampled nothing this step + means no drafts at all (an ordinary one-frame step). ``draft_token=None`` + repeats each request's last sampled id (the CUDA runner embeds it at frame 0). + """ + rows: list[list[int]] = [] + for index in range(num_reqs): + sampled = None + if isinstance(valid_sampled_token_ids, list) and index < len(valid_sampled_token_ids): + sampled = valid_sampled_token_ids[index] + if not sampled: + _log_block_once( + "a request in this batch sampled nothing; no drafts this step", + f"request {index} of {num_reqs}", + ) + return [[] for _ in range(num_reqs)] + token = int(sampled[-1]) if draft_token is None else int(draft_token) + rows.append([token] * (frames - 1)) + return rows + + +_STASH_ATTR = "_talker_frames_sampler_output" +_BUILTIN_LOGITSPROCS = (MinTokensLogitsProcessor, LogitBiasLogitsProcessor, MinPLogitsProcessor) +_SAMPLER: Sampler | None = None + + +@dataclass(frozen=True) +class Decline: + """Why a step runs one frame; logs dedupe on ``kind``, ``detail`` has numbers.""" + + kind: str + detail: str = "" + + def __str__(self) -> str: + return f"{self.kind} ({self.detail})" if self.detail else self.kind + + +_LOGGED: set[tuple[str, str]] = set() + + +def _log_once(site: str, decline: Decline) -> None: + """Log each distinct decline once per site, with its first numbers.""" + key = (site, decline.kind) + if key not in _LOGGED: + _LOGGED.add(key) + logger.info("[minicpmo] CUDA multi-frame Talker decode: %s: %s", site, decline) + + +def _codec_sampler() -> Sampler: + # Same default ``Sampler()`` the Talker's own sample() uses for frame 0. + global _SAMPLER + if _SAMPLER is None: + _SAMPLER = Sampler() + return _SAMPLER + + +def supports(model: Any) -> bool: + return bool(getattr(model, "supports_multi_frame_decode", False)) and callable( + getattr(model, "plan_codec_frames", None) + ) + + +def codec_vocab_size(model: Any) -> int: + """Width of the Talker's codec head, or 0 when the model does not say.""" + return int(getattr(model, "codec_vocab_size", 0) or 0) + + +def ensure_codec_vocab(runner: Any) -> None: + """Make stage 1 report the codec head's width (6562) as its vocab size. + + The remote-code ``MiniCPMTTSConfig`` has no ``vocab_size``, so vLLM reports + 0: ``RejectionSampler.parse_output`` then empties every multi-token row and + ``InputBatch.add_request`` silently drops the stage's ``top_k``. Called at + ``load_model``, before ``initialize_kv_cache`` rebuilds the ``InputBatch`` + from ``model_config``. + """ + model = getattr(runner, "model", None) + head = codec_vocab_size(model) + if not supports(model) or head <= 0: + return + model_config = getattr(runner, "model_config", None) + reported = int(model_config.get_vocab_size() or 0) if model_config is not None else head + arch = getattr(model_config, "model_arch_config", None) + if reported < head and arch is not None: + arch.vocab_size = head + batch = getattr(runner, "input_batch", None) + batch_vocab = int(getattr(batch, "vocab_size", 0) or 0) if batch is not None else head + if batch is not None and batch_vocab < head: + batch.vocab_size = head + if reported < head or batch_vocab < head: + logger.info( + "[minicpmo] Talker stage reported vocab_size %d (input batch %d); using the %d-wide codec head, " + "so multi-token steps keep their ids and the stage's top_k applies", + reported, + batch_vocab, + head, + ) + + +def _stop_ids(sampling_params: Any) -> set[int]: + """The ids ``check_stop`` ends a request on (EOS and stop ids).""" + ids = {int(t) for t in (getattr(sampling_params, "stop_token_ids", None) or ())} + eos = getattr(sampling_params, "eos_token_id", None) + if eos is not None: + ids.add(int(eos)) + return ids + + +def vocab_decline(runner: Any, head_vocab: int) -> Decline | None: + """Decline when ``parse_output`` would drop the frames (batch vocab narrower than the head).""" + batch_vocab = int(getattr(runner.input_batch, "vocab_size", 0) or 0) + if head_vocab > 0 and batch_vocab < head_vocab: + return Decline( + "input batch vocab is narrower than the codec head", + f"input batch vocab {batch_vocab}, codec head {head_vocab}", + ) + return None + + +def ineligible_reason(runner: Any, eos_token_id: int) -> Decline | None: + """Why this batch cannot run several frames per step, or None.""" + if getattr(runner, "use_async_scheduling", False): + # The drafts are this step's last ids and the accepted count must be + # known before the next step is scheduled. + return Decline("async scheduling is on", "stage 1 needs async_scheduling: false") + batch = runner.input_batch + md = batch.sampling_metadata + if md.max_num_logprobs is not None or md.logprob_token_ids: + return Decline("logprobs requested", f"max_num_logprobs {md.max_num_logprobs}") + if md.bad_words_token_ids or md.allowed_token_ids_mask is not None: + return Decline("bad words / allowed token ids") + holder = getattr(md, "thinking_budget_state_holder", None) + if holder is not None and holder.has_tracked_requests(): + return Decline("thinking budget") + custom = [type(proc).__name__ for proc in md.logitsprocs.all if not isinstance(proc, _BUILTIN_LOGITSPROCS)] + if custom: + return Decline("custom logits processor", ", ".join(custom)) + for req_id in batch.req_ids: + req_state = runner.requests.get(req_id) + params = getattr(req_state, "sampling_params", None) + if params is None: + return Decline("request without sampling params", f"request {req_id}") + if params.frequency_penalty or params.presence_penalty: + return Decline( + "frequency/presence penalty", + f"request {req_id}: frequency {params.frequency_penalty}, presence {params.presence_penalty}", + ) + if params.structured_outputs is not None: + return Decline("structured output", f"request {req_id}") + if getattr(params, "repetition_detection", None) is not None: + return Decline("repetition detection", f"request {req_id}") + stops = _stop_ids(params) + if eos_token_id not in stops: + return Decline( + "codec EOS is not a stop id", + f"request {req_id}: codec EOS {eos_token_id}, stop ids {sorted(stops)}", + ) + return None + + +def draft_decline(runner: Any, model: Any) -> Decline | None: + """What ``maybe_run`` would decline on that is known before drafting (a drafted + step forwards K rows per request whether or not the loop runs).""" + if not bool(getattr(model, "requires_request_sample_eligibility", False)): + # The runner builds request_max_tokens_remaining only for such models. + return Decline("no per-request token budget", "the model does not declare requires_request_sample_eligibility") + return vocab_decline(runner, codec_vocab_size(model)) or ineligible_reason(runner, int(model.codec_eos_token_id)) + + +def propose_drafts(runner: Any, sampled_token_ids: Any) -> list[list[int]] | None: + """The Talker's next-step drafts (last id K-1 times, whole batch or none), or + None for any other stage.""" + model: Any = getattr(runner, "model", None) + if not supports(model): + return None + frames = drafts_this_step(runner) + num_reqs = runner.input_batch.num_reqs + if frames <= 1: + return None + decline = draft_decline(runner, model) + if decline is None and not isinstance(sampled_token_ids, list): + decline = Decline("sampled ids are not a list", type(sampled_token_ids).__name__) + if decline is not None: + _log_once("no drafts", decline) + return [[] for _ in range(num_reqs)] + return constant_drafts(sampled_token_ids, frames, num_reqs, draft_token=None) + + +def take_sampler_output(runner: Any) -> SamplerOutput | None: + output = getattr(runner, _STASH_ATTR, None) + setattr(runner, _STASH_ATTR, None) + return output + + +@dataclass +class FrameControls: + """Device tensors driving frames 1..K-1 (``(K, B)`` rows are per frame).""" + + eos_token_id: int + force_eos: torch.Tensor # (K, B) bool + mask_eos: torch.Tensor # (K, B) bool + allowed: torch.Tensor # (K, B) bool: frame k is inside the token budget + min_tokens: torch.Tensor # (K, B) bool: vLLM min_tokens masks stop ids + min_tokens_stop: torch.Tensor | None # (B, V) bool: the stop ids it masks + stop_ids: torch.Tensor # (B, S) long, -1 padded: check_stop's ids + window: torch.Tensor # (B, W) long, right-aligned, padded with V + penalties: torch.Tensor | None # (B,) repetition penalty, None when off + forced_row: torch.Tensor # (V,) float: -inf except 0.0 at EOS + eos_col: torch.Tensor # (V,) bool + + +def _upload(tensors: list[torch.Tensor], device: torch.device) -> list[torch.Tensor]: + """One pinned host-to-device copy for all of a step's small int tensors.""" + flat = torch.cat([t.reshape(-1).to(torch.long) for t in tensors]) + if device.type == "cuda": + flat = flat.pin_memory().to(device, non_blocking=True) + else: + flat = flat.to(device) + out, offset = [], 0 + for t in tensors: + n = t.numel() + out.append(flat[offset : offset + n].view(t.shape)) + offset += n + return out + + +def build_controls( + plan: Any, + *, + frames: int, + vocab_size: int, + budgets: list[int], + stop_ids: list[set[int]], + min_tokens_state: dict[int, tuple[Any, ...]], + penalties: torch.Tensor | None, + device: torch.device, + extra: list[torch.Tensor] | None = None, +) -> tuple[FrameControls, list[torch.Tensor]]: + """Build the loop controls from host state; ``extra`` rides the same upload.""" + from vllm_omni.model_executor.models.minicpmo_4_5.talker_frame_plan import CODEC_PENALTY_WINDOW as WINDOW + + num_reqs = len(budgets) + force = torch.tensor(plan.force_eos, dtype=torch.bool) + mask = torch.tensor(plan.mask_eos, dtype=torch.bool) + allowed = torch.tensor([[k < int(b) for b in budgets] for k in range(frames)], dtype=torch.bool) + # vLLM's MinTokensLogitsProcessor state: batch index -> (min_tokens, + # output ids, stop ids, structured). Its mask is computed once per step + # from len(output ids); frame k has k more outputs than that. + min_mask = torch.zeros((frames, num_reqs), dtype=torch.bool) + pairs: list[tuple[int, int]] = [] + for index, (min_toks, out_ids, stop_tok_ids, _) in min_tokens_state.items(): + if index >= num_reqs: + continue + for k in range(frames): + min_mask[k, index] = len(out_ids) + k < int(min_toks) + pairs.extend((index, int(t)) for t in stop_tok_ids if 0 <= int(t) < vocab_size) + width = max(1, max((len(s) for s in stop_ids), default=1)) + stop_table = torch.full((num_reqs, width), -1, dtype=torch.long) + for index, ids in enumerate(stop_ids): + kept = sorted(t for t in ids if 0 <= t < vocab_size) + stop_table[index, : len(kept)] = torch.tensor(kept, dtype=torch.long) + window = torch.full((num_reqs, WINDOW), vocab_size, dtype=torch.long) + for index, codes in enumerate(plan.recent_codes): + codes = codes[-WINDOW:] + if codes: + window[index, WINDOW - len(codes) :] = torch.tensor(codes, dtype=torch.long) + pair_tensor = torch.tensor(pairs, dtype=torch.long).reshape(-1, 2) + host = [force, mask, allowed, min_mask, stop_table, window, pair_tensor, *(extra or [])] + dev = _upload(host, device) + force_d, mask_d, allowed_d, min_d, stop_d, window_d, pairs_d = dev[:7] + min_stop = None + if pairs: + min_stop = torch.zeros((num_reqs, vocab_size), dtype=torch.bool, device=device) + min_stop[pairs_d[:, 0], pairs_d[:, 1]] = True + eos = int(plan.eos_token_id) + eos_col = torch.arange(vocab_size, device=device) == eos + forced_row = torch.where( + eos_col, + torch.zeros((), dtype=torch.float32, device=device), + torch.full((), float("-inf"), dtype=torch.float32, device=device), + ) + controls = FrameControls( + eos_token_id=eos, + force_eos=force_d.bool(), + mask_eos=mask_d.bool(), + allowed=allowed_d.bool(), + min_tokens=min_d.bool(), + min_tokens_stop=min_stop, + stop_ids=stop_d, + window=window_d, + penalties=penalties, + forced_row=forced_row, + eos_col=eos_col, + ) + return controls, dev[7:] + + +def adjust_logits(logits: torch.Tensor, k: int, window: torch.Tensor, c: FrameControls) -> torch.Tensor: + """Frame k's logits as the single-frame path shapes them before ``Sampler``: + EOS force/mask, the windowed repetition penalty (same ops as + ``_apply_batched_repetition_penalty``, bit-identical), then min_tokens.""" + neg_inf = float("-inf") + logits = torch.where(c.force_eos[k].unsqueeze(1), c.forced_row, logits) + logits = torch.where(c.mask_eos[k].unsqueeze(1) & c.eos_col, neg_inf, logits) + if c.penalties is not None: + num_reqs, vocab = logits.shape + counts = torch.zeros((num_reqs, vocab + 1), dtype=torch.long, device=logits.device) + counts.scatter_add_(1, window, torch.ones_like(window)) + alpha = torch.pow(c.penalties.to(logits.dtype).unsqueeze(1), counts[:, :vocab].to(logits.dtype)) + logits = torch.where(logits < 0, logits * alpha, logits / alpha) + if c.min_tokens_stop is not None: + logits = logits.masked_fill(c.min_tokens_stop & c.min_tokens[k].unsqueeze(1), neg_inf) + return logits + + +def decode_frames( + first: torch.Tensor, + frames: int, + controls: FrameControls, + *, + embed: Callable[[torch.Tensor], torch.Tensor], + forward: Callable[[int, torch.Tensor], torch.Tensor], + codec_logits: Callable[[torch.Tensor], torch.Tensor], + sample: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + """Run frames 1..frames-1 after frame 0 sampled ``first`` ((B,) long). + + Returns ``(sampled, emitted)``, both ``(B, frames)``; a frame is emitted + until a stop id or the token budget, where ``check_stop`` would end it. + Nothing here reads the device. + """ + c = controls + sampled = [first] + emitted = [torch.ones_like(first, dtype=torch.bool)] + window = c.window + prev = first + for k in range(1, frames): + stopped = (prev.unsqueeze(1) == c.stop_ids).any(dim=1) + emitted.append(emitted[-1] & ~stopped & c.allowed[k]) + window = torch.cat([window[:, 1:], prev.unsqueeze(1)], dim=1) + logits = adjust_logits(codec_logits(forward(k, embed(prev))), k, window, c) + prev = torch.where(c.force_eos[k], c.eos_token_id, sample(logits).to(torch.long)) + sampled.append(prev) + return torch.stack(sampled, dim=1), torch.stack(emitted, dim=1) + + +def frame_sampling_metadata(md: Any) -> Any: + """Frames 1..K-1 metadata: min_tokens moves to ``adjust_logits`` and vLLM's + penalty pass is dropped (the Talker scores its own window; eligibility + requires zero frequency/presence penalties).""" + procs = LogitsProcessors(p for p in md.logitsprocs.all if not isinstance(p, MinTokensLogitsProcessor)) + return replace(md, no_penalties=True, logitsprocs=procs) + + +def sample_frame(logits: torch.Tensor, md: Any) -> torch.Tensor: + return _codec_sampler()(logits, md).sampled_token_ids[:, 0] + + +def frame_budgets(runner: Any, model_kwargs_extra: dict[str, Any], num_reqs: int) -> tuple[list[int], Decline | None]: + """Frames each request may still emit (``max_tokens``, else the context).""" + given = model_kwargs_extra.get("request_max_tokens_remaining") + if given is None or len(given) != num_reqs: + got = "none" if given is None else len(given) + return [], Decline("no per-request token budget", f"{got} budgets for {num_reqs} requests") + budgets: list[int] = [] + for req_id, budget in zip(runner.input_batch.req_ids[:num_reqs], given): + if budget is None: + req_state = runner.requests.get(req_id) + num_tokens = getattr(req_state, "num_tokens", None) + if num_tokens is None: + return [], Decline("no per-request token budget", f"request {req_id} has no max_tokens or state") + budget = max(int(runner.max_model_len) - int(num_tokens), 0) + budgets.append(int(budget)) + return budgets, None + + +def _hidden(output: Any) -> torch.Tensor: + if isinstance(output, OmniOutput): + return output.text_hidden_states + if isinstance(output, tuple): + return output[0] + return output + + +def maybe_run( + runner: Any, + model_output: Any, + *, + run_model: Callable[[], Any], + inputs_embeds: torch.Tensor | None, + model_kwargs_extra: dict[str, Any], +) -> Any: + """Run a Talker K-frame step on CUDA; any other step returns unchanged.""" + setattr(runner, _STASH_ATTR, None) + model: Any = getattr(runner, "model", None) + if not supports(model) or drafts_this_step(runner) <= 1: + return model_output + frames = applies(model, model_kwargs_extra, require_sampled_frame=False) + if frames <= 1: + if is_multi_token_decode(model, model_kwargs_extra): + spans = [(int(s), int(e)) for s, e in model_kwargs_extra.get("request_token_spans") or ()] + raise RuntimeError( + "MiniCPM-o scheduled a multi-token Talker decode step the CUDA " + f"multi-frame loop cannot run (request token spans {spans}); " + "see the '[minicpmo]' log for the reason" + ) + return model_output + if not isinstance(model_output, OmniOutput) or inputs_embeds is None: + raise RuntimeError("MiniCPM-o CUDA multi-frame decode needs the Talker's OmniOutput and inputs_embeds") + + spans = [(int(s), int(e)) for s, e in model_kwargs_extra["request_token_spans"]] + infos = model_kwargs_extra["model_intermediate_buffer"] + num_reqs = len(spans) + device = inputs_embeds.device + hidden = _hidden(model_output) + rows_host = torch.tensor([[s + k for s, _ in spans] for k in range(frames)], dtype=torch.long) + + # Frame 0: the single-frame step, verbatim. + rows0 = rows_host[0].to(device, non_blocking=True) + logits0 = model.compute_logits(hidden.index_select(0, rows0)) + md = runner.input_batch.sampling_metadata + vocab = int(logits0.shape[-1]) + sanitize_min_tokens_stop_ids(md.logitsprocs, vocab) + first = runner._sample(logits0, None).sampled_token_ids[:, 0].to(torch.long) + + eos_id = int(model.codec_eos_token_id) + decline = vocab_decline(runner, vocab) or ineligible_reason(runner, eos_id) + budgets: list[int] = [] + if decline is None: + budgets, decline = frame_budgets(runner, model_kwargs_extra, num_reqs) + if decline is not None: + _log_once("one frame this step", decline) + # Frame 0 is the whole step, as (B, 1): the one-token bookkeeping branch + # skips parse_output's vocab filter and the scheduler rejects the drafts. + # An empty row would roll the span back and re-send the codec id. + output = SamplerOutput(sampled_token_ids=first.to(torch.int32).unsqueeze(1), logprobs_tensors=None) + setattr(runner, _STASH_ATTR, output) + return model_output + + plan = model.plan_codec_frames(infos, frames) + req_ids = runner.input_batch.req_ids[:num_reqs] + min_state: dict[int, tuple[Any, ...]] = {} + for proc in md.logitsprocs.non_argmax_invariant: + if isinstance(proc, MinTokensLogitsProcessor): + min_state.update(proc.min_toks) + penalties = None if md.no_penalties else md.repetition_penalties[:num_reqs] + controls, (rows,) = build_controls( + plan, + frames=frames, + vocab_size=vocab, + budgets=budgets, + stop_ids=[_stop_ids(runner.requests[r].sampling_params) for r in req_ids], + min_tokens_state=min_state, + penalties=penalties, + device=device, + extra=[rows_host], + ) + frame_md = frame_sampling_metadata(md) + last = [model_output] + + def forward(k: int, embeds: torch.Tensor) -> torch.Tensor: + inputs_embeds.index_copy_(0, rows[k], embeds.to(inputs_embeds.dtype)) + last[0] = run_model() + return _hidden(last[0]).index_select(0, rows[k]) + + def replay() -> torch.Tensor: + last[0] = run_model() + return _hidden(last[0]) + + # Opt-in model hook: the model may run frames 1..K-1 itself (MiniCPM-o + # replays captured per-frame tails) and must then return exactly what + # decode_frames would; None, or no hook, keeps the loop below. + model_frames = getattr(model, "kstep_decode_frames", None) + graphed = None + if callable(model_frames): + graphed = model_frames( + first=first, + frames=frames, + controls=controls, + rows=rows, + inputs_embeds=inputs_embeds, + hidden=hidden, + replay=replay, + embed=model.embed_input_ids, + codec_logits=model.compute_logits, + sampler=_codec_sampler(), + sampling_metadata=frame_md, + ) + if graphed is not None: + sampled, emitted = graphed + else: + sampled, emitted = decode_frames( + first, + frames, + controls, + embed=model.embed_input_ids, + forward=forward, + codec_logits=model.compute_logits, + sample=lambda logits: sample_frame(logits, frame_md), + ) + ids = torch.where(emitted, sampled, -1) + ids_host = ids.cpu().tolist() # the step's one device read + forwarded = [] + for row in ids_host: + n = sum(1 for t in row if t >= 0) + forwarded.append(row[: n - 1]) + flags = model.commit_codec_frames(plan, forwarded) + + multimodal: Any = model_output.multimodal_outputs + audio = multimodal["codes"]["audio"] + finished = multimodal["meta"]["finished"] + for index, codes in enumerate(forwarded): + if not codes: + continue + tail = sampled[index, : len(codes)].reshape(-1, 1).to(device=audio[index].device) + audio[index] = torch.cat([audio[index].to(torch.long), tail], dim=0) + if flags[index]: + finished[index] = torch.tensor(True, dtype=torch.bool) + setattr(runner, _STASH_ATTR, SamplerOutput(sampled_token_ids=ids.to(torch.int32), logprobs_tensors=None)) + return model_output._replace(text_hidden_states=_hidden(last[0])) From d6f62a9015b78ebca68df2080ea4295b7622980a Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:30 +0800 Subject: [PATCH 78/89] [Perf][MiniCPM-o] Replay Talker K-step per-frame sampling tail as CUDA graphs Replay Talker K-step per-frame sampling tail as CUDA graphs Capture the per-frame logits adjustment and sampling tail into private CUDA graphs interleaved between Talker forward replays, eliminating per-frame host launch overhead. Signed-off-by: BeatSeat --- .../minicpmo_4_5/test_talker_kstep_graph.py | 481 ++++++++++++++ .../models/minicpmo_4_5/talker_kstep_graph.py | 595 ++++++++++++++++++ 2 files changed, 1076 insertions(+) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_talker_kstep_graph.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/talker_kstep_graph.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_graph.py b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_graph.py new file mode 100644 index 00000000000..640802f8b2b --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_kstep_graph.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Talker K-step frame graphs (``talker_kstep_graph``) vs. the eager ``decode_frames``. + +On CPU the captured tail is emulated: either the tail functions run eagerly +(no graph factory), or a fake factory records each tail and replays it by +calling it again, which is exactly what a CUDA graph does with static buffers. +Both must give ``decode_frames``'s ids, emit flags, final generator states and +final ``inputs_embeds`` bit for bit. The real capture is +``test_talker_kstep_graph_cuda.py``. +""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +from vllm.v1.sample.logits_processor import LogitsProcessors +from vllm.v1.sample.logits_processor.builtin import LogitBiasLogitsProcessor, MinPLogitsProcessor +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.models.minicpmo_4_5 import talker_kstep_graph as kg +from vllm_omni.model_executor.models.minicpmo_4_5.talker_frame_plan import CodecFramePlan +from vllm_omni.worker import gpu_talker_multiframe as mf + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +VOCAB, HIDDEN, EOS, WINDOW = 16, 8, 15, 16 +PAD_ROWS = 3 # the talker graph's padded token count exceeds B * K + + +def native_sampler() -> Sampler: + """vLLM's sampler on its PyTorch top-k/top-p path (what CUDA runs with seeded rows).""" + sampler = Sampler() + sampler.topk_topp_sampler.forward = sampler.topk_topp_sampler.forward_native + return sampler + + +class Toy: + """Codec embedding/head plus a causal 'talker graph' with static buffers. + + ``replay()`` recomputes every row of every request span from + ``inputs_embeds`` and writes the static output in place, like a FULL-graph + replay; ``fresh_at`` makes the given replays return a new tensor instead. + """ + + def __init__(self, batch: int, frames: int, seed: int = 0): + gen = torch.Generator().manual_seed(seed) + self.batch, self.frames = batch, frames + self.emb = nn.Embedding(VOCAB, HIDDEN).requires_grad_(False) + self.head = nn.Linear(HIDDEN, VOCAB).requires_grad_(False) + with torch.no_grad(): + self.emb.weight.copy_(torch.randn(VOCAB, HIDDEN, generator=gen)) + self.head.weight.copy_(torch.randn(VOCAB, HIDDEN, generator=gen)) + self.head.bias.copy_(torch.randn(VOCAB, generator=gen)) + self.head.bias[EOS] = 1.5 + self.proj = torch.randn(HIDDEN, HIDDEN, generator=gen) + self.prefix = torch.randn(batch, HIDDEN, generator=gen) + self.inputs_embeds = torch.randn(batch * frames + PAD_ROWS, HIDDEN, generator=gen) + self.out = torch.zeros(batch * frames + PAD_ROWS, HIDDEN) + self.replays = 0 + self.fresh_at: set[int] = set() + + def replay(self) -> torch.Tensor: + self.replays += 1 + hidden = torch.zeros_like(self.out) + for i in range(self.batch): + state = self.prefix[i] + for j in range(self.frames): + row = i * self.frames + j + state = 0.7 * state + self.inputs_embeds[row] + hidden[row] = torch.tanh(state @ self.proj) * 3.0 + if self.replays in self.fresh_at: + return hidden + self.out.copy_(hidden) + return self.out + + def embed(self, ids: torch.Tensor) -> torch.Tensor: + return self.emb(ids) + + def codec_logits(self, hidden: torch.Tensor) -> torch.Tensor: + return self.head(hidden).float() + + +def metadata(batch: int, mode: str) -> SamplingMetadata: + """Frame metadata (``frame_sampling_metadata`` shape: no penalties, no min_tokens).""" + greedy = mode == "greedy" + if mode == "seeded": + generators = {i: torch.Generator().manual_seed(100 + i) for i in range(batch)} + elif mode == "mixed": + generators = {i: torch.Generator().manual_seed(100 + i) for i in range(0, batch, 2)} + else: + generators = {} + return SamplingMetadata( + temperature=None if greedy else torch.full((batch,), 0.8), + all_greedy=greedy, + all_random=not greedy, + top_p=None if greedy else torch.full((batch,), 0.85), + top_k=None if greedy else torch.full((batch,), 6, dtype=torch.int32), + generators=generators, + max_num_logprobs=None, + no_penalties=True, + prompt_token_ids=None, + frequency_penalties=torch.zeros(batch), + presence_penalties=torch.zeros(batch), + repetition_penalties=torch.ones(batch), + output_token_ids=[], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ) + + +def controls_for(batch: int, frames: int, step: int, scenario: str): + """One step's ``build_controls`` output; ``full`` exercises every control.""" + gen = torch.Generator().manual_seed(1000 + step) + force = [[False] * batch for _ in range(frames)] + mask = [[False] * batch for _ in range(frames)] + budgets = [frames + 5] * batch + min_state: dict[int, tuple] = {} + stop_ids = [{EOS} for _ in range(batch)] + penalties = None + if scenario == "full": + for k in range(1, frames): + mask[k][0] = k < frames - 2 # the Talker's own min_tokens holds EOS back + if batch > 1 and k == frames - 1 - (step % 2): + force[k][1] = True # chunk limit: forced EOS mid-loop + budgets[-1] = max(1, frames - 2 - step % 3) # vLLM max_tokens truncates mid-loop + min_state = {0: (frames + 3, [1, 2], {EOS, 4}, False)} + stop_ids[0] = {EOS, 7} + penalties = torch.full((batch,), 1.3) + recent = [torch.randint(0, VOCAB - 1, (int(n),), generator=gen).tolist() for n in (0, 3, 20, 9)[:batch]] + recent += [[] for _ in range(batch - len(recent))] + plan = CodecFramePlan( + request_ids=[f"r{i}" for i in range(batch)], + eos_token_id=EOS, + force_eos=force, + mask_eos=mask, + recent_codes=recent, + ) + rows_host = torch.tensor([[i * frames + k for i in range(batch)] for k in range(frames)], dtype=torch.long) + controls, (rows,) = mf.build_controls( + plan, + frames=frames, + vocab_size=VOCAB, + budgets=budgets, + stop_ids=stop_ids, + min_tokens_state=min_state, + penalties=penalties, + device=torch.device("cpu"), + extra=[rows_host], + ) + first = torch.randint(0, VOCAB - 1, (batch,), generator=gen) + return controls, rows, first + + +def eager_frames(toy: Toy, first, frames, controls, rows, sampler, md): + """``maybe_run``'s eager loop, verbatim.""" + + def forward(k: int, embeds: torch.Tensor) -> torch.Tensor: + toy.inputs_embeds.index_copy_(0, rows[k], embeds.to(toy.inputs_embeds.dtype)) + return toy.replay().index_select(0, rows[k]) + + return mf.decode_frames( + first, + frames, + controls, + embed=toy.embed, + forward=forward, + codec_logits=toy.codec_logits, + sample=lambda logits: sampler(logits, md).sampled_token_ids[:, 0], + ) + + +def graph_frames(graphs, toy: Toy, first, frames, controls, rows, sampler, md): + return graphs( + first=first, + frames=frames, + controls=controls, + rows=rows, + inputs_embeds=toy.inputs_embeds, + hidden=toy.out, + replay=toy.replay, + embed=toy.embed, + codec_logits=toy.codec_logits, + sampler=sampler, + sampling_metadata=md, + ) + + +class FakeGraph: + """Records a tail; ``replay()`` reruns it, as a CUDA graph reruns its kernels.""" + + made: list[FakeGraph] = [] + + def __init__(self, fn, generators): + self.fn, self.generators, self.replays = fn, list(generators), 0 + FakeGraph.made.append(self) + + def replay(self) -> None: + self.replays += 1 + self.fn() + + +@pytest.fixture(autouse=True) +def _fresh_fake_graphs(): + FakeGraph.made = [] + yield + + +def _states(md: SamplingMetadata) -> dict[int, torch.Tensor]: + return {i: g.get_state() for i, g in md.generators.items()} + + +def _run_pair(batch, frames, mode, scenario, steps, factory, fresh_at=()): + """Drive both paths over ``steps`` steps from identical state; assert equality.""" + ref, new = Toy(batch, frames), Toy(batch, frames) + new.fresh_at = set(fresh_at) + md_ref, md_new = metadata(batch, mode), metadata(batch, mode) + sampler = native_sampler() + graphs = kg.TalkerKStepFrameGraphs(graph_factory=factory) + results = [] + for step in range(steps): + controls, rows, first = controls_for(batch, frames, step, scenario) + torch.manual_seed(5000 + step) + want = eager_frames(ref, first, frames, controls, rows, sampler, md_ref) + rng_ref = torch.get_rng_state() + torch.manual_seed(5000 + step) + got = graph_frames(graphs, new, first, frames, controls, rows, sampler, md_new) + results.append(got) + if got is None: + continue + # Unseeded rows draw from the default generator: same draws, same end state. + assert torch.equal(torch.get_rng_state(), rng_ref), f"step {step}: default generator state" + assert torch.equal(got[0], want[0]), f"step {step}: sampled ids differ" + assert torch.equal(got[1], want[1]), f"step {step}: emit flags differ" + assert torch.equal(new.inputs_embeds, ref.inputs_embeds), f"step {step}: inputs_embeds differ" + assert torch.equal(new.out, ref.out) + for i, state in _states(md_ref).items(): + assert torch.equal(md_new.generators[i].get_state(), state), f"step {step}: generator {i} state" + return graphs, results, (ref, new) + + +@pytest.mark.parametrize("frames", [2, 3, 8]) +@pytest.mark.parametrize("batch", [1, 3]) +@pytest.mark.parametrize("mode", ["greedy", "seeded", "mixed", "unseeded"]) +@pytest.mark.parametrize("scenario", ["plain", "full"]) +def test_eager_tail_matches_decode_frames(scenario, mode, batch, frames): + """No graph factory on CPU: the tail functions run eagerly, every step.""" + graphs, results, _ = _run_pair(batch, frames, mode, scenario, steps=3, factory=None) + assert all(r is not None for r in results) + assert graphs.stats.eager_steps == 3 and graphs.stats.graph_steps == 0 + assert not FakeGraph.made + + +@pytest.mark.parametrize("frames", [2, 3, 8]) +@pytest.mark.parametrize("batch", [1, 4]) +@pytest.mark.parametrize("mode", ["greedy", "seeded", "mixed"]) +def test_replayed_tails_match_decode_frames(mode, batch, frames): + """First step eager, then one capture (first/mid/last), then replays only.""" + graphs, results, (_, new) = _run_pair(batch, frames, mode, "full", steps=4, factory=FakeGraph) + assert all(r is not None for r in results) + assert graphs.stats.eager_steps == 1 and graphs.stats.graph_steps == 3 + assert graphs.stats.captures == 1 + names = 3 if frames > 2 else 2 + assert len(FakeGraph.made) == names + first, *rest = FakeGraph.made + assert first.replays == 3 and not first.generators + # mid replays frames - 2 times a step, last once; both own the seeded rows' slots. + assert sum(g.replays for g in rest) == 3 * (frames - 1) + seeded = len(metadata(batch, mode).generators) + assert all(len(g.generators) == seeded for g in rest) + assert new.replays == 4 * (frames - 1) + + +def test_a_moved_forward_buffer_finishes_eagerly_then_declines(): + """A replay that returns a new tensor (not the captured output): that step + finishes on what it returned, identically; the shape then stays eager.""" + batch, frames = 2, 8 + # Step 0 eager + capture; step 1 replays; the 3rd forward of step 2 moves. + fresh = {2 * (frames - 1) + 3} + graphs, results, _ = _run_pair(batch, frames, "seeded", "full", steps=4, factory=FakeGraph, fresh_at=fresh) + assert [r is None for r in results] == [False, False, False, True] + assert graphs.stats.unstable == 1 + + +def test_capture_failure_keeps_the_eager_tail(): + def broken(fn, generators): + raise RuntimeError("operation not permitted when stream is capturing") + + graphs, results, _ = _run_pair(1, 4, "seeded", "full", steps=3, factory=broken) + assert results[0] is not None and results[1] is None and results[2] is None + assert graphs.stats.captures == 0 + + +def test_capture_does_not_move_generators(): + """Capturing records the RNG ops without running them; a capture that does + consume randoms must still leave every slot where eager sampling expects it.""" + + def draws_while_capturing(fn, generators): + for gen in generators: + torch.rand(4, generator=gen) + torch.rand(4) # the default generator, which unseeded rows draw from + return FakeGraph(fn, generators) + + graphs, results, _ = _run_pair(3, 4, "mixed", "plain", steps=3, factory=draws_while_capturing) + assert all(r is not None for r in results) and graphs.stats.graph_steps == 2 + + +def test_shapes_are_keyed_by_rows_and_frames(): + graphs = kg.TalkerKStepFrameGraphs(graph_factory=FakeGraph) + sampler = native_sampler() + for batch in (1, 2, 1, 2): + toy = Toy(batch, 4) + md = metadata(batch, "seeded") + controls, rows, first = controls_for(batch, 4, 0, "plain") + assert graph_frames(graphs, toy, first, 4, controls, rows, sampler, md) is not None + # Each toy has its own buffers, so every call is a new shape: eager + capture. + assert graphs.stats.captures == 4 and graphs.stats.graph_steps == 0 + + +# --------------------------------------------------------------------------- +# Declines: what a graph cannot reproduce goes back to decode_frames. +# --------------------------------------------------------------------------- + + +def _controls(width: int = 1): + return SimpleNamespace(stop_ids=torch.zeros((1, width), dtype=torch.long)) + + +def test_flashinfer_sampling_without_generators_declines(): + md = metadata(2, "unseeded") + flashinfer = SimpleNamespace(topk_topp_sampler=SimpleNamespace(forward=lambda: None, use_fp64_gumbel=False)) + flashinfer.topk_topp_sampler.forward.__name__ = "forward_cuda" + assert "FlashInfer" in kg.graph_decline(flashinfer, md, _controls()) + # forward_cuda takes forward_native with per-request generators, or with no top-k/top-p. + assert kg.graph_decline(flashinfer, metadata(2, "mixed"), _controls()) is None + assert kg.graph_decline(flashinfer, replace(md, top_k=None, top_p=None), _controls()) is None + assert kg.graph_decline(flashinfer, metadata(2, "greedy"), _controls()) is None + assert kg.graph_decline(native_sampler(), md, _controls()) is None + cpu = Sampler() + if cpu.topk_topp_sampler.forward.__name__ == "forward_cpu": + assert kg.graph_decline(cpu, md, _controls()) is not None + + +def test_active_min_p_or_logit_bias_and_wide_stop_sets_decline(): + md = metadata(2, "seeded") + min_p = MinPLogitsProcessor.__new__(MinPLogitsProcessor) + min_p.min_p_count = 1 + bias = LogitBiasLogitsProcessor.__new__(LogitBiasLogitsProcessor) + bias.biases = {0: {3: 1.0}} + assert "MinPLogitsProcessor" in kg.graph_decline( + native_sampler(), replace(md, logitsprocs=LogitsProcessors([min_p])), _controls() + ) + assert "LogitBias" in kg.graph_decline( + native_sampler(), replace(md, logitsprocs=LogitsProcessors([bias])), _controls() + ) + min_p.min_p_count = 0 + bias.biases = {} + idle = replace(md, logitsprocs=LogitsProcessors([min_p, bias])) + assert kg.graph_decline(native_sampler(), idle, _controls()) is None + assert "stop ids" in kg.graph_decline(native_sampler(), md, _controls(kg.STOP_WIDTH + 1)) + assert kg.graph_decline(native_sampler(), md, _controls(kg.STOP_WIDTH)) is None + + +def test_a_non_full_talker_forward_declines(monkeypatch): + monkeypatch.setattr(kg, "_forward_graph_mode", lambda: "PIECEWISE") + toy, md = Toy(1, 4), metadata(1, "seeded") + controls, rows, first = controls_for(1, 4, 0, "plain") + graphs = kg.TalkerKStepFrameGraphs(graph_factory=FakeGraph) + assert graph_frames(graphs, toy, first, 4, controls, rows, native_sampler(), md) is None + assert toy.replays == 0 and "PIECEWISE" in next(iter(graphs.stats.declined)) + + +# --------------------------------------------------------------------------- +# Wiring: the MiniCPM-o wrapper hook, maybe_run's opt-in hook, the overlay. +# --------------------------------------------------------------------------- + + +def test_wrapper_hook_is_off_by_default(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + wrapper = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + nn.Module.__init__(wrapper) + assert wrapper.kstep_decode_frames(first=None) is None + calls = [] + wrapper._kstep_frame_graphs = lambda **kw: calls.append(kw) or "ran" + assert wrapper.kstep_decode_frames(first=1) == "ran" and calls == [{"first": 1}] + + +@pytest.mark.parametrize("greedy", [True, False], ids=["greedy", "seeded"]) +def test_maybe_run_with_the_hook_matches_single_frame(monkeypatch, greedy: bool): + """The real maybe_run / scheduler emulation of test_gpu_talker_multiframe with the + hook installed (the harness's forward allocates a new output every call, so the + hook finishes each step eagerly and then declines; the stable case is above).""" + from tests.worker import test_gpu_talker_multiframe as base + + monkeypatch.setattr(mf, "_SAMPLER", native_sampler()) + results: list[bool] = [] + + class Hooked(base._Model): + def __init__(self, talker): + super().__init__(talker) + self.graphs = kg.TalkerKStepFrameGraphs(graph_factory=FakeGraph) + + def kstep_decode_frames(self, **kwargs): + out = self.graphs(**kwargs) + results.append(out is not None) + return out + + def prepare(runner): + runner.model = Hooked(runner.model.talker) + + ref = base._single_frame(greedy) + got, widths = base._engine(greedy, 4, batch_vocab=base.VOCAB, prepare=prepare) + base._assert_same(got, ref) + assert set(widths[1:]) == {4} + assert results and results[0] + # Each step's new output buffer is a new shape that finishes eagerly; after + # MAX_UNSTABLE_PER_SHAPE of them the (rows, frames) size stays on decode_frames. + assert sum(results) <= kg.MAX_UNSTABLE_PER_SHAPE and not any(results[kg.MAX_UNSTABLE_PER_SHAPE :]) + + +def test_maybe_run_hook_returning_none_keeps_decode_frames(monkeypatch): + from tests.worker import test_gpu_talker_multiframe as base + + seen: list[dict] = [] + + class Declines(base._Model): + def kstep_decode_frames(self, **kwargs): + seen.append(kwargs) + return None + + ref = base._single_frame(False) + got, _ = base._engine( + False, 4, batch_vocab=base.VOCAB, prepare=lambda r: setattr(r, "model", Declines(r.model.talker)) + ) + base._assert_same(got, ref) + assert seen and set(seen[0]) == { + "first", + "frames", + "controls", + "rows", + "inputs_embeds", + "hidden", + "replay", + "embed", + "codec_logits", + "sampler", + "sampling_metadata", + } + assert seen[0]["sampler"] is mf._codec_sampler() + + +def test_s1graph_overlay_adds_only_the_stage1_switch(): + import vllm_omni + from vllm_omni.config.stage_config import resolve_deploy_yaml + + deploy = Path(vllm_omni.__file__).parent / "deploy" + raw = resolve_deploy_yaml(deploy / "minicpmo_4_5_dxsched_on_s1graph.yaml") + base = resolve_deploy_yaml(deploy / "minicpmo_4_5_dxsched_on.yaml") + stage1 = next(s for s in raw["stages"] if s["stage_id"] == 1) + base1 = next(s for s in base["stages"] if s["stage_id"] == 1) + assert stage1["hf_overrides"] == {**base1.get("hf_overrides", {}), "talker_kstep_graph_sampling": True} + assert "talker_kstep_graph_sampling" not in base1.get("hf_overrides", {}) + strip = {**stage1} + strip.pop("hf_overrides") + base_strip = {**base1} + base_strip.pop("hf_overrides", None) + assert strip == base_strip + assert [s for s in raw["stages"] if s["stage_id"] != 1] == [s for s in base["stages"] if s["stage_id"] != 1] + assert {k: v for k, v in raw.items() if k not in ("stages", "base_config")} == { + k: v for k, v in base.items() if k not in ("stages", "base_config") + } diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/talker_kstep_graph.py b/vllm_omni/model_executor/models/minicpmo_4_5/talker_kstep_graph.py new file mode 100644 index 00000000000..d79131faae9 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/talker_kstep_graph.py @@ -0,0 +1,595 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA graphs for the per-frame tail of the MiniCPM-o Talker K-step decode. + +``gpu_talker_multiframe.decode_frames`` runs frames 1..K-1 of a K-step step as +one talker FULL-graph replay per frame plus an eager tail of ~40 small kernels: +the codec head, the EOS force/mask and windowed repetition penalty, vLLM's +``Sampler`` (temperature, Triton top-k/top-p, exponential race), the stop/emit +bookkeeping and the next frame's input embedding. On H100 that tail keeps the +GPU busy for ~0.15 ms but costs ~1.7 ms of host time per frame. + +Opt-in (stage-1 ``hf_overrides: {talker_kstep_graph_sampling: true}``), this +module captures the tail once per step shape and replays it between the talker +graph replays (``R k`` = the runner's own replay of frame k):: + + first (in 1) -> R 1 -> mid (out 1, in 2) -> R 2 -> ... -> R K-1 -> last (out K-1) + +``in k`` writes the id sampled at frame k-1 into frame k's ``inputs_embeds`` +row; ``out k`` reads frame k's row of the talker graph's static output and +samples. The frame index lives on the device, so one ``mid`` graph serves every +frame. The talker forward itself stays vLLM's FULL graph: a graph cannot be +launched inside another capture, and capturing the forward here would bake in +this step's attention metadata. + +Lossless (see ``notes/s1graph.md``): the captured ops are ``decode_frames``'s +ops in the same order on the same values. Frame k's controls are selected with +a device index instead of a Python one; the repetition penalty and the +min_tokens mask always run, with exact-identity values (penalty 1.0, an +all-false mask) on steps that have none. Seeded rows sample from per-row +Philox *slot* generators registered with the graphs: before the frames each +slot takes its request generator's (seed, offset), afterwards the request +generator takes the slot's back, so every frame consumes the Philox counters +the eager sampler would have. Unseeded rows draw from the default generator, +which every capture registers. The sampled ids are therefore bit-identical to +the eager K-step path; anything the graphs cannot reproduce declines to it. +""" + +from __future__ import annotations + +import os +import time +from collections import OrderedDict +from collections.abc import Callable +from dataclasses import dataclass, field, replace +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +#: Stop ids a row may carry in the captured tail (-1 padded); wider rows decline. +STOP_WIDTH = 8 +#: Captured step shapes kept (least recently used dropped); each holds three small graphs. +MAX_SHAPES = 32 +#: Captures one (rows, frames) shape may take; more means its buffers keep moving. +MAX_CAPTURES_PER_SHAPE = 4 +#: Steps of one (rows, frames) shape whose forward returned a fresh buffer before it stays eager. +MAX_UNSTABLE_PER_SHAPE = 2 + +_LOGGED: set[str] = set() + +#: ``(fn, generators) -> graph`` with ``graph.replay()`` rerunning ``fn``. +GraphFactory = Callable[[Callable[[], None], list[torch.Generator]], Any] + + +def _log_once(kind: str, message: str, *args: Any) -> None: + if kind in _LOGGED: + return + _LOGGED.add(kind) + logger.info("[minicpmo] Talker K-step frame graphs: " + message, *args) + + +def _ptr(tensor: torch.Tensor | None) -> tuple[int, tuple[int, ...]] | None: + return None if tensor is None else (tensor.data_ptr(), tuple(tensor.shape)) + + +def _native_sampling(sampler: Any, md: Any) -> bool: + """Whether vLLM's sampler takes its PyTorch path for this batch. + + That path (Triton top-k/top-p masks, softmax, an exponential race on the + request or default generator) is capturable. FlashInfer's sampler reads the + generator state on the host, so a replay would reuse frozen randoms. + """ + if md.all_greedy: + return True + topk_topp = getattr(sampler, "topk_topp_sampler", None) + forward = getattr(getattr(topk_topp, "forward", None), "__name__", "") + if forward == "forward_native": + return True + if forward == "forward_cuda": + # forward_cuda falls back to forward_native for these. + return ( + bool(md.generators) + or (md.top_k is None and md.top_p is None) + or bool(getattr(topk_topp, "use_fp64_gumbel", False)) + ) + return False + + +def graph_decline(sampler: Any, md: Any, controls: Any) -> str | None: + """Why this step's frames cannot replay captured tails (they run eagerly), or None. + + ``gpu_talker_multiframe.ineligible_reason`` already declined logprobs, bad + words, allowed ids, custom logits processors and penalties other than the + Talker's own; this only adds what a graph cannot reproduce. + """ + if not _native_sampling(sampler, md): + return "FlashInfer sampling without per-request generators (host-side RNG state)" + for proc in md.logitsprocs.all: + if getattr(proc, "min_p_count", 0) or getattr(proc, "biases", None): + return f"active {type(proc).__name__} (a Python branch per step)" + if int(controls.stop_ids.shape[1]) > STOP_WIDTH: + return f"{int(controls.stop_ids.shape[1])} stop ids per row (captured width {STOP_WIDTH})" + return None + + +def _forward_graph_mode() -> str | None: + """The forward context's cudagraph mode name, or None outside a forward context.""" + try: + from vllm.forward_context import get_forward_context, is_forward_context_available + except ImportError: # pragma: no cover - vLLM without a forward context + return None + if not is_forward_context_available(): + return None + mode = getattr(get_forward_context(), "cudagraph_runtime_mode", None) + return getattr(mode, "name", None) + + +def _triton_sampler_buffers() -> list[Any]: + """vLLM's cached Triton top-k/top-p scratch buffers, kept alive by each capture. + + A graph bakes their addresses in; if vLLM ever replaced or cleared its cache, + these references keep the captured scratch memory from being reused. + """ + try: + from vllm.v1.sample.ops import topk_topp_triton as ops + except ImportError: # pragma: no cover - vLLM without the Triton sampler + return [] + held: list[Any] = [] + for name in ("_TRITON_BUFFER_CACHE", "_TRITON_TABLE_CACHE", "_TRITON_SPLIT_CACHE"): + cache = getattr(ops, name, None) + if isinstance(cache, dict): + held.extend(cache.values()) + return held + + +class FrameTail: + """Static buffers of one step shape and the per-frame functions over them. + + Every function reads and writes only these buffers (plus the talker's + ``inputs_embeds`` and output rows), so running it eagerly and replaying its + capture give the same results. Rows are ``(B,)``, frames ``(K, B)``. + """ + + def __init__( + self, + *, + batch: int, + frames: int, + vocab: int, + window: int, + eos: int, + hidden: torch.Tensor, + inputs_embeds: torch.Tensor, + ) -> None: + device = hidden.device + long, flag = torch.long, torch.bool + self.batch, self.frames, self.vocab, self.eos = batch, frames, vocab, int(eos) + self.hidden = hidden + self.inputs_embeds = inputs_embeds + self.k = torch.ones(1, dtype=long, device=device) + self.prev = torch.zeros(batch, dtype=long, device=device) + self.alive = torch.ones(batch, dtype=flag, device=device) + self.sampled = torch.zeros((batch, frames), dtype=long, device=device) + self.emitted = torch.ones((batch, frames), dtype=flag, device=device) + self.window = torch.full((batch, window), vocab, dtype=long, device=device) + self.force_eos = torch.zeros((frames, batch), dtype=flag, device=device) + self.mask_eos = torch.zeros((frames, batch), dtype=flag, device=device) + self.allowed = torch.zeros((frames, batch), dtype=flag, device=device) + self.min_tokens = torch.zeros((frames, batch), dtype=flag, device=device) + self.min_tokens_stop = torch.zeros((batch, vocab), dtype=flag, device=device) + self.stop_ids = torch.full((batch, STOP_WIDTH), -1, dtype=long, device=device) + self.penalties = torch.ones(batch, dtype=torch.float32, device=device) + self.rows = torch.zeros((frames, batch), dtype=long, device=device) + # build_controls' constants, for this shape's EOS. + self.eos_col = torch.arange(vocab, device=device) == self.eos + self.forced_row = torch.where( + self.eos_col, + torch.zeros((), dtype=torch.float32, device=device), + torch.full((), float("-inf"), dtype=torch.float32, device=device), + ) + + def load(self, first: torch.Tensor, controls: Any, rows: torch.Tensor) -> None: + """Stage one step: frame 0's ids and the step's controls.""" + self.k.fill_(1) + self.prev.copy_(first) + self.alive.fill_(True) + self.sampled[:, 0].copy_(first) + self.window.copy_(controls.window) + self.force_eos.copy_(controls.force_eos) + self.mask_eos.copy_(controls.mask_eos) + self.allowed.copy_(controls.allowed) + self.min_tokens.copy_(controls.min_tokens) + if controls.min_tokens_stop is None: + self.min_tokens_stop.zero_() + else: + self.min_tokens_stop.copy_(controls.min_tokens_stop) + width = int(controls.stop_ids.shape[1]) + self.stop_ids.fill_(-1) + self.stop_ids[:, :width].copy_(controls.stop_ids) + if controls.penalties is None: + self.penalties.fill_(1.0) + else: + self.penalties.copy_(controls.penalties) + self.rows.copy_(rows) + + def _at_frame(self, per_frame: torch.Tensor) -> torch.Tensor: + return per_frame.index_select(0, self.k).squeeze(0) + + def frame_in(self, embed: Callable[[torch.Tensor], torch.Tensor]) -> None: + """``decode_frames`` up to the forward: emit flags, the window and frame k's input row.""" + prev = self.prev + stopped = (prev.unsqueeze(1) == self.stop_ids).any(dim=1) + alive = self.alive & ~stopped & self._at_frame(self.allowed) + self.alive.copy_(alive) + self.emitted.index_copy_(1, self.k, alive.unsqueeze(1)) + self.window.copy_(torch.cat([self.window[:, 1:], prev.unsqueeze(1)], dim=1)) + embeds = embed(prev) + self.inputs_embeds.index_copy_(0, self._at_frame(self.rows), embeds.to(self.inputs_embeds.dtype)) + + def adjust_logits(self, logits: torch.Tensor) -> torch.Tensor: + """``gpu_talker_multiframe.adjust_logits`` at the device frame index.""" + neg_inf = float("-inf") + logits = torch.where(self._at_frame(self.force_eos).unsqueeze(1), self.forced_row, logits) + logits = torch.where(self._at_frame(self.mask_eos).unsqueeze(1) & self.eos_col, neg_inf, logits) + num_reqs, vocab = logits.shape + counts = torch.zeros((num_reqs, vocab + 1), dtype=torch.long, device=logits.device) + counts.scatter_add_(1, self.window, torch.ones_like(self.window)) + alpha = torch.pow(self.penalties.to(logits.dtype).unsqueeze(1), counts[:, :vocab].to(logits.dtype)) + logits = torch.where(logits < 0, logits * alpha, logits / alpha) + return logits.masked_fill(self.min_tokens_stop & self._at_frame(self.min_tokens).unsqueeze(1), neg_inf) + + def frame_out( + self, + hidden: torch.Tensor, + codec_logits: Callable[[torch.Tensor], torch.Tensor], + sample: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + """``decode_frames`` after the forward: frame k's logits, sample and forced EOS.""" + logits = codec_logits(hidden.index_select(0, self._at_frame(self.rows))) + logits = self.adjust_logits(logits) + prev = torch.where(self._at_frame(self.force_eos), self.eos, sample(logits).to(torch.long)) + self.prev.copy_(prev) + self.sampled.index_copy_(1, self.k, prev.unsqueeze(1)) + + def next_frame(self) -> None: + self.k.add_(1) + + +@dataclass +class _Shape: + tail: FrameTail + graphs: dict[str, Any] | None = None + held: list[Any] = field(default_factory=list) + replays: int = 0 + + +@dataclass +class _Stats: + graph_steps: int = 0 + eager_steps: int = 0 + declined: dict[str, int] = field(default_factory=dict) + captures: int = 0 + capture_ms: float = 0.0 + unstable: int = 0 + + +class TalkerKStepFrameGraphs: + """Runs a K-step step's frames 1..K-1 with captured tails (see the module doc). + + ``__call__`` is ``gpu_talker_multiframe.maybe_run``'s ``kstep_decode_frames`` + hook: it returns ``decode_frames``'s ``(sampled, emitted)`` or None to let the + eager loop run the step. The first step of a shape runs the tail functions + eagerly (that is also the capture warm-up) and captures afterwards. + """ + + def __init__(self, *, max_shapes: int = MAX_SHAPES, graph_factory: GraphFactory | None = None) -> None: + """``graph_factory(fn, generators)`` returns an object whose ``replay()`` + reruns ``fn`` (tests); by default tails are CUDA-graph captured, on CUDA only.""" + self.max_shapes = int(max_shapes) + self._graph_factory = graph_factory + self._shapes: OrderedDict[tuple[Any, ...], _Shape] = OrderedDict() + self._captures_per_shape: dict[tuple[int, int], int] = {} + self._unstable_per_shape: dict[tuple[int, int], int] = {} + self._failed: set[tuple[Any, ...]] = set() + self._slots: list[torch.Generator] = [] + self._pool: Any = None + self._stream: Any = None + self.stats = _Stats() + self._stats_enabled = os.environ.get("VLLM_OMNI_TALKER_KSTEP_STATS") == "1" + + # -- public hook -------------------------------------------------------- + + def __call__( + self, + *, + first: torch.Tensor, + frames: int, + controls: Any, + rows: torch.Tensor, + inputs_embeds: torch.Tensor, + hidden: torch.Tensor, + replay: Callable[[], torch.Tensor], + embed: Callable[[torch.Tensor], torch.Tensor], + codec_logits: Callable[[torch.Tensor], torch.Tensor], + sampler: Any, + sampling_metadata: Any, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + md = sampling_metadata + decline = graph_decline(sampler, md, controls) + if decline is None: + mode = _forward_graph_mode() + if mode is not None and mode != "FULL": + decline = f"the talker forward runs {mode}, not one FULL graph" + if ( + decline is None + and first.device.type == "cuda" + and not hasattr(torch.cuda.CUDAGraph, "register_generator_state") + ): + decline = "torch.cuda.CUDAGraph has no register_generator_state" + batch = int(first.shape[0]) + vocab = int(controls.forced_row.shape[0]) + key = ( + batch, + int(frames), + vocab, + int(controls.eos_token_id), + _ptr(hidden), + hidden.dtype, + _ptr(inputs_embeds), + inputs_embeds.dtype, + bool(md.all_greedy), + bool(md.all_random), + _ptr(md.temperature), + _ptr(md.top_k), + _ptr(md.top_p), + tuple(sorted(int(i) for i in md.generators)), + ) + if decline is None and key in self._failed: + decline = "capture failed for this step shape" + if decline is not None: + self._count_decline(decline) + return None + shape = self._shape(key, batch, frames, vocab, controls, hidden, inputs_embeds) + if shape is None: + self._count_decline("capture budget for this step shape is spent (its buffers keep moving)") + return None + tail = shape.tail + seeded = {int(i): gen for i, gen in md.generators.items()} + slots = {i: self._slot(i, first.device) for i in seeded} + slot_md = replace(md, generators=slots) + + def sample(logits: torch.Tensor) -> torch.Tensor: + return sampler(logits, slot_md).sampled_token_ids[:, 0] + + tail.load(first, controls, rows) + for i, gen in seeded.items(): + slots[i].set_state(gen.get_state()) + try: + if shape.graphs is not None: + stable = self._replay(shape, frames, replay, embed, codec_logits, sample) + else: + stable = self._eager(tail, frames, replay, embed, codec_logits, sample) + finally: + # The request generators end where frames-1 eager samples leave them. + for i, gen in seeded.items(): + gen.set_state(slots[i].get_state()) + if shape.graphs is not None: + self.stats.graph_steps += 1 + else: + self.stats.eager_steps += 1 + if not stable: + self._unstable(key) + elif shape.graphs is None and (first.device.type == "cuda" or self._graph_factory is not None): + self._capture(key, shape, embed, codec_logits, sample, list(slots.values())) + self._maybe_log_stats() + return tail.sampled.clone(), tail.emitted.clone() + + # -- frames ------------------------------------------------------------- + + @staticmethod + def _same_output(out: torch.Tensor, tail: FrameTail) -> bool: + return out.data_ptr() == tail.hidden.data_ptr() and out.shape == tail.hidden.shape + + def _eager(self, tail: FrameTail, frames: int, replay, embed, codec_logits, sample) -> bool: + """The tail functions, eagerly; True if every forward returned the static output.""" + stable = True + tail.frame_in(embed) + for k in range(1, frames): + out = replay() + stable = stable and self._same_output(out, tail) + tail.frame_out(out, codec_logits, sample) + if k < frames - 1: + tail.next_frame() + tail.frame_in(embed) + return stable + + def _replay(self, shape: _Shape, frames: int, replay, embed, codec_logits, sample) -> bool: + tail, graphs = shape.tail, shape.graphs + assert graphs is not None + shape.replays += 1 + graphs["first"].replay() + for k in range(1, frames): + out = replay() + if not self._same_output(out, tail): + # The forward did not return the captured output buffer: finish + # this step with the same functions, eagerly, on what it returned. + tail.frame_out(out, codec_logits, sample) + for _ in range(k + 1, frames): + tail.next_frame() + tail.frame_in(embed) + tail.frame_out(replay(), codec_logits, sample) + return False + graphs["mid" if k < frames - 1 else "last"].replay() + return True + + # -- shapes, generators, capture ---------------------------------------- + + def _shape(self, key, batch, frames, vocab, controls, hidden, inputs_embeds) -> _Shape | None: + shape = self._shapes.get(key) + if shape is not None: + self._shapes.move_to_end(key) + return shape + size = (batch, int(frames)) + if ( + self._captures_per_shape.get(size, 0) >= MAX_CAPTURES_PER_SHAPE + or self._unstable_per_shape.get(size, 0) >= MAX_UNSTABLE_PER_SHAPE + ): + return None + tail = FrameTail( + batch=batch, + frames=int(frames), + vocab=vocab, + window=int(controls.window.shape[1]), + eos=int(controls.eos_token_id), + hidden=hidden, + inputs_embeds=inputs_embeds, + ) + shape = _Shape(tail=tail) + self._shapes[key] = shape + while len(self._shapes) > self.max_shapes: + self._shapes.popitem(last=False) + return shape + + def _slot(self, index: int, device: torch.device) -> torch.Generator: + while len(self._slots) <= index: + self._slots.append(torch.Generator(device=device)) + return self._slots[index] + + def _unstable(self, key: tuple[Any, ...]) -> None: + self.stats.unstable += 1 + size = (int(key[0]), int(key[1])) + self._unstable_per_shape[size] = self._unstable_per_shape.get(size, 0) + 1 + self._shapes.pop(key, None) + self._failed.add(key) + _log_once( + "unstable", + "the talker forward returned a new buffer inside a K-step step (rows %d); " + "that shape runs eagerly from now on", + key[0], + ) + + def _capture(self, key, shape: _Shape, embed, codec_logits, sample, slots: list[torch.Generator]) -> None: + tail = shape.tail + size = (tail.batch, tail.frames) + self._captures_per_shape[size] = self._captures_per_shape.get(size, 0) + 1 + device = tail.hidden.device + + def mid() -> None: + tail.frame_out(tail.hidden, codec_logits, sample) + tail.next_frame() + tail.frame_in(embed) + + steps: list[tuple[str, Callable[[], None], list[torch.Generator]]] = [ + ("first", lambda: tail.frame_in(embed), []), + ("last", lambda: tail.frame_out(tail.hidden, codec_logits, sample), slots), + ] + if tail.frames > 2: + steps.insert(1, ("mid", mid, slots)) + generators = list(slots) + if device.type == "cuda": + index = device.index if device.index is not None else torch.accelerator.current_device_index() + generators.append(torch.cuda.default_generators[index]) + else: + generators.append(torch.default_generator) + # Capturing records the RNG ops without running them; restore every + # generator anyway so a capture can never shift what eager would draw. + saved = [(gen, gen.get_state()) for gen in generators] + reserved = torch.cuda.memory_reserved(device) if device.type == "cuda" else 0 + start = time.perf_counter() + try: + if self._graph_factory is not None: + graphs = {name: self._graph_factory(fn, gens) for name, fn, gens in steps} + else: + graphs = self._cuda_capture(steps, device) + except Exception as exc: # noqa: BLE001 - a failed capture keeps the eager path + self._failed.add(key) + self._shapes.pop(key, None) + logger.warning( + "[minicpmo] Talker K-step frame graphs: capture failed for rows %d, frames %d (%s: %s); " + "this step shape keeps the eager tail", + tail.batch, + tail.frames, + type(exc).__name__, + exc, + ) + return + finally: + for gen, state in saved: + gen.set_state(state) + elapsed = (time.perf_counter() - start) * 1e3 + shape.graphs = graphs + shape.held = _triton_sampler_buffers() + self.stats.captures += 1 + self.stats.capture_ms += elapsed + grown = (torch.cuda.memory_reserved(device) - reserved) / 2**20 if device.type == "cuda" else 0.0 + logger.info( + "[minicpmo] Talker K-step frame graphs captured: rows %d, frames %d, %d graphs in %.1f ms, " + "pool reserved %+.1f MiB (%d shapes)", + tail.batch, + tail.frames, + len(graphs), + elapsed, + grown, + len(self._shapes), + ) + + def _cuda_capture(self, steps, device: torch.device) -> dict[str, Any]: + """Capture each tail on a side stream into one private pool. + + Not ``torch.cuda.graph``: it synchronizes and empties the allocator cache. + The pool is private because vLLM's shared pool counts the talker graph's + (weak-ref) output as free, and a tail temporary there could overwrite it. + """ + if self._pool is None: + self._pool = torch.cuda.graph_pool_handle() + self._stream = torch.cuda.Stream(device=device) + current = torch.cuda.current_stream(device) + self._stream.wait_stream(current) + graphs: dict[str, Any] = {} + with torch.cuda.stream(self._stream): + for name, fn, generators in steps: + graph = torch.cuda.CUDAGraph() + for gen in generators: + graph.register_generator_state(gen) + graph.capture_begin(pool=self._pool, capture_error_mode="thread_local") + try: + fn() + except BaseException: + try: + graph.capture_end() + except Exception: # noqa: BLE001 - the capture error is the one to report + pass + raise + graph.capture_end() + graphs[name] = graph + current.wait_stream(self._stream) + return graphs + + # -- diagnostics ---------------------------------------------------------- + + def _count_decline(self, reason: str) -> None: + _log_once(reason, "eager tail: %s", reason) + self.stats.declined[reason] = self.stats.declined.get(reason, 0) + 1 + self._maybe_log_stats() + + def _maybe_log_stats(self) -> None: + if not self._stats_enabled: + return + s = self.stats + total = s.graph_steps + s.eager_steps + sum(s.declined.values()) + if total in (1, 100, 1000) or total % 10000 == 0: + logger.info( + "[minicpmo] Talker K-step frame graph stats: steps=%d graph=%d eager=%d declined=%s " + "captures=%d (%.1f ms) unstable=%d shapes=%d", + total, + s.graph_steps, + s.eager_steps, + dict(s.declined), + s.captures, + s.capture_ms, + s.unstable, + len(self._shapes), + ) From 171654b8bc3456820e8e1b9a8d4db0b14b610ff3 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:30 +0800 Subject: [PATCH 79/89] [Perf][MiniCPM-o] Stage 0 streaming audio encoder batching, in-place KV, and CUDA graphs Stage 0 streaming audio encoder batching, in-place KV, and CUDA graphs Batch streaming audio encoding across sessions, maintain KV caches in place, capture audio encoder CUDA graphs with resident cache slots, and add incremental fbank. Signed-off-by: BeatSeat --- .../minicpmo_4_5/test_incremental_fbank.py | 308 ++++++++ .../test_native_duplex_handoff_trim.py | 118 +++ .../test_native_duplex_sampling.py | 491 ++++++++++++ .../test_native_duplex_sampling_cuda.py | 75 ++ .../test_stage0_audio_graph_deferred.py | 347 ++++++++ .../minicpmo_4_5/test_stage0_fbank_stats.py | 69 ++ .../test_streaming_audio_batch.py | 563 +++++++++++++ .../test_streaming_audio_graph.py | 511 ++++++++++++ .../test_streaming_audio_resident.py | 222 ++++++ .../minicpmo_4_5/test_turn_audio_host_sync.py | 80 ++ .../minicpmo_4_5/duplex/incremental_fbank.py | 266 +++++++ .../models/minicpmo_4_5/duplex/plugin.py | 37 + .../models/minicpmo_4_5/duplex/stage0.py | 615 ++++++++++++++- .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 536 ++++++++++++- .../minicpmo_4_5/streaming_audio_encoder.py | 415 ++++++++++ .../streaming_audio_encoder_graph.py | 743 ++++++++++++++++++ 16 files changed, 5345 insertions(+), 51 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_incremental_fbank.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_native_duplex_handoff_trim.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling_cuda.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_stage0_audio_graph_deferred.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_stage0_fbank_stats.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_streaming_audio_batch.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_streaming_audio_graph.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_streaming_audio_resident.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_turn_audio_host_sync.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/duplex/incremental_fbank.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder_graph.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_incremental_fbank.py b/tests/model_executor/models/minicpmo_4_5/test_incremental_fbank.py new file mode 100644 index 00000000000..c7c3c5786ef --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_incremental_fbank.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""``duplex_incremental_fbank``: the incremental Stage-0 streaming fbank is bitwise equal to the full one. + +The reference is the remote-code ``StreamingMelProcessorExact._extract_full`` +(``processing_minicpmo.py`` of openbmb/MiniCPM-o-4_5), imported from the +transformers dynamic-module cache; without it those tests skip. +""" + +from __future__ import annotations + +import copy +import importlib +import wave +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch + +import vllm_omni.model_executor.models.minicpmo_4_5.duplex.incremental_fbank as incremental_fbank +import vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 as stage0_module +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.incremental_fbank import enable_incremental_fbank +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import MiniCPMO45Stage0DuplexRuntime, _FbankStats + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_ASSETS = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5" +_SR = 16000 +_HOP = 160 + + +@pytest.fixture(scope="module") +def remote() -> ModuleType: + """The remote ``processing_minicpmo`` module (the newest one in the transformers module cache).""" + from transformers.utils import HF_MODULES_CACHE + + candidates = [ + path + for path in Path(HF_MODULES_CACHE, "transformers_modules").glob("*/*/processing_minicpmo.py") + if "class StreamingMelProcessorExact" in path.read_text() + ] + if not candidates: + pytest.skip("MiniCPM-o 4.5 remote processing code is not in the transformers module cache") + from transformers.dynamic_module_utils import init_hf_modules + + init_hf_modules() + path = max(candidates, key=lambda candidate: candidate.stat().st_mtime) + return importlib.import_module(".".join(path.relative_to(HF_MODULES_CACHE).with_suffix("").parts)) + + +def _mel(remote: ModuleType, *, incremental: bool) -> object: + # Stage 0's streaming configuration (_configure_streaming_processor). + mel = remote.StreamingMelProcessorExact( + feature_extractor=remote.MiniCPMAAudioProcessor(), + chunk_ms=1000, + first_chunk_ms=1035, + cnn_redundancy_ms=20, + enable_sliding_window=True, + slide_trigger_seconds=30.0, + slide_stride_seconds=10.0, + ) + if incremental: + assert enable_incremental_fbank(mel) + return mel + + +def _remote_cls(mel: object) -> type: + return next(cls for cls in type(mel).__mro__ if cls.__name__ == "StreamingMelProcessorExact") + + +def _original(mel: object) -> torch.Tensor: + return _remote_cls(mel)._extract_full(mel) + + +def _speech() -> np.ndarray: + parts = [] + for name in ("soft_interrupt_16k.wav", "response_required_16k.wav"): + with wave.open(str(_ASSETS / name)) as wav: + assert wav.getframerate() == _SR and wav.getsampwidth() == 2 and wav.getnchannels() == 1 + parts.append(np.frombuffer(wav.readframes(wav.getnframes()), dtype=np.int16).astype(np.float32) / 32768) + parts.append(np.zeros(_SR // 2, dtype=np.float32)) + speech = np.concatenate(parts) + return np.concatenate([speech, speech, speech[: 10 * _SR]]) # ~72 s: two slides + + +def _noise(rng: np.random.Generator, seconds: float) -> np.ndarray: + audio = rng.standard_normal(int(seconds * _SR)) * 0.05 + audio[_SR * 3 : _SR * 7] *= 20 # loud stretch: clipping-level peaks for the dynamic range + audio[_SR * 12 : _SR * 14] = 0.0 # digital silence (clamped at 1e-10) + return audio.astype(np.float32) + + +def _record(mel: object, extract) -> list[torch.Tensor]: + # Capture every full mel the processor computes inside process(). + seen: list[torch.Tensor] = [] + + def wrapper() -> torch.Tensor: + seen.append(extract(mel)) + return seen[-1] + + mel._extract_full = wrapper + return seen + + +def _run_streams(remote: ModuleType, audio: np.ndarray, chunk_sizes) -> dict[str, int]: + ref, inc = _mel(remote, incremental=False), _mel(remote, incremental=True) + ref_full = _record(ref, _remote_cls(ref)._extract_full) + inc_full = _record(inc, type(inc)._extract_full) + position, calls, slides, short = 0, 0, 0, 0 + while position < len(audio): + size = ref.get_chunk_size() if chunk_sizes is None else next(chunk_sizes) + chunk = audio[position : position + size] + position += size + dropped = ref.left_samples_dropped + if len(ref.buffer) + len(chunk) < 400: + with pytest.raises(ValueError): + ref.process(chunk) + with pytest.raises(ValueError): + inc.process(chunk) + continue + expected, expected_info = ref.process(chunk) + got, got_info = inc.process(chunk) + assert torch.equal(inc_full[-1], ref_full[-1]), f"call {calls}: full mel differs" + assert torch.equal(got, expected) and got_info == expected_info + calls += 1 + slides += ref.left_samples_dropped != dropped + short += len(ref.buffer) < 5 * _SR + return {"calls": calls, "slides": slides, "short": short} + + +@pytest.mark.parametrize("kind", ["below_5s", "above_5s", "slide"]) +def test_random_buffers_are_bitwise_equal(remote, kind: str) -> None: + # 80 cases per kind: a cached buffer, then a longer one (slid by a random number of hops for + # "slide"); lengths in samples, not hop aligned. "above_5s" also crosses the 5 s switch. + rng = np.random.default_rng({"below_5s": 1, "above_5s": 2, "slide": 3}[kind]) + inc = _mel(remote, incremental=True) + for _ in range(80): + total = int(rng.integers(2_000, 5 * _SR) if kind == "below_5s" else rng.integers(5 * _SR, 640_000)) + audio = (rng.standard_normal(total) * rng.choice([1e-4, 0.01, 0.1, 0.9])).astype(np.float32) + if rng.random() < 0.3: + audio[: int(rng.integers(0, total))] = 0.0 + prev_end = int(rng.integers(400, min(total, 480_000) + 1)) + inc.reset() + inc.buffer, inc.left_samples_dropped = audio[:prev_end], 0 + inc._extract_full() + dropped = int(rng.integers(1, max(2, (prev_end - 400) // _HOP + 1))) * _HOP if kind == "slide" else 0 + dropped = min(dropped, (prev_end - 400) // _HOP * _HOP) + start = max(prev_end, 5 * _SR) if kind == "above_5s" else prev_end + end = int(rng.integers(min(start, total), min(total, dropped + 480_000) + 1)) + inc.buffer, inc.left_samples_dropped = audio[dropped:end], dropped + got = inc._extract_full() + expected = _original(inc) + assert torch.equal(got, expected), (prev_end, dropped, end, (got - expected).abs().max().item()) + + +@pytest.mark.parametrize("source", ["speech", "noise"]) +@pytest.mark.parametrize("chunking", ["units", "random"]) +def test_streams_are_bitwise_equal(remote, source: str, chunking: str) -> None: + rng = np.random.default_rng(7) + audio = _speech() if source == "speech" else _noise(rng, 66.0) + sizes = None if chunking == "units" else iter(lambda: int(rng.integers(100, 24_000)), None) + stats = _run_streams(remote, audio, sizes) + assert stats["slides"] >= 2 and stats["short"] >= 1 and stats["calls"] >= 60, stats + + +def test_reset_restore_and_copies_stay_exact(remote) -> None: + rng = np.random.default_rng(3) + audio = _noise(rng, 60.0) + ref, inc = _mel(remote, incremental=False), _mel(remote, incremental=True) + + def feed(start: int, stop: int, *pairs) -> None: + for offset in range(start, stop, _SR): + for a, b in pairs: + expected, _ = a.process(audio[offset : offset + _SR]) + got, _ = b.process(audio[offset : offset + _SR]) + assert torch.equal(got, expected) + assert torch.equal(_original(b), b._extract_full()) + + feed(0, 8 * _SR, (ref, inc)) + snapshot_ref, snapshot_inc = ref.get_snapshot(), inc.get_snapshot() + feed(8 * _SR, 14 * _SR, (ref, inc)) + ref.restore_snapshot(snapshot_ref) + inc.restore_snapshot(snapshot_inc) + assert inc._ifb_valid == 0 + feed(20 * _SR, 45 * _SR, (ref, inc)) # another continuation, past the slide + ref_copy, inc_copy = copy.deepcopy(ref), copy.deepcopy(inc) + assert type(inc_copy) is type(inc) + feed(0, 3 * _SR, (ref, inc)) + feed(5 * _SR, 9 * _SR, (ref_copy, inc_copy)) + ref.reset() + inc.reset() + assert inc._ifb_valid == 0 + feed(10 * _SR, 18 * _SR, (ref, inc)) + + +def test_unsupported_buffers_take_the_original_path(remote) -> None: + inc = _mel(remote, incremental=True) + inc.buffer = np.zeros(399, dtype=np.float32) + with pytest.raises(ValueError): + inc._extract_full() + audio = (np.random.default_rng(5).standard_normal(31 * _SR) * 0.1).astype(np.float32) + for buffer in (audio, audio[: 6 * _SR].astype(np.float64)): # truncated at 30 s; float64 cast by Whisper + inc.buffer = buffer + assert torch.equal(inc._extract_full(), _original(inc)) + assert inc._ifb_valid == 0 + + +def test_self_check_rejects_a_different_original(remote) -> None: + class ShiftedMel(remote.StreamingMelProcessorExact): + def _extract_full(self): + return super()._extract_full() + 1e-6 + + mel = ShiftedMel(feature_extractor=remote.MiniCPMAAudioProcessor(), chunk_ms=1000) + assert enable_incremental_fbank(mel) is False + assert type(mel) is ShiftedMel + + +def test_enable_rejects_other_processors() -> None: + assert enable_incremental_fbank(SimpleNamespace(buffer=np.zeros(0, dtype=np.float32))) is False + + +def _session_processor(remote: ModuleType) -> object: + """The remote MiniCPMOProcessor's streaming surface (its methods, without the tokenizer and image side).""" + names = ( + "_init_streaming_processor", + "set_streaming_mode", + "reset_streaming", + "get_streaming_chunk_size", + "process_audio_streaming", + ) + cls = type("StreamingProcessor", (), {name: getattr(remote.MiniCPMOProcessor, name) for name in names}) + processor = cls() + processor.audio_processor = remote.MiniCPMAAudioProcessor() + processor._streaming_mel_processor = None + return processor + + +def _runtime(processor: object, *, incremental: bool, stats: _FbankStats | None = None): + runtime = object.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.processor = processor + runtime.stage_model = runtime.thinker = SimpleNamespace() + runtime._incremental_fbank = incremental + runtime._fbank_stats = stats + return runtime + + +def test_stage0_sessions_stream_identical_features_with_fbank_stats(remote, monkeypatch) -> None: + logger = MagicMock() + monkeypatch.setattr(stage0_module, "logger", logger) + audio = _speech()[: 40 * _SR] + off = _runtime(_session_processor(remote), incremental=False) + on = _runtime(_session_processor(remote), incremental=True, stats=_FbankStats(every=10, incremental=True)) + off_processor = off._configure_streaming_processor(SimpleNamespace(streaming_processor=None)) + on_processor = on._configure_streaming_processor(SimpleNamespace(streaming_processor=None)) + assert type(off_processor._streaming_mel_processor) is remote.StreamingMelProcessorExact + assert isinstance(on_processor._streaming_mel_processor, remote.StreamingMelProcessorExact) + assert type(on_processor._streaming_mel_processor) is not remote.StreamingMelProcessorExact + position, index = 0, 0 + while position < len(audio): + size = off._streaming_chunk_size(off_processor) + chunk = audio[position : position + size] + position += size + expected = off._process_streaming_audio(chunk, index, processor=off_processor) + got = on._process_streaming_audio(chunk, index, processor=on_processor) + assert torch.equal(got["audio_features"], expected["audio_features"]) + index += 1 + assert on._fbank_stats.calls == index + assert logger.info.call_count == index // 10 + assert "(incremental)" in logger.info.call_args.args[0] + + +@pytest.mark.parametrize("incremental", [False, True]) +def test_configure_streaming_processor_enables_only_when_on(monkeypatch, incremental: bool) -> None: + enable = MagicMock(return_value=False) + logger = MagicMock() + monkeypatch.setattr(incremental_fbank, "enable_incremental_fbank", enable) + monkeypatch.setattr(stage0_module, "logger", logger) + + class Processor: + _streaming_mel_processor = None + + def set_streaming_mode(self, **kwargs) -> None: + self._streaming_mel_processor = SimpleNamespace(**kwargs) + + def reset_streaming(self) -> None: + pass + + runtime = _runtime(Processor(), incremental=incremental) + processor = runtime._configure_streaming_processor(SimpleNamespace(streaming_processor=None)) + if incremental: + enable.assert_called_once_with(processor._streaming_mel_processor) + logger.warning_once.assert_called_once() # not supported: stays on the full recompute + else: + enable.assert_not_called() + + +@pytest.mark.parametrize(("value", "expected"), [(True, True), (False, False), ("true", False), (None, False)]) +def test_runtime_reads_the_switch_from_the_hf_config(value: object, expected: bool) -> None: + stage_model = SimpleNamespace( + config=SimpleNamespace(duplex_incremental_fbank=value, duplex_fbank_stats=True), + processor=SimpleNamespace(), + ) + runtime = MiniCPMO45Stage0DuplexRuntime(stage_model, device="cpu") + assert runtime._incremental_fbank is expected + assert ("(incremental)" in runtime._fbank_stats._message) is expected diff --git a/tests/model_executor/models/minicpmo_4_5/test_native_duplex_handoff_trim.py b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_handoff_trim.py new file mode 100644 index 00000000000..516f5f956af --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_handoff_trim.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Guard for the native-duplex Talker handoff when decision tokens outlive their rows. + +A segment the session consumes without a Talker handoff (a listen decision, or +a ``speak`` whose turn ended with nothing to say) advances the request's +cumulative output, but the next streaming update folds those tokens into the +rewritten prompt and the forwarded hidden ledger starts over. The next handoff +then slices rows for tokens decoded before the rewrite -- rows that are gone -- +and the ``missing own-token hidden states`` ValueError fails the request (the +serving crash logged at orchestrator.py:2513). These tests pin the trim that +skips the stale leading decision run, and that it never touches normal inputs. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import ( + _native_duplex_trim_unbacked_decisions, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +#: listen / speak / turn_eos / chunk_eos of MiniCPM-o 4.5 (decoded from the +#: crash log: unit=[151706, 151717, 151705, 151706, 576, 883, 151718]). +LISTEN_ID = 151705 +SPEAK_ID = 151706 +TURN_EOS_ID = 151717 +CHUNK_EOS_ID = 151718 +TTS_BOS_ID = 151703 + +SPECIAL_TOKEN_IDS = { + "listen_token_id": LISTEN_ID, + "speak_token_id": SPEAK_ID, + "turn_eos_token_id": TURN_EOS_ID, + "chunk_eos_token_id": CHUNK_EOS_ID, + "chunk_tts_eos_token_id": 151719, + "tts_bos_token_id": TTS_BOS_ID, +} + + +def _ledger_mm_output(ids: list[int], positions: list[int]) -> dict[str, object]: + """A handoff ``multimodal_output`` carrying only the row ledger.""" + return { + "latent_input_ids": torch.tensor(ids, dtype=torch.long).reshape(-1, 1), + "latent_positions": torch.tensor(positions, dtype=torch.long).reshape(-1, 1), + } + + +def test_crash_shape_trims_stale_leading_decisions() -> None: + """The exact unit from the serving crash, against its actual ledger tail. + + The ledger only backs the trailing ``[speak, 576, 883]`` decode rows; the + leading ``[speak, turn_eos, listen]`` were folded into the rewritten + prompt. The trim drops exactly those control decisions, leaving a handoff + the ledger can satisfy. + """ + ledger_ids = [128244] * 18 + [151670] + [151698] * 10 + [SPEAK_ID, 576, 883] + ledger_positions = list(range(len(ledger_ids))) + mm_output = _ledger_mm_output(ledger_ids, ledger_positions) + unit = [SPEAK_ID, TURN_EOS_ID, LISTEN_ID, SPEAK_ID, 576, 883, CHUNK_EOS_ID] + + trimmed = _native_duplex_trim_unbacked_decisions(unit, mm_output, SPECIAL_TOKEN_IDS, request_id="req-crash") + + assert trimmed == [SPEAK_ID, 576, 883, CHUNK_EOS_ID] + + +def test_normal_unit_is_untouched() -> None: + """A unit whose rows are present is returned unchanged (no trim, no warning).""" + ledger_ids = [128244] * 18 + [SPEAK_ID, 576, 883] + mm_output = _ledger_mm_output(ledger_ids, list(range(len(ledger_ids)))) + unit = [SPEAK_ID, 576, 883, CHUNK_EOS_ID] + + trimmed = _native_duplex_trim_unbacked_decisions(unit, mm_output, SPECIAL_TOKEN_IDS, request_id="req-ok") + + assert trimmed == unit + + +def test_listen_run_without_rows_is_skipped_when_remainder_backed() -> None: + """Accumulated listen decisions ahead of a speak are skippable too.""" + ledger_ids = [151698] * 10 + [SPEAK_ID, 42] + mm_output = _ledger_mm_output(ledger_ids, list(range(len(ledger_ids)))) + unit = [LISTEN_ID, LISTEN_ID, SPEAK_ID, 42, CHUNK_EOS_ID] + + trimmed = _native_duplex_trim_unbacked_decisions(unit, mm_output, SPECIAL_TOKEN_IDS, request_id="req-listen") + + assert trimmed == [SPEAK_ID, 42, CHUNK_EOS_ID] + + +def test_text_tokens_are_never_trimmed() -> None: + """The trim stops at the first non-control token: stale text is an error, not a trim.""" + ledger_ids = [SPEAK_ID, 883] + mm_output = _ledger_mm_output(ledger_ids, list(range(len(ledger_ids)))) + unit = [SPEAK_ID, TURN_EOS_ID, 576, 883, CHUNK_EOS_ID] + + trimmed = _native_duplex_trim_unbacked_decisions(unit, mm_output, SPECIAL_TOKEN_IDS, request_id="req-text") + + assert trimmed == unit + + +def test_missing_ledger_returns_input_unchanged() -> None: + unit = [SPEAK_ID, 576, 883] + trimmed = _native_duplex_trim_unbacked_decisions(unit, {}, SPECIAL_TOKEN_IDS, request_id="req-nol") + assert trimmed == unit + + +def test_unbackable_input_returns_input_unchanged() -> None: + """When no suffix is backed either, the input is left for the explicit error.""" + ledger_ids = [151698] * 4 + mm_output = _ledger_mm_output(ledger_ids, list(range(len(ledger_ids)))) + unit = [SPEAK_ID, 576, 883, CHUNK_EOS_ID] + + trimmed = _native_duplex_trim_unbacked_decisions(unit, mm_output, SPECIAL_TOKEN_IDS, request_id="req-bad") + + assert trimmed == unit diff --git a/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling.py b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling.py new file mode 100644 index 00000000000..b1ce7f6f75b --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling.py @@ -0,0 +1,491 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stage-0 native duplex sampling: batched parameter reads and the vectorized repetition penalty.""" + +import copy +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import MiniCPMO45OmniForConditionalGeneration + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +Model = MiniCPMO45OmniForConditionalGeneration + + +@pytest.mark.parametrize( + "value", + [torch.tensor([0.5, 0.9, 1.3]), torch.tensor([0.25]), torch.tensor(0.7), torch.tensor([]), 0.4, None], +) +def test_batched_rows_match_per_row_reads(value): + md = SimpleNamespace(temperature=value) + rows = Model._sampling_metadata_rows(md, "temperature", 4, 0.8) + assert rows == [Model._sampling_metadata_value(md, "temperature", row, 0.8) for row in range(4)] + + +def _row_model(generated: list[int]) -> Model: + model = Model.__new__(Model) + model._minicpmo45_duplex_state_for_row = lambda row: SimpleNamespace(generated_tokens=list(generated)) + model._record_minicpmo45_duplex_generation_token = lambda row, token: None + model._maybe_cut_minicpmo45_native_duplex_text_chunk = lambda sampled, recent, token_ids: sampled + model._finalize_minicpmo45_native_duplex_sample = lambda row, sampled, token_ids: sampled + model._minicpmo45_native_forbidden_token_ids = lambda token_ids: [] + return model + + +def _reference_sample(logits, generated, temperature, top_k, top_p, generator): + """A straight-line stage-2 draw: penalty, temperature, candidates, multinomial. + + Uses the same candidate-space primitives the production paths use, so the + seeded-equality tests below stay meaningful for the batching/vectorization + they cover; the candidate distribution itself is pinned against the + full-vocabulary filter by test_candidate_distribution_matches_the_filter. + """ + logits = logits.clone() + for token_id in set(generated): + if 0 <= token_id < logits.shape[-1]: + logits[0, token_id] /= 1.05 + probs, indices = Model._duplex_top_k_top_p_candidates(logits / temperature, top_k=top_k, top_p=top_p) + draw = torch.multinomial(probs[0], 1, generator=generator) + return int(indices[0].gather(0, draw).item()) + + +@pytest.mark.parametrize("seed", range(4)) +@pytest.mark.parametrize( + ("top_k", "top_p"), + [(20, 0.85), (0, 0.9), (10, 1.0), (0, 1.0), (5, 0.5), (1, 0.8)], +) +def test_candidate_distribution_matches_the_filter(seed, top_k, top_p): + """Candidate-space sampling is the same categorical as filter + softmax over the vocabulary.""" + vocab = 4096 + logits = torch.randn(1, vocab, generator=torch.Generator().manual_seed(seed)) * 4 + filtered = torch.softmax(Model._top_k_top_p_filter(logits.clone(), top_k=top_k, top_p=top_p), dim=-1) + probs, indices = Model._duplex_top_k_top_p_candidates(logits.clone(), top_k=top_k, top_p=top_p) + dense = torch.zeros_like(filtered) + dense[0].scatter_(0, indices[0], probs[0]) + support = filtered[0] > 0 + assert torch.equal(support, dense[0] > 0) + assert torch.allclose(filtered[0][support], dense[0][support], atol=1e-5) + assert abs(dense[0].sum().item() - 1.0) < 1e-5 + + +@pytest.mark.parametrize("seed", range(4)) +def test_boundary_chunk_eos_probability_matches_full_softmax(seed): + """exp(l_eos - logsumexp(l)) is softmax(l)[chunk_eos], the boundary Bernoulli's p.""" + vocab = 2048 + logits = torch.randn(1, vocab, generator=torch.Generator().manual_seed(seed)) * 3 + logits[0, 7] += 4.0 # make chunk_eos competitive, not a denormal corner + chunk_eos = 7 + full = torch.softmax(logits, dim=-1)[0, chunk_eos] + fast = Model._duplex_boundary_chunk_eos_probs(logits, chunk_eos)[0] + assert torch.allclose(full, fast, atol=1e-6) + assert 0.0 < float(fast) < 1.0 + + +@pytest.mark.parametrize("seed", range(6)) +def test_vectorized_repetition_penalty_samples_like_the_per_token_loop(seed): + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy + + vocab = 64 + logits = torch.randn(1, vocab, generator=torch.Generator().manual_seed(seed)) * 3 + generated = [ + int(t) for t in torch.randint(0, vocab + 4, (40,), generator=torch.Generator().manual_seed(100 + seed)) + ] + generated += [3, 3, 3] # repeats are penalized once + history = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE + md = SimpleNamespace( + all_greedy=False, + generators={0: torch.Generator().manual_seed(7)}, + output_token_ids=[[]], + ) + sampled = _row_model(generated)._sample_minicpmo45_native_duplex_row( + logits.clone(), md, row_idx=0, token_ids={}, params=(0.9, 20, 0.85) + ) + expected = _reference_sample(logits, generated[-history:], 0.9, 20, 0.85, torch.Generator().manual_seed(7)) + assert sampled == expected + + +def test_row_without_params_still_reads_the_metadata(): + md = SimpleNamespace( + all_greedy=False, + generators={0: torch.Generator().manual_seed(3)}, + output_token_ids=[[]], + temperature=torch.tensor([0.8]), + top_k=torch.tensor([10]), + top_p=torch.tensor([0.9]), + ) + logits = torch.randn(1, 32, generator=torch.Generator().manual_seed(1)) + model = _row_model([]) + a = model._sample_minicpmo45_native_duplex_row(logits.clone(), md, row_idx=0, token_ids={}) + md.generators = {0: torch.Generator().manual_seed(3)} + b = model._sample_minicpmo45_native_duplex_row(logits.clone(), md, row_idx=0, token_ids={}, params=(0.8, 10, 0.9)) + assert a == b + + +# --- _sample_minicpmo45_native_duplex_rows: the batched Stage-0 core ------- +# +# The reference is `_sample_minicpmo45_native_duplex_row`, driven one row at +# a time like the pre-batching Stage-0 loop: the returned tokens and the host +# state it mutates must match. + +TOKEN_IDS = { + "unit_token_id": 1, + "chunk_eos_token_id": 5, + "listen_token_id": 6, + "tts_bos_token_id": 7, + "turn_eos_token_id": 8, + "chunk_tts_eos_token_id": 9, + "tts_pad_token_id": 10, +} +VOCAB = 40 + + +def _rows_model( + states: dict[int, SimpleNamespace], + *, + forbidden: list[int] | None = None, + max_new_speak_tokens_per_chunk: int | None = None, +) -> Model: + """A model with only the host-data lookups mocked; the samplers themselves are real.""" + model = Model.__new__(Model) + model._minicpmo45_duplex_state_for_row = lambda row_idx: states.get(row_idx) + model._minicpmo45_duplex_payload_for_row = lambda row_idx: None + model._minicpmo45_duplex_row_request_max_tokens = lambda row_idx: None + model._minicpmo45_native_forbidden_token_ids = lambda token_ids: list(forbidden or []) + # Avoid a tokenizer dependency; the chunk-boundary/forced-chunk_eos paths + # under test do not depend on chunk-cutting by decoded text length. + model._maybe_cut_minicpmo45_native_duplex_text_chunk = lambda sampled, recent, token_ids: sampled + if max_new_speak_tokens_per_chunk is not None: + model.max_new_speak_tokens_per_chunk = max_new_speak_tokens_per_chunk + return model + + +def _mixed_scenario(seed: int): + """Six rows exercising every branch: plain sampling, temperature<=0, + a forced chunk_eos (speak-length cap), a boundary-forced chunk_eos, and + a repetition-penalized row -- built twice (reference vs. vectorized) + with independent but identically-seeded state so the two runs are + directly comparable.""" + torch.manual_seed(seed) + logits = torch.randn(6, VOCAB) * 2 + # Row 3: make chunk_eos overwhelmingly likely so the boundary stage picks + # it on its own, without ever reaching stage 2. + logits[3, TOKEN_IDS["chunk_eos_token_id"]] += 50.0 + + output_token_ids = [ + [1, 2, 3], + [4, 4, 5, 6], + [1, 1], # row 2: length 2 trips the max_new_speak_tokens_per_chunk=3 cap + [2, 3], + [7, 8, 9], + [11, 11, 11, 12, 13], + ] + row_params = [ + (0.8, 10, 0.9), # row 0: plain sampled + (1.1, 0, 1.0), # row 1: sampled, top_k/top_p disabled + (0.7, 100, 0.8), # row 2: forced chunk_eos, params unused + (0.7, 100, 0.8), # row 3: boundary-forced chunk_eos, params unused + (0.0, 100, 0.8), # row 4: row-level greedy despite all_greedy=False + (0.9, 20, 0.85), # row 5: repetition-penalized sampled row + ] + seeds = [11, 22, 33, 44, 55, 66] + + def build(): + states = { + row_idx: SimpleNamespace(generated_tokens=list(tokens) if row_idx == 5 else []) + for row_idx, tokens in enumerate([[], [], [], [], [], [11, 11, 12]]) + } + model = _rows_model(states, forbidden=[TOKEN_IDS["tts_pad_token_id"]], max_new_speak_tokens_per_chunk=3) + md = SimpleNamespace( + all_greedy=False, + generators={row_idx: torch.Generator().manual_seed(s) for row_idx, s in enumerate(seeds)}, + output_token_ids=[list(tokens) for tokens in output_token_ids], + ) + return model, states, md + + return logits.clone(), row_params, build + + +def _reference_rows(model, logits, md, row_idxs, row_params): + """The pre-batching Stage-0 loop body: one `_sample_minicpmo45_native_duplex_row` call per row.""" + out = [] + for row_idx in row_idxs: + row_logits = logits[row_idx : row_idx + 1].clone() + sampled = model._sample_minicpmo45_native_duplex_row( + row_logits, + md, + row_idx=row_idx, + token_ids=TOKEN_IDS, + params=row_params[row_idx], + ) + out.append(sampled) + return out + + +@pytest.mark.parametrize("seed", range(4)) +def test_batched_rows_match_the_reference_per_row_loop(seed): + logits, row_params, build = _mixed_scenario(seed) + row_idxs = list(range(6)) + + ref_model, ref_states, ref_md = build() + expected = _reference_rows(ref_model, logits, ref_md, row_idxs, row_params) + + vec_model, vec_states, vec_md = build() + actual = vec_model._sample_minicpmo45_native_duplex_rows( + logits, vec_md, row_idxs=row_idxs, token_ids=TOKEN_IDS, row_params=row_params + ) + + assert actual == expected + # Row 2 and row 3 resolve without ever reaching stage 2, so nothing is + # recorded for them; every other row's generation history must match. + for row_idx in row_idxs: + assert vec_states[row_idx].generated_tokens == ref_states[row_idx].generated_tokens + + +def test_batched_rows_match_the_reference_when_all_greedy(): + logits, row_params, build = _mixed_scenario(seed=0) + row_idxs = list(range(6)) + + ref_model, ref_states, ref_md = build() + ref_md.all_greedy = True + expected = _reference_rows(ref_model, logits, ref_md, row_idxs, row_params) + + vec_model, vec_states, vec_md = build() + vec_md.all_greedy = True + actual = vec_model._sample_minicpmo45_native_duplex_rows( + logits, vec_md, row_idxs=row_idxs, token_ids=TOKEN_IDS, row_params=row_params + ) + + assert actual == expected + for row_idx in row_idxs: + assert vec_states[row_idx].generated_tokens == ref_states[row_idx].generated_tokens + + +def test_batched_rows_match_the_reference_without_chunk_eos(): + """When the model has no chunk_eos token, every row skips the boundary + stage entirely and goes straight to stage 2 (same as the reference).""" + logits, row_params, build = _mixed_scenario(seed=1) + row_idxs = list(range(6)) + token_ids_no_chunk_eos = {k: v for k, v in TOKEN_IDS.items() if k != "chunk_eos_token_id"} + + ref_model, ref_states, ref_md = build() + expected = [ + ref_model._sample_minicpmo45_native_duplex_row( + logits[row_idx : row_idx + 1].clone(), + ref_md, + row_idx=row_idx, + token_ids=token_ids_no_chunk_eos, + params=row_params[row_idx], + ) + for row_idx in row_idxs + ] + + vec_model, vec_states, vec_md = build() + actual = vec_model._sample_minicpmo45_native_duplex_rows( + logits, vec_md, row_idxs=row_idxs, token_ids=token_ids_no_chunk_eos, row_params=row_params + ) + assert actual == expected + + +# --- The deferred (device-decided) Stage-0 path ------------------------------ +# +# Reference: the synchronous path of `_sample_minicpmo45_native_duplex_stage0` +# (the deferred path forced off). Both run the same two consecutive steps; the +# tokens, the host state after the commit, and every generator's position must +# match, so a rewound extra draw can not change any later token. + + +class _CharTokenizer: + """Token t decodes to (t % 3) + 1 copies of one CJK character; special ids decode to "".""" + + clean_up_tokenization_spaces = False + bad_token_ids: list[int] = [] + + def __init__(self, special: set[int], vocab: int): + self.all_special_ids = sorted(special) + self._special = special + self._vocab = vocab + + def __len__(self): + return self._vocab + + def _piece(self, token_id: int, skip_special_tokens: bool) -> str: + if token_id in self._special: + return "" if skip_special_tokens else f"<{token_id}>" + return chr(0x4E00 + token_id) * (token_id % 3 + 1) + + def decode(self, ids, skip_special_tokens=False): + return "".join(self._piece(int(t), skip_special_tokens) for t in ids) + + def batch_decode(self, batch, skip_special_tokens=False): + return [self.decode(ids, skip_special_tokens=skip_special_tokens) for ids in batch] + + +def _stage0_model(states: dict[int, SimpleNamespace], *, max_chars: int) -> Model: + model = Model.__new__(Model) + model.model_stage = "llm" + model._minicpmo45_native_duplex_token_ids_cache = dict(TOKEN_IDS) + model._minicpmo45_tokenizer_cache = _CharTokenizer(set(TOKEN_IDS.values()), VOCAB) + model._minicpmo45_duplex_state_for_row = lambda row_idx: states.get(row_idx) + model._minicpmo45_duplex_payload_for_row = lambda row_idx: None + model._minicpmo45_duplex_row_request_max_tokens = lambda row_idx: None + model.max_new_speak_tokens_per_chunk = 5 + model.max_speak_chars_per_chunk = max_chars + return model + + +def _stage0_scenario(seed: int, *, all_greedy: bool = False): + torch.manual_seed(seed) + steps = [] + for _ in range(2): + logits = torch.randn(6, VOCAB) * 2 + logits[3, TOKEN_IDS["chunk_eos_token_id"]] += 50.0 # row 3: boundary picks chunk_eos + logits[1, TOKEN_IDS["listen_token_id"]] += 40.0 # row 1: listen, rewritten to tts_bos + steps.append(logits) + histories = [[11, 12], [2], [1, 2, 3, 4], [2, 3], [7, 8, 9], [20, 21, 22]] + temperature = torch.tensor([0.8, 1.1, 0.7, 0.7, 0.0, 0.9]) + + def build(): + states = { + row_idx: SimpleNamespace( + generated_tokens=[11, 11, 12] if row_idx == 5 else [], + current_turn_ended=row_idx != 1, + pending_speech_context=False, + ) + for row_idx in range(6) + } + md = SimpleNamespace( + all_greedy=all_greedy, + generators={row_idx: torch.Generator().manual_seed(100 + row_idx) for row_idx in range(6)}, + output_token_ids=[list(h) for h in histories], + temperature=temperature.clone(), + top_k=torch.tensor([10, 0, 100, 100, 100, 20]), + top_p=torch.tensor([0.9, 1.0, 0.8, 0.8, 0.8, 0.85]), + ) + return _stage0_model(states, max_chars=6), states, md + + return steps, build + + +def _run_two_steps(model, md, steps): + tokens = [] + for logits in steps: + out = model._sample_minicpmo45_native_duplex_stage0(logits.clone(), md, duplex_rows=list(range(6))) + step_tokens = out.sampled_token_ids.squeeze(-1).tolist() + tokens.append(step_tokens) + model._commit_minicpmo45_duplex_pending_samples() + for row_idx, token in enumerate(step_tokens): + md.output_token_ids[row_idx].append(token) + return tokens + + +def _snapshot(states, md): + return ( + {row: copy.deepcopy(vars(state)) for row, state in states.items()}, + {row: gen.get_state().clone() for row, gen in md.generators.items()}, + ) + + +@pytest.mark.parametrize("all_greedy", [False, True]) +@pytest.mark.parametrize("seed", range(4)) +def test_deferred_stage0_matches_the_synchronous_path(seed, all_greedy, monkeypatch): + steps, build = _stage0_scenario(seed, all_greedy=all_greedy) + + sync_model, sync_states, sync_md = build() + monkeypatch.setattr(sync_model, "_sample_minicpmo45_native_duplex_rows_deferred", lambda *a, **k: None) + expected = _run_two_steps(sync_model, sync_md, steps) + + model, states, md = build() + calls = {"deferred": 0} + real = model._sample_minicpmo45_native_duplex_rows_deferred + + def counting(*args, **kwargs): + out = real(*args, **kwargs) + calls["deferred"] += out is not None + return out + + monkeypatch.setattr(model, "_sample_minicpmo45_native_duplex_rows_deferred", counting) + actual = _run_two_steps(model, md, steps) + + assert calls["deferred"] == 2 + assert actual == expected + got_states, got_gens = _snapshot(states, md) + want_states, want_gens = _snapshot(sync_states, sync_md) + assert got_states == want_states + for row in want_gens: + assert torch.equal(got_gens[row], want_gens[row]), f"generator {row} advanced differently" + + +def test_deferred_stage0_cuts_and_rewrites_like_the_host(): + """Row 1's listen becomes tts_bos (its turn is open); a long chunk is cut to chunk_eos.""" + steps, build = _stage0_scenario(0, all_greedy=True) + model, states, md = build() + md.output_token_ids[0] = [14, 14] # 3 + 3 chars already in the chunk: any text token reaches 6 + logits = steps[0].clone() + logits[0, 20] += 40.0 # row 0 samples a text token + out = model._sample_minicpmo45_native_duplex_stage0(logits, md, duplex_rows=list(range(6))) + tokens = out.sampled_token_ids.squeeze(-1).tolist() + assert tokens[1] == TOKEN_IDS["tts_bos_token_id"] + assert tokens[0] == TOKEN_IDS["chunk_eos_token_id"] + assert tokens[3] == TOKEN_IDS["chunk_eos_token_id"] + + +def test_deferred_stage0_defers_the_host_state_until_commit(): + steps, build = _stage0_scenario(1) + model, states, md = build() + before = _snapshot(states, md)[0] + model._sample_minicpmo45_native_duplex_stage0(steps[0].clone(), md, duplex_rows=list(range(6))) + assert _snapshot(states, md)[0] == before + # An append of an unrelated session does not wait for the step. + model._minicpmo45_duplex_row_sessions = {row: f"s{row}" for row in range(6)} + model._minicpmo45_duplex_pending_samples.row_sessions = {row: f"s{row}" for row in range(6)} + model._commit_minicpmo45_duplex_pending_samples(session_ids={"other"}) + assert model._minicpmo45_duplex_pending_samples is not None + model._commit_minicpmo45_duplex_pending_samples(session_ids={"s2"}) + assert model._minicpmo45_duplex_pending_samples is None + assert _snapshot(states, md)[0] != before + + +def test_deferred_stage0_falls_back_for_the_process_generator(): + steps, build = _stage0_scenario(2) + model, states, md = build() + del md.generators[0] # row 0 samples from the process-wide generator + assert ( + model._sample_minicpmo45_native_duplex_rows_deferred( + steps[0].clone(), + md, + row_idxs=list(range(6)), + token_ids=TOKEN_IDS, + row_params=model._minicpmo45_duplex_row_params(md, 6), + ) + is None + ) + + +def test_deferred_stage0_falls_back_inside_a_character(): + steps, build = _stage0_scenario(3) + model, states, md = build() + tokenizer = model._minicpmo45_tokenizer_cache + original = tokenizer.decode + tokenizer.decode = lambda ids, skip_special_tokens=False: original(ids, skip_special_tokens) + "�" + assert ( + model._sample_minicpmo45_native_duplex_rows_deferred( + steps[0].clone(), + md, + row_idxs=list(range(6)), + token_ids=TOKEN_IDS, + row_params=model._minicpmo45_duplex_row_params(md, 6), + ) + is None + ) + + +def test_row_params_prefer_the_host_copies(monkeypatch): + model = Model.__new__(Model) + md = SimpleNamespace(temperature=torch.tensor([0.5, 0.9]), top_k=None, top_p=torch.tensor([0.7, 0.8])) + model._minicpmo45_duplex_row_sampling_host = {0: (0.5, 3, 0.7), 1: (0.9, 4, 0.8)} + monkeypatch.setattr(torch.Tensor, "tolist", lambda self: pytest.fail("device read")) + # top_k has no tensor: the default, as the device read would give. + assert model._minicpmo45_duplex_row_params(md, 2) == [(0.5, 100, 0.7), (0.9, 100, 0.8)] diff --git a/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling_cuda.py b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling_cuda.py new file mode 100644 index 00000000000..ca2565832d7 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling_cuda.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The deferred Stage-0 duplex sampler on CUDA: no host sync, and the synchronous path's tokens and generator offsets.""" + +import pytest +import torch + +from tests.model_executor.models.minicpmo_4_5.test_native_duplex_sampling import ( + _run_two_steps, + _snapshot, + _stage0_scenario, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _on_cuda(build): + def build_cuda(): + model, states, md = build() + md.generators = { + row: torch.Generator(device="cuda").manual_seed(100 + row) for row in md.generators + } + md.temperature, md.top_k, md.top_p = (t.cuda() for t in (md.temperature, md.top_k, md.top_p)) + model._minicpmo45_duplex_row_sampling_host = { + row: (float(md.temperature[row].cpu()), int(md.top_k[row].cpu()), float(md.top_p[row].cpu())) + for row in range(6) + } + return model, states, md + + return build_cuda + + +@requires_cuda +@pytest.mark.parametrize("all_greedy", [False, True]) +@pytest.mark.parametrize("seed", range(3)) +def test_deferred_stage0_on_cuda_matches_the_synchronous_path(seed, all_greedy, monkeypatch): + steps, build = _stage0_scenario(seed, all_greedy=all_greedy) + steps = [logits.cuda() for logits in steps] + build = _on_cuda(build) + + sync_model, sync_states, sync_md = build() + monkeypatch.setattr(sync_model, "_sample_minicpmo45_native_duplex_rows_deferred", lambda *a, **k: None) + expected = _run_two_steps(sync_model, sync_md, steps) + + model, states, md = build() + actual = _run_two_steps(model, md, steps) + + assert actual == expected + got_states, _ = _snapshot(states, md) + want_states, _ = _snapshot(sync_states, sync_md) + assert got_states == want_states + for row in md.generators: + assert md.generators[row].get_offset() == sync_md.generators[row].get_offset() + + +@requires_cuda +@pytest.mark.parametrize("all_greedy", [False, True]) +def test_deferred_stage0_samples_without_a_host_sync(all_greedy): + steps, build = _stage0_scenario(0, all_greedy=all_greedy) + model, states, md = _on_cuda(build)() + logits = steps[0].cuda() + # Warm the cached device tables and the pinned allocator outside the guard. + model._sample_minicpmo45_native_duplex_stage0(logits.clone(), md, duplex_rows=list(range(6))) + model._commit_minicpmo45_duplex_pending_samples() + torch.accelerator.synchronize() + + torch.cuda.set_sync_debug_mode("error") + try: + out = model._sample_minicpmo45_native_duplex_stage0(logits.clone(), md, duplex_rows=list(range(6))) + finally: + torch.cuda.set_sync_debug_mode("default") + assert model._minicpmo45_duplex_pending_samples is not None + assert out.sampled_token_ids.is_cuda diff --git a/tests/model_executor/models/minicpmo_4_5/test_stage0_audio_graph_deferred.py b/tests/model_executor/models/minicpmo_4_5/test_stage0_audio_graph_deferred.py new file mode 100644 index 00000000000..1cc6581ca9a --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_stage0_audio_graph_deferred.py @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The Stage-0 streaming audio encoder's CUDA graphs are captured after the weights load, not at construction. + +The duplex runtime is built in the model's ``__init__`` (#8146), where the +encoder is still in training mode with its initial weights; capturing there was +refused on every deployment ("stays per-session: training mode"). + +The stand-in tests run on CPU; ``test_post_load_capture_matches_eager_batched`` +captures real (tiny) encoder graphs and needs a GPU. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import MiniCPMO45Stage0DuplexRuntime +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import MiniCPMO45OmniForConditionalGeneration +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, + MiniCPMWhisperEncoder, + MultiModalProjector, +) +from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import ( + StreamingAudioChunk, + StreamingAudioKVCache, + encode_streaming_audio_batch, +) + +pytestmark = [pytest.mark.core_model] + +_cpu = pytest.mark.cpu +_cuda = pytest.mark.cuda + + +class _Processor: + """Streaming processor stand-in: 10 ms mel hop; unconfigured it streams 100 ms chunks.""" + + def __init__(self) -> None: + self._streaming_mel_processor = SimpleNamespace(buffer=np.zeros(0, dtype=np.float32), sample_rate=16000) + self.mode: dict | None = None + self.units = 0 + + def set_streaming_mode(self, **kwargs) -> None: + self.mode = kwargs + + def reset_streaming(self) -> None: + self.units = 0 + + def get_streaming_chunk_size(self) -> int: + if self.mode is None: + return 1600 + return self.mode["first_chunk_ms"] * 16 if self.units == 0 else self.mode["chunk_ms"] * 16 + + def process_audio_streaming(self, audio_chunk, *, reset: bool, return_batch_feature: bool = False): + self.units += 1 + return {"audio_features": np.zeros((1, 80, len(audio_chunk) // 160), dtype=np.float32)} + + +class _Thinker: + def __init__(self, *, cuda_graph: bool | None = None) -> None: + self.apm = torch.nn.Linear(1, 1) + self.apm.train() + self.calls: list[tuple[int, bool]] = [] + if cuda_graph is not None: + self.config = SimpleNamespace(duplex_audio_encoder_cuda_graph=cuda_graph) + + def build_streaming_audio_graph_encoder(self, *, unit_frames: int) -> bool: + self.calls.append((unit_frames, self.apm.training)) + return True + + +def _runtime(*, cuda_graph: bool | None = None) -> tuple[MiniCPMO45Stage0DuplexRuntime, _Thinker]: + thinker = _Thinker(cuda_graph=cuda_graph) + stage_model = SimpleNamespace(config=SimpleNamespace(), processor=_Processor(), thinker=thinker) + return MiniCPMO45Stage0DuplexRuntime(stage_model, device="cpu"), thinker + + +@_cpu +def test_construction_does_not_capture() -> None: + _, thinker = _runtime() + assert thinker.calls == [] + + +@_cpu +def test_capture_runs_once_after_load_in_eval_mode() -> None: + runtime, thinker = _runtime() + assert runtime.build_audio_cuda_graph() is True + assert thinker.calls == [(100, False)] + assert runtime.build_audio_cuda_graph() is False + assert len(thinker.calls) == 1 + + +@_cpu +def test_probe_uses_a_session_configured_copy_of_the_processor() -> None: + # The steady unit is 1000 ms (100 mel frames), not the unconfigured 100 ms default. + runtime, thinker = _runtime() + runtime.build_audio_cuda_graph() + assert thinker.calls[0][0] == 100 + assert runtime.processor.mode is None # the shared processor stays untouched + + +@_cpu +def test_model_post_load_hook_builds_the_runtime_graph() -> None: + runtime, thinker = _runtime() + model = SimpleNamespace( + _minicpmo45_duplex_data_plane_helper=runtime, + thinker=thinker.apm, + _module_device=lambda module: torch.device("cpu"), + ) + MiniCPMO45OmniForConditionalGeneration.omni_post_load(model) + assert thinker.calls == [(100, False)] + # Stages without the duplex runtime (Talker, turn mode) have nothing to do. + MiniCPMO45OmniForConditionalGeneration.omni_post_load(SimpleNamespace(thinker=None)) + + +@_cpu +def test_disabled_switch_skips_the_probe_and_the_capture(monkeypatch: pytest.MonkeyPatch) -> None: + # --hf-overrides '{"duplex_audio_encoder_cuda_graph": false}': nothing is probed or captured. + runtime, thinker = _runtime(cuda_graph=False) + probes: list[object] = [] + monkeypatch.setattr(runtime, "_configure_streaming_processor", lambda state=None: probes.append(state)) + assert runtime.build_audio_cuda_graph() is False + assert thinker.calls == [] + assert probes == [] + + +@_cpu +def test_runner_runs_post_load_once_at_the_start_of_profile_run(monkeypatch: pytest.MonkeyPatch) -> None: + # Not inside load_model: vLLM's Worker.load_model scopes max_split_size_mb=20 + # around it, under which the graph pool cannot reuse blocks across captures. + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + + from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + + events: list[object] = [] + + class _Model(torch.nn.Module): + def omni_post_load(self) -> None: + events.append(("post_load", self.training)) + + model = _Model() + assert model.training + + def load_model(runner, *args, **kwargs) -> None: + events.append("load") + runner.model = model.eval() # vLLM's loader returns ``model.eval()`` + + monkeypatch.setattr(GPUModelRunner, "load_model", load_model) + monkeypatch.setattr(GPUModelRunner, "profile_run", lambda runner: events.append("profile")) + for name in ( + "_snapshot_prefix_cache_model_policy", + "_maybe_enable_output_token_ids_for_model_sampler", + "_init_talker_mtp", + "_prewarm_attention_capture_workspaces", + "_report_model_local_kv", + "_warn_unexposed_stage_hooks", + ): + monkeypatch.setattr(OmniGPUModelRunner, name, lambda *args, **kwargs: None) + runner = object.__new__(OmniGPUModelRunner) + runner.load_model() + assert events == ["load"] + runner.profile_run() + runner.profile_run() + assert events == ["load", ("post_load", False), "profile", "profile"] + + +@_cpu +def test_generation_runner_profile_run_runs_post_load_first(monkeypatch: pytest.MonkeyPatch) -> None: + # GPUGenerationModelRunner replaces profile_run without calling up. + from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner + + events: list[str] = [] + runner = object.__new__(GPUGenerationModelRunner) + runner._omni_post_load = lambda: events.append("post_load") + runner.supports_mm_inputs = False + + def dummy_run(*args, **kwargs): + events.append("dummy") + raise RuntimeError("stop after the first dummy run") + + monkeypatch.setattr(GPUGenerationModelRunner, "_dummy_run", dummy_run, raising=False) + runner.max_num_tokens = 1 + with pytest.raises(RuntimeError, match="stop after the first dummy run"): + runner.profile_run() + assert events[:2] == ["post_load", "dummy"] + + +# --- CUDA: the post-load capture of a real (tiny) encoder vs. the eager batched encoder ---------- + +N_MELS = 16 +POOL = 5 +FRAMES = 24 # steady unit: unit_length 10 (see test_streaming_audio_graph.py) + + +class _GraphThinker: + """The audio half of the thinker, built like the model: in training mode until the post-load hook.""" + + supports_streaming_audio_batch = MiniCPMO45OmniLLMForConditionalGeneration.supports_streaming_audio_batch + build_streaming_audio_graph_encoder = MiniCPMO45OmniLLMForConditionalGeneration.build_streaming_audio_graph_encoder + get_audio_embedding_streaming_batch = MiniCPMO45OmniLLMForConditionalGeneration.get_audio_embedding_streaming_batch + + def __init__(self) -> None: + from transformers.models.whisper.modeling_whisper import WhisperConfig + + torch.manual_seed(0) + config = WhisperConfig( + num_mel_bins=N_MELS, + d_model=32, + encoder_layers=2, + encoder_attention_heads=2, + encoder_ffn_dim=128, + max_source_positions=40, + dropout=0.0, + attention_dropout=0.0, + ) + config._attn_implementation = "sdpa" + self.apm = MiniCPMWhisperEncoder(config).to(device="cuda", dtype=torch.bfloat16) + self.apm.train() + self.audio_projection_layer = MultiModalProjector(in_dim=config.encoder_ffn_dim // 4, out_dim=24).to( + device="cuda", dtype=torch.bfloat16 + ) + self.audio_avg_pooler = torch.nn.AvgPool1d(POOL, stride=POOL) + self.audio_encoder_layer = -1 + self.config = SimpleNamespace( + audio_pool_step=POOL, + duplex_audio_kv_page_positions=16, + duplex_audio_encoder_cuda_graph_batch_sizes=[1, 2, 4], + duplex_audio_encoder_cuda_graph_cache_buckets=[16, 32], + ) + + +def _mel(seed: int) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randn((1, N_MELS, FRAMES), generator=generator).to("cuda") + + +def _clone_cache(cache: StreamingAudioKVCache) -> StreamingAudioKVCache: + clone = StreamingAudioKVCache( + num_layers=cache.num_layers, + embed_dim=cache.embed_dim, + num_heads=cache.num_heads, + max_positions=cache.max_positions, + page_positions=cache.page_positions, + ) + history = cache.history(cache.length) + clone.reserve(cache.length, dtype=history.dtype, device=history.device) + clone.commit(0, history) + clone.length = cache.length + return clone + + +@_cuda +def test_post_load_capture_matches_eager_batched() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA graph capture needs a GPU") + thinker = _GraphThinker() + # The stage's chunk sizes make a steady unit FRAMES mel frames (10 ms hop). + stage_model = SimpleNamespace( + config=SimpleNamespace(), + processor=_Processor(), + thinker=thinker, + chunk_ms=FRAMES * 10, + first_chunk_ms=FRAMES * 10, + ) + runtime = MiniCPMO45Stage0DuplexRuntime(stage_model, device="cuda") + assert thinker.apm.training + assert getattr(thinker, "_duplex_audio_cuda_graph_encoder", None) is None + model = SimpleNamespace( + _minicpmo45_duplex_data_plane_helper=runtime, + thinker=thinker.apm, + _module_device=MiniCPMO45OmniForConditionalGeneration._module_device, + ) + MiniCPMO45OmniForConditionalGeneration.omni_post_load(model) + graph = thinker._duplex_audio_cuda_graph_encoder + assert graph is not None and not thinker.apm.training + assert graph.unit_frames == FRAMES + assert sorted(graph._graphs) == [(1, 16), (1, 32), (2, 16), (2, 32), (4, 16), (4, 32)] + replays: list[int] = [] + replay = graph._replay + + def counting_replay(group, *args): + replays.append(len(group)) + return replay(group, *args) + + graph._replay = counting_replay + + for batch in (1, 3, 4): + with torch.inference_mode(): + first = [ + StreamingAudioChunk( + features=_mel(10 * batch + s), cache=None, prefix_extra_frames=0, suffix_extra_frames=2 + ) + for s in range(batch) + ] + _, eager_caches = encode_streaming_audio_batch( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + first, + pool_step=POOL, + page_positions=16, + ) + graph_caches = [_clone_cache(cache) for cache in eager_caches] + replays.clear() + # Two steady rounds: past 11 (bucket 16) then 21 (bucket 32). + for round_index in range(2): + seeds = [1000 * (round_index + 1) + 10 * batch + s for s in range(batch)] + with torch.inference_mode(): + eager_outputs, eager_caches = encode_streaming_audio_batch( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + [ + StreamingAudioChunk( + features=_mel(seed), cache=cache, prefix_extra_frames=2, suffix_extra_frames=2 + ) + for seed, cache in zip(seeds, eager_caches, strict=True) + ], + pool_step=POOL, + page_positions=16, + ) + graph_outputs, graph_caches = thinker.get_audio_embedding_streaming_batch( + [ + StreamingAudioChunk( + features=_mel(seed), cache=cache, prefix_extra_frames=2, suffix_extra_frames=2 + ) + for seed, cache in zip(seeds, graph_caches, strict=True) + ] + ) + for eager_out, graph_out in zip(eager_outputs, graph_outputs, strict=True): + assert eager_out is not None and graph_out is not None + # The tolerance tier of test_streaming_audio_graph.py's padded-bucket cases. + torch.testing.assert_close(graph_out.float(), eager_out.float(), rtol=2e-2, atol=2e-2) + for eager_cache, graph_cache in zip(eager_caches, graph_caches, strict=True): + assert graph_cache.length == eager_cache.length + torch.testing.assert_close( + graph_cache.history(graph_cache.length).float(), + eager_cache.history(eager_cache.length).float(), + rtol=2e-2, + atol=2e-2, + ) + assert replays == [batch, batch] diff --git a/tests/model_executor/models/minicpmo_4_5/test_stage0_fbank_stats.py b/tests/model_executor/models/minicpmo_4_5/test_stage0_fbank_stats.py new file mode 100644 index 00000000000..f0b0c695008 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_stage0_fbank_stats.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""``duplex_fbank_stats``: timing the Stage-0 streaming mel front end without touching its output.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest + +import vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 as stage0_module +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import MiniCPMO45Stage0DuplexRuntime, _FbankStats + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Processor: + def __init__(self) -> None: + self._streaming_mel_processor = SimpleNamespace(buffer=np.zeros(32000, dtype=np.float32), sample_rate=16000) + self.calls = 0 + + def process_audio_streaming(self, audio_chunk, *, reset: bool, return_batch_feature: bool): + self.calls += 1 + return {"audio_features": audio_chunk * 2, "calls": self.calls} + + +def _runtime(stats: _FbankStats | None) -> MiniCPMO45Stage0DuplexRuntime: + runtime = object.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.processor = _Processor() + runtime._fbank_stats = stats + return runtime + + +def test_fbank_stats_off_returns_the_processor_output_untimed(monkeypatch) -> None: + logger = MagicMock() + monkeypatch.setattr(stage0_module, "logger", logger) + runtime = _runtime(None) + chunk = np.ones(16000, dtype=np.float32) + result = runtime._process_streaming_audio(chunk, 3) + np.testing.assert_array_equal(result["audio_features"], chunk * 2) + logger.info.assert_not_called() + + +def test_fbank_stats_on_logs_every_n_calls_and_keeps_the_output(monkeypatch) -> None: + logger = MagicMock() + monkeypatch.setattr(stage0_module, "logger", logger) + stats = _FbankStats(every=3) + runtime = _runtime(stats) + chunk = np.ones(16000, dtype=np.float32) + for index in range(5): + result = runtime._process_streaming_audio(chunk, index) + np.testing.assert_array_equal(result["audio_features"], chunk * 2) + assert result["calls"] == index + 1 + assert stats.calls == 5 + assert logger.info.call_count == 1 + message, count, *_rest, buffer_p50, buffer_max = logger.info.call_args.args + assert "fbank" in message + assert count == 3 + assert buffer_p50 == pytest.approx(2.0) and buffer_max == pytest.approx(2.0) + + +@pytest.mark.parametrize(("value", "expected"), [(True, True), (False, False), ("true", False), (None, False)]) +def test_runtime_reads_the_switch_from_the_hf_config(value: object, expected: bool) -> None: + # Same read as duplex_audio_encoder_pinned_h2d: an explicit True only. + stage_model = SimpleNamespace(config=SimpleNamespace(duplex_fbank_stats=value), processor=_Processor()) + runtime = MiniCPMO45Stage0DuplexRuntime(stage_model, device="cpu") + assert (runtime._fbank_stats is not None) is expected diff --git a/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_batch.py b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_batch.py new file mode 100644 index 00000000000..02d0e259ee5 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_batch.py @@ -0,0 +1,563 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Cross-session batched streaming audio encoder against the per-session path. + +A tiny random Whisper-like encoder runs several sessions through both paths: +the per-session ``get_audio_embedding_streaming`` (``EncoderDecoderCache``) and +one ``encode_streaming_audio_batch`` call per round. Ragged mel lengths, rows +that sit a round out, units the pooler cannot divide, the 30 s position-bound +reset and the Stage-0 prefill staging are all covered. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from torch import nn +from transformers.models.whisper.modeling_whisper import WhisperConfig + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, + MiniCPMWhisperEncoder, + MultiModalProjector, +) +from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import ( + StreamingAudioChunk, + StreamingAudioKVCache, + streaming_batch_unsupported_reason, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +N_MELS = 16 +POOL = 5 + + +class _Thinker: + """The audio half of the thinker: real modules, the real methods under test.""" + + get_audio_embedding_streaming = MiniCPMO45OmniLLMForConditionalGeneration.get_audio_embedding_streaming + get_audio_embedding_streaming_batch = MiniCPMO45OmniLLMForConditionalGeneration.get_audio_embedding_streaming_batch + supports_streaming_audio_batch = MiniCPMO45OmniLLMForConditionalGeneration.supports_streaming_audio_batch + _get_feat_extract_output_lengths = MiniCPMO45OmniLLMForConditionalGeneration._get_feat_extract_output_lengths + + def __init__(self, *, attn_implementation: str = "sdpa", max_positions: int = 40, dtype=torch.float32) -> None: + torch.manual_seed(0) + config = WhisperConfig( + num_mel_bins=N_MELS, + d_model=32, + encoder_layers=2, + encoder_attention_heads=2, + encoder_ffn_dim=128, # the projection takes ffn // 4 == d_model, as in the checkpoint + max_source_positions=max_positions, + dropout=0.0, + attention_dropout=0.0, + ) + config._attn_implementation = attn_implementation + self.apm = MiniCPMWhisperEncoder(config).eval().to(dtype) + self.audio_projection_layer = MultiModalProjector(in_dim=config.encoder_ffn_dim // 4, out_dim=24).to(dtype) + self.audio_avg_pooler = nn.AvgPool1d(POOL, stride=POOL) + self.audio_encoder_layer = -1 + self.audio_past_key_values = None + self.config = SimpleNamespace(audio_pool_step=POOL, duplex_audio_kv_page_positions=16) + + +def _mel(frames: int, seed: int) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randn((1, N_MELS, frames), generator=generator) + + +def _sequential(thinker: _Thinker, cache, mel: torch.Tensor, prefix: int, *, use_extra_context: bool = True): + """One unit through the per-session path, with that session's cache swapped in.""" + thinker.audio_past_key_values = cache + with torch.no_grad(): + nested = thinker.get_audio_embedding_streaming( + {"audio_features": mel, "audio_feature_lens": [torch.tensor([mel.shape[-1]])]}, + use_extra_context=use_extra_context, + prefix_extra_frames=prefix, + suffix_extra_frames=2, + ) + embeds = torch.cat([t for row in nested for t in row]) if nested else None + return embeds, thinker.audio_past_key_values + + +def _legacy_layer(cache, layer: int) -> tuple[torch.Tensor, torch.Tensor]: + layer_cache = cache.self_attention_cache.layers[layer] + return layer_cache.keys, layer_cache.values + + +# Per round, the mel frames of each session's unit (None: sits the round out). +# 22 frames at chunk 0 and 24 later mirror 1020 ms / 1040 ms at 1/10 scale; the +# odd lengths make conv2 read a padded column and 25 frames give an 11-position +# unit that the pooler cannot divide. max_positions=40 forces a reset per cycle. +_SCHEDULE = [ + [22, 22, 23, 21], + [24, None, 24, 25], + [24, 24, None, 24], + [24, 23, 24, 24], + [24, 24, 24, None], + [25, 24, 23, 24], +] + + +@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"]) +def test_batched_rounds_match_sequential_sessions(attn_implementation: str) -> None: + thinker = _Thinker(attn_implementation=attn_implementation) + sessions = len(_SCHEDULE[0]) + legacy_caches: list[object | None] = [None] * sessions + batched_caches: list[StreamingAudioKVCache | None] = [None] * sessions + chunk_index = [0] * sessions + resets = 0 + + for round_index, frames_per_session in enumerate(_SCHEDULE): + chunks, active, expected = [], [], [] + for session, frames in enumerate(frames_per_session): + if frames is None: + continue + mel = _mel(frames, seed=100 * round_index + session) + prefix = 0 if chunk_index[session] == 0 else 2 + before = batched_caches[session] + embeds, legacy_caches[session] = _sequential(thinker, legacy_caches[session], mel, prefix) + expected.append(embeds) + chunks.append( + StreamingAudioChunk(features=mel, cache=before, prefix_extra_frames=prefix, suffix_extra_frames=2) + ) + active.append(session) + chunk_index[session] += 1 + + idle = {s: (batched_caches[s], batched_caches[s].length if batched_caches[s] else 0) for s in range(sessions)} + outputs, caches = thinker.get_audio_embedding_streaming_batch(chunks) + + for session, output, cache, reference in zip(active, outputs, caches, expected, strict=True): + assert output is not None and reference is not None + torch.testing.assert_close(output, reference, rtol=1e-5, atol=1e-5) + if cache is not batched_caches[session] and batched_caches[session] is not None: + resets += 1 + batched_caches[session] = cache + legacy_length = legacy_caches[session].self_attention_cache.get_seq_length() + assert cache.length == legacy_length + for layer in range(len(thinker.apm.layers)): + keys, values = _legacy_layer(legacy_caches[session], layer) + torch.testing.assert_close(cache.head_view(cache.keys(layer)[: cache.length]), keys) + torch.testing.assert_close(cache.head_view(cache.values(layer)[: cache.length]), values) + for session in set(range(sessions)) - set(active): + # Not in the batch: the same cache object, the same committed length. + cache, length = idle[session] + assert batched_caches[session] is cache + assert (cache.length if cache is not None else 0) == length + + assert resets >= 2, "the schedule must cross the max_source_positions reset" + + +def test_ragged_rows_without_extra_context_see_zeros_past_their_end() -> None: + """Untrimmed odd-length rows: conv2's last column reads past the row's frames.""" + thinker = _Thinker() + legacy: list[object | None] = [None] * 3 + batched: list[StreamingAudioKVCache | None] = [None] * 3 + for index, lengths in enumerate([(21, 25, 23), (19, 25, 24)]): + chunks, expected = [], [] + for session, frames in enumerate(lengths): + mel = _mel(frames, seed=20 * index + session) + reference, legacy[session] = _sequential(thinker, legacy[session], mel, 0, use_extra_context=False) + expected.append(reference) + chunks.append(StreamingAudioChunk(features=mel, cache=batched[session], use_extra_context=False)) + outputs, batched = thinker.get_audio_embedding_streaming_batch(chunks) + for output, reference in zip(outputs, expected, strict=True): + torch.testing.assert_close(output, reference, rtol=1e-5, atol=1e-5) + + +def test_single_row_matches_the_per_session_path_bit_for_bit() -> None: + thinker = _Thinker(attn_implementation="eager") + legacy_cache, cache = None, None + for index, frames in enumerate([22, 24, 24]): + mel = _mel(frames, seed=index) + prefix = 0 if index == 0 else 2 + reference, legacy_cache = _sequential(thinker, legacy_cache, mel, prefix) + (output,), (cache,) = thinker.get_audio_embedding_streaming_batch( + [StreamingAudioChunk(features=mel, cache=cache, prefix_extra_frames=prefix, suffix_extra_frames=2)] + ) + # Same shapes, same kernels on CPU: batch size one is exact. + assert torch.equal(output, reference) + + +def test_bf16_batch_stays_within_bf16_rounding_of_the_per_session_path() -> None: + thinker = _Thinker(attn_implementation="sdpa", dtype=torch.bfloat16) + legacy = [None, None, None] + batched = [None, None, None] + for index in range(3): + chunks, expected = [], [] + for session, frames in enumerate([22 if index == 0 else 24, 23 if index == 0 else 24, 21]): + mel = _mel(frames, seed=10 * index + session) + prefix = 0 if index == 0 else 2 + reference, legacy[session] = _sequential(thinker, legacy[session], mel, prefix) + expected.append(reference) + chunks.append( + StreamingAudioChunk( + features=mel, cache=batched[session], prefix_extra_frames=prefix, suffix_extra_frames=2 + ) + ) + outputs, batched = thinker.get_audio_embedding_streaming_batch(chunks) + for output, reference in zip(outputs, expected, strict=True): + torch.testing.assert_close(output.float(), reference.float(), rtol=3e-2, atol=3e-2) + + +def test_unit_empty_after_trimming_returns_none_and_keeps_the_cache() -> None: + thinker = _Thinker() + (first,), (cache,) = thinker.get_audio_embedding_streaming_batch( + [StreamingAudioChunk(features=_mel(22, 0), cache=None, prefix_extra_frames=0, suffix_extra_frames=2)] + ) + assert first is not None + length = cache.length + # 4 frames -> 2 CNN positions, both trimmed as extra context. + (empty, full), (same, other) = thinker.get_audio_embedding_streaming_batch( + [ + StreamingAudioChunk(features=_mel(4, 1), cache=cache, prefix_extra_frames=2, suffix_extra_frames=2), + StreamingAudioChunk(features=_mel(24, 2), cache=None, prefix_extra_frames=0, suffix_extra_frames=2), + ] + ) + assert empty is None and full is not None + assert same is cache and cache.length == length + assert other is not None and other.length > 0 + + +def test_unit_longer_than_the_position_table_repeats_its_last_row() -> None: + thinker = _Thinker(max_positions=8) + mel = _mel(24, 3) # 12 CNN positions -> 10 after trimming, past the 8-row table. + reference, _ = _sequential(thinker, None, mel, prefix=2) + (output,), _ = thinker.get_audio_embedding_streaming_batch( + [StreamingAudioChunk(features=mel, cache=None, prefix_extra_frames=2, suffix_extra_frames=2)] + ) + torch.testing.assert_close(output, reference) + + +def test_kv_cache_grows_by_pages_up_to_the_position_bound_and_keeps_history() -> None: + cache = StreamingAudioKVCache(num_layers=2, embed_dim=8, num_heads=2, max_positions=40, page_positions=16) + cache.reserve(10, dtype=torch.float32, device=torch.device("cpu")) + assert cache.capacity == 16 + cache.keys(1)[:10] = torch.arange(80, dtype=torch.float32).reshape(10, 8) + cache.length = 10 + cache.reserve(20, dtype=torch.float32, device=torch.device("cpu")) + assert cache.capacity == 32 + torch.testing.assert_close(cache.keys(1)[:10], torch.arange(80, dtype=torch.float32).reshape(10, 8)) + cache.reserve(33, dtype=torch.float32, device=torch.device("cpu")) + assert cache.capacity == 40 + assert cache.head_view(cache.keys(0)[:5]).shape == (1, 2, 5, 4) + + +def test_legacy_cache_conversion_continues_the_same_history() -> None: + thinker = _Thinker() + legacy_cache, cache = None, None + for index in range(2): + mel = _mel(22 if index == 0 else 24, index) + _, legacy_cache = _sequential(thinker, legacy_cache, mel, 0 if index == 0 else 2) + _, (cache,) = thinker.get_audio_embedding_streaming_batch( + [ + StreamingAudioChunk( + features=mel, cache=cache, prefix_extra_frames=0 if index == 0 else 2, suffix_extra_frames=2 + ) + ] + ) + mel = _mel(24, 7) + reference, _ = _sequential(thinker, legacy_cache, mel, 2) + converted, _ = _sequential(thinker, cache.to_legacy_cache(), mel, 2) + torch.testing.assert_close(converted, reference) + + +def test_support_gate_names_the_configurations_it_leaves_per_session() -> None: + thinker = _Thinker() + assert streaming_batch_unsupported_reason(thinker.apm, audio_encoder_layer=-1) is None + assert thinker.supports_streaming_audio_batch() + assert "final encoder layer" in streaming_batch_unsupported_reason(thinker.apm, audio_encoder_layer=3) + thinker.config.duplex_audio_encoder_batching = False + assert not thinker.supports_streaming_audio_batch() + half = _Thinker(dtype=torch.float16) + assert "dtype" in streaming_batch_unsupported_reason(half.apm, audio_encoder_layer=-1) + flash = _Thinker() + flash.apm.config._attn_implementation = "flash_attention_2" + assert "flash_attention_2" in streaming_batch_unsupported_reason(flash.apm, audio_encoder_layer=-1) + + +# --- Stage-0 prefill staging ----------------------------------------------------------------- + + +class _Feature(dict): + """A BatchFeature-like mapping that takes the chunk attributes Stage 0 sets.""" + + +class _MelProcessor: + """Deterministic stand-in for the remote-code streaming mel processor (copied per session).""" + + def __init__(self) -> None: + self.calls = 0 + + def get_streaming_chunk_size(self) -> int: + return 240 + + def process_audio_streaming(self, audio_chunk, reset=False, return_batch_feature=True): + del reset, return_batch_feature + self.calls += 1 + frames = 22 if self.calls == 1 else 24 + seed = int(abs(float(np.asarray(audio_chunk).sum())) * 1000) % 10_000 + self.calls + mel = _mel(frames, seed) + return _Feature(audio_features=mel, audio_feature_lens=[torch.tensor([frames])]) + + +class _StageModel(nn.Module): + def __init__(self, thinker: _Thinker) -> None: + super().__init__() + self.embed = nn.Embedding(256, 24) + self._thinker = thinker + + def get_input_embeddings(self): + return self.embed + + +def _runtime(batching: bool): + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import MiniCPMO45Stage0DuplexRuntime + + thinker = _Thinker() + thinker.config.duplex_audio_encoder_batching = batching + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel(thinker) + runtime.thinker = thinker + thinker.llm = SimpleNamespace(model=SimpleNamespace(embed_tokens=runtime.stage_model.embed)) + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [], + ) + runtime.processor = _MelProcessor() + runtime.device = "cpu" + runtime._init_token_ids() + return runtime + + +def _audio(seed: int, samples: int) -> np.ndarray: + return np.random.default_rng(seed).standard_normal(samples).astype(np.float32) + + +@pytest.mark.parametrize("units_in_first_append", [1, 2]) +def test_stage0_batched_prefill_matches_per_session_prefill(units_in_first_append: int) -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState + + sequential, batched = _runtime(batching=False), _runtime(batching=True) + assert sequential._batched_audio_encoder() is None + assert batched._batched_audio_encoder() is not None + states = { + name: [_MiniCPMO45Stage0SessionState(session_id=f"s{i}") for i in range(3)] + for name in ("sequential", "batched") + } + # The runner's preprocess runs under inference mode; so do both paths here. + with torch.inference_mode(): + for seq in range(1, 5): + appends = [ + (index, _audio(10 * seq + index, 240 * (units_in_first_append if seq == 1 else 1))) + for index in range(3) + if not (seq == 3 and index == 1) # session 1 sends nothing this step + ] + expected = [ + sequential._stage_prefill_embeddings_only(states["sequential"][index], audio, seq=seq, epoch=0) + for index, audio in appends + ] + batch = [(states["batched"][index], audio, {"seq": seq, "epoch": 0}) for index, audio in appends] + batched.stage_prefill_batch(batch) + for (index, audio), reference in zip(appends, expected, strict=True): + state = states["batched"][index] + assert state.staged_prefill is not None + result = batched._stage_prefill_embeddings_only(state, audio, seq=seq, epoch=0) + assert state.staged_prefill is None + assert result["success"] is reference["success"] is True + assert result["input_token_ids"] == reference["input_token_ids"] + torch.testing.assert_close(result["inputs_embeds"], reference["inputs_embeds"], rtol=1e-5, atol=1e-5) + assert isinstance(state.audio_past_key_values, StreamingAudioKVCache) + assert not batched.needs_prefill(state, 0, seq) + # Every append's processor ran exactly as often on both paths. + for mine, theirs in zip(states["batched"], states["sequential"], strict=True): + assert mine.streaming_processor.calls == theirs.streaming_processor.calls > 0 + + +def test_stage0_staged_result_is_only_taken_by_its_own_append() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState + + runtime = _runtime(batching=True) + a, b = (_MiniCPMO45Stage0SessionState(session_id=name) for name in "ab") + runtime.stage_prefill_batch( + [(a, _audio(1, 240), {"seq": 1, "epoch": 0}), (b, _audio(2, 240), {"seq": 1, "epoch": 0})] + ) + staged = a.staged_prefill[1] + assert runtime._stage_prefill_embeddings_only(a, _audio(1, 240), seq=1, epoch=0) is staged + # A stale staged result (another identity) is dropped, not returned. + b.staged_prefill = ((0, 9), {"success": False}) + result = runtime._stage_prefill_embeddings_only(b, _audio(3, 240), seq=2, epoch=0) + assert result["success"] is True and b.staged_prefill is None + with pytest.raises(ValueError, match="one append per session"): + runtime.stage_prefill_batch([(a, _audio(4, 240), {"seq": 3}), (a, _audio(5, 240), {"seq": 4})]) + + +def test_preprocess_batch_stages_only_new_appends_of_distinct_sessions() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + staged: list[list[tuple[object, dict]]] = [] + prepared = {("done", 0, 1)} + sessions: dict[str, object] = {} + helper = SimpleNamespace( + sessions=sessions, + batches_audio_encoder=lambda: True, + needs_prefill=lambda state, epoch, seq: (getattr(state, "session_id", None), epoch, seq) not in prepared, + stage_prefill_batch=lambda appends: staged.append([(state, kwargs) for state, _, kwargs in appends]), + _decode_audio_payload=lambda payload: np.zeros(4, dtype=np.float32), + _decode_video_frames_payload=lambda payload: [], + prefetch_vision=lambda appends: None, + take_prefetched_vision=lambda duplex: None, + ) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model._minicpmo45_duplex_data_plane_helper = helper + model._minicpmo45_duplex_session_state = lambda helper, session_id, duplex: sessions.setdefault( + session_id, SimpleNamespace(session_id=session_id) + ) + sessions["done"] = SimpleNamespace(session_id="done") + + def append(session_id: str, seq: object) -> dict: + return {"duplex": {"data_plane": True, "session_id": session_id, "epoch": 0, "seq": seq, "payload": {}}} + + buffer = { + "r1": append("a", 1), + "r2": append("b", 5), + "r3": append("a", 2), # second append of "a": left to preprocess + "r4": append("done", 1), # already built + "r5": append("c", None), # no identity + "r6": {"duplex": {"data_plane": False}}, + "r7": append("d", 1), + } + model.preprocess_batch(req_ids=list(buffer), model_intermediate_buffer=buffer, device=torch.device("cpu")) + assert [[(state.session_id, kwargs["seq"]) for state, kwargs in call] for call in staged] == [ + [("a", 1), ("b", 5), ("d", 1)] + ] + + staged.clear() + single = {"r1": append("e", 1), "r2": append("done", 1)} + model.preprocess_batch(req_ids=list(single), model_intermediate_buffer=single, device=torch.device("cpu")) + assert staged == [] # one new append: preprocess builds it itself + + +def test_stage0_hands_a_batched_history_to_the_per_session_path() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + _MiniCPMO45AudioEncodeRequest, + _MiniCPMO45Stage0SessionState, + ) + + runtime = _runtime(batching=True) + reference = _runtime(batching=False) + state = _MiniCPMO45Stage0SessionState(session_id="s") + expected = _MiniCPMO45Stage0SessionState(session_id="s") + with torch.inference_mode(): + for index, frames in enumerate([22, 24]): + mel = _mel(frames, index) + feature = _Feature(audio_features=mel, audio_feature_lens=[torch.tensor([frames])]) + feature.chunk_idx = index + (embeds,) = runtime._stage_audio_embeddings_batch([_MiniCPMO45AudioEncodeRequest(feature, state)]) + (reference_embeds,) = reference._stage_audio_embeddings_batch( + [_MiniCPMO45AudioEncodeRequest(feature, expected)] + ) + torch.testing.assert_close(embeds, reference_embeds, rtol=1e-5, atol=1e-5) + assert isinstance(state.audio_past_key_values, StreamingAudioKVCache) + # The per-session path taking a unit of a batched session continues its history. + feature = _Feature(audio_features=_mel(24, 9), audio_feature_lens=torch.tensor([[24]])) + feature.chunk_idx = 2 + converted = runtime._stage_audio_embeddings(feature, state=state) + continued = reference._stage_audio_embeddings(feature, state=expected) + torch.testing.assert_close(converted, continued, rtol=1e-5, atol=1e-5) + assert not isinstance(state.audio_past_key_values, StreamingAudioKVCache) + + +def test_preprocess_batch_leaves_everything_per_request_without_a_batched_encoder() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + def unexpected(*args, **kwargs): + raise AssertionError("preprocess_batch staged work with batching disabled") + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace( + batches_audio_encoder=lambda: False, + needs_prefill=unexpected, + stage_prefill_batch=unexpected, + prefetch_vision=lambda appends: None, + ) + buffer = { + rid: {"duplex": {"data_plane": True, "session_id": rid, "epoch": 0, "seq": 1, "payload": {}}} + for rid in ("r1", "r2") + } + model.preprocess_batch(req_ids=["r1", "r2"], model_intermediate_buffer=buffer, device=torch.device("cpu")) + + runtime = _runtime(batching=False) + assert not runtime.batches_audio_encoder() + + +def test_preprocess_batch_hands_prefetched_frames_to_the_batched_prefill() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + staged: list[list[dict]] = [] + prefetched: dict[str, list] = {} + sessions: dict[str, object] = {} + + def undecodable(payload): + raise AssertionError("frames prefetched this step were decoded again") + + helper = SimpleNamespace( + sessions=sessions, + batches_audio_encoder=lambda: True, + needs_prefill=lambda state, epoch, seq: True, + stage_prefill_batch=lambda appends: staged.append([kwargs for _, _, kwargs in appends]), + _decode_audio_payload=lambda payload: np.zeros(4, dtype=np.float32), + _decode_video_frames_payload=undecodable, + prefetch_vision=lambda appends: prefetched.update( + {a["session_id"]: [["emb", a["session_id"]]] for a in appends} + ), + take_prefetched_vision=lambda duplex: prefetched.pop(duplex["session_id"], None), + ) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model._minicpmo45_duplex_data_plane_helper = helper + model._duplex_data_plane_helper = lambda: helper + model._minicpmo45_duplex_session_state = lambda helper, session_id, duplex: sessions.setdefault( + session_id, SimpleNamespace(session_id=session_id) + ) + buffer = { + rid: { + "duplex": { + "data_plane": True, + "session_id": rid, + "epoch": 0, + "seq": 1, + "payload": {"video_frames": ["jpeg"]}, + } + } + for rid in ("a", "b") + } + model.preprocess_batch(req_ids=list(buffer), model_intermediate_buffer=buffer, device=torch.device("cpu")) + assert [[kwargs.get("encoded_frames") for kwargs in call] for call in staged] == [[[["emb", "a"]], [["emb", "b"]]]] + assert all("video_frames" not in kwargs for kwargs in staged[0]) + assert prefetched == {} diff --git a/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_graph.py b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_graph.py new file mode 100644 index 00000000000..12ea680d1b6 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_graph.py @@ -0,0 +1,511 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA-graph capture of the streaming Whisper encoder's steady unit shape. + +The bucket/eligibility helpers are pure Python and run on CPU; the graph +itself needs CUDA (``torch.cuda.graph``), so ``test_graph_matches_eager_*`` +are skipped without a GPU. Both halves check the same thing the module +promises: batch and history padding never change a real row's result. +""" + +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, + MiniCPMWhisperEncoder, + MultiModalProjector, +) +from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import ( + StreamingAudioChunk, + StreamingAudioKVCache, + encode_streaming_audio_batch, +) +from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder_graph import ( + StreamingAudioGraphEncoder, + normalize_buckets, + row_is_steady, + select_bucket, + steady_pooled_length, + steady_unit_length, +) + +pytestmark = [pytest.mark.core_model] + +_cpu = pytest.mark.cpu +_cuda = pytest.mark.cuda + +N_MELS = 16 +POOL = 5 +FRAMES = 24 # steady unit: conv_length 12, trimmed by 1 + 1 -> unit_length 10 (POOL-aligned) + + +# --- pure bucket/eligibility logic: CPU, no CUDA needed ----------------------------------------- + + +@_cpu +def test_select_bucket_picks_the_smallest_covering_value() -> None: + assert select_bucket(0, (500, 1000, 1500)) == 500 + assert select_bucket(500, (500, 1000, 1500)) == 500 + assert select_bucket(501, (500, 1000, 1500)) == 1000 + assert select_bucket(1500, (500, 1000, 1500)) == 1500 + assert select_bucket(1501, (500, 1000, 1500)) is None + + +@_cpu +def test_normalize_buckets_dedupes_sorts_and_drops_non_positive() -> None: + assert normalize_buckets([1000, 500, 500, 0, -5, 1000]) == (500, 1000) + + +@_cpu +def test_steady_unit_length_matches_the_conv_and_trim_arithmetic() -> None: + # frames=24 -> conv_length (24-1)//2+1 = 12; trim(2) = (2+1)//2 = 1 each side. + assert steady_unit_length(24) == 10 + # frames=22 -> conv_length 11; still trimmed by 1 + 1. + assert steady_unit_length(22) == 9 + + +def _chunk( + *, + frames: int = FRAMES, + cache=object(), + prefix: int = 2, + suffix: int = 2, + use_extra_context=True, + feature_length=None, +): + return StreamingAudioChunk( + features=torch.zeros(1, N_MELS, frames), + cache=cache, + prefix_extra_frames=prefix, + suffix_extra_frames=suffix, + use_extra_context=use_extra_context, + feature_length=feature_length, + ) + + +@_cpu +def test_row_is_steady_accepts_only_the_stage0_steady_shape() -> None: + assert row_is_steady(_chunk(), expected_frames=FRAMES) + # A session's first unit: no cache yet, and prefix_extra_frames=0. + assert not row_is_steady(_chunk(cache=None), expected_frames=FRAMES) + assert not row_is_steady(_chunk(prefix=0), expected_frames=FRAMES) + # Any other trim, frame count, or a short final unit falls back too. + assert not row_is_steady(_chunk(suffix=0), expected_frames=FRAMES) + assert not row_is_steady(_chunk(use_extra_context=False), expected_frames=FRAMES) + assert not row_is_steady(_chunk(frames=FRAMES + 1), expected_frames=FRAMES) + assert not row_is_steady(_chunk(feature_length=FRAMES - 3), expected_frames=FRAMES) + assert row_is_steady(_chunk(feature_length=FRAMES), expected_frames=FRAMES) + + +@_cpu +def test_steady_pooled_length_takes_the_smaller_of_cap_and_actual() -> None: + # frames=24 -> unit_length 10: cap ((24-1)//2+1-5)//5+1=2, actual 10//5=2. + assert steady_pooled_length(24, 10, 5) == 2 + # Contrived numbers where the feature_length-derived cap binds below the + # unit_length-derived actual pooled count: still take the min, not actual. + assert steady_pooled_length(unit_frames=10, unit_length=100, pool_step=3) == 1 + + +# --- CUDA: graph replay vs. eager, bucketed and padded ---------------------------------------- + + +class _Thinker: + """The audio half of the thinker: real modules, the real methods under test.""" + + get_audio_embedding_streaming_batch = MiniCPMO45OmniLLMForConditionalGeneration.get_audio_embedding_streaming_batch + supports_streaming_audio_batch = MiniCPMO45OmniLLMForConditionalGeneration.supports_streaming_audio_batch + + def __init__(self, *, max_positions: int, dtype=torch.bfloat16, device="cuda") -> None: + from types import SimpleNamespace + + from transformers.models.whisper.modeling_whisper import WhisperConfig + + torch.manual_seed(0) + config = WhisperConfig( + num_mel_bins=N_MELS, + d_model=32, + encoder_layers=2, + encoder_attention_heads=2, + encoder_ffn_dim=128, + max_source_positions=max_positions, + dropout=0.0, + attention_dropout=0.0, + ) + config._attn_implementation = "sdpa" + self.apm = MiniCPMWhisperEncoder(config).eval().to(device=device, dtype=dtype) + self.audio_projection_layer = MultiModalProjector(in_dim=config.encoder_ffn_dim // 4, out_dim=24).to( + device=device, dtype=dtype + ) + self.audio_avg_pooler = nn.AvgPool1d(POOL, stride=POOL) + self.audio_encoder_layer = -1 + self.config = SimpleNamespace(audio_pool_step=POOL, duplex_audio_kv_page_positions=16) + + +def _mel(frames: int, seed: int, device: str) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randn((1, N_MELS, frames), generator=generator).to(device) + + +def _clone_cache(cache: StreamingAudioKVCache) -> StreamingAudioKVCache: + clone = StreamingAudioKVCache( + num_layers=cache.num_layers, + embed_dim=cache.embed_dim, + num_heads=cache.num_heads, + max_positions=cache.max_positions, + page_positions=cache.page_positions, + ) + if cache.length: + history = cache.history(cache.length) + clone.reserve(cache.length, dtype=history.dtype, device=history.device) + clone.commit(0, history) + clone.length = cache.length + return clone + + +def _start_session(thinker: _Thinker, seed: int, *, frames: int = FRAMES) -> StreamingAudioKVCache: + """One session's first (non-steady) unit, eager: the real cache-in state.""" + (embeds,), (cache,) = encode_streaming_audio_batch( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + [ + StreamingAudioChunk( + features=_mel(frames, seed, "cuda"), cache=None, prefix_extra_frames=0, suffix_extra_frames=2 + ) + ], + pool_step=POOL, + ) + assert embeds is not None + return cache + + +@pytest.fixture +def cuda_thinker(): + if not torch.cuda.is_available(): + pytest.skip("CUDA graph capture needs a GPU") + return _Thinker(max_positions=40) + + +@_cuda +def test_graph_matches_eager_across_batch_padding_and_a_cache_bucket_boundary(cuda_thinker) -> None: + thinker = cuda_thinker + graph = StreamingAudioGraphEncoder( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + unit_frames=FRAMES, + pool_step=POOL, + batch_sizes=(1, 2, 4), + cache_buckets=(16, 32, 40), + page_positions=16, + ) + with torch.inference_mode(): + graph.capture() + assert graph.unit_length == steady_unit_length(FRAMES) == 10 + + sessions = 3 + eager_caches = [_start_session(thinker, seed=session) for session in range(sessions)] + graph_caches = [_clone_cache(cache) for cache in eager_caches] + + # 4 rounds: past grows 11 -> 21 -> 31 (crossing the 16 bucket into 32); + # by round 3 some sessions' past + unit_length (>= 40) trips the eager + # reset path (past back to 0) on both sides, at different rounds for + # different sessions. Round 1 drops a session (batch padding: 2 of 3 + # sessions ready that round). + present_by_round = [[0, 1, 2], [0, 2], [0, 1, 2], [0, 1, 2]] + with torch.inference_mode(): + for round_index, present in enumerate(present_by_round): + eager_chunks = [ + StreamingAudioChunk( + features=_mel(FRAMES, 100 * round_index + s, "cuda"), + cache=eager_caches[s], + prefix_extra_frames=2, + suffix_extra_frames=2, + ) + for s in present + ] + eager_outputs, eager_updated = encode_streaming_audio_batch( + thinker.apm, thinker.audio_projection_layer, thinker.audio_avg_pooler, eager_chunks, pool_step=POOL + ) + for s, cache in zip(present, eager_updated, strict=True): + eager_caches[s] = cache + + graph_chunks = [ + StreamingAudioChunk( + features=_mel(FRAMES, 100 * round_index + s, "cuda"), + cache=graph_caches[s], + prefix_extra_frames=2, + suffix_extra_frames=2, + ) + for s in present + ] + graph_outputs, graph_updated = graph.encode(graph_chunks) + for s, cache in zip(present, graph_updated, strict=True): + graph_caches[s] = cache + + for eager_out, graph_out in zip(eager_outputs, graph_outputs, strict=True): + assert eager_out is not None and graph_out is not None + torch.testing.assert_close(graph_out.float(), eager_out.float(), rtol=2e-2, atol=2e-2) + for s in present: + assert graph_caches[s].length == eager_caches[s].length + torch.testing.assert_close( + graph_caches[s].history(graph_caches[s].length).float(), + eager_caches[s].history(eager_caches[s].length).float(), + rtol=2e-2, + atol=2e-2, + ) + + +def _cache_with_history(thinker: _Thinker, *, past: int, seed: int) -> StreamingAudioKVCache: + """A cache whose committed history is exactly ``past`` (a multiple of the + steady unit_length, 10) positions of real, reproducible content. + + Starts from an empty (not ``None``) cache and only ever runs steady + (``prefix_extra_frames=2``) units through it -- ``encode_streaming_audio_batch`` + resets on ``cache is None``, not on an empty-but-real cache, so this never + takes the session's-first-unit path and each unit adds exactly 10. + """ + assert past % 10 == 0 + cache = StreamingAudioKVCache( + num_layers=len(thinker.apm.layers), + embed_dim=int(thinker.apm.config.d_model), + num_heads=int(thinker.apm.config.encoder_attention_heads), + max_positions=int(thinker.apm.embed_positions.weight.shape[0]), + page_positions=16, + ) + for step in range(past // 10): + mel = _mel(FRAMES, seed * 1000 + step, "cuda") + with torch.inference_mode(): + (_out,), (cache,) = encode_streaming_audio_batch( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + [StreamingAudioChunk(features=mel, cache=cache, prefix_extra_frames=2, suffix_extra_frames=2)], + pool_step=POOL, + ) + assert cache.length == past, f"expected past={past}, got {cache.length}" + return cache + + +@_cuda +def test_alternating_buckets_share_storage_without_leaking(cuda_thinker) -> None: + """Every graph takes a prefix view of one shared storage (see the module + docstring): a small-batch/small-cache-bucket graph and a large one + physically overlap in that storage. Replay them interleaved, out of + capture order, and check neither ever reads the other's leftovers.""" + thinker = cuda_thinker + graph = StreamingAudioGraphEncoder( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + unit_frames=FRAMES, + pool_step=POOL, + batch_sizes=(1, 4), + cache_buckets=(16, 32), + page_positions=16, + ) + with torch.inference_mode(): + graph.capture() + # (1, 16) and (4, 32) are both prefixes of the same shared storage; (1, + # 16)'s whole footprint sits inside (4, 32)'s (same base pointer). + assert graph._views(1, 16)[1].data_ptr() == graph._views(4, 32)[1].data_ptr() + + # Fixed pasts (10 and 20: both multiples of the unit_length-10 steady + # step, so they land exactly, no drift) that resolve to distinct keys: + # 1 session at past 10 -> (1, 16); 4 sessions at past 20 -> (4, 32). + eager_small = [_cache_with_history(thinker, past=10, seed=11)] + eager_big = [_cache_with_history(thinker, past=20, seed=200 + i) for i in range(4)] + + def replay_round(base_caches: list[StreamingAudioKVCache], *, seed_base: int): + # Fresh clones every round, on both sides: the point is to replay the + # *same* (batch, cache bucket) key repeatedly, interleaved with the + # other key, so past must not drift between rounds. + eager_caches = [_clone_cache(c) for c in base_caches] + graph_caches = [_clone_cache(c) for c in base_caches] + + def chunks(caches: list[StreamingAudioKVCache]) -> list[StreamingAudioChunk]: + return [ + StreamingAudioChunk( + features=_mel(FRAMES, seed_base + s, "cuda"), + cache=caches[s], + prefix_extra_frames=2, + suffix_extra_frames=2, + ) + for s in range(len(caches)) + ] + + with torch.inference_mode(): + eager_outputs, eager_updated = encode_streaming_audio_batch( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + chunks(eager_caches), + pool_step=POOL, + ) + graph_outputs, graph_updated = graph.encode(chunks(graph_caches)) + for s in range(len(base_caches)): + torch.testing.assert_close(graph_outputs[s].float(), eager_outputs[s].float(), rtol=2e-2, atol=2e-2) + assert graph_updated[s].length == eager_updated[s].length + torch.testing.assert_close( + graph_updated[s].history(graph_updated[s].length).float(), + eager_updated[s].history(eager_updated[s].length).float(), + rtol=2e-2, + atol=2e-2, + ) + + # Interleave (1, 16) and (4, 32) replays, out of the order they were + # captured in, several times: if shared storage leaked between them, one + # of these rounds would see the other key's stale content instead of its + # own freshly refilled inputs. + for round_index in range(4): + replay_round(eager_small, seed_base=1000 + round_index) + replay_round(eager_big, seed_base=2000 + round_index) + replay_round(eager_big, seed_base=3000 + round_index) + replay_round(eager_small, seed_base=4000 + round_index) + + +@_cuda +def test_graph_batch_size_one_is_near_bit_exact_with_eager(cuda_thinker) -> None: + """No batch padding and no history padding (cache bucket == true past): tightest case.""" + thinker = cuda_thinker + graph = StreamingAudioGraphEncoder( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + unit_frames=FRAMES, + pool_step=POOL, + batch_sizes=(1,), + cache_buckets=(11, 21, 31), + page_positions=16, + ) + with torch.inference_mode(): + graph.capture() + eager_cache = _start_session(thinker, seed=0) # length 11 + graph_cache = _clone_cache(eager_cache) + + with torch.inference_mode(): + eager_chunk = StreamingAudioChunk( + features=_mel(FRAMES, 1, "cuda"), cache=eager_cache, prefix_extra_frames=2, suffix_extra_frames=2 + ) + (eager_out,), (eager_cache,) = encode_streaming_audio_batch( + thinker.apm, thinker.audio_projection_layer, thinker.audio_avg_pooler, [eager_chunk], pool_step=POOL + ) + graph_chunk = StreamingAudioChunk( + features=_mel(FRAMES, 1, "cuda"), cache=graph_cache, prefix_extra_frames=2, suffix_extra_frames=2 + ) + (graph_out,), (graph_cache,) = graph.encode([graph_chunk]) + + torch.testing.assert_close(graph_out, eager_out) + assert graph_cache.length == eager_cache.length + torch.testing.assert_close(graph_cache.history(graph_cache.length), eager_cache.history(eager_cache.length)) + + +@_cuda +def test_graph_leaves_ineligible_rows_to_the_eager_fallback(cuda_thinker) -> None: + thinker = cuda_thinker + graph = StreamingAudioGraphEncoder( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + unit_frames=FRAMES, + pool_step=POOL, + batch_sizes=(2,), + cache_buckets=(32,), + page_positions=16, + ) + with torch.inference_mode(): + graph.capture() + with torch.inference_mode(): + # A session's first unit (cache=None, prefix=0): not steady, must + # still succeed via the eager fallback inside encode(). + first_chunk = StreamingAudioChunk( + features=_mel(FRAMES, 5, "cuda"), cache=None, prefix_extra_frames=0, suffix_extra_frames=2 + ) + (out,), (cache,) = graph.encode([first_chunk]) + assert out is not None + assert isinstance(cache, StreamingAudioKVCache) and cache.length > 0 + + +@_cuda +def test_graph_matches_eager_with_a_non_pool_step_aligned_unit_length(cuda_thinker) -> None: + """Real checkpoints do not guarantee unit_length % pool_step == 0 (see + ``steady_pooled_length``'s cap/actual ``min``); this checks the graph + still matches eager when it is not (unit_frames=25 -> unit_length 11).""" + thinker = cuda_thinker + odd_frames = FRAMES + 1 + graph = StreamingAudioGraphEncoder( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + unit_frames=odd_frames, + pool_step=POOL, + batch_sizes=(1,), + cache_buckets=(16,), + page_positions=16, + ) + with torch.inference_mode(): + graph.capture() + assert graph.unit_length == steady_unit_length(odd_frames) == 11 + assert graph.pooled_length == steady_pooled_length(odd_frames, 11, POOL) == 2 + + eager_cache = _start_session(thinker, seed=0, frames=odd_frames) + graph_cache = _clone_cache(eager_cache) + with torch.inference_mode(): + eager_chunk = StreamingAudioChunk( + features=_mel(odd_frames, 1, "cuda"), cache=eager_cache, prefix_extra_frames=2, suffix_extra_frames=2 + ) + (eager_out,), _ = encode_streaming_audio_batch( + thinker.apm, thinker.audio_projection_layer, thinker.audio_avg_pooler, [eager_chunk], pool_step=POOL + ) + graph_chunk = StreamingAudioChunk( + features=_mel(odd_frames, 1, "cuda"), cache=graph_cache, prefix_extra_frames=2, suffix_extra_frames=2 + ) + (graph_out,), _ = graph.encode([graph_chunk]) + assert eager_out.shape[0] == 2 + torch.testing.assert_close(graph_out, eager_out) + + +def _steady_round(thinker: _Thinker, graph: StreamingAudioGraphEncoder, cache, seed: int): + chunk = StreamingAudioChunk( + features=_mel(FRAMES, seed, "cpu"), cache=cache, prefix_extra_frames=2, suffix_extra_frames=2 + ) + with torch.inference_mode(): + (out,), (cache,) = graph.encode([chunk]) + return out, cache + + +@_cuda +def test_pinned_h2d_is_bitwise_identical_and_does_not_block(cuda_thinker) -> None: + """Host mel staged through pinned memory: same bits, and no blocking copy in the replay.""" + thinker = cuda_thinker + graphs = [] + for pinned in (False, True): + graph = StreamingAudioGraphEncoder( + thinker.apm, + thinker.audio_projection_layer, + thinker.audio_avg_pooler, + unit_frames=FRAMES, + pool_step=POOL, + batch_sizes=(1,), + cache_buckets=(11, 21, 31), + page_positions=16, + pinned_h2d=pinned, + ) + with torch.inference_mode(): + graph.capture() + graphs.append(graph) + base_cache = _start_session(thinker, seed=0) + blocking_out, blocking_cache = _steady_round(thinker, graphs[0], _clone_cache(base_cache), seed=7) + torch.accelerator.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + pinned_out, pinned_cache = _steady_round(thinker, graphs[1], _clone_cache(base_cache), seed=7) + finally: + torch.cuda.set_sync_debug_mode("default") + torch.accelerator.synchronize() + assert torch.equal(pinned_out, blocking_out) + assert torch.equal(pinned_cache.history(pinned_cache.length), blocking_cache.history(blocking_cache.length)) diff --git a/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_resident.py b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_resident.py new file mode 100644 index 00000000000..5e7a33b393c --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_resident.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Resident mode of the streaming audio encoder graphs (``resident_slots > 0``). + +Sessions keep their cache in a slot pool the graphs read and write in place. +The parity tests replay the same session schedule through the copy-in graph +encoder (the default) and the resident one and require identical bits: the +resident graph gathers the history into a temporary laid out exactly like the +shared storage, so attention sees the same values at every unmasked position. +""" + +from __future__ import annotations + +import gc + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMWhisperEncoder, + MultiModalProjector, +) +from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import ( + StreamingAudioChunk, + StreamingAudioKVCache, + encode_streaming_audio_batch, +) +from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder_graph import ( + PooledStreamingAudioKVCache, + StreamingAudioGraphEncoder, + StreamingAudioKVSlotPool, +) + +pytestmark = [pytest.mark.core_model] + +N_MELS = 16 +POOL = 5 +FRAMES = 24 # steady unit_length 10; a first unit (prefix 0) is 11 +MAX_POSITIONS = 40 + + +class _Audio: + def __init__(self, implementation: str, device: str = "cuda") -> None: + from transformers.models.whisper.modeling_whisper import WhisperConfig + + torch.manual_seed(0) + config = WhisperConfig( + num_mel_bins=N_MELS, + d_model=32, + encoder_layers=2, + encoder_attention_heads=2, + encoder_ffn_dim=128, + max_source_positions=MAX_POSITIONS, + dropout=0.0, + attention_dropout=0.0, + ) + config._attn_implementation = implementation + self.apm = MiniCPMWhisperEncoder(config).eval().to(device=device, dtype=torch.bfloat16) + self.proj = MultiModalProjector(in_dim=config.encoder_ffn_dim // 4, out_dim=24).to( + device=device, dtype=torch.bfloat16 + ) + self.pooler = nn.AvgPool1d(POOL, stride=POOL) + + def graph(self, *, resident_slots: int, batch_sizes, cache_buckets) -> StreamingAudioGraphEncoder: + graph = StreamingAudioGraphEncoder( + self.apm, + self.proj, + self.pooler, + unit_frames=FRAMES, + pool_step=POOL, + batch_sizes=batch_sizes, + cache_buckets=cache_buckets, + page_positions=16, + resident_slots=resident_slots, + ) + with torch.inference_mode(): + graph.capture() + return graph + + +def _mel(frames: int, seed: int) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randn((1, N_MELS, frames), generator=generator) + + +def _chunk(cache, *, seed: int, frames: int = FRAMES) -> StreamingAudioChunk: + first = cache is None + return StreamingAudioChunk( + features=_mel(frames, seed), + cache=cache, + prefix_extra_frames=0 if first else 2, + suffix_extra_frames=2, + ) + + +def _history(cache: StreamingAudioKVCache) -> torch.Tensor: + return cache.history(cache.length) + + +# Rounds of (session, frames): sessions drop in and out (batch padding), a +# session starts late (first unit in a later round), one round carries a +# non-steady unit (frames + 2: eager; round 7's, at past 30, resets eagerly), +# and enough rounds pass for every session to cross the cache buckets, go past +# the largest one (eager in place) and reset (past back to 0) in the graph. +_SCHEDULE: list[list[tuple[int, int]]] = [ + [(0, FRAMES), (1, FRAMES), (2, FRAMES)], + [(0, FRAMES), (2, FRAMES)], + [(0, FRAMES), (1, FRAMES), (2, FRAMES), (3, FRAMES)], + [(0, FRAMES), (1, FRAMES + 2), (2, FRAMES), (3, FRAMES)], + [(0, FRAMES), (1, FRAMES), (3, FRAMES)], + [(0, FRAMES), (1, FRAMES), (2, FRAMES), (3, FRAMES)], + [(1, FRAMES + 2), (2, FRAMES), (3, FRAMES)], + [(0, FRAMES), (1, FRAMES), (2, FRAMES + 2), (3, FRAMES)], + [(0, FRAMES), (1, FRAMES), (2, FRAMES), (3, FRAMES)], + [(0, FRAMES), (2, FRAMES), (3, FRAMES)], +] + + +def _run(graph: StreamingAudioGraphEncoder, sessions: int = 4): + caches: list[StreamingAudioKVCache | None] = [None] * sessions + rounds = [] + with torch.inference_mode(): + for round_index, present in enumerate(_SCHEDULE): + chunks = [_chunk(caches[s], seed=1000 * round_index + s, frames=frames) for s, frames in present] + outputs, updated = graph.encode(chunks) + for (s, _), cache in zip(present, updated, strict=True): + caches[s] = cache + rounds.append( + ( + [out.clone() for out in outputs], + {s: (caches[s].length, _history(caches[s]).clone()) for s, _ in present}, + ) + ) + torch.accelerator.synchronize() + return rounds, caches + + +@pytest.fixture(params=["sdpa", "eager"]) +def audio(request): + if not torch.cuda.is_available(): + pytest.skip("CUDA graph capture needs a GPU") + return _Audio(request.param) + + +@pytest.mark.cuda +@pytest.mark.parametrize( + "cache_buckets", + [(16, 32), (16,)], + ids=["two-buckets", "past-the-largest-bucket-goes-eager"], +) +def test_resident_graph_is_bitwise_identical_to_the_copy_in_graph(audio, cache_buckets) -> None: + batch_sizes = (1, 2, 4) + copy_in = audio.graph(resident_slots=0, batch_sizes=batch_sizes, cache_buckets=cache_buckets) + resident = audio.graph(resident_slots=4, batch_sizes=batch_sizes, cache_buckets=cache_buckets) + assert copy_in._cache_storage is not None and resident._cache_storage is None + + expected, _ = _run(copy_in) + got, caches = _run(resident) + assert all(isinstance(cache, PooledStreamingAudioKVCache) for cache in caches) + lengths = set() + for (exp_out, exp_caches), (got_out, got_caches) in zip(expected, got, strict=True): + for e, g in zip(exp_out, got_out, strict=True): + assert torch.equal(e, g) + assert exp_caches.keys() == got_caches.keys() + for s, (length, history) in exp_caches.items(): + assert got_caches[s][0] == length + assert torch.equal(got_caches[s][1], history) + lengths.add(length) + # Graph resets (a steady unit from past 0: length 10) and eager resets (a + # non-steady unit from past 0: 11, also every first unit) both happened. + assert 10 in lengths and 11 in lengths + + +@pytest.mark.cuda +def test_resident_slots_are_released_with_their_session(audio) -> None: + resident = audio.graph(resident_slots=3, batch_sizes=(1, 2, 4), cache_buckets=(16, 32)) + pool = resident._pool + assert isinstance(pool, StreamingAudioKVSlotPool) and pool.free_slots == 3 + _, caches = _run(resident, sessions=4) + # Four sessions, three slots: the fourth to start got a paged cache. + pooled = [c for c in caches if isinstance(c, PooledStreamingAudioKVCache)] + assert len(pooled) == 3 and pool.free_slots == 0 + assert sorted(c.slot for c in pooled) == [0, 1, 2] + assert sum(type(c) is StreamingAudioKVCache for c in caches) == 1 + del pooled, caches + gc.collect() + assert pool.free_slots == 3 + + +@pytest.mark.cuda +def test_paged_rows_under_a_full_pool_match_the_eager_batch(audio) -> None: + """A session that found the pool full runs eagerly, as encode_streaming_audio_batch would.""" + resident = audio.graph(resident_slots=1, batch_sizes=(1, 2), cache_buckets=(16, 32)) + with torch.inference_mode(): + (_, _), (held, paged) = resident.encode([_chunk(None, seed=1), _chunk(None, seed=2)]) + assert isinstance(held, PooledStreamingAudioKVCache) and type(paged) is StreamingAudioKVCache + reference = StreamingAudioKVCache( + num_layers=paged.num_layers, + embed_dim=paged.embed_dim, + num_heads=paged.num_heads, + max_positions=paged.max_positions, + page_positions=16, + ) + reference.reserve(paged.length, dtype=torch.bfloat16, device=torch.device("cuda")) + reference.commit(0, _history(paged)) + reference.length = paged.length + (got,), _ = resident.encode([_chunk(paged, seed=3)]) + (want,), _ = encode_streaming_audio_batch( + audio.apm, audio.proj, audio.pooler, [_chunk(reference, seed=3)], pool_step=POOL + ) + assert torch.equal(got, want) + assert torch.equal(_history(paged), _history(reference)) + + +@pytest.mark.cuda +def test_resident_storage_replaces_the_shared_history_storage(audio) -> None: + resident = audio.graph(resident_slots=4, batch_sizes=(1, 2, 4), cache_buckets=(16, 32)) + storage = resident._pool.storage + # [slots, layers, 2, max_positions + unit_length tail, d] + assert tuple(storage.shape) == (4, 2, 2, MAX_POSITIONS + resident.unit_length, 32) + assert resident._cache_storage is None diff --git a/tests/model_executor/models/minicpmo_4_5/test_turn_audio_host_sync.py b/tests/model_executor/models/minicpmo_4_5/test_turn_audio_host_sync.py new file mode 100644 index 00000000000..266f40c7705 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_turn_audio_host_sync.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Turn-mode (full-utterance) Stage-0 audio: ``get_audio_hidden_states`` over a ragged batch. + +The assembly loop reads the pooled lengths back once (``.tolist()``) instead of +slicing with a device scalar per row; each row must keep its predicted length. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from transformers.models.whisper.modeling_whisper import WhisperConfig + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, + MiniCPMWhisperEncoder, + MultiModalProjector, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +N_MELS = 16 +POOL = 5 + + +class _Thinker: + """The audio half of the thinker: real modules, the real method under test.""" + + get_audio_hidden_states = MiniCPMO45OmniLLMForConditionalGeneration.get_audio_hidden_states + subsequent_chunk_mask = MiniCPMO45OmniLLMForConditionalGeneration.subsequent_chunk_mask + _get_feat_extract_output_lengths = MiniCPMO45OmniLLMForConditionalGeneration._get_feat_extract_output_lengths + + def __init__(self, *, attn_implementation: str = "sdpa", chunk_length: float = 0.0) -> None: + torch.manual_seed(0) + config = WhisperConfig( + num_mel_bins=N_MELS, + d_model=32, + encoder_layers=2, + encoder_attention_heads=2, + encoder_ffn_dim=128, # the projection takes ffn // 4 == d_model, matching the checkpoint + max_source_positions=200, + dropout=0.0, + attention_dropout=0.0, + ) + config._attn_implementation = attn_implementation + self.apm = MiniCPMWhisperEncoder(config).eval() + self.audio_projection_layer = MultiModalProjector(in_dim=config.encoder_ffn_dim // 4, out_dim=24) + self.audio_avg_pooler = nn.AvgPool1d(POOL, stride=POOL) + self.audio_encoder_layer = -1 + self.config = SimpleNamespace(audio_pool_step=POOL, audio_chunk_length=chunk_length) + + +def _feats(frames: int, seed: int) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randn((N_MELS, frames), generator=generator) + + +@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"]) +def test_batched_multi_request_matches_per_row_reference(attn_implementation: str) -> None: + """Several ragged-length audios in one call (mirrors several concurrent + requests grouped into one encoder call) assemble to rows of the length + ``_get_feat_extract_output_lengths`` predicts.""" + thinker = _Thinker(attn_implementation=attn_implementation) + lens = [131, 214, 98, 172] + feats = [_feats(n, seed=i) for i, n in enumerate(lens)] + data = { + "audio_features": feats, + "audio_feature_lens": [torch.tensor([n]) for n in lens], + } + with torch.no_grad(): + outputs = thinker.get_audio_hidden_states(data) + + assert len(outputs) == len(lens) + for out, n in zip(outputs, lens, strict=True): + _, expected_len = thinker._get_feat_extract_output_lengths(torch.tensor([n])) + assert out.shape[0] == int(expected_len.item()) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/incremental_fbank.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/incremental_fbank.py new file mode 100644 index 00000000000..6282110741f --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/incremental_fbank.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Incremental Stage-0 streaming fbank (HF override ``duplex_incremental_fbank``, default off). + +The remote-code ``StreamingMelProcessorExact._extract_full`` runs the Whisper +log-mel front end (``MiniCPMAAudioProcessor._torch_extract_fbank_features``, +on the CPU) over the session's whole streaming buffer, up to its 30 s slide +trigger, on every unit, although only the newest unit's frames change. The +incremental version keeps the power spectrum ``|STFT|^2`` of the frames it +already computed and recomputes only the right edge (the frames whose window +reaches the new samples or the reflect pad) and, after a slide, the new left +edge. The mel projection, clamp, log10, dynamic range (global max) and +scaling still run over the whole buffer with the original's shapes and memory +layout, so the features are bitwise equal to ``_extract_full``: + +* The mel projection is not row independent: MKL computes the rows of a + partial register block with another kernel, so the original itself gives + one frame different mel values at different buffer lengths. It is not + cached; it is ~10% of the original's time. +* The STFT and ``abs() ** 2`` are frame independent, except that a batch's + last frames or elements may take a remainder path. A recompute therefore + starts on a multiple of ``_ALIGN_FRAMES`` (its remainder then lines up + with the original's) and its last ``_ALIGN_FRAMES`` frames are not kept. +""" + +from __future__ import annotations + +import copy +from typing import Any + +import numpy as np +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Recompute start granularity and the trailing frames of a recompute that are +# not kept (>= the widest remainder of the FFT batch and the abs() vector loop); +# also the left-edge recompute after a slide (>= the 2 frames reading the pad). +_ALIGN_FRAMES = 16 +_FEATURE_EXTRACTOR_ATTRS = ( + "n_fft", + "hop_length", + "n_samples", + "sampling_rate", + "mel_filters", + "set_spac_log_norm", + "dynamic_log_norm", + "dynamic_range_db", + "log_floor_db", +) +_INCREMENTAL_CLASSES: dict[type, type] = {} +_SELF_CHECKED: dict[tuple[type, type], bool] = {} + + +def _stft_power( + buffer: np.ndarray, + start: int, + stop: int, + n_fft: int, + hop: int, + window: torch.Tensor, +) -> np.ndarray: + """``|STFT|^2`` of frames ``[start, stop)`` of ``buffer``, framed as the original's centered STFT. + + The batch is ``[start, stop]`` with its last frame dropped, like the + original's ``stft(...)[..., :-1].abs() ** 2``. The reflect padding at a + buffer edge is built explicitly (``np.pad`` reflect is ``torch.stft``'s). + Returns ``(stop - start, n_fft // 2 + 1)`` rows. + """ + half = n_fft // 2 + lo, hi = start * hop - half, stop * hop + half + length = len(buffer) + segment = buffer[max(lo, 0) : min(hi, length)] + if lo < 0 or hi > length: + segment = np.pad(segment, (max(0, -lo), max(0, hi - length)), mode="reflect") + spec = torch.stft( + torch.from_numpy(segment)[None], + n_fft=n_fft, + hop_length=hop, + window=window, + center=False, + return_complex=True, + ) + magnitudes = spec[..., :-1].abs() ** 2 # [1, n_freq, stop - start], frame-major in memory + return magnitudes[0].T.numpy() + + +def _log_mel(power: torch.Tensor, mel_filters: torch.Tensor, feature_extractor: Any) -> torch.Tensor: + """``_torch_extract_fbank_features`` from the ``|STFT|^2`` rows, op for op and layout for layout.""" + magnitudes = power.unsqueeze(0).transpose(1, 2) # [1, n_freq, T], strides of stft(...)[..., :-1] + mel_spec = mel_filters.T @ magnitudes + log_spec = torch.clamp(mel_spec, min=1e-10).log10() + if feature_extractor.dynamic_log_norm: + max_val_t = log_spec.max(dim=2, keepdim=True)[0] + max_val_bt = max_val_t.max(dim=1, keepdim=True)[0] + log_spec = torch.maximum(log_spec, max_val_bt - feature_extractor.dynamic_range_db) + else: + floor_tensor = torch.tensor(feature_extractor.log_floor_db, dtype=log_spec.dtype) + log_spec = torch.maximum(log_spec, floor_tensor) + return (log_spec + 4.0) / 4.0 + + +class _IncrementalFbankMixin: + """Overrides of the remote ``StreamingMelProcessorExact``, mixed in by ``enable_incremental_fbank``.""" + + #: ``|STFT|^2`` rows, frame t of the current buffer at row t; rows + #: ``[0, _ifb_valid)`` are exact for the current buffer. + _ifb_power: np.ndarray | None = None + _ifb_valid = 0 + #: ``left_samples_dropped`` and the buffer length when the rows were kept. + _ifb_dropped = 0 + _ifb_length = 0 + #: The ``n_fft`` samples ending where the last kept frame's window ends. + _ifb_guard: np.ndarray | None = None + _ifb_window: torch.Tensor | None = None + _ifb_mel_filters: torch.Tensor | None = None + + def reset(self) -> None: + super().reset() + self._ifb_valid = 0 + + def restore_snapshot(self, snapshot: dict) -> None: + super().restore_snapshot(snapshot) + self._ifb_valid = 0 + + def _extract_full(self) -> torch.Tensor: + fe = self.feature_extractor + buffer = self.buffer + if not ( + isinstance(buffer, np.ndarray) + and buffer.ndim == 1 + and buffer.dtype == np.float32 + and fe.n_fft <= len(buffer) <= fe.n_samples + and getattr(fe, "dither", 0.0) == 0.0 + and fe.sampling_rate == self.sample_rate + ): + # Too short (the original raises), truncated, dithered or misconfigured: the original. + self._ifb_valid = 0 + return super()._extract_full() + # The original's log-floor switch, side effect on the feature extractor included. + if len(buffer) < 5 * self.sample_rate: + fe.set_spac_log_norm(log_floor_db=-10) + else: + fe.set_spac_log_norm(dynamic_range_db=8) + return _log_mel(self._ifb_update_power(buffer, fe), self._ifb_mel_filters, fe) + + def _ifb_update_power(self, buffer: np.ndarray, fe: Any) -> torch.Tensor: + n_fft, hop = int(fe.n_fft), int(fe.hop_length) + if self._ifb_window is None: + self._ifb_window = torch.hann_window(n_fft) + self._ifb_mel_filters = torch.from_numpy(fe.mel_filters).to(torch.float32) + power = self._ifb_power + if power is None or power.ctypes.data % 64: + # Rows from a torch allocation: 64-byte aligned like the original's fresh + # |X|^2 tensor, which the GEMM reads (a deepcopy's numpy array may not be). + aligned = torch.empty(int(fe.n_samples) // hop + 1, n_fft // 2 + 1, dtype=torch.float32).numpy() + if power is not None: + aligned[: self._ifb_valid] = power[: self._ifb_valid] + self._ifb_power = power = aligned + frames = len(buffer) // hop # the original drops the STFT's last frame + valid = self._ifb_rebase(buffer, n_fft, hop) + start = valid // _ALIGN_FRAMES * _ALIGN_FRAMES + power[start:frames] = _stft_power(buffer, start, frames, n_fft, hop, self._ifb_window) + # Keep frames whose window lies inside the buffer, minus the recompute's tail. + stable = (len(buffer) - n_fft // 2) // hop + 1 + valid = min(stable, frames - _ALIGN_FRAMES) + guard_end = (valid - 1) * hop + n_fft // 2 + if guard_end < n_fft: + valid = 0 + else: + self._ifb_guard = buffer[guard_end - n_fft : guard_end].copy() + self._ifb_valid = valid + self._ifb_dropped = int(self.left_samples_dropped) + self._ifb_length = len(buffer) + return torch.from_numpy(power[:frames]) + + def _ifb_rebase(self, buffer: np.ndarray, n_fft: int, hop: int) -> int: + """Rows still exact for ``buffer``: the kept rows (moved down after a slide), or 0.""" + valid = self._ifb_valid + if valid == 0: + return 0 + shift = int(self.left_samples_dropped) - self._ifb_dropped + guard_end = (valid - 1) * hop + n_fft // 2 - shift + if ( + shift < 0 + or shift % hop + or len(buffer) + shift < self._ifb_length + or guard_end < n_fft + or not np.array_equal(buffer[guard_end - n_fft : guard_end], self._ifb_guard) + ): + # Not the same stream (or not a hop-aligned slide of it): start over. + return 0 + if shift == 0: + return valid + drop = shift // hop + valid -= drop + if valid <= _ALIGN_FRAMES: + return 0 + power = self._ifb_power + power[:valid] = power[drop : drop + valid] + # The new left edge's frames read the reflect pad. + power[:_ALIGN_FRAMES] = _stft_power(buffer, 0, _ALIGN_FRAMES, n_fft, hop, self._ifb_window) + return valid + + +def _self_check(mel_processor: Any, incremental_cls: type) -> bool: # Any: remote-code processor + """Compare the incremental and the original features on a probe stream (with a slide) once per class.""" + fe = mel_processor.feature_extractor + saved = {name: getattr(fe, name) for name in ("dynamic_log_norm", "dynamic_range_db", "log_floor_db")} + probe = copy.copy(mel_processor) + probe.__class__ = incremental_cls + hop = int(fe.hop_length) + audio = (0.1 * np.random.default_rng(0).standard_normal(8 * int(mel_processor.sample_rate))).astype(np.float32) + try: + # Below and above the 5 s dynamic-range switch, then a slide of 100 hops. + for dropped, end in ((0, len(audio) // 3 + 37), (0, len(audio) - 11), (100 * hop, len(audio))): + probe.buffer, probe.left_samples_dropped = audio[dropped:end], dropped + expected = type(mel_processor)._extract_full(probe) + if not torch.equal(incremental_cls._extract_full(probe), expected): + return False + return True + except Exception: + logger.warning("MiniCPM-o incremental fbank self-check raised", exc_info=True) + return False + finally: + for name, value in saved.items(): + setattr(fe, name, value) + + +def enable_incremental_fbank(mel_processor: Any) -> bool: # Any: remote-code StreamingMelProcessorExact + """Switch one session's streaming mel processor to the incremental ``_extract_full`` (idempotent). + + The instance's class becomes a subclass of its remote class with + ``_IncrementalFbankMixin`` in front, so copies keep the behavior. Returns + False, leaving the instance on the full recompute, for a processor of + another shape or one whose features differ from the original's on the + probe stream. + """ + if isinstance(mel_processor, _IncrementalFbankMixin): + return True + cls = type(mel_processor) + fe = getattr(mel_processor, "feature_extractor", None) + if not callable(getattr(cls, "_extract_full", None)) or not all( + hasattr(fe, name) for name in _FEATURE_EXTRACTOR_ATTRS + ): + return False + if not all(hasattr(mel_processor, name) for name in ("buffer", "sample_rate", "left_samples_dropped")): + return False + incremental_cls = _INCREMENTAL_CLASSES.get(cls) + if incremental_cls is None: + incremental_cls = type(f"Incremental{cls.__name__}", (_IncrementalFbankMixin, cls), {"__module__": __name__}) + _INCREMENTAL_CLASSES[cls] = incremental_cls + key = (cls, type(fe)) + if key not in _SELF_CHECKED: + _SELF_CHECKED[key] = _self_check(mel_processor, incremental_cls) + if not _SELF_CHECKED[key]: + logger.warning( + "MiniCPM-o incremental fbank differs from %s._extract_full on the probe stream; " + "the streaming fbank stays on the full recompute", + cls.__name__, + ) + if not _SELF_CHECKED[key]: + return False + mel_processor.__class__ = incremental_cls + return True diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py index 061ca14f996..1af659fbe88 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py @@ -535,6 +535,11 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin): """MiniCPM-owned sampling policy, append planning, session state and output projection.""" plugin_id = "minicpmo45" + # Stage 0 builds the session context once and keeps it in its per-session + # state. Re-sending the conversation and reference-audio base64 on every + # one-second unit only adds host/IPC work; ModelChannel resends it when a + # request starts or the session config generation changes. + append_configs_are_session_static = True private_runtime_config_keys = PRIVATE_RUNTIME_CONFIG_KEYS silence_continuation_samples = 16000 @@ -664,6 +669,38 @@ def decide_output( metadata=metadata, ) + def stage0_naturally_ended( + self, + *, + segment_finished: bool, + segment_token_ids: tuple[int, ...], + segment_output_metadata: Mapping[str, object], + output: object, + ) -> bool: + if not segment_finished: + return False + completion = _first_completion(output) + output_metadata = _multimodal_output(output, completion) + special_token_ids = _special_token_ids(segment_output_metadata) + special_token_ids.update(_special_token_ids(output_metadata)) + token_ids = set(segment_token_ids) + token_ids.update(_completion_token_ids(completion)) + natural_end_ids = { + special_token_ids.get("turn_eos_token_id", -1), + special_token_ids.get("tts_eos_token_id", -1), + } + if token_ids.intersection(natural_end_ids): + return True + return any( + _coerce_int(value) == 1 + for value in ( + output_metadata.get("turn_end"), + output_metadata.get("end_of_turn"), + output_metadata.get("meta.turn_end"), + output_metadata.get("meta.end_of_turn"), + ) + ) + # ---- session policy ---- def create_session_state(self) -> MiniCPMO45ServingSessionState: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py index c65c8dadb80..25a31e6b307 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py @@ -5,15 +5,19 @@ import base64 import copy +import os import time -from collections.abc import Callable +from collections.abc import Callable, Generator +from concurrent.futures import ThreadPoolExecutor from contextlib import suppress from dataclasses import dataclass, field from threading import Lock +from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast import numpy as np from numpy.typing import NDArray +from vllm.logger import init_logger from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy @@ -22,9 +26,88 @@ from PIL import Image from torch import nn +logger = init_logger(__name__) + _MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS _MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS _MINICPMO45_PROCESSOR_LOAD_LOCK = Lock() +_MINICPMO45_FRAME_POOL_LOCK = Lock() +_MINICPMO45_FRAME_POOL: ThreadPoolExecutor | None = None + + +class _FbankStats: + """``duplex_fbank_stats`` (HF override, default off): time the streaming mel front end. + + ``process_audio_streaming`` recomputes the STFT/mel over the session's whole + streaming buffer (up to its 30 s slide trigger) on every unit, on the CPU + (only the new frames' STFT with ``duplex_incremental_fbank``, then tagged + "incremental"). Every ``every`` calls one line reports the call time and + buffer length; the returned features are untouched. + """ + + def __init__(self, every: int = 200, *, incremental: bool = False) -> None: + self.every = every + self._message = ( + "Stage-0 streaming fbank (incremental)" if incremental else "Stage-0 streaming fbank" + ) + ": %d calls p50 %.2f ms p90 %.2f ms max %.2f ms; buffer p50 %.1f s max %.1f s" + self._lock = Lock() + self._ms: list[float] = [] + self._buffer_s: list[float] = [] + self.calls = 0 + + def record(self, elapsed_s: float, buffer_s: float | None) -> None: + with self._lock: + self.calls += 1 + self._ms.append(elapsed_s * 1000.0) + if buffer_s is not None: + self._buffer_s.append(buffer_s) + if len(self._ms) < self.every: + return + ms, buf = sorted(self._ms), sorted(self._buffer_s) + self._ms, self._buffer_s = [], [] + logger.info( + self._message, + len(ms), + ms[len(ms) // 2], + ms[min(len(ms) - 1, int(len(ms) * 0.9))], + ms[-1], + buf[len(buf) // 2] if buf else -1.0, + buf[-1] if buf else -1.0, + ) + + +def _streaming_buffer_seconds(processor: Any) -> float | None: # Any: remote-code processor + mel_processor = getattr(processor, "_streaming_mel_processor", None) + buffer = getattr(mel_processor, "buffer", None) + if buffer is None: + return None + try: + samples = len(buffer) + except TypeError: + return None + return samples / float(getattr(mel_processor, "sample_rate", 16000) or 16000) + + +def _frame_preprocess_pool() -> ThreadPoolExecutor | None: + """Shared pool for camera-frame JPEG decode + slicing, or ``None`` to stay inline. + + PIL decode/resize and the numpy/torch patch reshape release the GIL, so a + step's frames preprocess in parallel instead of back to back on the + worker thread. ``MINICPMO_VISION_PREPROCESS_THREADS`` sizes it (<=1: inline). + """ + global _MINICPMO45_FRAME_POOL + if _MINICPMO45_FRAME_POOL is None: + try: + workers = int(os.environ.get("MINICPMO_VISION_PREPROCESS_THREADS", "4")) + except ValueError: + workers = 4 + workers = min(workers, os.cpu_count() or 1) + if workers <= 1: + return None + with _MINICPMO45_FRAME_POOL_LOCK: + if _MINICPMO45_FRAME_POOL is None: + _MINICPMO45_FRAME_POOL = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="minicpmo45-frames") + return _MINICPMO45_FRAME_POOL @dataclass @@ -55,8 +138,13 @@ class _MiniCPMO45Stage0SessionState: prepared_inputs_embeds: torch.Tensor | None = None prepared_input_token_ids: list[int] = field(default_factory=list) prepared_result: dict[str, object] = field(default_factory=dict) - #: Opaque transformers cache, handed back to the model unchanged. + #: Streaming-encoder state: a ``StreamingAudioKVCache`` on the batched + #: path, else the opaque transformers cache handed back to the model. audio_past_key_values: object | None = None + #: Append identity of the last prefill that did work (successful or not). + attempted_append_identity: tuple[int | None, int] | None = None + #: This step's batched prefill result, taken by the append's own preprocess. + staged_prefill: tuple[tuple[int | None, int], dict[str, object]] | None = None pending_terminator_token: int | None = None last_terminator_token: int | None = None pending_speech_context: bool = False @@ -65,6 +153,18 @@ class _MiniCPMO45Stage0SessionState: generated_tokens: list[int] = field(default_factory=list) +@dataclass(frozen=True, slots=True) +class _MiniCPMO45AudioEncodeRequest: + """One streaming unit a prefill needs encoded; ``_stage_prefill_steps`` yields it.""" + + #: Any: transformers BatchFeature (or the dict fallback of _process_streaming_audio). + batch_feature: Any + state: _MiniCPMO45Stage0SessionState | None + + +_PrefillSteps = Generator[_MiniCPMO45AudioEncodeRequest, "torch.Tensor | None", dict[str, object]] + + class MiniCPMO45Stage0DuplexRuntime: """Build scheduler-owned MiniCPM-o 4.5 Stage0 duplex inputs.""" @@ -87,6 +187,9 @@ def __init__(self, stage_model: nn.Module, *, model_path: str | None = None, dev self.model_path = model_path self.device = device self.sessions: dict[tuple[str, int], _MiniCPMO45Stage0SessionState] = {} + # (session_id, epoch, seq) -> (frame payload, per-frame embeddings), + # filled by prefetch_vision for the current runner step. + self._prefetched_vision: dict[tuple[str, int | None, int], tuple[tuple[object, ...], list[Any]]] = {} self.thinker = getattr(stage_model, "thinker", None) or getattr(stage_model, "model", None) or stage_model self.processor = ( getattr(stage_model, "processor", None) @@ -101,6 +204,77 @@ def __init__(self, stage_model: nn.Module, *, model_path: str | None = None, dev self._init_token_ids() if self.tokenizer is not None: self._require_special_token_ids() + hf_config = getattr(stage_model, "config", None) + # ``duplex_incremental_fbank`` (HF override, default off): each session's + # streaming mel front end recomputes only the new frames (incremental_fbank.py). + self._incremental_fbank = getattr(hf_config, "duplex_incremental_fbank", False) is True + self._fbank_stats = ( + _FbankStats(incremental=self._incremental_fbank) + if getattr(hf_config, "duplex_fbank_stats", False) is True + else None + ) + # The runtime is built with the model, before its weights load and + # before the loader switches it to eval mode: the streaming encoder's + # CUDA graphs are captured later, by build_audio_cuda_graph(). + self._audio_graph_attempted = False + + def build_audio_cuda_graph(self) -> bool: + """Capture the batched streaming encoder's CUDA graphs once the weights are loaded (idempotent). + + The model runner calls this once the weights are loaded, at the start + of its profile run (``omni_post_load``). At + runtime construction the audio encoder still has its initial weights + and is in training mode, which the batched encoder refuses ("stays + per-session: training mode"), so the capture never happened there. + ``--hf-overrides '{"duplex_audio_encoder_cuda_graph": false}'`` skips + the probe and the capture; the batched encoder then stays eager. + """ + if self._audio_graph_attempted: + return False + self._audio_graph_attempted = True + if not bool(getattr(getattr(self.thinker, "config", None), "duplex_audio_encoder_cuda_graph", True)): + return False + apm = getattr(self.thinker, "apm", None) + if getattr(apm, "training", False) is True and callable(getattr(apm, "eval", None)): + # Inference only; the loader's model.eval() covers it as well. + apm.eval() + return self._maybe_build_audio_cuda_graph() + + def _maybe_build_audio_cuda_graph(self) -> bool: + """Capture the batched streaming encoder's CUDA graphs (see ``build_audio_cuda_graph``). + + A silent probe through a session-configured copy of the processor + gives the mel-frame count of one steady unit (every unit but a + session's first). Any failure leaves the thinker on the eager batched + path. + """ + build = getattr(self.thinker, "build_streaming_audio_graph_encoder", None) + if not callable(build) or self.processor is None: + return False + try: + # Probe a processor configured like a session's (exact mode, the + # stage's chunk_ms / first_chunk_ms): the shared processor is not + # configured until a session uses it, and its default chunk is not + # a model unit (H100: unit_frames=10, rejected by the graph). + probe = self._configure_streaming_processor(SimpleNamespace(streaming_processor=None)) + if probe is None: + return False + # The first unit has its own chunk size; consume it to reach the steady one. + first_chunk_size = self._streaming_chunk_size(probe) + probe.process_audio_streaming(np.zeros(first_chunk_size, dtype=np.float32), reset=False) + steady_chunk_size = self._streaming_chunk_size(probe) + feature = probe.process_audio_streaming(np.zeros(steady_chunk_size, dtype=np.float32), reset=False) + # A BatchFeature-like mapping or a (features, metadata) tuple, by processor version. + audio_features = feature["audio_features"] if isinstance(feature, dict) else feature[0] + frames = int(audio_features.shape[-1]) + except Exception: + logger.info( + "MiniCPM-o could not derive the steady streaming-audio unit shape; " + "the streaming audio encoder stays eager", + exc_info=True, + ) + return False + return bool(build(unit_frames=frames)) def _stage_runtime_ready(self) -> bool: return self.processor is not None and self.tokenizer is not None and self.thinker is not None @@ -138,6 +312,7 @@ def _configure_streaming_processor( reset_streaming = getattr(processor, "reset_streaming", None) if callable(reset_streaming): reset_streaming() + self._maybe_enable_incremental_fbank(processor) return processor configure_streaming = getattr(processor, "configure_streaming", None) if callable(configure_streaming): @@ -147,8 +322,24 @@ def _configure_streaming_processor( slide_trigger_seconds=30.0, slide_stride_seconds=10.0, ) + self._maybe_enable_incremental_fbank(processor) return processor + def _maybe_enable_incremental_fbank(self, processor: Any) -> None: # Any: remote-code processor + """``duplex_incremental_fbank``: put the processor's streaming mel front end on the incremental path.""" + if not getattr(self, "_incremental_fbank", False): + return + mel_processor = getattr(processor, "_streaming_mel_processor", None) + if mel_processor is None: + return + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.incremental_fbank import enable_incremental_fbank + + if not enable_incremental_fbank(mel_processor): + logger.warning_once( + "MiniCPM-o duplex_incremental_fbank: %s is not supported; the streaming fbank recomputes in full", + type(mel_processor).__name__, + ) + def _prepare_session_context( self, state: _MiniCPMO45Stage0SessionState, @@ -209,23 +400,111 @@ def _prepare_session_context( state.pending_speech_response_open = True def _stage_prefill_embeddings_only( + self, + state: _MiniCPMO45Stage0SessionState, + audio_waveform: NDArray[np.float32] | None, + **kwargs: Any, + ) -> dict[str, object]: + """Build scheduler-owned Stage0 input embeddings for one audio append. + + Unlike the legacy worker-control path, this method never calls an eager + model forward. The normal vLLM runner consumes the returned embeddings + and owns attention metadata, block tables, KV cache, and sampling. + Keyword arguments are those of :meth:`_stage_prefill_steps`; an append + :meth:`stage_prefill_batch` already built this step is returned as is. + """ + staged = self.take_staged_prefill(state, kwargs.get("epoch"), kwargs.get("seq")) + if staged is not None: + return staged + return self._run_prefill_steps([self._stage_prefill_steps(state, audio_waveform, **kwargs)])[0] + + @staticmethod + def take_staged_prefill( + state: _MiniCPMO45Stage0SessionState, epoch: int | None, seq: int | None + ) -> dict[str, object] | None: + """Pop the append :meth:`stage_prefill_batch` built for ``(epoch, seq)`` this step, if any.""" + staged = state.staged_prefill + if staged is None: + return None + state.staged_prefill = None + if seq is None or staged[0] != (epoch, seq): + return None + return staged[1] + + def batches_audio_encoder(self) -> bool: + """Whether the thinker can encode several sessions' audio units in one call.""" + return self._batched_audio_encoder() is not None + + def needs_prefill(self, state: _MiniCPMO45Stage0SessionState | None, epoch: int | None, seq: int | None) -> bool: + """Whether an append has not been built (or attempted) for this session yet.""" + if seq is None: + return False + if state is None: + return True + identity = (epoch, seq) + if state.prepared_append_identity == identity and state.prepared_inputs_embeds is not None: + return False + return state.attempted_append_identity != identity + + def stage_prefill_batch( + self, + appends: list[tuple[_MiniCPMO45Stage0SessionState, NDArray[np.float32] | None, dict[str, Any]]], + ) -> None: + """Build several sessions' appends with one streaming-encoder pass per unit. + + Each append runs the per-session prefill; only the sessions' ``k``-th + unit encodes are batched. Results are staged on the session for the + append's own :meth:`_stage_prefill_embeddings_only` call this step. + Sessions must be distinct and every append must carry a ``seq``. + """ + if len({id(state) for state, _, _ in appends}) != len(appends): + raise ValueError("stage_prefill_batch needs one append per session") + steps = [self._stage_prefill_steps(state, audio, **kwargs) for state, audio, kwargs in appends] + for (state, _, kwargs), result in zip(appends, self._run_prefill_steps(steps), strict=True): + state.staged_prefill = ((kwargs.get("epoch"), kwargs["seq"]), result) + + def _run_prefill_steps(self, steps: list[_PrefillSteps]) -> list[dict[str, object]]: + """Drive prefills to completion, encoding each round's pending units together.""" + results: list[dict[str, object] | None] = [None] * len(steps) + pending: list[tuple[int, _MiniCPMO45AudioEncodeRequest]] = [] + + def advance(index: int, value: torch.Tensor | None, *, first: bool = False) -> None: + try: + request = next(steps[index]) if first else steps[index].send(value) + except StopIteration as stop: + results[index] = stop.value + return + pending.append((index, request)) + + for index in range(len(steps)): + advance(index, None, first=True) + while pending: + round_, pending = pending, [] + embeds = self._stage_audio_embeddings_batch([request for _, request in round_]) + for (index, _), embed in zip(round_, embeds, strict=True): + advance(index, embed) + return cast("list[dict[str, object]]", results) + + def _stage_prefill_steps( self, state: _MiniCPMO45Stage0SessionState, audio_waveform: NDArray[np.float32] | None, *, video_frames: list[Image.Image] | None = None, + encoded_frames: list[list[torch.Tensor]] | None = None, epoch: int | None = None, seq: int | None = None, is_speech: bool = False, final: bool = False, stage0_window: dict[str, object] | None = None, stage0_reanchor: dict[str, object] | None = None, - ) -> dict[str, object]: - """Build scheduler-owned Stage0 input embeddings for one audio append. + ) -> _PrefillSteps: + """The prefill of one append, yielding each unit's encoder request. - Unlike the legacy worker-control path, this method never calls an eager - model forward. The normal vLLM runner consumes the returned embeddings - and owns attention metadata, block tables, KV cache, and sampling. + The caller sends back that unit's audio embeddings (or ``None``). + ``encoded_frames`` holds the append's frames already encoded by + ``prefetch_vision`` (then ``video_frames`` may be omitted); without it + ``video_frames`` are encoded here. """ start_time = time.time() processor = self._configure_streaming_processor(state) @@ -239,6 +518,7 @@ def _stage_prefill_embeddings_only( result["inputs_embeds"] = state.prepared_inputs_embeds result["input_token_ids"] = list(state.prepared_input_token_ids) return result + state.attempted_append_identity = append_identity self._require_special_token_ids() if isinstance(stage0_window, dict): completed_ids = stage0_window.get("completed_token_ids") @@ -252,11 +532,12 @@ def _stage_prefill_embeddings_only( state.pending_terminator_token = completed_terminator if audio_waveform is None or len(audio_waveform) == 0: return self._stage_prefill_result(False, start_time, "empty audio") + num_frames = len(encoded_frames) if encoded_frames is not None else len(video_frames or ()) state.audio_buffer = np.concatenate([state.audio_buffer, np.asarray(audio_waveform, dtype=np.float32)]) chunk_size = self._streaming_chunk_size(processor) self._pad_first_audio_chunk_if_needed(state, processor) if len(state.audio_buffer) < chunk_size: - if video_frames: + if num_frames: # Frames belong to the unit this append closes, and the # scheduler already reserved their vision tokens against this # append. A frame arriving before its unit can close has @@ -265,7 +546,7 @@ def _stage_prefill_embeddings_only( return self._stage_prefill_result( False, start_time, - f"{len(video_frames)} video frame(s) on an append that closes no model unit " + f"{num_frames} video frame(s) on an append that closes no model unit " f"(need {chunk_size} samples, only {len(state.audio_buffer)}); " "send frames on unit boundaries", ) @@ -278,10 +559,11 @@ def _stage_prefill_embeddings_only( # carry them, mirroring official streaming_prefill (feed , then # every frame_list entry as its own block, then the audio). frame_blocks: list[list[torch.Tensor]] = [] - if video_frames: + if num_frames: self._require_vision_token_ids() - encoded_frames = self._stage_vision_embeddings(video_frames) - if encoded_frames is None or len(encoded_frames) != len(video_frames): + if encoded_frames is None: + encoded_frames = self._stage_vision_embeddings(video_frames or []) + if encoded_frames is None or len(encoded_frames) != num_frames: return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") if any(not slices for slices in encoded_frames): return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") @@ -312,7 +594,7 @@ def _stage_prefill_embeddings_only( ): with suppress(Exception): setattr(batch_feature, name, value) - audio_embeds = self._stage_audio_embeddings(batch_feature, state=state) + audio_embeds = yield _MiniCPMO45AudioEncodeRequest(batch_feature, state) if audio_embeds is None: if units_built == 0: return self._stage_prefill_result(False, start_time, "streaming audio embedding returned empty") @@ -623,7 +905,9 @@ def _as_2d_tensor(value: torch.Tensor) -> torch.Tensor: def _embed_token(self, token_id: int) -> torch.Tensor: import torch - token = torch.tensor([int(token_id)], dtype=torch.long, device=self._model_device()) + # A device-side fill, not ``torch.tensor(..., device=...)``: that is a + # blocking host-to-device copy, and every unit embeds 2-3 tokens. + token = torch.full((1,), int(token_id), dtype=torch.long, device=self._model_device()) embedder = self._token_embedder() embeds = embedder(token) return self._as_2d_tensor(embeds) @@ -746,10 +1030,15 @@ def _process_streaming_audio( processor = processor or self.processor process = getattr(processor, "process_audio_streaming", None) if callable(process): + stats = getattr(self, "_fbank_stats", None) + started = time.perf_counter() if stats is not None else 0.0 try: - return process(audio_chunk, reset=False, return_batch_feature=True) + result = process(audio_chunk, reset=False, return_batch_feature=True) except TypeError: - return process(audio_chunk, chunk_idx=chunk_idx) + result = process(audio_chunk, chunk_idx=chunk_idx) + if stats is not None: + stats.record(time.perf_counter() - started, _streaming_buffer_seconds(processor)) + return result return {"audio_features": audio_chunk, "audio_feature_lens": [[len(audio_chunk)]]} def _stage_audio_embeddings( @@ -759,7 +1048,11 @@ def _stage_audio_embeddings( state: _MiniCPMO45Stage0SessionState | None = None, ) -> torch.Tensor | None: if hasattr(batch_feature, "to"): - batch_feature = batch_feature.to(self.device) + batch_feature = self._audio_feature_to_device(batch_feature) + if state is not None and hasattr(state.audio_past_key_values, "to_legacy_cache"): + # A batched-path session whose unit this path has to take: hand the + # history over as the transformers cache the per-session path grows. + state.audio_past_key_values = state.audio_past_key_values.to_legacy_cache() self._ensure_dynamic_cache_compat() has_audio_cache = state is not None and hasattr(self.thinker, "audio_past_key_values") previous_audio_past_key_values = ( @@ -799,6 +1092,108 @@ def _stage_audio_embeddings( if has_audio_cache: self.thinker.audio_past_key_values = previous_audio_past_key_values + def _stage_audio_embeddings_batch( + self, + requests: list[_MiniCPMO45AudioEncodeRequest], + ) -> list[torch.Tensor | None]: + """Encode one unit per request; sessions the batched encoder takes share one pass. + + A request stays on :meth:`_stage_audio_embeddings` when the thinker has + no batched encoder (or it is disabled), when its session already holds a + per-session transformers cache, or when the unit is not the processor's + standard single-audio log-mel chunk. Requests must be of distinct sessions. + """ + results: list[torch.Tensor | None] = [None] * len(requests) + encoder = self._batched_audio_encoder() + batched: list[tuple[int, _MiniCPMO45AudioEncodeRequest, Any]] = [] + for index, request in enumerate(requests): + chunk = self._batched_audio_chunk(request) if encoder is not None else None + if chunk is None: + results[index] = self._stage_audio_embeddings(request.batch_feature, state=request.state) + else: + batched.append((index, request, chunk)) + if batched: + assert encoder is not None + embeds, caches = encoder([chunk for _, _, chunk in batched]) + for (index, request, _), embed, cache in zip(batched, embeds, caches, strict=True): + assert request.state is not None + request.state.audio_past_key_values = cache + results[index] = embed + return results + + def _batched_audio_encoder(self) -> Callable[[list[Any]], tuple[list[Any], list[Any]]] | None: + encoder = getattr(self, "_batched_audio_encoder_fn", False) + if encoder is False: + thinker = getattr(self, "thinker", None) + batch = getattr(thinker, "get_audio_embedding_streaming_batch", None) + supported = getattr(thinker, "supports_streaming_audio_batch", None) + encoder = batch if callable(batch) and callable(supported) and supported() else None + self._batched_audio_encoder_fn = encoder + return encoder + + @staticmethod + def _batched_audio_chunk(request: _MiniCPMO45AudioEncodeRequest) -> Any | None: # Any: StreamingAudioChunk + import torch + + from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import ( + StreamingAudioChunk, + StreamingAudioKVCache, + ) + + state = request.state + if state is None: + return None + cache = state.audio_past_key_values + if cache is not None and not isinstance(cache, StreamingAudioKVCache): + return None + feature = request.batch_feature + try: + features = feature["audio_features"] + lengths = feature["audio_feature_lens"] + except (KeyError, TypeError): + return None + if not isinstance(features, torch.Tensor) or features.ndim != 3 or features.shape[0] != 1: + return None + if isinstance(lengths, (list, tuple)) and len(lengths) == 1: + lengths = lengths[0] + if not isinstance(lengths, torch.Tensor) or lengths.numel() != 1: + return None + # Same extra-context frames the per-session call passes. + return StreamingAudioChunk( + features=features, + cache=cache, + prefix_extra_frames=0 if int(getattr(feature, "chunk_idx", 0)) == 0 else 2, + suffix_extra_frames=2, + feature_length=int(lengths.reshape(-1)[0]), + ) + + def _audio_feature_to_device(self, batch_feature: Any) -> Any: # Any: transformers BatchFeature + """Move one unit's log-mel features to the encoder without a host sync. + + On CUDA the float features go through pinned memory and the integer + lengths stay on the host (``get_audio_embedding_streaming`` only reads + them as slice bounds); otherwise ``BatchFeature.to``. + """ + import torch + + data = getattr(batch_feature, "data", None) + streaming = any( + callable(getattr(target, "get_audio_embedding_streaming", None)) + for target in (self.stage_model, self.thinker) + ) + if torch.device(self.device).type != "cuda" or not isinstance(data, dict) or not streaming: + return batch_feature.to(self.device) + moved: dict[str, object] = {} + for key, value in data.items(): + if isinstance(value, torch.Tensor) and torch.is_floating_point(value): + if value.device.type == "cpu": + value = value.pin_memory().to(self.device, non_blocking=True) + else: + value = value.to(self.device) + moved[key] = value + batch_feature.data = moved + return batch_feature + @staticmethod def _decode_ref_audio_from_session_config(session_config: dict[str, object]) -> NDArray[np.float32] | None: from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import decode_native_ref_audio_from_config @@ -1092,7 +1487,17 @@ def _official_max_slice_nums(self, frame_count: int) -> list[int]: return [1] return [2] + [1] * (frame_count - 1) - def _encode_processed_vision(self, processed: Any) -> list[torch.Tensor] | None: # Any: BatchFeature + def _encode_processed_vision_batch( + self, + processed_frames: list[Any], # Any: BatchFeature, one per frame + ) -> list[list[torch.Tensor]] | None: + """Encode every slice of ``processed_frames`` in one vision-tower call; embeddings per frame. + + One call lets the tower group equal patch grids across frames (and + sessions, via ``prefetch_vision``); each slice's embedding is the same + as encoding its frame alone. ``tgt_sizes`` stays on the host so the + encoder reads the grids without a device sync. + """ targets = (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)) for target in targets: if target is None: @@ -1106,57 +1511,185 @@ def _encode_processed_vision(self, processed: Any) -> list[torch.Tensor] | None: vpm_param = next(vpm.parameters()) device, dtype = vpm_param.device, vpm_param.dtype - pixel_nested = processed["pixel_values"] - tgt_nested = processed["tgt_sizes"] flat_pixels: list[torch.Tensor] = [] flat_tgt: list[torch.Tensor] = [] - for image_slices, image_tgt in zip(pixel_nested, tgt_nested): - for slice_pixels in image_slices: - flat_pixels.append(slice_pixels.to(device=device, dtype=dtype)) - tgt_tensor = image_tgt if hasattr(image_tgt, "reshape") else torch.tensor(image_tgt) - flat_tgt.append(tgt_tensor.reshape(-1, 2)) + slices_per_frame: list[int] = [] + for processed in processed_frames: + num_slices = 0 + for image_slices, image_tgt in zip(processed["pixel_values"], processed["tgt_sizes"]): + for slice_pixels in image_slices: + flat_pixels.append(slice_pixels.to(device=device, dtype=dtype)) + num_slices += 1 + tgt_tensor = image_tgt if hasattr(image_tgt, "reshape") else torch.tensor(image_tgt) + flat_tgt.append(torch.as_tensor(tgt_tensor).reshape(-1, 2)) + if num_slices == 0: + return None + slices_per_frame.append(num_slices) if not flat_pixels: return None - tgt_sizes = torch.cat(flat_tgt, dim=0).to(device=device, dtype=torch.int64) + tgt_sizes = torch.cat(flat_tgt, dim=0).to(device="cpu", dtype=torch.int64) with torch.no_grad(): hidden = get_hidden({"pixel_values": flat_pixels, "tgt_sizes": tgt_sizes}) except Exception: # noqa: BLE001 - prefill fails with a reason return None expected = MiniCPMO45DuplexPolicy.VISION_EMBEDS_PER_FRAME embed_dtype = self._token_embedding_dtype() - out: list[torch.Tensor] = [] + blocks: list[torch.Tensor] = [] for block in hidden: block_2d = self._as_2d_tensor(block) if int(block_2d.shape[0]) != expected: return None if embed_dtype is not None and block_2d.dtype != embed_dtype: block_2d = block_2d.to(dtype=embed_dtype) - out.append(block_2d) + blocks.append(block_2d) + if len(blocks) != sum(slices_per_frame): + return None + out: list[list[torch.Tensor]] = [] + offset = 0 + for num_slices in slices_per_frame: + out.append(blocks[offset : offset + num_slices]) + offset += num_slices return out return None - def _stage_vision_embeddings(self, frames: list[Image.Image]) -> list[list[torch.Tensor]] | None: - """Encode camera frames for omni duplex via the loaded vision tower. + def _preprocess_frames(self, frames: list[Image.Image]) -> list[Any] | None: # Any: BatchFeature + """CPU slicing of each frame with its official ``max_slice_nums``; ``None`` on failure. - Official ``streaming_prefill`` encodes each ``frame_list`` entry as its - own image (source + optional HD slices) and never stacks audio: the - unit still carries one second of soundtrack. A stacked pair uses - ``max_slice_nums=[2, 1]`` so the current frame keeps the elevator - display readable and the composite stays a single 448-normalized tile. Frames are processed one at a time because ``process_image([a, b])`` packs both PILs into one batch item. """ process_image = getattr(self.processor, "process_image", None) if not callable(process_image): return None - out: list[list[torch.Tensor]] = [] + processed: list[Any] = [] for frame, max_slices in zip(frames, self._official_max_slice_nums(len(frames))): try: - processed = process_image([frame], max_slice_nums=max_slices) + processed.append(process_image([frame], max_slice_nums=max_slices)) except Exception: # noqa: BLE001 - prefill fails with a reason return None - encoded = self._encode_processed_vision(processed) - if encoded is None: + return processed + + def _stage_vision_embeddings(self, frames: list[Image.Image]) -> list[list[torch.Tensor]] | None: + """Encode camera frames for omni duplex via the loaded vision tower. + + Official ``streaming_prefill`` encodes each ``frame_list`` entry as its + own image (source + optional HD slices) and never stacks audio: the + unit still carries one second of soundtrack. A stacked pair uses + ``max_slice_nums=[2, 1]`` so the current frame keeps the elevator + display readable and the composite stays a single 448-normalized tile. + Each frame is sliced on its own, then all slices share one encoder call. + """ + processed = self._preprocess_frames(frames) + if processed is None: + return None + return self._encode_processed_vision_batch(processed) + + @staticmethod + def _vision_prefetch_key(duplex: dict[str, Any]) -> tuple[str, int | None, int] | None: + """``(session_id, epoch, seq)`` of an append, parsed like ``preprocess`` does; ``None`` if unkeyed.""" + session_id = str(duplex.get("session_id") or "") + try: + seq = int(duplex["seq"]) if duplex.get("seq") is not None else None + except (TypeError, ValueError): + seq = None + try: + epoch = int(duplex["epoch"]) if duplex.get("epoch") is not None else None + except (TypeError, ValueError): + epoch = None + if not session_id or seq is None: + return None + return session_id, epoch, seq + + def _vision_prefetch_chunk_size(self) -> int: + """Jobs per pipeline chunk in ``prefetch_vision``: the checkpoint's ``vision_batch_size``, else 16.""" + for target in (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)): + size = getattr(getattr(target, "config", None), "vision_batch_size", None) + if isinstance(size, int) and size > 0: + return size + return 16 + + def _finish_vision_prefetch_chunk( + self, + chunk: list[tuple[tuple[str, int | None, int], tuple[object, ...]]], + prepared: list[list[Any] | None], + ) -> None: + """Encode one chunk's ready preprocessed frames and cache their embeddings by append key.""" + ready = [(job, processed) for job, processed in zip(chunk, prepared) if processed is not None] + if not ready: + return + encoded = self._encode_processed_vision_batch([item for _, processed in ready for item in processed]) + if encoded is None: + return + offset = 0 + for (key, frames_payload), processed in ready: + self._prefetched_vision[key] = (frames_payload, encoded[offset : offset + len(processed)]) + offset += len(processed) + + def prefetch_vision(self, appends: list[dict[str, Any]]) -> None: + """Encode the camera frames of this runner step's new appends in batched vision-tower calls. + + JPEG decode and slicing run on the frame pool; every slice goes through + ``get_vision_hidden_states`` in ``vision_batch_size`` chunks, the next + chunk's CPU work overlapping this chunk's GPU encode. Results are keyed + by append identity and checked against the frame payload on use; + anything that fails is left for the per-request path. + """ + self._prefetched_vision.clear() + jobs: list[tuple[tuple[str, int | None, int], tuple[object, ...]]] = [] + for duplex in appends: + payload = duplex.get("payload") + key = self._vision_prefetch_key(duplex) + if key is None or not isinstance(payload, dict): + continue + frames = payload.get("video_frames") + if not isinstance(frames, list) or not frames: + continue + state = self.sessions.get(key[0]) + if ( + state is not None + and state.prepared_append_identity == (key[1], key[2]) + and state.prepared_inputs_embeds is not None + ): + continue # already built; preprocess replays the cached append + jobs.append((key, tuple(frames))) + if not jobs or not callable(getattr(self.processor, "process_image", None)): + return + + def preprocess(frames_payload: tuple[object, ...]) -> list[Any] | None: + try: + decoded = self._decode_video_frames_payload({"video_frames": list(frames_payload)}) + except ValueError: return None - out.append(encoded) - return out + return self._preprocess_frames(decoded) if decoded else None + + pool = _frame_preprocess_pool() if len(jobs) > 1 else None + chunk_size = self._vision_prefetch_chunk_size() + chunks = [jobs[begin : begin + chunk_size] for begin in range(0, len(jobs), chunk_size)] + + if pool is None: + for chunk in chunks: + self._finish_vision_prefetch_chunk(chunk, [preprocess(frames) for _, frames in chunk]) + return + + pending = [pool.submit(preprocess, frames) for _, frames in chunks[0]] + for idx, chunk in enumerate(chunks): + # Submit the next chunk before this chunk's encode so they overlap. + next_pending = ( + [pool.submit(preprocess, frames) for _, frames in chunks[idx + 1]] if idx + 1 < len(chunks) else None + ) + prepared = [future.result() for future in pending] + self._finish_vision_prefetch_chunk(chunk, prepared) + pending = next_pending + + def take_prefetched_vision(self, duplex: dict[str, Any]) -> list[list[torch.Tensor]] | None: + """Pop this append's prefetched per-frame embeddings if its frame payload still matches.""" + key = self._vision_prefetch_key(duplex) + entry = self._prefetched_vision.pop(key, None) if key is not None else None + if entry is None: + return None + frames_payload, encoded = entry + payload = duplex.get("payload") + frames = payload.get("video_frames") if isinstance(payload, dict) else None + if not isinstance(frames, list) or tuple(frames) != frames_payload: + return None + return encoded diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 2781fcc3041..3015d2324ca 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -136,6 +136,11 @@ def _encode_tokens(tokenizer: Any, prompt: str) -> list[int]: from vllm.utils.tensor_schema import TensorSchema, TensorShape +from vllm_omni.model_executor.models.minicpmo_4_5.vision_fused import ( + VisionGraphEncoder, + encode_packed_fused, + supports_fused_layers, +) from vllm_omni.model_executor.models.model_local_kv import ( ModelLocalKVScope, ModelLocalKVSpec, @@ -403,6 +408,23 @@ def __init__( super().__init__(**kwargs) +def _resolve_vision_attention_implementation(vision_config: Any, model_config: Any) -> str: + """Pick the SigLIP attention kernel: a legacy engine override, then the vision config, then SDPA. + + SDPA computes the same softmax attention as the eager matmul path (fp32 + softmax, native-dtype matmuls) without materializing the ``(B, H, L, L)`` + score tensor. Set ``vision_config._attn_implementation`` (e.g. through + ``hf_overrides``) to ``"eager"`` or ``"flash_attention_2"`` to select those. + """ + implementation = getattr(model_config, "_attn_implementation", None) + if implementation is None: + implementation = getattr(vision_config, "_attn_implementation", None) or "sdpa" + if implementation == "flash_attention_2" and not is_flash_attn_2_available(): + logger.warning("flash_attention_2 requested for the MiniCPM-o vision tower but unavailable; using sdpa") + return "sdpa" + return implementation + + # ============== Image Processing Classes ============== @@ -950,6 +972,8 @@ def __init__(self, config: SiglipVisionConfig): self.num_patches = self.num_patches_per_side**2 self.num_positions = self.num_patches self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + # (h, w, device) -> position ids of that patch grid, see layout_position_ids. + self._layout_position_ids: dict[tuple[int, int, str], torch.Tensor] = {} def _create_grid_position_ids( self, @@ -1016,6 +1040,27 @@ def _create_position_ids( return position_ids.to(device=device) + def layout_position_ids(self, patch_grid_height: int, patch_grid_width: int, device: torch.device) -> torch.Tensor: + """Position ids of one ``(h, w)`` patch grid, cached on ``device``. + + Built by the same host-side bucketize as ``_create_position_ids``, so + the ids are identical to the padded path's valid positions. + """ + cache = self._layout_position_ids + key = (int(patch_grid_height), int(patch_grid_width), str(device)) + ids = cache.get(key) + if ids is None: + if len(cache) >= 1024: + cache.clear() + boundaries = torch.arange( + 1 / self.num_patches_per_side, + 1.0, + 1 / self.num_patches_per_side, + ) + ids = self._create_grid_position_ids(key[0], key[1], boundaries).to(device=device) + cache[key] = ids + return ids + def forward( self, pixel_values: torch.FloatTensor, @@ -1110,6 +1155,78 @@ def forward( return attn_output, attn_weights + def _attend_unpadded(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """Attention over ``(n, L, heads, head_dim)`` items with no padding; same layout out.""" + query, key, value = (x.transpose(1, 2) for x in (query, key, value)) + attn_weights = torch.matmul(query, key.transpose(2, 3)) * self.scale + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + return torch.matmul(attn_weights, value).transpose(1, 2) + + def forward_packed(self, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """Self-attention over packed ``(tokens, embed)`` items, each attending only to itself. + + ``seq_groups`` holds contiguous ``(start_token, num_items, seq_len)`` + runs. Every item is unpadded, so no mask is built and each group runs + one mask-free attention call; the projections are shared token-wise. + """ + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + attn_output = torch.empty_like(query_states) + for start, num_items, seq_len in seq_groups: + end = start + num_items * seq_len + shape = (num_items, seq_len, self.num_heads, self.head_dim) + attn_output[start:end].view(shape).copy_( + self._attend_unpadded( + query_states[start:end].view(shape), + key_states[start:end].view(shape), + value_states[start:end].view(shape), + ) + ) + return self.out_proj(attn_output) + + +class SiglipSdpaAttention(SiglipAttention): + """SigLIP attention through ``F.scaled_dot_product_attention`` (adapted from vllm-omni#5385). + + Same math as the eager path -- scaled QK^T, fp32 softmax, PV in the model + dtype -- but the kernel never materializes the ``(B, H, L, L)`` scores. + """ + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + output_attentions: bool | None = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if output_attentions: + return super().forward(hidden_states, attention_mask=attention_mask, output_attentions=output_attentions) + + batch_size, q_len, _ = hidden_states.size() + shape = (batch_size, q_len, self.num_heads, self.head_dim) + query_states = self.q_proj(hidden_states).view(shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(shape).transpose(1, 2) + attn_output = F.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=self.dropout if self.training else 0.0, + scale=self.scale, + ) + attn_output = attn_output.transpose(1, 2).reshape(batch_size, q_len, self.embed_dim) + return self.out_proj(attn_output), None + + def _attend_unpadded(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + return F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + dropout_p=self.dropout if self.training else 0.0, + scale=self.scale, + ).transpose(1, 2) + class SiglipFlashAttention2(SiglipAttention): """ @@ -1327,6 +1444,10 @@ def _upad_input( (max_seqlen_in_batch_q, max_seqlen_in_batch_k), ) + def _attend_unpadded(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + # Flash attention consumes the (n, L, heads, head_dim) layout directly. + return self._flash_attention_forward(query, key, value, None, query.shape[1]) + # Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip class SiglipMLP(nn.Module): @@ -1350,7 +1471,12 @@ def __init__(self, config: SiglipVisionConfig): super().__init__() self.embed_dim = config.hidden_size self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" - self.self_attn = SiglipAttention(config) if not self._use_flash_attention_2 else SiglipFlashAttention2(config) + if self._use_flash_attention_2: + self.self_attn = SiglipFlashAttention2(config) + elif config._attn_implementation == "sdpa": + self.self_attn = SiglipSdpaAttention(config) + else: + self.self_attn = SiglipAttention(config) self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.mlp = SiglipMLP(config) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) @@ -1402,6 +1528,15 @@ def forward( return outputs + def forward_packed(self, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """``forward`` over packed ``(tokens, embed)`` items; see ``SiglipAttention.forward_packed``.""" + residual = hidden_states + hidden_states = self.self_attn.forward_packed(self.layer_norm1(hidden_states), seq_groups) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp(self.layer_norm2(hidden_states)) + return residual + hidden_states + class SiglipPreTrainedModel(PreTrainedModel): """ @@ -1568,8 +1703,14 @@ def forward( class SiglipVisionTransformer(SiglipPreTrainedModel): config_class = SiglipVisionConfig main_input_name = "pixel_values" + # Transformers 5 spells the flash flag ``_supports_flash_attn``; SDPA is + # the default vision kernel (see _resolve_vision_attention_implementation). + _supports_flash_attn = True _supports_flash_attn_2 = True + _supports_sdpa = True _no_split_modules = [] + # ``forward_packed`` through ``vision_fused.encode_packed_fused``; set by the model. + fused_layers = False def __init__(self, config: SiglipVisionConfig): super().__init__(config) @@ -1657,6 +1798,42 @@ def forward( attentions=encoder_outputs.attentions, ) + def forward_packed(self, pixel_values: torch.Tensor, layouts: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """Encode items packed along the patch axis, without padding; returns ``(tokens, embed)``. + + ``pixel_values`` is ``(1, C, patch, tokens * patch)``: every item's + ``(C, patch, h * w * patch)`` strip concatenated in order, which the + stride-``patch`` patch embedding cuts exactly at item boundaries. + ``layouts`` lists contiguous ``(h, w, count)`` runs of equal patch + grids. Per item this is the padded ``forward`` minus the padding: the + same position ids, and attention restricted to the item's own patches. + """ + embeddings = self.embeddings + device = embeddings.position_embedding.weight.device + position_ids = [ + embeddings.layout_position_ids(height, width, device).repeat(count) for height, width, count in layouts + ] + position_ids = position_ids[0] if len(position_ids) == 1 else torch.cat(position_ids) + patch_embeds = embeddings.patch_embedding(pixel_values).flatten(2).transpose(1, 2)[0] + hidden_states = patch_embeds + embeddings.position_embedding(position_ids) + + seq_groups: list[tuple[int, int, int]] = [] + start = 0 + for height, width, count in layouts: + seq_len = height * width + if seq_groups and seq_groups[-1][2] == seq_len: + # Equal-length grids share one attention call. + seq_groups[-1] = (seq_groups[-1][0], seq_groups[-1][1] + count, seq_len) + else: + seq_groups.append((start, count, seq_len)) + start += count * seq_len + + if self.fused_layers and hidden_states.is_cuda and supports_fused_layers(self): + return encode_packed_fused(self, hidden_states, seq_groups) + for encoder_layer in self.encoder.layers: + hidden_states = encoder_layer.forward_packed(hidden_states, seq_groups) + return self.post_layernorm(hidden_states) + # ============== Resampler Classes ============== # coding=utf-8 @@ -1843,6 +2020,175 @@ def forward(self, x, tgt_sizes=None): def _repeat(self, query, N: int): return query.unsqueeze(1).repeat(1, N, 1) + def supports_packed(self) -> bool: + attn = self.attn + return ( + not attn.batch_first + and attn._qkv_same_embed_dim + and attn.in_proj_weight is not None + and attn.bias_k is None + and attn.bias_v is None + and not attn.add_zero_attn + ) + + def forward_packed(self, x: torch.Tensor, layouts: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """``forward`` for packed ``(tokens, kv_dim)`` features in ``(h, w, count)`` runs; returns ``(B, Q, D)``. + + Per item this is ``forward`` without key padding: the same 2D sin-cos + positions and the same ``nn.MultiheadAttention`` weights, with the + learned queries projected once and each run of equal grids attending + in one mask-free call. Grids come from host ints, so no device sync. + """ + max_h = max(height for height, _, _ in layouts) + max_w = max(width for _, width, _ in layouts) + if max_h > self.max_size[0] or max_w > self.max_size[1]: + self.max_size = [max(max_h, self.max_size[0]), max(max_w, self.max_size[1])] + self._set_2d_pos_cache(self.max_size, x.device) + if self.pos_embed.device != x.device: + self.pos_embed = self.pos_embed.to(x.device) + + x = self.ln_kv(self.kv_proj(x)) + key_input = torch.empty_like(x) + start = 0 + for height, width, count in layouts: + end = start + count * height * width + pos_embed = self.pos_embed[:height, :width, :].reshape(height * width, -1).to(x.dtype) + torch.add( + x[start:end].view(count, height * width, -1), + pos_embed, + out=key_input[start:end].view(count, height * width, -1), + ) + start = end + + attn = self.attn + num_heads = attn.num_heads + head_dim = self.embed_dim // num_heads + w_q, w_k, w_v = attn.in_proj_weight.chunk(3) + b_q, b_k, b_v = attn.in_proj_bias.chunk(3) if attn.in_proj_bias is not None else (None, None, None) + query = F.linear(self.ln_q(self.query), w_q, b_q) + keys = F.linear(key_input, w_k, b_k) + values = F.linear(x, w_v, b_v) + num_queries = query.shape[0] + query_heads = query.view(num_queries, num_heads, head_dim).transpose(0, 1) + eager = getattr(self, "attn_implementation", "sdpa") == "eager" + + outputs: list[torch.Tensor] = [] + start = 0 + for height, width, count in layouts: + seq_len = height * width + end = start + count * seq_len + key_heads = keys[start:end].view(count, seq_len, num_heads, head_dim).transpose(1, 2) + value_heads = values[start:end].view(count, seq_len, num_heads, head_dim).transpose(1, 2) + q = query_heads.unsqueeze(0).expand(count, -1, -1, -1).contiguous() + if eager: + # Mirrors nn.MultiheadAttention's need_weights path used by ``forward``. + weights = torch.matmul(q / math.sqrt(head_dim), key_heads.transpose(-2, -1)) + out = torch.matmul(F.softmax(weights, dim=-1), value_heads) + else: + out = F.scaled_dot_product_attention(q, key_heads, value_heads) + outputs.append(out.transpose(1, 2).reshape(count, num_queries, self.embed_dim)) + start = end + + out = outputs[0] if len(outputs) == 1 else torch.cat(outputs) + out = self.ln_post(attn.out_proj(out)) + return out @ self.proj + + +def _packed_vision_layouts( + pixel_values: Sequence[torch.Tensor], + tgt_sizes: torch.Tensor | Sequence[Sequence[int]], + patch_size: int, +) -> list[tuple[int, int]] | None: + """Host ``(h, w)`` patch grid per slice, or ``None`` if the pixels do not match the grids. + + This is the only device sync of the packed vision path (none when + ``tgt_sizes`` is on the host, as in duplex). + """ + sizes = tgt_sizes.tolist() if isinstance(tgt_sizes, torch.Tensor) else [list(size) for size in tgt_sizes] + if len(sizes) != len(pixel_values) or not sizes: + return None + layouts: list[tuple[int, int]] = [] + for pixels, size in zip(pixel_values, sizes): + if len(size) != 2: + return None + height, width = int(size[0]), int(size[1]) + if ( + height <= 0 + or width <= 0 + or pixels.ndim != 3 + or pixels.shape[-2] != patch_size + or pixels.shape[-1] != height * width * patch_size + ): + return None + layouts.append((height, width)) + return layouts + + +def _plan_packed_vision_chunks( + layouts: Sequence[tuple[int, int]], max_items: int +) -> list[tuple[list[int], list[tuple[int, int, int]]]]: + """Order slices so equal grids are adjacent, then cut ``max_items``-slice chunks. + + Returns per chunk the original slice indices in packed order and the + ``(h, w, count)`` runs. Sorting by ``h * w`` first also puts equal + sequence lengths next to each other so they share an attention call. + """ + order = sorted(range(len(layouts)), key=lambda i: (layouts[i][0] * layouts[i][1], layouts[i])) + chunks: list[tuple[list[int], list[tuple[int, int, int]]]] = [] + for begin in range(0, len(order), max(1, max_items)): + indices = order[begin : begin + max(1, max_items)] + runs: list[tuple[int, int, int]] = [] + for index in indices: + height, width = layouts[index] + if runs and runs[-1][:2] == (height, width): + runs[-1] = (height, width, runs[-1][2] + 1) + else: + runs.append((height, width, 1)) + chunks.append((indices, runs)) + return chunks + + +def _encode_vision_packed( + vpm: "SiglipVisionTransformer", + resampler: Resampler, + pixel_values: Sequence[torch.Tensor], + layouts: Sequence[tuple[int, int]], + max_items: int, + graph_encoder: VisionGraphEncoder | None = None, +) -> torch.Tensor: + """SigLIP + resampler over unpadded slices; ``(num_slices, query_num, embed)`` in input order. + + Slices are grouped by patch grid and packed up to ``max_items`` per + encoder call (the checkpoint's ``vision_batch_size``), across images and + requests alike. No slice is padded or masked, so its output does not + depend on the size of the other slices in its batch. A single-grid chunk + replays ``graph_encoder``'s CUDA graph when it has one. + """ + weight = vpm.embeddings.patch_embedding.weight + outputs: list[torch.Tensor] = [] + packed_order: list[int] = [] + for indices, runs in _plan_packed_vision_chunks(layouts, max_items): + pixels = ( + torch.cat([pixel_values[index] for index in indices], dim=-1) + if len(indices) > 1 + else pixel_values[indices[0]] + ) + pixels = pixels.to(device=weight.device, dtype=weight.dtype).unsqueeze(0) + out = None + if graph_encoder is not None and len(runs) == 1: + out = graph_encoder.encode(pixels, *runs[0]) + if out is None: + out = resampler.forward_packed(vpm.forward_packed(pixels, runs), runs) + outputs.append(out) + packed_order.extend(indices) + out = outputs[0] if len(outputs) == 1 else torch.cat(outputs) + if packed_order != list(range(len(packed_order))): + inverse = [0] * len(packed_order) + for position, index in enumerate(packed_order): + inverse[index] = position + out = out.index_select(0, torch.tensor(inverse, dtype=torch.long).to(device=out.device, non_blocking=True)) + return out + class MultiheadAttention(nn.MultiheadAttention): def __init__( @@ -3909,6 +4255,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = config self.multimodal_config = multimodal_config + # Duplex session cap of this deployment; sizes the audio-encoder graph + # grid when ``duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions``. + self._duplex_max_sessions = int(getattr(vllm_config.model_config, "duplex_max_sessions", 1) or 1) # Initialize image processor self.image_processor = MiniCPMVImageProcessor( @@ -3921,11 +4270,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # Initialize vision encoder (SigLIP) if multimodal_config.get_limit_per_prompt("image"): - # Set attention implementation - if hasattr(vllm_config, "model_config") and hasattr(vllm_config.model_config, "_attn_implementation"): - config.vision_config._attn_implementation = vllm_config.model_config._attn_implementation - else: - config.vision_config._attn_implementation = "eager" + config.vision_config._attn_implementation = _resolve_vision_attention_implementation( + config.vision_config, + getattr(vllm_config, "model_config", None), + ) self.vpm = SiglipVisionTransformer(config.vision_config) # Drop last layer if configured @@ -3977,8 +4325,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): kv_dim=vision_dim, # Vision encoder dimension adaptive=True, # Enable adaptive resampling ) + self.resampler.attn_implementation = config.vision_config._attn_implementation else: self.resampler = None + # Unpadded, grid-grouped vision encoding (see _encode_vision_packed); + # False keeps the padded batch (the parity tests' reference). + self.vision_packed_encode = True + # Fused packed SigLIP layers and CUDA graphs of single-grid chunks + # (vision_fused.py). Off unless HF config vision_fused_layers / + # vision_cuda_graph (``--hf-overrides``) turn them on. + if self.vpm is not None: + self.vpm.fused_layers = bool(getattr(config, "vision_fused_layers", False)) + self.vision_cuda_graph = bool(getattr(config, "vision_cuda_graph", False)) + self._vision_graph_encoder: VisionGraphEncoder | None = None # Initialize audio encoder (APM) and audio projection if getattr(config, "init_audio", True) and hasattr(config, "audio_config") and config.audio_config is not None: @@ -4252,6 +4611,29 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens pixel_values = data["pixel_values"] tgt_sizes = data["tgt_sizes"] + vpm = self.vpm + if ( + getattr(self, "vision_packed_encode", False) + and hasattr(vpm, "forward_packed") + and hasattr(self.resampler, "forward_packed") + and self.resampler.supports_packed() + ): + layouts = _packed_vision_layouts(pixel_values, tgt_sizes, vpm.embeddings.patch_size) + if layouts is not None: + graph_encoder = None + if getattr(self, "vision_cuda_graph", False) and pixel_values[0].is_cuda: + graph_encoder = self._vision_graph_encoder + if graph_encoder is None: + graph_encoder = self._vision_graph_encoder = VisionGraphEncoder(vpm, self.resampler) + return _encode_vision_packed( + vpm, + self.resampler, + pixel_values, + layouts, + max(1, int(self.config.vision_batch_size)), + graph_encoder, + ) + B = len(pixel_values) P = pixel_values[0].shape[-2] L = max(item.shape[-1] for item in pixel_values) @@ -4439,7 +4821,8 @@ def get_audio_hidden_states(self, data: MiniCPMOAudioFeatureInputs) -> list[torc _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens) - num_audio_tokens = feature_lens_after_pooling + # One host read of every length instead of an implicit .item() per slice bound. + num_audio_tokens = feature_lens_after_pooling.tolist() final_audio_embeds = list[torch.Tensor]() idx = 0 @@ -4517,18 +4900,25 @@ def get_audio_embedding_streaming( device=wavforms.device, ) + # The final layer is ``last_hidden_state`` (the same tensor as + # ``hidden_states[-1]``); asking for every layer's output only to pick + # the last one keeps 25 activations alive per unit (as in #5587). + selects_final_layer = self.audio_encoder_layer == -1 audio_outputs = self.apm( wavforms, past_key_values=self.audio_past_key_values, use_cache=True, - output_hidden_states=True, + output_hidden_states=not selects_final_layer, attention_mask=audio_attention_mask, use_extra_context=use_extra_context, prefix_extra_frames=prefix_extra_frames, suffix_extra_frames=suffix_extra_frames, cnn_min_length=cnn_min_length, ) - audio_states = audio_outputs.hidden_states[self.audio_encoder_layer] + if selects_final_layer: + audio_states = audio_outputs.last_hidden_state + else: + audio_states = audio_outputs.hidden_states[self.audio_encoder_layer] self.audio_past_key_values = audio_outputs.past_key_values audio_embeds = self.audio_projection_layer(audio_states) @@ -4537,16 +4927,142 @@ def get_audio_embedding_streaming( audio_embeds = audio_embeds.transpose(1, 2) _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens) + # One host read of the lengths (free when they stay on the host) + # instead of a device sync per slice bound. + pooled_lens = feature_lens_after_pooling.tolist() final_audio_embeds: list[list[torch.Tensor]] = [] idx = 0 for i in range(len(audio_feature_lens_raw)): target_audio_embeds = [] for _ in range(len(audio_feature_lens_raw[i])): - target_audio_embeds.append(audio_embeds[idx, : feature_lens_after_pooling[idx], :]) + target_audio_embeds.append(audio_embeds[idx, : pooled_lens[idx], :]) idx += 1 final_audio_embeds.append(target_audio_embeds) return final_audio_embeds + def supports_streaming_audio_batch(self) -> bool: + """Whether ``get_audio_embedding_streaming_batch`` can serve this encoder. + + ``duplex_audio_encoder_batching`` on the HF config (``--hf-overrides``) + forces it on (``true``) or off (``false``, every session keeps the + per-session ``EncoderDecoderCache`` path). Unset, it is on for CUDA and + CPU, the devices its attention layout has been checked on. + """ + if self.apm is None or self.audio_projection_layer is None or self.audio_avg_pooler is None: + return False + from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import ( + streaming_batch_unsupported_reason, + ) + + requested = getattr(self.config, "duplex_audio_encoder_batching", None) + if requested is not None and not bool(requested): + return False + reason = streaming_batch_unsupported_reason(self.apm, audio_encoder_layer=self.audio_encoder_layer) + if reason is not None: + logger.info("MiniCPM-o streaming audio encoder stays per-session: %s", reason) + return False + return requested is not None or self.apm.conv1.weight.device.type in ("cuda", "cpu") + + def build_streaming_audio_graph_encoder(self, *, unit_frames: int) -> bool: + """Capture CUDA graphs of the streaming encoder's steady unit of ``unit_frames`` mel frames. + + Called once after the weights load, in eval mode (the Stage-0 duplex + runtime's ``build_audio_cuda_graph``, from the runner's + ``omni_post_load`` hook at the start of its profile run). HF config + ``duplex_audio_encoder_cuda_graph`` (default on, ~3 GiB at batch 32) + gates it; ``duplex_audio_encoder_cuda_graph_batch_sizes`` / + ``..._cache_buckets`` override the buckets, and + ``duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions: true`` + (default off) derives the batch grid from duplex ``max_sessions`` + instead (powers of two below it plus the cap itself). + ``duplex_audio_encoder_pinned_h2d`` (default off) stages the host mel + through pinned memory instead of a blocking copy. + ``duplex_audio_encoder_resident_slots`` (default unset = off) keeps up to + that many sessions' streaming caches resident in a slot pool the graphs + read and write in place, instead of a shared copy-in storage plus + per-session paged buffers; size it to ``duplex_session.max_sessions``. + Never raises: on any failure the batched encoder stays eager. + """ + self._duplex_audio_cuda_graph_encoder = None + if not self.supports_streaming_audio_batch(): + return False + if not bool(getattr(self.config, "duplex_audio_encoder_cuda_graph", True)): + return False + if self.apm.conv1.weight.device.type != "cuda": + logger.info("MiniCPM-o streaming audio encoder CUDA graph needs a CUDA device; staying eager") + return False + from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import DEFAULT_KV_PAGE_POSITIONS + from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder_graph import ( + DEFAULT_GRAPH_BATCH_SIZES, + DEFAULT_GRAPH_CACHE_BUCKETS, + StreamingAudioGraphEncoder, + batch_sizes_for_sessions, + ) + + batch_sizes = getattr(self.config, "duplex_audio_encoder_cuda_graph_batch_sizes", None) + if getattr(self.config, "duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions", False) is True: + # Overrides the explicit list: the grid follows duplex max_sessions. + batch_sizes = list(batch_sizes_for_sessions(getattr(self, "_duplex_max_sessions", 1))) + cache_buckets = getattr(self.config, "duplex_audio_encoder_cuda_graph_cache_buckets", None) + page = getattr(self.config, "duplex_audio_kv_page_positions", None) + allocated = torch.cuda.memory_allocated(self.apm.conv1.weight.device) + try: + wrapper = StreamingAudioGraphEncoder( + self.apm, + self.audio_projection_layer, + self.audio_avg_pooler, + unit_frames=int(unit_frames), + pool_step=int(self.config.audio_pool_step), + batch_sizes=[int(b) for b in batch_sizes] if batch_sizes else DEFAULT_GRAPH_BATCH_SIZES, + cache_buckets=[int(c) for c in cache_buckets] if cache_buckets else DEFAULT_GRAPH_CACHE_BUCKETS, + page_positions=int(page) if page else DEFAULT_KV_PAGE_POSITIONS, + pinned_h2d=getattr(self.config, "duplex_audio_encoder_pinned_h2d", False) is True, + resident_slots=int(getattr(self.config, "duplex_audio_encoder_resident_slots", None) or 0), + ) + wrapper.capture() + except Exception: + logger.exception("MiniCPM-o streaming audio encoder CUDA graph capture failed; staying eager") + return False + self._duplex_audio_cuda_graph_encoder = wrapper + logger.info( + "Captured MiniCPM-o streaming audio encoder CUDA graphs: unit_frames=%d batch_sizes=%s " + "cache_buckets=%s pinned_h2d=%s resident_slots=%d, took %.2f GiB", + unit_frames, + wrapper.batch_sizes, + wrapper.cache_buckets, + wrapper.pinned_h2d, + wrapper.resident_slots, + (torch.cuda.memory_allocated(self.apm.conv1.weight.device) - allocated) / (1 << 30), + ) + return True + + def get_audio_embedding_streaming_batch( + self, + chunks: Sequence[Any], + ) -> tuple[list[torch.Tensor | None], list[Any]]: + """Batched ``get_audio_embedding_streaming`` over ``StreamingAudioChunk`` rows. + + Returns the embeddings and each session's updated cache; the captured + graphs take the steady rows when present. + """ + from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import ( + DEFAULT_KV_PAGE_POSITIONS, + encode_streaming_audio_batch, + ) + + graph = getattr(self, "_duplex_audio_cuda_graph_encoder", None) + if graph is not None: + return graph.encode(chunks) + page = getattr(self.config, "duplex_audio_kv_page_positions", None) + return encode_streaming_audio_batch( + self.apm, + self.audio_projection_layer, + self.audio_avg_pooler, + chunks, + pool_step=int(self.config.audio_pool_step), + page_positions=int(page) if page else DEFAULT_KV_PAGE_POSITIONS, + ) + def _process_audio_input( self, audio_input: MiniCPMOAudioInputs, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder.py b/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder.py new file mode 100644 index 00000000000..2bae402f04e --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder.py @@ -0,0 +1,415 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One forward of MiniCPM-o's streaming Whisper encoder for many duplex sessions. + +The per-session path (``get_audio_embedding_streaming``) runs the encoder at +batch one against an ``EncoderDecoderCache`` that ``torch.cat``-grows every +layer on each one-second unit. This module encodes every session's ready unit +in one pass: + +* convolutions run on a zero-padded ``[B, mels, frames]`` batch, everything + after them on the packed ``[sum(n_b), d]`` rows, so padding is never attended; +* attention runs per row against that row's own cache, with the per-session + shapes, mask and kernel; +* each session owns a :class:`StreamingAudioKVCache` written in place. Its + ``length`` only advances after the whole forward, and a history at the + position bound is replaced rather than cleared, so a failed forward leaves + committed state untouched. + +Results match the per-session path up to floating-point reassociation (not +bit-exact). :func:`streaming_batch_unsupported_reason` names configurations +this does not handle, which keep the per-session path. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Sequence +from dataclasses import dataclass + +import torch +from torch import nn + +#: Growth step of a session buffer, in encoder positions (a 1 s unit is 50): +#: at most five reallocations per 1500-position cycle, ~12 MiB average slack. +DEFAULT_KV_PAGE_POSITIONS = 250 + +_SUPPORTED_ATTENTION = ("sdpa", "eager") +_SUPPORTED_DTYPES = (torch.bfloat16, torch.float32) + + +class StreamingAudioKVCache: + """One session's streaming-encoder self-attention cache, written in place. + + Layout ``[layers, 2 (key/value), capacity, embed_dim]``: token-major, so + appending a unit is one contiguous copy per layer and attention reads a + strided ``[1, heads, length, head_dim]`` view without materializing it. + Capacity grows in ``page_positions`` steps up to ``max_positions``. + """ + + __slots__ = ("_buffer", "embed_dim", "length", "max_positions", "num_heads", "num_layers", "page_positions") + + def __init__( + self, + *, + num_layers: int, + embed_dim: int, + num_heads: int, + max_positions: int, + page_positions: int = DEFAULT_KV_PAGE_POSITIONS, + ) -> None: + if min(num_layers, embed_dim, num_heads, max_positions, page_positions) <= 0: + raise ValueError("streaming audio KV cache dimensions must be positive") + self.num_layers = num_layers + self.embed_dim = embed_dim + self.num_heads = num_heads + self.max_positions = max_positions + self.page_positions = page_positions + self.length = 0 + self._buffer: torch.Tensor | None = None + + @property + def capacity(self) -> int: + return 0 if self._buffer is None else int(self._buffer.shape[2]) + + def reserve(self, positions: int, *, dtype: torch.dtype, device: torch.device) -> None: + """Make room for ``positions`` entries, keeping the committed prefix.""" + if positions <= self.capacity: + return + pages = math.ceil(positions / self.page_positions) * self.page_positions + capacity = max(positions, min(self.max_positions, pages)) + buffer = torch.empty((self.num_layers, 2, capacity, self.embed_dim), dtype=dtype, device=device) + if self._buffer is not None and self.length: + buffer[:, :, : self.length].copy_(self._buffer[:, :, : self.length]) + self._buffer = buffer + + def keys(self, layer: int) -> torch.Tensor: + """``[capacity, embed_dim]`` key rows of ``layer``.""" + assert self._buffer is not None + return self._buffer[layer, 0] + + def values(self, layer: int) -> torch.Tensor: + """``[capacity, embed_dim]`` value rows of ``layer``.""" + assert self._buffer is not None + return self._buffer[layer, 1] + + def head_view(self, rows: torch.Tensor) -> torch.Tensor: + """``[positions, embed_dim]`` rows as a ``[1, heads, positions, head_dim]`` view.""" + positions = int(rows.shape[0]) + return rows.view(positions, self.num_heads, self.embed_dim // self.num_heads).transpose(0, 1).unsqueeze(0) + + def history(self, length: int) -> torch.Tensor: + """``[layers, 2, length, embed_dim]`` rows ``[0, length)``, every layer at once.""" + assert self._buffer is not None + return self._buffer[:, :, :length] + + def commit(self, offset: int, new_kv: torch.Tensor) -> None: + """Write ``[layers, 2, positions, embed_dim]`` rows at ``offset``; the caller advances ``length``.""" + assert self._buffer is not None + positions = int(new_kv.shape[2]) + self._buffer[:, :, offset : offset + positions].copy_(new_kv) + + def layer_views(self, start: int, end: int) -> tuple[tuple[torch.Tensor, ...], ...]: + """Per-layer views of positions ``[start, end)``, built with a few ops for all layers. + + Returns ``(key_rows, value_rows, key_heads, value_heads)``: per layer the + ``[end - start, embed_dim]`` rows (write targets) and the + ``[1, heads, end, head_dim]`` attention inputs over ``[0, end)``. + """ + assert self._buffer is not None + head_dim = self.embed_dim // self.num_heads + heads = ( + self._buffer[:, :, :end] + .view(self.num_layers, 2, end, self.num_heads, head_dim) + .transpose(2, 3) + .unsqueeze(2) + ) + rows = self._buffer[:, :, start:end] + return rows[:, 0].unbind(0), rows[:, 1].unbind(0), heads[:, 0].unbind(0), heads[:, 1].unbind(0) + + def to_legacy_cache(self): + """The committed history as the ``EncoderDecoderCache`` the per-session path grows.""" + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + + self_attention = DynamicCache() + if self.length: + for layer in range(self.num_layers): + self_attention.update( + self.head_view(self.keys(layer)[: self.length]).contiguous(), + self.head_view(self.values(layer)[: self.length]).contiguous(), + layer, + ) + return EncoderDecoderCache(self_attention, DynamicCache()) + + +@dataclass(frozen=True, slots=True) +class StreamingAudioChunk: + """One session's ready unit: log-mel frames plus the state it continues.""" + + #: ``[n_mels, frames]`` or ``[1, n_mels, frames]`` log-mel features, host or device. + features: torch.Tensor + #: Session cache; ``None`` starts a new history. + cache: StreamingAudioKVCache | None + prefix_extra_frames: int = 0 + suffix_extra_frames: int = 0 + use_extra_context: bool = True + #: ``audio_feature_lens`` of the unit when it differs from ``features.shape[-1]``. + feature_length: int | None = None + + +@dataclass(slots=True) +class _Row: + index: int + frames: int + start: int + length: int + pooled_length: int + cache: StreamingAudioKVCache + past: int = 0 + offset: int = 0 + + @property + def total(self) -> int: + return self.past + self.length + + +def _extra_context_trim(frames: int) -> int: + return (int(frames) + 1) // 2 if frames > 0 else 0 + + +def streaming_batch_unsupported_reason(encoder: nn.Module, *, audio_encoder_layer: int | None) -> str | None: + """Why ``encoder`` cannot take the batched path, or ``None`` when it can.""" + if audio_encoder_layer != -1: + return "only the final encoder layer output is batched" + if encoder.training: + return "training mode (dropout / layerdrop)" + weight = getattr(getattr(encoder, "conv1", None), "weight", None) + if weight is None or weight.dtype not in _SUPPORTED_DTYPES: + # float16 needs the per-session inf/nan clamp, which is per tensor. + return f"unsupported encoder dtype {getattr(weight, 'dtype', None)}" + layers = getattr(encoder, "layers", None) + if not layers: + return "encoder has no layers" + for layer in layers: + attention = getattr(layer, "self_attn", None) + if attention is None or not all( + hasattr(attention, name) for name in ("q_proj", "k_proj", "v_proj", "out_proj", "scaling") + ): + return "encoder layers are not Whisper self-attention layers" + implementation = getattr(encoder.config, "_attn_implementation", None) or "eager" + if implementation not in _SUPPORTED_ATTENTION: + return f"attention implementation {implementation!r} is not batched" + return None + + +def _attend( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor, + implementation: str, +) -> torch.Tensor: + """Whisper self-attention of one row, exactly as the per-session path calls it. + + The query already carries ``scaling`` (Whisper scales after ``q_proj``), so + the kernel runs with ``scale=1``. Returns ``[1, heads, positions, head_dim]``. + """ + if implementation == "sdpa": + return torch.nn.functional.scaled_dot_product_attention( + query, key, value, attn_mask=mask, dropout_p=0.0, scale=1.0 + ) + weights = torch.matmul(query, key.transpose(2, 3)) + mask + return torch.matmul(torch.nn.functional.softmax(weights, dim=-1), value) + + +def _stack_features(rows: list[_Row], features: list[torch.Tensor], device: torch.device) -> torch.Tensor: + """Zero-padded ``[B, mels, max_frames]`` on ``device`` with one host-to-device copy.""" + max_frames = max(row.frames for row in rows) + on_host = all(feature.device.type == "cpu" for feature in features) + pin = on_host and device.type == "cuda" + batch = torch.zeros( + (len(rows), int(features[0].shape[0]), max_frames), + dtype=features[0].dtype, + device="cpu" if on_host else device, + pin_memory=pin, + ) + for slot, (row, feature) in enumerate(zip(rows, features, strict=True)): + batch[slot, :, : row.frames].copy_(feature) + # Pinned staging keeps the host from waiting on queued GPU work; the + # caching host allocator holds the buffer until the copy completes. + return batch.to(device, non_blocking=pin) + + +def _position_rows(encoder: nn.Module, rows: list[_Row]) -> torch.Tensor: + """Learned positions ``past .. past + n`` per row; past the table, repeat its last row.""" + table = encoder.embed_positions.weight + capacity = int(table.shape[0]) + pieces: list[torch.Tensor] = [] + for row in rows: + end = row.past + row.length + if end <= capacity: + pieces.append(table[row.past : end]) + else: + pieces.append(table[row.past :]) + pieces.append(table[-1:].expand(end - max(row.past, capacity), -1)) + return pieces[0] if len(pieces) == 1 else torch.cat(pieces) + + +@torch.inference_mode() +def encode_streaming_audio_batch( + encoder: nn.Module, + projection: nn.Module, + pooler: nn.Module, + chunks: Sequence[StreamingAudioChunk], + *, + pool_step: int, + page_positions: int = DEFAULT_KV_PAGE_POSITIONS, + new_cache: Callable[[StreamingAudioKVCache | None], StreamingAudioKVCache] | None = None, +) -> tuple[list[torch.Tensor | None], list[StreamingAudioKVCache | None]]: + """Encode one streaming unit per session in one encoder pass. + + Mirrors ``MiniCPMO45OmniLLMForConditionalGeneration.get_audio_embedding_streaming`` + row by row: same CNN extra-context trim, same reset when the history would + reach ``max_source_positions``, same learned positions, same pooled-length + clip. Returns, per chunk, the ``[tokens, hidden]`` embeddings (``None`` when + the unit is empty after trimming, which leaves its cache untouched) and the + cache now holding that session's history. + + ``new_cache`` (default: a fresh paged :class:`StreamingAudioKVCache`) builds + the cache of a new session or of a reset, from the row's current cache; the + resident graph encoder passes one that hands out slot-pool caches. + """ + implementation = getattr(encoder.config, "_attn_implementation", None) or "eager" + weight = encoder.conv1.weight + dtype, device = weight.dtype, weight.device + max_positions = int(encoder.embed_positions.weight.shape[0]) + num_layers = len(encoder.layers) + embed_dim = int(encoder.config.d_model) + num_heads = int(encoder.config.encoder_attention_heads) + + outputs: list[torch.Tensor | None] = [None] * len(chunks) + caches: list[StreamingAudioKVCache | None] = [chunk.cache for chunk in chunks] + rows: list[_Row] = [] + features: list[torch.Tensor] = [] + for index, chunk in enumerate(chunks): + feature = chunk.features + if feature.ndim == 3: + if feature.shape[0] != 1: + raise ValueError("a streaming chunk carries exactly one audio") + feature = feature[0] + frames = int(feature.shape[-1]) + conv_length = (frames - 1) // 2 + 1 + prefix = _extra_context_trim(chunk.prefix_extra_frames) if chunk.use_extra_context else 0 + suffix = _extra_context_trim(chunk.suffix_extra_frames) if chunk.use_extra_context else 0 + length = conv_length - prefix - suffix + if frames <= 0 or length <= 0: + continue + feature_length = frames if chunk.feature_length is None else int(chunk.feature_length) + pooled_length = ((feature_length - 1) // 2 + 1 - pool_step) // pool_step + 1 + cache = chunk.cache + # A new session, or the per-session path's reset: same bound, same + # moment. The old cache is dropped only when the caller commits this one. + if new_cache is not None and (cache is None or cache.length + length >= max_positions): + cache = new_cache(cache) + elif cache is None or cache.length + length >= max_positions: + cache = StreamingAudioKVCache( + num_layers=num_layers, + embed_dim=embed_dim, + num_heads=num_heads, + max_positions=max_positions, + page_positions=page_positions, + ) + rows.append( + _Row( + index=index, + frames=frames, + start=prefix, + length=length, + pooled_length=pooled_length, + cache=cache, + past=cache.length, + ) + ) + features.append(feature) + if not rows: + return outputs, caches + + offset = 0 + for row in rows: + row.offset = offset + offset += row.length + row.cache.reserve(row.total, dtype=dtype, device=device) + + # --- CNN front end on the zero-padded batch ------------------------------- + mel = _stack_features(rows, features, device).to(dtype=dtype) + hidden = nn.functional.gelu(encoder.conv1(mel)) + max_frames = int(mel.shape[-1]) + for slot, row in enumerate(rows): + if row.frames < max_frames: + # conv2 must see zeros past each row's end, as its own padding + # would in the per-session call. + hidden[slot, :, row.frames :].zero_() + hidden = nn.functional.gelu(encoder.conv2(hidden)).permute(0, 2, 1) + pieces = [hidden[slot, row.start : row.start + row.length] for slot, row in enumerate(rows)] + hidden = pieces[0] if len(pieces) == 1 else torch.cat(pieces) + hidden = hidden + _position_rows(encoder, rows) + + # --- Transformer layers on the packed rows -------------------------------- + # Every per-row view is built once here, not per layer: at batch sizes + # that matter the host cost of view ops, not the GPU, bounds this loop. + views = [row.cache.layer_views(row.past, row.total) for row in rows] + lengths = [row.length for row in rows] + total_rows = sum(lengths) + head_dim = embed_dim // num_heads + # The per-session path passes an all-zero additive mask; keep it so each + # row dispatches to the same attention kernel. + masks = [torch.zeros((1, 1, row.length, row.total), dtype=dtype, device=device) for row in rows] + for layer_index, layer in enumerate(encoder.layers): + attention = layer.self_attn + residual = hidden + normed = layer.self_attn_layer_norm(hidden) + query = attention.q_proj(normed) * attention.scaling + key = attention.k_proj(normed) + value = attention.v_proj(normed) + destinations = [view[0][layer_index] for view in views] + [view[1][layer_index] for view in views] + torch._foreach_copy_(destinations, [*key.split(lengths), *value.split(lengths)]) + queries = query.view(1, total_rows, num_heads, head_dim).transpose(1, 2).split(lengths, dim=2) + attended = [ + _attend(row_query, view[2][layer_index], view[3][layer_index], mask, implementation) + for row_query, view, mask in zip(queries, views, masks, strict=True) + ] + # [1, heads, rows, head_dim] per row -> packed [rows, embed_dim]. + attended_rows = (attended[0] if len(attended) == 1 else torch.cat(attended, dim=2)).transpose(1, 2) + hidden = residual + attention.out_proj(attended_rows.reshape(total_rows, embed_dim)) + residual = hidden + hidden = layer.activation_fn(layer.fc1(layer.final_layer_norm(hidden))) + hidden = residual + layer.fc2(hidden) + hidden = encoder.layer_norm(hidden) + + # --- Projection and pooling ------------------------------------------------ + embeds = projection(hidden) + if all(row.length % pool_step == 0 for row in rows): + pooled = pooler(embeds.transpose(0, 1).unsqueeze(0)).squeeze(0).transpose(0, 1) + for row in rows: + start = row.offset // pool_step + count = min(row.length // pool_step, row.pooled_length) + outputs[row.index] = pooled[start : start + count] + else: + for row in rows: + segment = embeds[row.offset : row.offset + row.length].transpose(0, 1).unsqueeze(0) + pooled = pooler(segment).squeeze(0).transpose(0, 1) + outputs[row.index] = pooled[: row.pooled_length] + + for row in rows: + row.cache.length = row.total + caches[row.index] = row.cache + return outputs, caches + + +__all__ = [ + "DEFAULT_KV_PAGE_POSITIONS", + "StreamingAudioChunk", + "StreamingAudioKVCache", + "encode_streaming_audio_batch", + "streaming_batch_unsupported_reason", +] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder_graph.py b/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder_graph.py new file mode 100644 index 00000000000..efa2a842ba6 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder_graph.py @@ -0,0 +1,743 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA graphs for MiniCPM-o's streaming Whisper encoder, for the steady unit shape. + +``encode_streaming_audio_batch`` still launches one attention kernel per row +per layer, so launches scale with the batch; a CUDA graph replays the whole +forward with one launch. + +A Stage-0 duplex *steady* unit (every unit but a session's first) always has +``prefix_extra_frames=2``, ``suffix_extra_frames=2`` and one processor chunk of +frames, so its post-trim length ``unit_length`` is constant. Only the batch size +and each row's committed history vary, so graphs are captured per +``(batch bucket, cache-length bucket)``: + +* padding rows replay stale buffer content, but their output is never read and + rows never attend across the batch; +* a row's history is copied in right-aligned against the bucket + (``bucket - past .. bucket``) and the unused prefix is masked to ``-inf``; + the new unit is always written at ``bucket .. bucket + unit_length``. + +Real positions are placed and masked exactly as the eager batch places them. +Rows this module cannot describe (a first unit, a short final unit, history +past the largest bucket) go to the eager batch. + +Only one graph replays at a time and every replay refills its inputs, so all +graphs take prefix views of one shared storage per static input, sized for the +largest bucket (like ``WholeEulerExecutionArena``). Graphs are captured once, at +load (``build_streaming_audio_graph_encoder``), never while serving. + +Resident mode (``resident_slots > 0``, default off): instead of a shared +``[layers, 2, max_batch, bucket + unit, d]`` history storage that every replay +refills from the session caches (and copies the new unit back out of), every +session's cache lives in one slot of a :class:`StreamingAudioKVSlotPool` and +the graph reads its history from, and writes its new unit into, that slot. +Inside the graph each layer gathers the history into a temporary with the +shared storage's exact layout (right-aligned, same strides), so attention sees +the same values at every unmasked position as the copy-in path. The pool +replaces both the shared storage and the per-session paged buffers. +""" + +from __future__ import annotations + +import logging +import weakref +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from torch import nn + +from .streaming_audio_encoder import ( + DEFAULT_KV_PAGE_POSITIONS, + StreamingAudioChunk, + StreamingAudioKVCache, + _attend, + _extra_context_trim, + encode_streaming_audio_batch, +) + +if TYPE_CHECKING: + from torch.cuda import CUDAGraph + +_NEG_INF = float("-inf") + +logger = logging.getLogger(__name__) + +#: ``stage0.py`` builds every unit but a session's first with these trims +#: (``("prefix_extra_frames", 0 if state.audio_chunk_idx == 0 else 2)``, +#: ``("suffix_extra_frames", 2)``); only that steady shape is captured. +STEADY_PREFIX_EXTRA_FRAMES = 2 +STEADY_SUFFIX_EXTRA_FRAMES = 2 + +DEFAULT_GRAPH_BATCH_SIZES: tuple[int, ...] = (1, 2, 4, 8, 16, 32) +#: Shared storage scales with the *largest* bucket only. 900 keeps it under +#: 3 GiB at batch 32 on MiniCPM-o 4.5's encoder (bf16, measured 2.8-2.9 GiB); +#: history past it (the reset bound is ~1450) falls back to eager. +DEFAULT_GRAPH_CACHE_BUCKETS: tuple[int, ...] = (250, 500, 750, 900) + + +def batch_sizes_for_sessions(max_sessions: int) -> tuple[int, ...]: + """Powers of two below ``max_sessions`` plus ``max_sessions`` itself. + + Each live session contributes at most one steady row per step, so a larger + bucket is never selected; 16 -> (1, 2, 4, 8, 16), 20 -> (1, 2, 4, 8, 16, 20). + """ + top = max(1, int(max_sessions)) + sizes = [1 << i for i in range(top.bit_length()) if (1 << i) < top] + return normalize_buckets([*sizes, top]) + + +def select_bucket(value: int, buckets: Sequence[int]) -> int | None: + """The smallest bucket ``>= value``, or ``None`` past the largest one.""" + return min((bucket for bucket in buckets if value <= bucket), default=None) + + +def normalize_buckets(values: Sequence[int]) -> tuple[int, ...]: + """Sorted, deduplicated, positive bucket values.""" + return tuple(sorted({int(v) for v in values if int(v) > 0})) + + +def steady_unit_length(unit_frames: int) -> int: + """Post-trim encoder positions of one steady unit of ``unit_frames`` mel frames.""" + conv_length = (int(unit_frames) - 1) // 2 + 1 + trim = _extra_context_trim(STEADY_PREFIX_EXTRA_FRAMES) + _extra_context_trim(STEADY_SUFFIX_EXTRA_FRAMES) + return conv_length - trim + + +def steady_pooled_length(unit_frames: int, unit_length: int, pool_step: int) -> int: + """The pooled length ``encode_streaming_audio_batch`` emits for a steady row. + + The smaller of its ``pooled_length`` cap (from the raw frame count) and the + pooled post-trim length, like its ``pooled[: row.pooled_length]``. + """ + nominal_cap = ((unit_frames - 1) // 2 + 1 - pool_step) // pool_step + 1 + return min(nominal_cap, unit_length // pool_step) + + +def row_is_steady(chunk: StreamingAudioChunk, *, expected_frames: int) -> bool: + """Whether ``chunk`` is a steady unit of ``expected_frames`` frames: the captured shape.""" + if chunk.cache is None or not chunk.use_extra_context: + return False + if int(chunk.prefix_extra_frames) != STEADY_PREFIX_EXTRA_FRAMES: + return False + if int(chunk.suffix_extra_frames) != STEADY_SUFFIX_EXTRA_FRAMES: + return False + feature = chunk.features + frames = int(feature.shape[-1]) + if frames != expected_frames: + return False + return chunk.feature_length is None or int(chunk.feature_length) == frames + + +class PooledStreamingAudioKVCache(StreamingAudioKVCache): + """A :class:`StreamingAudioKVCache` whose buffer is one slot of a :class:`StreamingAudioKVSlotPool`. + + Capacity is the full ``max_positions`` from the start (``reserve`` never + reallocates); the slot goes back to the pool when this object is collected, + unless :meth:`StreamingAudioKVSlotPool.take_over` handed it to a successor. + """ + + __slots__ = ("__weakref__", "_finalizer", "pool", "slot") + + def __init__(self, pool: StreamingAudioKVSlotPool, slot: int) -> None: + super().__init__( + num_layers=pool.num_layers, + embed_dim=pool.embed_dim, + num_heads=pool.num_heads, + max_positions=pool.max_positions, + page_positions=pool.max_positions, + ) + self.pool = pool + self.slot = int(slot) + self._buffer = pool.storage[self.slot, :, :, : pool.max_positions] + self._finalizer = weakref.finalize(self, pool._release, self.slot) + + +class StreamingAudioKVSlotPool: + """Fixed homes for the streaming caches of up to ``num_slots`` sessions. + + ``storage`` is ``[num_slots, layers, 2, max_positions + tail, d]``: positions + ``[0, max_positions)`` of a slot are its session's history (a + :class:`PooledStreamingAudioKVCache` view); the ``tail`` positions after them + are never history and take the resident graph's padding-row writes. + """ + + def __init__( + self, + *, + num_slots: int, + num_layers: int, + embed_dim: int, + num_heads: int, + max_positions: int, + tail_positions: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + if min(num_slots, num_layers, embed_dim, num_heads, max_positions, tail_positions) <= 0: + raise ValueError("streaming audio KV slot pool dimensions must be positive") + self.num_slots = int(num_slots) + self.num_layers = int(num_layers) + self.embed_dim = int(embed_dim) + self.num_heads = int(num_heads) + self.max_positions = int(max_positions) + self.tail_positions = int(tail_positions) + self.storage = torch.zeros( + (self.num_slots, self.num_layers, 2, self.max_positions + self.tail_positions, self.embed_dim), + dtype=dtype, + device=device, + ) + self._free: list[int] = list(range(self.num_slots - 1, -1, -1)) + self._warned_exhausted = False + + @property + def free_slots(self) -> int: + return len(self._free) + + def acquire(self) -> PooledStreamingAudioKVCache | None: + """An empty cache in a free slot, or ``None`` when every slot is taken.""" + if not self._free: + if not self._warned_exhausted: + self._warned_exhausted = True + logger.warning( + "MiniCPM-o resident audio KV pool exhausted (%d slots); further sessions use paged " + "caches and the eager batched encoder", + self.num_slots, + ) + return None + return PooledStreamingAudioKVCache(self, self._free.pop()) + + def take_over(self, cache: PooledStreamingAudioKVCache) -> PooledStreamingAudioKVCache: + """An empty cache in ``cache``'s slot (a reset); ``cache`` no longer frees the slot. + + The slot's memory is reused in place, so unlike a paged reset a forward + that fails after this leaves ``cache``'s history partly overwritten. + """ + assert cache.pool is self + cache._finalizer.detach() + return PooledStreamingAudioKVCache(self, cache.slot) + + def _release(self, slot: int) -> None: + self._free.append(slot) + + +@dataclass(slots=True) +class _Group: + """One steady row this call will replay through a graph.""" + + index: int + cache_in: StreamingAudioKVCache + cache_out: StreamingAudioKVCache + past: int + + +class StreamingAudioGraphEncoder: + """Pre-captured CUDA graphs of the streaming encoder's steady unit shape. + + :meth:`capture` builds one graph per ``(batch size, cache bucket)``; + :meth:`encode` replays them and hands every other row to + :func:`encode_streaming_audio_batch`. + """ + + def __init__( + self, + encoder: nn.Module, + projection: nn.Module, + pooler: nn.Module, + *, + unit_frames: int, + pool_step: int, + batch_sizes: Sequence[int] = DEFAULT_GRAPH_BATCH_SIZES, + cache_buckets: Sequence[int] = DEFAULT_GRAPH_CACHE_BUCKETS, + page_positions: int = DEFAULT_KV_PAGE_POSITIONS, + pinned_h2d: bool = False, + resident_slots: int = 0, + ) -> None: + weight = encoder.conv1.weight + if weight.device.type != "cuda": + raise ValueError("StreamingAudioGraphEncoder requires a CUDA encoder") + self.encoder = encoder + self.projection = projection + self.pooler = pooler + self.device = weight.device + self.dtype = weight.dtype + self.pool_step = int(pool_step) + self.page_positions = int(page_positions) + # Stage host mel through pinned memory (see ``_stage_mel``). + self.pinned_h2d = bool(pinned_h2d) + self.num_layers = len(encoder.layers) + self.embed_dim = int(encoder.config.d_model) + self.num_heads = int(encoder.config.encoder_attention_heads) + self.num_mels = int(encoder.config.num_mel_bins) + self.max_positions = int(encoder.embed_positions.weight.shape[0]) + self.implementation = getattr(encoder.config, "_attn_implementation", None) or "eager" + self.unit_frames = int(unit_frames) + self.unit_length = steady_unit_length(self.unit_frames) + if self.unit_length <= 0: + raise ValueError(f"unit_frames={unit_frames} is too small for the steady prefix/suffix trim") + self.pooled_length = steady_pooled_length(self.unit_frames, self.unit_length, self.pool_step) + if self.pooled_length <= 0: + raise ValueError( + f"unit_frames={unit_frames} pool_step={pool_step} leaves no pooled output for the steady unit" + ) + self.batch_sizes = normalize_buckets(batch_sizes) + self.cache_buckets = normalize_buckets(cache_buckets) + if not self.batch_sizes or not self.cache_buckets: + raise ValueError("StreamingAudioGraphEncoder needs at least one batch size and one cache bucket") + self.max_batch = max(self.batch_sizes) + self.max_cache_bucket = max(self.cache_buckets) + self.max_total = self.max_cache_bucket + self.unit_length + # Resident mode (see the module docstring); 0 keeps the copy-in storage. + self.resident_slots = max(0, int(resident_slots)) + + self._graphs: dict[tuple[int, int], CUDAGraph] = {} + self._pooled: dict[tuple[int, int], torch.Tensor] = {} + # Shared static-input storage (see the module docstring), allocated by capture(). + self._mel_storage: torch.Tensor | None = None + self._cache_storage: torch.Tensor | None = None + self._mask_storage: torch.Tensor | None = None + self._position_ids_storage: torch.Tensor | None = None + # Resident mode: the slot pool and per-row ``[slot, history pad, write offset]``. + self._pool: StreamingAudioKVSlotPool | None = None + self._row_meta_storage: torch.Tensor | None = None + self._iota: torch.Tensor | None = None + + @property + def resident(self) -> bool: + return self.resident_slots > 0 + + def _allocate_storage(self) -> None: + if self.resident: + self._allocate_resident_storage() + return + self._mel_storage = torch.zeros( + (self.max_batch, self.num_mels, self.unit_frames), dtype=self.dtype, device=self.device + ) + self._cache_storage = torch.zeros( + (self.num_layers, 2, self.max_batch, self.max_total, self.embed_dim), + dtype=self.dtype, + device=self.device, + ) + self._mask_storage = torch.zeros( + (self.max_batch, 1, self.unit_length, self.max_total), dtype=self.dtype, device=self.device + ) + self._position_ids_storage = torch.zeros( + (self.max_batch, self.unit_length), dtype=torch.long, device=self.device + ) + + def _allocate_resident_storage(self) -> None: + self._mel_storage = torch.zeros( + (self.max_batch, self.num_mels, self.unit_frames), dtype=self.dtype, device=self.device + ) + self._pool = StreamingAudioKVSlotPool( + num_slots=self.resident_slots, + num_layers=self.num_layers, + embed_dim=self.embed_dim, + num_heads=self.num_heads, + max_positions=self.max_positions, + tail_positions=self.unit_length, + dtype=self.dtype, + device=self.device, + ) + self._mask_storage = torch.zeros( + (self.max_batch, 1, self.unit_length, self.max_total), dtype=self.dtype, device=self.device + ) + self._position_ids_storage = torch.zeros( + (self.max_batch, self.unit_length), dtype=torch.long, device=self.device + ) + meta = torch.zeros((self.max_batch, 3), dtype=torch.long, device=self.device) + # Until the first replay every row is a padding row: slot 0's tail. + meta[:, 2] = self.max_positions + self._row_meta_storage = meta + self._iota = torch.arange(self.max_cache_bucket, dtype=torch.long, device=self.device) + + def new_cache(self, cache: StreamingAudioKVCache | None) -> StreamingAudioKVCache: + """The cache a new session or a reset starts from: a pool slot when one is free.""" + pool = self._pool + if pool is not None: + if isinstance(cache, PooledStreamingAudioKVCache) and cache.pool is pool: + return pool.take_over(cache) + pooled = pool.acquire() + if pooled is not None: + return pooled + return StreamingAudioKVCache( + num_layers=self.num_layers, + embed_dim=self.embed_dim, + num_heads=self.num_heads, + max_positions=self.max_positions, + page_positions=self.page_positions, + ) + + def _views(self, batch: int, cache_len: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Prefix views of the shared storage for ``(batch, cache_len)``: the memory its graph captured.""" + assert self._cache_storage is not None + total = cache_len + self.unit_length + mel = self._mel_storage[:batch] + cache = self._cache_storage[:, :, :batch, :total, :] + mask = self._mask_storage[:batch, :, :, :total] + position_ids = self._position_ids_storage[:batch] + return mel, cache, mask, position_ids + + # --- capture ------------------------------------------------------------- + + def capture(self) -> None: + """Capture every ``(batch, cache-length)`` graph. Call once, at load.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("cannot capture the streaming audio encoder graph mid-capture") + from vllm.platforms import current_platform + + self._allocate_storage() + pool = current_platform.get_global_graph_pool() + for batch in self.batch_sizes: + for cache_len in self.cache_buckets: + self._capture_one(batch, cache_len, pool) + + def _capture_one(self, batch: int, cache_len: int, pool: object) -> None: + key = (batch, cache_len) + if self.resident: + mel, mask, position_ids, meta = self._resident_views(batch, cache_len) + self._capture_forward(key, lambda: self._forward_resident(mel, mask, position_ids, meta, cache_len), pool) + return + mel, cache, mask, position_ids = self._views(batch, cache_len) + + current_stream = torch.cuda.current_stream(self.device) + warmup_stream = torch.cuda.Stream(device=self.device) + warmup_stream.wait_stream(current_stream) + with torch.cuda.stream(warmup_stream), torch.inference_mode(): + for _ in range(3): + warmup = self._forward(mel, cache, mask, position_ids, cache_len) + current_stream.wait_stream(warmup_stream) + del warmup + + graph = torch.cuda.CUDAGraph() + with torch.inference_mode(), torch.cuda.graph(graph, pool=pool): + pooled = self._forward(mel, cache, mask, position_ids, cache_len) + + self._graphs[key] = graph + self._pooled[key] = pooled + + def _capture_forward(self, key: tuple[int, int], forward, pool: object) -> None: + current_stream = torch.cuda.current_stream(self.device) + warmup_stream = torch.cuda.Stream(device=self.device) + warmup_stream.wait_stream(current_stream) + with torch.cuda.stream(warmup_stream), torch.inference_mode(): + for _ in range(3): + warmup = forward() + current_stream.wait_stream(warmup_stream) + del warmup + + graph = torch.cuda.CUDAGraph() + with torch.inference_mode(), torch.cuda.graph(graph, pool=pool): + pooled = forward() + + self._graphs[key] = graph + self._pooled[key] = pooled + + def _resident_views( + self, batch: int, cache_len: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + assert self._mask_storage is not None and self._row_meta_storage is not None + total = cache_len + self.unit_length + return ( + self._mel_storage[:batch], + self._mask_storage[:batch, :, :, :total], + self._position_ids_storage[:batch], + self._row_meta_storage[:batch], + ) + + def _forward_resident( + self, + mel: torch.Tensor, + mask: torch.Tensor, + position_ids: torch.Tensor, + meta: torch.Tensor, + cache_len: int, + ) -> torch.Tensor: + """:meth:`_forward` reading/writing each row's pool slot instead of the shared storage. + + ``meta[:, 0]`` is the row's slot, ``meta[:, 1]`` the history pad + (``cache_len - past``), ``meta[:, 2]`` the write offset (0, or + ``max_positions`` for a padding row, whose new unit lands in the tail). + Per layer the history is gathered into a ``[2, batch, max_total, d]`` + temporary laid out like the shared storage, then the new unit is + written into the slot. + """ + assert self._pool is not None and self._iota is not None + storage = self._pool.storage + encoder = self.encoder + batch = mel.shape[0] + total = cache_len + self.unit_length + slots = meta[:, 0:1] + # History index j of a row reads slot position j - pad (pad positions + # read position 0: finite, masked to -inf like the shared storage's stale prefix). + history_positions = (self._iota[:cache_len].unsqueeze(0) - meta[:, 1:2]).clamp_(min=0) if cache_len else None + write_positions = position_ids + meta[:, 2:3] + hidden = nn.functional.gelu(encoder.conv1(mel)) + hidden = nn.functional.gelu(encoder.conv2(hidden)).permute(0, 2, 1) + prefix = _extra_context_trim(STEADY_PREFIX_EXTRA_FRAMES) + hidden = hidden[:, prefix : prefix + self.unit_length, :] + hidden = hidden + encoder.embed_positions(position_ids) + head_dim = self.embed_dim // self.num_heads + for layer_index, layer in enumerate(encoder.layers): + attention = layer.self_attn + residual = hidden + normed = layer.self_attn_layer_norm(hidden) + query = attention.q_proj(normed) * attention.scaling + key = attention.k_proj(normed) + value = attention.v_proj(normed) + kv = torch.empty((2, batch, self.max_total, self.embed_dim), dtype=self.dtype, device=self.device) + if history_positions is not None: + kv[0, :, :cache_len].copy_(storage[slots, layer_index, 0, history_positions]) + kv[1, :, :cache_len].copy_(storage[slots, layer_index, 1, history_positions]) + kv[0, :, cache_len:total].copy_(key) + kv[1, :, cache_len:total].copy_(value) + storage[slots, layer_index, 0, write_positions] = key + storage[slots, layer_index, 1, write_positions] = value + q = query.view(batch, self.unit_length, self.num_heads, head_dim).transpose(1, 2) + k_all = kv[0, :, :total].view(batch, -1, self.num_heads, head_dim).transpose(1, 2) + v_all = kv[1, :, :total].view(batch, -1, self.num_heads, head_dim).transpose(1, 2) + attended = _attend(q, k_all, v_all, mask, self.implementation) + attended = attended.transpose(1, 2).reshape(batch, self.unit_length, self.embed_dim) + hidden = residual + attention.out_proj(attended) + residual = hidden + hidden = layer.activation_fn(layer.fc1(layer.final_layer_norm(hidden))) + hidden = residual + layer.fc2(hidden) + hidden = encoder.layer_norm(hidden) + embeds = self.projection(hidden) + return self.pooler(embeds.transpose(1, 2)).transpose(1, 2) + + def _forward( + self, + mel: torch.Tensor, + cache: torch.Tensor, + mask: torch.Tensor, + position_ids: torch.Tensor, + cache_len: int, + ) -> torch.Tensor: + """``encode_streaming_audio_batch``'s per-row math with equal row shapes, as batched ops. + + Padding rows and history positions are inert only through the + caller's mask (see the module docstring). + """ + encoder = self.encoder + batch = mel.shape[0] + hidden = nn.functional.gelu(encoder.conv1(mel)) + hidden = nn.functional.gelu(encoder.conv2(hidden)).permute(0, 2, 1) + prefix = _extra_context_trim(STEADY_PREFIX_EXTRA_FRAMES) + hidden = hidden[:, prefix : prefix + self.unit_length, :] + hidden = hidden + encoder.embed_positions(position_ids) + head_dim = self.embed_dim // self.num_heads + for layer_index, layer in enumerate(encoder.layers): + attention = layer.self_attn + residual = hidden + normed = layer.self_attn_layer_norm(hidden) + query = attention.q_proj(normed) * attention.scaling + key = attention.k_proj(normed) + value = attention.v_proj(normed) + cache[layer_index, 0, :, cache_len : cache_len + self.unit_length, :].copy_(key) + cache[layer_index, 1, :, cache_len : cache_len + self.unit_length, :].copy_(value) + q = query.view(batch, self.unit_length, self.num_heads, head_dim).transpose(1, 2) + k_all = cache[layer_index, 0].view(batch, -1, self.num_heads, head_dim).transpose(1, 2) + v_all = cache[layer_index, 1].view(batch, -1, self.num_heads, head_dim).transpose(1, 2) + attended = _attend(q, k_all, v_all, mask, self.implementation) + attended = attended.transpose(1, 2).reshape(batch, self.unit_length, self.embed_dim) + hidden = residual + attention.out_proj(attended) + residual = hidden + hidden = layer.activation_fn(layer.fc1(layer.final_layer_norm(hidden))) + hidden = residual + layer.fc2(hidden) + hidden = encoder.layer_norm(hidden) + embeds = self.projection(hidden) + return self.pooler(embeds.transpose(1, 2)).transpose(1, 2) + + # --- replay ---------------------------------------------------------------- + + def encode( + self, + chunks: Sequence[StreamingAudioChunk], + ) -> tuple[list[torch.Tensor | None], list[StreamingAudioKVCache | None]]: + """``encode_streaming_audio_batch``'s contract: covered steady rows replay a graph, the rest run eagerly.""" + outputs: list[torch.Tensor | None] = [None] * len(chunks) + caches: list[StreamingAudioKVCache | None] = [chunk.cache for chunk in chunks] + steady: list[_Group] = [] + eager_indices: list[int] = [] + for index, chunk in enumerate(chunks): + if not row_is_steady(chunk, expected_frames=self.unit_frames): + eager_indices.append(index) + continue + cache_in = chunk.cache + assert cache_in is not None + if self.resident and not ( + isinstance(cache_in, PooledStreamingAudioKVCache) and cache_in.pool is self._pool + ): + # A paged cache (pool exhausted when it started): eager, in place. + eager_indices.append(index) + continue + reset = cache_in.length + self.unit_length >= self.max_positions + past = 0 if reset else cache_in.length + if past > self.max_cache_bucket: + eager_indices.append(index) + continue + # Resident: a reset reuses the row's slot in place (past is 0). + cache_out = ( + StreamingAudioKVCache( + num_layers=self.num_layers, + embed_dim=self.embed_dim, + num_heads=self.num_heads, + max_positions=self.max_positions, + page_positions=self.page_positions, + ) + if reset and not self.resident + else cache_in + ) + steady.append(_Group(index=index, cache_in=cache_in, cache_out=cache_out, past=past)) + + for start in range(0, len(steady), self.max_batch): + group = steady[start : start + self.max_batch] + if self.resident: + self._replay_resident(group, chunks, outputs, caches) + else: + self._replay(group, chunks, outputs, caches) + + if eager_indices: + sub_chunks = [chunks[i] for i in eager_indices] + sub_outputs, sub_caches = encode_streaming_audio_batch( + self.encoder, + self.projection, + self.pooler, + sub_chunks, + pool_step=self.pool_step, + page_positions=self.page_positions, + new_cache=self.new_cache if self.resident else None, + ) + for local, index in enumerate(eager_indices): + outputs[index] = sub_outputs[local] + caches[index] = sub_caches[local] + + return outputs, caches + + def _stage_mel(self, dst: torch.Tensor, feature: torch.Tensor) -> None: + """Write one unit's log-mel into its static slot. + + The processor builds the features on the host. A pageable + ``.to(device)`` is a blocking copy that first waits for all queued work + on the stream -- under async scheduling, the previous step's forward -- + so the next step's host preparation stalls in the middle. With + ``pinned_h2d`` the dtype conversion stays on the host exactly as the + blocking copy does it (an H2D copy converts on the source side), and + only the transfer becomes a pinned, stream-ordered copy: same values, + no host wait. + """ + if self.pinned_h2d and feature.device.type == "cpu": + dst.copy_(feature.to(dtype=self.dtype).pin_memory(), non_blocking=True) + else: + dst.copy_(feature.to(device=self.device, dtype=self.dtype)) + + @torch.inference_mode() + def _replay( + self, + group: list[_Group], + chunks: Sequence[StreamingAudioChunk], + outputs: list[torch.Tensor | None], + caches: list[StreamingAudioKVCache | None], + ) -> None: + batch = select_bucket(len(group), self.batch_sizes) + cache_len = select_bucket(max((row.past for row in group), default=0), self.cache_buckets) + assert batch is not None and cache_len is not None + key = (batch, cache_len) + + mel, cache, mask, position_ids = self._views(batch, cache_len) + + mask.fill_(0.0) + # Dummy/padding rows: mask their whole history away so only the + # (unread) new unit is attendable, which can never be all -inf. + if len(group) < batch: + mask[len(group) :, :, :, :cache_len].fill_(_NEG_INF) + position_ids[len(group) :].copy_(torch.arange(self.unit_length, device=self.device)) + + for slot, row in enumerate(group): + feature = chunks[row.index].features + if feature.ndim == 3: + feature = feature[0] + self._stage_mel(mel[slot], feature) + pad = cache_len - row.past + if pad: + mask[slot, :, :, :pad].fill_(_NEG_INF) + if row.past: + cache[:, :, slot, pad:cache_len, :].copy_(row.cache_in.history(row.past)) + position_ids[slot].copy_(torch.arange(row.past, row.past + self.unit_length, device=self.device)) + + self._graphs[key].replay() + + pooled = self._pooled[key] + for slot, row in enumerate(group): + outputs[row.index] = pooled[slot, : self.pooled_length].clone() + row.cache_out.reserve(row.past + self.unit_length, dtype=self.dtype, device=self.device) + row.cache_out.commit(row.past, cache[:, :, slot, cache_len : cache_len + self.unit_length, :]) + row.cache_out.length = row.past + self.unit_length + caches[row.index] = row.cache_out + + @torch.inference_mode() + def _replay_resident( + self, + group: list[_Group], + chunks: Sequence[StreamingAudioChunk], + outputs: list[torch.Tensor | None], + caches: list[StreamingAudioKVCache | None], + ) -> None: + """:meth:`_replay` for resident rows: no history copy-in, no new-unit copy-out.""" + batch = select_bucket(len(group), self.batch_sizes) + cache_len = select_bucket(max((row.past for row in group), default=0), self.cache_buckets) + assert batch is not None and cache_len is not None + key = (batch, cache_len) + + mel, mask, position_ids, meta = self._resident_views(batch, cache_len) + host_meta = torch.empty((batch, 3), dtype=torch.long, pin_memory=True) + mask.fill_(0.0) + if len(group) < batch: + mask[len(group) :, :, :, :cache_len].fill_(_NEG_INF) + position_ids[len(group) :].copy_(torch.arange(self.unit_length, device=self.device)) + host_meta[len(group) :, 0] = 0 + host_meta[len(group) :, 1] = cache_len + host_meta[len(group) :, 2] = self.max_positions + + for slot, row in enumerate(group): + feature = chunks[row.index].features + if feature.ndim == 3: + feature = feature[0] + self._stage_mel(mel[slot], feature) + pad = cache_len - row.past + if pad: + mask[slot, :, :, :pad].fill_(_NEG_INF) + position_ids[slot].copy_(torch.arange(row.past, row.past + self.unit_length, device=self.device)) + cache = row.cache_in + assert isinstance(cache, PooledStreamingAudioKVCache) + host_meta[slot, 0] = cache.slot + host_meta[slot, 1] = pad + host_meta[slot, 2] = 0 + meta.copy_(host_meta, non_blocking=True) + + self._graphs[key].replay() + + pooled = self._pooled[key] + for slot, row in enumerate(group): + outputs[row.index] = pooled[slot, : self.pooled_length].clone() + row.cache_out.length = row.past + self.unit_length + caches[row.index] = row.cache_out + + +__all__ = [ + "DEFAULT_GRAPH_BATCH_SIZES", + "DEFAULT_GRAPH_CACHE_BUCKETS", + "PooledStreamingAudioKVCache", + "STEADY_PREFIX_EXTRA_FRAMES", + "STEADY_SUFFIX_EXTRA_FRAMES", + "StreamingAudioGraphEncoder", + "StreamingAudioKVSlotPool", + "normalize_buckets", + "row_is_steady", + "select_bucket", + "steady_pooled_length", + "steady_unit_length", +] From e6fcbc72dccfbaf0ed77a55e4ffbaaba332af61d Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:30 +0800 Subject: [PATCH 80/89] [Perf][MiniCPM-o] Stage 0 fused vision encoder and CUDA graphs Stage 0 fused vision encoder and CUDA graphs Add packed SDPA vision encoding, cross-session frame batching, fused layer norm ops, and vision CUDA graph capture for multi-modal duplex serving. Signed-off-by: BeatSeat --- .../duplex/test_stage0_vision_batching.py | 278 ++++++++++++++++++ .../test_fused_residual_layer_norm.py | 126 ++++++++ .../models/minicpmo_4_5/test_vision_fused.py | 127 ++++++++ .../test_vision_packed_encoder.py | 272 +++++++++++++++++ .../models/minicpmo_4_5/ops/__init__.py | 7 + .../minicpmo_4_5/ops/fused_qkv_layer_norm.py | 205 +++++++++++++ .../ops/fused_residual_layer_norm.py | 273 +++++++++++++++++ .../models/minicpmo_4_5/vision_fused.py | 262 +++++++++++++++++ 8 files changed, 1550 insertions(+) create mode 100644 tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_vision_batching.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_fused_residual_layer_norm.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_vision_fused.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_vision_packed_encoder.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/vision_fused.py diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_vision_batching.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_vision_batching.py new file mode 100644 index 00000000000..f129ad88bd1 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_vision_batching.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniCPM-o 4.5 omni-duplex camera frames: one encoder call per append and across sessions.""" + +from __future__ import annotations + +import base64 +from io import BytesIO +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from PIL import Image + +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, +) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import MiniCPMO45OmniForConditionalGeneration +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, + MiniCPMVImageProcessor, + Resampler, + SiglipVisionConfig, + SiglipVisionTransformer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Processor: + """``process_image`` of the checkpoint's MiniCPMOProcessor, over the in-tree image processor.""" + + tokenizer = None + + def __init__(self) -> None: + self.image_processor = MiniCPMVImageProcessor(max_slice_nums=9, scale_resolution=448, patch_size=14) + self.calls: list[int] = [] + + def process_image(self, images, max_slice_nums=1): + self.calls.append(int(max_slice_nums)) + return self.image_processor.preprocess(images, max_slice_nums=max_slice_nums, return_tensors="pt") + + +class _Thinker: + """Real tiny SigLIP + resampler behind the Thinker's ``get_vision_hidden_states``.""" + + def __init__(self) -> None: + torch.manual_seed(0) + config = SiglipVisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + image_size=980, + patch_size=14, + ) + config._attn_implementation = "sdpa" + self.vpm = SiglipVisionTransformer(config).eval() + self.resampler = Resampler(num_queries=64, embed_dim=64, num_heads=4, kv_dim=32, adaptive=True).eval() + for param in self.resampler.parameters(): + torch.nn.init.normal_(param, std=0.2) + self.config = SimpleNamespace(vision_batch_size=16) + self.vision_packed_encode = True + self.llm = SimpleNamespace(model=SimpleNamespace(embed_tokens=torch.nn.Embedding(8, 64))) + self.encoder_calls: list[int] = [] + + def get_vision_hidden_states(self, data): + self.encoder_calls.append(len(data["pixel_values"])) + assert data["tgt_sizes"].device.type == "cpu" # grids stay on the host: no device sync + return MiniCPMO45OmniLLMForConditionalGeneration.get_vision_hidden_states(self, data) + + def parameters(self): + return self.vpm.parameters() + + +def _runtime() -> MiniCPMO45Stage0DuplexRuntime: + stage_model = SimpleNamespace(thinker=_Thinker(), processor=_Processor()) + return MiniCPMO45Stage0DuplexRuntime(stage_model, device="cpu") + + +def _jpeg(width: int, height: int, seed: int) -> str: + generator = torch.Generator().manual_seed(seed) + pixels = torch.randint(0, 256, (height, width, 3), generator=generator, dtype=torch.uint8).numpy() + buffer = BytesIO() + Image.fromarray(pixels).save(buffer, format="JPEG") + return base64.b64encode(buffer.getvalue()).decode("ascii") + + +def _append(session_id: str, seq: int, frames: list[str], epoch: int = 0) -> dict[str, Any]: + return { + "data_plane": True, + "session_id": session_id, + "epoch": epoch, + "seq": seq, + "payload": {"audio": "", "video_frames": frames}, + } + + +def _assert_same_blocks(actual: list[list[torch.Tensor]], expected: list[list[torch.Tensor]]) -> None: + assert [len(frame) for frame in actual] == [len(frame) for frame in expected] + for actual_frame, expected_frame in zip(actual, expected): + for actual_block, expected_block in zip(actual_frame, expected_frame): + assert actual_block.shape == (64, 64) + torch.testing.assert_close(actual_block, expected_block, rtol=1e-5, atol=1e-5) + + +def _per_frame_reference(runtime: MiniCPMO45Stage0DuplexRuntime, frames: list[Image.Image]) -> list[list[torch.Tensor]]: + """What the per-frame loop computed: each frame sliced and encoded alone.""" + reference = [] + for frame, max_slices in zip(frames, runtime._official_max_slice_nums(len(frames))): + processed = runtime.processor.process_image([frame], max_slice_nums=max_slices) + encoded = runtime._encode_processed_vision_batch([processed]) + assert encoded is not None + reference.append(encoded[0]) + return reference + + +def test_stacked_pair_is_encoded_in_one_call_and_matches_per_frame() -> None: + runtime = _runtime() + frames = runtime._decode_video_frames_payload({"video_frames": [_jpeg(960, 540, 1), _jpeg(640, 480, 2)]}) + expected = _per_frame_reference(runtime, frames) + runtime.thinker.encoder_calls.clear() + + with torch.inference_mode(): + encoded = runtime._stage_vision_embeddings(frames) + + assert runtime.thinker.encoder_calls == [4] # HD base frame: 1 source + 2 slices, composite: 1 + assert encoded is not None + _assert_same_blocks(encoded, expected) + + +def test_prefetch_encodes_all_sessions_in_one_call() -> None: + runtime = _runtime() + appends = [ + _append("a", 3, [_jpeg(640, 480, 10)]), + _append("b", 7, [_jpeg(960, 540, 11), _jpeg(640, 480, 12)]), + _append("c", 1, [_jpeg(480, 640, 13)]), + ] + expected = { + duplex["session_id"]: _per_frame_reference(runtime, runtime._decode_video_frames_payload(duplex["payload"])) + for duplex in appends + } + runtime.thinker.encoder_calls.clear() + + with torch.inference_mode(): + runtime.prefetch_vision(appends) + + assert runtime.thinker.encoder_calls == [1 + 4 + 1] + for duplex in appends: + encoded = runtime.take_prefetched_vision(duplex) + assert encoded is not None + assert len(encoded) == len(duplex["payload"]["video_frames"]) + _assert_same_blocks(encoded, expected[duplex["session_id"]]) + assert runtime.take_prefetched_vision(duplex) is None # consumed once + + +def test_prefetch_chunk_size_follows_vision_batch_size_config() -> None: + runtime = _runtime() + assert runtime._vision_prefetch_chunk_size() == 16 # from _Thinker.config.vision_batch_size + + runtime.thinker.config.vision_batch_size = 3 + assert runtime._vision_prefetch_chunk_size() == 3 + + # Missing/invalid config falls back to 16. + runtime.thinker.config.vision_batch_size = None + assert runtime._vision_prefetch_chunk_size() == 16 + del runtime.thinker.config.vision_batch_size + assert runtime._vision_prefetch_chunk_size() == 16 + + +def test_prefetch_pipelines_chunks_without_changing_results() -> None: + """Small ``vision_batch_size`` forces multiple pipelined chunks; results must match single-chunk encoding.""" + runtime = _runtime() + runtime.thinker.config.vision_batch_size = 2 # 5 one-slice jobs -> chunks of [2, 2, 1] + appends = [_append(chr(ord("a") + i), i, [_jpeg(640, 480, 100 + i)]) for i in range(5)] + expected = { + duplex["session_id"]: _per_frame_reference(runtime, runtime._decode_video_frames_payload(duplex["payload"])) + for duplex in appends + } + runtime.thinker.encoder_calls.clear() + + with torch.inference_mode(): + runtime.prefetch_vision(appends) + + assert runtime.thinker.encoder_calls == [2, 2, 1] # one encode call per pipelined chunk + for duplex in appends: + encoded = runtime.take_prefetched_vision(duplex) + assert encoded is not None + _assert_same_blocks(encoded, expected[duplex["session_id"]]) + + +def test_prefetch_single_job_skips_pool_and_still_encodes() -> None: + runtime = _runtime() + appends = [_append("solo", 1, [_jpeg(640, 480, 200)])] + expected = _per_frame_reference(runtime, runtime._decode_video_frames_payload(appends[0]["payload"])) + runtime.thinker.encoder_calls.clear() + + with torch.inference_mode(): + runtime.prefetch_vision(appends) + + assert runtime.thinker.encoder_calls == [1] + encoded = runtime.take_prefetched_vision(appends[0]) + assert encoded is not None + _assert_same_blocks(encoded, expected) + + +def test_prefetch_skips_prepared_appends_and_rejects_changed_payloads() -> None: + runtime = _runtime() + prepared = _MiniCPMO45Stage0SessionState(session_id="a") + prepared.prepared_append_identity = (0, 3) + prepared.prepared_inputs_embeds = torch.zeros(1, 64) + runtime.sessions["a"] = prepared + appends = [_append("a", 3, [_jpeg(640, 480, 20)]), _append("b", 4, [_jpeg(640, 480, 21)])] + + with torch.inference_mode(): + runtime.prefetch_vision(appends) + + assert runtime.thinker.encoder_calls == [1] # session a replays its cached append + assert runtime.take_prefetched_vision(appends[0]) is None + changed = _append("b", 4, [_jpeg(640, 480, 22)]) + assert runtime.take_prefetched_vision(changed) is None + + +def test_prefetch_leaves_bad_frames_to_the_per_request_path() -> None: + runtime = _runtime() + good = _append("a", 1, [_jpeg(640, 480, 30)]) + bad = _append("b", 1, ["bm90IGEganBlZw=="]) # valid base64, not an image + + with torch.inference_mode(): + runtime.prefetch_vision([good, bad]) + + assert runtime.take_prefetched_vision(bad) is None + assert runtime.take_prefetched_vision(good) is not None + with pytest.raises(ValueError, match="invalid omni duplex video frame payload"): + runtime._decode_video_frames_payload(bad["payload"]) + + +def test_prefetch_drops_unconsumed_entries_next_step() -> None: + runtime = _runtime() + duplex = _append("a", 1, [_jpeg(640, 480, 40)]) + with torch.inference_mode(): + runtime.prefetch_vision([duplex]) + runtime.prefetch_vision([]) + assert runtime.take_prefetched_vision(duplex) is None + + +def test_omni_preprocess_batch_prefetches_only_for_multiple_frame_appends() -> None: + calls: list[list[str]] = [] + helper = SimpleNamespace( + prefetch_vision=lambda appends: calls.append([a["session_id"] for a in appends]), + batches_audio_encoder=lambda: False, + ) + model = SimpleNamespace( + model_stage="llm", + _minicpmo45_duplex_data_plane_helper=helper, + _duplex_data_plane_helper=lambda: helper, + ) + buffer = { + "r0": {"duplex": _append("a", 1, ["x"])}, + "r1": {"duplex": _append("b", 1, [])}, # audio-only append + "r2": {"duplex": _append("c", 1, ["y"])}, + "r3": {"prompt_token_ids": [1, 2]}, # ordinary request + } + + MiniCPMO45OmniForConditionalGeneration.preprocess_batch( + model, req_ids=["r0", "r1", "r2", "r3"], model_intermediate_buffer=buffer, device=torch.device("cpu") + ) + MiniCPMO45OmniForConditionalGeneration.preprocess_batch( + model, req_ids=["r0", "r3"], model_intermediate_buffer=buffer, device=torch.device("cpu") + ) + model.model_stage = "tts" + MiniCPMO45OmniForConditionalGeneration.preprocess_batch( + model, req_ids=["r0", "r2"], model_intermediate_buffer=buffer, device=torch.device("cpu") + ) + + assert calls == [["a", "c"], []] diff --git a/tests/model_executor/models/minicpmo_4_5/test_fused_residual_layer_norm.py b/tests/model_executor/models/minicpmo_4_5/test_fused_residual_layer_norm.py new file mode 100644 index 00000000000..5654a0c0982 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_fused_residual_layer_norm.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for the fused gated residual + LayerNorm operators. + +``residual_layer_norm`` and ``qkv_head_layer_norm`` are checked against the +eager PyTorch chains they replace, on the native path (any device) and on the +Triton path (CUDA). +""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm_omni.model_executor.models.minicpmo_4_5.ops import qkv_head_layer_norm, residual_layer_norm + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +def _conv_taps(history: torch.Tensor, conv: nn.Conv1d) -> torch.Tensor: + """``history @ [W_0; ...; W_{K-1}]^T``, the tap-stacked GEMM ``residual_layer_norm`` reduces.""" + weight = conv.weight.permute(2, 0, 1).reshape(-1, conv.in_channels) + return F.linear(history, weight) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("channels", [64, 512, 300]) +def test_gated_residual_modulated_layer_norm(device: str, channels: int) -> None: + torch.manual_seed(0) + x = torch.randn(3, 7, channels, device=device) + y = torch.randn(3, 7, channels, device=device) + gate, scale, shift = (torch.randn(channels, device=device) for _ in range(3)) + + expected_residual = x + gate * y + expected = F.layer_norm(expected_residual, (channels,), eps=1e-6) * (1 + scale) + shift + + residual = x.clone() + out = residual_layer_norm(residual, y, gate=gate, weight=1 + scale, bias=shift, eps=1e-6, residual_out=residual) + torch.testing.assert_close(residual, expected_residual, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5) + + plain = residual_layer_norm(x, weight=1 + scale, bias=shift, eps=1e-6) + torch.testing.assert_close( + plain, F.layer_norm(x, (channels,), eps=1e-6) * (1 + scale) + shift, rtol=1e-5, atol=1e-5 + ) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_causal_conv_taps_layer_norm_mish_into_strided_history(device: str) -> None: + """Conv tap sum + bias + LayerNorm + Mish, stored behind a history's cached frames.""" + torch.manual_seed(1) + channels, frames, width = 32, 9, 2 + conv = nn.Conv1d(channels, channels, width + 1).to(device) + norm = nn.LayerNorm(channels).to(device) + with torch.no_grad(): + norm.weight.normal_() + norm.bias.normal_() + history = torch.randn(4, width + frames, channels, device=device) + with torch.no_grad(): + expected = F.mish(norm(conv(history.transpose(1, 2)).transpose(1, 2))) + destination = torch.full((4, width + frames, channels), float("nan"), device=device) + residual_layer_norm( + None, + _conv_taps(history, conv), + taps=width + 1, + y_bias=conv.bias, + weight=norm.weight, + bias=norm.bias, + eps=norm.eps, + activation="mish", + out=destination[:, width:], + ) + torch.testing.assert_close(destination[:, width:], expected, rtol=1e-5, atol=1e-5) + assert torch.isnan(destination[:, :width]).all() + + +@pytest.mark.parametrize("device", _DEVICES) +def test_causal_conv_taps_gated_residual(device: str) -> None: + """``x += gate * conv(history)`` followed by the modulated LayerNorm of the result.""" + torch.manual_seed(2) + channels, frames, width = 48, 5, 2 + conv = nn.Conv1d(channels, channels, width + 1).to(device) + history = torch.randn(2, width + frames, channels, device=device) + x = torch.randn(2, frames, channels, device=device) + gate, scale, shift = (torch.randn(channels, device=device) for _ in range(3)) + with torch.no_grad(): + expected_x = x + gate * conv(history.transpose(1, 2)).transpose(1, 2) + expected = F.layer_norm(expected_x, (channels,), eps=1e-6) * (1 + scale) + shift + residual = x.clone() + out = residual_layer_norm( + residual, + _conv_taps(history, conv), + taps=width + 1, + y_bias=conv.bias, + gate=gate, + weight=1 + scale, + bias=shift, + eps=1e-6, + residual_out=residual, + ) + torch.testing.assert_close(residual, expected_x, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_qkv_head_layer_norm_writes_q_and_interleaved_kv(device: str) -> None: + torch.manual_seed(4) + batch, frames, heads, head_dim, cached = 3, 6, 4, 16, 5 + q_norm, k_norm = nn.LayerNorm(head_dim).to(device), nn.LayerNorm(head_dim, eps=1e-6).to(device) + with torch.no_grad(): + for norm in (q_norm, k_norm): + norm.weight.normal_() + norm.bias.normal_() + qkv = torch.randn(batch, frames, 3 * heads * head_dim, device=device) + kv = torch.full((batch, heads, frames + cached, 2 * head_dim), float("nan"), device=device) + with torch.no_grad(): + q = qkv_head_layer_norm(qkv, kv, num_heads=heads, head_dim=head_dim, q_norm=q_norm, k_norm=k_norm) + split = qkv.view(batch, frames, 3, heads, head_dim).transpose(1, 3) # (b, h, 3, t, d) + torch.testing.assert_close(q, q_norm(split[:, :, 0]), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(kv[:, :, :frames, :head_dim], k_norm(split[:, :, 1]), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(kv[:, :, :frames, head_dim:], split[:, :, 2]) + # The cached frames behind the current chunk are never written. + assert torch.isnan(kv[:, :, frames:]).all() diff --git a/tests/model_executor/models/minicpmo_4_5/test_vision_fused.py b/tests/model_executor/models/minicpmo_4_5/test_vision_fused.py new file mode 100644 index 00000000000..9843bbe997d --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_vision_fused.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused packed SigLIP layers and vision CUDA graphs match the eager packed encode.""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5 import vision_fused +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + Resampler, + SiglipVisionConfig, + SiglipVisionTransformer, + _encode_vision_packed, +) + +_PATCH = 2 + +pytestmark = [pytest.mark.core_model] + + +def _towers(*, dtype: torch.dtype = torch.float32, device: str = "cpu") -> tuple[SiglipVisionTransformer, Resampler]: + torch.manual_seed(0) + config = SiglipVisionConfig( + hidden_size=32, + intermediate_size=72, + num_hidden_layers=3, + num_attention_heads=4, + image_size=28, + patch_size=_PATCH, + attention_dropout=0.0, + ) + config._attn_implementation = "sdpa" + vpm = SiglipVisionTransformer(config).eval() + for module in vpm.modules(): + if isinstance(module, torch.nn.LayerNorm): + # Non-trivial affines, so a LayerNorm fused into the wrong layer shows. + torch.nn.init.normal_(module.weight, mean=1.0, std=0.2) + torch.nn.init.normal_(module.bias, std=0.2) + resampler = Resampler(num_queries=4, embed_dim=64, num_heads=4, kv_dim=32, adaptive=True, max_size=(8, 8)).eval() + for param in resampler.parameters(): + torch.nn.init.normal_(param, std=0.2) + return vpm.to(device=device, dtype=dtype), resampler.to(device=device, dtype=dtype) + + +def _embeddings(vpm: SiglipVisionTransformer, tokens: int, *, seed: int = 1) -> torch.Tensor: + weight = vpm.embeddings.patch_embedding.weight + generator = torch.Generator().manual_seed(seed) + return torch.randn((tokens, 32), generator=generator).to(device=weight.device, dtype=weight.dtype) + + +def _eager(vpm: SiglipVisionTransformer, hidden: torch.Tensor, seq_groups) -> torch.Tensor: + for layer in vpm.encoder.layers: + hidden = layer.forward_packed(hidden, seq_groups) + return vpm.post_layernorm(hidden) + + +@pytest.mark.cpu +def test_fused_layers_match_eager_packed_layers() -> None: + vpm, _ = _towers() + # Two runs of different sequence length, one with several items. + seq_groups = [(0, 3, 16), (48, 2, 15)] + hidden = _embeddings(vpm, 78) + with torch.inference_mode(): + expected = _eager(vpm, hidden.clone(), seq_groups) + fused = vision_fused.encode_packed_fused(vpm, hidden.clone(), seq_groups) + torch.testing.assert_close(fused, expected, rtol=1e-5, atol=1e-5) + + +@pytest.mark.cpu +def test_explicit_eager_attention_keeps_the_unfused_layers() -> None: + vpm, _ = _towers() + assert vision_fused.supports_fused_layers(vpm) + vpm.config._attn_implementation = "eager" + assert not vision_fused.supports_fused_layers(vpm) + + +@pytest.mark.cpu +def test_packed_qkv_repoints_the_projections_without_changing_them() -> None: + vpm, _ = _towers() + attn = vpm.encoder.layers[0].self_attn + before = [p.weight.detach().clone() for p in (attn.q_proj, attn.k_proj, attn.v_proj)] + weight, bias = vision_fused.packed_qkv(attn) + assert weight.shape == (96, 32) and bias is not None + for index, projection in enumerate((attn.q_proj, attn.k_proj, attn.v_proj)): + assert projection.weight.data_ptr() == weight[index * 32].data_ptr() + torch.testing.assert_close(projection.weight, before[index], rtol=0, atol=0) + assert vision_fused.packed_qkv(attn)[0] is weight + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_fused_layers_on_cuda_stay_as_close_to_fp32_as_eager(dtype: torch.dtype) -> None: + reference, _ = _towers(device="cuda") + vpm, _ = _towers(dtype=dtype, device="cuda") + seq_groups = [(0, 4, 16)] + hidden = _embeddings(reference, 64) + with torch.inference_mode(): + expected = _eager(reference, hidden.clone(), seq_groups) + eager = _eager(vpm, hidden.to(dtype), seq_groups).float() + fused = vision_fused.encode_packed_fused(vpm, hidden.to(dtype), seq_groups).float() + eager_error = (eager - expected).abs().max().item() + fused_error = (fused - expected).abs().max().item() + assert fused_error <= 1.5 * eager_error + 1e-5, (fused_error, eager_error) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("fused", [False, True]) +def test_vision_graph_matches_eager_and_ignores_batch_padding(monkeypatch, fused: bool) -> None: + monkeypatch.setattr(vision_fused, "_CAPTURE_AFTER", 0) + vpm, resampler = _towers(device="cuda") + vpm.fused_layers = fused + height, width = 4, 4 + generator = torch.Generator().manual_seed(2) + pixels = [torch.randn((3, _PATCH, height * width * _PATCH), generator=generator).cuda() for _ in range(3)] + layouts = [(height, width)] * 3 + encoder = vision_fused.VisionGraphEncoder(vpm, resampler, batch_sizes=(4,)) + with torch.inference_mode(): + expected = _encode_vision_packed(vpm, resampler, pixels, layouts, 16) + # Three slices replay the batch-4 graph with one zero slice of padding. + replayed = _encode_vision_packed(vpm, resampler, pixels, layouts, 16, encoder) + # A second replay with other pixels in the same graph. + again = _encode_vision_packed(vpm, resampler, pixels[:1], layouts[:1], 16, encoder) + assert len(encoder._graphs) == 1 + torch.testing.assert_close(replayed, expected, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(again, expected[:1], rtol=1e-4, atol=1e-4) diff --git a/tests/model_executor/models/minicpmo_4_5/test_vision_packed_encoder.py b/tests/model_executor/models/minicpmo_4_5/test_vision_packed_encoder.py new file mode 100644 index 00000000000..f842fc4132d --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_vision_packed_encoder.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Parity of the unpadded, grid-grouped MiniCPM-o 4.5 vision path with the padded batch.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.minicpmo_4_5 import minicpmo_4_5_omni_llm as minicpmo_llm +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, + Resampler, + SiglipAttention, + SiglipEncoderLayer, + SiglipSdpaAttention, + SiglipVisionConfig, + SiglipVisionTransformer, + _plan_packed_vision_chunks, + _resolve_vision_attention_implementation, +) + +_PATCH = 2 +# Mixed grids: equal sequence length with different shapes, and repeats. +_RAGGED_LAYOUTS = [(4, 4), (2, 8), (3, 5), (4, 4), (5, 3), (1, 16), (4, 4)] + + +@dataclass +class _Config: + vision_batch_size: int + + +def _vision_config(attn_implementation: str, *, hidden_size: int = 32, heads: int = 4) -> SiglipVisionConfig: + config = SiglipVisionConfig( + hidden_size=hidden_size, + intermediate_size=2 * hidden_size, + num_hidden_layers=2, + num_attention_heads=heads, + image_size=28, + patch_size=_PATCH, + attention_dropout=0.0, + ) + config._attn_implementation = attn_implementation + return config + + +def _build_towers( + attn_implementation: str, *, seed: int = 0, dtype: torch.dtype = torch.float32, device: str = "cpu" +) -> tuple[SiglipVisionTransformer, Resampler]: + torch.manual_seed(seed) + vpm = SiglipVisionTransformer(_vision_config(attn_implementation)).eval() + resampler = Resampler(num_queries=4, embed_dim=64, num_heads=4, kv_dim=32, adaptive=True, max_size=(8, 8)).eval() + for param in resampler.parameters(): + # Resampler parameters default to zeros/constant: randomize for a meaningful parity check. + torch.nn.init.normal_(param, std=0.2) + resampler.attn_implementation = attn_implementation + return vpm.to(device=device, dtype=dtype), resampler.to(device=device, dtype=dtype) + + +def _copy_towers( + source: tuple[SiglipVisionTransformer, Resampler], attn_implementation: str +) -> tuple[SiglipVisionTransformer, Resampler]: + weight = source[0].embeddings.patch_embedding.weight + vpm, resampler = _build_towers(attn_implementation, dtype=weight.dtype, device=str(weight.device)) + vpm.load_state_dict(source[0].state_dict()) + resampler.load_state_dict(source[1].state_dict()) + return vpm, resampler + + +def _model(towers: tuple[SiglipVisionTransformer, Resampler], *, packed: bool, vision_batch_size: int = 16): + return SimpleNamespace( + vpm=towers[0], + resampler=towers[1], + config=_Config(vision_batch_size=vision_batch_size), + vision_packed_encode=packed, + ) + + +def _pixels(layouts, *, seed: int = 1, dtype: torch.dtype = torch.float32, device: str = "cpu") -> list[torch.Tensor]: + generator = torch.Generator().manual_seed(seed) + return [ + torch.randn((3, _PATCH, h * w * _PATCH), generator=generator).to(device=device, dtype=dtype) for h, w in layouts + ] + + +def _encode(model, pixel_values, layouts) -> torch.Tensor: + tgt_sizes = torch.tensor(layouts, dtype=torch.int32) + with torch.inference_mode(): + return MiniCPMO45OmniLLMForConditionalGeneration.get_vision_hidden_states( + model, + {"pixel_values": pixel_values, "tgt_sizes": tgt_sizes}, + ) + + +def _per_item(model, pixel_values, layouts) -> torch.Tensor: + return torch.cat([_encode(model, [pixels], [layout]) for pixels, layout in zip(pixel_values, layouts)]) + + +pytestmark = [pytest.mark.core_model] + + +@pytest.mark.cpu +def test_vision_attention_defaults_to_sdpa() -> None: + assert _resolve_vision_attention_implementation(SimpleNamespace(), SimpleNamespace()) == "sdpa" + unset = SimpleNamespace(_attn_implementation=None) + assert _resolve_vision_attention_implementation(unset, SimpleNamespace()) == "sdpa" + + +@pytest.mark.cpu +def test_explicit_vision_attention_override_is_preserved(monkeypatch) -> None: + vision_config = SimpleNamespace(_attn_implementation="eager") + assert _resolve_vision_attention_implementation(vision_config, SimpleNamespace()) == "eager" + + flash_config = SimpleNamespace(_attn_implementation="flash_attention_2") + monkeypatch.setattr(minicpmo_llm, "is_flash_attn_2_available", lambda: True) + assert _resolve_vision_attention_implementation(flash_config, SimpleNamespace()) == "flash_attention_2" + monkeypatch.setattr(minicpmo_llm, "is_flash_attn_2_available", lambda: False) + assert _resolve_vision_attention_implementation(flash_config, SimpleNamespace()) == "sdpa" + + +@pytest.mark.cpu +def test_legacy_engine_attention_override_wins() -> None: + vision_config = SimpleNamespace(_attn_implementation="sdpa") + model_config = SimpleNamespace(_attn_implementation="eager") + assert _resolve_vision_attention_implementation(vision_config, model_config) == "eager" + + +@pytest.mark.cpu +def test_sdpa_encoder_builds_and_selects_sdpa_attention() -> None: + vpm = SiglipVisionTransformer(_vision_config("sdpa")) + assert all(type(layer.self_attn) is SiglipSdpaAttention for layer in vpm.encoder.layers) + eager = SiglipVisionTransformer(_vision_config("eager")) + assert all(type(layer.self_attn) is SiglipAttention for layer in eager.encoder.layers) + + +@pytest.mark.cpu +def test_sdpa_layer_matches_eager_layer_with_padding_mask() -> None: + torch.manual_seed(7) + eager = SiglipEncoderLayer(_vision_config("eager")).eval() + sdpa = SiglipEncoderLayer(_vision_config("sdpa")).eval() + sdpa.load_state_dict(eager.state_dict()) + hidden_states = torch.randn(2, 5, 32) + attention_mask = torch.zeros(2, 1, 5, 5) + attention_mask[1, :, :, -2:] = torch.finfo(torch.float32).min + + with torch.inference_mode(): + eager_output = eager(hidden_states, attention_mask)[0] + sdpa_output = sdpa(hidden_states, attention_mask)[0] + + torch.testing.assert_close(sdpa_output, eager_output, rtol=1e-5, atol=1e-6) + + +@pytest.mark.cpu +def test_plan_groups_equal_grids_and_bounds_chunks() -> None: + chunks = _plan_packed_vision_chunks(_RAGGED_LAYOUTS, max_items=16) + assert len(chunks) == 1 + indices, runs = chunks[0] + assert sorted(indices) == list(range(len(_RAGGED_LAYOUTS))) + # (5, 3) and (3, 5) share L=15 and are adjacent; the three (4, 4) grids form one run. + assert runs == [(3, 5, 1), (5, 3, 1), (1, 16, 1), (2, 8, 1), (4, 4, 3)] + + chunks = _plan_packed_vision_chunks(_RAGGED_LAYOUTS, max_items=3) + assert [len(indices) for indices, _ in chunks] == [3, 3, 1] + assert sorted(i for indices, _ in chunks for i in indices) == list(range(len(_RAGGED_LAYOUTS))) + for indices, runs in chunks: + assert sum(count for _, _, count in runs) == len(indices) + + +@pytest.mark.cpu +@pytest.mark.parametrize("vision_batch_size", [1, 3, 16]) +@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"]) +def test_packed_matches_padded_and_per_item(vision_batch_size: int, attn_implementation: str) -> None: + reference = _build_towers("eager") + towers = _copy_towers(reference, attn_implementation) + pixel_values = _pixels(_RAGGED_LAYOUTS) + + padded = _encode( + _model(reference, packed=False, vision_batch_size=vision_batch_size), pixel_values, _RAGGED_LAYOUTS + ) + per_item = _per_item(_model(reference, packed=False), pixel_values, _RAGGED_LAYOUTS) + packed = _encode(_model(towers, packed=True, vision_batch_size=vision_batch_size), pixel_values, _RAGGED_LAYOUTS) + + assert packed.shape == (len(_RAGGED_LAYOUTS), 4, 64) + torch.testing.assert_close(packed, per_item, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(packed, padded, rtol=1e-5, atol=1e-5) + + +@pytest.mark.cpu +def test_packed_output_does_not_depend_on_batch_composition() -> None: + towers = _build_towers("sdpa") + model = _model(towers, packed=True) + pixel_values = _pixels(_RAGGED_LAYOUTS) + + alone = _encode(model, pixel_values[:1], _RAGGED_LAYOUTS[:1]) + mixed = _encode(model, pixel_values, _RAGGED_LAYOUTS) + + torch.testing.assert_close(mixed[:1], alone, rtol=1e-6, atol=1e-6) + + +@pytest.mark.cpu +def test_packed_grows_resampler_position_cache_like_padded_path() -> None: + reference = _build_towers("eager") + towers = _copy_towers(reference, "sdpa") + layouts = [(3, 10), (2, 2)] # w=10 exceeds the resampler's 8x8 sin-cos cache + pixel_values = _pixels(layouts) + + packed = _encode(_model(towers, packed=True), pixel_values, layouts) + padded = _encode(_model(reference, packed=False), pixel_values, layouts) + + assert list(towers[1].max_size) == [8, 10] + torch.testing.assert_close(packed, padded, rtol=1e-5, atol=1e-5) + + +@pytest.mark.cpu +def test_packed_falls_back_to_padded_path_on_inconsistent_grids(mocker) -> None: + towers = _build_towers("sdpa") + pixel_values = _pixels([(2, 2)]) + packed_spy = mocker.spy(minicpmo_llm, "_encode_vision_packed") + + # The pixels hold 4 patches, the grid claims 6: the packed path defers to + # the padded path, which keeps its own (failing) behavior for such input. + with torch.inference_mode(), pytest.raises(RuntimeError): + MiniCPMO45OmniLLMForConditionalGeneration.get_vision_hidden_states( + _model(towers, packed=True), + {"pixel_values": pixel_values, "tgt_sizes": torch.tensor([[2, 3]])}, + ) + + packed_spy.assert_not_called() + + +@pytest.mark.cpu +def test_packed_position_ids_match_padded_embeddings() -> None: + vpm, _ = _build_towers("sdpa") + embeddings = vpm.embeddings + for height, width in _RAGGED_LAYOUTS: + mask = torch.ones((1, height, width), dtype=torch.bool) + expected = embeddings._create_position_ids(mask, torch.tensor([[height, width]]), device=torch.device("cpu")) + actual = embeddings.layout_position_ids(height, width, torch.device("cpu")) + assert torch.equal(actual, expected[0]) + assert embeddings.layout_position_ids(height, width, torch.device("cpu")) is actual + + +def _cuda_or_skip() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + + +@hardware_test(res={"cuda": "L4"}) +def test_packed_bf16_error_is_no_worse_than_padded_eager_bf16() -> None: + """bf16 SDPA packed path vs the shipped eager padded path, both measured against fp32.""" + _cuda_or_skip() + reference = _build_towers("eager", device="cuda") + layouts = [(28, 37), (32, 32), (28, 37), (37, 28), (16, 64)] + pixel_values = _pixels(layouts, device="cuda") + fp32 = _encode(_model(reference, packed=False), pixel_values, layouts) + + eager_bf16 = _copy_towers(reference, "eager") + eager_bf16 = (eager_bf16[0].to(torch.bfloat16), eager_bf16[1].to(torch.bfloat16)) + sdpa_bf16 = _copy_towers(reference, "sdpa") + sdpa_bf16 = (sdpa_bf16[0].to(torch.bfloat16), sdpa_bf16[1].to(torch.bfloat16)) + pixel_values_bf16 = [pixels.to(torch.bfloat16) for pixels in pixel_values] + + baseline = _encode(_model(eager_bf16, packed=False), pixel_values_bf16, layouts).float() + packed = _encode(_model(sdpa_bf16, packed=True), pixel_values_bf16, layouts).float() + + baseline_error = (baseline - fp32).abs().max().item() + packed_error = (packed - fp32).abs().max().item() + assert packed_error <= 1.5 * baseline_error + 1e-3, (packed_error, baseline_error) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py b/vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py new file mode 100644 index 00000000000..6ba466ae5c8 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from .fused_qkv_layer_norm import qkv_head_layer_norm +from .fused_residual_layer_norm import residual_layer_norm + +__all__ = ["qkv_head_layer_norm", "residual_layer_norm"] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py b/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py new file mode 100644 index 00000000000..aafc18a3c35 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# ruff: noqa: N803 + +"""Fused attention prologue for a packed QKV projection with LayerNorm q/k norms. + +``omnivoice/fused_qkv_rope.py`` fuses the Qwen3-style prologue (split the +packed QKV, RMSNorm q and k per head, RoPE, transpose to the attention +layout). DiT attention blocks with ``qk_norm=True`` and ``nn.LayerNorm`` norms +and no RoPE run the same steps with a LayerNorm, and a streaming block also +writes the new keys and values into its attention cache. Eagerly that is a +transpose copy and a LayerNorm for each of q and k (``F.layer_norm`` needs a +contiguous input) plus a copy of each of k and v into the cache. + +``qkv_head_layer_norm`` does it in one pass: each program reads one token's +packed ``[q | k | v]`` row, normalizes the q and k heads, and stores q in the +``(N, H, T, D)`` attention layout and k and v straight into a key/value +buffer whose last axis is ``[k | v]`` (``2 * D`` wide, the layout of the CFM +estimator's attention cache). Statistics are accumulated in fp32. Other +devices run the native PyTorch ops. + +The key/value buffer may also be a pool of resident caches: ``rows`` maps +each of the ``N`` rows to its pool row and ``positions`` each token to the +frame it is written to (negative: not written, e.g. query padding). +""" + +import torch +from vllm.triton_utils import tl, triton + +from .fused_residual_layer_norm import _use_triton + + +@triton.jit +def _head_layer_norm(x, mask, dims, weight_ptr, bias_ptr, eps, HEAD_DIM: tl.constexpr, HAS_AFFINE: tl.constexpr): + """LayerNorm over the last axis of ``(heads, dims)`` in fp32, with an optional per-dim affine.""" + x = x.to(tl.float32) + mean = tl.sum(x, axis=1) / HEAD_DIM + centered = tl.where(mask, x - mean[:, None], 0.0) + variance = tl.sum(centered * centered, axis=1) / HEAD_DIM + x = centered * tl.rsqrt(variance + eps)[:, None] + if HAS_AFFINE: + x = x * tl.load(weight_ptr + dims, mask=dims < HEAD_DIM, other=1.0).to(tl.float32)[None, :] + x = x + tl.load(bias_ptr + dims, mask=dims < HEAD_DIM, other=0.0).to(tl.float32)[None, :] + return x + + +@triton.jit +def _qkv_head_layer_norm_kernel( + q_out_ptr, + kv_ptr, + qkv_ptr, + q_weight_ptr, + q_bias_ptr, + k_weight_ptr, + k_bias_ptr, + rows_ptr, + positions_ptr, + frames, + stride_qkv_n, + stride_qkv_t, + stride_qn, + stride_qh, + stride_qt, + stride_kvn, + stride_kvh, + stride_kvt, + q_eps, + k_eps, + HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + HAS_Q_AFFINE: tl.constexpr, + HAS_K_AFFINE: tl.constexpr, + HAS_ROWS: tl.constexpr, + HAS_POSITIONS: tl.constexpr, +): + row = tl.program_id(0) + n = row // frames + t = row % frames + heads = tl.arange(0, BLOCK_H) + dims = tl.arange(0, BLOCK_D) + mask = (heads[:, None] < HEADS) & (dims[None, :] < HEAD_DIM) + inner = HEADS * HEAD_DIM + src = qkv_ptr + n * stride_qkv_n + t * stride_qkv_t + heads[:, None] * HEAD_DIM + dims[None, :] + + q = tl.load(src, mask=mask, other=0.0) + q = _head_layer_norm(q, mask, dims, q_weight_ptr, q_bias_ptr, q_eps, HEAD_DIM, HAS_Q_AFFINE) + tl.store(q_out_ptr + n * stride_qn + heads[:, None] * stride_qh + t * stride_qt + dims[None, :], q, mask=mask) + + k = tl.load(src + inner, mask=mask, other=0.0) + k = _head_layer_norm(k, mask, dims, k_weight_ptr, k_bias_ptr, k_eps, HEAD_DIM, HAS_K_AFFINE) + kv_row = n + if HAS_ROWS: + kv_row = tl.load(rows_ptr + n).to(tl.int64) + frame = t + if HAS_POSITIONS: + frame = tl.load(positions_ptr + row).to(tl.int64) + mask = mask & (frame >= 0) + dst = kv_ptr + kv_row * stride_kvn + heads[:, None] * stride_kvh + frame * stride_kvt + dims[None, :] + tl.store(dst, k, mask=mask) + + v = tl.load(src + 2 * inner, mask=mask, other=0.0) + tl.store(dst + HEAD_DIM, v, mask=mask) + + +def qkv_head_layer_norm( + qkv: torch.Tensor, + kv: torch.Tensor, + *, + num_heads: int, + head_dim: int, + q_norm: torch.nn.LayerNorm, + k_norm: torch.nn.LayerNorm, + q_out: torch.Tensor | None = None, + rows: torch.Tensor | None = None, + positions: torch.Tensor | None = None, +) -> torch.Tensor: + """Split ``qkv`` (N, T, 3*H*D), LayerNorm q/k per head, store q and ``kv``. + + ``kv`` is an ``(N, H, >= T, 2 * D)`` buffer (any strides with a + contiguous last axis); frame ``t`` of it receives ``[k_t | v_t]``. + With ``rows`` (``(N,)`` integers) row ``n`` goes to ``kv[rows[n]]`` of an + ``(R, H, F, 2 * D)`` pool, and with ``positions`` (``(N, T)`` integers) + token ``t`` to frame ``positions[n, t]``, or nowhere when that is + negative. Returns q as ``(N, H, T, D)`` (``q_out`` when given). + """ + if not isinstance(q_norm, torch.nn.LayerNorm) or not isinstance(k_norm, torch.nn.LayerNorm): + raise TypeError("qkv_head_layer_norm expects nn.LayerNorm q/k norms") + batch, frames, width = (int(dim) for dim in qkv.shape) + inner = num_heads * head_dim + if width != 3 * inner or qkv.stride(2) != 1: + raise ValueError(f"qkv_head_layer_norm: qkv must be (N, T, 3*{inner}) with contiguous channels") + if ( + kv.dim() != 4 + or (rows is None and int(kv.shape[0]) != batch) + or int(kv.shape[1]) != num_heads + or (positions is None and int(kv.shape[2]) < frames) + or int(kv.shape[3]) != 2 * head_dim + or kv.stride(3) != 1 + ): + raise ValueError(f"qkv_head_layer_norm: kv must be (N, H, >= T, 2*D), got {tuple(kv.shape)}") + if rows is not None and tuple(rows.shape) != (batch,): + raise ValueError(f"qkv_head_layer_norm: rows must be ({batch},), got {tuple(rows.shape)}") + if positions is not None and (tuple(positions.shape) != (batch, frames) or not positions.is_contiguous()): + raise ValueError(f"qkv_head_layer_norm: positions must be contiguous ({batch}, {frames})") + if q_out is None: + q_out = torch.empty((batch, num_heads, frames, head_dim), device=qkv.device, dtype=qkv.dtype) + + if not _use_triton(qkv): + heads = qkv.view(batch, frames, 3, num_heads, head_dim) + q_out.copy_(q_norm(heads[:, :, 0].transpose(1, 2))) + keys = k_norm(heads[:, :, 1].transpose(1, 2)) + values = heads[:, :, 2].transpose(1, 2) + if rows is None and positions is None: + kv[:, :, :frames, :head_dim].copy_(keys) + kv[:, :, :frames, head_dim:].copy_(values) + return q_out + pool_rows = rows.tolist() if rows is not None else range(batch) + for n, pool_row in enumerate(pool_rows): + frame_index = positions[n] if positions is not None else torch.arange(frames, device=kv.device) + written = frame_index >= 0 + destination = frame_index[written].to(kv.device) + kv[pool_row, :, destination, :head_dim] = keys[n][:, written] + kv[pool_row, :, destination, head_dim:] = values[n][:, written] + return q_out + if batch * frames == 0: + return q_out + q_weight, q_bias, q_eps = q_norm.weight, q_norm.bias, float(q_norm.eps) + k_weight, k_bias, k_eps = k_norm.weight, k_norm.bias, float(k_norm.eps) + if (q_weight is None) != (q_bias is None) or (k_weight is None) != (k_bias is None): + raise ValueError("qkv_head_layer_norm: LayerNorm affine needs both weight and bias") + _qkv_head_layer_norm_kernel[(batch * frames,)]( + q_out, + kv, + qkv, + q_weight if q_weight is not None else qkv, + q_bias if q_bias is not None else qkv, + k_weight if k_weight is not None else qkv, + k_bias if k_bias is not None else qkv, + rows if rows is not None else qkv, + positions if positions is not None else qkv, + frames, + qkv.stride(0), + qkv.stride(1), + q_out.stride(0), + q_out.stride(1), + q_out.stride(2), + kv.stride(0), + kv.stride(1), + kv.stride(2), + q_eps, + k_eps, + HEADS=num_heads, + HEAD_DIM=head_dim, + BLOCK_H=triton.next_power_of_2(num_heads), + BLOCK_D=triton.next_power_of_2(head_dim), + HAS_Q_AFFINE=q_weight is not None, + HAS_K_AFFINE=k_weight is not None, + HAS_ROWS=rows is not None, + HAS_POSITIONS=positions is not None, + num_warps=4, + ) + return q_out diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py b/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py new file mode 100644 index 00000000000..5194af6d904 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# ruff: noqa: N803 + +"""Fused gated residual + LayerNorm (+ adaLN modulation, + Mish) operator. + +adaLN DiT blocks run the same chain three times per block:: + + x = x + gate * y # gated residual + h = layer_norm(x) * (1 + scale) + shift # the next sub-block's input + +In eager PyTorch that is five full passes over the activation (mul, add, +layer_norm, mul, add) plus a ``1 + scale`` op. ``residual_layer_norm`` does +the whole chain in one pass per row, the same shape of kernel as Qwen-Image's +``_residual_layernorm_select01_kernel`` (``diffusion/layers/ +qwen_select01_modulation.py``) minus its two-branch select. The modulation is +the LayerNorm's affine: pass ``weight = 1 + scale`` and ``bias = shift``. + +``y`` may also be the output of a causal convolution computed as one GEMM +over its taps (``taps > 1``): for a ``K``-tap convolution over a history of +``K - 1 + T`` frames, ``Y = history @ [W_0; ...; W_{K-1}]^T`` has ``K * C`` +columns and the convolution output is ``sum_k Y[t + k, k*C:(k+1)*C] + y_bias``, +which the kernel reduces on the fly. That lets a convolution, its bias, the +gated residual and the next LayerNorm share one pass, and lets a convolution +followed by ``LayerNorm + Mish`` (no residual) finish in one pass too. + +The output may be a strided view (e.g. the frames of a causal-conv history +buffer behind its cached prefix), so the result lands where the next GEMM +reads it without a ``torch.cat``. Statistics are accumulated in fp32; storage +dtype follows the tensors. On devices without Triton the native PyTorch chain +runs instead, so callers never need a platform check. +""" + +import torch +import torch.nn.functional as F +from vllm.triton_utils import HAS_TRITON, tl, triton + +_ACTIVATIONS = {None: 0, "mish": 1} + + +def _use_triton(tensor: torch.Tensor) -> bool: + """Triton on CUDA tensors; everything else runs the native path.""" + return HAS_TRITON and tensor.is_cuda + + +@triton.jit +def _residual_layer_norm_kernel( + out_ptr, + residual_out_ptr, + x_ptr, + y_ptr, + gate_ptr, + y_bias_ptr, + weight_ptr, + bias_ptr, + frames, + channels, + stride_on, + stride_ot, + stride_rn, + stride_rt, + stride_xn, + stride_xt, + stride_yn, + stride_yt, + eps, + HAS_X: tl.constexpr, + HAS_Y: tl.constexpr, + TAPS: tl.constexpr, + HAS_GATE: tl.constexpr, + HAS_Y_BIAS: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + STORE_RESIDUAL: tl.constexpr, + ACTIVATION: tl.constexpr, + BLOCK_C: tl.constexpr, +): + row = tl.program_id(0) + n = row // frames + t = row % frames + cols = tl.arange(0, BLOCK_C) + mask = cols < channels + + value = tl.zeros((BLOCK_C,), dtype=tl.float32) + if HAS_Y: + y = tl.zeros((BLOCK_C,), dtype=tl.float32) + for k in tl.static_range(TAPS): + y += tl.load( + y_ptr + n * stride_yn + (t + k) * stride_yt + k * channels + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + if HAS_Y_BIAS: + y += tl.load(y_bias_ptr + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_GATE: + y = tl.load(gate_ptr + cols, mask=mask, other=0.0).to(tl.float32) * y + value = y + if HAS_X: + value += tl.load(x_ptr + n * stride_xn + t * stride_xt + cols, mask=mask, other=0.0).to(tl.float32) + if STORE_RESIDUAL: + tl.store(residual_out_ptr + n * stride_rn + t * stride_rt + cols, value, mask=mask) + + mean = tl.sum(value, axis=0) / channels + centered = tl.where(mask, value - mean, 0.0) + variance = tl.sum(centered * centered, axis=0) / channels + normalized = centered * tl.rsqrt(variance + eps) + if HAS_WEIGHT: + normalized = normalized * tl.load(weight_ptr + cols, mask=mask, other=1.0).to(tl.float32) + if HAS_BIAS: + normalized = normalized + tl.load(bias_ptr + cols, mask=mask, other=0.0).to(tl.float32) + if ACTIVATION == 1: + # Mish = v * tanh(softplus(v)), with tanh(log(1 + e)) = e(e + 2) / (e(e + 2) + 2) + # for e = exp(v): no log/tanh, exact for very negative v, and + # PyTorch's softplus threshold (v > 20 -> v) makes the factor 1 there. + e = tl.exp(tl.minimum(normalized, 20.0)) + numerator = e * (e + 2.0) + factor = tl.where(normalized > 20.0, 1.0, numerator / (numerator + 2.0)) + normalized = normalized * factor + tl.store(out_ptr + n * stride_on + t * stride_ot + cols, normalized, mask=mask) + + +def _check_rows(name: str, tensor: torch.Tensor, rows: int, frames: int, width: int) -> None: + if tensor.dim() != 3 or int(tensor.shape[0]) != rows or int(tensor.shape[1]) < frames: + raise ValueError(f"residual_layer_norm: {name} must be (N, >= T, C), got {tuple(tensor.shape)}") + if int(tensor.shape[2]) != width or (width > 1 and tensor.stride(2) != 1): + raise ValueError(f"residual_layer_norm: {name} needs {width} contiguous channels, got {tuple(tensor.shape)}") + + +def _native( + *, + residual, + y, + gate, + weight, + bias, + eps, + taps, + y_bias, + activation, + residual_out, + out, + frames, + channels, +): + value = None + if y is not None: + value = sum(y[:, k : k + frames, k * channels : (k + 1) * channels] for k in range(taps)) + if y_bias is not None: + value = value + y_bias + if gate is not None: + value = gate * value + if residual is not None: + value = residual[:, :frames] if value is None else residual[:, :frames] + value + if residual_out is not None: + residual_out.copy_(value) + result = F.layer_norm(value, (channels,), None, None, eps) + if weight is not None: + result = result * weight + if bias is not None: + result = result + bias + if activation == "mish": + result = F.mish(result) + out.copy_(result) + return out + + +def residual_layer_norm( + residual: torch.Tensor | None, + y: torch.Tensor | None = None, + *, + gate: torch.Tensor | None = None, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + eps: float = 1e-6, + taps: int = 1, + y_bias: torch.Tensor | None = None, + activation: str | None = None, + residual_out: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """``act(layer_norm(residual + gate * y) * weight + bias)`` in one pass per row. + + Args: + residual: ``(N, T, C)`` residual stream, or ``None`` to normalize + ``gate * y`` alone. + y: ``(N, T, C)`` sub-block output, or with ``taps = K > 1`` the + ``(N, K - 1 + T, K * C)`` GEMM of a ``K``-tap causal convolution + over its history (see the module docstring). ``None`` normalizes + ``residual`` alone. + gate: ``(C,)`` multiplier of ``y`` (after ``y_bias``). + weight / bias: ``(C,)`` LayerNorm affine; for adaLN pass + ``1 + scale`` and ``shift``. + taps: convolution taps folded into ``y``. + y_bias: ``(C,)`` added to the tap sum (the convolution bias). + activation: ``None`` or ``"mish"``, applied after the affine. + residual_out: where ``residual + gate * y`` is stored, may be + ``residual`` itself (in place). ``None`` skips the store. + out: ``(N, T, C)`` destination, may be a strided view with + contiguous channels. Allocated when ``None``. + """ + if activation not in _ACTIVATIONS: + raise ValueError(f"residual_layer_norm: unsupported activation {activation!r}") + if residual is None and y is None: + raise ValueError("residual_layer_norm needs a residual or a y") + reference = residual if residual is not None else y + rows = int(reference.shape[0]) + if residual is not None: + frames, channels = int(residual.shape[1]), int(residual.shape[2]) + _check_rows("residual", residual, rows, frames, channels) + else: + channels = int(y.shape[2]) // taps + frames = int(y.shape[1]) - (taps - 1) + if y is not None: + _check_rows("y", y, rows, frames + taps - 1, taps * channels) + if residual_out is not None: + _check_rows("residual_out", residual_out, rows, frames, channels) + if out is None: + out = torch.empty((rows, frames, channels), device=reference.device, dtype=reference.dtype) + _check_rows("out", out, rows, frames, channels) + + if not _use_triton(reference): + return _native( + residual=residual, + y=y, + gate=gate, + weight=weight, + bias=bias, + eps=eps, + taps=taps, + y_bias=y_bias, + activation=activation, + residual_out=residual_out, + out=out, + frames=frames, + channels=channels, + ) + if rows * frames == 0: + return out + dummy = reference + _residual_layer_norm_kernel[(rows * frames,)]( + out, + residual_out if residual_out is not None else dummy, + residual if residual is not None else dummy, + y if y is not None else dummy, + gate if gate is not None else dummy, + y_bias if y_bias is not None else dummy, + weight if weight is not None else dummy, + bias if bias is not None else dummy, + frames, + channels, + out.stride(0), + out.stride(1), + residual_out.stride(0) if residual_out is not None else 0, + residual_out.stride(1) if residual_out is not None else 0, + residual.stride(0) if residual is not None else 0, + residual.stride(1) if residual is not None else 0, + y.stride(0) if y is not None else 0, + y.stride(1) if y is not None else 0, + float(eps), + HAS_X=residual is not None, + HAS_Y=y is not None, + TAPS=int(taps), + HAS_GATE=gate is not None, + HAS_Y_BIAS=y_bias is not None, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + STORE_RESIDUAL=residual_out is not None, + ACTIVATION=_ACTIVATIONS[activation], + BLOCK_C=triton.next_power_of_2(channels), + num_warps=4 if channels <= 1024 else 8, + ) + return out diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/vision_fused.py b/vllm_omni/model_executor/models/minicpmo_4_5/vision_fused.py new file mode 100644 index 00000000000..c508f92cb04 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/vision_fused.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused SigLIP encoder layers and CUDA graphs for MiniCPM-o 4.5's packed vision encode. + +``SiglipVisionTransformer.forward_packed`` runs each of the 27 layers as +LayerNorm, three projections, attention, a copy of the attention output into +a packed buffer, the output projection, a residual add, LayerNorm, fc1, GELU, +fc2 and another residual add: 13 kernels per layer, six of them full passes +over the activations. Here a layer is six kernels: + +* one packed QKV GEMM (the q/k/v weights stacked once; the modules keep views + of the stacked weight, so no memory is added); +* attention per run of equal grids, reading q/k/v as strided views of the + packed projection and returning its output in the token layout the next + GEMM reads, with no copy for the single-run case; +* the output projection without its bias, then one pass that adds the bias + and the residual, stores the new residual in place and applies + ``layer_norm2`` (``residual_layer_norm``); +* fc1 with its bias and the tanh GELU as the cuBLASLt epilogue + (``torch._addmm_activation``); +* fc2 without its bias, then the same residual pass into the next layer's + ``layer_norm1`` (the encoder's ``post_layernorm`` after the last layer). + +The math per token is the eager layer's; the residual add and the LayerNorm +statistics are computed in fp32 inside one pass instead of rounding the sum +to the model dtype first. + +A duplex camera frame always has one patch grid, and a runner step encodes +up to ``vision_batch_size`` of them, so ``VisionGraphEncoder`` captures the +packed tower + resampler per ``(grid, batch bucket)`` and replays it: at +batch 1..4 most of an eager encode is kernel launch time. Batches pad up to +their bucket with zero slices; a slice attends only to its own patches in +both the tower and the resampler, so padding never reaches a real slice. +""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Sequence + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.logger import init_logger +from vllm.platforms import current_platform + +from .ops import residual_layer_norm + +logger = init_logger(__name__) + +# Batch buckets of the vision graphs; the checkpoint's vision_batch_size (16) +# bounds a chunk, so it is the largest bucket. +DEFAULT_GRAPH_BATCH_SIZES = (1, 2, 4, 8, 16) +# Distinct (grid, bucket) graphs kept at once; duplex uses one grid. +_MAX_GRAPHS = 16 +# A (grid, bucket) is captured on its third encode, so one-off shapes (the +# startup profile run, a single turn-mode image) stay eager. +_CAPTURE_AFTER = 2 + + +def supports_fused_layers(vpm: nn.Module) -> bool: + """Whether ``encode_packed_fused`` covers ``vpm``'s layers (SigLIP with a GELU MLP). + + An explicit ``eager`` attention implementation (the A/B reference) keeps + the unfused layers, since the fused ones always attend through SDPA. + """ + layers = getattr(getattr(vpm, "encoder", None), "layers", None) + if not layers or getattr(vpm, "post_layernorm", None) is None: + return False + if getattr(getattr(vpm, "config", None), "_attn_implementation", None) == "eager": + return False + layer = layers[0] + attn = getattr(layer, "self_attn", None) + mlp = getattr(layer, "mlp", None) + return ( + attn is not None + and mlp is not None + and all(getattr(attn, name, None) is not None for name in ("q_proj", "k_proj", "v_proj", "out_proj")) + and getattr(mlp.config, "hidden_act", None) in ("gelu_pytorch_tanh", "gelu") + ) + + +def packed_qkv(attn: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]: + """``(3E, E)`` q/k/v weight (and bias) of ``attn``, stacked once. + + The q/k/v parameters are then re-pointed at row views of the stacked + tensors, so the old storage is freed and the eager path keeps working. + """ + cached = attn.__dict__.get("_packed_qkv") + projections = (attn.q_proj, attn.k_proj, attn.v_proj) + if cached is not None and all( + p.weight.data_ptr() == cached[0][i * p.weight.shape[0]].data_ptr() for i, p in enumerate(projections) + ): + return cached + with torch.no_grad(): + weight = torch.cat([p.weight for p in projections]) + biases = [p.bias for p in projections] + bias = torch.cat(biases) if all(b is not None for b in biases) else None + rows = int(projections[0].weight.shape[0]) + for index, projection in enumerate(projections): + projection.weight.data = weight[index * rows : (index + 1) * rows] + if bias is not None: + projection.bias.data = bias[index * rows : (index + 1) * rows] + attn.__dict__["_packed_qkv"] = (weight, bias) + return weight, bias + + +def _attend(attn: nn.Module, qkv: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor: + """``(tokens, E)`` attention output of packed ``(tokens, 3E)`` q/k/v, each item attending to itself.""" + heads, head_dim = attn.num_heads, attn.head_dim + embed = heads * head_dim + outputs = [] + for start, count, seq_len in seq_groups: + run = qkv[start : start + count * seq_len].view(count, seq_len, 3, heads, head_dim) + query, key, value = (run[:, :, index].transpose(1, 2) for index in range(3)) + out = F.scaled_dot_product_attention(query, key, value, scale=attn.scale) + outputs.append(out.transpose(1, 2).reshape(count * seq_len, embed)) + return outputs[0] if len(outputs) == 1 else torch.cat(outputs) + + +def _mlp_up(mlp: nn.Module, hidden: torch.Tensor) -> torch.Tensor: + fc1 = mlp.fc1 + if mlp.config.hidden_act == "gelu_pytorch_tanh" and fc1.bias is not None and hidden.is_cuda: + # cuBLASLt's GELU epilogue is the tanh approximation. + return torch._addmm_activation(fc1.bias, hidden, fc1.weight.t(), use_gelu=True) + return mlp.activation_fn(fc1(hidden)) + + +def encode_packed_fused( + vpm: nn.Module, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]] +) -> torch.Tensor: + """The encoder layers and ``post_layernorm`` of ``vpm.forward_packed`` over ``(tokens, E)`` embeddings. + + The residual stream is updated in place in a contiguous copy of + ``hidden_states`` (the patch embedding hands over a transposed view). + """ + layers = vpm.encoder.layers + residual = hidden_states.contiguous().unsqueeze(0) + first = layers[0].layer_norm1 + normed = residual_layer_norm(residual, weight=first.weight, bias=first.bias, eps=first.eps) + for index, layer in enumerate(layers): + attn = layer.self_attn + qkv_weight, qkv_bias = packed_qkv(attn) + attended = _attend(attn, F.linear(normed[0], qkv_weight, qkv_bias), seq_groups) + norm = layer.layer_norm2 + normed = residual_layer_norm( + residual, + F.linear(attended, attn.out_proj.weight).unsqueeze(0), + y_bias=attn.out_proj.bias, + weight=norm.weight, + bias=norm.bias, + eps=norm.eps, + residual_out=residual, + ) + mlp = layer.mlp + last = index + 1 == len(layers) + norm = vpm.post_layernorm if last else layers[index + 1].layer_norm1 + normed = residual_layer_norm( + residual, + F.linear(_mlp_up(mlp, normed[0]), mlp.fc2.weight).unsqueeze(0), + y_bias=mlp.fc2.bias, + weight=norm.weight, + bias=norm.bias, + eps=norm.eps, + residual_out=None if last else residual, + ) + return normed[0] + + +def _bucket(count: int, sizes: Sequence[int]) -> int | None: + for size in sizes: + if count <= size: + return size + return None + + +class VisionGraphEncoder: + """CUDA graphs of the packed SigLIP tower + resampler for single-grid chunks. + + ``encode(pixels, height, width, count)`` returns ``(count, queries, D)`` + for ``count`` slices of one ``(height, width)`` patch grid packed as + ``(1, C, patch, count * height * width * patch)``, or ``None`` when no + graph covers the chunk (the caller runs it eagerly). A ``(grid, bucket)`` + is captured once it recurs (``_CAPTURE_AFTER``) and kept in LRU order. + """ + + def __init__(self, vpm: nn.Module, resampler: nn.Module, *, batch_sizes: Sequence[int] = DEFAULT_GRAPH_BATCH_SIZES): + self.vpm = vpm + self.resampler = resampler + self.batch_sizes = tuple(sorted({int(b) for b in batch_sizes if int(b) > 0})) + self._graphs: OrderedDict[tuple[int, int, int], tuple[torch.Tensor, torch.Tensor, torch.cuda.CUDAGraph]] = ( + OrderedDict() + ) + self._failed: set[tuple[int, int, int]] = set() + self._seen: dict[tuple[int, int, int], int] = {} + self.enabled = True + + def encode(self, pixels: torch.Tensor, height: int, width: int, count: int) -> torch.Tensor | None: + if not self.enabled or torch.cuda.is_current_stream_capturing(): + return None + bucket = _bucket(count, self.batch_sizes) + if bucket is None: + return None + key = (int(height), int(width), bucket) + entry = self._graphs.get(key) + if entry is None: + if key in self._failed: + return None + seen = self._seen.get(key, 0) + if seen < _CAPTURE_AFTER: + if len(self._seen) >= 4 * _MAX_GRAPHS: + self._seen.clear() + self._seen[key] = seen + 1 + return None + entry = self._capture(key, pixels) + if entry is None: + return None + self._graphs.move_to_end(key) + static_pixels, static_out, graph = entry + width_px = int(pixels.shape[-1]) + static_pixels[..., :width_px].copy_(pixels) + if width_px < int(static_pixels.shape[-1]): + static_pixels[..., width_px:].zero_() + graph.replay() + return static_out[:count].clone() + + def _run(self, pixels: torch.Tensor, runs: list[tuple[int, int, int]]) -> torch.Tensor: + return self.resampler.forward_packed(self.vpm.forward_packed(pixels, runs), runs) + + def _capture(self, key: tuple[int, int, int], like: torch.Tensor) -> tuple | None: + height, width, bucket = key + runs = [(height, width, bucket)] + patch = int(like.shape[-2]) + static_pixels = torch.zeros( + (*like.shape[:-1], bucket * height * width * patch), device=like.device, dtype=like.dtype + ) + try: + stream = torch.cuda.Stream(device=like.device) + stream.wait_stream(torch.cuda.current_stream(like.device)) + with torch.cuda.stream(stream), torch.inference_mode(): + for _ in range(2): + self._run(static_pixels, runs) + torch.cuda.current_stream(like.device).wait_stream(stream) + graph = torch.cuda.CUDAGraph() + # thread_local: other threads of the worker may use CUDA meanwhile. + with ( + torch.inference_mode(), + torch.cuda.graph( + graph, pool=current_platform.get_global_graph_pool(), capture_error_mode="thread_local" + ), + ): + static_out = self._run(static_pixels, runs) + except Exception: + logger.warning("MiniCPM-o vision CUDA graph capture failed for %s; staying eager", key, exc_info=True) + self._failed.add(key) + return None + if len(self._graphs) >= _MAX_GRAPHS: + _, (_, _, evicted) = self._graphs.popitem(last=False) + evicted.reset() + self._graphs[key] = (static_pixels, static_out, graph) + logger.info("Captured MiniCPM-o vision CUDA graph grid=%dx%d batch=%d", height, width, bucket) + return self._graphs[key] From 11565f252084d296c6617a38e73c9339a5c1f95c Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:30 +0800 Subject: [PATCH 81/89] [Perf][MiniCPM-o] Code2Wav Whole-Euler slot pool and fused DiT body Code2Wav Whole-Euler slot pool and fused DiT body Implement Whole-Euler slot pool memory management for high-concurrency duplex Code2Wav, fused DiT block execution with tiled TF32 attention, and single-projection speaker embeddings. Signed-off-by: BeatSeat --- .../models/minicpmo_4_5/test_cfm_attention.py | 62 ++ .../minicpmo_4_5/test_code2wav_batching.py | 269 +++++- .../test_code2wav_bucket_stats.py | 99 +++ .../minicpmo_4_5/test_code2wav_model_dir.py | 33 + .../models/minicpmo_4_5/test_dit_fused.py | 435 ++++++++++ .../models/minicpmo_4_5/test_llm2tts.py | 17 +- .../test_minicpmo_4_5_omni.py | 23 +- .../test_minicpmo_4_5_transport.py | 199 +++++ .../models/minicpmo_4_5/batched_token2wav.py | 769 ++++++++++++++++-- .../models/minicpmo_4_5/cfm_attention.py | 149 ++++ .../models/minicpmo_4_5/dit_fused.py | 405 +++++++++ .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 479 ++++++++++- .../minicpmo_4_5/minicpmo_4_5_token2wav.py | 6 + .../minicpmo_4_5_omni.py | 208 +++-- 14 files changed, 2989 insertions(+), 164 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_cfm_attention.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_code2wav_bucket_stats.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_dit_fused.py create mode 100644 tests/model_executor/stage_input_processors/test_minicpmo_4_5_transport.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_cfm_attention.py b/tests/model_executor/models/minicpmo_4_5/test_cfm_attention.py new file mode 100644 index 00000000000..68188b745b0 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_cfm_attention.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"), +] + + +@pytest.mark.parametrize( + "batch,queries,keys,dim", [(2, 16, 16, 16), (1, 17, 17, 64), (3, 50, 350, 64), (8, 64, 368, 64)] +) +@pytest.mark.parametrize("masked", [False, True]) +@pytest.mark.parametrize("strided_dim", [False, True]) +def test_fp32_tiled_attention_mask_and_graph(batch, queries, keys, dim, masked, strided_dim): + from vllm_omni.model_executor.models.minicpmo_4_5.cfm_attention import cfm_attention + + torch.manual_seed(773) + q = torch.randn(batch, queries, 8, dim, device="cuda").transpose(1, 2) + # Like unpacking the model's interleaved K/V history cache. + kv = torch.randn(batch, 8, keys, dim * 2, device="cuda") + k, v = kv.chunk(2, dim=-1) + mask = torch.rand(batch, queries, keys, device="cuda") > 0.3 if masked else None + if mask is not None: + mask[:, 0] = False + if strided_dim: + q = q.transpose(-1, -2).contiguous().transpose(-1, -2) + k = k.transpose(-1, -2).contiguous().transpose(-1, -2) + v = v.transpose(-1, -2).contiguous().transpose(-1, -2) + if mask is not None: + mask = mask.transpose(-1, -2).contiguous().transpose(-1, -2) + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + expected = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=mask[:, None] if masked else None + ) + actual = cfm_attention(q, k, v, mask) + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = cfm_attention(q, k, v, mask) + if masked: + mask.logical_not_() + q.mul_(0.3) + expected = cfm_attention(q, k, v, mask) + graph.replay() + torch.testing.assert_close(captured, expected, atol=0, rtol=0) + + +def test_masked_prefix_keeps_later_negative_logits_normalized(): + from vllm_omni.model_executor.models.minicpmo_4_5.cfm_attention import cfm_attention + + q = torch.full((1, 2, 4, 64), 10.0, device="cuda") + k = torch.full((1, 2, 130, 64), -10.0, device="cuda") + v = torch.ones_like(k) + mask = torch.zeros(1, 4, 130, dtype=torch.bool, device="cuda") + mask[:, 1:, 128:] = True + expected = torch.ones_like(q) + expected[:, :, 0] = 0 + torch.testing.assert_close(cfm_attention(q, k, v, mask), expected, atol=2e-5, rtol=2e-5) diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index e97eb8e5a50..bb54f8f782b 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py @@ -13,14 +13,20 @@ import torch.nn as nn import vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav as batched_token2wav_module +import vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav as code2wav_module from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( BatchedToken2Wav, + BatchedToken2WavState, _token2wav_sdpa_context, plan_token2wav_encode_slices, relpos_encode_token_budget, + row_offset_signature, + state_shape_signature, ) +from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import AttSlotPool from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( MiniCPMO45Code2Wav, + _BucketStats, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -173,7 +179,7 @@ def _config(minimum: int = 1, initial: int = 0, *, runtime_prompt_cache_size: in @pytest.mark.parametrize( - ("extra", "max_num_seqs", "micro"), [({}, 6, 6), ({}, 64, 16), ({"micro_batch_size": 2}, 6, 2)] + ("extra", "max_num_seqs", "micro"), [({}, 6, 6), ({}, 64, 64), ({"micro_batch_size": 2}, 6, 6)] ) def test_whole_euler_micro_batch_defaults_to_stage_concurrency(extra, max_num_seqs, micro): config = _config() @@ -183,6 +189,47 @@ def test_whole_euler_micro_batch_defaults_to_stage_concurrency(extra, max_num_se model = MiniCPMO45Code2Wav(vllm_config=config) assert model._cfm_graph_config["micro_batch_size"] == micro + assert model._cfm_graph_config["max_graph_batch"] == micro + + +def test_stage2_memory_switches_are_opt_in(): + config = _config() + config.model_config.stage_connector_config["extra"].update( + { + "cfm_prompt_att_sharing": True, + "cfm_arena_rows_from_graph_grid": True, + } + ) + model = MiniCPMO45Code2Wav(vllm_config=config) + assert model._cfm_graph_config["prompt_att_sharing"] is True + assert model._cfm_graph_config["arena_rows_from_graph_grid"] is True + + default = MiniCPMO45Code2Wav(vllm_config=_config()) + assert default._cfm_graph_config["prompt_att_sharing"] is False + assert default._cfm_graph_config["arena_rows_from_graph_grid"] is False + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_precapture_empty_cache_is_opt_in_and_runs_once_after_precapture(enabled, monkeypatch): + config = _config() + if enabled: + config.model_config.stage_connector_config["extra"]["cfm_precapture_empty_cache"] = True + model = MiniCPMO45Code2Wav(vllm_config=config) + assert model._precapture_empty_cache is enabled + assert MiniCPMO45Code2Wav(vllm_config=_config())._precapture_empty_cache is False + calls = [] + model.backend = SimpleNamespace() + model._precapture_pending = True + + def precapture(): + calls.append("precapture") + model._precapture_pending = False + + monkeypatch.setattr(model, "_precapture_default_prompt", precapture) + monkeypatch.setattr(model, "_release_precapture_cache", lambda: calls.append("release")) + model.forward(input_ids=torch.empty(0, dtype=torch.long)) + model.forward(input_ids=torch.empty(0, dtype=torch.long)) + assert calls == (["precapture", "release"] if enabled else ["precapture"]) def _model(initial: int = 0, minimum: int = 1, *, runtime_prompt_cache_size: int = 4, setup_cache_size: int = 1): @@ -1809,3 +1856,223 @@ def test_padding_does_not_change_the_valid_frames(): assert int(padded_x.shape[2]) == int(exact_x.shape[2]) == mel_frames assert torch.allclose(padded_x, exact_x, atol=1e-5) + + +@pytest.mark.parametrize( + ("config", "env", "mode"), + [ + ({}, None, "off"), + ({"code2wav_allow_tf32": True}, None, "all"), + ({"code2wav_allow_tf32": "flow"}, None, "flow"), + ({"code2wav_allow_tf32": True}, "0", "off"), + ({}, "flow", "flow"), + ({}, "1", "all"), + ({"token2wav_allow_tf32": True}, None, "all"), + ({"token2wav_allow_tf32": True}, "0", "all"), + ], +) +def test_code2wav_tf32_mode(monkeypatch, config, env, mode): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import _tf32_mode + + if env is None: + monkeypatch.delenv("MINICPMO_CODE2WAV_TF32", raising=False) + else: + monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", env) + assert _tf32_mode(config) == mode + + +def test_precapture_scopes_flow_tf32_to_whole_euler_graphs_only(): + # A graph replays the kernels it was captured with, so the flow TF32 scope + # must cover the Whole-Euler precapture (otherwise every serving replay + # runs fp32 SGEMM) and must not leak into the HiFT capture. + backend = BatchedToken2Wav(_FakeToken2Wav(), cfm_tf32=True) + seen = {} + + class _Wrapper: + enabled = True + device = torch.device("cpu") + graph: dict = {} + + def precapture(self, **kwargs): + seen["whole_euler"] = torch.backends.cuda.matmul.allow_tf32 + return 1 + + def capture(self): + seen["hift"] = torch.backends.cuda.matmul.allow_tf32 + + backend._whole_euler_graph_wrapper = _Wrapper() + backend.hift_graph_wrapper = _Wrapper() + backend.flow.spk_embed_affine_layer = nn.Linear(1, 3) + before = torch.backends.cuda.matmul.allow_tf32 + try: + torch.backends.cuda.matmul.allow_tf32 = False + assert backend.precapture_whole_euler(SimpleNamespace(mels=torch.zeros(1, 4, 2))) == 1 + backend.precapture_hift() + assert torch.backends.cuda.matmul.allow_tf32 is False + finally: + torch.backends.cuda.matmul.allow_tf32 = before + assert seen == {"whole_euler": True, "hift": False} + + +def _att_pool(slots: int = 4) -> AttSlotPool: + return AttSlotPool( + n_timesteps=1, + depth=1, + heads=1, + width=2, + slots=slots, + frames=16, + suffix=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + +def _resident_state(pool: AttSlotPool | None, frames: int, *, mel_cache: int = 1) -> BatchedToken2WavState: + """A continuation state whose estimator cache is resident in ``pool`` (a tensor when ``pool`` is None).""" + estimator_att = torch.zeros(1, 1, 2, 1, frames, 2) + return BatchedToken2WavState( + flow_cache={ + "conformer_cnn_cache": torch.zeros(1, 1, 1), + "conformer_att_cache": torch.zeros(1, 1, 1, frames, 1), + "estimator_cnn_cache": torch.zeros(1, 1, 2, 1, 1), + "estimator_att_cache": estimator_att if pool is None else pool.adopt(estimator_att), + }, + hift_cache={"mel": torch.zeros(1, 1, mel_cache), "source": torch.zeros(1, 1, 2), "speech": torch.zeros(1, 2)}, + ) + + +def _continuation(state: BatchedToken2WavState | None) -> SimpleNamespace: + previous = None if state is None else SimpleNamespace(token2wav=state) + return SimpleNamespace(previous=previous, prompt_cache_id="p", prompt_wav="w", cache_epoch=0) + + +def test_row_offset_merge_is_opt_in_and_drops_only_slot_cache_lengths_from_the_key(): + model = MiniCPMO45Code2Wav(vllm_config=_config()) + assert model._row_offset_merge is False + assert model._cfm_graph_config["row_offset_merge"] is False + assert BatchedToken2Wav(_FakeToken2Wav())._row_offset_merge is False + config = _config() + config.model_config.stage_connector_config["extra"]["cfm_row_offset_merge"] = True + assert MiniCPMO45Code2Wav(vllm_config=config)._cfm_graph_config["row_offset_merge"] is True + assert BatchedToken2Wav(_FakeToken2Wav(), cfm_graph_config={"row_offset_merge": True})._row_offset_merge is True + + pool = _att_pool() + second, steady = _continuation(_resident_state(pool, 5)), _continuation(_resident_state(pool, 7)) + key = MiniCPMO45Code2Wav._bucket_key + # Off: the historical key, which holds each cache length. + assert key(second) == ("p", "w", state_shape_signature(second.previous.token2wav), 0) + assert key(second) != key(steady) + assert key(second, row_offsets=True) == key(steady, row_offsets=True) + signature = row_offset_signature(steady.previous.token2wav) + assert key(steady, row_offsets=True) == ("p", "w", ("row_offsets", signature), 0) + assert key(steady, cross_turn=True, row_offsets=True) == ("p", "w", ("row_offsets", signature)) + flow = dict(signature[0]) + assert flow["estimator_att_cache"][0] == (1, 1, 2, 1, -1, 2) + assert flow["conformer_att_cache"][0] == (1, 1, 1, -1, 1) + assert flow["estimator_cnn_cache"][0] == (1, 1, 2, 1, 1) + # Every other cache shape stays in the key. + assert key(_continuation(_resident_state(pool, 7, mel_cache=2)), row_offsets=True) != key(steady, row_offsets=True) + # Fresh streams and tensor caches (the arena path) keep their keys. + fresh = _continuation(None) + assert key(fresh, row_offsets=True) == key(fresh) + tensor_row = _continuation(_resident_state(None, 5)) + assert key(tensor_row, row_offsets=True) == key(tensor_row) + + +def test_row_offset_merge_needs_resident_rows_of_one_slot_pool(): + pool = _att_pool() + states = [_resident_state(pool, 5), _resident_state(pool, 7)] + backend = BatchedToken2Wav(_FakeToken2Wav(), cfm_graph_config={"row_offset_merge": True}) + # No ragged Whole-Euler graphs, no merge. + assert backend.can_merge_row_offsets(states, [2, 2]) is False + backend._whole_euler_graph_wrapper = SimpleNamespace(enabled=True, ragged_body=object(), slot_pool=pool) + assert backend.can_merge_row_offsets(states, [2, 2]) is True + assert backend.can_merge_row_offsets(states[:1], [2]) is False + assert backend.can_merge_row_offsets([states[0], _resident_state(None, 7)], [2, 2]) is False + assert backend.can_merge_row_offsets([states[0], _resident_state(_att_pool(), 7)], [2, 2]) is False + backend._row_offset_merge = False + assert backend.can_merge_row_offsets(states, [2, 2]) is False + + +def _second_chunk_next_to_steady(model): + """Streams ``a`` and ``b`` on their third chunk next to ``c``'s second: the last forward's output.""" + _forward(model, [_info("a", 0, [1, 2]), _info("b", 0, [3, 4])]) + _forward(model, [_info("a", 1, [5, 6]), _info("b", 1, [7, 8])]) + _forward(model, [_info("c", 0, [9, 10])]) + return _forward(model, [_info("a", 2, [11, 12]), _info("b", 2, [13, 14]), _info("c", 1, [15, 16])]) + + +def _row_offset_model(monkeypatch, *, mergeable: bool): + """A fake-backend model with ``cfm_row_offset_merge`` on; continuation rows count as slot-pool rows.""" + model, token2wav = _model() + model._row_offset_merge = True + model.backend._row_offset_merge = True + model._bucket_stats = _BucketStats(every=1000) + monkeypatch.setattr(code2wav_module, "_slot_resident", lambda item: item.previous is not None) + if mergeable: + monkeypatch.setattr(model.backend, "can_merge_row_offsets", lambda states, counts: True) + return model, token2wav + + +def test_row_offset_merge_decodes_a_second_chunk_with_steady_rows(monkeypatch): + split_model, split_t2w = _model() + split = _second_chunk_next_to_steady(split_model) + # Off: the second chunk (cache 4 + 1 frames) and the steady rows (4 + 2) are two buckets. + assert split_t2w.hift.calls[-2:] == [2, 1] + assert [state.token2wav.flow_cache["estimator_att_cache"].shape[4] for state in split_model._states.values()] == [ + 7, + 7, + 6, + ] + + merged_model, merged_t2w = _row_offset_model(monkeypatch, mergeable=True) + merged = _second_chunk_next_to_steady(merged_model) + # On: one decode and one vocoder call; the encoder still runs once per conformer cache length. + assert merged_t2w.hift.calls[-1:] == [3] + assert merged_t2w.flow.encoder.calls[-2:] == [2, 1] + stats = merged_model._bucket_stats + assert (stats.forwards, stats.decodes, stats.decoded_rows, stats.reasons) == (4, 4, 8, {}) + for got, want in zip( + merged.multimodal_outputs["model_outputs"], split.multimodal_outputs["model_outputs"], strict=True + ): + torch.testing.assert_close(got, want, rtol=0, atol=1e-6) + for request_id, state in split_model._states.items(): + merged_state = merged_model._states[request_id] + for name, expected in {**state.token2wav.flow_cache, **state.token2wav.hift_cache}.items(): + actual = {**merged_state.token2wav.flow_cache, **merged_state.token2wav.hift_cache}[name] + torch.testing.assert_close(actual, expected, rtol=0, atol=1e-6, msg=f"{request_id} {name}") + + +def test_row_offset_bucket_the_backend_refuses_splits_back_into_signature_buckets(monkeypatch): + split_model, _ = _model() + split = _second_chunk_next_to_steady(split_model) + + model, token2wav = _row_offset_model(monkeypatch, mergeable=False) + output = _second_chunk_next_to_steady(model) + assert token2wav.hift.calls[-2:] == [2, 1] + assert model._bucket_stats.reasons == {"row_offset_merge_rejected": 1} + for got, want in zip( + output.multimodal_outputs["model_outputs"], split.multimodal_outputs["model_outputs"], strict=True + ): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +def test_projected_speaker_does_not_depend_on_the_batch_size(): + torch.manual_seed(0) + layer = nn.Linear(6, 4) + rows_seen: list[int] = [] + layer.register_forward_pre_hook(lambda _module, args: rows_seen.append(int(args[0].shape[0]))) + backend = SimpleNamespace(flow=SimpleNamespace(spk_embed_affine_layer=layer)) + speaker = torch.randn(1, 6) + + with torch.no_grad(): + single = BatchedToken2Wav._project_speakers(backend, speaker.expand(1, -1)) + batched = BatchedToken2Wav._project_speakers(backend, speaker.expand(5, -1)) + distinct = BatchedToken2Wav._project_speakers(backend, torch.randn(3, 6)) + + # One prompt speaker expanded over the batch projects as a single row. + assert rows_seen == [1, 1, 3] + assert batched.shape == (5, 4) and batched.is_contiguous() + assert torch.equal(batched, single.expand(5, -1)) + assert distinct.shape == (3, 4) diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_bucket_stats.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_bucket_stats.py new file mode 100644 index 00000000000..aa424f95243 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_bucket_stats.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Code2Wav decode buckets: ``code2wav_bucket_stats`` and ``cfm_cross_turn_buckets`` (connector extras).""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +import vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav as code2wav_module +from tests.model_executor.models.minicpmo_4_5.test_code2wav_batching import _forward, _info, _model +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( + MiniCPMO45Code2Wav, + _BucketStats, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _two_turn_rows() -> list[dict]: + # Two fresh streams of the same voice that are on different turns. + return [_info("a", 0, [1, 2], cache_epoch=0), _info("b", 0, [3, 4], cache_epoch=1)] + + +def test_switches_default_off_and_keep_the_historical_key() -> None: + model, _ = _model() + assert model._cross_turn_buckets is False + assert model._bucket_stats is None + item = SimpleNamespace(previous=None, prompt_cache_id="p", prompt_wav="w", cache_epoch=3) + assert MiniCPMO45Code2Wav._bucket_key(item) == ("p", "w", ("uninitialized",), 3) + assert MiniCPMO45Code2Wav._bucket_key(item, cross_turn=True) == ("p", "w", ("uninitialized",)) + key = ("p", "w", ("sig",), 3) + assert MiniCPMO45Code2Wav._onset_group_key(key) == ("p", "w", 3) + assert MiniCPMO45Code2Wav._onset_group_key(key[:3], cross_turn=True) == ("p", "w") + + +def test_rows_of_different_turns_split_by_default_and_merge_when_enabled() -> None: + split_model, split_t2w = _model() + split = _forward(split_model, _two_turn_rows()) + assert split_t2w.hift.calls == [1, 1] + + merged_model, merged_t2w = _model() + merged_model._cross_turn_buckets = True + merged = _forward(merged_model, _two_turn_rows()) + assert merged_t2w.hift.calls == [2] + for got, want in zip( + merged.multimodal_outputs["model_outputs"], split.multimodal_outputs["model_outputs"], strict=True + ): + torch.testing.assert_close(got, want, rtol=0, atol=1e-6) + # Each stream keeps its own turn's state. + assert merged_model._states["a"].cache_epoch == 0 + assert merged_model._states["b"].cache_epoch == 1 + + +def test_bucket_stats_count_rows_buckets_decodes_and_split_reasons(monkeypatch) -> None: + logger = MagicMock() + monkeypatch.setattr(code2wav_module, "logger", logger) + model, _ = _model() + model._bucket_stats = _BucketStats(every=2) + _forward(model, _two_turn_rows()) + assert logger.info.call_count == 0 + _forward(model, [_info("c", 0, [5, 6])]) + assert logger.info.call_count == 1 + message, forwards, rows, buckets, multi_buckets, multi_forwards, decodes, rows_per_decode, reasons = ( + logger.info.call_args.args + ) + assert "Code2Wav buckets" in message + assert forwards == 2 + assert rows == {1: 1, 2: 1} + assert buckets == pytest.approx(1.5) + assert (multi_buckets, multi_forwards) == (pytest.approx(2.0), 1) + assert decodes == pytest.approx(1.5) + assert rows_per_decode == pytest.approx(1.0) + assert reasons == {"cache_epoch": 1} + # The window restarts after each line. + assert model._bucket_stats.forwards == 0 + + +def test_split_reason_names_prompt_onset_resident_and_tensor_signatures() -> None: + stats = _BucketStats() + + def row(prompt: str, att_cache: object | None, epoch: int = 0) -> SimpleNamespace: + previous = ( + None + if att_cache is None + else SimpleNamespace(token2wav=SimpleNamespace(flow_cache={"estimator_att_cache": att_cache})) + ) + return SimpleNamespace(previous=previous, prompt_cache_id=prompt, prompt_wav="w", cache_epoch=epoch) + + tensor = torch.zeros(1) + assert stats.split_reason(row("p", tensor), row("q", tensor), cross_turn=False) == "prompt" + # A fresh stream next to a continuation is the onset split, not a cache layout one. + assert stats.split_reason(row("p", None), row("p", object()), cross_turn=True) == "onset" + assert stats.split_reason(row("p", tensor), row("p", None), cross_turn=True) == "onset" + assert stats.split_reason(row("p", object()), row("p", tensor), cross_turn=True) == "signature_resident" + assert stats.split_reason(row("p", tensor), row("p", torch.zeros(2)), cross_turn=True) == "signature_tensor" diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py index fced5ae46ea..b8f76aac019 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py @@ -86,3 +86,36 @@ def test_default_prompt_wav_follows_resolved_model_path(monkeypatch, tmp_path): model = MiniCPMO45Code2Wav(vllm_config=config) model.model_path = "/resolved/snapshot" assert model._default_prompt_wav == "/resolved/snapshot/assets/HT_ref_audio.wav" + + +def _config_with_extra(extra: dict, *, max_num_seqs=None): + scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) if max_num_seqs is not None else None + return SimpleNamespace( + model_config=SimpleNamespace( + model="/fake/model", + stage_connector_config={"extra": extra}, + ), + scheduler_config=scheduler_config, + ) + + +@pytest.mark.parametrize( + ("extra", "max_num_seqs", "sizes"), + [ + ({}, None, [1, 2, 4, 8, 16, 32]), + ({}, 6, [1, 2, 4, 6]), + ({"hift_graph_capture_batch_sizes": [1, 3]}, None, [1, 3]), + ], +) +def test_hift_capture_batch_sizes(monkeypatch, tmp_path, extra, max_num_seqs, sizes): + """By default every vocoder batch up to min(max_num_seqs, 32) rounds up to a captured size.""" + _no_hub(monkeypatch, tmp_path) + config = _config_with_extra({"enable_hift_graph": True, **extra}, max_num_seqs=max_num_seqs) + assert MiniCPMO45Code2Wav(vllm_config=config)._hift_graph_config["capture_batch_sizes"] == sizes + + +@pytest.mark.parametrize(("extra", "frames"), [({}, []), ({"hift_graph_codec_chunk_frames": [25, 75]}, [25, 75])]) +def test_hift_graph_codec_chunk_frames_reach_the_connector_config(monkeypatch, tmp_path, extra, frames): + _no_hub(monkeypatch, tmp_path) + config = _config_with_extra({"enable_hift_graph": True, **extra}) + assert MiniCPMO45Code2Wav(vllm_config=config)._connector_config["hift_graph_codec_chunk_frames"] == frames diff --git a/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py b/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py new file mode 100644 index 00000000000..5a76274ba75 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The fused Code2Wav DiT body computes the ragged body's function (``dit_fused.py``).""" + +import gc + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav +from vllm_omni.model_executor.models.minicpmo_4_5.dit_fused import ( + blocks_forward_chunk_fused, + dit_modulation, + supports_fused_body, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +def _upstream_dit(device: str) -> nn.Module: + """The shipped DiT architecture at toy width with non-trivial adaLN weights.""" + decoder_dit = pytest.importorskip("stepaudio2.cosyvoice2.flow.decoder_dit") + torch.manual_seed(0) + estimator = decoder_dit.DiT(in_channels=320, out_channels=80, depth=2, num_heads=4, head_dim=16, hidden_size=64) + with torch.no_grad(): + # The adaLN-Zero init makes every block an identity; any weights will do here. + for parameter in estimator.parameters(): + parameter.normal_(0.0, 0.08) + return estimator.eval().to(device) + + +def _chunk(estimator: nn.Module, device: str, *, rows: int, frames: int, cached: int): + torch.manual_seed(1) + depth = len(estimator.blocks) + channels = estimator.blocks[0].conv.block[1].in_channels + attn = estimator.blocks[0].attn + estimator_input = torch.randn(rows, 320, frames, device=device) + time_embedding = estimator.t_embedder(torch.rand((), device=device).expand(rows)).unsqueeze(1) + mask = torch.rand(rows, frames, cached + frames, device=device) > 0.3 + mask[..., 0] = True + cnn = [torch.randn(rows, 2 * channels, 2, device=device) for _ in range(depth)] + att = [torch.randn(rows, attn.num_heads, cached, 2 * attn.head_dim, device=device) for _ in range(depth)] + return estimator_input, time_embedding, mask, cnn, att + + +def _run(body, estimator, chunk, lengths, **kwargs): + """``body`` over ``chunk`` (``_chunk``'s tuple; ``None`` caches for a first chunk) into fresh output caches.""" + estimator_input, time_embedding, mask, cnn, att = chunk + rows, _, frames = estimator_input.shape + depth = len(estimator.blocks) + attn = estimator.blocks[0].attn + channels = estimator.blocks[0].conv.block[1].in_channels + cached = 0 if att[0] is None else int(att[0].shape[2]) + cnn_out = torch.zeros((depth, rows, 2 * channels, 2), device=estimator_input.device) + att_out = torch.zeros( + (depth, rows, attn.num_heads, cached + frames, 2 * attn.head_dim), device=estimator_input.device + ) + with torch.no_grad(): + out = body(estimator, estimator_input, time_embedding, mask, cnn, att, cnn_out, att_out, lengths, **kwargs) + return out, cnn_out, att_out + + +def _solve_inputs(batch_size: int, width: int, seed: int) -> dict[str, torch.Tensor]: + """One chunk's Whole-Euler solve inputs (``WholeEulerCFMGraphWrapper.replay``).""" + torch.manual_seed(seed) + return { + "x": torch.randn(batch_size, 80, width, device="cuda"), + "mu_cfg": torch.randn(2 * batch_size, 80, width, device="cuda"), + "speakers_cfg": torch.randn(2 * batch_size, 80, device="cuda"), + "cond_cfg": torch.randn(2 * batch_size, 80, width, device="cuda"), + } + + +@pytest.fixture +def fp32_graphs(monkeypatch: pytest.MonkeyPatch) -> None: + """A private graph pool, and fp32 cuDNN convolutions (they default to TF32). + + Both bodies then differ only in reduction order across the 10 Euler steps. + """ + from vllm.platforms import current_platform + + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_fused_body_matches_ragged_body(device: str) -> None: + estimator = _upstream_dit(device) + assert supports_fused_body(estimator) + chunk = _chunk(estimator, device, rows=4, frames=6, cached=5) + lengths = [6, 4] + + expected = _run(BatchedToken2Wav._blocks_forward_chunk_ragged, estimator, chunk, lengths) + # On CUDA the fused body attends with the tiled tf32x3 kernel, the reference with SDPA. + actual = _run(blocks_forward_chunk_fused, estimator, chunk, lengths) + for want, got in zip(expected, actual, strict=True): + torch.testing.assert_close(got, want, rtol=1e-5, atol=2e-5) + + # A precomputed per-timestep table is the one computed on the fly. + with torch.no_grad(): + table = dit_modulation(estimator, chunk[1][:1]) + tabled = _run(blocks_forward_chunk_fused, estimator, chunk, lengths, modulation=table) + torch.testing.assert_close(tabled[0], actual[0], rtol=0, atol=0) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_fused_body_first_chunk_without_caches(device: str) -> None: + estimator = _upstream_dit(device) + estimator_input, time_embedding, _, _, _ = _chunk(estimator, device, rows=2, frames=8, cached=0) + depth = len(estimator.blocks) + mask = torch.ones(2, 8, 8, dtype=torch.bool, device=device) + chunk = (estimator_input, time_embedding, mask, [None] * depth, [None] * depth) + expected = _run(BatchedToken2Wav._blocks_forward_chunk_ragged, estimator, chunk, [8]) + actual = _run(blocks_forward_chunk_fused, estimator, chunk, [8]) + for want, got in zip(expected, actual, strict=True): + torch.testing.assert_close(got, want, rtol=1e-5, atol=2e-5) + + + +def test_modulation_table_layout() -> None: + estimator = _upstream_dit("cpu") + embedding = estimator.t_embedder(torch.tensor([0.3])).unsqueeze(1) + with torch.no_grad(): + table = dit_modulation(estimator, embedding) + block = estimator.blocks[1].adaLN_modulation(embedding).reshape(9, -1) + final_shift, final_scale = estimator.final_layer.adaLN_modulation(embedding).reshape(2, -1) + depth = len(estimator.blocks) + assert table.shape == (depth + 1, 9, 64) + torch.testing.assert_close(table[1, 0], block[0]) + torch.testing.assert_close(table[1, 1], block[1] + 1) # scale_msa -> 1 + scale + torch.testing.assert_close(table[1, 2], block[2]) # gates stay as they are + torch.testing.assert_close(table[depth, 0], final_shift) + torch.testing.assert_close(table[depth, 1], final_scale + 1) + + +def test_unsupported_estimator_is_refused() -> None: + estimator = _upstream_dit("cpu") + estimator.blocks[0].attn.q_norm = nn.Identity() + assert not supports_fused_body(estimator) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_graph_with_fused_body_matches_ragged_body(fp32_graphs: None) -> None: + """Captured with the arena's per-timestep modulation, the fused body replays the ragged body's solve.""" + from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import WholeEulerCFMGraphWrapper + + estimator = _upstream_dit("cuda") + batch_size, width = 3, 8 + results = [] + for body, modulation_fn in ( + (BatchedToken2Wav._blocks_forward_chunk_ragged, None), + (blocks_forward_chunk_fused, lambda embedding: dit_modulation(estimator, embedding)), + ): + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=8, + query_bucket_frames=16, + ragged_body=body, + modulation_fn=modulation_fn, + ) + first = wrapper.replay(**_solve_inputs(batch_size, width, 2), cnn_cache=None, att_cache=None) + assert first is not None and wrapper.stats_snapshot()["eager"] == 0 + second = wrapper.replay( + **_solve_inputs(batch_size, width, 3), cnn_cache=first[1], att_cache=first[2], valid_lengths=[8, 5, 3] + ) + assert second is not None + results.append((first, second)) + if modulation_fn is not None: + # One table per timestep, shared by every graph. + assert wrapper.arena._modulation.shape[:2] == (10, len(estimator.blocks) + 1) + + for want, got in zip(results[0], results[1], strict=True): + torch.testing.assert_close(got[0], want[0], rtol=1e-4, atol=2e-4) + torch.testing.assert_close(got[1], want[1], rtol=1e-4, atol=2e-4) + want_att = want[2] if isinstance(want[2], torch.Tensor) else torch.cat([a.flatten() for a in want[2]]) + got_att = got[2] if isinstance(got[2], torch.Tensor) else torch.cat([a.flatten() for a in got[2]]) + torch.testing.assert_close(got_att, want_att, rtol=1e-4, atol=2e-4) + + +def test_slot_pool_adopt_materialize_round_trip() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import AttSlotPool + + pool = AttSlotPool( + n_timesteps=2, + depth=3, + heads=2, + width=4, + slots=2, + frames=16, + suffix=3, + device=torch.device("cpu"), + dtype=torch.float32, + ) + cache = torch.randn(2, 3, 2, 2, 7, 4) + first, second = pool.adopt(cache), pool.adopt(cache) + assert first is not None and second is not None and pool.free_slots == 0 + assert pool.adopt(cache) is None + # The suffix sits at the fixed frames, the rest in order behind it. + assert first.physical_frames() == [3, 4, 5, 6, 0, 1, 2] + assert first.shape == cache.shape + torch.testing.assert_close(first.materialize(), cache, rtol=0, atol=0) + del first + gc.collect() + assert pool.free_slots == 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_slot_pool_replay_matches_arena_replay(fp32_graphs: None) -> None: + """Resident caches give the arena path's chunks, CNN caches and trimmed attention caches.""" + from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + ResidentAttCache, + WholeEulerCFMGraphWrapper, + _materialize_att_rows, + ) + + estimator = _upstream_dit("cuda") + depth, heads, att_width = ( + len(estimator.blocks), + estimator.blocks[0].attn.num_heads, + 2 * estimator.blocks[0].attn.head_dim, + ) + batch_size, width, prompt_len, suffix = 3, 8, 12, 6 + keep = (prompt_len, suffix) + torch.manual_seed(5) + # One shared prompt state, as ``setup_batch`` hands every fresh request. + prompt = torch.randn(10, depth, 2, heads, prompt_len, att_width, device="cuda") + cnn = torch.randn(10, depth, 2 * batch_size, 2 * 64, 2, device="cuda") + lengths = [None, [8, 5, 3], None, [4, 8, 6]] + results = [] + for slots in (0, 4): + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=16, + query_bucket_frames=16, + ragged_body=blocks_forward_chunk_fused, + modulation_fn=lambda embedding: dit_modulation(estimator, embedding), + att_slots=slots, + ) + att: list = [prompt] * batch_size + cnn_cache = cnn + outputs = [] + for step, valid in enumerate(lengths): + mask = None + if valid is not None: + # ``_decode_cfm``'s ragged mask: each row's own queries and current keys, all of its cache. + offset = int(att[0].shape[4]) + cfg = torch.tensor(valid * 2, device="cuda") + queries = torch.arange(width, device="cuda").unsqueeze(0) < cfg.unsqueeze(1) + keys = torch.cat((queries, torch.ones(2 * batch_size, offset, dtype=torch.bool, device="cuda")), dim=1) + mask = queries.unsqueeze(2) & keys.unsqueeze(1) + out = wrapper.replay( + **_solve_inputs(batch_size, width, 10 + step), + cnn_cache=cnn_cache, + att_cache=att, + attn_mask=mask, + valid_lengths=valid, + att_keep=keep, + ) + assert out is not None + mel, cnn_cache, att = out + if valid is not None: + # Each row keeps its own frames; the next chunk needs one cache length again. + shortest = min(int(row.shape[4]) for row in att) + assert all(int(row.shape[4]) == shortest for row in att) + outputs.append((mel, cnn_cache.clone(), [row.clone() for row in _materialize_att_rows(att)])) + stats = wrapper.stats_snapshot() + assert stats["eager"] == 0 + if slots: + assert stats["slot_replays"] == len(lengths) + assert all(isinstance(row, ResidentAttCache) for row in att) + # Each request holds one slot; dropping its state frees it. + assert wrapper.slot_pool.free_slots == slots - batch_size + del att, out + gc.collect() + assert wrapper.slot_pool.free_slots == slots + results.append(outputs) + + for step, (want, got) in enumerate(zip(results[0], results[1], strict=True)): + + def close(actual: torch.Tensor, expected: torch.Tensor, step: int = step) -> None: + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=2e-4, msg=lambda text: f"chunk {step}: {text}") + + close(got[0], want[0]) + close(got[1], want[1]) + for want_att, got_att in zip(want[2], got[2], strict=True): + assert got_att.shape == want_att.shape, f"chunk {step}" + close(got_att, want_att) + + +def _cfg_rows(tensor: torch.Tensor, batch_size: int, rows: list[int]) -> torch.Tensor: + """Requests ``rows`` of a CFG-stacked ``[cond x B | uncond x B]`` tensor (axis 0).""" + return tensor[[*rows, *(batch_size + row for row in rows)]] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +@pytest.mark.parametrize("grid", ["4", ""], ids=["one-graph-batch", "graph-grid"]) +def test_slot_pool_replay_merges_rows_of_different_cache_lengths( + fp32_graphs: None, monkeypatch: pytest.MonkeyPatch, grid: str +) -> None: + """A stream's second chunk and steady chunks in one slot-pool replay (``row_offsets``). + + Each row attends its own slot frames and the columns past a shorter cache + are masked, so on one graph batch the merged solve is bit-identical to + the separate replays; across graph batches only the GEMM tiling differs. + """ + from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + WholeEulerCFMGraphWrapper, + _materialize_att_rows, + ) + + monkeypatch.setenv("VLLM_OMNI_CFM_GRAPH_GRID", grid) + estimator = _upstream_dit("cuda") + depth, heads, att_width = ( + len(estimator.blocks), + estimator.blocks[0].attn.num_heads, + 2 * estimator.blocks[0].attn.head_dim, + ) + # A slot keeps the last ``suffix`` frames fixed, so the prompt holds at least that many. + width, prompt_len, suffix = 8, 20, 16 + keep = (prompt_len, suffix) + torch.manual_seed(5) + prompt = torch.randn(10, depth, 2, heads, prompt_len, att_width, device="cuda") + # Requests 0 and 1 run two chunks (steady cache 20 + 16 frames), request 2 + # one (second-chunk cache 20 + 8); then all three take a chunk together. + cnn = [torch.randn(10, depth, 2, 2 * 64, 2, device="cuda") for _ in range(3)] + final = _solve_inputs(3, width, 30) + + def stacked_cnn(rows: list) -> torch.Tensor: + return torch.cat([*(row[:, :, 0:1] for row in rows), *(row[:, :, 1:2] for row in rows)], dim=2) + + def solve(wrapper, requests: list[int], att: list, cnn_rows: list, inputs: dict, **kwargs): + out = wrapper.replay(**inputs, cnn_cache=stacked_cnn(cnn_rows), att_cache=att, att_keep=keep, **kwargs) + assert out is not None + mel, out_cnn, out_att = out + batch_size = len(requests) + return mel, [out_cnn[:, :, [row, batch_size + row]] for row in range(batch_size)], out_att + + def history(row_offsets: bool): + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=16, + micro_batch_size=4, + pad_max_rows=3, + query_bucket_frames=16, + ragged_body=blocks_forward_chunk_fused, + modulation_fn=lambda embedding: dit_modulation(estimator, embedding), + att_slots=4, + row_offsets=row_offsets, + ) + att: list = [prompt, prompt, prompt] + cnn_rows = list(cnn) + for step, requests in enumerate(([0, 1], [0, 1], [2])): + inputs = _solve_inputs(len(requests), width, 10 + step) + _, new_cnn, new_att = solve( + wrapper, requests, [att[r] for r in requests], [cnn_rows[r] for r in requests], inputs + ) + for request, row_cnn, row_att in zip(requests, new_cnn, new_att, strict=True): + cnn_rows[request], att[request] = row_cnn, row_att + assert [int(row.shape[4]) for row in att] == [prompt_len + suffix, prompt_len + suffix, prompt_len + width] + return wrapper, att, cnn_rows + + def subset(rows: list[int]) -> dict: + return { + "x": final["x"][rows], + **{name: _cfg_rows(final[name], 3, rows) for name in ("mu_cfg", "speakers_cfg", "cond_cfg")}, + } + + # Switch off: rows of different cache lengths never reach the slot pool. + wrapper, att, cnn_rows = history(row_offsets=False) + layouts = [list(row.layout) for row in att] + assert ( + wrapper.replay(**final, cnn_cache=stacked_cnn(cnn_rows), att_cache=[att[2], att[0], att[1]], att_keep=keep) + is None + ) + assert [list(row.layout) for row in att] == layouts + + # Separate buckets: the second chunk alone, then the steady pair. + split_mel, split_cnn, split_att = {}, {}, {} + for requests in ([2], [0, 1]): + rows = [{2: 0, 0: 1, 1: 2}[r] for r in requests] + mel, new_cnn, new_att = solve( + wrapper, requests, [att[r] for r in requests], [cnn_rows[r] for r in requests], subset(rows) + ) + for index, request in enumerate(requests): + split_mel[request], split_cnn[request] = mel[index], new_cnn[index] + split_att[request] = _materialize_att_rows([new_att[index]])[0].clone() + + # One replay: rows [second chunk, steady, steady] with ``_decode_cfm``'s row-offset mask. + merged_wrapper, merged_att, merged_cnn_rows = history(row_offsets=True) + order = [2, 0, 1] + offsets = [int(merged_att[r].shape[4]) for r in order] + lengths = torch.tensor([width] * 6, device="cuda") + queries = torch.arange(width, device="cuda").unsqueeze(0) < lengths.unsqueeze(1) + cached = torch.arange(max(offsets), device="cuda").unsqueeze(0) < torch.tensor( + offsets * 2, device="cuda" + ).unsqueeze(1) + mask = queries.unsqueeze(2) & torch.cat((queries, cached), dim=1).unsqueeze(1) + replays = merged_wrapper.stats_snapshot()["slot_replays"] + mel, new_cnn, new_att = solve( + merged_wrapper, + order, + [merged_att[r] for r in order], + [merged_cnn_rows[r] for r in order], + final, + attn_mask=mask, + valid_lengths=[width] * 3, + ) + stats = merged_wrapper.stats_snapshot() + assert stats["slot_replays"] == replays + 1 and stats["eager"] == 0 + # Each row keeps its own cache length and slot layout. + assert [int(row.shape[4]) for row in new_att] == [prompt_len + suffix] * 3 + assert [row.layout for row in new_att] == [att[r].layout for r in order] + + exact = grid == "4" + for index, request in enumerate(order): + + def close(actual: torch.Tensor, expected: torch.Tensor, what: str, request: int = request) -> None: + if exact: + assert torch.equal(actual, expected), f"request {request} {what}" + else: + torch.testing.assert_close( + actual, expected, rtol=1e-4, atol=2e-4, msg=lambda t: f"{request} {what}: {t}" + ) + + close(mel[index], split_mel[request], "mel") + close(new_cnn[index], split_cnn[request], "cnn") + close(_materialize_att_rows([new_att[index]])[0], split_att[request], "att") diff --git a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py index d93755fec08..1088a16b28d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py +++ b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py @@ -22,6 +22,7 @@ import pytest import torch +from vllm_omni.engine.duplex.intermediate import get_tts_handoff, unpack_transport_tensor from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( MiniCPMO45OmniForConditionalGeneration, ) @@ -149,7 +150,7 @@ def test_request_level_multimodal_output_feeds_orchestrator_bridge(self) -> None out = llm2tts([thinker_output], prompt=None) assert out[0]["prompt_token_ids"] == [0, 0, 0] - assert torch.equal(torch.tensor(out[0]["model_intermediate_buffer"]["hidden_states"]["tts"]), hidden[1:]) + assert torch.equal(get_tts_handoff(out[0]["model_intermediate_buffer"])[1], hidden[1:]) def test_returns_one_entry_per_input(self) -> None: hidden = torch.zeros((3, _HIDDEN_DIM)) @@ -210,7 +211,7 @@ def test_strict_talker_receives_canonical_nested_handoff(self) -> None: buffer = result["model_intermediate_buffer"] assert result["prompt_token_ids"] == [0, 0, 0, 0] assert buffer["ids"]["tts"] == [20, 21] - assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[2:4]) + assert torch.equal(get_tts_handoff(buffer)[1], hidden[2:4]) def test_latent_in_multimodal_output_takes_precedence(self) -> None: # When both ``multimodal_output["latent"]`` and ``hidden_states`` are @@ -234,7 +235,7 @@ def test_latent_in_multimodal_output_takes_precedence(self) -> None: ) buffer = result[0]["model_intermediate_buffer"] # latent (ones) won over hidden_states (zeros) - assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), latent[3:4].to(torch.float32)) + assert torch.equal(get_tts_handoff(buffer)[1], latent[3:4].to(torch.float32)) class TestTtsRegionDetection: @@ -269,7 +270,7 @@ def test_4_5_markers_detected(self, tts_end_id: int) -> None: # 4.5 BOS at idx 2 -> slice starts at 3; end at idx 5 -> slice ends at 5. buffer, hidden = self._run([10, 11], [151703, 30, 31, tts_end_id, 40]) assert buffer["ids"]["tts"] == [30, 31] - assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5]) + assert torch.equal(get_tts_handoff(buffer)[1], hidden[3:5]) def test_bos_without_eos_runs_to_end(self) -> None: # When BOS is found but EOS is missing (typical for an in-flight or @@ -278,12 +279,12 @@ def test_bos_without_eos_runs_to_end(self) -> None: # sequence: [10, 11, 151703, 30, 31] buffer, hidden = self._run([10, 11], [151703, 30, 31]) assert buffer["ids"]["tts"] == [30, 31] - assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5]) + assert torch.equal(get_tts_handoff(buffer)[1], hidden[3:5]) def test_plain_chat_without_tts_markers_uses_assistant_span(self) -> None: buffer, hidden = self._run([10, 11], [20, 21, 22]) assert buffer["ids"]["tts"] == [20, 21, 22] - assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[2:5]) + assert torch.equal(get_tts_handoff(buffer)[1], hidden[2:5]) class TestPromptAndMultiModal: @@ -325,7 +326,7 @@ def test_reference_audio_is_added_to_stage_handoff(self) -> None: ) intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"] - assert intermediate_buffer["codes"]["ref"] == [0.25, -0.5] + assert unpack_transport_tensor(intermediate_buffer["codes"]["ref"]).tolist() == [0.25, -0.5] assert intermediate_buffer["meta"]["ref_audio_sr"] == 16000 def test_serving_reference_audio_is_added_to_stage_handoff(self) -> None: @@ -338,7 +339,7 @@ def test_serving_reference_audio_is_added_to_stage_handoff(self) -> None: ) intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"] - assert intermediate_buffer["codes"]["ref"] == [0.125, -0.25] + assert unpack_transport_tensor(intermediate_buffer["codes"]["ref"]).tolist() == [0.125, -0.25] assert intermediate_buffer["meta"]["ref_audio_sr"] == 24000 def test_missing_reference_audio_does_not_add_stage_handoff_fields(self) -> None: diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py index 47b0f32c15c..a91a2f8c250 100644 --- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py @@ -6,6 +6,7 @@ import pytest import torch +from vllm_omni.engine.duplex.intermediate import get_tts_handoff, unpack_transport_tensor from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import ( _extract_first_audio_ref, llm2tts, @@ -70,7 +71,7 @@ def test_plain_chat_handoff_owns_talker_prompt_contract() -> None: info = converted["model_intermediate_buffer"] assert info["ids"]["tts"] == output_ids - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[2:4]) + assert torch.equal(get_tts_handoff(info)[1], latent[2:4]) assert converted["prompt_token_ids"] == [0, 0, 0, 0] assert info["meta"]["replace_streaming_prompt"] is True assert info["meta"]["next_stage_prompt_len"] == 4 @@ -97,7 +98,7 @@ def test_llm2tts_carries_request_ref_audio() -> None: )[0] info = converted["model_intermediate_buffer"] - assert info["codes"]["ref"] == ref_waveform.tolist() + assert unpack_transport_tensor(info["codes"]["ref"]).tolist() == ref_waveform.tolist() assert info["meta"]["ref_audio_sr"] == 22050 assert info["ids"]["tts"] == [11, 12] @@ -178,7 +179,7 @@ def test_native_duplex_uses_forwarded_row_ledger_after_later_audio_prefill() -> info = converted["model_intermediate_buffer"] assert info["ids"]["tts"] == [21, 22] - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + assert torch.equal(get_tts_handoff(info)[1], latent[3:5]) def test_native_duplex_selects_latest_contiguous_forwarded_span() -> None: @@ -211,7 +212,7 @@ def test_native_duplex_selects_latest_contiguous_forwarded_span() -> None: converted = llm2tts([source], prompt=[{}], _streaming_context=context)[0] info = converted["model_intermediate_buffer"] - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[5:7]) + assert torch.equal(get_tts_handoff(info)[1], latent[5:7]) def test_native_duplex_rejects_missing_forwarded_terminal_token() -> None: @@ -290,7 +291,7 @@ def test_native_duplex_ledger_tolerates_async_lookahead_row_after_terminator() - info = _native_handoff(source) assert info["ids"]["tts"] == [21, 22] - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + assert torch.equal(get_tts_handoff(info)[1], latent[3:5]) def test_native_duplex_ledger_without_forwarded_terminator() -> None: @@ -304,7 +305,7 @@ def test_native_duplex_ledger_without_forwarded_terminator() -> None: info = _native_handoff(source) assert info["ids"]["tts"] == [21, 22] - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + assert torch.equal(get_tts_handoff(info)[1], latent[3:5]) def test_native_duplex_turn_eos_row_is_handed_to_talker() -> None: @@ -320,7 +321,7 @@ def test_native_duplex_turn_eos_row_is_handed_to_talker() -> None: info = _native_handoff(source) assert info["ids"]["tts"] == [21, 9310] - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + assert torch.equal(get_tts_handoff(info)[1], latent[3:5]) assert info["meta"]["turn_end"] is True @@ -338,7 +339,7 @@ def test_native_duplex_ledger_anchors_on_unit_after_reinjected_listens() -> None info = _native_handoff(source) assert info["ids"]["tts"] == [21, 22] - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[9:11]) + assert torch.equal(get_tts_handoff(info)[1], latent[9:11]) @pytest.mark.parametrize("previous_terminator_forwarded", [False, True]) @@ -359,7 +360,7 @@ def test_native_duplex_ledger_prefers_latest_repeat_of_the_unit( info = _native_handoff(source) assert info["ids"]["tts"] == [21, 22] - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[current_start + 1 : current_start + 3]) + assert torch.equal(get_tts_handoff(info)[1], latent[current_start + 1 : current_start + 3]) def test_native_duplex_ledger_rejects_non_contiguous_positions() -> None: @@ -384,7 +385,7 @@ def test_native_duplex_mid_turn_tts_bos_slices_after_boundary() -> None: info = _native_handoff(source) assert info["ids"]["tts"] == [21, 22] - assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + assert torch.equal(get_tts_handoff(info)[1], latent[3:5]) @pytest.mark.parametrize("folded_decisions", [0, 1, 2]) @@ -420,7 +421,7 @@ def test_native_duplex_tts_bos_aligns_after_window_rebuild(folded_decisions, bos converted = llm2tts([source], prompt=[{}], _streaming_context=SimpleNamespace(bridge_states={}))[0] info = converted["model_intermediate_buffer"] assert info["ids"]["tts"] == [21, 22] - torch.testing.assert_close(torch.as_tensor(info["hidden_states"]["tts"]), latent[-3:-1]) + torch.testing.assert_close(get_tts_handoff(info)[1], latent[-3:-1]) def test_native_duplex_continuation_appends_only_new_talker_condition() -> None: diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_transport.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_transport.py new file mode 100644 index 00000000000..0ff0d330dfc --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_transport.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Host-side transport of the MiniCPM-o 4.5 stage handoffs. + +The bridges must hand the Talker and Code2Wav bit-identical values while +keeping per-handoff work off the orchestrator thread. +""" + +import base64 +from collections import defaultdict +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from vllm.v1.request import RequestStatus + +from vllm_omni.engine.duplex.intermediate import get_tts_handoff, pack_transport_tensor, unpack_transport_tensor +from vllm_omni.model_executor.models.minicpmo_4_5.duplex import input as duplex_input +from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import ( + _extract_native_runtime_ref_audio, + llm2tts, + tts2code2wav_async_chunk, + tts2code2wav_full_payload, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_SPECIAL = { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, +} + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + return tensor.contiguous().view(torch.int32) + + +def _native_source(output_ids: list[int], request_id: str = "req-1"): + prompt_ids = [101, 102] + row_ids = [*prompt_ids, *output_ids] + latent = torch.randn(len(row_ids), 8) + completion = SimpleNamespace( + token_ids=output_ids, + text="hello", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": prompt_ids, + "latent_input_ids": torch.tensor(row_ids).reshape(-1, 1), + "latent_positions": torch.arange(len(row_ids)).reshape(-1, 1), + "meta": dict(_SPECIAL), + }, + ) + return SimpleNamespace(request_id=request_id, prompt_token_ids=prompt_ids, outputs=[completion]), latent + + +def _ref_config(samples: np.ndarray, sample_rate: int = 16000) -> dict: + return { + "ref_audio_data": base64.b64encode(samples.astype(" SimpleNamespace: + return SimpleNamespace(bridge_states={"duplex": {"epoch": 3, "model_turn_id": 7, "runtime_config": runtime_config}}) + + +@pytest.fixture() +def decode_calls(monkeypatch): + calls = [] + real_decode = duplex_input.decode_native_ref_audio_from_config + + def counting_decode(session_config): + calls.append(session_config) + return real_decode(session_config) + + monkeypatch.setattr(duplex_input, "decode_native_ref_audio_from_config", counting_decode) + return calls + + +def test_native_duplex_handoff_is_packed_bit_exact() -> None: + source, latent = _native_source([9304, 21, 22, 9308]) + + info = llm2tts([source], prompt=[{}], _streaming_context=_context({}))[0]["model_intermediate_buffer"] + + packed = info["hidden_states"]["tts"] + assert isinstance(packed["data"], bytes) + token_ids, hidden = get_tts_handoff(info) + assert token_ids == [21, 22] + assert hidden.dtype == torch.float32 + # Same values the legacy tolist()/as_tensor transport produced. + assert torch.equal(_bits(hidden), _bits(torch.as_tensor(latent[3:5].tolist(), dtype=torch.float32))) + assert info["meta"]["next_stage_prompt_len"] == 3 + + +def test_reference_voice_is_decoded_once_per_session_config(decode_calls) -> None: + samples = np.linspace(-1.0, 1.0, 1600, dtype=np.float32) + context = _context(_ref_config(samples)) + + refs = [] + output_ids: list[int] = [] + for unit in range(3): + # Each unit extends the session's cumulative Stage-0 output. + output_ids = [*output_ids, 9304, 21 + unit, 40 + unit, 9308] + source, _ = _native_source(output_ids) + info = llm2tts([source], prompt=[{}], _streaming_context=context)[0]["model_intermediate_buffer"] + assert info["ids"]["tts"] == [21 + unit, 40 + unit] + refs.append(unpack_transport_tensor(info["codes"]["ref"])) + assert info["meta"]["ref_audio_sr"] == 16000 + + assert len(decode_calls) == 1 + for ref in refs: + assert torch.equal(_bits(ref), _bits(torch.from_numpy(samples))) + + # A new voice (or rate) is a new key: decode again, never serve stale audio. + reversed_samples = samples[::-1].copy() + context.bridge_states["duplex"]["runtime_config"] = _ref_config(reversed_samples, sample_rate=24000) + source, _ = _native_source([*output_ids, 9304, 30, 31, 9308]) + info = llm2tts([source], prompt=[{}], _streaming_context=context)[0]["model_intermediate_buffer"] + assert len(decode_calls) == 2 + assert torch.equal(_bits(unpack_transport_tensor(info["codes"]["ref"])), _bits(torch.from_numpy(reversed_samples))) + assert info["meta"]["ref_audio_sr"] == 24000 + + +def test_listen_only_unit_skips_reference_voice_decode(decode_calls) -> None: + samples = np.ones(160, dtype=np.float32) + source, _ = _native_source([9303]) + + assert llm2tts([source], prompt=[{}], _streaming_context=_context(_ref_config(samples))) == [] + assert decode_calls == [] + + +def test_reference_voice_without_bridge_state_decodes_every_call(decode_calls) -> None: + metadata = {"runtime_config": _ref_config(np.full(160, 0.5, dtype=np.float32))} + + first = _extract_native_runtime_ref_audio(metadata, None) + second = _extract_native_runtime_ref_audio(metadata, SimpleNamespace(bridge_states=None)) + + assert len(decode_calls) == 2 + assert torch.equal(first[0], second[0]) and first[1] == second[1] == 16000 + + +def test_reference_voice_path_is_still_rejected_on_every_call(decode_calls) -> None: + config = {**_ref_config(np.ones(16, dtype=np.float32)), "ref_audio_path": "/tmp/voice.wav"} + context = _context(config) + + for _ in range(2): + with pytest.raises(ValueError, match="ref_audio_path"): + _extract_native_runtime_ref_audio({"runtime_config": config}, context) + assert len(decode_calls) == 2 + + +def _manager(): + return SimpleNamespace( + connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 3}}), + code_prompt_token_ids=defaultdict(list), + request_payload={}, + put_req_chunk=defaultdict(int), + ) + + +def _talker_request(ref: object): + request = SimpleNamespace(external_req_id="req", request_id="req", status=RequestStatus.RUNNING) + request.is_finished = lambda: RequestStatus.is_finished(request.status) + request.model_intermediate_buffer = {"codes": {"ref": ref}, "meta": {"ref_audio_sr": 16000}} + return request + + +@pytest.mark.parametrize("packed", [True, False]) +def test_code2wav_first_chunk_reads_packed_or_list_reference(packed) -> None: + ref = torch.randn(64) + request = _talker_request(pack_transport_tensor(ref) if packed else ref.tolist()) + delta = {"codes": {"audio": torch.arange(7, dtype=torch.long).reshape(-1, 1)}, "meta": {}} + + payload = tts2code2wav_async_chunk(_manager(), delta, request, True) + + assert payload is not None + assert torch.equal(_bits(payload.codes.ref), _bits(ref)) + assert payload.meta.ref_audio_sr == 16000 + + +@pytest.mark.parametrize("packed", [True, False]) +def test_code2wav_full_payload_reads_packed_or_list_reference(packed) -> None: + ref = torch.randn(64) + request = _talker_request(pack_transport_tensor(ref) if packed else ref.tolist()) + + payload = tts2code2wav_full_payload( + _manager(), + {"codes.audio": torch.arange(7, dtype=torch.long).reshape(-1, 1)}, + request, + ) + + assert torch.equal(_bits(payload.codes.ref), _bits(ref)) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index 88afc919f36..5ea5cded45f 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -4,6 +4,9 @@ from __future__ import annotations +import functools +import math +import time from collections import OrderedDict from collections.abc import Mapping from contextlib import nullcontext @@ -15,16 +18,35 @@ import torch.nn.functional as F from vllm.logger import init_logger +from vllm_omni.utils.device_copy import index_to_device + from .cuda_graph_wrapper import ( CFMGraphWrapper, HiFTGraphWrapper, + ResidentAttCache, + SharedPromptAttCache, WholeEulerCFMGraphWrapper, _att_keep_ranges, _copy_frame_segments, + _device_used_bytes, _euler_step, _euler_timeline, + _format_memory_delta, + _materialize_att_rows, + _memory_snapshot, _zero_padded_cnn_cache, + empty_hift_outputs, +) +from .dit_fused import ( + DitGemm, + blocks_forward_chunk_fused, + causal_cache_index, + dit_modulation, + gather_causal_cache, + supports_fused_body, + tiled_attention_supported, ) +from .flow_encoder_graph import FlowEncoderGraphs logger = init_logger(__name__) @@ -37,10 +59,25 @@ # Streaming cache trim (stepaudio2 ``Token2wav.stream``): past ``prompt_len + # 100`` frames, keep the first ``prompt_len`` frames and the last 100. _CACHE_TRIM_SUFFIX = 100 +_TIMESTEP_MAX_PERIOD = 10000 + + +def _timestep_frequencies(dim: int) -> torch.Tensor: + """Host frequency table of ``TimestepEmbedder.timestep_embedding``. + + Built with the same ops, in the same order, as the ``cosyvoice2`` DiT + embedder, so moving it to the device once gives the table that embedder + would otherwise rebuild and copy on every call. + """ + half = dim // 2 + return torch.exp(-math.log(_TIMESTEP_MAX_PERIOD) * torch.arange(start=0, end=half) / half) def _trim_streaming_cache(cache: torch.Tensor, prompt_len: int) -> torch.Tensor: """The streaming trim of a cache whose frame axis is ``-2``.""" + if isinstance(cache, (ResidentAttCache, SharedPromptAttCache)): + # Its slot-pool solve already trimmed it (``WholeEulerCFMGraphWrapper._replay_slots``). + return cache ranges = _att_keep_ranges(int(cache.shape[-2]), (prompt_len, _CACHE_TRIM_SUFFIX)) if len(ranges) == 1: return cache @@ -207,6 +244,25 @@ def state_shape_signature(state: BatchedToken2WavState) -> tuple[Any, ...]: return flow, hift +# Cache frame axes a row-offset merge (``row_offset_merge``) keeps per row: the +# encoder runs once per conformer cache length (``decode_ragged_batch``), and a +# slot-pool Whole-Euler solve attends each estimator cache in its own slot. +_ROW_OFFSET_FRAME_AXES = {"conformer_att_cache": 3, "estimator_att_cache": 4} + + +def row_offset_signature(state: BatchedToken2WavState) -> tuple[Any, ...]: + """``state_shape_signature`` without the cache lengths a row-offset merge keeps per row.""" + flow = [] + for name in sorted(state.flow_cache): + shape, dtype, device = tensor_signature(state.flow_cache[name]) + axis = _ROW_OFFSET_FRAME_AXES.get(name) + if axis is not None: + shape = (*shape[:axis], -1, *shape[axis + 1 :]) + flow.append((name, (shape, dtype, device))) + hift = tuple((name, tensor_signature(state.hift_cache[name])) for name in sorted(state.hift_cache)) + return tuple(flow), hift + + @dataclass(frozen=True) class PromptFeatures: cache_key: tuple[str, str] @@ -240,6 +296,27 @@ def _undecorate_dynamo(module: nn.Module, method: str) -> None: logger.info("Bypassed TorchDynamo wrapper on %s.%s", type(module).__name__, method) +def _cfm_matmul_tf32(fn): + """Run the flow CFM with TF32 matmuls when the backend asks for it (``cfm_tf32``). + + Only the flow DiT is scoped: HiFT and everything else keep fp32. CFM + graphs are captured inside this call, so they are captured with TF32. + """ + + @functools.wraps(fn) + def wrapper(self, *args, **kwargs): + if not getattr(self, "_cfm_tf32", False): + return fn(self, *args, **kwargs) + previous = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = True + try: + return fn(self, *args, **kwargs) + finally: + torch.backends.cuda.matmul.allow_tf32 = previous + + return wrapper + + class BatchedToken2Wav(nn.Module): """Drive Token2wav's modules with dynamically-sized, request-owned caches. @@ -258,11 +335,23 @@ def __init__( cfm_graph_config: Mapping[str, Any] | None = None, bfloat16_attention_cache: bool = False, setup_cache_size: int = 1, + cfm_tf32: bool = False, + encoder_graph_config: Mapping[str, Any] | None = None, ): super().__init__() if setup_cache_size < 0: raise ValueError("setup_cache_size must be >= 0") self._token2wav = token2wav + self._cfm_tf32 = bool(cfm_tf32) + cfm_graph_cfg = dict(cfm_graph_config or {}) + self._prompt_att_sharing = bool(cfm_graph_cfg.get("prompt_att_sharing", False)) + self._prompt_att_caches: dict[tuple[str, str], torch.Tensor] = {} + self._prompt_att_verified: set[tuple[str, str]] = set() + self._prompt_att_rejected: set[tuple[str, str]] = set() + # Ragged solves of rows whose estimator caches differ in length (a + # stream's second chunk next to steady ones): each row keeps its own + # cache offset, in one slot-pool Whole-Euler replay when it can. + self._row_offset_merge = bool(cfm_graph_cfg.get("row_offset_merge", False)) # Optional TrtDiTStepper (step_audio2_dit_trt): replaces only the # per-timestep DiT estimator call; encoder and HiFT stay on torch. self._trt_stepper = trt_stepper @@ -274,6 +363,7 @@ def __init__( encoder = getattr(self.flow, "encoder", None) if encoder is not None: _undecorate_dynamo(encoder, "forward_chunk") + self._encoder_graphs = self._build_encoder_graphs(dict(encoder_graph_config or {}), connector_config) hift_parameter = next(self.hift.parameters(), None) if hift_parameter is not None and hift_parameter.device.type == "cuda": # Prime the CUDA state used by HiFT during backend construction. @@ -336,12 +426,14 @@ def __init__( capture_batch_sizes=capture_batch_sizes, max_serial_batch=max_serial_batch, ) - with torch.inference_mode(), _autocast_disabled(hift_parameter.device): - self.hift_graph_wrapper.capture() - logger.info("HiFT CUDA Graph captured successfully") + # Captured by the first forward (``precapture_hift``): graphs + # captured while vLLM loads the weights hold far more memory. self._cfm_graph_wrapper: CFMGraphWrapper | None = None self._whole_euler_graph_wrapper: WholeEulerCFMGraphWrapper | None = None - cfm_graph_cfg = dict(cfm_graph_config or {}) + # Whether the ragged DiT body is the fused one (``dit_fused.py``). + self._ragged_fused_body = False + # The fused body's GEMM backend (``cfm_dit_gemm``; ``None``: cuBLAS). + self._dit_gemm: DitGemm | None = None # Graph capture can be provided by the platform instead of the CUDA # wrapper above (on NPU it is the platform graph runner), so keep the # request flag separate: bucketing and padding key off this. @@ -356,7 +448,51 @@ def __init__( max_graph_batch = int(max_graph_batch_cfg) if max_graph_batch_cfg is not None else None micro_batch_size_cfg = cfm_graph_cfg.get("micro_batch_size") micro_batch_size = int(micro_batch_size_cfg) if micro_batch_size_cfg is not None else 4 - if bool(cfm_graph_cfg.get("enable_whole_euler", True)) and self._trt_stepper is None: + enable_whole_euler = bool(cfm_graph_cfg.get("enable_whole_euler", True)) + if bool(cfm_graph_cfg.get("fused_body", False)): + # A reduced-precision attention cache keeps the ragged body, + # which attends in the activation dtype through upstream. + self._ragged_fused_body = ( + supports_fused_body(estimator) and self._estimator_att_cache_dtype == torch.float32 + ) + if not self._ragged_fused_body: + logger.warning( + "cfm_fused_body requested but unsupported here (DiT blocks or a non-fp32 attention cache)" + ) + dit_gemm = cfm_graph_cfg.get("dit_gemm") + precision = cfm_graph_cfg.get("dit_gemm_precision") or None + if precision not in (None, "tf32", "tf32x3"): + raise ValueError(f"cfm_dit_gemm_precision must be tf32 or tf32x3, got {precision!r}") + if dit_gemm not in (None, "", "cublas"): + if dit_gemm not in ("triton", "triton_windows"): + raise ValueError(f"cfm_dit_gemm must be cublas, triton or triton_windows, got {dit_gemm!r}") + if self._ragged_fused_body and tiled_attention_supported(flow_parameter.get_device()): + self._dit_gemm = DitGemm(conv_windows=dit_gemm == "triton_windows", precision=precision) + logger.info( + "CFM fused body GEMMs on Triton tf32_linear (%s, precision=%s)", + dit_gemm, + precision or "follows the TF32 scope", + ) + else: + logger.warning("cfm_dit_gemm=%s needs cfm_fused_body on NVIDIA SM80+; using cuBLAS", dit_gemm) + elif precision is not None: + logger.warning("cfm_dit_gemm_precision=%s needs cfm_dit_gemm=triton; ignored", precision) + # Resident request caches need the fused body and its tiled + # attention, which reads keys/values out of the pool in place. + slot_pool = bool(cfm_graph_cfg.get("slot_pool", False)) + if slot_pool and self._prompt_att_sharing: + # Slot rows are mutable full-cache storage. Keep the + # shared-prefix path on the arena, whose input copy can + # address prompt and tail separately, until a slot layout + # with shared immutable rows is available. + logger.info("Disabling CFM slot pool while prompt attention sharing is enabled") + slot_pool = False + if slot_pool and not ( + self._ragged_fused_body and tiled_attention_supported(flow_parameter.get_device()) + ): + logger.warning("cfm_slot_pool requested but needs cfm_fused_body on NVIDIA SM80+; using the arena") + slot_pool = False + if enable_whole_euler and self._trt_stepper is None: capture_frames = cfm_graph_cfg.get("capture_frames") if capture_frames is None and connector_config is not None: # The steady chunk's mel width, derived like HiFTGraphWrapper's @@ -365,7 +501,12 @@ def __init__( capture_frames = int(connector_config["codec_chunk_frames"]) * int( getattr(self.flow, "token_mel_ratio", 2) ) - query_bucket_frames = int(capture_frames if capture_frames is not None else 64) + if isinstance(capture_frames, (list, tuple)): + # Several widths (e.g. a 25-token duplex unit and a + # 75-token turn chunk) each get their own graphs. + query_bucket_frames: int | tuple[int, ...] = tuple(int(w) for w in capture_frames) + else: + query_bucket_frames = int(capture_frames if capture_frames is not None else 64) self._whole_euler_graph_wrapper = WholeEulerCFMGraphWrapper( estimator=estimator, n_timesteps=self.n_timesteps, @@ -377,24 +518,37 @@ def __init__( micro_batch_size=micro_batch_size, query_bucket_frames=query_bucket_frames, pad_max_rows=cfm_graph_cfg.get("pad_max_rows"), - # Looked up at capture time, so a replaced ragged kernel is the one captured. - ragged_body=( - (lambda *args: self._blocks_forward_chunk_ragged(*args)) - if _supports_ragged_kernel(estimator) + offset_bucket_frames=int(cfm_graph_cfg.get("offset_bucket_frames", 0)), + # Resident request caches, a slot per stream plus a few for streams + # that end while new ones start (``AttSlotPool``). + att_slots=(micro_batch_size + max(2, micro_batch_size // 8)) if slot_pool else 0, + # ``_ragged_body`` picks the configured body on every call, so a + # replaced body is the one a capture records. + ragged_body=self._ragged_body if _supports_ragged_kernel(estimator) else None, + # The fused body takes each timestep's modulation, computed once outside the graphs. + modulation_fn=( + (lambda time_embedding: dit_modulation(estimator, time_embedding)) + if self._ragged_fused_body else None ), + arena_rows_from_graph_grid=bool(cfm_graph_cfg.get("arena_rows_from_graph_grid", False)), + row_offsets=self._row_offset_merge, ) logger.info( "Whole-Euler CFM CUDA Graph enabled " "(max_graphs=%d, max_serial_batch=%d, max_graph_batch=%s, " - "micro_batch_size=%d, query_bucket_frames=%d)", + "micro_batch_size=%d, arena_rows=%d, query_bucket_frames=%s, offset_bucket_frames=%d, " + "fused_body=%s)", max_graphs, max_serial_batch, str(max_graph_batch), micro_batch_size, - query_bucket_frames, + self._whole_euler_graph_wrapper._arena_rows, + str(query_bucket_frames), + self._whole_euler_graph_wrapper.offset_bucket_frames, + self._ragged_fused_body, ) - elif self._trt_stepper is not None and bool(cfm_graph_cfg.get("enable_whole_euler", True)): + elif enable_whole_euler: logger.info("Whole-Euler CFM CUDA Graph disabled because TensorRT stepper is configured") self._cfm_graph_wrapper = CFMGraphWrapper( graph_fn=estimator.blocks_forward_chunk, @@ -424,6 +578,8 @@ def __init__( "CFM CUDA Graph bucketing enabled (bucket_frames=%d)", self._cfm_graph_bucket_frames, ) + # Device copies of the DiT timestep-embedding frequency table. + self._timestep_freqs: dict[tuple[Any, ...], torch.Tensor] = {} self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} self._setup_cache_size = setup_cache_size self._setup_cache: OrderedDict[tuple[tuple[str, str], int], BatchedToken2WavState] = OrderedDict() @@ -433,6 +589,9 @@ def _hift_inference( mel: torch.Tensor, source_cache: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + if int(mel.shape[2]) == 0: + # No frames, no samples: HiFT itself cannot run on an empty mel. + return empty_hift_outputs(mel) if self.hift_graph_wrapper is None: return self.hift.inference(mel, source_cache) return self.hift_graph_wrapper.replay(mel, source_cache) @@ -464,6 +623,9 @@ def evict_prompt(self, prompt_cache_id: str, prompt_wav: str) -> None: """Release all cached artifacts associated with one prompt.""" prompt_key = (prompt_cache_id, prompt_wav) self._prompt_features.pop(prompt_key, None) + self._prompt_att_caches.pop(prompt_key, None) + self._prompt_att_verified.discard(prompt_key) + self._prompt_att_rejected.discard(prompt_key) for setup_key in list(self._setup_cache): if setup_key[0] == prompt_key: self._setup_cache.pop(setup_key) @@ -476,6 +638,19 @@ def _repeat_prompt(features: PromptFeatures, batch_size: int) -> tuple[torch.Ten features.mels.expand(batch_size, -1, -1), ) + def _project_speakers(self, speakers: torch.Tensor) -> torch.Tensor: + """``spk_embed_affine_layer(normalize(speakers))``, independent of the batch size. + + Every row of one call holds the prompt's speaker (``expand`` of one + row), so one row is projected and expanded. The GEMM then always runs + with M=1: an M=1 and an M>=2 GEMM round differently (~6e-8), and that + moved a row's CFM mel by ~1e-3 whenever its decode batch changed size. + """ + if speakers.shape[0] > 1 and speakers.stride(0) == 0: + projected = self.flow.spk_embed_affine_layer(F.normalize(speakers[:1], dim=1)) + return projected.expand(speakers.shape[0], -1).contiguous() + return self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + def _autocast(self, device: torch.device): if device.type != "cuda": return nullcontext() @@ -547,6 +722,155 @@ def _encode_chunk( ) return self.flow.encoder_proj(hidden), new_cnn, new_att + def _encode_continuation( + self, + tokens: torch.Tensor, + cnn_cache: torch.Tensor, + att_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``_encode_chunk`` of a continuation chunk (both caches, not the last): what an encoder graph records.""" + hidden, new_cnn, new_att = self.flow.encoder.forward_chunk( + xs=self.flow.input_embedding(tokens), + last_chunk=False, + cnn_cache=cnn_cache, + att_cache=att_cache, + ) + return self.flow.encoder_proj(hidden), new_cnn, new_att + + def _flow_on_cuda(self) -> bool: + flow_parameter = next(self.flow.parameters(), None) + return flow_parameter is not None and flow_parameter.device.type == "cuda" + + @staticmethod + def _default_encoder_token_widths(connector_config: Mapping[str, Any] | None) -> list[int]: + """The steady duplex chunk's codec tokens: the left context plus one unit of new frames.""" + connector = dict(connector_config or {}) + frames = int(connector.get("initial_codec_chunk_frames", 0) or 0) or int( + connector.get("codec_chunk_frames", 25) + ) + return [int(connector.get("codec_left_context_frames", 3)) + frames] + + def _build_encoder_graphs( + self, + config: dict[str, Any], + connector_config: Mapping[str, Any] | None, + ) -> FlowEncoderGraphs | None: + """``cfm_encoder_cuda_graph``: exact-shape graphs of the flow encoder's continuation chunk (default off).""" + if not config.get("enabled", False): + return None + encoder = getattr(self.flow, "encoder", None) + if not self._flow_on_cuda(): + logger.info("Flow encoder CUDA Graph is disabled off CUDA") + return None + embed = getattr(encoder, "embed", None) + up_embed = getattr(encoder, "up_embed", None) + lookahead = self._pre_lookahead_len() + if ( + not callable(getattr(encoder, "forward_chunk", None)) + or lookahead is None + or getattr(encoder, "up_layer", None) is None + or not isinstance(getattr(getattr(embed, "pos_enc", None), "pe", None), torch.Tensor) + # ``_ensure_relpos_pe`` would grow such a table between calls. + or callable(getattr(embed, "extend_pe", None)) + ): + logger.warning("Flow encoder CUDA Graph needs the CosyVoice2 upsample-conformer encoder; staying eager") + return None + widths = config.get("token_widths") or self._default_encoder_token_widths(connector_config) + rows = config.get("rows") or [1] + + def relpos_tables() -> tuple[torch.Tensor, ...]: + # Every ``position_encoding`` slice a graph reads lives in these tables. + return tuple( + module.pos_enc.pe + for module in (embed, up_embed) + if isinstance(getattr(getattr(module, "pos_enc", None), "pe", None), torch.Tensor) + ) + + graphs = FlowEncoderGraphs( + self._encode_continuation, + rows=rows, + token_widths=widths, + lookahead=lookahead, + upsample=self._upsample_stride(), + held_tensors=relpos_tables, + ) + logger.info( + "Flow encoder CUDA Graph enabled for continuation chunks (rows %s, tokens %s), captured at first forward", + list(graphs.rows), + list(graphs.token_widths), + ) + return graphs + + def _graph_encode( + self, + tokens: torch.Tensor, + states: list[BatchedToken2WavState], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """A continuation chunk's encoder outputs from its captured graph, or ``None`` (run it eager). + + The outputs are the graphs' shared result views, rewritten by the next + encoder replay. + """ + graphs = getattr(self, "_encoder_graphs", None) + if graphs is None: + return None + return graphs.run( + tokens, + [state.flow_cache["conformer_cnn_cache"] for state in states], + [state.flow_cache["conformer_att_cache"] for state in states], + ) + + @torch.inference_mode() + def precapture_flow_encoder(self, features: PromptFeatures) -> int: + """Capture the encoder graphs this prompt's streams replay (``cfm_encoder_cuda_graph``). + + Every continuation chunk of the configured rows and token widths, at + each conformer cache length from the prompt state to the trimmed steady + ``prompt + 100`` frames. Returns the number of graphs captured. + """ + graphs = getattr(self, "_encoder_graphs", None) + if graphs is None or graphs.graphs: + return 0 + (state,) = self.setup_batch(features, 1) + cnn = state.flow_cache["conformer_cnn_cache"] + att = state.flow_cache["conformer_att_cache"] + keys = graphs.keys_for( + start=int(att.shape[3]), + prompt_len=int(features.mels.shape[1]), + suffix=_CACHE_TRIM_SUFFIX, + ) + device = att.device + started = time.perf_counter() + if device.type == "cuda": + torch.accelerator.synchronize(device) + memory_before = _memory_snapshot(device) + used_before = _device_used_bytes(device) + with self._autocast(device): + captured = graphs.capture( + keys, + cnn_shape=(int(cnn.shape[1]), int(cnn.shape[2])), + att_layout=(int(att.shape[0]), int(att.shape[2]), int(att.shape[4])), + hidden_dim=int(self.flow.encoder_proj.out_features), + dtype=att.dtype, + device=device, + ) + if device.type == "cuda": + torch.accelerator.synchronize(device) + used_after = _device_used_bytes(device) + logger.info( + "Captured %d flow encoder CUDA Graphs (rows %s, tokens %s, cache frames %s) in %.1f s: " + "shared buffers %.1f MiB%s, device memory %s", + captured, + list(graphs.rows), + list(graphs.token_widths), + sorted({key[2] for key in keys}), + time.perf_counter() - started, + graphs.storage_bytes() / 2**20, + _format_memory_delta(memory_before, _memory_snapshot(device)), + "n/a" if used_before is None or used_after is None else f"+{(used_after - used_before) / 2**20:.1f} MiB", + ) + return captured + @staticmethod def _estimator_buffers( estimator: nn.Module, @@ -567,6 +891,34 @@ def _estimator_buffers( att = x.new_empty((depth, batch_size, heads, old_att_len + chunk_size, att_width)) return cnn, att + def _time_embedding(self, estimator: nn.Module, time: torch.Tensor) -> torch.Tensor: + """``estimator.t_embedder(time)`` without a host sync per call. + + ``TimestepEmbedder.timestep_embedding`` rebuilds its frequency table on + the host and moves it with a pageable ``.to(t)``, which blocks the host + until the queued GPU work drains -- once per CFM step. The table only + depends on the width, dtype and device, so it is built the same way + once and kept on the device; the remaining ops are the embedder's own, + so the result is bitwise identical. + """ + embedder = estimator.t_embedder + dim = getattr(embedder, "frequency_embedding_size", None) + scale = getattr(embedder, "scale", None) + mlp = getattr(embedder, "mlp", None) + if not isinstance(dim, int) or scale is None or mlp is None: + return embedder(time) + # The host table takes the default dtype before ``.to(t)``, as upstream. + key = (time.device, time.dtype, torch.get_default_dtype(), dim) + freqs = self._timestep_freqs.get(key) + if freqs is None: + freqs = _timestep_frequencies(dim).to(time) + self._timestep_freqs[key] = freqs + args = (time * scale)[:, None] * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return mlp(embedding) + def _estimator_step( self, estimator: nn.Module, @@ -581,7 +933,6 @@ def _estimator_step( attn_mask: torch.Tensor | None = None, valid_lengths: list[int] | None = None, valid_frames: int | None = None, - time_embedding: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if self._trt_stepper is not None and valid_lengths is None: out, new_cnn, new_att = self._trt_stepper.step( @@ -594,8 +945,7 @@ def _estimator_step( att_cache=att_cache, ) return out.to(mu.dtype), new_cnn, new_att - if time_embedding is None: - time_embedding = estimator.t_embedder(time).unsqueeze(1) + time_embedding = self._time_embedding(estimator, time).unsqueeze(1) width = int(x.shape[-1]) speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) if valid_frames is not None and valid_frames < width: @@ -630,7 +980,7 @@ def _estimator_step( if valid_lengths is not None: if not hasattr(estimator, "in_proj"): raise RuntimeError('MiniCPMO45Code2WavBatchError {"reason":"ragged_kernel_unavailable"}') - result = self._blocks_forward_chunk_ragged( + result = self._ragged_body( estimator, estimator_input, time_embedding, @@ -653,15 +1003,21 @@ def _estimator_step( ) return result, cnn_out, att_out - @staticmethod - def _gather_causal_cache( - history: torch.Tensor, - valid_lengths: torch.Tensor, - width: int, + def _ragged_body( + self, + *args: Any, + modulation: torch.Tensor | None = None, + slots: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: - """The ``width`` frames of ``history`` (b, width + t, c) ending each row's valid frames, as (b, c, width).""" - indices = valid_lengths[:, None] + torch.arange(width, device=history.device)[None, :] - return history.gather(1, indices[:, :, None].expand(-1, -1, int(history.shape[2]))).transpose(1, 2) + """The configured ragged DiT body; ``modulation`` (fused body only) comes from the Whole-Euler arena. + + Looked up per call (and so per capture), so a replaced body is the one a graph records. + """ + estimator_input = args[1] + if slots is not None or (self._ragged_fused_body and estimator_input.dtype == torch.float32): + # Slot-pool graphs are only built over the fused body (``slot_pool`` above). + return blocks_forward_chunk_fused(*args, modulation=modulation, slots=slots, gemm=self._dit_gemm) + return self._blocks_forward_chunk_ragged(*args) @staticmethod def _causal_conv_frames_major(conv: nn.Conv1d, history: torch.Tensor) -> torch.Tensor: @@ -774,6 +1130,8 @@ def _blocks_forward_chunk_ragged( device=estimator_input.device, dtype=torch.long, ) + # Causal cache gather index per conv width, shared by the blocks. + cache_index: dict[int, torch.Tensor] = {} x = estimator.in_proj(estimator_input.transpose(1, 2)) for block_index, block in enumerate(estimator.blocks): ( @@ -812,11 +1170,13 @@ def _blocks_forward_chunk_ragged( ) ) width = int(old_cnn.shape[2]) + if width not in cache_index: + cache_index[width] = causal_cache_index(lengths, width) history = torch.cat((old_cnn[:, : conv.in_channels].transpose(1, 2), conv_input), dim=1) - new_cnn1 = BatchedToken2Wav._gather_causal_cache(history, lengths, width) + new_cnn1 = gather_causal_cache(history, cache_index[width]) hidden = conv.block[4](conv.block[3](BatchedToken2Wav._causal_conv_frames_major(conv.block[1], history))) history = torch.cat((old_cnn[:, conv.in_channels :].transpose(1, 2), hidden), dim=1) - new_cnn2 = BatchedToken2Wav._gather_causal_cache(history, lengths, width) + new_cnn2 = gather_causal_cache(history, cache_index[width]) x = x + gate_conv * BatchedToken2Wav._causal_conv_frames_major(conv.block[6], history) x = x + gate_mlp * block.mlp(block.norm2(x) * (1 + scale_mlp) + shift_mlp) @@ -838,6 +1198,7 @@ def _get_timeline(self, device: torch.device, dtype: torch.dtype) -> tuple[torch self._timeline_cache[key] = cached return cached + @_cfm_matmul_tf32 def _decode_cfm( self, mu: torch.Tensor, @@ -848,6 +1209,8 @@ def _decode_cfm( att_cache: torch.Tensor | list[torch.Tensor] | None, valid_lengths: list[int] | None = None, att_keep: tuple[int, int] | None = None, + prompt_att_cache: torch.Tensor | None = None, + prompt_cache_key: tuple[str, str] | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -868,6 +1231,13 @@ def _decode_cfm( offset = int(att_rows[0].shape[4]) else: offset = int(att_cache.shape[4]) if att_cache is not None else 0 + # ``row_offset_merge``: ragged rows whose caches differ in length each + # start from the noise at their own offset and attend their own cache. + row_offsets = None + if self._row_offset_merge and att_rows is not None and valid_lengths is not None: + row_offsets = [int(row.shape[4]) for row in att_rows] + if len(set(row_offsets)) == 1: + row_offsets = None mel_frames = int(mu.shape[2]) # Padding only pays off while replay is active: whole-Euler replay, # or the CFM wrapper (`_disable` keeps the object alive, so check its @@ -895,14 +1265,18 @@ def _decode_cfm( # carry less of a step change into the attention and CNN caches. mu = torch.nn.functional.pad(mu, (0, pad_frames), mode="replicate") cond = torch.nn.functional.pad(cond, (0, pad_frames), mode="replicate") - end = offset + int(mu.shape[2]) + end = (offset if row_offsets is None else max(row_offsets)) + int(mu.shape[2]) if end > int(decoder.rand_noise.shape[2]): raise RuntimeError( "MiniCPMO45Code2WavBatchError " f'{{"reason":"noise_capacity","required":{end},' f'"available":{int(decoder.rand_noise.shape[2])}}}' ) - x = decoder.rand_noise[:, :, offset:end].expand(batch_size, -1, -1).clone() + if row_offsets is None: + x = decoder.rand_noise[:, :, offset:end].expand(batch_size, -1, -1).clone() + else: + width = int(mu.shape[2]) + x = torch.cat([decoder.rand_noise[:, :, start : start + width] for start in row_offsets], dim=0) if pad_frames: # The padded columns would otherwise carry real noise values. They # are excluded from this chunk's attention by ``attn_mask`` below @@ -918,18 +1292,21 @@ def _decode_cfm( if valid_lengths is not None: if len(valid_lengths) != batch_size: raise ValueError(f"valid length count {len(valid_lengths)} != batch {batch_size}") - cfg_lengths = torch.tensor( - (*valid_lengths, *valid_lengths), - device=mu.device, - ) + cfg_lengths = index_to_device((*valid_lengths, *valid_lengths), mu.device) positions = torch.arange(int(mu.shape[2]), device=mu.device) valid_queries = positions.unsqueeze(0) < cfg_lengths.unsqueeze(1) current_keys = valid_queries.unsqueeze(1).expand(-1, int(mu.shape[2]), -1) - old_keys = torch.ones( - (2 * batch_size, int(mu.shape[2]), offset), - dtype=torch.bool, - device=mu.device, - ) + if row_offsets is None: + old_keys = torch.ones( + (2 * batch_size, int(mu.shape[2]), offset), + dtype=torch.bool, + device=mu.device, + ) + else: + # Cache columns span the longest cache; a shorter one's tail stays masked. + cfg_offsets = index_to_device((*row_offsets, *row_offsets), mu.device) + cached = torch.arange(max(row_offsets), device=mu.device).unsqueeze(0) < cfg_offsets.unsqueeze(1) + old_keys = cached.unsqueeze(1).expand(-1, int(mu.shape[2]), -1) attn_mask = valid_queries.unsqueeze(2) & torch.cat((current_keys, old_keys), dim=2) elif pad_frames: # Mask the padded keys instead of only zeroing their content: a @@ -961,16 +1338,41 @@ def _decode_cfm( att_keep=att_keep, ) if whole_euler_result is not None: + if ( + prompt_att_cache is not None + and prompt_cache_key is not None + and isinstance(whole_euler_result[2], list) + ): + shared_rows = self._share_prompt_att_rows( + prompt_cache_key, + prompt_att_cache, + whole_euler_result[2], + ) + whole_euler_result = ( + whole_euler_result[0], + whole_euler_result[1], + shared_rows, + ) return whole_euler_result + if row_offsets is not None: + # No single solve took them (graphs off, or the slot pool refused). + return self._decode_cfm_per_offset( + mu, + speakers, + cond, + cnn_cache=cnn_cache, + att_rows=att_rows, + valid_lengths=valid_lengths, + att_keep=att_keep, + prompt_att_cache=prompt_att_cache, + prompt_cache_key=prompt_cache_key, + ) if att_rows is not None: att_cache = self._stack_estimator_att(att_rows) next_cnn: list[torch.Tensor] = [] next_att_cache: torch.Tensor | None = None ragged_att_cache: list[torch.Tensor] | None = None - time_embeddings = [ - estimator.t_embedder(timeline[s].expand(2 * batch_size)).unsqueeze(1) for s in range(self.n_timesteps) - ] with _token2wav_sdpa_context(mu.device): for step in range(self.n_timesteps): old_cnn = cnn_cache[step] if cnn_cache is not None else None @@ -988,7 +1390,6 @@ def _decode_cfm( attn_mask=attn_mask, valid_lengths=valid_lengths, valid_frames=mel_frames if pad_frames else None, - time_embedding=time_embeddings[step], ) if pad_frames: _zero_padded_cnn_cache(step_cnn, estimator, pad_frames) @@ -1046,12 +1447,60 @@ def _decode_cfm( assert next_att_cache is not None return x, torch.stack(next_cnn), next_att_cache + def _decode_cfm_per_offset( + self, + mu: torch.Tensor, + speakers: torch.Tensor, + cond: torch.Tensor, + *, + cnn_cache: torch.Tensor | None, + att_rows: list, + valid_lengths: list[int], + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor, list]: + """``_decode_cfm`` once per estimator cache length, reassembled in row order. + + The fallback of a row-offset merge (``row_offset_merge``) that no + single solve took; each group then decodes as its own bucket would. + """ + batch_size = len(att_rows) + groups: dict[int, list[int]] = {} + for row, cache in enumerate(att_rows): + groups.setdefault(int(cache.shape[4]), []).append(row) + x: torch.Tensor | None = None + cnn: torch.Tensor | None = None + att: list = [None] * batch_size + for rows in groups.values(): + index = index_to_device(rows, mu.device) + # CFG rows are ``[cond x B | uncond x B]``. + cfg_index = index_to_device([*rows, *(batch_size + row for row in rows)], mu.device) + group_x, group_cnn, group_att = self._decode_cfm( + mu.index_select(0, index), + speakers.index_select(0, index), + cond.index_select(0, index), + cnn_cache=None if cnn_cache is None else cnn_cache.index_select(2, cfg_index), + att_cache=[att_rows[row] for row in rows], + valid_lengths=[valid_lengths[row] for row in rows], + **kwargs, + ) + if x is None or cnn is None: + x = group_x.new_empty((batch_size, *group_x.shape[1:])) + cnn = group_cnn.new_empty((*group_cnn.shape[:2], 2 * batch_size, *group_cnn.shape[3:])) + x.index_copy_(0, index, group_x) + cnn.index_copy_(2, cfg_index, group_cnn) + for row, cache in zip(rows, group_att, strict=True): + att[row] = cache + assert x is not None and cnn is not None + return x, cnn, att + def _split_flow_cache( self, cache: dict[str, Any], batch_size: int, *, estimator_att_keep: tuple[int, int] | None = None, + prompt_att_cache: torch.Tensor | None = None, + prompt_cache_key: tuple[str, str] | None = None, ) -> list[dict[str, torch.Tensor]]: """Per-request flow caches from a batch's stacked caches. @@ -1094,11 +1543,55 @@ def _split_flow_cache( "estimator_att_cache": request_att, } ) + if stacked_att and prompt_att_cache is not None and prompt_cache_key is not None: + shared_rows = self._share_prompt_att_rows( + prompt_cache_key, + prompt_att_cache, + [row["estimator_att_cache"] for row in result], + ) + for row, shared in zip(result, shared_rows, strict=True): + row["estimator_att_cache"] = shared return result + def _share_prompt_att_rows( + self, + cache_key: tuple[str, str], + prompt: torch.Tensor, + rows: list[torch.Tensor | SharedPromptAttCache], + ) -> list[torch.Tensor | SharedPromptAttCache]: + """Keep the immutable prompt suffix after a bit-exact suffix check.""" + if not self._prompt_att_sharing or cache_key in self._prompt_att_rejected: + return rows + if any(isinstance(row, SharedPromptAttCache) for row in rows): + return rows + shared_len = min(_CACHE_TRIM_SUFFIX, int(prompt.shape[4])) + if any( + tuple(prompt.shape[:4]) != tuple(row.shape[:4]) + or prompt.shape[5] != row.shape[5] + or int(row.shape[4]) < shared_len + for row in rows + ): + self._prompt_att_rejected.add(cache_key) + return rows + prompt_segment = prompt[..., -shared_len:, :] + if cache_key not in self._prompt_att_verified: + if any(not torch.equal(row[..., -shared_len:, :], prompt_segment) for row in rows): + logger.warning("CFM prompt attention suffix differs; disabling sharing for %s", cache_key[0]) + self._prompt_att_rejected.add(cache_key) + return rows + self._prompt_att_verified.add(cache_key) + logger.info( + "CFM prompt attention suffix verified bit-exact (prompt=%s rows=%d frames=%d)", + cache_key[0], + len(rows), + shared_len, + ) + return [SharedPromptAttCache(row[..., :-shared_len, :].detach().clone(), prompt_segment) for row in rows] + @staticmethod def _stack_estimator_att(rows: list[torch.Tensor]) -> torch.Tensor: """Stack per-request ``(n_t, depth, 2, ...)`` caches into CFG order ``[cond x B | uncond x B]``.""" + rows = _materialize_att_rows(rows) return torch.cat((*[row[:, :, 0:1] for row in rows], *[row[:, :, 1:2] for row in rows]), dim=2) def _stack_flow_cache( @@ -1106,19 +1599,26 @@ def _stack_flow_cache( states: list[BatchedToken2WavState], *, include_estimator_att: bool = True, + include_conformer: bool = True, ) -> dict[str, torch.Tensor | None]: """Stack per-request flow caches for one batched call. The estimator attention cache is by far the largest (~0.5 GiB per request in fp32). The Whole-Euler path reads it per request, so it - skips stacking it (``include_estimator_att=False``). + skips stacking it (``include_estimator_att=False``). A caller that + already ran the encoder skips the conformer caches + (``include_conformer=False``). """ flows = [state.flow_cache for state in states] conditional_cnn = [flow["estimator_cnn_cache"][:, :, 0:1] for flow in flows] unconditional_cnn = [flow["estimator_cnn_cache"][:, :, 1:2] for flow in flows] return { - "conformer_cnn_cache": torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0), - "conformer_att_cache": torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1), + "conformer_cnn_cache": ( + torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0) if include_conformer else None + ), + "conformer_att_cache": ( + torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1) if include_conformer else None + ), "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2), "estimator_att_cache": ( self._stack_estimator_att([flow["estimator_att_cache"] for flow in flows]) @@ -1135,6 +1635,88 @@ def _whole_euler_ragged_active(self) -> bool: """Whether Whole-Euler graphs also solve ragged (per-row length) batches.""" return self._whole_euler_active() and self._whole_euler_graph_wrapper.ragged_body is not None + def can_merge_state_shapes(self, states: list[BatchedToken2WavState], token_counts: list[int]) -> bool: + """Whether mixed onset/continuation rows fit ragged Whole-Euler.""" + if len(states) < 2 or len(states) != len(token_counts) or not self._whole_euler_ragged_active(): + return False + if len({tuple(state.flow_cache["estimator_cnn_cache"].shape) for state in states}) != 1: + return False + estimator_att = [state.flow_cache.get("estimator_att_cache") for state in states] + if any(not isinstance(cache, torch.Tensor) for cache in estimator_att): + return False + if len({(tuple(cache.shape[:4]), cache.dtype, cache.device) for cache in estimator_att}) != 1: + return False + return all( + count <= self._max_encode_token_frames([state]) for count, state in zip(token_counts, states, strict=True) + ) + + def can_merge_row_offsets(self, states: list[BatchedToken2WavState], token_counts: list[int]) -> bool: + """Whether continuation rows whose caches differ only in length fit one slot-pool ragged solve. + + ``row_offset_merge``: every estimator cache must be resident in the + Whole-Euler slot pool (``ResidentAttCache``), where each row attends its + own frames; ``decode_ragged_batch`` runs the encoder once per conformer + cache length. + """ + if ( + not self._row_offset_merge + or len(states) < 2 + or len(states) != len(token_counts) + or not self._whole_euler_ragged_active() + ): + return False + pool = self._whole_euler_graph_wrapper.slot_pool + if pool is None: + return False + for state in states: + cache = state.flow_cache.get("estimator_att_cache") + if not isinstance(cache, ResidentAttCache) or cache.pool is not pool: + return False + if len({tuple(state.flow_cache["estimator_cnn_cache"].shape) for state in states}) != 1: + return False + return all( + count <= self._max_encode_token_frames([state]) for count, state in zip(token_counts, states, strict=True) + ) + + @torch.inference_mode() + def precapture_hift(self) -> int: + """Capture every configured HiFT CUDA Graph bucket (prompt-independent) before serving. + + Called from the first forward; a second call is a no-op. Returns the + number of graphs captured. + """ + wrapper = self.hift_graph_wrapper + if wrapper is None or wrapper.graph: + return 0 + with _autocast_disabled(wrapper.device): + wrapper.capture() + return len(wrapper.graph) + + @torch.inference_mode() + @_cfm_matmul_tf32 + def precapture_whole_euler(self, features: PromptFeatures) -> int: + """Capture the Whole-Euler graphs this prompt's streams will replay. + + Every cache length a stream passes through, from the prompt to the + steady ``prompt + 100`` frames, on the wrapper's offset grid. Runs in + the precision ``_decode_cfm`` captures in (``_cfm_matmul_tf32``, + ``_autocast``), since a graph replays the kernels it was captured with. + Returns the number of graphs captured. + """ + if not self._whole_euler_active(): + return 0 + wrapper = self._whole_euler_graph_wrapper + prompt_len = int(features.mels.shape[1]) + steady = prompt_len + _CACHE_TRIM_SUFFIX + with self._autocast(wrapper.device): + return wrapper.precapture( + offsets=range(prompt_len, steady + 1), + steady=steady, + channels=int(features.mels.shape[2]), + spk_dim=int(self.flow.spk_embed_affine_layer.out_features), + keep=(prompt_len, _CACHE_TRIM_SUFFIX), + ) + def _create_initial_states( self, features: PromptFeatures, @@ -1153,7 +1735,7 @@ def _create_initial_states( cnn_cache=None, att_cache=None, ) - projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + projected_speakers = self._project_speakers(speakers) # The prompt cache never reaches the trim length; ``att_keep`` only # tells the Whole-Euler arena the steady cache length to reserve. _, estimator_cnn, estimator_att = self._decode_cfm( @@ -1207,8 +1789,12 @@ def setup_batch(self, features: PromptFeatures, batch_size: int) -> list[Batched state = self._setup_cache.get(cache_key) if state is not None: self._setup_cache.move_to_end(cache_key) + if self._prompt_att_sharing: + self._prompt_att_caches[features.cache_key] = state.flow_cache["estimator_att_cache"] else: (state,) = self._create_initial_states(features, 1) + if self._prompt_att_sharing: + self._prompt_att_caches[features.cache_key] = state.flow_cache["estimator_att_cache"] if self._setup_cache_size > 0: self._setup_cache[cache_key] = state while len(self._setup_cache) > self._setup_cache_size: @@ -1315,18 +1901,28 @@ def _decode_batch_once( # Whole-Euler reads and writes each request's estimator cache directly, # so the batch never holds a stacked copy of it. per_request_att = self._whole_euler_active() - flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_request_att) + # A captured encoder graph stacks the conformer caches into its own inputs. + graphed = None if last_chunk or flush_encoder else self._graph_encode(tokens, states) + flow_cache = self._stack_flow_cache( + states, include_estimator_att=not per_request_att, include_conformer=graphed is None + ) prompt_len = int(features.mels.shape[1]) att_keep = (prompt_len, _CACHE_TRIM_SUFFIX) + prompt_att_cache = self._prompt_att_caches.get(features.cache_key) speakers = features.speaker_embedding.expand(batch_size, -1) with self._autocast(tokens.device): - hidden, conformer_cnn, conformer_att = self._encode_chunk( - tokens, - last_chunk=last_chunk or flush_encoder, - cnn_cache=flow_cache["conformer_cnn_cache"], - att_cache=flow_cache["conformer_att_cache"], - ) - projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + if graphed is None: + hidden, conformer_cnn, conformer_att = self._encode_chunk( + tokens, + last_chunk=last_chunk or flush_encoder, + cnn_cache=flow_cache["conformer_cnn_cache"], + att_cache=flow_cache["conformer_att_cache"], + ) + else: + # Shared graph results: read below (mu, cond) and copied per row by + # _split_flow_cache before any other encoder call. + hidden, conformer_cnn, conformer_att = graphed + projected_speakers = self._project_speakers(speakers) cond = torch.zeros_like(hidden).transpose(1, 2).contiguous() chunk_mel, estimator_cnn, estimator_att = self._decode_cfm( hidden.transpose(1, 2).contiguous(), @@ -1339,6 +1935,8 @@ def _decode_batch_once( else flow_cache["estimator_att_cache"] ), att_keep=att_keep, + prompt_att_cache=prompt_att_cache, + prompt_cache_key=features.cache_key, ) conformer_att = _trim_streaming_cache(conformer_att, prompt_len) @@ -1351,6 +1949,8 @@ def _decode_batch_once( }, batch_size, estimator_att_keep=att_keep, + prompt_att_cache=prompt_att_cache, + prompt_cache_key=features.cache_key, ) old_mel = torch.cat([state.hift_cache["mel"] for state in states], dim=0) old_source = torch.cat([state.hift_cache["source"] for state in states], dim=0) @@ -1446,25 +2046,41 @@ def decode_ragged_batch( audios[row] = group_audio[group_row] next_states[row] = group_states[group_row] return self._require_complete_ragged_outputs(audios, next_states) - encoder_groups: dict[tuple[int, bool], list[int]] = {} - for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): - encoder_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row) + encoder_groups: dict[tuple[Any, ...], list[int]] = {} + for row, (row_tokens, state, last_chunk) in enumerate(zip(tokens, states, last_chunks, strict=True)): + # Onset and continuation rows may have different conformer cache + # lengths. Group those encoder calls separately; estimator caches + # remain per-request rows in the ragged Whole-Euler solve. + key = ( + tuple(state.flow_cache["conformer_cnn_cache"].shape), + tuple(state.flow_cache["conformer_att_cache"].shape), + int(row_tokens.numel()), + last_chunk, + ) + encoder_groups.setdefault(key, []).append(row) hidden_rows: list[torch.Tensor | None] = [None] * batch_size conformer_cnn_rows: list[torch.Tensor | None] = [None] * batch_size conformer_att_rows: list[torch.Tensor | None] = [None] * batch_size - for (_, last_chunk), rows in encoder_groups.items(): + for (*_, last_chunk), rows in encoder_groups.items(): group_states = [states[row] for row in rows] - # The encoder reads only the conformer caches. - group_cache = self._stack_flow_cache(group_states, include_estimator_att=False) group_tokens = torch.stack([tokens[row] for row in rows], dim=0) - with self._autocast(group_tokens.device): - hidden, conformer_cnn, conformer_att = self._encode_chunk( - group_tokens, - last_chunk=last_chunk, - cnn_cache=group_cache["conformer_cnn_cache"], - att_cache=group_cache["conformer_att_cache"], - ) + graphed = None if last_chunk else self._graph_encode(group_tokens, group_states) + if graphed is not None: + hidden, conformer_cnn, conformer_att = graphed + if len(encoder_groups) > 1: + # The next group's replay rewrites the shared graph results. + hidden, conformer_cnn, conformer_att = (value.clone() for value in graphed) + else: + # The encoder reads only the conformer caches. + group_cache = self._stack_flow_cache(group_states, include_estimator_att=False) + with self._autocast(group_tokens.device): + hidden, conformer_cnn, conformer_att = self._encode_chunk( + group_tokens, + last_chunk=last_chunk, + cnn_cache=group_cache["conformer_cnn_cache"], + att_cache=group_cache["conformer_att_cache"], + ) for group_row, row in enumerate(rows): hidden_rows[row] = hidden[group_row : group_row + 1] conformer_cnn_rows[row] = conformer_cnn[group_row : group_row + 1] @@ -1484,11 +2100,22 @@ def decode_ragged_batch( # A graphed ragged solve reads and writes each request's estimator # cache directly (trimmed on the way out), like _decode_batch_once. - flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_request_att) + # Rows of a row-offset merge hold caches of different lengths, which + # never stack: their solve takes them per row as well. The encoder + # groups above already read the conformer caches, whose lengths such + # rows also mix. + per_row_att = per_request_att or ( + self._row_offset_merge + and len({int(state.flow_cache["estimator_att_cache"].shape[4]) for state in states}) > 1 + ) + flow_cache = self._stack_flow_cache( + states, include_estimator_att=not per_row_att, include_conformer=not self._row_offset_merge + ) prompt_len = int(features.mels.shape[1]) + prompt_att_cache = self._prompt_att_caches.get(features.cache_key) speakers = features.speaker_embedding.expand(batch_size, -1) with self._autocast(padded_hidden.device): - projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + projected_speakers = self._project_speakers(speakers) cond = torch.zeros_like(padded_hidden).transpose(1, 2).contiguous() chunk_mel, estimator_cnn, estimator_att = self._decode_cfm( padded_hidden.transpose(1, 2).contiguous(), @@ -1497,11 +2124,13 @@ def decode_ragged_batch( cnn_cache=flow_cache["estimator_cnn_cache"], att_cache=( [state.flow_cache["estimator_att_cache"] for state in states] - if per_request_att + if per_row_att else flow_cache["estimator_att_cache"] ), valid_lengths=hidden_lengths, att_keep=(prompt_len, _CACHE_TRIM_SUFFIX), + prompt_att_cache=prompt_att_cache, + prompt_cache_key=features.cache_key, ) assert isinstance(estimator_att, list) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py b/vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py new file mode 100644 index 00000000000..1c2af229666 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Tiled FP32 attention for the short streaming CFM windows on NVIDIA CUDA.""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _attention( + q, + k, + v, + mask, + kv_rows, + out, + qb: tl.constexpr, + qh: tl.constexpr, + qt: tl.constexpr, + qd: tl.constexpr, + kb: tl.constexpr, + kh: tl.constexpr, + kt: tl.constexpr, + kd: tl.constexpr, + vb: tl.constexpr, + vh: tl.constexpr, + vt: tl.constexpr, + vd: tl.constexpr, + mb: tl.constexpr, + mt: tl.constexpr, + mk: tl.constexpr, + heads: tl.constexpr, + nq: tl.constexpr, + nk: tl.constexpr, + dim: tl.constexpr, + has_mask: tl.constexpr, + has_kv_rows: tl.constexpr, + scale: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_d: tl.constexpr, +): + b, h, tile = tl.program_id(0), tl.program_id(1), tl.program_id(2) + # Keys and values of a pooled cache live in row ``kv_rows[b]`` of the pool. + kv_b = b + if has_kv_rows: + kv_b = tl.load(kv_rows + b).to(tl.int64) + rows = tile * block_m + tl.arange(0, block_m) + ds = tl.arange(0, block_d) + cols = tl.arange(0, block_n) + queries = tl.load( + q + b * qb + h * qh + rows[:, None] * qt + ds[None, :] * qd, (rows[:, None] < nq) & (ds[None, :] < dim), 0 + ) + maximum = tl.full((block_m,), -float("inf"), tl.float32) + denominator = tl.full((block_m,), 0, tl.float32) + accumulator = tl.full((block_m, block_d), 0, tl.float32) + for start in range(tl.cdiv(nk, block_n)): + keys = start * block_n + cols + kval = tl.load( + k + kv_b * kb + h * kh + keys[None, :] * kt + ds[:, None] * kd, + (keys[None, :] < nk) & (ds[:, None] < dim), + 0, + ) + score = tl.dot(queries, kval, input_precision="tf32x3") * scale + valid = (rows[:, None] < nq) & (keys[None, :] < nk) + if has_mask: + keep = tl.load(mask + b * mb + rows[:, None] * mt + keys[None, :] * mk, valid, 0) + valid = valid & (keep != 0) + score = tl.where(valid, score, -float("inf")) + new_max = tl.maximum(maximum, tl.max(score, 1)) + # Fully masked rows produce zero, matching PyTorch SDPA. + # Keep the running maximum at -inf until a valid tile arrives. A + # permanent zero would underflow a later tile with very negative logits. + safe_max = tl.where(new_max == -float("inf"), 0.0, new_max) + correction = tl.exp(maximum - safe_max) + prob = tl.exp(score - safe_max[:, None]) + denominator = denominator * correction + tl.sum(prob, 1) + values = tl.load( + v + kv_b * vb + h * vh + keys[:, None] * vt + ds[None, :] * vd, + (keys[:, None] < nk) & (ds[None, :] < dim), + 0, + ) + accumulator = accumulator * correction[:, None] + tl.dot(prob, values, input_precision="tf32x3") + maximum = new_max + result = accumulator / tl.where(denominator > 0, denominator, 1.0)[:, None] + tl.store( + out + ((b * nq + rows[:, None]) * heads + h) * dim + ds[None, :], + result, + (rows[:, None] < nq) & (ds[None, :] < dim), + ) + + +@torch.library.custom_op("vllm_omni::cfm_tiled_attention", mutates_args=()) +def cfm_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + mask: torch.Tensor | None = None, + kv_rows: torch.Tensor | None = None, +) -> torch.Tensor: + """FP32 Q/K/V, bool [B,Q,K] mask, [B,H,Q,D] result (BQHD storage). + + Uses three TF32 products per FP32 dot; no BF16 casts or cache changes. + Numerical equivalence is tolerance-based, not bitwise. With ``kv_rows`` + (``[B]`` integers) K/V are ``[R,H,K,D]`` pools and query row ``b`` + attends to pool row ``kv_rows[b]``. + """ + b, heads, nq, dim = q.shape + nk = k.shape[2] + if q.dtype != torch.float32 or k.dtype != torch.float32 or v.dtype != torch.float32: + raise ValueError("CFM tiled attention requires float32 Q/K/V") + # Strided FP32 tiles can exceed the 99 KiB per-block limit on L4. + # Measured on A800 (SM80) across the deployed shapes (2..32 rows, Q<=150, + # K<=524, D=64): two stages beat three everywhere, e.g. 265 vs 356 us at + # (16, 8, 150, 524); the extra stage only paid for SRAM capacity, not + # latency. + num_stages = 2 + out = torch.empty((b, nq, heads, dim), device=q.device, dtype=q.dtype) + _attention[(b, heads, triton.cdiv(nq, 32))]( + q, + k, + v, + mask, + kv_rows, + out, + *q.stride(), + *k.stride(), + *v.stride(), + *(mask.stride() if mask is not None else (0, 0, 0)), + heads, + nq, + nk, + dim, + mask is not None, + kv_rows is not None, + dim**-0.5, + 32, + 64, + triton.next_power_of_2(dim), + num_warps=4, + num_stages=num_stages, + ) + return out.transpose(1, 2) + + +@cfm_attention.register_fake +def _cfm_attention_fake(q, k, v, mask=None, kv_rows=None): + b, heads, nq, dim = q.shape + return q.new_empty((b, nq, heads, dim)).transpose(1, 2) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py b/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py new file mode 100644 index 00000000000..c88cf399de8 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused forward of the Code2Wav CFM DiT blocks (Stage 2). + +The estimator's block is ``attention -> causal conv block -> MLP``, each behind +an adaLN-modulated LayerNorm and a gated residual. The ragged body +(``BatchedToken2Wav._blocks_forward_chunk_ragged``) runs it op by op: about 47 +kernels per block, most of them full passes over the activation (the +modulation, the gates, transposes, ``torch.cat`` of the conv histories, and the +layout conversions cuDNN adds around every convolution). This body computes +the same function with about 15: + +* The modulation depends only on the Euler timestep (``c = t_embedder(t)``, + identical for every row), so the shift/scale/gate vectors of all blocks are + computed once per timestep (``dit_modulation``) instead of by 16 SiLU+GEMM + pairs over every row on every step. +* ``gated residual -> LayerNorm -> (1 + scale), shift`` is one pass + (``residual_layer_norm``); the modulation is the LayerNorm affine. +* q/k/v come from one packed GEMM; the q/k LayerNorms and the key/value + writes into the attention cache are one pass (``qkv_head_layer_norm``). +* Each 3-tap causal convolution is one GEMM over its history with the taps + stacked on the output axis; its tap sum and bias are reduced by the pass + that follows it (``LayerNorm + Mish`` after the first, the gated residual + and the MLP's LayerNorm after the second). The pass before each convolution + writes straight into the history buffer behind its cached frames. +* The MLP's GELU (tanh) runs in the GEMM epilogue. +* Attention is the tiled fp32 (tf32x3) kernel (``cfm_attention.py``) on + NVIDIA SM80+, fp32-equivalent and faster than SDPA at these short windows; + SDPA elsewhere. + +Everything is fp32 storage with the matmul precision the caller configured +(the ``flow`` TF32 scope in serving), so results match the ragged body to fp32 +rounding, not bitwise: reductions and the tap sums associate differently. + +With ``gemm=DitGemm(...)`` (``cfm_dit_gemm``) the GEMMs run on +``tf32_linear`` (Triton) instead of cuBLAS, with a tile shape and a +deterministic split-K picked per row count (``dit_gemm_config``) so the small +streaming batches fill the device; TF32 products under the ``flow`` scope, +three TF32 products (fp32-level) outside it, or as ``DitGemm.precision`` says. ``conv_windows`` also computes +each causal convolution as one GEMM over its overlapping im2col windows of the +history (``conv_window_view``) instead of the tap-stacked GEMM, so the GEMM +writes ``C`` instead of ``K * C`` columns over ``T`` instead of ``K - 1 + T`` +frames and the following pass reads a third of the data. +""" + +from __future__ import annotations + +import functools +from typing import Any, NamedTuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .ops import qkv_head_layer_norm, residual_layer_norm + +# adaLN chunk order of ``DiTBlock.adaLN_modulation``. +_SHIFT_MSA, _SCALE_MSA, _GATE_MSA, _SHIFT_MLP, _SCALE_MLP, _GATE_MLP, _SHIFT_CONV, _SCALE_CONV, _GATE_CONV = range(9) + + +class DitGemm(NamedTuple): + """The fused body's GEMM backend (``cfm_dit_gemm``; default cuBLAS).""" + + conv_windows: bool = False + precision: str | None = None + + +def _linear( + a: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + gemm: DitGemm | None = None, + *, + activation: str | None = None, + view=None, +) -> torch.Tensor: + """``act(A @ weight.T + bias)`` on cuBLAS.""" + if activation == "gelu_tanh": + if bias is not None and a.is_cuda: + # cuBLASLt applies bias + tanh-GELU in the GEMM epilogue. + return torch._addmm_activation(bias, a, weight.t(), use_gelu=True) + return F.gelu(F.linear(a, weight, bias), approximate="tanh") + return F.linear(a, weight, bias) + + +def supports_fused_body(estimator: nn.Module) -> bool: + """Whether ``blocks_forward_chunk_fused`` computes this estimator's blocks.""" + blocks = getattr(estimator, "blocks", None) + final_layer = getattr(estimator, "final_layer", None) + if not blocks or not hasattr(estimator, "in_proj") or final_layer is None: + return False + if not isinstance(getattr(final_layer, "norm_final", None), nn.LayerNorm) or not hasattr(final_layer, "linear"): + return False + for block in blocks: + attn = getattr(block, "attn", None) + conv_block = getattr(getattr(block, "conv", None), "block", None) + mlp = getattr(block, "mlp", None) + if attn is None or conv_block is None or mlp is None: + return False + if not all(hasattr(attn, name) for name in ("to_q", "to_k", "to_v", "proj", "num_heads", "head_dim")): + return False + if not (isinstance(attn.q_norm, nn.LayerNorm) and isinstance(attn.k_norm, nn.LayerNorm)): + return False + if not all(isinstance(getattr(block, name, None), nn.LayerNorm) for name in ("norm1", "norm2", "norm3")): + return False + if len(conv_block) < 7 or not isinstance(conv_block[3], nn.LayerNorm) or not isinstance(conv_block[4], nn.Mish): + return False + for conv in (conv_block[1], conv_block[6]): + if not isinstance(conv, nn.Conv1d) or conv.stride != (1,) or conv.dilation != (1,) or conv.groups != 1: + return False + if conv.padding != (0,) or conv.kernel_size[0] < 2: + return False + if not (hasattr(mlp, "fc1") and hasattr(mlp, "fc2") and isinstance(mlp.act, nn.GELU)): + return False + if not isinstance(getattr(mlp, "norm", nn.Identity()), nn.Identity): + return False + return True + + +def dit_modulation(estimator: nn.Module, time_embedding: torch.Tensor) -> torch.Tensor: + """The adaLN vectors of every block and of the final layer for one timestep. + + ``time_embedding`` is one row, ``(1, 1, C)`` (every row of a solve shares + the timestep). Returns ``(depth + 1, 9, C)``: row ``i < depth`` is block + ``i``'s chunks in ``adaLN_modulation`` order with each scale replaced by + ``1 + scale``; row ``depth`` holds the final layer's ``shift`` and + ``1 + scale`` in its first two slots. + """ + embedding = time_embedding.reshape(1, 1, -1) + rows = [] + for block in estimator.blocks: + chunks = block.adaLN_modulation(embedding).reshape(9, -1) + rows.append(chunks) + table = torch.stack(rows, dim=0) + table[:, (_SCALE_MSA, _SCALE_MLP, _SCALE_CONV)] += 1.0 + final_shift, final_scale = estimator.final_layer.adaLN_modulation(embedding).reshape(2, -1) + final = table.new_zeros((1, 9, int(table.shape[2]))) + # Where a block keeps its MSA shift/scale, so the last block's closing LayerNorm reads it like a next block. + final[0, _SHIFT_MSA] = final_shift + final[0, _SCALE_MSA] = final_scale + 1.0 + return torch.cat((table, final), dim=0).contiguous() + + +def _cached(module: nn.Module, key: str, sources: tuple[torch.Tensor, ...], build) -> Any: + version = tuple((t.data_ptr(), t._version, t.dtype) for t in sources) + cached = module.__dict__.get(key) + if cached is None or cached[0] != version: + with torch.no_grad(): + cached = (version, build()) + module.__dict__[key] = cached + return cached[1] + + +def packed_qkv(attn: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]: + """``[to_q; to_k; to_v]`` weight and bias as one projection, rebuilt only when a source changes.""" + sources = (attn.to_q.weight, attn.to_k.weight, attn.to_v.weight) + + def build(): + weight = torch.cat(sources, dim=0).contiguous() + bias = None + if attn.to_q.bias is not None: + bias = torch.cat((attn.to_q.bias, attn.to_k.bias, attn.to_v.bias), dim=0).contiguous() + return weight, bias + + return _cached(attn, "_packed_qkv", sources, build) + + +def _conv_taps_weight(conv: nn.Conv1d) -> torch.Tensor: + """``(C_out, C_in, K)`` -> ``(K * C_out, C_in)``: tap ``k`` is output rows ``[k*C_out, (k+1)*C_out)``.""" + return _cached( + conv, + "_fused_taps_weight", + (conv.weight,), + lambda: conv.weight.detach().permute(2, 0, 1).reshape(-1, conv.in_channels).contiguous(), + ) + + + +def _conv_history( + old_cache: torch.Tensor | None, rows: int, width: int, frames: int, channels: int, like: torch.Tensor +) -> torch.Tensor: + """``(N, width + frames, C)``: the cached frames, then room for the current ones.""" + history = like.new_empty((rows, width + frames, channels)) + if old_cache is None: + history[:, :width].zero_() + else: + history[:, :width].copy_(old_cache.transpose(1, 2)) + return history + + +def causal_cache_index(lengths: torch.Tensor, width: int) -> torch.Tensor: + """``(N, width)`` history frames that end each row's valid frames, for ``gather_causal_cache``.""" + return lengths[:, None] + torch.arange(width, device=lengths.device)[None, :] + + +def gather_causal_cache(history: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """The ``index`` frames of ``history`` (N, width + T, C), as the (N, C, width) causal cache.""" + return history.gather(1, index[:, :, None].expand(-1, -1, int(history.shape[2]))).transpose(1, 2) + + +def _mlp(mlp: nn.Module, hidden: torch.Tensor, gemm: DitGemm | None = None) -> torch.Tensor: + rows = hidden.reshape(-1, int(hidden.shape[-1])) + if mlp.act.approximate == "tanh" and rows.is_cuda: + inner = _linear(rows, mlp.fc1.weight, mlp.fc1.bias, gemm, activation="gelu_tanh") + else: + inner = mlp.act(F.linear(rows, mlp.fc1.weight, mlp.fc1.bias)) + return _linear(inner, mlp.fc2.weight, mlp.fc2.bias, gemm).view(*hidden.shape[:-1], -1) + + +@functools.cache +def tiled_attention_supported(device_index: int) -> bool: + """``cfm_attention``'s tf32x3 dots need NVIDIA tensor cores with TF32 (SM80+).""" + return torch.version.hip is None and torch.cuda.get_device_capability(device_index) >= (8, 0) + + +def _attention( + attn: nn.Module, + hidden: torch.Tensor, + att_cache: torch.Tensor | None, + attn_mask: torch.Tensor | None, + kv: torch.Tensor, + slots: tuple[torch.Tensor, torch.Tensor] | None = None, + gemm: DitGemm | None = None, +) -> torch.Tensor: + """``attn.forward_chunk`` over ``hidden`` that writes the new keys/values into ``kv`` (``[current | cache]``). + + With ``slots = (rows, positions)`` ``kv`` is an ``AttSlotPool`` block + (``(2 * slots, heads, frames, 2 * head_dim)``): row ``n`` writes its + keys/values to frames ``positions[n]`` of pool row ``rows[n]`` and attends + to that row under ``attn_mask`` over the pool frames. Attention is the + tiled tf32x3 kernel (``cfm_attention``) on NVIDIA SM80+ and always for a + pool, SDPA otherwise. + """ + batch, frames, _ = hidden.shape + weight, bias = packed_qkv(attn) + qkv = _linear(hidden.reshape(batch * frames, -1), weight, bias, gemm).view(batch, frames, -1) + rows, positions = slots if slots is not None else (None, None) + if slots is None: + # ``att_cache`` behind the current frames: left where it is when it + # already sits there (the Whole-Euler arena's aliasing), else copied once. + cached = 0 if att_cache is None else int(att_cache.shape[2]) + kv = kv[:, :, : frames + cached] + behind = kv[:, :, frames:] + if cached and (att_cache.data_ptr() != behind.data_ptr() or att_cache.stride() != behind.stride()): + behind.copy_(att_cache) + q = qkv_head_layer_norm( + qkv, + kv, + num_heads=int(attn.num_heads), + head_dim=int(attn.head_dim), + q_norm=attn.q_norm, + k_norm=attn.k_norm, + rows=rows, + positions=positions, + ) + head_dim = int(q.shape[-1]) + keys, values = kv[..., :head_dim], kv[..., head_dim:] + if rows is not None or ( + kv.is_cuda and kv.dtype == q.dtype == torch.float32 and tiled_attention_supported(kv.get_device()) + ): + from .cfm_attention import cfm_attention + + out = cfm_attention(q, keys, values, attn_mask, rows) + else: + out = F.scaled_dot_product_attention( + q, keys, values, attn_mask=None if attn_mask is None else attn_mask.unsqueeze(1) + ) + out = out.transpose(1, 2).reshape(batch * frames, -1) + return _linear(out, attn.proj.weight, attn.proj.bias, gemm).view(batch, frames, -1) + + +def blocks_forward_chunk_fused( + estimator: nn.Module, + estimator_input: torch.Tensor, + time_embedding: torch.Tensor, + attn_mask: torch.Tensor | None, + cnn_cache: Any, + att_cache: Any, + cnn_cache_buffer: torch.Tensor, + att_cache_buffer: torch.Tensor, + valid_lengths: list[int] | torch.Tensor, + modulation: torch.Tensor | None = None, + slots: tuple[torch.Tensor, torch.Tensor] | None = None, + gemm: DitGemm | None = None, +) -> torch.Tensor: + """``BatchedToken2Wav._blocks_forward_chunk_ragged``, fused (see the module docstring). + + Same arguments and results; ``modulation`` is this timestep's + ``dit_modulation`` table (computed from ``time_embedding[:1]`` when + ``None``). ``slots`` attends to an ``AttSlotPool`` in place (see ``_attention``); + ``att_cache`` is then unused and ``att_cache_buffer`` is the pool. + ``gemm`` selects the GEMM backend (``DitGemm``; ``None``: cuBLAS). + """ + if isinstance(valid_lengths, torch.Tensor): + lengths = valid_lengths + else: + lengths = torch.tensor((*valid_lengths, *valid_lengths), device=estimator_input.device, dtype=torch.long) + if modulation is None: + modulation = dit_modulation(estimator, time_embedding[:1]) + blocks = estimator.blocks + depth = len(blocks) + # Causal cache gather index per conv width, shared by the blocks. + cache_index: dict[int, torch.Tensor] = {} + + x = estimator.in_proj(estimator_input.transpose(1, 2)).contiguous() + rows, frames, _ = x.shape + first = modulation[0] + hidden = residual_layer_norm(x, weight=first[_SCALE_MSA], bias=first[_SHIFT_MSA], eps=blocks[0].norm1.eps) + for block_index, block in enumerate(blocks): + mod = modulation[block_index] + x_att = _attention( + block.attn, + hidden, + att_cache[block_index], + attn_mask, + att_cache_buffer[block_index], + slots, + gemm, + ) + + # CausalConvBlock: [conv1, LayerNorm, Mish, conv2] on (N, width + T, C) histories. + conv_layers = block.conv.block + conv1, conv_norm, conv2 = conv_layers[1], conv_layers[3], conv_layers[6] + width = int(conv1.kernel_size[0]) - 1 + old_cnn = cnn_cache[block_index] + history = _conv_history( + None if old_cnn is None else old_cnn[:, : conv1.in_channels], + rows, + width, + frames, + conv1.in_channels, + x, + ) + # x += gate_msa * attention; the conv block's modulated LayerNorm lands in the history. + residual_layer_norm( + x, + x_att, + gate=mod[_GATE_MSA], + weight=mod[_SCALE_CONV], + bias=mod[_SHIFT_CONV], + eps=block.norm3.eps, + residual_out=x, + out=history[:, width:], + ) + taps1 = _linear(history.view(-1, conv1.in_channels), _conv_taps_weight(conv1), None, gemm) + history2 = _conv_history( + None if old_cnn is None else old_cnn[:, conv1.in_channels :], + rows, + width, + frames, + conv2.in_channels, + x, + ) + residual_layer_norm( + None, + taps1.view(rows, width + frames, -1), + taps=width + 1, + y_bias=conv1.bias, + weight=conv_norm.weight, + bias=conv_norm.bias, + eps=conv_norm.eps, + activation="mish", + out=history2[:, width:], + ) + taps2 = _linear(history2.view(-1, conv2.in_channels), _conv_taps_weight(conv2), None, gemm) + # x += gate_conv * conv2; the MLP's modulated LayerNorm. + mlp_in = residual_layer_norm( + x, + taps2.view(rows, width + frames, -1), + taps=width + 1, + y_bias=conv2.bias, + gate=mod[_GATE_CONV], + weight=mod[_SCALE_MLP], + bias=mod[_SHIFT_MLP], + eps=block.norm2.eps, + residual_out=x, + ) + mlp_out = _mlp(block.mlp, mlp_in, gemm) + # x += gate_mlp * MLP, then the next block's (or the final layer's) modulated LayerNorm: + # table row ``depth`` holds the final layer's shift and 1 + scale where a block keeps its MSA pair. + last = block_index + 1 == depth + hidden = residual_layer_norm( + x, + mlp_out, + gate=mod[_GATE_MLP], + weight=modulation[block_index + 1][_SCALE_MSA], + bias=modulation[block_index + 1][_SHIFT_MSA], + eps=estimator.final_layer.norm_final.eps if last else blocks[block_index + 1].norm1.eps, + residual_out=None if last else x, + ) + + if width not in cache_index: + cache_index[width] = causal_cache_index(lengths, width) + cnn_out = cnn_cache_buffer[block_index] + cnn_out[:, : conv1.in_channels].copy_(gather_causal_cache(history, cache_index[width])) + cnn_out[:, conv1.in_channels :].copy_(gather_causal_cache(history2, cache_index[width])) + + final = estimator.final_layer.linear + if gemm is None: + return final(hidden).transpose(1, 2) + return ( + _linear(hidden.reshape(rows * frames, -1), final.weight, final.bias, gemm) + .view(rows, frames, -1) + .transpose(1, 2) + ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 3f93f0f3021..dce3daa1e1a 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -7,13 +7,14 @@ import json import os import tempfile +import time from collections import OrderedDict from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from functools import lru_cache from hashlib import sha256 from pathlib import Path -from typing import Any +from typing import Any, cast import soundfile as sf import torch @@ -27,8 +28,10 @@ from .batched_token2wav import ( BatchedToken2Wav, BatchedToken2WavState, + row_offset_signature, state_shape_signature, ) +from .cuda_graph_wrapper import ResidentAttCache, _memory_snapshot logger = init_logger(__name__) @@ -47,6 +50,28 @@ def _resolve_model_dir(model_ref: str, revision: str | None = None) -> str: return hf_api().snapshot_download(model_ref, revision=revision, allow_patterns=["assets/*"]) +def _tf32_mode(extra: Mapping[str, Any]) -> str: + """Stage-2 TF32 scope: ``"off"`` (default), ``"flow"`` (the CFM DiT only) or ``"all"``. + + Set by the connector extra ``code2wav_allow_tf32`` (false / "flow" / true), + overridden by ``MINICPMO_CODE2WAV_TF32=0|flow|1``; ``token2wav_allow_tf32: + true`` (upstream's switch) always means ``"all"``. Off by default: on A800 + TF32 everywhere moves the log-mel by 0.7-1.4 dB (L1) against fp32, 10-30x + fp32's own seed-to-seed spread. ``"flow"`` keeps HiFT in fp32 and speeds up + the DiT GEMMs, the top Stage-2 kernel. + """ + if bool(extra.get("token2wav_allow_tf32", False)): + return "all" + raw = os.environ.get("MINICPMO_CODE2WAV_TF32") + value: Any = raw if raw not in (None, "") else extra.get("code2wav_allow_tf32", False) + if isinstance(value, str): + value = value.strip().lower() + if value == "flow": + return "flow" + return "all" if value in ("1", "true", "all", "yes") else "off" + return "all" if bool(value) else "off" + + def _batch_error(reason: str, **details: Any) -> RuntimeError: payload = {"reason": reason, **details} return RuntimeError(f"MiniCPMO45Code2WavBatchError {json.dumps(payload, sort_keys=True)}") @@ -225,6 +250,102 @@ class _WorkItem: has_payload: bool = True +def _resident_signature(item: _WorkItem) -> bool: + """Whether a continuation row's estimator attention cache lives in the slot pool (not a tensor).""" + if item.previous is None: + return False + cache = item.previous.token2wav.flow_cache.get("estimator_att_cache") + return cache is not None and not isinstance(cache, torch.Tensor) + + +def _slot_resident(item: _WorkItem) -> bool: + """Whether a continuation row's estimator attention cache is a Whole-Euler slot (``ResidentAttCache``).""" + if item.previous is None: + return False + return isinstance(item.previous.token2wav.flow_cache.get("estimator_att_cache"), ResidentAttCache) + + +class _BucketStats: + """``code2wav_bucket_stats`` (connector extra, default off): how the scheduler's rows split into decodes. + + One line per ``every`` forwards: row-count distribution, exact-shape buckets + and decode calls per forward, rows per decode, and why rows of one forward + did not share a bucket (each extra bucket is compared with the forward's + first: prompt, cache epoch, a fresh stream next to a continuation, or the + slot-pool / tensor cache signature of two continuations) or + why a mixed onset group or a row-offset bucket fell back to its + signature buckets. + """ + + def __init__(self, every: int = 200) -> None: + self.every = every + self._reset() + + def _reset(self) -> None: + self.forwards = 0 + self.rows: dict[int, int] = {} + self.buckets = 0 + self.decodes = 0 + self.decoded_rows = 0 + self.multi_row_forwards = 0 + self.multi_row_buckets = 0 + self.reasons: dict[str, int] = {} + + def split_reason(self, first: _WorkItem, other: _WorkItem, *, cross_turn: bool, row_offsets: bool = False) -> str: + if (first.prompt_cache_id, first.prompt_wav) != (other.prompt_cache_id, other.prompt_wav): + return "prompt" + if not cross_turn and first.cache_epoch != other.cache_epoch: + same_signature = MiniCPMO45Code2Wav._bucket_key( + first, cross_turn=True, row_offsets=row_offsets + ) == MiniCPMO45Code2Wav._bucket_key(other, cross_turn=True, row_offsets=row_offsets) + if same_signature: + return "cache_epoch" + if (first.previous is None) != (other.previous is None): + # A fresh stream next to a continuation: the onset split, whatever the cache layout. + return "onset" + if _resident_signature(first) or _resident_signature(other): + return "signature_resident" + return "signature_tensor" + + def reason(self, name: str) -> None: + self.reasons[name] = self.reasons.get(name, 0) + 1 + + def record_forward( + self, rows: int, buckets: list[list[_WorkItem]], *, cross_turn: bool, row_offsets: bool = False + ) -> None: + self.forwards += 1 + self.rows[rows] = self.rows.get(rows, 0) + 1 + self.buckets += len(buckets) + if rows > 1: + self.multi_row_forwards += 1 + self.multi_row_buckets += len(buckets) + if len(buckets) > 1: + first = buckets[0][0] + for bucket in buckets[1:]: + self.reason(self.split_reason(first, bucket[0], cross_turn=cross_turn, row_offsets=row_offsets)) + + def record_decode(self, rows: int) -> None: + self.decodes += 1 + self.decoded_rows += rows + + def maybe_log(self) -> None: + if self.forwards < self.every: + return + logger.info( + "Code2Wav buckets: %d forwards rows=%s buckets/forward %.2f (multi-row %.2f over %d) " + "decodes/forward %.2f rows/decode %.2f split=%s", + self.forwards, + dict(sorted(self.rows.items())), + self.buckets / self.forwards, + self.multi_row_buckets / self.multi_row_forwards if self.multi_row_forwards else 0.0, + self.multi_row_forwards, + self.decodes / self.forwards, + self.decoded_rows / self.decodes if self.decodes else 0.0, + dict(sorted(self.reasons.items())), + ) + self._reset() + + class MiniCPMO45Code2Wav(nn.Module): """LLM_GENERATION model with request-owned state and compatible batching.""" @@ -265,12 +386,40 @@ def __init__( raise ValueError("MiniCPM-o Code2Wav setup cache capacity must be >= 0") self._connector_config = { "codec_chunk_frames": int(extra.get("codec_chunk_frames", 25)), + "initial_codec_chunk_frames": int(extra.get("initial_codec_chunk_frames", 0)), "codec_left_context_frames": int(extra.get("codec_left_context_frames", 3)), + "hift_max_lazy_graphs": int(extra.get("hift_max_lazy_graphs", 8)), + # Additional codec chunk sizes whose first and continuation + # vocoder shapes get HiFT graphs, e.g. [25] for full-duplex units + # under a 75-token turn chunk. + "hift_graph_codec_chunk_frames": list(extra.get("hift_graph_codec_chunk_frames") or ()), + # Exact-shape HiFT graphs (default off): inclusive [first, last] + # codec-frame ranges of first chunks (no source cache) and of + # continuations, captured at hift_graph_exact_batch_sizes (default + # [1]) and replayed unpadded instead of running eager. + "hift_graph_first_chunk_frames": list(extra.get("hift_graph_first_chunk_frames") or ()), + "hift_graph_continuation_frames": list(extra.get("hift_graph_continuation_frames") or ()), + "hift_graph_exact_batch_sizes": list(extra.get("hift_graph_exact_batch_sizes") or (1,)), } if self._connector_config["codec_chunk_frames"] <= 0 or self._connector_config["codec_left_context_frames"] < 0: raise ValueError(f"Invalid MiniCPM-o connector chunk configuration: {self._connector_config}") + # The most requests any Stage-2 batch call ever holds; both the HiFT + # graph batch ladder and the Whole-Euler micro-batch size are sized + # off it below. + max_num_seqs = getattr(getattr(vllm_config, "scheduler_config", None), "max_num_seqs", None) raw_capture_batch_sizes = extra.get("hift_graph_capture_batch_sizes") - capture_batch_sizes = [1] if raw_capture_batch_sizes is None else raw_capture_batch_sizes + if raw_capture_batch_sizes is None: + # Powers of two up to the scheduler's batch (capped at 32), so every + # vocoder batch rounds up to a graph (HiFTGraphWrapper.replay). + batch_cap = min(int(max_num_seqs), 32) if max_num_seqs else 32 + capture_batch_sizes = [] + size = 1 + while size < batch_cap: + capture_batch_sizes.append(size) + size *= 2 + capture_batch_sizes.append(batch_cap) + else: + capture_batch_sizes = raw_capture_batch_sizes max_serial_batch = extra.get("max_serial_batch") max_serial_batch = 4 if max_serial_batch is None else int(max_serial_batch) self._hift_graph_config = { @@ -279,26 +428,68 @@ def __init__( "max_serial_batch": max_serial_batch, } enable_whole_euler = extra.get("enable_whole_euler") - max_graph_batch_raw = extra.get("max_graph_batch") - max_graph_batch = int(max_graph_batch_raw) if max_graph_batch_raw is not None else None - micro_batch_size_raw = extra.get("micro_batch_size") - if micro_batch_size_raw is not None: - micro_batch_size = int(micro_batch_size_raw) - else: - # The Whole-Euler arena reserves one attention cache per micro-batch - # row, so size it for the most requests this stage ever batches. - max_num_seqs = getattr(getattr(vllm_config, "scheduler_config", None), "max_num_seqs", None) - micro_batch_size = min(int(max_num_seqs), max_graph_batch or 16) if max_num_seqs else None + # The Whole-Euler micro batch (one arena attention cache per row) and + # graph-path gate follow max_num_seqs alone: the old micro_batch_size / + # max_graph_batch overrides could contradict it and are not read. + micro_batch_size = int(max_num_seqs) if max_num_seqs else 16 self._cfm_graph_config = { "enabled": bool(extra.get("enable_cfm_graph", False)), "max_graphs": int(extra.get("cfm_max_graphs", 32)), "bucket_frames": int(extra.get("cfm_graph_bucket_frames", 0)), "capture_frames": extra.get("cfm_graph_capture_frames"), + # Cache lengths snap onto a grid of this many frames (0: exact), + # so first chunks of any size share a few graphs. + "offset_bucket_frames": int(extra.get("cfm_graph_offset_bucket_frames", 50)), "enable_whole_euler": enable_whole_euler is None or bool(enable_whole_euler), "max_serial_batch": max_serial_batch, - "max_graph_batch": max_graph_batch, + "max_graph_batch": micro_batch_size, "micro_batch_size": micro_batch_size, "pad_max_rows": extra.get("whole_euler_pad_max_rows"), + # Fused DiT block forward for the ragged / Whole-Euler solves + # (``dit_fused.py``): same function, ~1/3 of the kernels. + "fused_body": bool(extra.get("cfm_fused_body", False)), + # Request estimator caches resident in a Whole-Euler slot pool + # (``AttSlotPool``, needs fused_body): no per-replay cache copies. + "slot_pool": bool(extra.get("cfm_slot_pool", False)), + # Fused-body GEMM backend: cublas (default), or Triton TF32 tiles + # sized for the streaming row counts (``triton``; ``triton_windows`` + # also runs each causal conv as one GEMM over its im2col windows). + "dit_gemm": extra.get("cfm_dit_gemm"), + # Its products: tf32, or tf32x3 (three TF32 products, fp32-level + # error); unset follows MINICPMO_CODE2WAV_TF32 (flow: tf32). + "dit_gemm_precision": extra.get("cfm_dit_gemm_precision"), + # Keep one immutable prompt estimator-cache suffix per voice and + # only a request-owned prefix in each state. Disabled until a CUDA + # probe confirms the suffix is bit-identical. + "prompt_att_sharing": bool(extra.get("cfm_prompt_att_sharing", False)), + # Arena storage follows the largest explicitly captured graph tier + # instead of max_num_seqs; larger batches replay multiple tiers. + "arena_rows_from_graph_grid": bool(extra.get("cfm_arena_rows_from_graph_grid", False)), + # Slot-pool ragged solves keep each row's cache offset, so a + # stream's second chunk and steady ones share one replay. + "row_offset_merge": extra.get("cfm_row_offset_merge", False) is True, + } + # Exact-shape CUDA graphs of the flow encoder's continuation chunk + # (``cfm_encoder_cuda_graph``, default off; ``FlowEncoderGraphs``): + # one per (rows, tokens, conformer cache frames), captured at startup + # for the default prompt, replayed only for calls of exactly that + # shape. ``cfm_encoder_graph_rows`` is the largest row count (every + # count up to it) or a list of counts; ``cfm_encoder_graph_token_widths`` + # defaults to the duplex unit chunk (left context + one unit). + encoder_rows = extra.get("cfm_encoder_graph_rows", 8) + if isinstance(encoder_rows, (list, tuple)): + encoder_rows = sorted({int(r) for r in encoder_rows}) + else: + encoder_rows = list(range(1, int(encoder_rows) + 1)) + if max_num_seqs: + encoder_rows = [r for r in encoder_rows if r <= int(max_num_seqs)] + if any(r < 1 for r in encoder_rows): + raise ValueError("MiniCPM-o cfm_encoder_graph_rows must be positive") + encoder_widths = extra.get("cfm_encoder_graph_token_widths") + self._encoder_graph_config = { + "enabled": extra.get("cfm_encoder_cuda_graph", False) is True, + "rows": encoder_rows, + "token_widths": [int(w) for w in encoder_widths] if encoder_widths else None, } self._ref_max_seconds = float(extra.get("ref_audio_max_seconds", _REF_MAX_SECONDS)) if self._ref_max_seconds <= 0: @@ -314,6 +505,24 @@ def __init__( self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") self._default_prompt_normalized: tuple[str, str] | None = None + # A fresh stream normally has a different cache-shape signature from + # a continuation. Ragged Whole-Euler can keep estimator offsets per + # request, so an explicit opt-in may join those rows when supported. + self._onset_merge = os.environ.get("VLLM_OMNI_CODEC_ONSET_MERGE") == "1" + # Decode rows of different turns of one stream family together: + # cache_epoch only decides whether a row starts a fresh state + # (_parse_item); the decode itself never reads it. + self._cross_turn_buckets = extra.get("cfm_cross_turn_buckets", False) is True + # Decode a stream's second chunk with steady ones: slot-pool rows whose + # caches differ only in length share a bucket and one Whole-Euler + # replay (``BatchedToken2Wav.can_merge_row_offsets``). + self._row_offset_merge = self._cfm_graph_config["row_offset_merge"] + self._bucket_stats = _BucketStats() if extra.get("code2wav_bucket_stats", False) is True else None + # Return the caching allocator's unused blocks to the driver once the + # startup precapture is done (``cfm_precapture_empty_cache``, default + # off). Startup leaves the eager warmup/prompt solves' blocks cached; + # releasing them changes no tensor, only what the process holds. + self._precapture_empty_cache = extra.get("cfm_precapture_empty_cache", False) is True @property def _default_prompt_wav(self) -> str: @@ -658,12 +867,17 @@ def _parse_item( ) @staticmethod - def _bucket_key(item: _WorkItem) -> tuple[Any, ...]: + def _bucket_key(item: _WorkItem, *, cross_turn: bool = False, row_offsets: bool = False) -> tuple[Any, ...]: cache_signature: Any if item.previous is None: cache_signature = ("uninitialized",) + elif row_offsets and _slot_resident(item): + # A slot-pool row keeps its cache lengths per row (``cfm_row_offset_merge``). + cache_signature = ("row_offsets", row_offset_signature(item.previous.token2wav)) else: cache_signature = state_shape_signature(item.previous.token2wav) + if cross_turn: + return (item.prompt_cache_id, item.prompt_wav, cache_signature) return ( item.prompt_cache_id, item.prompt_wav, @@ -671,6 +885,44 @@ def _bucket_key(item: _WorkItem) -> tuple[Any, ...]: item.cache_epoch, ) + def _signature_groups( + self, bucket: list[_WorkItem], states: list[BatchedToken2WavState], *, row_offsets: bool + ) -> list[list[int]]: + """A bucket's rows grouped by cache signature, for a bucket that cannot decode as one. + + With ``row_offsets`` (``cfm_row_offset_merge``) the slot-pool groups + that differ only in their cache lengths stay one group when the + backend can solve them together. + """ + groups: dict[tuple[Any, ...], list[int]] = {} + for row, state in enumerate(states): + groups.setdefault(state_shape_signature(state), []).append(row) + if not row_offsets: + return list(groups.values()) + clusters: dict[tuple[Any, ...], list[list[int]]] = {} + for signature, rows in groups.items(): + if _slot_resident(bucket[rows[0]]): + signature = ("row_offsets", row_offset_signature(states[rows[0]])) + clusters.setdefault(signature, []).append(rows) + result: list[list[int]] = [] + for parts in clusters.values(): + rows = sorted(row for part in parts for row in part) + if len(parts) > 1 and self.backend.can_merge_row_offsets( + [states[row] for row in rows], [int(bucket[row].tokens.numel()) for row in rows] + ): + result.append(rows) + else: + result.extend(parts) + return result + + @staticmethod + def _onset_group_key(key: tuple[Any, ...], *, cross_turn: bool = False) -> tuple[Any, ...]: + """The onset-merge group of a bucket key: its prompt, plus its cache epoch unless rows merge across turns.""" + if cross_turn: + return key[:2] + prompt_cache_id, prompt_wav, _, cache_epoch = key + return (prompt_cache_id, prompt_wav, cache_epoch) + def _iter_limited_batches( self, buckets: Iterable[list[_WorkItem]], @@ -762,6 +1014,12 @@ def _forward_impl( **kwargs: Any, ) -> OmniOutput: del positions, intermediate_tensors, inputs_embeds + if getattr(self, "_precapture_pending", False) and self.backend is not None: + try: + self._precapture_default_prompt() + finally: + if getattr(self, "_precapture_empty_cache", False): + self._release_precapture_cache() ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.empty(0, dtype=torch.long) segments = self._split_segments(ids, kwargs.get("seq_token_counts")) empty = torch.empty(0, dtype=torch.float32, device=ids.device) @@ -845,9 +1103,16 @@ def _forward_impl( self._trim_runtime_prompts() raise _batch_error("empty_nonfinal_chunk", request_ids=invalid_empty) + cross_turn = getattr(self, "_cross_turn_buckets", False) is True + row_offsets = getattr(self, "_row_offset_merge", False) is True + bucket_stats = getattr(self, "_bucket_stats", None) buckets: dict[tuple[Any, ...], list[_WorkItem]] = {} for item in compute_items: - buckets.setdefault(self._bucket_key(item), []).append(item) + buckets.setdefault(self._bucket_key(item, cross_turn=cross_turn, row_offsets=row_offsets), []).append(item) + if bucket_stats is not None: + bucket_stats.record_forward( + len(compute_items), list(buckets.values()), cross_turn=cross_turn, row_offsets=row_offsets + ) undersized = [ { "size": len(bucket), @@ -921,20 +1186,99 @@ def _forward_impl( prompt_wav=item.prompt_wav, token2wav=state, ) - for bucket in self._iter_decode_batches(buckets.values()): + decode_buckets: Iterable[list[_WorkItem]] + if self._onset_merge and self.backend is not None and self.backend._whole_euler_ragged_active(): + # Keep ordinary exact-shape buckets unchanged, but offer rows from + # the same prompt/epoch to the ragged Whole-Euler path when a + # fresh onset and a continuation are ready together. + grouped: dict[tuple[Any, ...], list[list[_WorkItem]]] = {} + for key, bucket in buckets.items(): + grouped.setdefault(self._onset_group_key(key, cross_turn=cross_turn), []).append(bucket) + merged: list[list[_WorkItem]] = [] + for entries in grouped.values(): + combined = [item for bucket in entries for item in bucket] + mixed = any(item.previous is None for item in combined) and any( + item.previous is not None for item in combined + ) + merged.extend([combined] if mixed else entries) + decode_buckets = merged + else: + decode_buckets = buckets.values() + for bucket in self._iter_decode_batches(decode_buckets): batch_size = len(bucket) try: features = self.backend.prepare_prompt( bucket[0].prompt_cache_id, bucket[0].prompt_wav, ) - if bucket[0].previous is None: - states = self.backend.setup_batch(features, batch_size) - else: - states = [item.previous.token2wav for item in bucket if item.previous is not None] + fresh = [item for item in bucket if item.previous is None] + setup = iter(self.backend.setup_batch(features, len(fresh)) if fresh else []) + states = [next(setup) if item.previous is None else item.previous.token2wav for item in bucket] token_lengths = {int(item.tokens.numel()) for item in bucket} last_chunk_values = {item.last_chunk for item in bucket} - if len(token_lengths) > 1 or len(last_chunk_values) > 1: + onset_mix = ( + self._onset_merge + and any(item.previous is None for item in bucket) + and any(item.previous is not None for item in bucket) + ) + onset_ready = onset_mix and self.backend.can_merge_state_shapes( + states, [int(item.tokens.numel()) for item in bucket] + ) + row_groups: list[list[int]] | None = None + offset_mix = False + if onset_mix and not onset_ready: + # The opt-in grouping is deliberately conservative. If + # cache layouts do not satisfy ragged Whole-Euler, split + # back into the historical signature buckets. + if bucket_stats is not None: + bucket_stats.reason("onset_merge_rejected") + row_groups = self._signature_groups(bucket, states, row_offsets=row_offsets) + elif row_offsets and not onset_ready and batch_size > 1: + groups = self._signature_groups(bucket, states, row_offsets=True) + if len(groups) > 1: + # Slot-pool rows the backend cannot solve at once go + # back to their signature buckets. + if bucket_stats is not None: + bucket_stats.reason("row_offset_merge_rejected") + row_groups = groups + else: + # One decode, ragged when it holds rows of different cache lengths. + offset_mix = len({state_shape_signature(state) for state in states}) > 1 + if row_groups is not None: + audios_rows: list[torch.Tensor | None] = [None] * batch_size + next_states_rows: list[Any] = [None] * batch_size + for rows in row_groups: + if bucket_stats is not None: + bucket_stats.record_decode(len(rows)) + row_tokens = [bucket[row].tokens for row in rows] + row_states = [states[row] for row in rows] + row_last = [bucket[row].last_chunk for row in rows] + if ( + len({int(tokens.numel()) for tokens in row_tokens}) > 1 + or len(set(row_last)) > 1 + or (row_offsets and len({state_shape_signature(state) for state in row_states}) > 1) + ): + row_audios, row_next = self.backend.decode_ragged_batch( + row_tokens, + features, + row_states, + last_chunks=row_last, + ) + else: + row_audios, row_next = self.backend.decode_batch( + torch.stack(row_tokens, dim=0), + features, + row_states, + last_chunk=row_last[0], + ) + for row, audio, next_state in zip(rows, row_audios, row_next, strict=True): + audios_rows[row] = audio + next_states_rows[row] = next_state + audios = cast(list[torch.Tensor], audios_rows) + next_states = next_states_rows + elif onset_ready or offset_mix or len(token_lengths) > 1 or len(last_chunk_values) > 1: + if bucket_stats is not None: + bucket_stats.record_decode(batch_size) audios, next_states = self.backend.decode_ragged_batch( [item.tokens for item in bucket], features, @@ -942,6 +1286,8 @@ def _forward_impl( last_chunks=[item.last_chunk for item in bucket], ) else: + if bucket_stats is not None: + bucket_stats.record_decode(batch_size) tokens = torch.stack([item.tokens for item in bucket], dim=0) audios, next_states = self.backend.decode_batch( tokens, @@ -981,6 +1327,8 @@ def _forward_impl( ) ) + if bucket_stats is not None: + bucket_stats.maybe_log() self._commit_runtime_prompt_owners(items) for request_id, state in pending.items(): if state is None: @@ -1055,6 +1403,15 @@ def _build_backend(self) -> None: token2wav_path = Path(self.model_path) / "assets" / "token2wav" if not token2wav_path.is_dir(): raise FileNotFoundError(f"MiniCPM-o Code2Wav assets not found: {token2wav_path}") + # Token2wav runs in fp32, so without TF32 every flow-DiT GEMM runs on + # SIMT cores (the top Stage-2 kernel on A800); cuDNN convolutions + # already default to TF32. Stage 2 owns its process, so the global flag + # stays within it; it is set before any CUDA graph is captured. + tf32_mode = _tf32_mode(extra) if current_omni_platform.is_cuda() else "off" + if tf32_mode == "all": + torch.backends.cuda.matmul.allow_tf32 = True + if tf32_mode != "off": + logger.info("MiniCPM-o Code2Wav: TF32 matmul enabled (%s)", tf32_mode) use_float16 = bool(extra.get("token2wav_float16", False)) previous_dtype = torch.get_default_dtype() try: @@ -1062,10 +1419,20 @@ def _build_backend(self) -> None: # Token2wav contains fp32-only S3Tokenizer/HiFT modules, so build # its independent assets in their native precision. torch.set_default_dtype(torch.float32) + # Connector extra ``token2wav_s3tokenizer_device`` (default unset: + # same device as the vocoder). "cpu" frees ~472 MiB of fp32 + # S3Tokenizer weights; it only runs when a new reference voice is + # prepared (startup default prompt, runtime ``ref`` audio misses). + tokenizer_device = extra.get("token2wav_s3tokenizer_device") + tokenizer_kwargs = {"audio_tokenizer_device": str(tokenizer_device)} if tokenizer_device else {} token2wav = Token2wav( str(token2wav_path), float16=use_float16, n_timesteps=int(extra.get("token2wav_n_timesteps", 10)), + # The batched backend builds its estimator caches itself and + # never runs the upstream chunk path that owns these buffers. + drop_upstream_chunk_att_buffers=extra.get("cfm_drop_upstream_att_buffer", False) is True, + **tokenizer_kwargs, ) finally: torch.set_default_dtype(previous_dtype) @@ -1099,4 +1466,74 @@ def _build_backend(self) -> None: cfm_graph_config=self._cfm_graph_config, bfloat16_attention_cache=bool(extra.get("code2wav_bfloat16_attention_cache", False)), setup_cache_size=self._setup_cache_size, + cfm_tf32=tf32_mode == "flow", + encoder_graph_config=getattr(self, "_encoder_graph_config", None), ) + # Captured by the first forward (the engine's warmup run): graphs captured + # while vLLM loads the weights held GiBs each, a few MiB at forward time. + self._precapture_pending = bool(extra.get("cfm_graph_precapture", True)) + + def _release_precapture_cache(self) -> None: + """Release the allocator blocks startup left cached (``cfm_precapture_empty_cache``). + + Only unused cached segments go back to the driver: live tensors, the + slot pool and the CUDA-graph pool's segments (held while their graphs + live) stay where they are, so every later solve reads the same values. + """ + if not torch.cuda.is_available() or not torch.cuda.is_initialized(): + return + device = torch.device("cuda", torch.accelerator.current_device_index()) + torch.accelerator.synchronize(device) + before = _memory_snapshot(device) + torch.accelerator.empty_cache() + after = _memory_snapshot(device) + if before is not None and after is not None: + logger.info( + "MiniCPM-o Code2Wav: released %.0f MiB of cached memory after precapture " + "(allocated %.0f MiB, reserved %.0f -> %.0f MiB)", + (before[1] - after[1]) / 2**20, + after[0] / 2**20, + before[1] / 2**20, + after[1] / 2**20, + ) + + def _precapture_default_prompt(self) -> None: + """Capture the HiFT and default-voice Whole-Euler graphs before serving. + + Otherwise each graph is captured by the first chunk that needs it, + stalling every stream for 1-2 s per graph. + """ + self._precapture_pending = False + assert self.backend is not None + started = time.perf_counter() + try: + hift_captured = self.backend.precapture_hift() + except Exception: + logger.warning("MiniCPM-o Code2Wav: HiFT precapture failed; graphs capture lazily", exc_info=True) + hift_captured = 0 + if hift_captured: + logger.info( + "MiniCPM-o Code2Wav: precaptured %d HiFT CUDA Graph shapes in %.1f s", + hift_captured, + time.perf_counter() - started, + ) + prompt_wav, prompt_cache_id = self._normalized_default_prompt() + try: + features = self.backend.prepare_prompt(prompt_cache_id, prompt_wav) + captured = self.backend.precapture_whole_euler(features) + except Exception: + logger.warning("MiniCPM-o Code2Wav: Whole-Euler precapture failed; graphs capture lazily", exc_info=True) + return + if captured: + logger.info( + "MiniCPM-o Code2Wav: precaptured %d Whole-Euler graphs for the default prompt in %.1f s", + captured, + time.perf_counter() - started, + ) + if getattr(self.backend, "_encoder_graphs", None) is not None: + try: + self.backend.precapture_flow_encoder(features) + except Exception: + logger.warning( + "MiniCPM-o Code2Wav: flow encoder precapture failed; the encoder stays eager", exc_info=True + ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py index 9451c256061..ae7b6372c29 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py @@ -64,6 +64,8 @@ def __init__( float16: bool = False, n_timesteps: int = 10, device: str | torch.device | None = None, + drop_upstream_chunk_att_buffers: bool = False, + audio_tokenizer_device: str | None = None, ): self.float16 = float16 self.n_timesteps = n_timesteps @@ -73,6 +75,10 @@ def __init__( float16=float16, device=self.device, n_timesteps=n_timesteps, + # ``BatchedToken2Wav`` keeps its own estimator caches; only + # ``set_stream_cache`` / ``stream`` below read these buffers. + drop_upstream_chunk_att_buffers=drop_upstream_chunk_att_buffers, + audio_tokenizer_device=audio_tokenizer_device, ) # Eager-load so construction failures surface at init time (same as # the external Token2wav package), not on the first request. diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index 514ea84d95c..af910a28da3 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -14,6 +14,7 @@ build_duplex_intermediate_buffer, set_ref_audio, set_tts_handoff, + unpack_transport_tensor, ) from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.minicpmo_4_5 import ( @@ -25,6 +26,7 @@ logger = logging.getLogger(__name__) _MINICPMO45_ASYNC_STATE = "_minicpmo45_async_codec_state" _MINICPMO45_STREAM_RECORD = "_minicpmo45_async_stream_record" +_MINICPMO45_REF_AUDIO_CACHE = "minicpmo45_ref_audio" _MINICPMO45_SILENCE_CODE = 4218 _MINICPMO45_MIN_STREAM_BODY_FRAMES = 5 @@ -92,19 +94,40 @@ def _extract_prompt_reference_audio( return _extract_first_audio_ref({"audio": serving_reference_audio}) -def _extract_native_runtime_ref_audio(data_plane_metadata): +def _extract_native_runtime_ref_audio(data_plane_metadata, streaming_context=None): if not isinstance(data_plane_metadata, dict): - return None + data_plane_metadata = {} runtime_config = data_plane_metadata.get("runtime_config") + if not isinstance(runtime_config, dict): + bridge_states = getattr(streaming_context, "bridge_states", None) + duplex_state = bridge_states.get("duplex") if isinstance(bridge_states, dict) else None + runtime_config = duplex_state.get("runtime_config") if isinstance(duplex_state, dict) else None if not isinstance(runtime_config, dict): return None + # Every handoff of a session carries the same base64 reference voice. + # Decode it once per distinct config instead of once per ~1 s unit on the + # orchestrator thread; the key holds every input the decode reads. + cache_key = ( + runtime_config.get("ref_audio_data"), + runtime_config.get("ref_audio_format"), + runtime_config.get("ref_audio_sample_rate_hz"), + any(key in runtime_config for key in ("ref_audio_path", "tts_ref_audio_path")), + ) + bridge_states = getattr(streaming_context, "bridge_states", None) + cached = bridge_states.get(_MINICPMO45_REF_AUDIO_CACHE) if isinstance(bridge_states, dict) else None + if isinstance(cached, tuple) and cached[0] == cache_key: + return cached[1] from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import decode_native_ref_audio_from_config waveform = decode_native_ref_audio_from_config({"extra_body": runtime_config}) if waveform is None: - return None - sample_rate = runtime_config.get("ref_audio_sample_rate_hz") or 16000 - return torch.as_tensor(waveform, dtype=torch.float32).reshape(-1).cpu(), int(sample_rate) + ref_audio = None + else: + sample_rate = runtime_config.get("ref_audio_sample_rate_hz") or 16000 + ref_audio = torch.as_tensor(waveform, dtype=torch.float32).reshape(-1).cpu(), int(sample_rate) + if isinstance(bridge_states, dict): + bridge_states[_MINICPMO45_REF_AUDIO_CACHE] = (cache_key, ref_audio) + return ref_audio def _coerce_token_id_list(value): @@ -127,14 +150,6 @@ def _coerce_token_id_list(value): return out -def _to_transport_list(value): - if hasattr(value, "detach"): - value = value.detach().cpu() - if isinstance(value, torch.Tensor): - return value.tolist() - return value - - def _coerce_int(value): if hasattr(value, "detach"): flat = value.detach().cpu().reshape(-1) @@ -147,13 +162,20 @@ def _coerce_int(value): return None -def _codec_config(transfer_manager: Any) -> tuple[int, int]: +def _connector_extra(transfer_manager: Any) -> dict: connector = getattr(transfer_manager, "connector", None) raw_config = getattr(connector, "config", {}) or {} config = raw_config.get("extra", raw_config) if isinstance(raw_config, dict) else {} - config = config if isinstance(config, dict) else {} + return config if isinstance(config, dict) else {} + + +def _codec_config(transfer_manager: Any) -> tuple[int, int]: + config = _connector_extra(transfer_manager) chunk_frames = int(config.get("codec_chunk_frames", 25)) left_context_frames = int(config.get("codec_left_context_frames", 3)) + initial = int(config.get("initial_codec_chunk_frames", 0)) + if initial and not _MINICPMO45_MIN_STREAM_BODY_FRAMES <= initial <= chunk_frames: + raise ValueError("MiniCPM-o initial_codec_chunk_frames must be 0 or between 5 and codec_chunk_frames") if chunk_frames <= 0 or left_context_frames < 0: raise ValueError( "Invalid MiniCPM-o codec chunk config: " @@ -164,20 +186,8 @@ def _codec_config(transfer_manager: Any) -> tuple[int, int]: def _initial_chunk_frames(transfer_manager: Any, chunk_frames: int) -> int: - """First-payload window size in frames. - - Reads ``initial_codec_chunk_frames`` from the connector extra; falls back to - the steady-state value when absent, non-positive, or not smaller than the - steady ``codec_chunk_frames`` (which means the feature is off). - """ - connector = getattr(transfer_manager, "connector", None) - raw_config = getattr(connector, "config", {}) or {} - config = raw_config.get("extra", raw_config) if isinstance(raw_config, dict) else {} - config = config if isinstance(config, dict) else {} - value = int(config.get("initial_codec_chunk_frames", 0) or 0) - if value <= 0 or value >= chunk_frames: - return chunk_frames - return value + """First-payload window: ``initial_codec_chunk_frames`` (validated by ``_codec_config``; 0 = off).""" + return int(_connector_extra(transfer_manager).get("initial_codec_chunk_frames", 0) or 0) or chunk_frames def _request_intermediate_section(request: object, section: str) -> dict[str, object]: @@ -357,12 +367,12 @@ def tts2code2wav_async_chunk( request_finished = getattr(request, "is_finished", None) finished = bool(is_finished or (callable(request_finished) and request_finished())) chunk_frames, left_context_frames = _codec_config(transfer_manager) - # First-payload window: the payload with chunk_seq == 0 uses - # initial_codec_chunk_frames; later payloads return to the steady - # chunk_frames. When unset, first_window == chunk_frames and behavior is - # unchanged. + # Only a turn-mode stream's first payload uses initial_codec_chunk_frames; + # native duplex keeps its own short-unit threshold (hold_short_unit below). first_window = ( - _initial_chunk_frames(transfer_manager, chunk_frames) if int(record["chunk_seq"]) == 0 else chunk_frames + _initial_chunk_frames(transfer_manager, chunk_frames) + if not native_duplex and int(record["chunk_seq"]) == 0 + else chunk_frames ) flush_pending = finished last_chunk = bool(flush_pending and (not native_duplex or turn_end)) @@ -416,7 +426,7 @@ def tts2code2wav_async_chunk( if int(record["cache_epoch"]) == 0 and chunk_seq == 0: codes_info = _request_intermediate_section(request, "codes") meta_info = _request_intermediate_section(request, "meta") - raw_ref_audio = codes_info.get("ref") + raw_ref_audio = unpack_transport_tensor(codes_info.get("ref")) raw_ref_audio_sr = meta_info.get("ref_audio_sr") ref_audio_sr = _coerce_int(raw_ref_audio_sr) if raw_ref_audio is not None: @@ -476,7 +486,7 @@ def tts2code2wav_full_payload( meta_info = _request_intermediate_section(request, "meta") duplex_info = _request_intermediate_section(request, "duplex") - ref_audio = codes_info.get("ref") + ref_audio = unpack_transport_tensor(codes_info.get("ref")) finished = torch.tensor(True, dtype=torch.bool) return OmniPayloadStruct( codes=CodesStruct( @@ -783,6 +793,80 @@ def _latest_block(width: int) -> int | None: ) +def _native_duplex_ledger_row_ids(mm_output: Mapping[str, object]) -> tuple[list[int], list[int]] | None: + """Flatten the handoff's row ledger to ``(ids, positions)`` lists.""" + row_ids = mm_output.get("latent_input_ids") + row_positions = mm_output.get("latent_positions") + if not isinstance(row_ids, torch.Tensor) or not isinstance(row_positions, torch.Tensor): + return None + ids = [int(token_id) for token_id in row_ids.reshape(-1).tolist()] + positions = [int(position) for position in row_positions.reshape(-1).tolist()] + if len(ids) != len(positions): + return None + return ids, positions + + +def _native_duplex_block_base(ids: list[int], positions: list[int], block: list[int]) -> int | None: + """Latest contiguous, position-monotonic occurrence of ``block`` in the ledger.""" + width = len(block) + if width == 0: + return None + for base in range(len(ids) - width, -1, -1): + if ids[base : base + width] != block: + continue + first = positions[base] + if all(positions[base + offset] == first + offset for offset in range(width)): + return base + return None + + +def _native_duplex_trim_unbacked_decisions( + out_ids: list[int], + mm_output: Mapping[str, object], + special_token_ids: dict[str, int], + *, + request_id: str, +) -> list[int]: + """Drop leading decision tokens whose hidden rows are no longer in the ledger. + + A segment consumed without a Talker handoff (a listen, or a speak whose + turn ended empty) still advances the cumulative output, but the next + streaming update folds it into the rewritten prompt and the hidden ledger + restarts, so the next handoff would fail with ``missing own-token hidden + states``. Only control decision tokens are dropped, and only while the rest + is still a block the ledger backs; otherwise the input is returned unchanged + so the explicit error still fires. + """ + ledger = _native_duplex_ledger_row_ids(mm_output) + if ledger is None: + return out_ids + ids, positions = ledger + control_ids = {token_id for token_id in special_token_ids.values() if isinstance(token_id, int) and token_id >= 0} + + def backed(segment: list[int]) -> bool: + for candidate in (segment, segment[:-1]): + if candidate and _native_duplex_block_base(ids, positions, candidate) is not None: + return True + return False + + if backed(out_ids): + return out_ids + for trim in range(1, len(out_ids)): + if out_ids[trim - 1] not in control_ids: + break + if backed(out_ids[trim:]): + logger.warning( + "MiniCPM-o native duplex: skipped %d stale leading decision token(s) without " + "hidden rows in the current ledger for request_id=%s (unit=%s, ledger tail=%s)", + trim, + request_id, + out_ids, + ids[-32:], + ) + return out_ids[trim:] + return out_ids + + def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | None: bridge_states = getattr(streaming_context, "bridge_states", None) if not isinstance(bridge_states, dict): @@ -801,12 +885,20 @@ def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | turn_id = duplex_state.get("model_turn_id", duplex_state.get("turn_id")) if isinstance(turn_id, int): metadata["turn_id"] = turn_id - session_config = duplex_state.get("session_config") - if isinstance(session_config, dict): - metadata["session_config"] = dict(session_config) - runtime_config = duplex_state.get("runtime_config") - if isinstance(runtime_config, dict): - metadata["runtime_config"] = dict(runtime_config) + # MiniCPM Stage 0 and the native Talker keep session configuration in + # request-local state. Send the snapshot once per config generation so a + # large reference-audio base64 value does not cross the Stage0→Stage1 + # boundary for every one-second unit. + config_generation = duplex_state.get("config_generation") + sent_generation = duplex_state.get("_native_config_generation") + if sent_generation != config_generation: + session_config = duplex_state.get("session_config") + if isinstance(session_config, dict): + metadata["session_config"] = dict(session_config) + runtime_config = duplex_state.get("runtime_config") + if isinstance(runtime_config, dict): + metadata["runtime_config"] = dict(runtime_config) + duplex_state["_native_config_generation"] = config_generation return metadata @@ -979,6 +1071,12 @@ def llm2tts( out_ids = llm_output_ids listen_id = special_token_ids.get("listen_token_id") turn_eos_id = special_token_ids.get("turn_eos_token_id") + out_ids = _native_duplex_trim_unbacked_decisions( + out_ids, + mm_output, + special_token_ids, + request_id=str(llm_output.request_id), + ) unit_start = 0 while unit_start < len(out_ids) and out_ids[unit_start] == listen_id: unit_start += 1 @@ -1054,30 +1152,34 @@ def llm2tts( meta["turn_start"] = native_turn_start if native_segment_end: meta["segment_end"] = True + native_turn_end_handoff = False + if is_native_duplex_handoff: + turn_eos_id = special_token_ids.get("turn_eos_token_id") + native_turn_end_handoff = turn_eos_id is not None and handoff_ids is not None and turn_eos_id in handoff_ids + if not handoff_ids: + continue + # Resolve the reference voice only for a handoff that is submitted; + # listen-only units drop this buffer above. ref_audio = reference_audio_by_request_id.get(llm_output.request_id) if ref_audio is None: ref_audio = _extract_native_runtime_ref_audio( model_intermediate_buffer.get("duplex"), + _streaming_context, ) if ref_audio is not None: ref_waveform, ref_sr = ref_audio - set_ref_audio(model_intermediate_buffer, _to_transport_list(ref_waveform), ref_sr) - handoff_hidden = _to_transport_list(tts_hidden_slice) if tts_hidden_slice is not None else None - native_turn_end_handoff = False - if is_native_duplex_handoff: - turn_eos_id = special_token_ids.get("turn_eos_token_id") - native_turn_end_handoff = turn_eos_id is not None and handoff_ids is not None and turn_eos_id in handoff_ids - if not handoff_ids: - continue - set_tts_handoff(model_intermediate_buffer, handoff_ids, handoff_hidden) + set_ref_audio(model_intermediate_buffer, ref_waveform, ref_sr) + # Tensors are packed as raw bytes (bit-identical) instead of Python + # lists: see pack_transport_tensor. + set_tts_handoff(model_intermediate_buffer, handoff_ids, tts_hidden_slice) if native_turn_end_handoff: model_intermediate_buffer.setdefault("meta", {})["turn_end"] = True condition_sequence_state = None condition_sequence_value = None - if handoff_ids is not None and handoff_hidden is not None: + if handoff_ids is not None and tts_hidden_slice is not None: condition_suffix_length = 1 if is_native_duplex_handoff else 2 - condition_length = max(len(handoff_ids), len(handoff_hidden)) + condition_suffix_length + condition_length = max(len(handoff_ids), len(tts_hidden_slice)) + condition_suffix_length # Dummy ids only reserve scheduler slots; prefill overwrites the # embeddings. Talker.sample() blanks the whole prompt out of the # repetition penalty, so codec id 0 is not taxed from step one. From 94be6ab3f510e01c21905552fca667185828b845 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:30 +0800 Subject: [PATCH 82/89] [Perf][MiniCPM-o] Code2Wav exact-shape CUDA graphs for flow encoder, HiFT, and deadline scheduling Code2Wav exact-shape CUDA graphs for flow encoder, HiFT, and deadline scheduling Add exact-shape CUDA graphs for the flow encoder continuation chunks and HiFT vocoder, along with deadline-driven request scheduling to maximize concurrency and minimize latency. Signed-off-by: BeatSeat --- tests/core/sched/test_codec_deadline.py | 278 ++++ .../minicpmo_4_5/test_cfm_graph_bucketing.py | 87 ++ .../minicpmo_4_5/test_cuda_graph_wrapper.py | 333 ++++- .../minicpmo_4_5/test_flow_encoder_graph.py | 637 +++++++++ .../minicpmo_4_5/test_hift_exact_graphs.py | 248 ++++ .../models/minicpmo_4_5/codec_deadline.py | 481 +++++++ .../models/minicpmo_4_5/cuda_graph_wrapper.py | 1189 +++++++++++++++-- .../models/minicpmo_4_5/flow_encoder_graph.py | 330 +++++ 8 files changed, 3433 insertions(+), 150 deletions(-) create mode 100644 tests/core/sched/test_codec_deadline.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_hift_exact_graphs.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/codec_deadline.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py diff --git a/tests/core/sched/test_codec_deadline.py b/tests/core/sched/test_codec_deadline.py new file mode 100644 index 00000000000..aa5443ba08f --- /dev/null +++ b/tests/core/sched/test_codec_deadline.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Deadline batching policy for the codec stage (``additional_config.codec_deadline``), pure Python.""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.core.sched.codec_deadline import ( + ChunkMeta, + CodecDeadlineConfig, + CodecDeadlinePolicy, + LateBreaker, + ReadyChunk, + StepTimeModel, + StreamLedger, + build_codec_deadline_policy, + graph_step_tiers, + seed_step_s, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NOW = 100.0 + + +@pytest.fixture(autouse=True) +def _default_graph_grid(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("VLLM_OMNI_CFM_GRAPH_GRID", raising=False) + + +def _meta(chunk_seq: int = 3, *, frames: int = 25, last: bool = False, turn_end: bool = False, epoch=0, cache=0): + return ChunkMeta( + chunk_seq=chunk_seq, + cache_epoch=cache, + last_chunk=last, + turn_end=turn_end, + codec_chunk_frames=frames, + duplex_epoch=epoch, + ) + + +def _policy(**config: object) -> CodecDeadlinePolicy: + return CodecDeadlinePolicy(CodecDeadlineConfig(enabled=True, **config)) + + +def _stream(policy: CodecDeadlinePolicy, rid: str, *, play_end: float, anchor: float = 90.0, audio_s: float = 20.0): + policy.ledgers[rid] = StreamLedger(key=(0, 0), anchor=anchor, play_end=play_end, audio_s=audio_s, d1=1.0) + return ReadyChunk(rid, _meta()) + + +def test_config_defaults_validation_and_parsing() -> None: + config = CodecDeadlineConfig() + assert config.enabled is False + assert (config.margin_s, config.max_hold_s, config.max_rows, config.eager_rows) == (0.1, 0.35, 8, 4) + assert CodecDeadlineConfig.from_additional_config({}) is None + assert CodecDeadlineConfig.from_additional_config(None) is None + parsed = CodecDeadlineConfig.from_additional_config({"codec_deadline": {"enabled": True, "max_rows": 16}}) + assert parsed is not None and parsed.enabled and parsed.max_rows == 16 + with pytest.raises(ValueError, match="unknown keys"): + CodecDeadlineConfig.from_additional_config({"codec_deadline": {"enabld": True}}) + # The client buffer and the graph tiers are not configured here any more. + for removed in ("prebuffer_s", "step_tiers", "step_seed_s", "trim_split_batches"): + with pytest.raises(ValueError, match="unknown keys"): + CodecDeadlineConfig.from_additional_config({"codec_deadline": {removed: 1}}) + with pytest.raises(ValueError, match="must be a boolean"): + CodecDeadlineConfig(enabled=1) # type: ignore[arg-type] + + +def test_chunk_meta_reads_tensor_fields_and_ignores_tts_is_last_chunk() -> None: + info = { + "meta": { + "chunk_seq": torch.tensor(4), + "cache_epoch": torch.tensor([1]), + "last_chunk": torch.tensor(False), + "turn_end": False, + "codec_chunk_frames": 28, + "duplex_epoch": 2, + # True on every native duplex unit (flush_pending): not a last chunk. + "tts_is_last_chunk": True, + } + } + meta = ChunkMeta.from_info(info) + assert meta == ChunkMeta(4, 1, False, False, 28, 2) + assert ChunkMeta.from_info({}) is None + assert ChunkMeta.from_info({"meta": {"request_id": "x"}}) is None + + +@pytest.mark.parametrize( + "meta", + [ + _meta(0), + _meta(last=True), + _meta(turn_end=True), + _meta(frames=0), + _meta(epoch=None), + _meta(cache=1), # another turn than the ledger + None, + ], +) +def test_eager_chunks(meta) -> None: + policy = _policy() + _stream(policy, "a", play_end=NOW + 10) + assert policy._eager(ReadyChunk("a", meta), NOW) is True + + +def test_a_continuation_with_a_ledger_is_paced_but_not_without_one() -> None: + policy = _policy() + chunk = _stream(policy, "a", play_end=NOW + 10) + assert policy._eager(chunk, NOW) is False + assert policy._eager(ReadyChunk("b", _meta()), NOW) is True + # The client ran dry long ago: a new segment, go now. + policy.ledgers["a"].play_end = NOW - 2.0 + assert policy._eager(chunk, NOW) is True + + +def test_ledger_anchors_buffers_once_and_reanchors_late_chunks() -> None: + policy = _policy() + policy.on_scheduled("a", _meta(0), NOW) + policy.on_output("a", 0.12, NOW) + ledger = policy.ledgers["a"] + assert ledger.play_end == pytest.approx(NOW + 0.5 + 0.12) + assert (ledger.anchor, ledger.d1) == (NOW, 0.12) + # Arrives 0.3 s after the client ran out: re-anchor at arrival. + policy.on_scheduled("a", _meta(1), NOW + 0.9) + policy.on_output("a", 1.0, NOW + 0.92) + assert ledger.late_s == pytest.approx(0.3) + assert ledger.play_end == pytest.approx(NOW + 1.92) + assert ledger.audio_s == pytest.approx(1.12) + # The turn ends: the next turn re-anchors without the buffer. + policy.on_scheduled("a", _meta(2, last=True), NOW + 1.5) + policy.on_output("a", 0.5, NOW + 1.5) + assert "a" not in policy.ledgers + policy.on_scheduled("a", _meta(0, cache=1), NOW + 5) + policy.on_output("a", 1.0, NOW + 5) + assert policy.ledgers["a"].play_end == pytest.approx(NOW + 6.0) + policy.forget("a") + assert "a" not in policy.ledgers and "a" not in policy.buffered + + +def test_rate_deadline_is_tighter_after_a_long_first_chunk() -> None: + policy = _policy(max_hold_s=10.0) + policy.ledgers["a"] = StreamLedger(key=(0, 0), anchor=NOW, play_end=NOW + 1.34, audio_s=0.84, d1=0.84) + chunk = ReadyChunk("a", _meta(1, frames=25)) + # rate: NOW + 1.1 * (0.84 + 1.0 - 0.84) = NOW + 1.1 < zero-stall NOW + 1.34; + # minus the margin, its own one-row step and an eager step that may run first. + release = policy.release_by(chunk, NOW, 1) + assert release == pytest.approx(NOW + 1.1 - 0.1 - seed_step_s(1) - seed_step_s(4)) + + +def test_eager_step_takes_only_rows_due_within_a_step() -> None: + policy = _policy() + ready = [ + ReadyChunk("e", _meta(0)), + _stream(policy, "soon", play_end=NOW + 0.5), # release 100.1 <= now + T(1) + _stream(policy, "slack", play_end=NOW + 10), # release by max hold: 100.35 + ] + plan = policy.plan(ready, NOW) + assert plan.trigger == "eager" + assert plan.released == ["e", "soon"] + assert plan.held == {"slack"} + assert policy.next_release_in(NOW) == pytest.approx(0.35) + + +def _two_due_three_slack(policy: CodecDeadlinePolicy) -> list[ReadyChunk]: + ready = [_stream(policy, f"due{i}", play_end=NOW + 0.3) for i in range(2)] + ready += [_stream(policy, f"slack{i}", play_end=NOW + 10 + i) for i in range(3)] + return ready + + +def test_deadline_step_fills_the_graph_tier(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("VLLM_OMNI_CFM_GRAPH_GRID", "1,4,8,16,20") + policy = _policy() + plan = policy.plan(_two_due_three_slack(policy), NOW) + assert plan.trigger == "deadline" + assert len(plan.released) == 4 # 2 due -> tier 4 + assert plan.released[:2] == ["due0", "due1"] + assert len(plan.held) == 1 + + +def test_deadline_step_does_not_fill_while_a_first_chunk_is_pending(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("VLLM_OMNI_CFM_GRAPH_GRID", "1,4,8,16,20") + policy = _policy() + plan = policy.plan(_two_due_three_slack(policy), NOW, onset_pending=True) + assert plan.trigger == "deadline" + assert plan.released == ["due0", "due1"] + assert len(plan.held) == 3 + + +def test_graph_step_tiers_follow_the_code2wav_grid(monkeypatch: pytest.MonkeyPatch) -> None: + assert graph_step_tiers(8) == [1, 2, 4, 8] # the Code2Wav default: powers of two + monkeypatch.setenv("VLLM_OMNI_CFM_GRAPH_GRID", "1,4,8,16,20") + assert graph_step_tiers(8) == [1, 4, 8] + assert graph_step_tiers(16) == [1, 4, 8, 16] + monkeypatch.setenv("VLLM_OMNI_CFM_GRAPH_GRID", "3,5") + assert graph_step_tiers(8) == [3, 5, 8] + + +def test_holds_everything_until_a_release_time() -> None: + policy = _policy() + ready = [_stream(policy, "a", play_end=NOW + 10), _stream(policy, "b", play_end=NOW + 10)] + plan = policy.plan(ready, NOW) + assert plan.trigger == "hold" + assert plan.released == [] and plan.held == {"a", "b"} + assert policy.next_release_in(NOW) == pytest.approx(0.35) + # The max hold expires: both go. + plan = policy.plan(ready, NOW + 0.36) + assert plan.trigger == "deadline" and set(plan.released) == {"a", "b"} + + +def test_max_rows_caps_a_rows_trigger() -> None: + policy = _policy(max_rows=8) + ready = [_stream(policy, f"r{i}", play_end=NOW + 10 + i) for i in range(10)] + plan = policy.plan(ready, NOW) + assert plan.trigger == "rows" and len(plan.released) == 8 + assert len(plan.held) == 2 + + +def test_step_time_model_ewma() -> None: + model = StepTimeModel((1, 4, 8), 0.5) + assert model.mean == pytest.approx([seed_step_s(1), seed_step_s(4), seed_step_s(8)]) + model.mean = [0.1, 0.2, 0.3] + assert model.tier_capacity(2) == 4 and model.tier_capacity(30) == 8 + assert model.estimate(1) == pytest.approx(0.1) + model.observe(1, 0.2) + assert model.mean[0] == pytest.approx(0.15) + assert model.mad[0] == pytest.approx(0.05) + assert model.estimate(1) == pytest.approx(0.15 + 1.5 * 0.05) + model.observe(0, 1.0) + assert model.mean[0] == pytest.approx(0.15) + + +def test_breaker_opens_after_late_held_chunks_and_cools_off() -> None: + breaker = LateBreaker(window=10, late_frac=0.1, cooloff_s=5.0) + breaker.record(True, NOW) + assert not breaker.is_open(NOW) + breaker.record(True, NOW) + assert breaker.is_open(NOW + 1) and breaker.opens == 1 + assert not breaker.is_open(NOW + 5.1) + + policy = _policy(breaker_window=2, breaker_late_frac=0.4) + for step in range(2): + rid = "a" + _stream(policy, rid, play_end=NOW - 0.5) + policy.held_once.add(rid) + policy.on_scheduled(rid, _meta(3 + step), NOW) + policy.on_output(rid, 1.0, NOW) + assert policy.breaker.is_open(NOW) + assert policy._eager(_stream(policy, "b", play_end=NOW + 10), NOW) is True + + +def test_late_held_chunks_are_counted_by_the_trigger_that_released_them() -> None: + policy = _policy(stats=True) + for trigger in ("eager", "deadline", "deadline"): + _stream(policy, "a", play_end=NOW - 0.5) + policy.held_once.add("a") + policy._last_trigger = trigger + policy.on_scheduled("a", _meta(3), NOW) + policy.on_output("a", 1.0, NOW) + assert policy._stats.late == {"eager": 1, "deadline": 2} + + +@pytest.mark.parametrize( + ("kwargs", "enabled"), + [ + ({"async_scheduling": False, "native_data_plane": False, "idle_wait_s": 0.05}, True), + ({"async_scheduling": True, "native_data_plane": False, "idle_wait_s": 0.05}, False), + ({"async_scheduling": False, "native_data_plane": True, "idle_wait_s": 0.05}, False), + ({"async_scheduling": False, "native_data_plane": False, "idle_wait_s": 0.0}, False), + ], +) +def test_build_requires_sync_scheduling_the_chunk_adapter_and_the_idle_park(kwargs, enabled) -> None: + policy = build_codec_deadline_policy({"codec_deadline": {"enabled": True}}, **kwargs) + assert (policy is not None) is enabled + assert build_codec_deadline_policy({"codec_deadline": {"enabled": False}}, **kwargs) is None + assert build_codec_deadline_policy({}, **kwargs) is None diff --git a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py index c72e947c179..4f688aaff0e 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py @@ -19,8 +19,10 @@ _zero_padded_frames, ) from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + SharedPromptAttCache, _att_keep_ranges, _build_capture_mask, + _capture_offset, _capture_query_width, _whole_euler_att_segments, _zero_padded_cnn_cache, @@ -29,6 +31,17 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_shared_prompt_attention_cache_is_bit_exact_and_range_copyable(): + prefix = torch.arange(2 * 2 * 2 * 1 * 5 * 4, dtype=torch.float32).reshape(2, 2, 2, 1, 5, 4) + prompt = torch.arange(2 * 2 * 2 * 1 * 3 * 4, dtype=torch.float32).reshape(2, 2, 2, 1, 3, 4) + 1000 + shared = SharedPromptAttCache(prefix, prompt) + torch.testing.assert_close(shared.materialize(), torch.cat((prefix, prompt), dim=4), rtol=0, atol=0) + copied = torch.empty_like(shared.materialize()) + shared.copy_range_to(copied[..., :4, :], 0, 4) + shared.copy_range_to(copied[..., 4:, :], 4, 4) + torch.testing.assert_close(copied, shared.materialize(), rtol=0, atol=0) + + def test_pads_partial_chunk_up_to_bucket(): # 50 frames on a 16-frame grid -> 14 padding frames. assert _cfm_pad_frames(mel_frames=50, offset=0, noise_capacity=1000, bucket_frames=16, disabled=False) == 14 @@ -145,6 +158,80 @@ def test_varied_chunk_lengths_collapse_onto_few_widths(): assert len(padded) < len(varied) +def test_capture_query_width_tuple_takes_the_narrowest_width_that_fits(): + widths = (56, 150) + # A 25-token duplex unit (50 mel frames, 56 as a final chunk) keeps its own + # narrow graph instead of the 75-token turn chunk's 150. + assert _capture_query_width(50, widths) == 56 + assert _capture_query_width(56, widths) == 56 + assert _capture_query_width(57, widths) == 150 + assert _capture_query_width(150, widths) == 150 + assert _capture_query_width(156, widths) == 300 + assert _capture_query_width(50, (150,)) == _capture_query_width(50, 150) == 150 + assert _capture_query_width(50, ()) == 50 + + +def test_capture_offset_grid_is_anchored_at_the_steady_cache(): + steady = 400 # prompt (300) + the 100 frames the streaming trim keeps + # First chunks of any size land on one of three offsets; steady is exact. + assert {_capture_offset(o, 50, steady) for o in range(300, 401)} == {300, 350, 400} + assert _capture_offset(301, 50, steady) == 350 + assert _capture_offset(350, 50, steady) == 350 + assert _capture_offset(351, 50, steady) == 400 + assert _capture_offset(400, 50, steady) == 400 + # A 274-frame prompt: steady 374, still exact. + assert {_capture_offset(o, 50, 374) for o in range(274, 375)} == {274, 324, 374} + # The offset-0 prompt solve, no grid, and caches past steady. + assert _capture_offset(0, 50, steady) == 0 + assert _capture_offset(322, 0, steady) == 322 + assert _capture_offset(322, 1, steady) == 322 + assert _capture_offset(410, 50, steady) == 450 + assert _capture_offset(322, 50, 0) == 350 + + +def test_capture_mask_masks_the_offset_grid_padding(): + offset, offset_cap, query_cap, mel_width = 322, 350, 64, 50 + mask = _build_capture_mask( + attn_mask=None, + batch_size=2, + query_cap=query_cap, + offset=offset, + mel_width=mel_width, + mel_frames=mel_width, + device=torch.device("cpu"), + offset_cap=offset_cap, + ) + assert mask.shape == (4, query_cap, query_cap + offset_cap) + assert mask[:, :, query_cap : query_cap + offset].all() + assert not mask[:, :, query_cap + offset :].any() + exact = _build_capture_mask( + attn_mask=None, + batch_size=2, + query_cap=query_cap, + offset=offset, + mel_width=mel_width, + mel_frames=mel_width, + device=torch.device("cpu"), + ) + assert torch.equal(mask[..., : query_cap + offset], exact) + # A caller mask (padded chunk columns masked) moves into the grid layout too. + caller = torch.ones(4, mel_width, mel_width + offset, dtype=torch.bool) + caller[:, :, 40:mel_width] = False + gridded = _build_capture_mask( + attn_mask=caller, + batch_size=2, + query_cap=query_cap, + offset=offset, + mel_width=mel_width, + mel_frames=40, + device=torch.device("cpu"), + offset_cap=offset_cap, + ) + assert not gridded[:, :mel_width, 40:query_cap].any() + assert gridded[:, :mel_width, query_cap : query_cap + offset].all() + assert not gridded[:, :, query_cap + offset :].any() + + def test_capture_query_width_collapses_decode_grid_onto_one_bucket(): """Whole-Euler capture bucket is the #7416 pad helper with a 64-wide decode grid.""" assert _capture_query_width(16, 0) == 16 diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 7519b3b5f83..0cf187a0e1e 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -19,6 +19,7 @@ from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( CFMGraphWrapper, HiFTGraphWrapper, + SharedPromptAttCache, WholeEulerCFMGraphWrapper, ) @@ -103,6 +104,7 @@ def _fake_wrapper(monkeypatch: pytest.MonkeyPatch) -> HiFTGraphWrapper: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) wrapper = object.__new__(HiFTGraphWrapper) wrapper.capture_batch_sizes = [1] + wrapper._legit_shapes = {(7, 0), (9, 0)} wrapper.graph = {} wrapper.static_speech_inputs = {} wrapper.static_cache_source_inputs = {} @@ -166,6 +168,65 @@ def test_multibatch_serial_replay_with_batch1_graph(monkeypatch: pytest.MonkeyPa wrapper.decode_fn.assert_not_called() +def test_shape_outside_capture_buckets_never_captures(monkeypatch: pytest.MonkeyPatch) -> None: + """An arbitrary duplex first-chunk width must never trigger a capture. + + Such a width is seen once and never reused, so capturing it would stall + every stream on the device for a graph that pays for itself once. It must + run eagerly over the whole batch instead, even though the lazy budget has + room and the batch would otherwise qualify for serial fallback. + """ + wrapper = _fake_wrapper(monkeypatch) + speech_feat = torch.randn(3, 80, 14) # 14 frames is not in {7, 9} + cache_source = torch.zeros(3, 1, 0) + + result = wrapper.replay(speech_feat, cache_source) + + wrapper._capture.assert_not_called() + assert wrapper.lazy_graph_count == 0 + wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source) + assert result is wrapper.decode_fn.return_value + + +def test_shape_outside_capture_buckets_ignores_lazy_budget(monkeypatch: pytest.MonkeyPatch) -> None: + """Budget headroom must not matter: only bucket membership does.""" + wrapper = _fake_wrapper(monkeypatch) + wrapper.max_lazy_graphs = 100 + speech_feat = torch.randn(1, 80, 14) + cache_source = torch.zeros(1, 1, 0) + + wrapper.replay(speech_feat, cache_source) + + wrapper._capture.assert_not_called() + wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source) + + +def _bucket_wrapper(extra_codec_chunk_frames: list[int]) -> HiFTGraphWrapper: + wrapper = object.__new__(HiFTGraphWrapper) + wrapper.codec_chunk_frames = 75 + wrapper.initial_codec_chunk_frames = 25 + wrapper.codec_left_context_frames = 3 + wrapper.pre_lookahead_len = 3 + wrapper.flow_upsample_rate = 2 + wrapper.mel_cache_len = 8 + wrapper.source_cache_len = 3840 + wrapper.extra_codec_chunk_frames = extra_codec_chunk_frames + return wrapper + + +def test_hift_buckets_default_cover_turn_chunks_and_first_chunk() -> None: + frames, cache = _bucket_wrapper([]).derive_capture_bucket_size() + + assert list(zip(frames, cache, strict=True)) == [(150, 0), (158, 3840), (50, 0)] + + +def test_hift_buckets_cover_duplex_unit_continuations() -> None: + """A 25-token duplex unit after the first one carries the mel and source caches.""" + frames, cache = _bucket_wrapper([25, 75]).derive_capture_bucket_size() + + assert list(zip(frames, cache, strict=True)) == [(150, 0), (158, 3840), (50, 0), (58, 3840)] + + # --------------------------------------------------------------------------- # CFMGraphWrapper tests # --------------------------------------------------------------------------- @@ -1007,13 +1068,38 @@ def mock_fail(*args, **kwargs): 16: [(8, 8), (8, 8)], }, ), + # A pinned 12 tier absorbs staggered remainders as one padded replay. + ( + 16, + 3, + { + 9: [(12, 9)], + 11: [(12, 11)], + 13: [(16, 13)], + 17: [(16, 16), (1, 1)], + 21: [(16, 16), (8, 5)], + }, + ), + # The pinned grid stops at micro_batch_size (a batch-4 graph could never replay). + ( + 2, + 0, + { + 1: [(1, 1)], + 2: [(2, 2)], + 3: [(2, 2), (1, 1)], + 5: [(2, 2), (2, 2), (1, 1)], + }, + ), ], ) def test_whole_euler_plan_groups_uses_power_of_two_graph_batches( micro: int, pad_max_rows: int, plans: dict[int, list[tuple[int, int]]], + monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.setenv("VLLM_OMNI_CFM_GRAPH_GRID", "1,2,4,8,12,16") wrapper = WholeEulerCFMGraphWrapper( estimator=_WholeEulerDiT(), n_timesteps=10, @@ -1027,6 +1113,77 @@ def test_whole_euler_plan_groups_uses_power_of_two_graph_batches( assert sum(rows for _, rows in plan) == batch_size +@pytest.mark.parametrize( + ("micro", "grid", "sizes"), + [ + # The powers of two below micro, and micro: never a tier above it (a + # batch that large splits into micro-batches), even for micro 2 or 3. + (1, None, [1]), + (2, None, [1, 2]), + (3, None, [1, 2, 3]), + (16, None, [1, 2, 4, 8, 16]), + (20, None, [1, 2, 4, 8, 16, 20]), + # A pinned grid (the H7 and G2 lanes) drops the sizes above micro. + (20, "1,4,8,16,20", [1, 4, 8, 16, 20]), + (16, "1,4,8,16,20", [1, 4, 8, 16]), + (16, "1,4,8,16", [1, 4, 8, 16]), + (2, "4,8", [2]), + ], +) +def test_whole_euler_graph_batches_default_and_pinned_grid( + micro: int, grid: str | None, sizes: list[int], monkeypatch: pytest.MonkeyPatch +) -> None: + if grid is None: + monkeypatch.delenv("VLLM_OMNI_CFM_GRAPH_GRID", raising=False) + else: + monkeypatch.setenv("VLLM_OMNI_CFM_GRAPH_GRID", grid) + wrapper = WholeEulerCFMGraphWrapper( + estimator=_WholeEulerDiT(), n_timesteps=10, max_graphs=8, micro_batch_size=micro + ) + assert wrapper._graph_batches() == sizes + + +def test_whole_euler_arena_rows_can_follow_pinned_graph_grid(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("VLLM_OMNI_CFM_GRAPH_GRID", "1,4,8,16") + wrapper = WholeEulerCFMGraphWrapper( + estimator=_WholeEulerDiT(), + n_timesteps=10, + max_graphs=8, + max_graph_batch=20, + micro_batch_size=20, + arena_rows_from_graph_grid=True, + ) + assert wrapper._arena_rows == 16 + assert wrapper._plan_groups(20) == [(16, 16), (4, 4)] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_shared_prompt_cache_is_bit_exact(monkeypatch: pytest.MonkeyPatch) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + torch.manual_seed(17) + estimator = _WholeEulerDiT().eval().cuda() + prompt = torch.randn(10, 2, 2, 2, 12, 8, device="cuda") + shared = SharedPromptAttCache(prompt[..., :8, :], prompt[..., 8:, :]) + inputs = _whole_euler_chunk(1, 4) + + def run(att_cache): + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=4, + query_bucket_frames=16, + ) + result = wrapper.replay(**inputs, cnn_cache=None, att_cache=[att_cache], att_keep=(12, 4)) + assert result is not None + return result[0], result[1], result[2][0].materialize() if isinstance(result[2][0], SharedPromptAttCache) else result[2][0] + + tensor_result = run(prompt) + shared_result = run(shared) + for actual, expected in zip(shared_result, tensor_result, strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") @pytest.mark.parametrize(("pad_max_rows", "graph_batches"), [(1, {4}), (0, {1, 2})]) def test_whole_euler_multibatch_remainder_replay( @@ -2013,6 +2170,132 @@ def _mask(width: int, valid: int, offset: int) -> torch.Tensor: bucketed._flush() +def _run_stream(wrapper, widths, *, att_keep): + """Replay one stream's chunks (full-width, unpadded) and return every chunk's outputs.""" + torch.manual_seed(11) + chunks = [_whole_euler_chunk(1, width) for width in widths] + cnn = att = None + outputs = [] + for chunk in chunks: + result = wrapper.replay( + **{key: value.clone() for key, value in chunk.items()}, + cnn_cache=cnn, + att_cache=att, + att_keep=att_keep, + ) + assert result is not None + assert all(torch.isfinite(t).all() for t in result) + outputs.append(result) + cnn, att = result[1], result[2] + return outputs + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_offset_grid_matches_exact_offsets(monkeypatch: pytest.MonkeyPatch) -> None: + """Caches padded onto the offset grid solve like exact ones, on fewer graphs. + + The stream's first chunks leave cache lengths 16, 24 and 40 on the way to + the steady 64; on a 16-frame grid anchored at 64 those run as 16, 32 and + 48 with the padding masked out, and every output, cache included, matches. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + # The shipped DiT with the ragged body, as in serving: padded query frames + # then leave the CNN cache the unpadded solve would. + dit = _tiny_upstream_dit() + ragged = BatchedToken2Wav._blocks_forward_chunk_ragged + exact = WholeEulerCFMGraphWrapper( + estimator=dit, n_timesteps=10, max_graphs=16, query_bucket_frames=16, ragged_body=ragged + ) + gridded = WholeEulerCFMGraphWrapper( + estimator=dit, + n_timesteps=10, + max_graphs=16, + query_bucket_frames=16, + offset_bucket_frames=16, + ragged_body=ragged, + ) + widths = (16, 8, 16, 16, 16) + reference = _run_stream(exact, widths, att_keep=(16, 48)) + results = _run_stream(gridded, widths, att_keep=(16, 48)) + for ref_chunk, chunk in zip(reference, results, strict=True): + for ref_value, value in zip(ref_chunk, chunk, strict=True): + torch.testing.assert_close(value, ref_value, rtol=1e-4, atol=1e-5) + assert sorted({key[3] for key in exact._cache}) == [0, 16, 24, 40, 56] + assert sorted({key[3] for key in gridded._cache}) == [0, 16, 32, 48, 64] + exact._flush() + gridded._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_query_width_tuple_matches_one_wide_width(monkeypatch: pytest.MonkeyPatch) -> None: + """A narrow chunk on its own narrow graph solves like it does padded onto the wide one.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + # The shipped DiT with the ragged body, as in serving: padded query frames + # then leave the CNN cache the unpadded solve would. + dit = _tiny_upstream_dit() + ragged = BatchedToken2Wav._blocks_forward_chunk_ragged + wide = WholeEulerCFMGraphWrapper( + estimator=dit, n_timesteps=10, max_graphs=16, query_bucket_frames=32, ragged_body=ragged + ) + tiered = WholeEulerCFMGraphWrapper( + estimator=dit, n_timesteps=10, max_graphs=16, query_bucket_frames=(8, 32), ragged_body=ragged + ) + widths = (32, 8, 8, 32) + reference = _run_stream(wide, widths, att_keep=(32, 64)) + results = _run_stream(tiered, widths, att_keep=(32, 64)) + for ref_chunk, chunk in zip(reference, results, strict=True): + for ref_value, value in zip(ref_chunk, chunk, strict=True): + torch.testing.assert_close(value, ref_value, rtol=1e-4, atol=1e-5) + assert {key[2] for key in wide._cache} == {32} + assert {key[2] for key in tiered._cache} == {8, 32} + wide._flush() + tiered._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_precapture_leaves_nothing_to_capture_while_serving(monkeypatch: pytest.MonkeyPatch) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + # The shipped DiT with the ragged body, as in serving: padded query frames + # then leave the CNN cache the unpadded solve would. + dit = _tiny_upstream_dit() + ragged = BatchedToken2Wav._blocks_forward_chunk_ragged + wrapper = WholeEulerCFMGraphWrapper( + estimator=dit, + n_timesteps=10, + max_graphs=32, + micro_batch_size=2, + query_bucket_frames=(8, 16), + offset_bucket_frames=16, + ragged_body=ragged, + ) + chunk = _whole_euler_chunk(1, 8) + captured = wrapper.precapture( + offsets=range(16, 65), + steady=64, + channels=int(chunk["x"].shape[1]), + spk_dim=int(chunk["speakers_cfg"].shape[1]), + ) + # Graph batches 1 and 2, widths 8 and 16, offsets 16/32/48/64. + assert captured == 2 * 2 * 4 + storage = wrapper.arena._att + before = wrapper.stats_snapshot() + # The stream's first chunk (offset 0) is the prompt solve, captured on + # first use like the prompt features themselves; everything after replays. + widths = (16, 8, 16, 8, 16) + _run_stream(wrapper, widths, att_keep=(16, 48)) + after = wrapper.stats_snapshot() + assert after["captures"] - before["captures"] == 1 + assert after["flushes"] == before["flushes"] == 0 + assert wrapper.arena._att is storage + wrapper._flush() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") def test_whole_euler_skipped_when_trt_stepper_configured() -> None: from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( @@ -2218,6 +2501,49 @@ def test_whole_euler_ragged_rows_match_per_row_exact_solves(monkeypatch: pytest. wrapper._flush() +def test_whole_euler_padded_rows_do_not_keep_a_wider_graphs_lengths() -> None: + """A padded ragged replay gets in-range lengths on its padded rows. + + Every query width of one graph batch shares the lengths buffer. The ragged + body gathers each row's causal cache at ``length + arange(width)``, so a 56 + left by a 64-frame graph reads past a 32-frame graph's conv history (a + device-side assert in serving). + """ + graph_batch, rows, channels, query_cap = 4, 3, 2, 32 + + def zeros(*shape: int) -> torch.Tensor: + return torch.zeros(shape) + + statics = wrapper_module._WholeEulerStatics( + x=zeros(graph_batch, channels, query_cap), + mu_cfg=zeros(2 * graph_batch, channels, query_cap), + speakers_cfg=zeros(2 * graph_batch, 3), + cond_cfg=zeros(2 * graph_batch, channels, query_cap), + cnn_cache=zeros(1, 1, 2 * graph_batch, 1, 1), + att_cache=None, + attn_mask=torch.ones((2 * graph_batch, query_cap, query_cap), dtype=torch.bool), + time_embeddings=None, + lengths=torch.full((2 * graph_batch,), 56, dtype=torch.long), + ) + WholeEulerCFMGraphWrapper._fill_static_inputs( + statics, + graph_batch=graph_batch, + start=0, + stop=rows, + x=zeros(rows, channels, query_cap), + mu_cfg=zeros(2 * rows, channels, query_cap), + speakers_cfg=zeros(2 * rows, 3), + cond_cfg=zeros(2 * rows, channels, query_cap), + attn_mask=torch.ones((2 * rows, query_cap, query_cap), dtype=torch.bool), + cnn_cache=None, + att_cache=None, + att_rows=None, + offset=0, + lengths=torch.tensor((30, 20, 10, 30, 20, 10)), + ) + assert statics.lengths.tolist() == [30, 20, 10, 0, 30, 20, 10, 0] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") def test_whole_euler_padded_ragged_replay_after_a_wider_replay(monkeypatch: pytest.MonkeyPatch) -> None: """A padded ragged replay is exact after a wider replay of the same graph batch. @@ -2315,7 +2641,7 @@ def fixed_copy(cache: torch.Tensor) -> torch.Tensor: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_whole_euler_groups_are_acquired_again_after_a_flush(monkeypatch: pytest.MonkeyPatch) -> None: +def test_whole_euler_groups_are_acquired_again_after_a_flush() -> None: """Capturing a later group can retire an earlier one; no group replays until all are held at once.""" wrapper = WholeEulerCFMGraphWrapper(estimator=_tiny_upstream_dit(), n_timesteps=10, max_graphs=8) calls: list[int] = [] @@ -2326,13 +2652,12 @@ def entry(*, graph_batch: int, fill, **_: object) -> tuple: wrapper._stats["flushes"] += 1 return ("entry", graph_batch, len(calls)) - monkeypatch.setattr(wrapper, "_entry", entry) flush_always = False groups, fills = [(16, 16), (1, 1)], [None, None] - assert wrapper._group_entries(groups, fills) == [("entry", 16, 3), ("entry", 1, 4)] + assert wrapper._group_entries(entry, groups, fills) == [("entry", 16, 3), ("entry", 1, 4)] assert calls == [16, 1, 16, 1] flush_always = True - assert wrapper._group_entries(groups, fills) is None + assert wrapper._group_entries(entry, groups, fills) is None @pytest.mark.parametrize("fail", [False, True]) diff --git a/tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py b/tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py new file mode 100644 index 00000000000..b3c5eff109d --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py @@ -0,0 +1,637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exact-shape flow-encoder graphs (``cfm_encoder_cuda_graph``): keys, shared buffers, backend wiring. + +CPU only: ``FlowEncoderGraphs._record`` is replaced by a fake graph whose +replay runs the recorded body, so these tests cover everything around the CUDA +capture itself -- which shapes get a graph, how inputs reach the shared static +buffers, that results match the eager encoder exactly, and that the backend +copies shared results it keeps past the next replay. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +import vllm_omni.model_executor.models.minicpmo_4_5.flow_encoder_graph as graph_module +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav +from vllm_omni.model_executor.models.minicpmo_4_5.flow_encoder_graph import ( + FlowEncoderGraphs, + reachable_conformer_cache_frames, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DEPTH, _HEADS, _WIDTH = 2, 2, 4 +_CNN = (3, 2) +_LOOKAHEAD, _UPSAMPLE = 1, 2 + + +class _FakeGraph: + def __init__(self, graphs: FlowEncoderGraphs, entry) -> None: + self.graphs = graphs + self.entry = entry + self.replays = 0 + + def replay(self) -> None: + self.replays += 1 + self.graphs._body(self.entry) + + +@pytest.fixture(autouse=True) +def _fake_capture(monkeypatch: pytest.MonkeyPatch): + recorded: list[tuple[int, ...]] = [] + + def record(self, entry): + recorded.append((int(entry.tokens.shape[0]), int(entry.tokens.shape[1]), int(entry.att.shape[3]))) + return _FakeGraph(self, entry) + + monkeypatch.setattr(FlowEncoderGraphs, "_record", record) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return recorded + + +class _Encoder(nn.Module): + """Row-independent stand-in for ``UpsampleConformerEncoderV2.forward_chunk`` with its cache layouts. + + ``cnn_cache`` ``(rows, 3, 2)``, ``att_cache`` ``(depth, rows, heads, frames, width)``, + a 1-token lookahead and a 2x upsample; every output depends on the row's tokens and caches. + """ + + def __init__(self) -> None: + super().__init__() + self.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=_LOOKAHEAD) + self.up_layer = SimpleNamespace(stride=_UPSAMPLE) + self.calls: list[tuple[int, int, bool]] = [] + + def forward_chunk(self, xs, last_chunk=False, cnn_cache=None, att_cache=None): + rows, tokens, _ = xs.shape + self.calls.append((rows, tokens, last_chunk)) + if last_chunk: + xs = F.pad(xs, (0, 0, 0, _LOOKAHEAD)) + out = xs[:, :-_LOOKAHEAD].repeat_interleave(_UPSAMPLE, dim=1) + if cnn_cache is not None: + out = out + cnn_cache.sum(dim=(1, 2))[:, None, None] + if att_cache is not None: + out = out + 0.5 * att_cache.mean(dim=(0, 2, 3, 4))[:, None, None] + new_cnn = out[:, -_CNN[1] :, :1].transpose(1, 2).expand(-1, _CNN[0], -1).contiguous() + frames = out[:, :, 0].transpose(0, 1).reshape(1, -1, rows, 1, 1) + added = (frames.permute(0, 2, 3, 1, 4) * torch.arange(1, _WIDTH + 1, dtype=out.dtype)).expand( + _DEPTH, rows, _HEADS, -1, _WIDTH + ) + new_att = added.contiguous() if att_cache is None else torch.cat((att_cache, added), dim=3) + return out, new_cnn, new_att + + +class _Estimator(nn.Module): + def __init__(self) -> None: + super().__init__() + conv1 = SimpleNamespace(causal_padding=(1, 0)) + block = SimpleNamespace( + conv=SimpleNamespace(in_channels=1, out_channels=1, block=[None, conv1]), + attn=SimpleNamespace(num_heads=1, head_dim=1), + ) + self.blocks = [block] + self.in_proj = nn.Identity() + + def t_embedder(self, time): + return time[:, None] + + def blocks_forward_chunk(self, inputs, time, mask, cnn_cache, att_cache, cnn_out, att_out): + del time, mask, cnn_cache, att_cache + marker = inputs[:, 1, 0] + cnn_out.copy_(marker.reshape(1, -1, 1, 1).expand_as(cnn_out)) + att_out.copy_(marker.reshape(1, -1, 1, 1, 1).expand_as(att_out)) + return inputs[:, 1:2] + 0.25 * inputs[:, 0:1] + + +class _Decoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.estimator = _Estimator() + self.inference_cfg_rate = 0.7 + self.register_buffer("rand_noise", torch.linspace(-1, 1, 600).reshape(1, 1, 600), persistent=False) + + +class _Flow(nn.Module): + def __init__(self) -> None: + super().__init__() + self.encoder = _Encoder() + self.encoder_proj = nn.Linear(1, 1) + with torch.no_grad(): + self.encoder_proj.weight.fill_(0.5) + self.encoder_proj.bias.fill_(0.125) + self.decoder = _Decoder() + self.spk_embed_affine_layer = nn.Identity() + + def input_embedding(self, tokens): + return tokens.to(torch.float32).unsqueeze(-1) / 16 + + +class _HiFT(nn.Module): + def inference(self, mel, source): + del source + speech = mel[:, 0].repeat_interleave(3, dim=1) + return speech, speech[:, None] + + +class _Token2Wav: + def __init__(self) -> None: + self.flow = _Flow() + self.hift = _HiFT() + self.float16 = False + self.n_timesteps = 2 + self.mel_cache_len = 1 + self.source_cache_len = 2 + self.speech_window = torch.hamming_window(4, periodic=False) + + def _prepare_prompt(self, prompt_wav): + del prompt_wav + return ( + torch.tensor([[5, 6, 7, 8]], dtype=torch.long), + torch.tensor([4], dtype=torch.int32), + torch.ones(1, 1), + torch.ones(1, 8, 1), + torch.tensor([8], dtype=torch.int32), + ) + + +def _ragged_kernel(adapter: BatchedToken2Wav) -> None: + def kernel(estimator, estimator_input, time_embedding, attn_mask, cnn_cache, att_cache, cnn_buf, att_buf, lengths): + del lengths + return estimator.blocks_forward_chunk( + estimator_input, time_embedding, attn_mask, cnn_cache, att_cache, cnn_buf, att_buf + ) + + adapter._blocks_forward_chunk_ragged = kernel # type: ignore[method-assign] + + +def _backend(*, rows=(1, 2), widths=(6,)) -> BatchedToken2Wav: + adapter = BatchedToken2Wav(_Token2Wav()) + _ragged_kernel(adapter) + if rows: + adapter._encoder_graphs = FlowEncoderGraphs( + adapter._encode_continuation, + rows=rows, + token_widths=widths, + lookahead=_LOOKAHEAD, + upsample=_UPSAMPLE, + ) + return adapter + + +def _tokens(chunk: int, rows: int, width: int) -> torch.Tensor: + return torch.stack([torch.arange(width) * (row + 1) + 7 * chunk + row for row in range(rows)]) + + +def _assert_states_equal(actual, expected) -> None: + for got, want in zip(actual, expected, strict=True): + for name, value in want.flow_cache.items(): + assert torch.equal(got.flow_cache[name], value), name + for name, value in want.hift_cache.items(): + assert torch.equal(got.hift_cache[name], value), name + + +# --------------------------------------------------------------------------- +# Keys + + +@pytest.mark.parametrize( + ("growths", "expected"), + [ + ([50], [300, 350, 400]), # 6 s default voice, 28-token duplex units + ([150], [300, 400]), # 78-token turn chunks + ([50, 150], [300, 350, 400]), + ([30], [300, 330, 360, 390, 400]), + ([], [300]), + ], +) +def test_reachable_cache_frames_follow_the_streaming_trim(growths, expected) -> None: + assert reachable_conformer_cache_frames(start=300, prompt_len=300, suffix=100, growths=growths) == expected + + +def test_keys_cover_every_row_count_width_and_cache_length() -> None: + graphs = FlowEncoderGraphs(lambda *a: a, rows=[2, 1, 2], token_widths=[28], lookahead=3, upsample=2) + + assert graphs.rows == (1, 2) + assert graphs.output_frames(28) == 50 + assert graphs.keys_for(start=300, prompt_len=300, suffix=100) == [ + (1, 28, 300), + (1, 28, 350), + (1, 28, 400), + (2, 28, 300), + (2, 28, 350), + (2, 28, 400), + ] + + +def test_widths_without_new_frames_are_dropped() -> None: + graphs = FlowEncoderGraphs(lambda *a: a, rows=[1], token_widths=[3, 28], lookahead=3, upsample=2) + assert graphs.token_widths == (28,) + + +# --------------------------------------------------------------------------- +# FlowEncoderGraphs on a stand-in encoder + + +def _graphs(rows=(1, 2, 3), widths=(6,), frames=(8, 18, 28)): + encoder = _Encoder() + proj = nn.Linear(1, 1) + + def encode(tokens, cnn, att): + hidden, new_cnn, new_att = encoder.forward_chunk( + tokens.to(torch.float32).unsqueeze(-1), last_chunk=False, cnn_cache=cnn, att_cache=att + ) + return proj(hidden), new_cnn, new_att + + graphs = FlowEncoderGraphs(encode, rows=rows, token_widths=widths, lookahead=_LOOKAHEAD, upsample=_UPSAMPLE) + keys = [(r, w, f) for r in rows for w in widths for f in frames] + graphs.capture( + keys, cnn_shape=_CNN, att_layout=(_DEPTH, _HEADS, _WIDTH), hidden_dim=1, dtype=torch.float32, device="cpu" + ) + return graphs, encode + + +def _inputs(rows: int, width: int, frames: int, seed: int): + generator = torch.Generator().manual_seed(seed) + tokens = torch.randint(0, 50, (rows, width), generator=generator) + cnn = [torch.randn(1, *_CNN, generator=generator) for _ in range(rows)] + att = [torch.randn(_DEPTH, 1, _HEADS, frames, _WIDTH, generator=generator) for _ in range(rows)] + return tokens, cnn, att + + +@pytest.mark.parametrize(("rows", "frames"), [(1, 8), (2, 18), (3, 28), (3, 8)]) +def test_replay_matches_the_eager_encoder_exactly(rows, frames) -> None: + graphs, encode = _graphs() + tokens, cnn, att = _inputs(rows, 6, frames, seed=rows * 100 + frames) + + graphed = graphs.run(tokens, cnn, att) + expected = encode(tokens, torch.cat(cnn, dim=0), torch.cat(att, dim=1)) + + assert graphed is not None + for actual, reference in zip(graphed, expected, strict=True): + assert actual.shape == reference.shape + assert torch.equal(actual, reference) + assert graphs.stats["replays"] == 1 + + +@pytest.mark.parametrize( + ("rows", "width", "frames"), + [ + (4, 6, 8), # more rows than captured: never padded down or split + (1, 5, 8), # another token width + (1, 6, 10), # a cache length no prompt stream reaches + ], +) +def test_shapes_without_their_own_graph_run_eager(rows, width, frames) -> None: + graphs, _ = _graphs() + tokens, cnn, att = _inputs(rows, width, frames, seed=1) + + assert graphs.run(tokens, cnn, att) is None + assert graphs.stats["replays"] == 0 + + +def test_mismatched_cache_layouts_run_eager() -> None: + graphs, _ = _graphs() + tokens, cnn, att = _inputs(2, 6, 8, seed=2) + + assert graphs.run(tokens, cnn, [att[0], att[1].double()]) is None + assert graphs.run(tokens, [cnn[0], cnn[1][:, :2]], att) is None + assert graphs.run(tokens.int(), cnn, att) is None + assert graphs.run(tokens[:1], cnn, att) is None + assert graphs.run(tokens, cnn, att) is not None + + +def test_graphs_share_one_storage_per_buffer_and_capture_the_largest_first(_fake_capture) -> None: + graphs, _ = _graphs() + + assert _fake_capture[0] == (3, 6, 28) + assert _fake_capture[-1] == (1, 6, 8) + for name in ("tokens", "cnn", "att", "hidden", "new_cnn", "new_att"): + pointers = {getattr(entry, name).data_ptr() for entry in graphs.graphs.values()} + assert pointers == {graphs._storage[name].data_ptr()} + largest = graphs.graphs[(3, 6, 28)] + assert graphs._storage["att"].numel() == largest.att.numel() + assert graphs._storage["new_att"].numel() == largest.new_att.numel() + assert graphs.storage_bytes() == sum(t.numel() * t.element_size() for t in graphs._storage.values()) + + +def test_results_are_shared_views_valid_until_the_next_replay() -> None: + graphs, encode = _graphs() + first_inputs = _inputs(1, 6, 8, seed=3) + first = graphs.run(*first_inputs) + kept = [value.clone() for value in first] + + graphs.run(*_inputs(1, 6, 8, seed=4)) + + assert not torch.equal(first[0], kept[0]) # rewritten in place by the second replay + expected = encode(first_inputs[0], torch.cat(first_inputs[1]), torch.cat(first_inputs[2], dim=1)) + for value, reference in zip(kept, expected, strict=True): + assert torch.equal(value, reference) + + +def test_a_precision_change_after_capture_runs_eager(monkeypatch: pytest.MonkeyPatch) -> None: + graphs, _ = _graphs() + inputs = _inputs(1, 6, 8, seed=5) + monkeypatch.setattr(graph_module, "_precision_state", lambda: ("tf32",)) + + assert graphs.run(*inputs) is None + assert graphs.stats["eager_precision"] == 1 + assert graphs.enabled + + +def test_a_replaced_relpos_table_disables_the_graphs() -> None: + table = {"pe": torch.zeros(3)} + graphs = FlowEncoderGraphs( + lambda tokens, cnn, att: (tokens.float()[:, :, None].repeat(1, 2, 1)[:, :-2], cnn, att), + rows=[1], + token_widths=[2], + lookahead=1, + upsample=2, + held_tensors=lambda: (table["pe"],), + ) + graphs.capture( + [(1, 2, 4)], + cnn_shape=_CNN, + att_layout=(_DEPTH, _HEADS, _WIDTH), + hidden_dim=1, + dtype=torch.float32, + device="cpu", + ) + assert graphs._held == (table["pe"],) + + table["pe"] = torch.zeros(3) + tokens, cnn, att = _inputs(1, 2, 4, seed=6) + + assert graphs.run(tokens, cnn, att) is None + assert not graphs.enabled + + +def test_capture_is_all_or_nothing(monkeypatch: pytest.MonkeyPatch) -> None: + calls = {"n": 0} + + def record(self, entry): + calls["n"] += 1 + if calls["n"] == 2: + raise RuntimeError("capture failed") + return _FakeGraph(self, entry) + + monkeypatch.setattr(FlowEncoderGraphs, "_record", record) + graphs = FlowEncoderGraphs(lambda *a: a, rows=[1, 2], token_widths=[6], lookahead=1, upsample=2) + + with pytest.raises(RuntimeError, match="capture failed"): + graphs.capture( + [(1, 6, 8), (2, 6, 8)], + cnn_shape=_CNN, + att_layout=(_DEPTH, _HEADS, _WIDTH), + hidden_dim=1, + dtype=torch.float32, + device="cpu", + ) + assert graphs.graphs == {} + assert graphs._storage == {} + assert not graphs.enabled + + +def test_an_output_layout_mismatch_fails_the_capture() -> None: + graphs = FlowEncoderGraphs( + lambda tokens, cnn, att: (torch.zeros(1, 3, 1), cnn, att), + rows=[1], + token_widths=[6], + lookahead=1, + upsample=2, + ) + graphs._record = lambda entry: graphs._body(entry) # type: ignore[method-assign] + + with pytest.raises(RuntimeError, match="hidden output"): + graphs.capture( + [(1, 6, 8)], + cnn_shape=_CNN, + att_layout=(_DEPTH, _HEADS, _WIDTH), + hidden_dim=1, + dtype=torch.float32, + device="cpu", + ) + + +# --------------------------------------------------------------------------- +# Backend wiring + + +def test_default_width_is_the_duplex_unit_chunk() -> None: + widths = BatchedToken2Wav._default_encoder_token_widths + assert widths({"codec_chunk_frames": 75, "initial_codec_chunk_frames": 25, "codec_left_context_frames": 3}) == [28] + assert widths({"codec_chunk_frames": 25, "codec_left_context_frames": 3}) == [28] + assert widths(None) == [28] + + +def test_encoder_graphs_are_off_by_default_and_off_cuda() -> None: + assert BatchedToken2Wav(_Token2Wav())._encoder_graphs is None + # The stand-in flow lives on the CPU: the switch alone does not build graphs there. + assert BatchedToken2Wav(_Token2Wav(), encoder_graph_config={"enabled": True, "rows": [1]})._encoder_graphs is None + + +def test_build_requires_the_upsample_conformer_encoder(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(BatchedToken2Wav, "_flow_on_cuda", lambda self: True) + adapter = BatchedToken2Wav(_Token2Wav()) + # No RelPos table on the stand-in encoder: stays eager. + assert adapter._build_encoder_graphs({"enabled": True, "rows": [1]}, None) is None + + pe = torch.zeros(1, 9, 4) + adapter.flow.encoder.embed = SimpleNamespace(pos_enc=SimpleNamespace(pe=pe)) + adapter.flow.encoder.up_embed = SimpleNamespace(pos_enc=SimpleNamespace(pe=pe + 1)) + graphs = adapter._build_encoder_graphs( + {"enabled": True, "rows": [1, 2]}, {"initial_codec_chunk_frames": 25, "codec_left_context_frames": 3} + ) + assert graphs is not None + assert graphs.rows == (1, 2) + assert graphs.token_widths == (28,) + assert (graphs.lookahead, graphs.upsample) == (_LOOKAHEAD, _UPSAMPLE) + held = graphs._held_tensors() + assert len(held) == 2 + assert held[0] is pe and held[1] is adapter.flow.encoder.up_embed.pos_enc.pe + + adapter.flow.encoder.embed.extend_pe = lambda x: None + assert adapter._build_encoder_graphs({"enabled": True, "rows": [1]}, None) is None + + +def test_precapture_covers_the_prompt_stream_cache_lengths() -> None: + adapter = _backend(rows=(1, 2), widths=(6,)) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + + assert adapter.precapture_flow_encoder(prompt) == 2 * 11 + # Prompt state: 4 tokens + 1 lookahead -> 8 frames; +10 per 6-token chunk; trimmed at 8 + 100. + assert sorted({key[2] for key in adapter._encoder_graphs.graphs}) == list(range(8, 109, 10)) + assert adapter.precapture_flow_encoder(prompt) == 0 # once + + +def test_decode_batch_with_graphs_is_bitwise_the_eager_decode() -> None: + eager, graphed = _backend(rows=()), _backend(rows=(1, 2)) + prompts = [backend.prepare_prompt("shared", "/fake/prompt.wav") for backend in (eager, graphed)] + graphed.precapture_flow_encoder(prompts[1]) + states = [backend.setup_batch(prompt, 2) for backend, prompt in zip((eager, graphed), prompts, strict=True)] + + for chunk in range(14): # through the cache trim and well into the steady length + last = chunk == 13 + tokens = _tokens(chunk, 2, 6) + outputs = [ + backend.decode_batch(tokens, prompt, state, last_chunk=last) + for backend, prompt, state in zip((eager, graphed), prompts, states, strict=True) + ] + (eager_audio, eager_states), (graph_audio, graph_states) = outputs + for got, want in zip(graph_audio, eager_audio, strict=True): + assert torch.equal(got, want) + _assert_states_equal(graph_states, eager_states) + states = [eager_states, graph_states] + + # Every chunk but the last (lookahead padding) replayed a graph. + assert graphed._encoder_graphs.stats["replays"] == 13 + + +def test_one_row_streams_replay_the_one_row_graph() -> None: + eager, graphed = _backend(rows=()), _backend(rows=(1, 2)) + prompts = [backend.prepare_prompt("shared", "/fake/prompt.wav") for backend in (eager, graphed)] + graphed.precapture_flow_encoder(prompts[1]) + states = [backend.setup_batch(prompt, 1) for backend, prompt in zip((eager, graphed), prompts, strict=True)] + + for chunk in range(4): + tokens = _tokens(chunk, 1, 6) + (eager_audio, eager_states), (graph_audio, graph_states) = ( + backend.decode_batch(tokens, prompt, state, last_chunk=False) + for backend, prompt, state in zip((eager, graphed), prompts, states, strict=True) + ) + assert torch.equal(graph_audio[0], eager_audio[0]) + _assert_states_equal(graph_states, eager_states) + states = [eager_states, graph_states] + assert graphed._encoder_graphs.stats["replays"] == 4 + + +def _ragged_step(backends, prompts, states, rows, last_chunks): + (eager_audio, eager_states), (graph_audio, graph_states) = ( + backend.decode_ragged_batch(rows, prompt, state, last_chunks=last_chunks) + for backend, prompt, state in zip(backends, prompts, states, strict=True) + ) + for got, want in zip(graph_audio, eager_audio, strict=True): + assert torch.equal(got, want) + _assert_states_equal(graph_states, eager_states) + + +@pytest.mark.parametrize("warm_chunks", [0, 11]) +def test_ragged_batch_keeps_each_encoder_group_result_past_the_next_replay(warm_chunks) -> None: + """Two encoder groups (token widths 6 and 4) replay two graphs that share their result buffers.""" + backends = (_backend(rows=()), _backend(rows=(1, 2, 3), widths=(4, 6))) + prompts = [backend.prepare_prompt("shared", "/fake/prompt.wav") for backend in backends] + backends[1].precapture_flow_encoder(prompts[1]) + states = [backend.setup_batch(prompt, 3) for backend, prompt in zip(backends, prompts, strict=True)] + for chunk in range(warm_chunks): # 11 chunks: the steady, trimmed cache length + tokens = _tokens(chunk, 3, 6) + (eager_audio, eager_states), (graph_audio, graph_states) = ( + backend.decode_batch(tokens, prompt, state, last_chunk=False) + for backend, prompt, state in zip(backends, prompts, states, strict=True) + ) + _assert_states_equal(graph_states, eager_states) + states = [eager_states, graph_states] + stats = backends[1]._encoder_graphs.stats + before = stats["replays"] + + six = _tokens(warm_chunks, 2, 6) + _ragged_step(backends, prompts, states, [six[0], six[1], _tokens(warm_chunks, 1, 4)[0]], [False] * 3) + + assert stats["replays"] - before == 2 + + +def test_ragged_final_rows_stay_eager_next_to_graphed_ones() -> None: + backends = (_backend(rows=()), _backend(rows=(1, 2), widths=(4, 6))) + prompts = [backend.prepare_prompt("shared", "/fake/prompt.wav") for backend in backends] + backends[1].precapture_flow_encoder(prompts[1]) + states = [backend.setup_batch(prompt, 3) for backend, prompt in zip(backends, prompts, strict=True)] + + six = _tokens(0, 2, 6) + _ragged_step(backends, prompts, states, [six[0], six[1], _tokens(0, 1, 4)[0]], [False, False, True]) + + assert backends[1]._encoder_graphs.stats["replays"] == 1 + + +# --------------------------------------------------------------------------- +# Code2Wav connector extras + + +def _code2wav(extra: dict, max_num_seqs: int | None = 20): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import MiniCPMO45Code2Wav + + config = SimpleNamespace( + model_config=SimpleNamespace(model="/fake/model", stage_connector_config={"extra": dict(extra)}), + ) + if max_num_seqs is not None: + config.scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) + return MiniCPMO45Code2Wav(vllm_config=config) + + +@pytest.mark.parametrize( + ("extra", "max_num_seqs", "expected"), + [ + ({}, 20, {"enabled": False, "rows": list(range(1, 9)), "token_widths": None}), + ({"cfm_encoder_cuda_graph": True}, 20, {"enabled": True, "rows": list(range(1, 9)), "token_widths": None}), + ({"cfm_encoder_cuda_graph": True, "cfm_encoder_graph_rows": 4}, 2, {"enabled": True, "rows": [1, 2]}), + ({"cfm_encoder_graph_rows": [4, 1, 4]}, None, {"enabled": False, "rows": [1, 4]}), + ({"cfm_encoder_cuda_graph": True, "cfm_encoder_graph_token_widths": [28]}, 20, {"token_widths": [28]}), + ({"cfm_encoder_cuda_graph": "true"}, 20, {"enabled": False}), + ], +) +def test_encoder_graph_extras_reach_the_backend_config(extra, max_num_seqs, expected) -> None: + config = _code2wav(extra, max_num_seqs)._encoder_graph_config + assert {key: config[key] for key in expected} == expected + + +def test_encoder_graph_rows_must_be_positive() -> None: + with pytest.raises(ValueError, match="cfm_encoder_graph_rows"): + _code2wav({"cfm_encoder_graph_rows": [0, 1]}) + + +@pytest.mark.parametrize("fails", [False, True]) +def test_precapture_captures_the_encoder_graphs_after_whole_euler(fails) -> None: + model = _code2wav({"cfm_encoder_cuda_graph": True}) + calls: list[str] = [] + features = object() + + def encoder(got): + assert got is features + calls.append("encoder") + if fails: + raise RuntimeError("capture failed") + return 3 + + model.backend = SimpleNamespace( + precapture_hift=lambda: calls.append("hift") or 0, + prepare_prompt=lambda *args: features, + precapture_whole_euler=lambda got: calls.append("whole_euler") or 0, + precapture_flow_encoder=encoder, + _encoder_graphs=object(), + ) + model._default_prompt_normalized = ("/fake/prompt.wav", "shared") + + model._precapture_default_prompt() # a failed encoder capture only logs: the encoder stays eager + + assert calls == ["hift", "whole_euler", "encoder"] + + +def test_precapture_skips_the_encoder_without_graphs() -> None: + model = _code2wav({}) + calls: list[str] = [] + model.backend = SimpleNamespace( + precapture_hift=lambda: 0, + prepare_prompt=lambda *args: object(), + precapture_whole_euler=lambda got: 0, + precapture_flow_encoder=lambda got: calls.append("encoder"), + _encoder_graphs=None, + ) + model._default_prompt_normalized = ("/fake/prompt.wav", "shared") + + model._precapture_default_prompt() + + assert calls == [] diff --git a/tests/model_executor/models/minicpmo_4_5/test_hift_exact_graphs.py b/tests/model_executor/models/minicpmo_4_5/test_hift_exact_graphs.py new file mode 100644 index 00000000000..2b9f033072c --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_hift_exact_graphs.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exact-shape HiFT graphs: the duplex vocoder shapes outside the capture buckets (CPU, fake graphs).""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn + +import vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper as wrapper_module +from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import HiFTGraphWrapper, codec_frame_range + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MEL_CACHE = 8 +_SOURCE_CACHE = 3840 + + +class _HiFT(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv_pre = nn.Conv1d(80, 4, 1) + self.inference = Mock(return_value=(torch.tensor([[99.0]]), torch.tensor([[[98.0]]]))) + self.finalize_calls: list[tuple[int, ...]] = [] + + def _inference_pre_istft(self, speech_feat, cache_source): # pragma: no cover - fake graphs never run it + raise AssertionError("the fake capture never runs the graph function") + + def _finalize_decode(self, magnitude, phase): + self.finalize_calls.append(tuple(magnitude.shape)) + return magnitude + phase + + +def _token2wav() -> SimpleNamespace: + return SimpleNamespace( + hift=_HiFT(), + flow=SimpleNamespace( + encoder=SimpleNamespace(pre_lookahead_layer=SimpleNamespace(pre_lookahead_len=3)), + token_mel_ratio=2, + ), + mel_cache_len=_MEL_CACHE, + source_cache_len=_SOURCE_CACHE, + ) + + +def _connector(**extra) -> dict: + return { + "codec_chunk_frames": 75, + "initial_codec_chunk_frames": 25, + "codec_left_context_frames": 3, + "hift_graph_codec_chunk_frames": [25, 75], + **extra, + } + + +class _FakeGraph: + """Replays like the captured HiFT prefix: magnitude = 2 * mel[:, :1], phase = 0, source = mel[:, :1].""" + + def __init__(self, wrapper: HiFTGraphWrapper, key: tuple[int, int, int]) -> None: + self.wrapper = wrapper + self.key = key + self.replays = 0 + + def replay(self) -> None: + self.replays += 1 + mel = self.wrapper.static_speech_inputs[self.key] + self.wrapper.static_magnitude_outputs[self.key].copy_(2 * mel[:, :1]) + self.wrapper.static_phase_outputs[self.key].zero_() + self.wrapper.static_cache_source_outputs[self.key].copy_(mel[:, :1]) + + +def _wrapper(monkeypatch: pytest.MonkeyPatch, **extra) -> HiFTGraphWrapper: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = HiFTGraphWrapper(_token2wav(), _connector(**extra), capture_batch_sizes=[1, 2]) + + def capture(batch_size: int, num_frames: int, cache_len: int) -> None: + key = (batch_size, num_frames, cache_len) + if key in wrapper.graph: + return + wrapper.graph[key] = _FakeGraph(wrapper, key) + wrapper.static_speech_inputs[key] = torch.zeros(batch_size, 80, num_frames) + wrapper.static_cache_source_inputs[key] = torch.zeros(batch_size, 1, cache_len) + wrapper.static_magnitude_outputs[key] = torch.zeros(batch_size, 1, num_frames) + wrapper.static_phase_outputs[key] = torch.zeros(batch_size, 1, num_frames) + wrapper.static_cache_source_outputs[key] = torch.zeros(batch_size, 1, num_frames) + + wrapper._capture = Mock(side_effect=capture) + return wrapper + + +@pytest.mark.parametrize( + ("value", "expected"), + [(None, range(0)), ([], range(0)), ((), range(0)), ([1, 30], range(1, 31)), ([26, 26], range(26, 27))], +) +def test_codec_frame_range_parses_inclusive_ranges(value, expected) -> None: + assert codec_frame_range(value, name="x") == expected + + +@pytest.mark.parametrize("value", [[0, 3], [5, 2], [1], [1, 2, 3], 7, "1,30"]) +def test_codec_frame_range_rejects_malformed_ranges(value) -> None: + with pytest.raises(ValueError, match="hift_graph_first_chunk_frames"): + codec_frame_range(value, name="hift_graph_first_chunk_frames") + + +def test_exact_shapes_are_off_by_default(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _wrapper(monkeypatch) + + assert wrapper.exact_shapes == [] + assert wrapper._exact_keys == frozenset() + wrapper.capture() + # Only the buckets, at every capture batch size: the default is unchanged. + assert sorted(wrapper.graph) == sorted( + (b, frames, cache) for b in (1, 2) for frames, cache in ((150, 0), (158, 3840), (50, 0), (58, 3840)) + ) + assert wrapper.finalize_fn.__self__.finalize_calls == [] + + +def test_exact_shapes_cover_first_chunks_and_continuations_outside_the_buckets(monkeypatch) -> None: + wrapper = _wrapper( + monkeypatch, + hift_graph_first_chunk_frames=[1, 30], + hift_graph_continuation_frames=[26, 74], + ) + + first = [(2 * f, 0) for f in range(1, 31) if f != 25] # (50, 0) is the first-chunk bucket + continuation = [(_MEL_CACHE + 2 * f, _SOURCE_CACHE) for f in range(26, 75)] + assert wrapper.exact_shapes == first + continuation + assert (14, 0) in wrapper.exact_shapes # dx3's most frequent eager fallback (1, 14, 0) + assert (62, _SOURCE_CACHE) in wrapper.exact_shapes # and (1, 62, 3840) + assert wrapper._exact_keys == frozenset((1, frames, cache) for frames, cache in wrapper.exact_shapes) + + +def test_continuation_range_skips_the_bucket_shapes(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _wrapper(monkeypatch, hift_graph_continuation_frames=[24, 26]) + + # f=25 is the duplex unit continuation bucket (58, 3840). + assert wrapper.exact_shapes == [(56, _SOURCE_CACHE), (60, _SOURCE_CACHE)] + + +def test_capture_takes_exact_shapes_at_the_exact_batch_sizes_only(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _wrapper(monkeypatch, hift_graph_first_chunk_frames=[6, 8]) + hift = wrapper.finalize_fn.__self__ + + wrapper.capture() + + exact = {key for key in wrapper.graph if (key[1], key[2]) in {(12, 0), (14, 0), (16, 0)}} + assert exact == {(1, 12, 0), (1, 14, 0), (1, 16, 0)} + # Each exact shape builds its ISTFT envelope at capture time, not on its first live replay. + assert hift.finalize_calls == [(1, 1, 12), (1, 1, 14), (1, 1, 16)] + + +def test_exact_batch_sizes_are_configurable(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _wrapper(monkeypatch, hift_graph_first_chunk_frames=[7, 7], hift_graph_exact_batch_sizes=[1, 2]) + + wrapper.capture() + + assert {(1, 14, 0), (2, 14, 0)} <= set(wrapper.graph) + + +def test_exact_shape_replays_its_own_graph_unpadded(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _wrapper(monkeypatch, hift_graph_first_chunk_frames=[7, 7]) + wrapper.capture() + wrapper._capture.reset_mock() + speech_feat = torch.randn(1, 80, 14) + + speech, source = wrapper.replay(speech_feat, torch.zeros(1, 1, 0)) + + graph = wrapper.graph[(1, 14, 0)] + assert graph.replays == 1 + assert wrapper.static_speech_inputs[(1, 14, 0)].shape == (1, 80, 14) + torch.testing.assert_close(speech, 2 * speech_feat[:, :1], rtol=0, atol=0) + torch.testing.assert_close(source, speech_feat[:, :1], rtol=0, atol=0) + wrapper._capture.assert_not_called() + wrapper.decode_fn.assert_not_called() + + +def test_exact_shape_at_other_batch_sizes_stays_eager(monkeypatch: pytest.MonkeyPatch) -> None: + """Only the captured batch sizes replay: no padding to a larger graph, no lazy capture.""" + wrapper = _wrapper(monkeypatch, hift_graph_first_chunk_frames=[7, 7]) + wrapper.capture() + wrapper._capture.reset_mock() + speech_feat = torch.randn(2, 80, 14) + cache_source = torch.zeros(2, 1, 0) + + result = wrapper.replay(speech_feat, cache_source) + + wrapper._capture.assert_not_called() + wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source) + assert result is wrapper.decode_fn.return_value + + +def test_shapes_outside_the_ranges_still_run_eager(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _wrapper(monkeypatch, hift_graph_first_chunk_frames=[7, 7]) + wrapper.capture() + speech_feat = torch.randn(1, 80, 16) + cache_source = torch.zeros(1, 1, 0) + + wrapper.replay(speech_feat, cache_source) + + wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source) + + +def test_bucket_shapes_keep_their_padded_replay(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _wrapper(monkeypatch, hift_graph_first_chunk_frames=[1, 30]) + wrapper.capture() + speech_feat = torch.randn(1, 80, 50) + + wrapper.replay(speech_feat, torch.zeros(1, 1, 0)) + + assert wrapper.graph[(1, 50, 0)].replays == 1 + wrapper.decode_fn.assert_not_called() + + +def test_capture_exact_reports_memory_without_cuda(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _wrapper(monkeypatch, hift_graph_first_chunk_frames=[2, 3]) + assert wrapper_module._device_used_bytes(torch.device("cpu")) is None + + assert wrapper.capture_exact() == 2 + assert wrapper.capture_exact() == 0 + + +def test_exact_shape_extras_reach_the_connector_config() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import MiniCPMO45Code2Wav + + def connector(extra: dict) -> dict: + config = SimpleNamespace( + model_config=SimpleNamespace(model="/fake/model", stage_connector_config={"extra": extra}) + ) + return MiniCPMO45Code2Wav(vllm_config=config)._connector_config + + default = connector({}) + assert default["hift_graph_first_chunk_frames"] == [] + assert default["hift_graph_continuation_frames"] == [] + assert default["hift_graph_exact_batch_sizes"] == [1] + configured = connector( + { + "hift_graph_first_chunk_frames": [1, 30], + "hift_graph_continuation_frames": [26, 74], + "hift_graph_exact_batch_sizes": [1, 2], + } + ) + assert configured["hift_graph_first_chunk_frames"] == [1, 30] + assert configured["hift_graph_continuation_frames"] == [26, 74] + assert configured["hift_graph_exact_batch_sizes"] == [1, 2] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/codec_deadline.py b/vllm_omni/model_executor/models/minicpmo_4_5/codec_deadline.py new file mode 100644 index 00000000000..face2fa21aa --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/codec_deadline.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Deadline-batched Code2Wav steps for MiniCPM-o duplex streams (``additional_config.codec_deadline``). + +A Code2Wav decode costs nearly the same for one stream as for several (H100 +Whole-Euler graph replay, w50: 17.7 ms at B=1, 22.5 at B=4, 39.9 at B=16), but +with synchronous scheduling every chunk that arrives is decoded in the next +step, so 16 duplex sessions run ~1.25 rows per step. This policy holds a ready +continuation chunk while its stream can afford it and releases it together with +others, earliest deadline first. + +Each stream (one Stage-2 request) gets a ledger rebuilt from this stage's own +emit times -- the transport delay to the client shifts the emit times and the +client's arrival times alike, so it cancels: + +* zero-stall deadline: the end of the audio the client already holds + (``DUPLEX_CLIENT_PREBUFFER_S`` once, for the session's first segment, then + play on arrival); +* rate deadline: the arrival that keeps ``stream_rtf`` at ``rtf_cap`` if this + chunk were the stream's last. + +A chunk may wait until ``min(deadlines) - margin_s - T(eager_rows) - T(step)`` +and never longer than ``max_hold_s``. ``T`` is an EWMA of this stage's own +step time per Code2Wav graph tier (``step_tiers``, else powers of two, up to +``max_rows``); ``T(eager_rows)`` reserves the eager step that may +start just before the chunk falls due, since a synchronous step cannot be +preempted. First chunks (``chunk_seq == 0``), last chunks (``last_chunk`` / +``turn_end``; not ``tts_is_last_chunk``, which is true for every duplex unit) +and chunks without a ledger are eager: they release at once, with only the rows +due within one step. A deadline step fills its graph tier with not-yet-due rows +unless a stream is still waiting for its first chunk, which would otherwise +queue behind the wider step. A breaker stops holding for ``breaker_cooloff_s`` +once more than ``breaker_late_frac`` of the last ``breaker_window`` held chunks +arrived after their zero-stall deadline. + +Pure Python. The stage's generation scheduler drives it as its +``StepReleasePolicy``: the stage sets ``additional_config.step_release_policy`` +to ``build_codec_deadline_policy`` below, and a hold needs the event-driven idle +park (``additional_config.chunk_idle_wait_s`` > 0). +""" + +from __future__ import annotations + +import bisect +from collections import deque +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field, fields +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, StepReleasePolicy +from vllm_omni.distributed.omni_connectors.transfer_adapter.base import stage_idle_wait_s +from vllm_omni.engine.duplex.config import DUPLEX_CLIENT_PREBUFFER_S +from vllm_omni.engine.serialization import deserialize_additional_information + +logger = init_logger(__name__) + +#: Codec frames per second of the MiniCPM-o Talker codec. +CODEC_FRAME_RATE_HZ = 25.0 + + +def graph_step_tiers(max_rows: int, step_tiers: Iterable[int] | None = None) -> list[int]: + """Code2Wav graph batch sizes up to ``max_rows`` (``step_tiers``, else powers of two).""" + sizes = list(step_tiers) if step_tiers else [1 << power for power in range(max_rows.bit_length())] + return sorted({int(size) for size in sizes if 0 < int(size) < max_rows} | {max_rows}) + + +def seed_step_s(rows: int) -> float: + """Initial step-time estimate (H100 d16 serving: B=1 ~0.11 s, B=4 ~0.20 s, B=8 ~0.28 s).""" + return 0.09 + 0.024 * rows + + +@dataclass(frozen=True) +class CodecDeadlineConfig: + enabled: bool = False + margin_s: float = 0.1 + rtf_cap: float = 1.1 + max_hold_s: float = 0.35 + max_rows: int = 8 + eager_rows: int = 4 + step_alpha: float = 0.1 + gap_reanchor_s: float = 1.0 + breaker_window: int = 50 + breaker_late_frac: float = 0.1 + breaker_cooloff_s: float = 10.0 + #: Step-time tiers, i.e. the Code2Wav graph batch sizes (e.g. the connector's + #: ``cfm_graph_batch_sizes``); None means powers of two. + step_tiers: tuple[int, ...] | None = None + + def __post_init__(self) -> None: + if not isinstance(self.enabled, bool): + raise ValueError("codec_deadline.enabled must be a boolean") + if self.step_tiers is not None: + tiers = tuple(self.step_tiers) + if not tiers or any(isinstance(t, bool) or not isinstance(t, int) or t < 1 for t in tiers): + raise ValueError("codec_deadline.step_tiers must be a list of positive integers") + object.__setattr__(self, "step_tiers", tiers) + for name in ("margin_s", "max_hold_s", "gap_reanchor_s", "breaker_cooloff_s"): + if float(getattr(self, name)) < 0: + raise ValueError(f"codec_deadline.{name} must be non-negative") + if self.rtf_cap <= 0 or not 0 < self.step_alpha <= 1 or not 0 <= self.breaker_late_frac <= 1: + raise ValueError("codec_deadline: rtf_cap > 0, 0 < step_alpha <= 1, 0 <= breaker_late_frac <= 1") + if self.max_rows < 1 or self.eager_rows < 1 or self.breaker_window < 1: + raise ValueError("codec_deadline: max_rows, eager_rows and breaker_window must be >= 1") + + @classmethod + def from_additional_config(cls, additional_config: object) -> CodecDeadlineConfig | None: + raw = additional_config.get("codec_deadline") if isinstance(additional_config, Mapping) else None + if raw is None: + return None + if not isinstance(raw, Mapping): + raise ValueError("additional_config.codec_deadline must be a mapping") + known = {f.name for f in fields(cls)} + unknown = set(raw) - known + if unknown: + raise ValueError(f"additional_config.codec_deadline: unknown keys {sorted(unknown)}") + return cls(**dict(raw)) + + +def _scalar(value: Any) -> Any: + """A plain Python value from a connector meta field (tensor, array, list or scalar).""" + if isinstance(value, list | tuple): + return _scalar(value[0]) if value else None + numel = getattr(value, "numel", None) + if callable(numel): + return value.reshape(-1)[0].item() if numel() else None + item = getattr(value, "item", None) + if callable(item) and getattr(value, "size", 1) == 1: + return item() + return value + + +@dataclass(frozen=True, slots=True) +class ChunkMeta: + chunk_seq: int + cache_epoch: int + last_chunk: bool + turn_end: bool + codec_chunk_frames: int + duplex_epoch: int | None + + @classmethod + def from_info(cls, info: object) -> ChunkMeta | None: + """The producer metadata of a request's current chunk; None when it carries none.""" + meta = info.get("meta") if isinstance(info, Mapping) else None + if not isinstance(meta, Mapping) or meta.get("chunk_seq") is None: + return None + try: + duplex_epoch = _scalar(meta.get("duplex_epoch")) + return cls( + chunk_seq=int(_scalar(meta.get("chunk_seq"))), + cache_epoch=int(_scalar(meta.get("cache_epoch")) or 0), + last_chunk=bool(_scalar(meta.get("last_chunk"))), + turn_end=bool(_scalar(meta.get("turn_end"))), + codec_chunk_frames=int(_scalar(meta.get("codec_chunk_frames")) or 0), + duplex_epoch=None if duplex_epoch is None else int(duplex_epoch), + ) + except (TypeError, ValueError): + return None + + +@dataclass(slots=True) +class StreamLedger: + key: tuple[int | None, int] + anchor: float + play_end: float + audio_s: float + d1: float + late_s: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class ReadyChunk: + request_id: str + meta: ChunkMeta | None + + +@dataclass +class DeadlinePlan: + held: set[str] = field(default_factory=set) + released: list[str] = field(default_factory=list) + trigger: str = "none" + + +class StepTimeModel: + """EWMA of this stage's step wall time per graph tier: mean + 1.5 x mean absolute deviation.""" + + def __init__(self, tiers: Iterable[int], alpha: float) -> None: + self.tiers = list(tiers) + self.mean = [seed_step_s(tier) for tier in self.tiers] + self.mad = [0.0 for _ in self.tiers] + self.alpha = alpha + + def tier_index(self, rows: int) -> int: + return min(bisect.bisect_left(self.tiers, max(1, rows)), len(self.tiers) - 1) + + def tier_capacity(self, rows: int) -> int: + return self.tiers[self.tier_index(rows)] + + def estimate(self, rows: int) -> float: + index = self.tier_index(rows) + return self.mean[index] + 1.5 * self.mad[index] + + def observe(self, rows: int, elapsed_s: float) -> None: + if rows <= 0 or elapsed_s <= 0: + return + index = self.tier_index(rows) + error = elapsed_s - self.mean[index] + self.mean[index] += self.alpha * error + self.mad[index] += self.alpha * (abs(error) - self.mad[index]) + + +class LateBreaker: + """Stops holding for a cool-off once too many held chunks arrived late.""" + + def __init__(self, window: int, late_frac: float, cooloff_s: float) -> None: + self.outcomes: deque[bool] = deque(maxlen=window) + self.window = window + self.late_frac = late_frac + self.cooloff_s = cooloff_s + self.open_until = 0.0 + + def is_open(self, now: float) -> bool: + return now < self.open_until + + def record(self, late: bool, now: float) -> None: + self.outcomes.append(late) + if sum(self.outcomes) > self.late_frac * self.window: + self.open_until = now + self.cooloff_s + self.outcomes.clear() + logger.warning("Codec deadline batching: held chunks arrived late; holding off for %.1f s", self.cooloff_s) + + +def _audio_seconds(mm_output: Any) -> float | None: + """Seconds of mono audio in a Code2Wav output (``model_outputs`` at ``sr``), if unambiguous.""" + if not isinstance(mm_output, dict): + return None + audio = mm_output.get("model_outputs") + rate = mm_output.get("sr") + if not isinstance(audio, torch.Tensor) or audio.ndim != 1 or audio.numel() == 0: + return None + if isinstance(rate, torch.Tensor): + rate = int(rate.reshape(-1)[0]) if rate.numel() else 0 + if not isinstance(rate, int) or rate <= 0: + return None + return audio.numel() / rate + + +class CodecDeadlinePolicy(StepReleasePolicy): + """Per-step hold/release decisions and the per-stream playback ledgers (see module docstring).""" + + def __init__(self, config: CodecDeadlineConfig) -> None: + self.config = config + self.ledgers: dict[str, StreamLedger] = {} + #: Streams whose client already spent its jitter buffer (once per session). + self.buffered: set[str] = set() + self.step_time = StepTimeModel(graph_step_tiers(config.max_rows, config.step_tiers), config.step_alpha) + self.breaker = LateBreaker(config.breaker_window, config.breaker_late_frac, config.breaker_cooloff_s) + self.ready_since: dict[str, float] = {} + #: Requests whose current chunk was held at least once. + self.held_once: set[str] = set() + #: Chunk metadata, held flag and trigger at dispatch, consumed by ``on_output``. + self.dispatched: dict[str, tuple[ChunkMeta | None, bool, str]] = {} + self._last_trigger = "none" + self.next_release_at: float | None = None + self.last_schedule_t: float | None = None + #: Metadata of every chunk that was ready in the current step. + self._metas: dict[str, ChunkMeta | None] = {} + + # ---- planning ------------------------------------------------------ # + + def _eager(self, chunk: ReadyChunk, now: float) -> bool: + meta = chunk.meta + if self.breaker.is_open(now) or meta is None: + return True + if meta.chunk_seq == 0 or meta.last_chunk or meta.turn_end or meta.codec_chunk_frames <= 0: + return True + if meta.duplex_epoch is None: + return True + ledger = self.ledgers.get(chunk.request_id) + if ledger is None or ledger.key != (meta.duplex_epoch, meta.cache_epoch): + return True + return now - ledger.play_end > self.config.gap_reanchor_s + + def release_by(self, chunk: ReadyChunk, now: float, ready_rows: int) -> float: + """Latest dispatch time of a held chunk (``r_eff``).""" + meta = chunk.meta + ledger = self.ledgers[chunk.request_id] + assert meta is not None + chunk_s = meta.codec_chunk_frames / CODEC_FRAME_RATE_HZ + rate_deadline = ledger.anchor + self.config.rtf_cap * (ledger.audio_s + chunk_s - ledger.d1) + deadline = min(ledger.play_end, rate_deadline) + # Its own step, and an eager step that may start just before it falls due. + steps_s = self.step_time.estimate(ready_rows) + self.step_time.estimate(self.config.eager_rows) + release = deadline - self.config.margin_s - steps_s + since = self.ready_since.setdefault(chunk.request_id, now) + return min(release, since + self.config.max_hold_s) + + def plan(self, ready: list[ReadyChunk], now: float, *, onset_pending: bool = False) -> DeadlinePlan: + """This step's held and released chunks. + + ``onset_pending``: a stream is waiting for its first chunk; a deadline + step then releases only what is due instead of filling its graph tier. + """ + config = self.config + self.last_schedule_t = now + self.next_release_at = None + ready_ids = {chunk.request_id for chunk in ready} + for request_id in [rid for rid in self.ready_since if rid not in ready_ids]: + del self.ready_since[request_id] + if not ready: + return DeadlinePlan() + for chunk in ready: + self.ready_since.setdefault(chunk.request_id, now) + eager = [chunk.request_id for chunk in ready if self._eager(chunk, now)] + eager_set = set(eager) + release_by = { + chunk.request_id: self.release_by(chunk, now, len(ready)) + for chunk in ready + if chunk.request_id not in eager_set + } + paced = sorted(release_by, key=release_by.__getitem__) + one_step = self.step_time.estimate(1) + due = [rid for rid in paced if release_by[rid] <= now] + soon = [rid for rid in paced if now < release_by[rid] <= now + one_step] + if eager: + trigger = "eager" + # Protect the first chunk: only rows that are due anyway ride along. + selected = eager + due + soon[: max(0, config.eager_rows - len(eager) - len(due))] + elif due: + trigger = "deadline" + selected = due + soon + if not onset_pending: + target = min(config.max_rows, self.step_time.tier_capacity(len(selected))) + selected += [rid for rid in paced if rid not in selected][: max(0, target - len(selected))] + elif len(ready) >= config.max_rows: + trigger = "rows" + selected = paced[: config.max_rows] + else: + self.next_release_at = release_by[paced[0]] + plan = DeadlinePlan(held=set(ready_ids), trigger="hold") + self._note_plan(plan) + return plan + selected = selected[: max(config.max_rows, len(eager) + len(due))] + held = ready_ids - set(selected) + if held: + self.next_release_at = min(release_by[rid] for rid in held if rid in release_by) + plan = DeadlinePlan(held=held, released=selected, trigger=trigger) + self._note_plan(plan) + return plan + + def _note_plan(self, plan: DeadlinePlan) -> None: + self.held_once.update(plan.held) + self._last_trigger = plan.trigger + + # ---- StepReleasePolicy -------------------------------------------- # + + def held_requests(self, scheduler: OmniGenerationScheduler, now: float) -> set[str]: + """Which ready chunks this step holds back. + + Ready: a live request with unscheduled prompt tokens, no batch in flight + and (with a chunk adapter) a loaded chunk. A live request without any + prompt token yet is a stream waiting for its first chunk. + """ + adapter = getattr(scheduler, "chunk_transfer_adapter", None) + ready_marks = getattr(adapter, "requests_with_ready_chunks", None) if adapter is not None else None + ready: list[ReadyChunk] = [] + onset_pending = False + for request in (*scheduler.running, *scheduler.waiting): + request_id = request.request_id + if request_id not in scheduler.requests or request.num_in_flight_tokens > 0: + continue + prompt_len = len(request.prompt_token_ids or ()) + if prompt_len == 0: + onset_pending = True + continue + if prompt_len <= request.num_computed_tokens: + continue + if ready_marks is not None and request_id not in ready_marks: + continue + info = getattr(request, "additional_information", None) + if not isinstance(info, dict): + info = deserialize_additional_information(info) + ready.append(ReadyChunk(request_id, ChunkMeta.from_info(info))) + self._metas = {chunk.request_id: chunk.meta for chunk in ready} + return self.plan(ready, now, onset_pending=onset_pending).held + + def on_scheduled(self, request_ids: Iterable[str], now: float) -> None: + """The scheduler dispatched these requests' chunks.""" + for request_id in request_ids: + held = request_id in self.held_once + self.held_once.discard(request_id) + self.ready_since.pop(request_id, None) + self.dispatched[request_id] = (self._metas.get(request_id), held, self._last_trigger) + + def next_release_in(self, now: float) -> float | None: + if self.next_release_at is None: + return None + return max(0.0, self.next_release_at - now) + + # ---- outputs ------------------------------------------------------- # + + def on_output(self, request_id: str, mm_output: Any, now: float) -> None: + meta, held, _trigger = self.dispatched.pop(request_id, (None, False, "none")) + if meta is None: + return + audio_s = _audio_seconds(mm_output) + if audio_s: + key = (meta.duplex_epoch, meta.cache_epoch) + ledger = self.ledgers.get(request_id) + if ( + ledger is None + or meta.chunk_seq == 0 + or ledger.key != key + or now - ledger.play_end > self.config.gap_reanchor_s + ): + # The client buffers once per session: its first segment only. + first = request_id not in self.buffered and meta.duplex_epoch == 0 and meta.cache_epoch == 0 + buffer_s = DUPLEX_CLIENT_PREBUFFER_S if first else 0.0 + self.buffered.add(request_id) + self.ledgers[request_id] = StreamLedger( + key=key, anchor=now, play_end=now + buffer_s + audio_s, audio_s=audio_s, d1=audio_s + ) + else: + late = now > ledger.play_end + if late: + ledger.late_s += now - ledger.play_end + ledger.play_end = now + if held: + self.breaker.record(late, now) + ledger.play_end += audio_s + ledger.audio_s += audio_s + if meta.last_chunk or meta.turn_end: + # The next turn re-anchors (without buffering again). + self.ledgers.pop(request_id, None) + + def observe_step(self, rows: int, now: float) -> None: + if self.last_schedule_t is not None and rows > 0: + self.step_time.observe(rows, now - self.last_schedule_t) + + def forget(self, request_id: str) -> None: + self.ledgers.pop(request_id, None) + self.buffered.discard(request_id) + self.ready_since.pop(request_id, None) + self.held_once.discard(request_id) + self.dispatched.pop(request_id, None) + + +def build_codec_deadline_policy(scheduler: OmniGenerationScheduler) -> CodecDeadlinePolicy | None: + """``additional_config.step_release_policy`` factory: the policy, or None (reason logged).""" + additional_config = getattr(scheduler.vllm_config, "additional_config", None) + config = CodecDeadlineConfig.from_additional_config(additional_config) + if config is None or not config.enabled: + return None + reason = None + if bool(getattr(scheduler.scheduler_config, "async_scheduling", False)): + reason = "async_scheduling is on (the batch queue would delay each step's output by one step)" + elif bool(getattr(scheduler, "_native_data_plane", False)): + reason = "the native MRv2 data plane is not supported" + elif stage_idle_wait_s(additional_config) <= 0: + reason = "additional_config.chunk_idle_wait_s is unset (a hold needs the event-driven idle park)" + if reason is not None: + logger.warning("Codec deadline batching disabled: %s", reason) + return None + if getattr(scheduler, "_first_chunk_express", False): + logger.warning( + "Codec deadline batching is on: VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS only acts on the native data plane" + ) + policy = CodecDeadlinePolicy(config) + logger.info( + "Codec deadline batching enabled: margin %.2f s, max hold %.2f s, rtf cap %.2f, max rows %d, step tiers %s", + config.margin_s, + config.max_hold_s, + config.rtf_cap, + config.max_rows, + policy.step_time.tiers, + ) + return policy diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py index a37218fe537..6e969ccc8c3 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py @@ -1,16 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import os +import time +import weakref from collections import Counter -from collections.abc import Callable +from collections.abc import Callable, Iterable from typing import NamedTuple +import numpy as np import torch import torch.nn.functional as F from torch.cuda import CUDAGraph from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking + logger = init_logger(__name__) @@ -34,12 +40,37 @@ def _euler_timeline(n_timesteps: int, device: torch.device, dtype: torch.dtype) return timeline, (timeline[1:] - timeline[:-1]).tolist() +def codec_frame_range(value, *, name: str) -> range: + """An inclusive ``[first, last]`` codec-frame range from the connector extra (empty: off).""" + if value is None or (isinstance(value, (list, tuple)) and len(value) == 0): + return range(0) + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError(f"MiniCPM-o {name} must be [first, last] codec frames, got {value!r}") + first, last = (int(v) for v in value) + if first < 1 or last < first: + raise ValueError(f"MiniCPM-o {name} must satisfy 1 <= first <= last, got {value!r}") + return range(first, last + 1) + + +def empty_hift_outputs(speech_feat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """HiFT's ``(speech, source)`` for a zero-width mel, without touching the device. + + HiFT emits a fixed number of samples per mel frame (480 at 24 kHz), so no + frames means no samples. Running it instead fails in its first convolution, + which cannot pad an empty input up to its kernel. Zero-element tensors + allocate nothing and launch no kernel. + """ + batch_size = int(speech_feat.shape[0]) + return speech_feat.new_empty((batch_size, 0)), speech_feat.new_empty((batch_size, 1, 0)) + + class HiFTGraphWrapper: def __init__(self, token2wav, connector_config, capture_batch_sizes, max_serial_batch: int | None = None): self.decode_fn = token2wav.hift.inference self.graph_fn = token2wav.hift._inference_pre_istft self.finalize_fn = token2wav.hift._finalize_decode self.codec_chunk_frames = connector_config["codec_chunk_frames"] + self.initial_codec_chunk_frames = int(connector_config.get("initial_codec_chunk_frames", 0)) self.codec_left_context_frames = connector_config["codec_left_context_frames"] lookahead_layer = getattr(token2wav.flow.encoder, "pre_lookahead_layer", None) pre_lookahead_len = getattr(lookahead_layer, "pre_lookahead_len", None) @@ -48,7 +79,36 @@ def __init__(self, token2wav, connector_config, capture_batch_sizes, max_serial_ self.source_cache_len = int(token2wav.source_cache_len) self.mel_frames = int(token2wav.hift.conv_pre.in_channels) self.flow_upsample_rate = int(getattr(token2wav.flow, "token_mel_ratio", 2)) + self.extra_codec_chunk_frames = [int(c) for c in connector_config.get("hift_graph_codec_chunk_frames") or ()] self.capture_bucket_size, self.capture_source_cache_len = self.derive_capture_bucket_size() + # The only (mel_frames, cache_source_len) shapes a graph is ever + # captured for. Duplex first chunks carry arbitrary mel widths that are + # seen once, so capturing one would stall every stream on the device + # for a graph that pays for itself once; they run eager (``replay``). + self._legit_shapes = set(zip(self.capture_bucket_size, self.capture_source_cache_len, strict=True)) + # Exact-shape graphs (``hift_graph_first_chunk_frames`` / + # ``hift_graph_continuation_frames``, default off): the vocoder shapes + # outside the buckets -- a stream's short first chunk, a merged backlog + # continuation -- each get a graph of their own at + # ``hift_graph_exact_batch_sizes``, captured with the buckets and + # replayed without padding. The replay runs the kernels eager HiFT runs + # for that shape, so it changes no output; other batch sizes of these + # shapes still run eager (never padded, never captured lazily). + self.first_chunk_frames = codec_frame_range( + connector_config.get("hift_graph_first_chunk_frames"), name="hift_graph_first_chunk_frames" + ) + self.continuation_frames = codec_frame_range( + connector_config.get("hift_graph_continuation_frames"), name="hift_graph_continuation_frames" + ) + self.exact_batch_sizes = sorted( + {int(b) for b in connector_config.get("hift_graph_exact_batch_sizes") or (1,) if int(b) > 0} + ) + self.exact_shapes = self.derive_exact_shapes() + self._exact_keys = frozenset( + (batch_size, mel_frames, cache_len) + for batch_size in self.exact_batch_sizes + for mel_frames, cache_len in self.exact_shapes + ) self.capture_batch_sizes = capture_batch_sizes self.graph: dict[tuple[int, int, int], torch.cuda.CUDAGraph] = {} self.static_speech_inputs: dict[tuple[int, int, int], torch.Tensor] = {} @@ -59,19 +119,59 @@ def __init__(self, token2wav, connector_config, capture_batch_sizes, max_serial_ parameter = next(token2wav.hift.parameters()) self.device = parameter.device self.dtype = parameter.dtype - self.max_lazy_graphs = 8 + self.max_lazy_graphs = int(connector_config.get("hift_max_lazy_graphs", 8)) + if self.max_lazy_graphs < 0: + raise ValueError("MiniCPM-o hift_max_lazy_graphs must be >= 0") self.lazy_graph_count = 0 self.max_serial_batch = 4 if max_serial_batch is None else int(max_serial_batch) - def derive_capture_bucket_size(self): - chunk_mel_frames = ( - self.codec_chunk_frames + self.codec_left_context_frames - self.pre_lookahead_len - ) * self.flow_upsample_rate + def _chunk_mel_frames(self, codec_frames: int) -> int: + return (int(codec_frames) + self.codec_left_context_frames - self.pre_lookahead_len) * self.flow_upsample_rate - return [chunk_mel_frames, chunk_mel_frames + self.mel_cache_len], [ - 0, - self.source_cache_len, + def derive_capture_bucket_size(self): + chunk_mel_frames = self._chunk_mel_frames(self.codec_chunk_frames) + + frames = [chunk_mel_frames, chunk_mel_frames + self.mel_cache_len] + cache_lengths = [0, self.source_cache_len] + if self.initial_codec_chunk_frames and self.initial_codec_chunk_frames != self.codec_chunk_frames: + first_frames = self._chunk_mel_frames(self.initial_codec_chunk_frames) + if first_frames <= self.mel_cache_len: + raise ValueError("MiniCPM-o initial codec chunk must emit audio beyond the HiFT mel cache") + frames.append(first_frames) + cache_lengths.append(0) + # Extra codec chunk sizes that also stream as continuations. A + # full-duplex unit is one ``initial_codec_chunk_frames``-sized chunk + # after another, so its steady vocoder call (mel cache + source + # cache) is outside the buckets above and would run eager every unit. + shapes = list(zip(frames, cache_lengths, strict=True)) + for codec_frames in self.extra_codec_chunk_frames: + mel_frames = self._chunk_mel_frames(codec_frames) + if mel_frames <= self.mel_cache_len: + raise ValueError("MiniCPM-o hift_graph_codec_chunk_frames must emit audio beyond the HiFT mel cache") + for shape in ((mel_frames, 0), (mel_frames + self.mel_cache_len, self.source_cache_len)): + if shape not in shapes: + shapes.append(shape) + return [shape[0] for shape in shapes], [shape[1] for shape in shapes] + + def derive_exact_shapes(self) -> list[tuple[int, int]]: + """``(mel_frames, cache_source_len)`` of the exact-shape graphs, outside the buckets. + + A first chunk of ``f`` codec frames vocodes ``_chunk_mel_frames(f)`` + mel frames with no source cache; a continuation also carries the mel + cache in front and the source cache. + """ + buckets = set(zip(self.capture_bucket_size, self.capture_source_cache_len, strict=True)) + shapes: list[tuple[int, int]] = [] + candidates = [(self._chunk_mel_frames(f), 0) for f in getattr(self, "first_chunk_frames", ())] + candidates += [ + (self._chunk_mel_frames(f) + self.mel_cache_len, self.source_cache_len) + for f in getattr(self, "continuation_frames", ()) ] + for mel_frames, cache_len in candidates: + shape = (int(mel_frames), int(cache_len)) + if shape[0] > 0 and shape not in buckets and shape not in shapes: + shapes.append(shape) + return shapes def capture(self): for batch_size in self.capture_batch_sizes: @@ -81,6 +181,38 @@ def capture(self): strict=True, ): self._capture(batch_size, mel_frames, source_cache_len) + self.capture_exact() + + def capture_exact(self) -> int: + """Capture the exact-shape graphs (``derive_exact_shapes``); returns how many are new.""" + shapes = getattr(self, "exact_shapes", ()) + batch_sizes = getattr(self, "exact_batch_sizes", ()) + if not shapes or not batch_sizes: + return 0 + before = len(self.graph) + started = time.perf_counter() + memory_before = _memory_snapshot(self.device) + used_before = _device_used_bytes(self.device) + for batch_size in batch_sizes: + for mel_frames, source_cache_len in shapes: + key = (batch_size, mel_frames, source_cache_len) + self._capture(*key) + if key in self.graph: + # Build this frame count's ISTFT envelope now: the first build + # reads a value back to the host (``_istft_without_host_sync``). + self.finalize_fn(self.static_magnitude_outputs[key][:1], self.static_phase_outputs[key][:1]) + captured = len(self.graph) - before + used_after = _device_used_bytes(self.device) + logger.info( + "Captured %d exact-shape HiFT CUDA Graphs (batch sizes %s, %d shapes) in %.1f s%s, device memory %s", + captured, + list(batch_sizes), + len(shapes), + time.perf_counter() - started, + _format_memory_delta(memory_before, _memory_snapshot(self.device)), + "n/a" if used_before is None or used_after is None else f"+{(used_after - used_before) / 2**20:.1f} MiB", + ) + return captured def _capture( self, @@ -124,6 +256,8 @@ def _capture( logger.info("Captured HiFT CUDA Graph for shape %s", key) def replay(self, speech_feat, cache_source): + if int(speech_feat.shape[2]) == 0: + return empty_hift_outputs(speech_feat) if torch.cuda.is_current_stream_capturing(): logger.info("Falling back to eager HiFT inference during an active stream capture") return self.decode_fn(speech_feat, cache_source) @@ -131,6 +265,23 @@ def replay(self, speech_feat, cache_source): batch_size = speech_feat.shape[0] num_frames = speech_feat.shape[2] cache_source_len = cache_source.shape[2] + + exact_key = (batch_size, num_frames, cache_source_len) + if exact_key in getattr(self, "_exact_keys", ()) and exact_key in self.graph: + # An exact-shape graph: same shape as the eager call, no padding. + return self._replay_key(exact_key, speech_feat, cache_source) + + if (num_frames, cache_source_len) not in self._legit_shapes: + # Never lazily captured (see ``_legit_shapes``); one eager call over + # the whole batch, since none of its rows can graph anyway. + logger.info( + "Falling back to eager HiFT inference for shape (%d, %d, %d) outside the capture buckets", + batch_size, + num_frames, + cache_source_len, + ) + return self.decode_fn(speech_feat, cache_source) + target_b = next((b for b in sorted(self.capture_batch_sizes) if b >= batch_size), None) if target_b is None: @@ -155,6 +306,11 @@ def replay(self, speech_feat, cache_source): self._capture(*key) self.lazy_graph_count += 1 + return self._replay_key(key, speech_feat, cache_source) + + def _replay_key(self, key, speech_feat, cache_source): + """Replay the captured graph ``key`` on a batch of at most ``key[0]`` rows.""" + batch_size = speech_feat.shape[0] static_speech_inputs = self.static_speech_inputs[key].zero_() static_speech_inputs[:batch_size].copy_(speech_feat) static_cache_sources = self.static_cache_source_inputs[key].zero_() @@ -194,6 +350,21 @@ def _memory_snapshot(device: torch.device) -> tuple[int, int] | None: return None +def _device_used_bytes(device: torch.device) -> int | None: + """Device memory in use (``cudaMemGetInfo``), including what the caching allocator never sees. + + CUDA graphs also hold driver-side memory per captured graph, which the + allocator's ``memory_reserved`` does not count. + """ + if device.type != "cuda": + return None + try: + free, total = torch.cuda.mem_get_info(device) + except Exception: + return None + return int(total - free) + + def _format_memory_delta(before: tuple[int, int] | None, after: tuple[int, int] | None) -> str: if before is None or after is None: return "" @@ -208,13 +379,17 @@ def _format_memory_delta(before: tuple[int, int] | None, after: tuple[int, int] _ATT_FRAME_ALIGN = 16 +def _minus_ones(shape, *, device, dtype) -> torch.Tensor: + return torch.full(shape, -1, device=device, dtype=dtype) + + def _align_up(n: int, bucket: int) -> int: if n <= 0 or bucket <= 1: return max(0, n) return ((n + bucket - 1) // bucket) * bucket -def _capture_query_width(mel_width: int, bucket: int) -> int: +def _capture_query_width(mel_width: int, bucket: int | tuple[int, ...]) -> int: """Snap the query axis onto the Whole-Euler capture grid. ``_cfm_pad_frames`` already aligns onto ``bucket_frames=16``, which is why @@ -224,12 +399,36 @@ def _capture_query_width(mel_width: int, bucket: int) -> int: decode widths collapse onto one CUDA graph; first-chunk widths above the bucket (e.g. 304) align up separately so streaming replay stays on the small decode graph. + + A tuple of widths snaps to the narrowest one that holds the chunk (a + 25-token duplex unit and a 75-token turn chunk each keep their own width); + past the widest, the chunk aligns up on it. """ - if bucket <= 1: + widths = sorted(int(w) for w in (bucket if isinstance(bucket, tuple) else (bucket,)) if int(w) > 1) + if not widths: return int(mel_width) - if mel_width <= bucket: - return int(bucket) - return _align_up(int(mel_width), int(bucket)) + for width in widths: + if mel_width <= width: + return width + return _align_up(int(mel_width), widths[-1]) + + +def _capture_offset(offset: int, bucket: int, steady: int) -> int: + """Snap an estimator-cache length onto the Whole-Euler offset grid. + + Exact lengths made one graph per first-chunk size: under full duplex a + capture (1-2 s, stalling every stream) on almost every response. The grid + is anchored at the steady ``prompt + 100`` length so the steady solve is + never padded; shorter caches round up to it in ``bucket``-frame steps. The + padded frames are masked out of attention (``_build_capture_mask``), the + only place the DiT reads its cache, so no row's result changes. + """ + offset = int(offset) + if bucket <= 1 or offset <= 0: + return offset + if 0 < steady and offset <= steady: + return steady - ((steady - offset) // bucket) * bucket + return _align_up(offset, bucket) def _pad_query_for_capture( @@ -262,6 +461,7 @@ def _build_capture_mask( mel_width: int, mel_frames: int, device: torch.device, + offset_cap: int | None = None, ) -> torch.Tensor: """Always-on attention mask in the capture layout, so presence-of-mask is not a graph key. @@ -279,26 +479,41 @@ def _build_capture_mask( query_cap = int(query_cap) offset = int(offset) mel_width = int(mel_width) + # ``offset_cap`` columns of cache follow the query block; the ones past the + # request's own ``offset`` are ``_capture_offset`` padding and stay False. + offset_cap = offset if offset_cap is None else int(offset_cap) mask = torch.zeros( 2 * int(batch_size), query_cap, - query_cap + offset, + query_cap + offset_cap, dtype=torch.bool, device=device, ) if attn_mask is not None: mask[:, :mel_width, :mel_width] = attn_mask[:, :mel_width, :mel_width] if offset > 0: - mask[:, :mel_width, query_cap:] = attn_mask[:, :mel_width, mel_width : mel_width + offset] + mask[:, :mel_width, query_cap : query_cap + offset] = attn_mask[ + :, :mel_width, mel_width : mel_width + offset + ] else: mask[:, :mel_width, : int(mel_frames)] = True if offset > 0: - mask[:, :mel_width, query_cap:] = True + mask[:, :mel_width, query_cap : query_cap + offset] = True if query_cap > mel_width: mask[:, mel_width:] = mask[:, :1] return mask +def _pad_invalid_queries(mask: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: + """``mask`` with the queries past each CFG row's ``lengths`` taking that row's first query's keys. + + Their output is dropped, but no attention row may be empty (see + ``_build_capture_mask``). + """ + invalid = torch.arange(int(mask.shape[1]), device=mask.device).unsqueeze(0) >= lengths.unsqueeze(1) + return torch.where(invalid.unsqueeze(2), mask[:, :1], mask) + + def _att_keep_ranges(total: int, keep: tuple[int, int] | None) -> list[tuple[int, int]]: """``(start, length)`` frame ranges of a new estimator cache that survive the streaming trim. @@ -377,6 +592,74 @@ def _copy_frame_segments(dst: torch.Tensor, src: torch.Tensor, segments: list[tu position += length +class SharedPromptAttCache: + """A request-owned cache prefix plus a read-only prompt suffix. + + The logical cache remains ``(..., frames, width)`` so callers and graph + bucketing keep their existing shape contracts. Only the tail is retained + per request; graph staging can copy a range directly without materializing + the shared suffix. The current-first cache layout means the immutable + voice segment is at the logical tail after streaming trim. + """ + + __slots__ = ("prefix", "prompt") + + def __init__(self, prefix: torch.Tensor, prompt: torch.Tensor) -> None: + if prefix.ndim != 6 or prompt.ndim != 6 or tuple(prefix.shape[:4]) != tuple(prompt.shape[:4]): + raise ValueError("shared prompt attention cache tensors have incompatible shapes") + if prefix.shape[5] != prompt.shape[5] or prefix.device != prompt.device or prefix.dtype != prompt.dtype: + raise ValueError("shared prompt attention cache tensors must have the same device, dtype and width") + self.prefix = prefix + self.prompt = prompt + + @property + def shape(self) -> torch.Size: + return torch.Size((*self.prefix.shape[:4], self.prefix.shape[4] + self.prompt.shape[4], self.prefix.shape[5])) + + @property + def dtype(self) -> torch.dtype: + return self.prefix.dtype + + @property + def device(self) -> torch.device: + return self.prefix.device + + def data_ptr(self) -> int: + """Expose a stable pointer for cache-owner accounting without a full materialization.""" + return self.prefix.data_ptr() + + def copy_range_to(self, dst: torch.Tensor, start: int, length: int) -> None: + """Copy a logical frame range into ``dst`` without joining prompt and tail.""" + start, length = int(start), int(length) + if start < 0 or length < 0 or start + length > int(self.shape[4]): + raise ValueError(f"attention-cache range ({start}, {length}) is outside {tuple(self.shape)}") + prefix_len = int(self.prefix.shape[4]) + position = 0 + if start < prefix_len: + count = min(length, prefix_len - start) + dst[..., position : position + count, :].copy_(self.prefix[..., start : start + count, :]) + position += count + start += count + length -= count + if length: + prompt_start = start - prefix_len + dst[..., position : position + length, :].copy_(self.prompt[..., prompt_start : prompt_start + length, :]) + + def materialize(self) -> torch.Tensor: + """Materialize the logical cache for eager or unsupported graph paths.""" + return torch.cat((self.prefix, self.prompt), dim=4) + + +def _copy_att_cache_range(dst: torch.Tensor, src, start: int, length: int) -> None: + if isinstance(src, SharedPromptAttCache): + src.copy_range_to(dst, start, length) + elif isinstance(src, ResidentAttCache): + materialized = src.materialize() + dst.copy_(materialized[..., int(start) : int(start) + int(length), :]) + else: + dst.copy_(src[..., int(start) : int(start) + int(length), :]) + + def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]: """Rebuild zero tensors from a cache key (shape, dtype, device tuples). @@ -585,36 +868,22 @@ def _zero_padded_cnn_cache( can come from it are cleared; the valid part of the window is kept. ``cnn_cache`` is indexed by block on its first axis. """ - blocks = getattr(estimator, "blocks", None) + blocks = estimator.blocks if pad_frames > 0 and blocks and len(blocks) == cnn_cache.shape[0]: width = int(cnn_cache.shape[-1]) - if width > 0 and all( - hasattr(block, "conv") - and hasattr(block.conv, "block") - and len(block.conv.block) > 1 - and hasattr(block.conv.block[1], "causal_padding") - and int(block.conv.block[1].causal_padding[0]) == width - for block in blocks - ): + if width > 0 and all(int(block.conv.block[1].causal_padding[0]) == width for block in blocks): # _estimator_buffers packs equal-width blocks into one tensor. # Clear their shared tail with one write instead of one per block. cnn_cache[..., max(0, width - pad_frames) :] = 0.0 return - if blocks is not None: - for index, block in enumerate(blocks): - if index >= len(cnn_cache): - break - conv = getattr(block, "conv", None) - conv_block = getattr(conv, "block", None) if conv is not None else None - if conv_block is None or len(conv_block) <= 1 or not hasattr(conv_block[1], "causal_padding"): - continue - width = int(conv_block[1].causal_padding[0]) - if width <= 0: - continue - zero_from = max(0, width - pad_frames) - if zero_from < width: - cnn_cache[index][..., zero_from:] = 0.0 + for index, block in enumerate(blocks): + width = int(block.conv.block[1].causal_padding[0]) + if width <= 0: + continue + zero_from = max(0, width - pad_frames) + if zero_from < width: + cnn_cache[index][..., zero_from:] = 0.0 class WholeEulerExecutionArena: @@ -644,14 +913,14 @@ def __init__( device: torch.device, dtype: torch.dtype, att_cache_dtype: torch.dtype = torch.float32, - time_steps: list[torch.Tensor] | None = None, + timeline: torch.Tensor, ) -> None: self.estimator = estimator self.n_timesteps = int(n_timesteps) self.device = device self.dtype = dtype self.att_cache_dtype = att_cache_dtype - self.time_steps = list(time_steps) if time_steps is not None else [] + self.timeline = timeline blocks = estimator.blocks self.depth = len(blocks) @@ -665,6 +934,8 @@ def __init__( self._buffers: dict[tuple, torch.Tensor] = {} # The one attention cache storage; every graph takes a prefix view. self._att: torch.Tensor | None = None + # Per-timestep modulation table for a fused body (``get_modulation``). + self._modulation: torch.Tensor | None = None def _buffer(self, name: str, shape: tuple[int, ...], new=torch.zeros, dtype: torch.dtype | None = None): key = (name, shape) @@ -680,7 +951,7 @@ def get_time_embeddings(self, batch_size: int) -> torch.Tensor: if t_emb is None: t_emb = torch.stack( [ - self.estimator.t_embedder(self.time_steps[s].expand(2 * batch_size)).unsqueeze(1) + self.estimator.t_embedder(self.timeline[s].expand(2 * batch_size)).unsqueeze(1) for s in range(self.n_timesteps) ], dim=0, @@ -688,6 +959,22 @@ def get_time_embeddings(self, batch_size: int) -> torch.Tensor: self._buffers[key] = t_emb return t_emb + def get_modulation(self, modulation_fn: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor: + """``modulation_fn`` of each timestep's (one-row) time embedding, stacked on dim 0. + + Every row of a solve shares the timestep, so this is computed once, + outside any capture, and shared by every graph. + """ + if self._modulation is None: + self._modulation = torch.stack( + [ + modulation_fn(self.estimator.t_embedder(self.timeline[s].expand(1)).unsqueeze(1)) + for s in range(self.n_timesteps) + ], + dim=0, + ) + return self._modulation + def _cnn_cache_shape(self, batch_size: int) -> tuple[int, ...]: return (self.n_timesteps, self.depth, 2 * batch_size, self.cnn_channels, self.cnn_width) @@ -750,9 +1037,135 @@ def get_lengths_staging(self, batch_size: int) -> torch.Tensor: """Per-CFG-row valid query lengths read by a ragged graph body.""" return self._buffer("lengths", (2 * batch_size,), dtype=torch.long) + def att_bytes(self) -> int: + """Device bytes of the shared attention cache storage.""" + storage = self._att + return 0 if storage is None else storage.numel() * storage.element_size() + def clear(self) -> None: self._buffers.clear() self._att = None + self._modulation = None + + +class AttSlotPool: + """Resident estimator attention caches, one slot per request, for the Whole-Euler graphs. + + The arena path copies each request's cache (~0.5 GiB in fp32) into the + shared arena and the new one back out on every replay; here the cache + stays in its slot and the graph attends to it in place. + + The streaming trim (``_att_keep_ranges``) keeps a cache's first ``prefix`` + and last ``suffix`` frames, and chunks are prepended, so the last + ``suffix`` frames never change once present. They sit at physical frames + ``0 .. suffix - 1`` and the streaming frames anywhere behind them (a + ``ResidentAttCache`` lists where). The graph writes a chunk's keys/values + to free frames and a trim only releases frames, so nothing moves and the + cache offset is no graph key: attention runs over the whole slot under a + mask of the request's frames (no positional term reads the key order). + + Storage is ``(n_t, depth, slots, 2, heads, frames, width)``; request + ``s``'s CFG rows are rows ``2s`` and ``2s + 1`` of ``rows()``. + """ + + def __init__( + self, + *, + n_timesteps: int, + depth: int, + heads: int, + width: int, + slots: int, + frames: int, + suffix: int, + device: torch.device, + dtype: torch.dtype, + ) -> None: + self.n_timesteps, self.depth, self.heads, self.width = n_timesteps, depth, heads, width + self.slots, self.frames, self.suffix = int(slots), int(frames), int(suffix) + self.device, self.dtype = device, dtype + self.storage = torch.zeros( + (n_timesteps, depth, self.slots, 2, heads, self.frames, width), device=device, dtype=dtype + ) + # Lowest slot first, so a lightly loaded pool keeps touching the same pages. + self._free = list(range(self.slots - 1, -1, -1)) + + def rows(self) -> torch.Tensor: + """``(n_t, depth, 2 * slots, heads, frames, width)``: one row per request CFG half.""" + return self.storage.flatten(2, 3) + + @property + def free_slots(self) -> int: + return len(self._free) + + def bytes(self) -> int: + return self.storage.numel() * self.storage.element_size() + + def fits(self, frames: int) -> bool: + """Whether a logical cache of ``frames`` frames (after a chunk) fits one slot.""" + return self.suffix <= frames <= self.frames + + def adopt(self, cache: torch.Tensor) -> "ResidentAttCache | None": + """A slot holding the logical ``(n_t, depth, 2, heads, L, width)`` ``cache``; ``None`` when full. + + The last ``suffix`` frames go to the fixed frames, the rest in order + behind them. ``cache`` itself is not written, so a shared one (the + prompt state) can be adopted by every request that starts from it. + """ + frames = int(cache.shape[4]) + if not self._free or not self.fits(frames): + return None + slot = self._free.pop() + streaming = frames - self.suffix + _copy_frame_segments(self.storage[:, :, slot], cache, [(streaming, self.suffix), (0, streaming)]) + return ResidentAttCache(self, slot, list(range(self.suffix, self.suffix + streaming))) + + +class ResidentAttCache: + """A request's estimator attention cache resident in an ``AttSlotPool`` slot. + + Stands in for the ``(n_t, depth, 2, heads, L, width)`` tensor the + Whole-Euler path otherwise hands back per request: ``shape``, ``dtype`` + and ``device`` describe that logical tensor, and ``materialize()`` builds + it for the paths that need one. ``layout`` is the physical frame of each + streaming frame, in logical order. The slot returns to the pool when the + handle is collected; work queued on it before then runs first, since all + of it is on one stream. + """ + + __slots__ = ("pool", "slot", "layout", "__weakref__") + + def __init__(self, pool: AttSlotPool, slot: int, layout: list[int]) -> None: + self.pool = pool + self.slot = slot + self.layout = layout + weakref.finalize(self, pool._free.append, slot) + + @property + def shape(self) -> torch.Size: + pool = self.pool + return torch.Size((pool.n_timesteps, pool.depth, 2, pool.heads, len(self.layout) + pool.suffix, pool.width)) + + @property + def dtype(self) -> torch.dtype: + return self.pool.dtype + + @property + def device(self) -> torch.device: + return self.pool.device + + def physical_frames(self) -> list[int]: + """The physical frame of every logical frame, in logical order.""" + return [*self.layout, *range(self.pool.suffix)] + + def materialize(self) -> torch.Tensor: + """The logical cache as a tensor of its own.""" + index = index_to_device(self.physical_frames(), self.pool.device) + return self.pool.storage[:, :, self.slot].index_select(4, index) + + +def _materialize_att_rows(rows: list) -> list[torch.Tensor]: + return [row.materialize() if isinstance(row, (ResidentAttCache, SharedPromptAttCache)) else row for row in rows] class _WholeEulerStatics(NamedTuple): @@ -767,6 +1180,12 @@ class _WholeEulerStatics(NamedTuple): attn_mask: torch.Tensor time_embeddings: torch.Tensor lengths: torch.Tensor | None + # (n_timesteps, ...) per-timestep modulation for a body that takes it. + modulation: torch.Tensor | None = None + # Slot-pool graphs (``AttSlotPool``): each row's pool row, and the pool + # frame each query frame's key/value is written to (-1: not written). + slot_rows: torch.Tensor | None = None + slot_positions: torch.Tensor | None = None class WholeEulerCFMGraphWrapper: @@ -784,9 +1203,10 @@ class WholeEulerCFMGraphWrapper: ten times what one step-level graph holds, so memory is the constraint: graphs exist for the powers of two up to ``micro_batch_size``, all of them share the one attention cache storage of ``WholeEulerExecutionArena`` - (sized for ``micro_batch_size`` requests), and larger batches run as a - sequence of micro-batch replays that copy each request's cache in and out - of that storage directly. + (sized for ``micro_batch_size`` requests by default), and larger batches + run as a sequence of graph-tier replays that copy each request's cache in + and out of that storage directly. An opt-in grid-sized arena can use fewer + rows while retaining the scheduler's larger admission limit. """ def __init__( @@ -800,9 +1220,14 @@ def __init__( max_serial_batch: int | None = None, max_graph_batch: int | None = None, micro_batch_size: int = 4, - query_bucket_frames: int | None = None, + query_bucket_frames: int | tuple[int, ...] | None = None, pad_max_rows: int | None = None, ragged_body: Callable[..., torch.Tensor] | None = None, + offset_bucket_frames: int = 0, + modulation_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, + att_slots: int = 0, + arena_rows_from_graph_grid: bool = False, + row_offsets: bool = False, ) -> None: """``ragged_body(estimator, input, t_emb, mask, cnn, att, cnn_out, att_out, lengths)`` replaces ``estimator.blocks_forward_chunk`` in the captured solve when @@ -811,16 +1236,39 @@ def __init__( It lets one graph solve rows of different lengths, and lets a padded query axis keep the CNN cache the unpadded solve would produce instead of zeroing it. + + ``query_bucket_frames`` is one capture width or a tuple of them + (``_capture_query_width``); ``offset_bucket_frames`` > 1 snaps cache + lengths onto a grid (``_capture_offset``). ``modulation_fn`` maps a + one-row time embedding to the per-timestep table a fused body takes as + ``modulation=``. ``att_slots`` > 0 keeps request caches resident in an + ``AttSlotPool`` (fused body only): streaming chunks then replay graphs + keyed by batch and width alone, with no cache copies; the arena serves + the prompt solve and any chunk the pool cannot hold. ``row_offsets`` + lets such a replay take per-request caches of different lengths (a + stream's second chunk next to steady ones): each row attends its own + slot frames, so one graph solves them all. """ self.estimator = estimator self.ragged_body = ragged_body + self.modulation_fn = modulation_fn if ragged_body is not None else None self.n_timesteps = int(n_timesteps) self.inference_cfg_rate = float(inference_cfg_rate) self.att_cache_dtype = att_cache_dtype self.max_graphs = int(max_graphs) - self.query_bucket_frames = 0 if query_bucket_frames is None else int(query_bucket_frames) - self.max_serial_batch = 4 if max_serial_batch is None else int(max_serial_batch) + # Lazy captures may grow ``max_graphs`` up to 4x this (``_entry``). + self._configured_max_graphs = int(max_graphs) + if isinstance(query_bucket_frames, (tuple, list)): + widths = tuple(sorted({int(w) for w in query_bucket_frames if int(w) > 1})) + self.query_bucket_frames = widths[-1] if widths else 0 + else: + self.query_bucket_frames = 0 if query_bucket_frames is None else int(query_bucket_frames) + widths = (self.query_bucket_frames,) if self.query_bucket_frames > 1 else () + # The capture widths, narrowest first (``_capture_query_width``). + self.query_widths: tuple[int, ...] = widths + self.offset_bucket_frames = int(offset_bucket_frames) self.micro_batch_size = 4 if micro_batch_size is None else int(micro_batch_size) + # ``max_serial_batch`` only caps the default ``max_graph_batch``. if max_graph_batch is None: if max_serial_batch is not None and max_serial_batch < self.micro_batch_size: self.max_graph_batch = int(max_serial_batch) @@ -831,6 +1279,8 @@ def __init__( if pad_max_rows is None: pad_max_rows = self.micro_batch_size // 4 self.pad_max_rows = max(0, int(pad_max_rows)) + self.arena_rows_from_graph_grid = bool(arena_rows_from_graph_grid) + self._arena_rows = max(self._graph_batches()) if self.arena_rows_from_graph_grid else self.micro_batch_size # Largest steady cache length a caller announced; sizes the arena storage. self._att_capacity = 0 parameter = next(estimator.parameters(), None) @@ -850,30 +1300,43 @@ def __init__( "captures": 0, "flushes": 0, "eager": 0, + "slot_replays": 0, } self.timeline, self.dt_steps = _euler_timeline(self.n_timesteps, self.device, self.dtype) - self.time_steps = [self.timeline[i] for i in range(self.n_timesteps)] - self.arena = WholeEulerExecutionArena( - estimator=self.estimator, - n_timesteps=self.n_timesteps, - device=self.device, - dtype=self.dtype, - att_cache_dtype=self.att_cache_dtype, - time_steps=self.time_steps, - ) + arena_args = { + "estimator": self.estimator, + "n_timesteps": self.n_timesteps, + "device": self.device, + "dtype": self.dtype, + "att_cache_dtype": self.att_cache_dtype, + "timeline": self.timeline, + } + self.arena = WholeEulerExecutionArena(**arena_args) + # Resident request caches (``AttSlotPool``), sized once the prompt + # length is known; their graphs and staging buffers are kept apart + # from the arena's so an arena flush never retires them. + self.att_slots = int(att_slots) if ragged_body is not None and modulation_fn is not None else 0 + self.slot_pool: AttSlotPool | None = None + self.row_offsets = bool(row_offsets) and self.att_slots > 0 + self._slot_arena = WholeEulerExecutionArena(**arena_args) + self._slot_graphs: dict[tuple, tuple] = {} def stats_snapshot(self) -> dict[str, int]: """Bounded cumulative telemetry for the graph cache.""" - return {**self._stats, "cache_size": len(self._cache)} + return {**self._stats, "cache_size": len(self._cache) + len(self._slot_graphs)} + + @staticmethod + def _retire(graphs: dict[tuple, tuple]) -> None: + for entry in graphs.values(): + entry[4].reset() + graphs.clear() def _flush(self) -> None: """Retire every captured graph at once.""" torch.accelerator.synchronize(self.device) if self._cache: - for entry in self._cache.values(): - entry[4].reset() - self._cache.clear() + self._retire(self._cache) self._stats["flushes"] += 1 logger.info("Whole-Euler CFM graph cache flushed; stats=%s", self.stats_snapshot()) self.arena.clear() @@ -887,6 +1350,7 @@ def _disable(self, reason: str, key: tuple) -> None: ) self.enabled = False self._flush() + self._retire(self._slot_graphs) def _run_euler_loop( self, @@ -899,6 +1363,9 @@ def _run_euler_loop( cur_x = statics.x width = int(statics.mu_cfg.shape[2]) speaker_features = statics.speakers_cfg.unsqueeze(-1).expand(-1, -1, width) + # A slot-pool graph attends to the pool in place: no cache goes in. + slotted = statics.slot_rows is not None + no_cache = [None] * len(self.estimator.blocks) for step in range(self.n_timesteps): dt = self.dt_steps[step] @@ -909,13 +1376,18 @@ def _run_euler_loop( statics.time_embeddings[step], statics.attn_mask, statics.cnn_cache[step], - statics.att_cache[step], + no_cache if slotted else statics.att_cache[step], out_cnn_cache[step], out_att_cache[step], ) if statics.lengths is not None: assert self.ragged_body is not None - estimate = self.ragged_body(self.estimator, *args, statics.lengths) + kwargs = {} + if statics.modulation is not None: + kwargs["modulation"] = statics.modulation[step] + if slotted: + kwargs["slots"] = (statics.slot_rows, statics.slot_positions) + estimate = self.ragged_body(self.estimator, *args, statics.lengths, **kwargs) else: estimate = self.estimator.blocks_forward_chunk(*args) @@ -924,16 +1396,25 @@ def _run_euler_loop( return cur_x def _graph_batches(self) -> list[int]: - """Native graph batch sizes: the powers of two below ``micro_batch_size``, and it.""" + """Native graph batch sizes: the powers of two below ``micro_batch_size``, and it. + + ``VLLM_OMNI_CFM_GRAPH_GRID`` (comma-separated sizes) pins them instead; + sizes above ``micro_batch_size`` are dropped. + """ + micro = self.micro_batch_size + raw = os.environ.get("VLLM_OMNI_CFM_GRAPH_GRID") + if raw: + sizes = sorted({int(x) for x in raw.split(",") if x.strip()}) + return [b for b in sizes if 0 < b <= micro] or [micro] sizes = [] size = 1 - while size < self.micro_batch_size: + while size < micro: sizes.append(size) size *= 2 - sizes.append(self.micro_batch_size) + sizes.append(micro) return sizes - def _plan_groups(self, batch_size: int) -> list[tuple[int, int]] | None: + def _plan_groups(self, batch_size: int) -> list[tuple[int, int]]: """``(graph_batch, rows)`` replays that cover ``batch_size`` requests. Larger batches run whole micro-batches. The remainder takes the smallest @@ -941,16 +1422,15 @@ def _plan_groups(self, batch_size: int) -> list[tuple[int, int]] | None: (the padded rows are ignored), and otherwise the largest native size below it, repeating on what is left. On a saturated device a small graph's replay costs well over its share of a larger one, so a padded - row is cheaper than a further replay. Returns ``None`` when the plan - needs more replays than allowed. + row is cheaper than a further replay. """ - micro = self.micro_batch_size sizes = self._graph_batches() + largest = max(sizes) groups: list[tuple[int, int]] = [] remainder = batch_size - while remainder >= micro: - groups.append((micro, micro)) - remainder -= micro + while remainder >= largest: + groups.append((largest, largest)) + remainder -= largest while remainder > 0: up = next(size for size in sizes if size >= remainder) if up - remainder <= self.pad_max_rows: @@ -959,12 +1439,6 @@ def _plan_groups(self, batch_size: int) -> list[tuple[int, int]] | None: down = max(size for size in sizes if size <= remainder) groups.append((down, down)) remainder -= down - max_allowed = max( - self.max_serial_batch, - (self.max_graph_batch + micro - 1) // micro + (micro - 1), - ) - if len(groups) > max_allowed: - return None return groups @staticmethod @@ -1022,13 +1496,57 @@ def _fill_static_inputs( ) if offset <= 0: return - static_att_rows = statics.att_cache.unflatten(2, (2, graph_batch)) + # A graph on the ``_capture_offset`` grid holds at least ``offset`` + # cache frames; the rest are masked out and keep the finite values the + # zero-initialized storage last held, like padded rows do. + static_att_rows = statics.att_cache.unflatten(2, (2, graph_batch))[..., :offset, :] if att_rows is None: assert att_cache is not None static_att_rows[:, :, :, :rows].copy_(att_cache.unflatten(2, (2, batch_size))[:, :, :, start:stop]) else: for row in range(rows): - static_att_rows[:, :, :, row].copy_(att_rows[start + row]) + _copy_att_cache_range(static_att_rows[:, :, :, row], att_rows[start + row], 0, offset) + + def _group_fills(self, groups: list[tuple[int, int]], after=None, **inputs) -> list[Callable]: + """Per replay group, a ``fill(statics)`` that copies the group's requests in, then runs ``after``.""" + fills = [] + start = 0 + for graph_batch, rows in groups: + + def fill(statics, *, _graph_batch=graph_batch, _start=start, _stop=start + rows): + self._fill_static_inputs(statics, graph_batch=_graph_batch, start=_start, stop=_stop, **inputs) + if after is not None: + after(statics, _graph_batch, _start, _stop) + + fills.append(fill) + start += rows + return fills + + def _statics( + self, + arena: WholeEulerExecutionArena, + *, + graph_batch: int, + channels: int, + query_cap: int, + spk_dim: int, + att_cache: torch.Tensor, + key_frames: int, + **slots: torch.Tensor, + ) -> _WholeEulerStatics: + return _WholeEulerStatics( + x=arena.get_x_staging(graph_batch, channels, query_cap), + mu_cfg=arena.get_mu_staging(graph_batch, channels, query_cap), + speakers_cfg=arena.get_speakers_staging(graph_batch, spk_dim), + cond_cfg=arena.get_cond_staging(graph_batch, channels, query_cap), + cnn_cache=arena.get_cnn_cache_in(graph_batch), + att_cache=att_cache, + attn_mask=arena.get_mask_staging(graph_batch, query_cap, key_frames), + time_embeddings=arena.get_time_embeddings(graph_batch), + lengths=arena.get_lengths_staging(graph_batch) if self.ragged_body is not None else None, + modulation=arena.get_modulation(self.modulation_fn) if self.modulation_fn is not None else None, + **slots, + ) def _capture( self, @@ -1050,23 +1568,25 @@ def _capture( # adding its width on top of the steady offset would hold ~300 frames # per row that nothing uses. capacity = max(offset + query_cap, max(offset, self._att_capacity) + 2 * self.query_bucket_frames) - rows = self.micro_batch_size - if self._cache and not arena.att_cache_fits(graph_batch, offset + query_cap): + rows = self._arena_rows + if self.att_slots: + # Streaming chunks run from the slot pool; the arena only holds + # what reaches it (the prompt solve), not a steady batch. + capacity, rows = offset + query_cap, graph_batch + if self._cache and not arena.att_cache_fits(max(graph_batch, rows), offset + query_cap): # The captured graphs hold views of the current storage: retire # them, or they keep it alive beside the larger replacement. self._flush() try: - statics = _WholeEulerStatics( - x=arena.get_x_staging(graph_batch, channels, query_cap), - mu_cfg=arena.get_mu_staging(graph_batch, channels, query_cap), - speakers_cfg=arena.get_speakers_staging(graph_batch, spk_dim), - cond_cfg=arena.get_cond_staging(graph_batch, channels, query_cap), - cnn_cache=arena.get_cnn_cache_in(graph_batch), + statics = self._statics( + arena, + graph_batch=graph_batch, + channels=channels, + query_cap=query_cap, + spk_dim=spk_dim, # Right behind the current chunk in the output view below. att_cache=arena.att_cache_view(graph_batch, offset, start=query_cap, capacity=capacity, rows=rows), - attn_mask=arena.get_mask_staging(graph_batch, query_cap, offset + query_cap), - time_embeddings=arena.get_time_embeddings(graph_batch), - lengths=arena.get_lengths_staging(graph_batch) if self.ragged_body is not None else None, + key_frames=offset + query_cap, ) out_cnn_cache = arena.get_cnn_cache_out(graph_batch) # The output view aliases the input view; see WholeEulerExecutionArena. @@ -1076,7 +1596,27 @@ def _capture( logger.warning("Failed to allocate static buffers for Whole-Euler graph key: %s", key, exc_info=True) self._unsupported.add(key) return None + storage = arena._att + return self._record( + key, + statics, + out_cnn_cache, + out_att_cache, + graph_batch=graph_batch, + where=f"att storage rows/frames={int(storage.shape[2]) // 2}/{int(storage.shape[4])}", + ) + def _record( + self, + key: tuple, + statics: _WholeEulerStatics, + out_cnn_cache: torch.Tensor, + out_att_cache: torch.Tensor, + *, + graph_batch: int, + where: str, + ) -> tuple | None: + """Warm up and capture one graph over filled ``statics``; ``None`` (and disabled) on failure.""" static_x = statics.x # The fused Euler step integrates static_x in place, so every warmup # and the capture itself start from the same initial state. @@ -1102,12 +1642,11 @@ def _capture( self._stats["captures"] += 1 logger.info( - "Captured Whole-Euler CFM CUDA Graph for shape %s (cache=%d/%d, att storage rows/frames=%d/%d, stats=%s)%s", + "Captured Whole-Euler CFM CUDA Graph for shape %s (cache=%d/%d, %s, stats=%s)%s", key, - len(self._cache) + 1, + len(self._cache) + len(self._slot_graphs) + 1, self.max_graphs, - int(arena._att.shape[2]) // 2, - int(arena._att.shape[4]), + where, self.stats_snapshot(), _format_memory_delta(memory_before, _memory_snapshot(self.device)), ) @@ -1133,7 +1672,19 @@ def _entry( self._stats["hits"] += 1 return entry if len(self._cache) >= self.max_graphs: - self._flush() + # Serving lazily captures keys precapture cannot enumerate (the + # prompt solve). A flush would retire every graph and the arena, + # so grow the budget instead, up to 4x the configured one. + if len(self._cache) < 4 * self._configured_max_graphs: + self.max_graphs = len(self._cache) + 1 + logger.warning( + "Whole-Euler lazy capture grew max_graphs to %d (cache=%d)", + self.max_graphs, + len(self._cache), + ) + else: + self._unsupported.add(key) + return None entry = self._capture( key, graph_batch=graph_batch, @@ -1147,7 +1698,9 @@ def _entry( self._cache[key] = entry return entry - def _group_entries(self, groups: list[tuple[int, int]], fills: list, **entry_args) -> list[tuple] | None: + def _group_entries( + self, entry_fn: Callable[..., tuple | None], groups: list[tuple[int, int]], fills: list, **entry_args + ) -> list[tuple] | None: """Every group's graph, acquired before any of them replays. A replay may update request caches in place, so no group may fall back @@ -1159,7 +1712,7 @@ def _group_entries(self, groups: list[tuple[int, int]], fills: list, **entry_arg flushes = self._stats["flushes"] entries = [] for (graph_batch, _), fill in zip(groups, fills, strict=True): - entry = self._entry(graph_batch=graph_batch, fill=fill, **entry_args) + entry = entry_fn(graph_batch=graph_batch, fill=fill, **entry_args) if entry is None: return None entries.append(entry) @@ -1167,6 +1720,339 @@ def _group_entries(self, groups: list[tuple[int, int]], fills: list, **entry_arg return entries return None + def _ensure_slot_pool(self, keep: tuple[int, int]) -> AttSlotPool | None: + """The slot pool for streaming trim ``keep = (prefix, suffix)``, created on first use. + + A slot holds a steady cache (``prefix + suffix`` frames) plus the + widest capture width, the largest chunk a steady cache takes. + """ + if not self.att_slots or not self.enabled: + return None + prefix, suffix = int(keep[0]), int(keep[1]) + pool = self.slot_pool + if pool is None: + if prefix < 0 or suffix <= 0 or self.query_bucket_frames <= 1: + return None + arena = self.arena + frames = _align_up(prefix + suffix + self.query_bucket_frames, _ATT_FRAME_ALIGN) + try: + pool = AttSlotPool( + n_timesteps=self.n_timesteps, + depth=arena.depth, + heads=arena.heads, + width=arena.att_width, + slots=self.att_slots, + frames=frames, + suffix=suffix, + device=self.device, + dtype=self.att_cache_dtype, + ) + except torch.OutOfMemoryError: + logger.warning( + "Whole-Euler slot pool (%d slots x %d frames) does not fit; using the arena", self.att_slots, frames + ) + self.att_slots = 0 + return None + self.slot_pool = pool + logger.info( + "Whole-Euler slot pool: %d slots x %d frames (%.2f GiB)", pool.slots, pool.frames, pool.bytes() / 2**30 + ) + return pool if pool.suffix == suffix else None + + def _slot_entry(self, *, graph_batch: int, query_cap: int, channels: int, spk_dim: int, fill) -> tuple | None: + """The slot-pool graph of ``graph_batch`` rows at width ``query_cap``, captured on first use.""" + key = ("whole_euler_slots", graph_batch, query_cap) + entry = self._slot_graphs.get(key) + if entry is not None: + self._stats["hits"] += 1 + return entry + if key in self._unsupported: + return None + pool = self.slot_pool + assert pool is not None + arena = self._slot_arena + try: + statics = self._statics( + arena, + graph_batch=graph_batch, + channels=channels, + query_cap=query_cap, + spk_dim=spk_dim, + att_cache=pool.rows(), + key_frames=pool.frames, + slot_rows=arena._buffer("slot_rows", (2 * graph_batch,), dtype=torch.long), + # Nothing is written until a replay says where. + slot_positions=arena._buffer( + "slot_positions", (2 * graph_batch, query_cap), new=_minus_ones, dtype=torch.int32 + ), + ) + out_cnn_cache = arena.get_cnn_cache_out(graph_batch) + fill(statics) + except Exception: + logger.warning("Failed to allocate static buffers for Whole-Euler graph key: %s", key, exc_info=True) + self._unsupported.add(key) + return None + entry = self._record( + key, + statics, + out_cnn_cache, + pool.rows(), + graph_batch=graph_batch, + where=f"slot pool {pool.slots} x {pool.frames} frames", + ) + if entry is not None: + self._slot_graphs[key] = entry + return entry + + def _replay_slots( + self, + *, + x_cap: torch.Tensor, + mu_cap: torch.Tensor, + cond_cap: torch.Tensor, + speakers_cfg: torch.Tensor, + cnn_cache: torch.Tensor | None, + att_rows: list, + attn_mask: torch.Tensor | None, + mel_width: int, + query_cap: int, + row_lengths: list[int], + lengths: torch.Tensor | int, + att_keep: tuple[int, int], + groups: list[tuple[int, int]], + ) -> tuple[torch.Tensor, torch.Tensor, list[ResidentAttCache]] | None: + """``replay`` on resident caches (``AttSlotPool``); ``None`` leaves the batch to the arena. + + Each row's new frames take free frames of its slot, its logical + ``[current | cache]`` mask is scattered onto the slot's frames, the + graph writes the new keys/values there and attends in place, and the + streaming trim drops frames from the row's layout afterwards. Rows may + hold caches of different lengths (``row_offsets``): the mask's cache + columns then span the longest one, and a shorter row's extra columns + map past its slot and are dropped. + """ + pool = self._ensure_slot_pool(att_keep) + if pool is None: + return None + if attn_mask is None and any(length != mel_width for length in row_lengths): + # Without a mask the arena path lets every row attend all current + # frames, a row's padding included; a slot never holds padding. + return None + batch_size = len(att_rows) + offsets = [int(row.shape[4]) for row in att_rows] + # Everything that can refuse the pool is checked before a slot changes. + if min(offsets) < pool.suffix or any( + offset + length > pool.frames for offset, length in zip(offsets, row_lengths, strict=True) + ): + return None + seen: set[int] = set() + adopt = [] + for row in att_rows: + resident = isinstance(row, ResidentAttCache) and row.pool is pool and id(row) not in seen + if resident: + seen.add(id(row)) + elif row.dtype != pool.dtype: + return None + adopt.append(not resident) + if sum(adopt) > pool.free_slots: + return None + handles: list[ResidentAttCache] = [] + for row, needs in zip(att_rows, adopt, strict=True): + handle = pool.adopt(row.materialize() if isinstance(row, ResidentAttCache) else row) if needs else row + assert handle is not None + handles.append(handle) + + # Per CFG row, in one host buffer (one copy to the device): its pool + # row, the slot frame each query writes (-1: none), and the slot frame + # of each logical key (a row's current padding maps past the slot and + # is dropped). + rows_cfg, keys = 2 * batch_size, mel_width + max(offsets) + host = np.full((rows_cfg, 1 + query_cap + keys), -1, dtype=np.int64) + host[:, 1 + query_cap :] = pool.frames + new_frames: list[np.ndarray] = [] + for row, (handle, length) in enumerate(zip(handles, row_lengths, strict=True)): + free = np.ones(pool.frames, dtype=bool) + free[: pool.suffix] = False + free[handle.layout] = False + frames = np.flatnonzero(free)[:length] + new_frames.append(frames) + halves = [row, row + batch_size] + host[halves, 0] = [2 * handle.slot, 2 * handle.slot + 1] + host[halves, 1 : 1 + length] = frames + host[halves, 1 + query_cap : 1 + query_cap + length] = frames + cached = handle.physical_frames() + host[halves, 1 + query_cap + mel_width : 1 + query_cap + mel_width + len(cached)] = cached + device = x_cap.device + index = to_device_nonblocking(torch.from_numpy(host), device) + slot_rows, positions, columns = index.split((1, query_cap, keys), dim=1) + columns = columns.unsqueeze(1).expand(-1, mel_width, -1) + slot_mask = torch.zeros((rows_cfg, query_cap, pool.frames + 1), dtype=torch.bool, device=device) + if attn_mask is None: + slot_mask[:, :mel_width].scatter_(2, columns, True) + else: + slot_mask[:, :mel_width].scatter_(2, columns, attn_mask[:, :mel_width, :keys]) + cfg_lengths = lengths if isinstance(lengths, torch.Tensor) else index_to_device(row_lengths * 2, device) + slot_mask = _pad_invalid_queries(slot_mask[..., : pool.frames], cfg_lengths) + request_rows = slot_rows.squeeze(1).unflatten(0, (2, batch_size)) + request_positions = positions.unflatten(0, (2, batch_size)) + + def copy_slots(statics: _WholeEulerStatics, graph_batch: int, start: int, stop: int) -> None: + count = stop - start + static_rows = statics.slot_rows.unflatten(0, (2, graph_batch)) + static_positions = statics.slot_positions.unflatten(0, (2, graph_batch)) + static_rows[:, :count].copy_(request_rows[:, start:stop]) + static_positions[:, :count].copy_(request_positions[:, start:stop]) + if count < graph_batch: + # Padded rows only read (the first row's slot) and never write. + static_rows[:, count:].copy_(request_rows[:, start : start + 1].expand(-1, graph_batch - count)) + static_positions[:, count:].fill_(-1) + + fills = self._group_fills( + groups, + copy_slots, + x=x_cap, + mu_cfg=mu_cap, + speakers_cfg=speakers_cfg, + cond_cfg=cond_cap, + attn_mask=slot_mask, + cnn_cache=cnn_cache, + att_cache=None, + att_rows=None, + offset=0, + lengths=lengths, + ) + entries = self._group_entries( + self._slot_entry, + groups, + fills, + query_cap=query_cap, + channels=int(x_cap.shape[1]), + spk_dim=int(speakers_cfg.shape[1]), + ) + if entries is None: + return None + + chunk_mel = x_cap.new_empty((batch_size, int(x_cap.shape[1]), mel_width)) + out_cnn = torch.empty(self._slot_arena._cnn_cache_shape(batch_size), device=device, dtype=self.dtype) + out_cnn_rows = out_cnn.unflatten(2, (2, batch_size)) + start = 0 + for (graph_batch, rows), fill, entry in zip(groups, fills, entries, strict=True): + stop = start + rows + statics, static_final_x, out_cnn_cache, _, graph = entry + fill(statics) + graph.replay() + chunk_mel[start:stop].copy_(static_final_x[:rows, :, :mel_width]) + out_cnn_rows[:, :, :, start:stop].copy_(out_cnn_cache.unflatten(2, (2, graph_batch))[:, :, :, :rows]) + start = stop + + for handle, frames in zip(handles, new_frames, strict=True): + layout = [*frames.tolist(), *handle.layout] + # The streaming frames the trim keeps are its first range. + _, kept = _att_keep_ranges(len(layout) + pool.suffix, att_keep)[0] + handle.layout = layout[:kept] + self._stats["slot_replays"] += 1 + return chunk_mel, out_cnn, handles + + def _precapture_slots(self, *, channels: int, spk_dim: int) -> int: + """Capture every slot-pool graph (batch grid x capture widths); the cache offset is not a key.""" + before = self._stats["captures"] + for graph_batch in sorted(self._graph_batches(), reverse=True): + for query_cap in self.query_widths: + entry = self._slot_entry( + graph_batch=graph_batch, + query_cap=query_cap, + channels=channels, + spk_dim=spk_dim, + fill=self._precapture_fill, + ) + if entry is None: + return self._stats["captures"] - before + return self._stats["captures"] - before + + @staticmethod + def _precapture_fill(statics: _WholeEulerStatics) -> None: + # The zero-initialized staging buffers are finite; a ragged body also + # needs lengths within its query width for the causal cache gather. + if statics.lengths is not None: + statics.lengths.fill_(int(statics.x.shape[2])) + + def precapture( + self, + *, + offsets: Iterable[int], + steady: int, + channels: int, + spk_dim: int, + max_graph_bytes: int = 1 << 30, + keep: tuple[int, int] | None = None, + ) -> int: + """Capture the graphs streams of one prompt need, before serving them. + + A capture while serving stalls every stream on the device (warmup + solves plus the capture, 1-2 s), so every graph batch, capture width + and cache offset (``offsets``, snapped onto the ``_capture_offset`` + grid anchored at ``steady``) is captured here instead. The cross + product is bounded by ``max_graphs`` so no capture flushes another, + and it stops once one graph holds more than ``max_graph_bytes`` of + device memory beyond the shared arena (a graph normally holds a few + MiB; the rest then capture on first use). Returns the number of + graphs captured. + + With a slot pool (``att_slots``) and the streaming trim ``keep``, + the streaming graphs are the slot-pool ones, one per batch and width. + """ + if not self.enabled: + return 0 + if self.att_slots and keep is not None and self._ensure_slot_pool(keep) is not None: + return self._precapture_slots(channels=int(channels), spk_dim=int(spk_dim)) + self._att_capacity = max(self._att_capacity, int(steady)) + grid = sorted({_capture_offset(int(o), self.offset_bucket_frames, int(steady)) for o in offsets}, reverse=True) + widths = self.query_widths + if not widths: + return 0 + keys = [(b, w, o) for b in sorted(self._graph_batches(), reverse=True) for w in widths for o in grid] + room = self.max_graphs - len(self._cache) + if len(keys) > room: + # Grow the budget rather than truncate the sweep (max_graph_bytes + # still bounds it), with a slot per width for the offset-0 prompt + # solve the sweep cannot cover, so its lazy capture flushes nothing. + self.max_graphs = len(self._cache) + len(keys) + len(widths) + logger.warning( + "Whole-Euler precapture needs %d graphs; raising max_graphs to %d", + len(keys), + self.max_graphs, + ) + x = torch.empty((1, int(channels), 1), device=self.device, dtype=self.dtype) + before = self._stats["captures"] + for graph_batch, query_cap, offset in keys: + memory_before, arena_before = _memory_snapshot(self.device), self.arena.att_bytes() + entry = self._entry( + graph_batch=graph_batch, + query_cap=query_cap, + offset=offset, + x=x, + spk_dim=int(spk_dim), + fill=self._precapture_fill, + ) + if entry is None: + break + memory_after = _memory_snapshot(self.device) + if memory_before is not None and memory_after is not None: + # What the graph itself holds, apart from the shared arena. + held = memory_after[1] - memory_before[1] - (self.arena.att_bytes() - arena_before) + if held > max_graph_bytes: + # A capture failure disables every graph (see _capture), + # so stop well before one could run the device dry. + logger.warning( + "Whole-Euler precapture stopped at %s: the graph reserved %.2f GiB (limit %.2f GiB); " + "the remaining graphs capture on first use", + (graph_batch, query_cap, offset), + held / 2**30, + max_graph_bytes / 2**30, + ) + break + return self._stats["captures"] - before + def replay( self, *, @@ -1215,12 +2101,12 @@ def replay( if batch_size > self.max_graph_batch: return None groups = self._plan_groups(batch_size) - if groups is None: - return None att_rows = None if att_cache is None or isinstance(att_cache, torch.Tensor) else list(att_cache) + # Per-request caches of different lengths only replay from the slot pool (``row_offsets``). + mixed_offsets = att_rows is not None and len({int(row.shape[4]) for row in att_rows}) != 1 if att_rows is not None: - if len(att_rows) != batch_size or len({int(row.shape[4]) for row in att_rows}) != 1: + if len(att_rows) != batch_size or (mixed_offsets and not self.row_offsets): return None offset = int(att_rows[0].shape[4]) else: @@ -1237,8 +2123,35 @@ def replay( if att_keep is not None: self._att_capacity = max(self._att_capacity, sum(att_keep)) - query_cap = _capture_query_width(mel_width, self.query_bucket_frames) + query_cap = _capture_query_width(mel_width, self.query_widths) x_cap, mu_cap, cond_cap = _pad_query_for_capture(x, mu_cfg, cond_cfg, query_cap) + row_lengths = [mel_width] * batch_size if valid_lengths is None else [int(n) for n in valid_lengths] + lengths: torch.Tensor | int | None = None + if self.ragged_body is not None: + lengths = mel_width + if valid_lengths is not None: + lengths = torch.tensor((*row_lengths, *row_lengths), dtype=torch.long, device=x.device) + if self.att_slots and att_rows is not None and pad_frames == 0 and att_keep is not None: + slotted = self._replay_slots( + x_cap=x_cap, + mu_cap=mu_cap, + cond_cap=cond_cap, + speakers_cfg=speakers_cfg, + cnn_cache=cnn_cache, + att_rows=att_rows, + attn_mask=attn_mask, + mel_width=mel_width, + query_cap=query_cap, + row_lengths=row_lengths, + lengths=lengths, + att_keep=att_keep, + groups=groups, + ) + if slotted is not None: + return slotted + if mixed_offsets: + return None + offset_cap = _capture_offset(offset, self.offset_bucket_frames, sum(att_keep) if att_keep is not None else 0) mask_cap = _build_capture_mask( attn_mask=attn_mask, batch_size=batch_size, @@ -1247,17 +2160,12 @@ def replay( mel_width=mel_width, mel_frames=mel_frames, device=x.device, + offset_cap=offset_cap, ) - row_lengths = [mel_width] * batch_size if valid_lengths is None else [int(n) for n in valid_lengths] - lengths: torch.Tensor | int | None = None - if self.ragged_body is not None: - lengths = mel_width - if valid_lengths is not None: - lengths = torch.tensor((*row_lengths, *row_lengths), dtype=torch.long, device=x.device) - # Queries past a row's length take its first query's mask, as - # capture padding does, so no attention row is left empty. - invalid = torch.arange(query_cap, device=x.device).unsqueeze(0) >= lengths.unsqueeze(1) - mask_cap = torch.where(invalid.unsqueeze(2), mask_cap[:, :1, :], mask_cap) + if isinstance(lengths, torch.Tensor): + # Queries past a row's length take its first query's mask, as + # capture padding does. + mask_cap = _pad_invalid_queries(mask_cap, lengths) segments_by_length = { length: _whole_euler_att_segments(mel_width=length, query_cap=query_cap, offset=offset, keep=att_keep) for length in set(row_lengths) @@ -1267,11 +2175,7 @@ def replay( arena = self.arena chunk_mel = x.new_empty((batch_size, int(x.shape[1]), mel_frames)) - out_cnn = torch.empty( - (self.n_timesteps, arena.depth, 2 * batch_size, arena.cnn_channels, arena.cnn_width), - device=x.device, - dtype=self.dtype, - ) + out_cnn = torch.empty(arena._cnn_cache_shape(batch_size), device=x.device, dtype=self.dtype) out_cnn_rows = out_cnn.unflatten(2, (2, batch_size)) out_att: torch.Tensor | None = None out_att_rows = None @@ -1287,36 +2191,25 @@ def replay( # padding (``pad_frames``) is then cleared. Otherwise every padded position is. cnn_pad = pad_frames if self.ragged_body is not None else query_cap - mel_frames - fills = [] - start = 0 - for graph_batch, rows in groups: - stop = start + rows - - def fill(statics, *, _graph_batch=graph_batch, _start=start, _stop=stop): - self._fill_static_inputs( - statics, - graph_batch=_graph_batch, - start=_start, - stop=_stop, - x=x_cap, - mu_cfg=mu_cap, - speakers_cfg=speakers_cfg, - cond_cfg=cond_cap, - attn_mask=mask_cap, - cnn_cache=cnn_cache, - att_cache=att_cache if att_rows is None else None, - att_rows=att_rows, - offset=offset, - lengths=lengths, - ) - - fills.append(fill) - start = stop + fills = self._group_fills( + groups, + x=x_cap, + mu_cfg=mu_cap, + speakers_cfg=speakers_cfg, + cond_cfg=cond_cap, + attn_mask=mask_cap, + cnn_cache=cnn_cache, + att_cache=att_cache if att_rows is None else None, + att_rows=att_rows, + offset=offset, + lengths=lengths, + ) entries = self._group_entries( + self._entry, groups, fills, query_cap=query_cap, - offset=offset, + offset=offset_cap, x=x, spk_dim=int(speakers_cfg.shape[1]), ) @@ -1358,7 +2251,11 @@ def fill(statics, *, _graph_batch=graph_batch, _start=start, _stop=stop): # The row's old cache was copied into the arena by ``fill``. old = att_rows[start + row] if att_rows is not None else None request = None - if old is not None and old.dtype == self.att_cache_dtype and owners[old.data_ptr()] == 1: + if ( + isinstance(old, torch.Tensor) + and old.dtype == self.att_cache_dtype + and owners[old.data_ptr()] == 1 + ): request = _resized_frame_view(old, shape[4]) if request is None: # Room for the steady length, so the cache grows into it in place. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py b/vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py new file mode 100644 index 00000000000..46757cb1520 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exact-shape CUDA graphs of the Code2Wav flow encoder's streaming chunk. + +Every Code2Wav step runs the CosyVoice2 upsample-conformer encoder +(``UpsampleConformerEncoderV2.forward_chunk``: pre-lookahead conv, 6 conformer +layers, 2x upsample, 4 more) between the token embedding and +``encoder_proj``: ~320 small eager kernels whose launches, not their GPU time, +set the step's latency in serving. A CUDA graph replays them with one launch. + +The graphs change no output. One is captured per exact +``(rows, tokens, conformer attention-cache frames)`` and replays only calls of +exactly that shape, so it runs the kernels the eager call runs. Nothing is +padded: rounding 3 rows up to a 4-row graph (7cd254051) changes the M of every +GEMM, and with it cuBLAS's kernel and the FP32 rounding of every row; the TF32 +CFM and HiFT amplified that to 0.81 dB at B=3. Exact shapes stay few because +the streaming trim caps the conformer cache at ``prompt + 100`` frames: one +prompt's streams only pass through the cache lengths +:func:`reachable_conformer_cache_frames` lists (300, 350 and 400 frames for the +6 s default voice and 28-token duplex chunks). Calls of any other shape, the +last chunk (lookahead padding) and the prompt's own encode run eager. + +Graphs replay one at a time and every replay refills its inputs, so all of +them take prefix views of one storage per static input and per result, sized +for the largest key (like ``streaming_audio_encoder_graph``), and share one +private memory pool for their temporaries. A replay copies its results into +the shared result views: they stay valid until the next encoder replay. +""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +EncodeFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor]] + + +def reachable_conformer_cache_frames( + *, + start: int, + prompt_len: int, + suffix: int, + growths: Iterable[int], +) -> list[int]: + """Conformer attention-cache lengths a stream passes through, from its prompt state. + + A chunk adds ``growth`` frames (the encoder's output frames) to the cache; + past ``prompt_len + suffix`` frames the streaming trim keeps the first + ``prompt_len`` and the last ``suffix`` (``_trim_streaming_cache``). These + are the cache lengths a continuation chunk's encoder call can see, for + chunks of any mix of the ``growths``. + """ + steady = int(prompt_len) + int(suffix) + steps = sorted({int(g) for g in growths if int(g) > 0}) + seen = {int(start)} + frontier = [int(start)] + while frontier: + frames = frontier.pop() + for growth in steps: + following = frames + growth + if following > steady: + following = steady + if following not in seen: + seen.add(following) + frontier.append(following) + return sorted(seen) + + +def _precision_state() -> tuple[Any, ...]: + """The matmul/conv precision flags a captured graph's kernels were chosen under.""" + return ( + bool(torch.backends.cuda.matmul.allow_tf32), + bool(torch.backends.cudnn.allow_tf32), + torch.get_float32_matmul_precision(), + ) + + +@dataclass +class _EncoderGraph: + graph: Any + tokens: torch.Tensor + cnn: torch.Tensor + att: torch.Tensor + hidden: torch.Tensor + new_cnn: torch.Tensor + new_att: torch.Tensor + + +class FlowEncoderGraphs: + """Exact-shape CUDA graphs of ``encode_fn`` (embedding, ``forward_chunk``, projection) for continuation chunks. + + ``encode_fn(tokens, cnn_cache, att_cache)`` is the eager continuation chunk + (both caches present, ``last_chunk=False``). Caches use the encoder's + layouts: ``cnn_cache`` ``(rows, channels, width)``, ``att_cache`` + ``(depth, rows, heads, frames, 2 * head_dim)``; a chunk of ``tokens`` + tokens emits ``upsample * (tokens - lookahead)`` frames. + """ + + def __init__( + self, + encode_fn: EncodeFn, + *, + rows: Iterable[int], + token_widths: Iterable[int], + lookahead: int, + upsample: int, + held_tensors: Callable[[], Sequence[torch.Tensor]] | None = None, + ) -> None: + self.encode_fn = encode_fn + self.rows = tuple(sorted({int(r) for r in rows if int(r) > 0})) + self.lookahead = int(lookahead) + self.upsample = int(upsample) + self.token_widths = tuple(sorted({int(w) for w in token_widths if int(w) > self.lookahead})) + # Tensors the graphs read in place without owning them (the RelPos + # tables): held so they stay alive, and checked so a replaced table + # sends calls back to eager instead of replaying stale reads. + self._held_tensors = held_tensors or (lambda: ()) + self._held: tuple[torch.Tensor, ...] = () + self._precision: tuple[Any, ...] | None = None + self.graphs: dict[tuple[int, int, int], _EncoderGraph] = {} + self.stats: Counter = Counter() + self.enabled = True + self._pool: Any = None + self._storage: dict[str, torch.Tensor] = {} + + def output_frames(self, tokens: int) -> int: + return self.upsample * (int(tokens) - self.lookahead) + + def keys_for(self, *, start: int, prompt_len: int, suffix: int) -> list[tuple[int, int, int]]: + """``(rows, tokens, cache frames)`` of every continuation chunk a prompt's streams can produce.""" + frames = reachable_conformer_cache_frames( + start=start, + prompt_len=prompt_len, + suffix=suffix, + growths=[self.output_frames(width) for width in self.token_widths], + ) + return [(rows, width, length) for rows in self.rows for width in self.token_widths for length in frames] + + # -- capture --------------------------------------------------------- + + def _views( + self, + key: tuple[int, int, int], + *, + cnn_shape: tuple[int, int], + att_layout: tuple[int, int, int], + hidden_dim: int, + ) -> dict[str, tuple[int, ...]]: + rows, width, frames = key + depth, heads, att_width = att_layout + out_frames = self.output_frames(width) + return { + "tokens": (rows, width), + "cnn": (rows, *cnn_shape), + "att": (depth, rows, heads, frames, att_width), + "hidden": (rows, out_frames, hidden_dim), + "new_cnn": (rows, *cnn_shape), + "new_att": (depth, rows, heads, frames + out_frames, att_width), + } + + def _view(self, name: str, shape: tuple[int, ...]) -> torch.Tensor: + numel = 1 + for size in shape: + numel *= int(size) + return self._storage[name][:numel].view(shape) + + def _body(self, entry: _EncoderGraph) -> None: + hidden, new_cnn, new_att = self.encode_fn(entry.tokens, entry.cnn, entry.att) + for name, value, target in ( + ("hidden", hidden, entry.hidden), + ("cnn", new_cnn, entry.new_cnn), + ("att", new_att, entry.new_att), + ): + if tuple(value.shape) != tuple(target.shape) or value.dtype != target.dtype: + raise RuntimeError( + f"flow encoder {name} output {tuple(value.shape)}/{value.dtype} does not match " + f"the captured layout {tuple(target.shape)}/{target.dtype}" + ) + target.copy_(value) + + def _record(self, entry: _EncoderGraph) -> Any: + """Warm up on a side stream, then capture ``_body`` into the shared private pool.""" + device = entry.att.device + current = torch.cuda.current_stream(device) + side = torch.cuda.Stream(device=device) + side.wait_stream(current) + with torch.cuda.stream(side): + for _ in range(2): + self._body(entry) + current.wait_stream(side) + torch.accelerator.synchronize(device) + if self._pool is None: + self._pool = torch.cuda.graph_pool_handle() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self._pool): + self._body(entry) + return graph + + def capture( + self, + keys: Iterable[tuple[int, int, int]], + *, + cnn_shape: tuple[int, int], + att_layout: tuple[int, int, int], + hidden_dim: int, + dtype: torch.dtype, + device: torch.device, + ) -> int: + """Capture ``keys``; returns the number of graphs held afterwards. + + ``cnn_shape`` is ``(channels, width)``, ``att_layout`` + ``(depth, heads, 2 * head_dim)``. Larger keys go first, so the private + pool grows once to the largest graph's temporaries and the others + reuse that memory. + """ + if self.graphs: + raise RuntimeError("flow encoder graphs are already captured") + keys = sorted( + {(int(r), int(w), int(f)) for r, w, f in keys}, + key=lambda key: (key[0] * (key[2] + self.output_frames(key[1])), key[0] * key[1]), + reverse=True, + ) + if not keys: + return 0 + shapes = [self._views(key, cnn_shape=cnn_shape, att_layout=att_layout, hidden_dim=hidden_dim) for key in keys] + for name in shapes[0]: + numel = 0 + for views in shapes: + count = 1 + for size in views[name]: + count *= int(size) + numel = max(numel, count) + # Ordinary tensors, not inference tensors: replays write them in + # place from inside and outside inference mode alike. + with torch.inference_mode(False): + self._storage[name] = torch.zeros(numel, dtype=torch.long if name == "tokens" else dtype, device=device) + self._precision = _precision_state() + try: + for key, views in zip(keys, shapes, strict=True): + with torch.inference_mode(False): + entry = _EncoderGraph( + graph=None, **{name: self._view(name, shape) for name, shape in views.items()} + ) + entry.graph = self._record(entry) + self.graphs[key] = entry + except Exception: + # All or nothing: a partial set would replay some shapes and not others. + self.graphs.clear() + self._storage.clear() + self._pool = None + self.enabled = False + raise + self._held = tuple(self._held_tensors()) + return len(self.graphs) + + def storage_bytes(self) -> int: + """Bytes of the shared static inputs and results (not the private pool).""" + return sum(t.numel() * t.element_size() for t in self._storage.values()) + + # -- replay ---------------------------------------------------------- + + def _layout_matches( + self, + entry: _EncoderGraph, + tokens: torch.Tensor, + cnn_rows: Sequence[torch.Tensor], + att_rows: Sequence[torch.Tensor], + ) -> bool: + if tokens.dtype != entry.tokens.dtype or tokens.device != entry.tokens.device: + return False + att_shape = (entry.att.shape[0], 1, *entry.att.shape[2:]) + cnn_shape = (1, *entry.cnn.shape[1:]) + for att, cnn in zip(att_rows, cnn_rows, strict=True): + if tuple(att.shape) != att_shape or att.dtype != entry.att.dtype or att.device != entry.att.device: + return False + if tuple(cnn.shape) != cnn_shape or cnn.dtype != entry.cnn.dtype or cnn.device != entry.cnn.device: + return False + return True + + def _held_changed(self) -> bool: + current = tuple(self._held_tensors()) + return len(current) != len(self._held) or any(a is not b for a, b in zip(current, self._held, strict=True)) + + def run( + self, + tokens: torch.Tensor, + cnn_rows: Sequence[torch.Tensor], + att_rows: Sequence[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """The continuation chunk of ``rows`` streams from its graph, or ``None`` (run it eager). + + ``cnn_rows`` / ``att_rows`` are each stream's caches, one row each, in + batch order. The results are the shared result views: copy what must + outlive the next encoder replay. + """ + if not self.enabled or not self.graphs: + return None + rows = len(att_rows) + if tokens.dim() != 2 or int(tokens.shape[0]) != rows or len(cnn_rows) != rows or rows == 0: + self.stats["eager"] += 1 + return None + first = att_rows[0] + if first.dim() != 5: + self.stats["eager"] += 1 + return None + entry = self.graphs.get((rows, int(tokens.shape[1]), int(first.shape[3]))) + if entry is None or not self._layout_matches(entry, tokens, cnn_rows, att_rows): + self.stats["eager"] += 1 + return None + if _precision_state() != self._precision or torch.cuda.is_current_stream_capturing(): + self.stats["eager_precision"] += 1 + return None + if self._held_changed(): + self.enabled = False + logger.warning("Flow encoder RelPos tables were replaced after capture; encoder graphs disabled") + return None + with torch.no_grad(): # a replay is not differentiable; serving runs in inference mode anyway + torch.cat(list(cnn_rows), dim=0, out=entry.cnn) + torch.cat(list(att_rows), dim=1, out=entry.att) + entry.tokens.copy_(tokens) + entry.graph.replay() + self.stats["replays"] += 1 + return entry.hidden, entry.new_cnn, entry.new_att From fbc0446f447f246a055d2b46847d36bc768fb118 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:30 +0800 Subject: [PATCH 83/89] [Bugfix][MiniCPM-o] HiFT zero-width mel safeguard and Stage 2 memory trim HiFT zero-width mel safeguard and Stage 2 memory trim Return empty HiFT outputs for zero-width mels to prevent conv padding failure, and drop unused upstream flow attention buffers in Step-Audio2 token2wav. Signed-off-by: BeatSeat --- .../minicpmo_4_5/test_hift_exact_widths.py | 99 +++++++++ .../minicpmo_4_5/test_memory_misc_switches.py | 188 ++++++++++++++++++ ...est_step_audio2_token2wav_chunk_buffers.py | 72 +++++++ .../step_audio2/step_audio2_token2wav.py | 60 +++++- 4 files changed, 417 insertions(+), 2 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_hift_exact_widths.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_memory_misc_switches.py create mode 100644 tests/model_executor/models/step_audio2/test_step_audio2_token2wav_chunk_buffers.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_hift_exact_widths.py b/tests/model_executor/models/minicpmo_4_5/test_hift_exact_widths.py new file mode 100644 index 00000000000..4412da32249 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_hift_exact_widths.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""HiFT vocoder shapes outside the capture buckets: zero-width mels. + +Mock-driven: no HiFT runs, so these hold on CPU. +""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav +from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + HiFTGraphWrapper, + empty_hift_outputs, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _no_device(*_args, **_kwargs): + raise AssertionError("a zero-width mel must not reach the device") + + +def _bare_backend(hift: object, wrapper: object | None) -> BatchedToken2Wav: + backend = object.__new__(BatchedToken2Wav) + nn.Module.__init__(backend) + backend.hift = hift + backend.hift_graph_wrapper = wrapper + return backend + + +def _mock_wrapper(monkeypatch: pytest.MonkeyPatch) -> HiFTGraphWrapper: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = object.__new__(HiFTGraphWrapper) + wrapper.capture_batch_sizes = [1] + wrapper._legit_shapes = {(50, 0)} + wrapper.graph = {} + wrapper.lazy_graph_count = 0 + wrapper.max_lazy_graphs = 8 + wrapper.max_serial_batch = 4 + wrapper.decode_fn = Mock(side_effect=_no_device) + wrapper._capture = Mock(side_effect=_no_device) + return wrapper + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_empty_hift_outputs_match_hift_output_layout(dtype: torch.dtype) -> None: + speech, source = empty_hift_outputs(torch.zeros(3, 80, 0, dtype=dtype)) + + assert speech.shape == (3, 0) + assert source.shape == (3, 1, 0) + assert speech.dtype == dtype and source.dtype == dtype + + +@pytest.mark.parametrize("cache_len", [0, 3840]) +def test_zero_width_mel_never_runs_hift_or_captures(monkeypatch: pytest.MonkeyPatch, cache_len: int) -> None: + wrapper = _mock_wrapper(monkeypatch) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", _no_device) + + speech, source = wrapper.replay(torch.zeros(2, 80, 0), torch.zeros(2, 1, cache_len)) + + assert speech.shape == (2, 0) + assert source.shape == (2, 1, 0) + wrapper.decode_fn.assert_not_called() + wrapper._capture.assert_not_called() + assert wrapper.lazy_graph_count == 0 + + +@pytest.mark.parametrize("with_wrapper", [False, True]) +def test_backend_zero_width_mel_skips_the_vocoder(with_wrapper: bool) -> None: + hift = Mock() + hift.inference = Mock(side_effect=_no_device) + wrapper = Mock() + wrapper.replay = Mock(side_effect=_no_device) + backend = _bare_backend(hift, wrapper if with_wrapper else None) + + speech, source = backend._hift_inference(torch.zeros(1, 80, 0), torch.zeros(1, 1, 0)) + + assert speech.shape == (1, 0) + assert source.shape == (1, 1, 0) + hift.inference.assert_not_called() + wrapper.replay.assert_not_called() + + +def test_backend_nonempty_mel_still_reaches_the_vocoder() -> None: + expected = (torch.zeros(1, 480 * 4), torch.zeros(1, 1, 480 * 4)) + hift = Mock() + hift.inference = Mock(return_value=expected) + backend = _bare_backend(hift, None) + mel = torch.zeros(1, 80, 4) + source = torch.zeros(1, 1, 0) + + assert backend._hift_inference(mel, source) is expected + hift.inference.assert_called_once_with(mel, source) diff --git a/tests/model_executor/models/minicpmo_4_5/test_memory_misc_switches.py b/tests/model_executor/models/minicpmo_4_5/test_memory_misc_switches.py new file mode 100644 index 00000000000..e863e5d7d4a --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_memory_misc_switches.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Opt-in memory switches: S3Tokenizer placement and the session-sized S0 audio graph grid. + +Both are off by default and must leave the default path unchanged. +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import soundfile as sf +import torch + +import vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav as code2wav_module +import vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm as omni_llm_module +import vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav as token2wav_module +import vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder_graph as graph_module +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import MiniCPMO45Code2Wav +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, +) +from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder_graph import batch_sizes_for_sessions +from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import StepAudio2Token2WavCore + +pytestmark = [pytest.mark.core_model] + + +# --- S0 audio graph grid --------------------------------------------------------- + + +@pytest.mark.cpu +@pytest.mark.parametrize( + ("sessions", "grid"), + [ + (1, (1,)), + (2, (1, 2)), + (4, (1, 2, 4)), + (12, (1, 2, 4, 8, 12)), + (16, (1, 2, 4, 8, 16)), + (20, (1, 2, 4, 8, 16, 20)), + (32, (1, 2, 4, 8, 16, 32)), + (0, (1,)), + ], +) +def test_batch_sizes_for_sessions(sessions, grid): + assert batch_sizes_for_sessions(sessions) == grid + + +class _FakeGraphEncoder: + built: list[dict] = [] + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + self.batch_sizes = tuple(kwargs["batch_sizes"]) + self.cache_buckets = tuple(kwargs["cache_buckets"]) + self.pinned_h2d = kwargs["pinned_h2d"] + _FakeGraphEncoder.built.append(kwargs) + + def capture(self) -> None: + pass + + +def _thinker(config: SimpleNamespace, max_sessions: int): + thinker = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration) + torch.nn.Module.__init__(thinker) + weight = SimpleNamespace(device=SimpleNamespace(type="cuda")) + object.__setattr__(thinker, "config", config) + object.__setattr__(thinker, "apm", SimpleNamespace(conv1=SimpleNamespace(weight=weight))) + object.__setattr__(thinker, "audio_projection_layer", None) + object.__setattr__(thinker, "audio_avg_pooler", None) + object.__setattr__(thinker, "_duplex_max_sessions", max_sessions) + object.__setattr__(thinker, "supports_streaming_audio_batch", lambda: True) + return thinker + + +def _build(monkeypatch, config: SimpleNamespace, max_sessions: int) -> dict: + _FakeGraphEncoder.built.clear() + monkeypatch.setattr(graph_module, "StreamingAudioGraphEncoder", _FakeGraphEncoder) + monkeypatch.setattr(omni_llm_module.torch.cuda, "memory_allocated", lambda *_: 0) + thinker = _thinker(config, max_sessions) + assert thinker.build_streaming_audio_graph_encoder(unit_frames=104) is True + assert len(_FakeGraphEncoder.built) == 1 + return _FakeGraphEncoder.built[0] + + +@pytest.mark.cpu +def test_grid_default_and_explicit_list_unchanged(monkeypatch): + kwargs = _build(monkeypatch, SimpleNamespace(audio_pool_step=5), max_sessions=16) + assert tuple(kwargs["batch_sizes"]) == graph_module.DEFAULT_GRAPH_BATCH_SIZES + + explicit = SimpleNamespace(audio_pool_step=5, duplex_audio_encoder_cuda_graph_batch_sizes=[1, 2, 4, 8, 16, 20]) + kwargs = _build(monkeypatch, explicit, max_sessions=16) + assert kwargs["batch_sizes"] == [1, 2, 4, 8, 16, 20] + + +@pytest.mark.cpu +@pytest.mark.parametrize(("sessions", "grid"), [(16, [1, 2, 4, 8, 16]), (20, [1, 2, 4, 8, 16, 20])]) +def test_grid_follows_max_sessions_when_enabled(monkeypatch, sessions, grid): + config = SimpleNamespace( + audio_pool_step=5, + duplex_audio_encoder_cuda_graph_batch_sizes=[1, 2, 4, 8, 16, 20], + duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions=True, + ) + kwargs = _build(monkeypatch, config, max_sessions=sessions) + assert kwargs["batch_sizes"] == grid + + +@pytest.mark.cpu +def test_grid_flag_needs_literal_true(monkeypatch): + config = SimpleNamespace(audio_pool_step=5, duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions="false") + kwargs = _build(monkeypatch, config, max_sessions=4) + assert tuple(kwargs["batch_sizes"]) == graph_module.DEFAULT_GRAPH_BATCH_SIZES + + +# --- S3Tokenizer placement ---------------------------------------------------- + + +def _code2wav_config(tmp_path: Path, extra: dict) -> SimpleNamespace: + (tmp_path / "assets" / "token2wav").mkdir(parents=True) + sf.write(tmp_path / "assets" / "HT_ref_audio.wav", np.zeros(1600, dtype=np.float32), 16000) + return SimpleNamespace(model_config=SimpleNamespace(model=str(tmp_path), stage_connector_config={"extra": extra})) + + +@pytest.mark.cpu +@pytest.mark.parametrize(("extra", "expected"), [({}, None), ({"token2wav_s3tokenizer_device": "cpu"}, "cpu")]) +def test_code2wav_passes_tokenizer_device_only_when_set(monkeypatch, tmp_path, extra, expected): + seen: list[dict] = [] + + class _FakeToken2wav: + def __init__(self, path, **kwargs): + seen.append(kwargs) + + monkeypatch.setattr(token2wav_module, "MiniCPMO45Token2wav", _FakeToken2wav) + monkeypatch.setattr(code2wav_module, "BatchedToken2Wav", lambda token2wav, **_: SimpleNamespace()) + model = MiniCPMO45Code2Wav(vllm_config=_code2wav_config(tmp_path, extra)) + model._build_backend() + + assert len(seen) == 1 + if expected is None: + assert "audio_tokenizer_device" not in seen[0] + else: + assert seen[0]["audio_tokenizer_device"] == expected + + +class _RecordingTokenizer: + def __init__(self) -> None: + self.devices: list[torch.device] = [] + + def quantize(self, mels: torch.Tensor, mels_lens: torch.Tensor): + self.devices.append(mels.device) + tokens = torch.arange(4, device=mels.device).unsqueeze(0) + return tokens, torch.tensor([4], device=mels.device) + + +def _core(tmp_path: Path, device: str, tokenizer_device: str | None): + core = StepAudio2Token2WavCore(str(tmp_path), device=device, audio_tokenizer_device=tokenizer_device) + tokenizer = _RecordingTokenizer() + core._audio_tokenizer = tokenizer + core._spk_model = lambda feat: torch.zeros(1, 192) + core._flow = SimpleNamespace(up_rate=20) + core._hift = None + core._models_loaded = True + wav = tmp_path / "ref.wav" + sf.write(wav, (0.1 * np.sin(np.linspace(0, 200, 16000))).astype(np.float32), 16000) + return core, tokenizer, str(wav) + + +@pytest.mark.cpu +def test_core_tokenizer_device_defaults_to_vocoder_device(tmp_path): + core, tokenizer, wav = _core(tmp_path, "cpu", None) + assert core.audio_tokenizer_device == core.device + tokens, lens, *_ = core._prepare_prompt(wav) + assert tokenizer.devices == [torch.device("cpu")] + assert tokens.device == torch.device("cpu") + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_core_cpu_tokenizer_returns_tokens_on_vocoder_device(tmp_path): + core, tokenizer, wav = _core(tmp_path, "cuda", "cpu") + tokens, lens, spk, mels, mels_lens = core._prepare_prompt(wav) + assert tokenizer.devices == [torch.device("cpu")] + assert tokens.device.type == "cuda" and lens.device.type == "cuda" + assert mels.device.type == "cuda" diff --git a/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_chunk_buffers.py b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_chunk_buffers.py new file mode 100644 index 00000000000..2896e45c8d6 --- /dev/null +++ b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_chunk_buffers.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dropping the upstream chunk-streaming attention buffers of the flow.""" + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import ( + StepAudio2Token2WavCore, + _StreamState, + drop_upstream_chunk_att_buffers, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Estimator(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("att_cache_buffer", torch.zeros(2, 3, 4), persistent=False) + self.register_buffer("cnn_cache_buffer", torch.zeros(5), persistent=False) + + +class _Decoder(nn.Module): + def __init__(self): + super().__init__() + self.estimator = _Estimator() + self.register_buffer("rand_noise", torch.randn(1, 2, 8), persistent=False) + self.register_buffer("att_cache_buffer", torch.zeros(4, 4), persistent=False) + + +class _Flow(nn.Module): + def __init__(self): + super().__init__() + self.decoder = _Decoder() + + +def test_drop_removes_only_the_attention_buffers(): + flow = _Flow() + noise = flow.decoder.rand_noise.clone() + released = drop_upstream_chunk_att_buffers(flow) + assert released == (2 * 3 * 4 + 4 * 4) * 4 + assert not hasattr(flow.decoder, "att_cache_buffer") + assert not hasattr(flow.decoder.estimator, "att_cache_buffer") + # The noise the batched backend draws from and the CNN buffers stay. + assert torch.equal(flow.decoder.rand_noise, noise) + assert flow.decoder.estimator.cnn_cache_buffer.shape == (5,) + assert drop_upstream_chunk_att_buffers(flow) == 0 + + +def test_drop_tolerates_a_flow_without_them(): + assert drop_upstream_chunk_att_buffers(nn.Module()) == 0 + + +@pytest.mark.parametrize("dropped", [False, True]) +def test_chunk_streaming_refuses_dropped_buffers(dropped): + core = StepAudio2Token2WavCore("/nonexistent", device="cpu", drop_upstream_chunk_att_buffers=dropped) + assert core._chunk_att_buffers_dropped is False + core._chunk_att_buffers_dropped = dropped + state = _StreamState() + state.stream_cache = object() + if dropped: + with pytest.raises(RuntimeError, match="drop_upstream_chunk_att_buffers"): + core.setup_stream_for("prompt.wav", state) + with pytest.raises(RuntimeError, match="drop_upstream_chunk_att_buffers"): + core.stream_chunk_for([1], "prompt.wav", False, state) + else: + # Without the switch the guard passes and the call reaches the prompt. + core.cache["prompt.wav"] = None + with pytest.raises(TypeError): + core.setup_stream_for("prompt.wav", state) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index 8e3c56b327d..e9fa6ff7af3 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -95,6 +95,33 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.classifier(self.condnet(x).transpose(1, 2)).squeeze(-1).abs() +# Upstream cosyvoice2 registers these non-persistent fp32 buffers for its own +# chunked streaming (``CausalConditionalCFM.solve_euler_chunk`` and +# ``DiT.forward_chunk``): (16, 16, 2, 8, 1000, 128) on the decoder (2,000 MiB) +# and (16, 2, 8, 1000, 128) on the estimator (125 MiB). Only +# ``flow.setup_cache`` / ``flow.inference_chunk`` reach them. +_UPSTREAM_CHUNK_ATT_BUFFER = "att_cache_buffer" + + +def drop_upstream_chunk_att_buffers(flow: nn.Module) -> int: + """Remove the upstream chunk-streaming attention buffers from ``flow``. + + Returns the bytes released. The attributes are gone afterwards, so any + upstream chunk path that still reaches them fails with an AttributeError + instead of computing on a stand-in. + """ + released = 0 + decoder = getattr(flow, "decoder", None) + for module in (decoder, getattr(decoder, "estimator", None)): + if not isinstance(module, nn.Module): + continue + buffer = module._buffers.get(_UPSTREAM_CHUNK_ATT_BUFFER) + if isinstance(buffer, torch.Tensor): + released += buffer.numel() * buffer.element_size() + del module._buffers[_UPSTREAM_CHUNK_ATT_BUFFER] + return released + + def _build_hift() -> HiFTGenerator: return HiFTGenerator( sampling_rate=24000, @@ -115,12 +142,22 @@ def __init__( float16: bool = False, device: str = "cuda", n_timesteps: int = DEFAULT_TOKEN2WAV_CONFIG.n_timesteps, + drop_upstream_chunk_att_buffers: bool = False, + audio_tokenizer_device: str | None = None, ): super().__init__() self.model_path = model_path + # A backend that never runs ``flow.setup_cache`` / ``flow.inference_chunk`` + # (MiniCPM-o's ``BatchedToken2Wav``) can drop their 2.1 GiB of buffers. + self._drop_chunk_att_buffers = bool(drop_upstream_chunk_att_buffers) + self._chunk_att_buffers_dropped = False self.float16 = float16 self.device = torch.device(device) self.n_timesteps = n_timesteps + # S3Tokenizer only runs in ``_prepare_prompt`` (once per new reference + # voice). ``None`` keeps it on ``device``; "cpu" keeps its fp32 weights + # (~472 MiB) off the accelerator at the cost of a CPU quantize. + self.audio_tokenizer_device = torch.device(audio_tokenizer_device) if audio_tokenizer_device else self.device self._models_loaded = False self._audio_tokenizer = None @@ -163,7 +200,9 @@ def _ensure_models_loaded(self): logger.info(f"Loading Token2Wav models from: {self.model_path}") self._audio_tokenizer = ( - s3tokenizer.load_model(f"{self.model_path}/speech_tokenizer_v2_25hz.onnx").to(self.device).eval() + s3tokenizer.load_model(f"{self.model_path}/speech_tokenizer_v2_25hz.onnx") + .to(self.audio_tokenizer_device) + .eval() ) option = onnxruntime.SessionOptions() @@ -176,6 +215,11 @@ def _ensure_models_loaded(self): with open(f"{self.model_path}/flow.yaml") as f: configs = load_hyperpyyaml(f) self._flow = configs["flow"] + if self._drop_chunk_att_buffers: + # Before ``.to(self.device)``: the buffers never reach the device. + released = drop_upstream_chunk_att_buffers(self._flow) + self._chunk_att_buffers_dropped = True + logger.info("Dropped the upstream flow chunk attention buffers (%.1f MiB)", released / 2**20) if self.float16: self._flow.half() self._flow.load_state_dict( @@ -251,8 +295,11 @@ def _prepare_prompt(self, prompt_wav: str): mels = s3tokenizer.log_mel_spectrogram(audio) mels, mels_lens = s3tokenizer.padding([mels]) prompt_speech_tokens, prompt_speech_tokens_lens = self.audio_tokenizer.quantize( - mels.to(self.device), mels_lens.to(self.device) + mels.to(self.audio_tokenizer_device), mels_lens.to(self.audio_tokenizer_device) ) + if self.audio_tokenizer_device != self.device: + prompt_speech_tokens = prompt_speech_tokens.to(self.device) + prompt_speech_tokens_lens = prompt_speech_tokens_lens.to(self.device) spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000) spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True) @@ -333,10 +380,18 @@ def forward( # ------------------------------------------------------------------ # Per-request streaming (does NOT mutate self – uses external state) + def _require_chunk_att_buffers(self) -> None: + if self._chunk_att_buffers_dropped: + raise RuntimeError( + "Token2Wav chunk streaming needs the upstream flow attention buffers, " + "which drop_upstream_chunk_att_buffers removed" + ) + # ------------------------------------------------------------------ def setup_stream_for(self, prompt_wav: str, state: _StreamState) -> None: """Initialise flow + HiFT caches into *state* (no self mutation).""" + self._require_chunk_att_buffers() if prompt_wav not in self.cache: self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) prompt_speech_tokens, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav] @@ -371,6 +426,7 @@ def stream_chunk_for( state: _StreamState, ) -> torch.Tensor: """Process one chunk using *state* (no self mutation except speech_window).""" + self._require_chunk_att_buffers() if state.stream_cache is None: raise ValueError("stream_cache not initialised – call setup_stream_for() first") From 78e087a68a02ece71b78bebadc1930a8e7a5f725 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:30 +0800 Subject: [PATCH 84/89] [Config][MiniCPM-o] Consolidated single recommended deployment profile Consolidated single recommended deployment profile Consolidate high-concurrency full duplex tuning into minicpmo_4_5_h100.yaml and update deploy config test expectations. Signed-off-by: BeatSeat --- .../duplex/test_duplex_deploy_config.py | 145 +++++++++++++++++- vllm_omni/deploy/minicpmo_4_5.yaml | 17 +- vllm_omni/deploy/minicpmo_4_5_dxsched.yaml | 73 +++++++++ vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml | 52 +++++++ .../minicpmo_4_5_dxsched_on_s1graph.yaml | 18 +++ vllm_omni/deploy/minicpmo_4_5_h100.yaml | 63 ++++++++ vllm_omni/deploy/minicpmo_4_5_kstep.yaml | 28 ++++ vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml | 6 + vllm_omni/deploy/minicpmo_4_5_pipe.yaml | 22 +++ 9 files changed, 416 insertions(+), 8 deletions(-) create mode 100644 vllm_omni/deploy/minicpmo_4_5_dxsched.yaml create mode 100644 vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml create mode 100644 vllm_omni/deploy/minicpmo_4_5_dxsched_on_s1graph.yaml create mode 100644 vllm_omni/deploy/minicpmo_4_5_h100.yaml create mode 100644 vllm_omni/deploy/minicpmo_4_5_kstep.yaml create mode 100644 vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml create mode 100644 vllm_omni/deploy/minicpmo_4_5_pipe.yaml diff --git a/tests/engine/duplex/test_duplex_deploy_config.py b/tests/engine/duplex/test_duplex_deploy_config.py index 8544c7d909a..1bfc04f2def 100644 --- a/tests/engine/duplex/test_duplex_deploy_config.py +++ b/tests/engine/duplex/test_duplex_deploy_config.py @@ -4,16 +4,23 @@ from __future__ import annotations -from dataclasses import FrozenInstanceError +from dataclasses import FrozenInstanceError, fields +from pathlib import Path import pytest +import vllm_omni from tests.helpers.stage_config import get_deploy_duplex_max_sessions from vllm_omni.config.stage_config import ( + DuplexPacingConfig, DuplexSessionRuntimeConfig, load_deploy_config, + resolve_deploy_yaml, ) +_DEPLOY_DIR = Path(vllm_omni.__file__).parent / "deploy" +_OMP_CAPS = {"OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", "OMP_WAIT_POLICY": "PASSIVE"} + pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -34,6 +41,7 @@ def test_duplex_session_runtime_defaults_are_typed_and_immutable(tmp_path) -> No assert deploy.duplex_session.max_sessions == 1 assert deploy.duplex_session.completed_append_cache_size == 256 assert deploy.duplex_session.server_vad_model_path is None + assert deploy.duplex_session.abort_on_model_listen is False with pytest.raises(FrozenInstanceError): deploy.duplex_session.idle_ttl_s = 1.0 # type: ignore[misc] @@ -100,6 +108,8 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str, ("minicpmo_4_5_duplex_h200.yaml", 16), ("minicpmo_4_5_8x4090.yaml", 1), ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 16), + ("minicpmo_4_5_dxsched.yaml", 16), + ("minicpmo_4_5_dxsched_on.yaml", 16), ], ) def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, expected: int) -> None: @@ -107,3 +117,136 @@ def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, e # or an overlay inheriting one from its base must surface here rather than # as another nightly duplex admission timeout. assert get_deploy_duplex_max_sessions(deploy_yaml) == expected + + +def _stage(raw: dict, stage_id: int) -> dict: + return next(stage for stage in raw["stages"] if stage["stage_id"] == stage_id) + + +def test_dxsched_profile_keeps_the_inherited_switches_and_adds_the_omp_caps() -> None: + # duplex_session and stage env / hf_overrides are replaced, not merged, by + # resolve_deploy_yaml: the profile has to restate what it inherits. + raw = resolve_deploy_yaml(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml") + assert raw["duplex_session"]["abort_on_model_listen"] is True + stage0, stage1, stage2 = (_stage(raw, stage_id) for stage_id in (0, 1, 2)) + assert stage0["hf_overrides"]["duplex_audio_encoder_pinned_h2d"] is True + assert stage0["hf_overrides"]["duplex_fbank_stats"] is False + assert stage0["hf_overrides"]["duplex_incremental_fbank"] is False + assert stage2["async_scheduling"] is False + assert stage2["env"]["VLLM_OMNI_STAGE_IDLE_WAIT_S"] == "0.05" + for stage in (stage0, stage1, stage2): + assert {key: stage["env"][key] for key in _OMP_CAPS} == _OMP_CAPS + # connectors are deep-merged: the base connector definition survives. + connector = raw["connectors"]["connector_of_shared_memory"] + assert connector["name"] == "SharedMemoryConnector" + assert connector["extra"]["hift_graph_codec_chunk_frames"] == [25, 75] + assert "connector_get_sleep_s" in connector["extra"] + + deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml") + assert deploy.duplex_session.abort_on_model_listen is True + assert deploy.duplex_session.pacing == DuplexPacingConfig() + assert deploy.duplex_session.barge_cut_on_model_yield is False + # Code2Wav bucket switches are listed off. + assert connector["extra"]["code2wav_bucket_stats"] is False + assert connector["extra"]["cfm_cross_turn_buckets"] is False + assert connector["extra"]["cfm_row_offset_merge"] is False + + +def test_dxsched_profile_values_match_the_code_defaults() -> None: + raw = resolve_deploy_yaml(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml")["duplex_session"] + runtime_defaults = DuplexSessionRuntimeConfig() + barge_names = {name for name in raw if name.startswith("barge_")} + assert barge_names == {f.name for f in fields(DuplexSessionRuntimeConfig) if f.name.startswith("barge_")} + for name in barge_names: + assert raw[name] == getattr(runtime_defaults, name), name + pacing_defaults = DuplexPacingConfig() + assert set(raw["pacing"]) == {f.name for f in fields(DuplexPacingConfig)} + for name, value in raw["pacing"].items(): + assert value == getattr(pacing_defaults, name), name + + +def test_duplex_session_pacing_parses_from_a_mapping(tmp_path) -> None: + deploy_path = tmp_path / "duplex.yaml" + deploy_path.write_text( + "duplex_session:\n pacing:\n enabled: true\n onset_lead_max_s: 0.9\n fire_grid_ms: 250\nstages: []\n", + encoding="utf-8", + ) + pacing = load_deploy_config(deploy_path).duplex_session.pacing + assert isinstance(pacing, DuplexPacingConfig) + assert pacing.enabled is True and pacing.onset_lead_max_s == 0.9 and pacing.fire_grid_ms == 250 + with pytest.raises(FrozenInstanceError): + pacing.enabled = False # type: ignore[misc] + + +@pytest.mark.parametrize( + ("body", "match"), + [ + ("pacing:\n onset_lead_max_s: 1.0", "onset_lead_max_s must be < 1.0"), + ("pacing:\n onset_lead_max_s: -0.1", "onset_lead_max_s must be a non-negative number"), + ("pacing:\n critical_stall_s: -1", "critical_stall_s must be a non-negative number"), + ("pacing:\n fire_grid_ms: 300", "fire_grid_ms must be 0 or a positive divisor of 1000"), + ("pacing:\n fire_grid_ms: -250", "fire_grid_ms must be 0 or a positive divisor of 1000"), + ("pacing:\n enabled: 1", "pacing.enabled must be a boolean"), + ("pacing:\n client_prebuffer_s: 0.5", "client_prebuffer_s"), + ("pacing: 3", "pacing must be a mapping"), + ("barge_cut_on_model_yield: 1", "barge_cut_on_model_yield must be a boolean"), + ("barge_arm_min_silence_s: -0.5", "barge_arm_min_silence_s must be a non-negative number"), + ("barge_arm_min_speech_s: -0.1", "barge_arm_min_speech_s must be a non-negative number"), + ], +) +def test_duplex_session_rejects_invalid_dxsched_values(tmp_path, body: str, match: str) -> None: + deploy_path = tmp_path / "duplex.yaml" + deploy_path.write_text(f"duplex_session:\n {body}\nstages: []\n", encoding="utf-8") + with pytest.raises((ValueError, TypeError), match=match): + load_deploy_config(deploy_path) + + +def test_dxsched_profile_passes_codec_deadline_to_stage2_with_the_code_defaults() -> None: + from vllm_omni.core.sched.codec_deadline import CodecDeadlineConfig + + deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml") + stage2 = next(stage for stage in deploy.stages if stage.stage_id == 2) + raw = stage2.engine_extras["additional_config"]["codec_deadline"] + assert CodecDeadlineConfig.from_additional_config({"codec_deadline": raw}) == CodecDeadlineConfig() + assert set(raw) == {f.name for f in fields(CodecDeadlineConfig)} + assert stage2.env["VLLM_OMNI_STAGE_IDLE_WAIT_S"] == "0.05" + + +def test_dxsched_on_profile_turns_on_the_measured_switches_only() -> None: + from vllm_omni.core.sched.codec_deadline import CodecDeadlineConfig + + raw = resolve_deploy_yaml(_DEPLOY_DIR / "minicpmo_4_5_dxsched_on.yaml") + base = resolve_deploy_yaml(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml") + stage0, stage1, stage2 = (_stage(raw, stage_id) for stage_id in (0, 1, 2)) + assert stage0["hf_overrides"] == { + **_stage(base, 0)["hf_overrides"], + "duplex_incremental_fbank": True, + "duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions": True, + "duplex_audio_encoder_resident_slots": 20, + "vision_fused_layers": True, + "vision_cuda_graph": True, + } + # Stage env and the S2 scheduling mode are inherited untouched. + for stage_id, stage in ((0, stage0), (1, stage1), (2, stage2)): + assert stage["env"] == _stage(base, stage_id)["env"] + assert stage2["async_scheduling"] is False + extra = raw["connectors"]["connector_of_shared_memory"]["extra"] + base_extra = base["connectors"]["connector_of_shared_memory"]["extra"] + assert extra == { + **base_extra, + "cfm_row_offset_merge": True, + "cfm_drop_upstream_att_buffer": True, + "cfm_precapture_empty_cache": True, + "token2wav_s3tokenizer_device": "cpu", + } + assert extra["cfm_cross_turn_buckets"] is False + + deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5_dxsched_on.yaml") + session = deploy.duplex_session + assert session.abort_on_model_listen is True + assert session.barge_cut_on_model_yield is True + assert session.pacing == DuplexPacingConfig( + enabled=True, onset_lead_max_s=0.9, onset_skip_response_wait=True, fire_grid_ms=250, idle_grid=True + ) + codec = next(stage for stage in deploy.stages if stage.stage_id == 2).engine_extras["additional_config"] + assert CodecDeadlineConfig.from_additional_config(codec) == CodecDeadlineConfig(enabled=True) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 7f4fed0ce24..8291bfbf1fa 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -28,17 +28,20 @@ connectors: codec_left_context_frames: 3 # HiFT capture buckets are derived from codec_chunk_frames and # codec_left_context_frames. (#5869) - # A future initial_codec_chunk_frames shape will use lazy capture unless - # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. + # The smaller first payload (initial_codec_chunk_frames) rides the lazy + # capture path unless HiFTGraphWrapper.derive_capture_bucket_size + # explicitly includes it. enable_hift_graph: true enable_cfm_graph: true # Whole-Euler CUDA graphs exist for batch 1, 2, 4, ... up to - # micro_batch_size, which defaults to Stage 2 max_num_seqs: all graphs - # share one attention-cache arena of that many requests (~0.6 GiB each - # in fp32, 10 GiB at 16). Larger batches replay micro-batches; - # B > max_graph_batch runs eager. - max_graph_batch: 16 + # micro_batch_size, which follows Stage 2 max_num_seqs (the old + # micro_batch_size / max_graph_batch overrides are no longer read). All + # graphs share one attention-cache arena of that many requests (~0.6 GiB + # each in fp32, 10 GiB at 16); larger batches replay micro-batches. cfm_max_graphs: 32 + # Opt-in, lossless memory experiments. Keep disabled for the baseline: + # cfm_prompt_att_sharing: true + # cfm_arena_rows_from_graph_grid: true stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml b/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml new file mode 100644 index 00000000000..b5ac38d6f08 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml @@ -0,0 +1,73 @@ +# dxsched measurement profile: pipe (R2) + OpenMP caps + every dxsched switch present and off. +# +# Merge semantics this file depends on (resolve_deploy_yaml): +# - duplex_session is replaced, not merged: restate every inherited field here. +# - Stage env / hf_overrides are replaced per key: restate the pipe values next +# to the new entries. engine_extras and connectors are deep-merged. +# Lanes copy this file and flip only the switch lines they test. +base_config: minicpmo_4_5_pipe.yaml +duplex_session: + max_sessions: 16 # restated from minicpmo_4_5.yaml + abort_on_model_listen: true # restated from kstep_best + barge_cut_on_model_yield: false + barge_arm_min_silence_s: 0.6 + barge_arm_min_speech_s: 0.4 # backchannel floor + barge_arm_min_playback_s: 1.0 + barge_cut_min_unplayed_s: 0.15 + pacing: + enabled: false + onset_lead_max_s: 0.0 # 0.9 when on + onset_lead_target_s: 1.1 + onset_skip_response_wait: false + critical_stall_s: 0.05 + fire_grid_ms: 0 # 250 when on + idle_grid: false + stats: false + +connectors: + connector_of_shared_memory: + extra: + hift_graph_codec_chunk_frames: [25, 75] # HiFT graphs for duplex unit continuations + code2wav_bucket_stats: false + cfm_cross_turn_buckets: false + cfm_row_offset_merge: false # second chunk + steady rows in one slot-pool replay + +stages: + - stage_id: 0 + hf_overrides: + duplex_audio_encoder_pinned_h2d: true # restated from pipe + duplex_fbank_stats: false + duplex_incremental_fbank: false + env: + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + OMP_WAIT_POLICY: PASSIVE + + - stage_id: 1 + env: + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + OMP_WAIT_POLICY: PASSIVE + + - stage_id: 2 + env: + VLLM_OMNI_STAGE_IDLE_WAIT_S: "0.05" # restated from pipe; codec_deadline needs it + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + OMP_WAIT_POLICY: PASSIVE + # engine_extras is deep-merged; the scheduler reads additional_config.codec_deadline. + engine_extras: + additional_config: + codec_deadline: + enabled: false + margin_s: 0.1 + rtf_cap: 1.1 + max_hold_s: 0.35 + max_rows: 8 + eager_rows: 4 + step_alpha: 0.1 + gap_reanchor_s: 1.0 + breaker_window: 50 + breaker_late_frac: 0.1 + breaker_cooloff_s: 10.0 + stats: false diff --git a/vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml b/vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml new file mode 100644 index 00000000000..39f57228894 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml @@ -0,0 +1,52 @@ +# dxsched recommended profile: minicpmo_4_5_dxsched.yaml with the measured switches on. +# +# Same merge semantics as the base profile: duplex_session is replaced, so every +# field is restated; stage hf_overrides are replaced per key; connectors and +# engine_extras are deep-merged. Diagnostics (pacing/codec_deadline stats, +# fbank and bucket stats) stay off; lanes turn them on. +base_config: minicpmo_4_5_dxsched.yaml +duplex_session: + max_sessions: 16 + abort_on_model_listen: true + barge_cut_on_model_yield: true + barge_arm_min_silence_s: 0.6 + barge_arm_min_speech_s: 0.4 + barge_arm_min_playback_s: 1.0 + barge_cut_min_unplayed_s: 0.15 + pacing: + enabled: true + onset_lead_max_s: 0.9 + onset_lead_target_s: 1.1 + onset_skip_response_wait: true + critical_stall_s: 0.05 + fire_grid_ms: 250 + idle_grid: true + stats: false + +connectors: + connector_of_shared_memory: + extra: + cfm_row_offset_merge: true # second chunk + steady rows in one slot-pool replay + cfm_drop_upstream_att_buffer: true # upstream chunk buffers the batched backend never reads + cfm_precapture_empty_cache: true # return startup allocator cache after precapture + token2wav_s3tokenizer_device: cpu # only used when a new reference voice is registered + +stages: + - stage_id: 0 + hf_overrides: + duplex_audio_encoder_pinned_h2d: true + duplex_fbank_stats: false + duplex_incremental_fbank: true + # Audio-encoder graph grid follows duplex max_sessions (powers of two + # below it plus the cap); the batch-32 graph is never captured. + duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions: true + # Sessions' streaming caches live in the graph's slot pool (>= max_sessions). + duplex_audio_encoder_resident_slots: 20 + vision_fused_layers: true + vision_cuda_graph: true + + - stage_id: 2 + engine_extras: + additional_config: + codec_deadline: + enabled: true diff --git a/vllm_omni/deploy/minicpmo_4_5_dxsched_on_s1graph.yaml b/vllm_omni/deploy/minicpmo_4_5_dxsched_on_s1graph.yaml new file mode 100644 index 00000000000..0613f71d77d --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_dxsched_on_s1graph.yaml @@ -0,0 +1,18 @@ +# dxsched_on + Stage-1 K-step frame graphs (opt/mcpo-dx3 measurement overlay). +# +# The CUDA K-step Talker decode (worker/gpu_talker_multiframe.py) replays the +# talker's FULL graph once per frame; with this switch the per-frame tail +# between those replays (codec head, EOS/penalty/min_tokens shaping, vLLM's +# sampler, stop/emit bookkeeping, next-frame embedding) also replays as a +# captured CUDA graph, one set per step shape, captured on its first step +# (models/minicpmo_4_5/talker_kstep_graph.py). Lossless: seeded rows sample +# from Philox slot generators that take over the request generators' state, +# so the codec ids are bit-identical to the eager tail. +# +# Stage hf_overrides are replaced per key, so this adds the one key to stage 1. +base_config: minicpmo_4_5_dxsched_on.yaml + +stages: + - stage_id: 1 + hf_overrides: + talker_kstep_graph_sampling: true diff --git a/vllm_omni/deploy/minicpmo_4_5_h100.yaml b/vllm_omni/deploy/minicpmo_4_5_h100.yaml new file mode 100644 index 00000000000..2ea11f669c8 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_h100.yaml @@ -0,0 +1,63 @@ +# H100 single-GPU duplex profile. +# +# This overlay contains the H7 graph path plus scheduling and host-runtime +# settings. The fused DiT body is the measured near-fp32 H7 path (mel 72.8 dB, +# same quality tier as SDPA). TF32 flow is retained for comparison with H7 +# (mel shift 0.67-1.41 dB); replace it after opt-num validates its alternative. +# +# MPS is started by the existing private daemon integration. It requires the +# default ``parallel_stage_init: false`` so physical-GPU initialization locks +# serialize the three local EngineCore processes. +base_config: minicpmo_4_5_kstep.yaml +cuda_mps: true + +duplex_session: + max_sessions: 20 + +connectors: + connector_of_shared_memory: + extra: + # H7's precapture buckets; graph shape selection is arithmetic-neutral. + cfm_max_graphs: 128 + cfm_graph_capture_frames: [50, 56, 150, 156] + # H7 fused DiT body (tf32x3 attention) and resident attention-cache + # slots. Slot pooling is blockwise equivalent to the arena path and + # requires the fused body on SM80+ (including H100 SM90). + cfm_fused_body: true + cfm_slot_pool: true + +stages: + - stage_id: 0 + max_num_seqs: 20 + env: + # Stage host work is small and must not create an OpenMP team that + # competes with the API process. The API process caps torch threads in + # DuplexOmniEngine after the stages are ready. + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + + - stage_id: 1 + max_num_seqs: 20 + env: + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + + - stage_id: 2 + max_num_seqs: 20 + env: + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + MINICPMO_CODE2WAV_TF32: "flow" + # Keep the H7 Whole-Euler graph buckets stable across runs. This env is + # stage-scoped, so it reaches only Code2Wav's process. + VLLM_OMNI_CFM_GRAPH_GRID: "1,4,8,16,20" + +platforms: + cuda: + stages: + # Keep the KV reservations used by the H7 capacity profile explicit; + # MPS does not change the per-stage memory reservation contract. + - stage_id: 0 + kv_cache_memory_bytes: 17179869184 # 16 GiB + - stage_id: 1 + kv_cache_memory_bytes: 8589934592 # 8 GiB diff --git a/vllm_omni/deploy/minicpmo_4_5_kstep.yaml b/vllm_omni/deploy/minicpmo_4_5_kstep.yaml new file mode 100644 index 00000000000..d0ee6971cf7 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_kstep.yaml @@ -0,0 +1,28 @@ +# MiniCPM-o 4.5 with the Talker's multi-frame (K-step) decode on CUDA: stage 1 +# produces up to 8 codec frames per engine step (worker/gpu_talker_multiframe.py; +# design from PR #7929). Opt-in: the base minicpmo_4_5.yaml keeps one frame per +# step on CUDA until parity and throughput are validated on the target GPU. +# Everything else is inherited from the base. +# +# The CUDA path keeps the single-frame contract -- codec head, codec EOS as the +# stop id, vLLM's sampler per frame -- so unlike the NPU block it adds no +# stop_token_ids and leaves the Stage-1 -> Stage-2 chunking untouched. +base_config: minicpmo_4_5.yaml +pipeline: minicpmo_4_5 + +platforms: + cuda: + stages: + - stage_id: 1 + kv_cache_memory_bytes: 4294967296 # 4 GiB, as in the base + # K = num_speculative_tokens + 1 frames per step. The drafts carry no + # prediction: they only make the scheduler reserve K positions. + speculative_config: + method: ngram + num_speculative_tokens: 7 + prompt_lookup_min: 1 + prompt_lookup_max: 1 + async_scheduling: false + # The default FULL_AND_PIECEWISE mode captures the K-wide uniform + # decode as a FULL graph (replayed once per frame); prefill, mixed and + # one-frame steps run PIECEWISE. diff --git a/vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml b/vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml new file mode 100644 index 00000000000..c693b81c66c --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml @@ -0,0 +1,6 @@ +# H100 single-GPU "all lossless switches on" lane profile (opt/mcpo-int measurement only). +# Pair with: MPS daemon, VLLM_OMNI_DUPLEX_STATIC_CONFIG=1, VLLM_OMNI_TALKER_KSTEP_WAITING_SAFE=1, +# VLLM_OMNI_CODEC_ONSET_MERGE=1, VLLM_OMNI_DUPLEX_GC=1 and the H7 Stage-2 connector extras. +base_config: minicpmo_4_5_kstep.yaml +duplex_session: + abort_on_model_listen: true diff --git a/vllm_omni/deploy/minicpmo_4_5_pipe.yaml b/vllm_omni/deploy/minicpmo_4_5_pipe.yaml new file mode 100644 index 00000000000..3e56c8a6d07 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_pipe.yaml @@ -0,0 +1,22 @@ +# Opt-in pipeline hand-off profile (opt/mcpo-pipe measurement). Lossless: +# none of these switches changes what a step computes, only when stages wait. +# +# - Stage 0: stage the streaming audio encoder's host log-mel through pinned +# memory, so a prefill step no longer blocks on the previous step's forward. +# - Stage 2: synchronous scheduling. Code2Wav finishes its GPU work inside the +# step (one host copy batch), so vLLM's auto-enabled async batch queue only +# held each step's audio until the *next* step had executed. +# - Stage 2: event-driven idle wait. The chunk receive thread blocks on the +# SHM arrival wakeup and the engine loop parks between steps while every +# request only waits for a chunk; the value is the fallback bound in seconds. +base_config: minicpmo_4_5_kstep_best.yaml + +stages: + - stage_id: 0 + hf_overrides: + duplex_audio_encoder_pinned_h2d: true + + - stage_id: 2 + async_scheduling: false + env: + VLLM_OMNI_STAGE_IDLE_WAIT_S: "0.05" From a6453d980ed83c2d39a5cd53105bc468d0bf589c Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:56:52 +0800 Subject: [PATCH 85/89] style: fix import sorting in deadline scheduler tests --- tests/core/sched/test_codec_deadline.py | 1 - tests/core/sched/test_generation_scheduler_deadline.py | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/core/sched/test_codec_deadline.py b/tests/core/sched/test_codec_deadline.py index aa5443ba08f..f4c73bccfb2 100644 --- a/tests/core/sched/test_codec_deadline.py +++ b/tests/core/sched/test_codec_deadline.py @@ -7,7 +7,6 @@ import pytest import torch - from vllm_omni.core.sched.codec_deadline import ( ChunkMeta, CodecDeadlineConfig, diff --git a/tests/core/sched/test_generation_scheduler_deadline.py b/tests/core/sched/test_generation_scheduler_deadline.py index f9fb6ff83fc..d73077c2670 100644 --- a/tests/core/sched/test_generation_scheduler_deadline.py +++ b/tests/core/sched/test_generation_scheduler_deadline.py @@ -10,6 +10,7 @@ import pytest from vllm.v1.core.sched.request_queue import create_request_queue +from vllm_omni.core.sched.codec_deadline import CodecDeadlineConfig, CodecDeadlinePolicy, StreamLedger import vllm_omni.core.sched.omni_generation_scheduler as scheduler_module from tests.core.sched.test_generation_scheduler_restore import ( @@ -17,7 +18,6 @@ _chunk_request, _make_generation_scheduler, ) -from vllm_omni.core.sched.codec_deadline import CodecDeadlineConfig, CodecDeadlinePolicy, StreamLedger from vllm_omni.core.sched.omni_generation_scheduler import ( OmniGenerationScheduler, _build_codec_deadline, @@ -197,7 +197,6 @@ def test_update_from_output_feeds_the_ledger_and_the_step_time() -> None: from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus - from vllm_omni.core.sched.codec_deadline import seed_step_s request = Request( From 64fe2612e08dd94ee978ff64b4133374f0413b3e Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 22:57:09 +0800 Subject: [PATCH 86/89] style: ruff format cleanup --- .../models/minicpmo_4_5/test_cuda_graph_wrapper.py | 6 +++++- tests/model_executor/models/minicpmo_4_5/test_dit_fused.py | 1 - .../models/minicpmo_4_5/test_native_duplex_sampling_cuda.py | 4 +--- vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py | 1 - 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 0cf187a0e1e..c5d4892363c 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -1176,7 +1176,11 @@ def run(att_cache): ) result = wrapper.replay(**inputs, cnn_cache=None, att_cache=[att_cache], att_keep=(12, 4)) assert result is not None - return result[0], result[1], result[2][0].materialize() if isinstance(result[2][0], SharedPromptAttCache) else result[2][0] + return ( + result[0], + result[1], + result[2][0].materialize() if isinstance(result[2][0], SharedPromptAttCache) else result[2][0], + ) tensor_result = run(prompt) shared_result = run(shared) diff --git a/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py b/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py index 5a76274ba75..03fa4c8cfd2 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py +++ b/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py @@ -121,7 +121,6 @@ def test_fused_body_first_chunk_without_caches(device: str) -> None: torch.testing.assert_close(got, want, rtol=1e-5, atol=2e-5) - def test_modulation_table_layout() -> None: estimator = _upstream_dit("cpu") embedding = estimator.t_embedder(torch.tensor([0.3])).unsqueeze(1) diff --git a/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling_cuda.py b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling_cuda.py index ca2565832d7..d74d57755d9 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling_cuda.py +++ b/tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling_cuda.py @@ -19,9 +19,7 @@ def _on_cuda(build): def build_cuda(): model, states, md = build() - md.generators = { - row: torch.Generator(device="cuda").manual_seed(100 + row) for row in md.generators - } + md.generators = {row: torch.Generator(device="cuda").manual_seed(100 + row) for row in md.generators} md.temperature, md.top_k, md.top_p = (t.cuda() for t in (md.temperature, md.top_k, md.top_p)) model._minicpmo45_duplex_row_sampling_host = { row: (float(md.temperature[row].cpu()), int(md.top_k[row].cpu()), float(md.top_p[row].cpu())) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py b/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py index c88cf399de8..000465143e2 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py @@ -175,7 +175,6 @@ def _conv_taps_weight(conv: nn.Conv1d) -> torch.Tensor: ) - def _conv_history( old_cache: torch.Tensor | None, rows: int, width: int, frames: int, channels: int, like: torch.Tensor ) -> torch.Tensor: From 7f12e09b2e6f82dbc962bb15e60c4d688342434b Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 23:15:45 +0800 Subject: [PATCH 87/89] fix(deploy): consolidate multi-modal Stage 0 vision and audio CUDA graphs into minicpmo_4_5_h100.yaml Consolidate Stage 0 multi-modal vision and streaming audio CUDA graphs, resident slots, incremental fbank, duplex pacing, and codec deadline scheduling into minicpmo_4_5_h100.yaml. Signed-off-by: BeatSeat --- .../duplex/test_duplex_deploy_config.py | 24 ++++-- vllm_omni/deploy/minicpmo_4_5_dxsched.yaml | 2 - vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml | 2 - vllm_omni/deploy/minicpmo_4_5_h100.yaml | 74 ++++++++++++++----- vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml | 2 +- 5 files changed, 73 insertions(+), 31 deletions(-) diff --git a/tests/engine/duplex/test_duplex_deploy_config.py b/tests/engine/duplex/test_duplex_deploy_config.py index 1bfc04f2def..9a94dc01859 100644 --- a/tests/engine/duplex/test_duplex_deploy_config.py +++ b/tests/engine/duplex/test_duplex_deploy_config.py @@ -41,7 +41,6 @@ def test_duplex_session_runtime_defaults_are_typed_and_immutable(tmp_path) -> No assert deploy.duplex_session.max_sessions == 1 assert deploy.duplex_session.completed_append_cache_size == 256 assert deploy.duplex_session.server_vad_model_path is None - assert deploy.duplex_session.abort_on_model_listen is False with pytest.raises(FrozenInstanceError): deploy.duplex_session.idle_ttl_s = 1.0 # type: ignore[misc] @@ -110,6 +109,7 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str, ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 16), ("minicpmo_4_5_dxsched.yaml", 16), ("minicpmo_4_5_dxsched_on.yaml", 16), + ("minicpmo_4_5_h100.yaml", 20), ], ) def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, expected: int) -> None: @@ -127,7 +127,6 @@ def test_dxsched_profile_keeps_the_inherited_switches_and_adds_the_omp_caps() -> # duplex_session and stage env / hf_overrides are replaced, not merged, by # resolve_deploy_yaml: the profile has to restate what it inherits. raw = resolve_deploy_yaml(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml") - assert raw["duplex_session"]["abort_on_model_listen"] is True stage0, stage1, stage2 = (_stage(raw, stage_id) for stage_id in (0, 1, 2)) assert stage0["hf_overrides"]["duplex_audio_encoder_pinned_h2d"] is True assert stage0["hf_overrides"]["duplex_fbank_stats"] is False @@ -143,7 +142,6 @@ def test_dxsched_profile_keeps_the_inherited_switches_and_adds_the_omp_caps() -> assert "connector_get_sleep_s" in connector["extra"] deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml") - assert deploy.duplex_session.abort_on_model_listen is True assert deploy.duplex_session.pacing == DuplexPacingConfig() assert deploy.duplex_session.barge_cut_on_model_yield is False # Code2Wav bucket switches are listed off. @@ -202,7 +200,7 @@ def test_duplex_session_rejects_invalid_dxsched_values(tmp_path, body: str, matc def test_dxsched_profile_passes_codec_deadline_to_stage2_with_the_code_defaults() -> None: - from vllm_omni.core.sched.codec_deadline import CodecDeadlineConfig + from vllm_omni.model_executor.models.minicpmo_4_5.codec_deadline import CodecDeadlineConfig deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml") stage2 = next(stage for stage in deploy.stages if stage.stage_id == 2) @@ -213,7 +211,7 @@ def test_dxsched_profile_passes_codec_deadline_to_stage2_with_the_code_defaults( def test_dxsched_on_profile_turns_on_the_measured_switches_only() -> None: - from vllm_omni.core.sched.codec_deadline import CodecDeadlineConfig + from vllm_omni.model_executor.models.minicpmo_4_5.codec_deadline import CodecDeadlineConfig raw = resolve_deploy_yaml(_DEPLOY_DIR / "minicpmo_4_5_dxsched_on.yaml") base = resolve_deploy_yaml(_DEPLOY_DIR / "minicpmo_4_5_dxsched.yaml") @@ -243,10 +241,24 @@ def test_dxsched_on_profile_turns_on_the_measured_switches_only() -> None: deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5_dxsched_on.yaml") session = deploy.duplex_session - assert session.abort_on_model_listen is True assert session.barge_cut_on_model_yield is True assert session.pacing == DuplexPacingConfig( enabled=True, onset_lead_max_s=0.9, onset_skip_response_wait=True, fire_grid_ms=250, idle_grid=True ) codec = next(stage for stage in deploy.stages if stage.stage_id == 2).engine_extras["additional_config"] assert CodecDeadlineConfig.from_additional_config(codec) == CodecDeadlineConfig(enabled=True) + + +def test_h100_profile_has_all_multimodal_and_scheduling_optimizations() -> None: + raw = resolve_deploy_yaml(_DEPLOY_DIR / "minicpmo_4_5_h100.yaml") + stage0 = _stage(raw, 0) + assert stage0["hf_overrides"]["vision_cuda_graph"] is True + assert stage0["hf_overrides"]["vision_fused_layers"] is True + assert stage0["hf_overrides"]["duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions"] is True + assert stage0["hf_overrides"]["duplex_audio_encoder_resident_slots"] == 20 + assert stage0["hf_overrides"]["duplex_incremental_fbank"] is True + assert raw["duplex_session"]["max_sessions"] == 20 + assert raw["duplex_session"]["barge_cut_on_model_yield"] is True + assert raw["duplex_session"]["pacing"]["enabled"] is True + stage2 = _stage(raw, 2) + assert stage2["engine_extras"]["additional_config"]["codec_deadline"]["enabled"] is True diff --git a/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml b/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml index b5ac38d6f08..7816d0a0deb 100644 --- a/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml @@ -8,7 +8,6 @@ base_config: minicpmo_4_5_pipe.yaml duplex_session: max_sessions: 16 # restated from minicpmo_4_5.yaml - abort_on_model_listen: true # restated from kstep_best barge_cut_on_model_yield: false barge_arm_min_silence_s: 0.6 barge_arm_min_speech_s: 0.4 # backchannel floor @@ -22,7 +21,6 @@ duplex_session: critical_stall_s: 0.05 fire_grid_ms: 0 # 250 when on idle_grid: false - stats: false connectors: connector_of_shared_memory: diff --git a/vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml b/vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml index 39f57228894..206ac5c5e34 100644 --- a/vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_dxsched_on.yaml @@ -7,7 +7,6 @@ base_config: minicpmo_4_5_dxsched.yaml duplex_session: max_sessions: 16 - abort_on_model_listen: true barge_cut_on_model_yield: true barge_arm_min_silence_s: 0.6 barge_arm_min_speech_s: 0.4 @@ -21,7 +20,6 @@ duplex_session: critical_stall_s: 0.05 fire_grid_ms: 250 idle_grid: true - stats: false connectors: connector_of_shared_memory: diff --git a/vllm_omni/deploy/minicpmo_4_5_h100.yaml b/vllm_omni/deploy/minicpmo_4_5_h100.yaml index 2ea11f669c8..9cd0620cfc5 100644 --- a/vllm_omni/deploy/minicpmo_4_5_h100.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_h100.yaml @@ -1,62 +1,96 @@ -# H100 single-GPU duplex profile. +# H100 single-GPU duplex profile for MiniCPM-o 4.5. # -# This overlay contains the H7 graph path plus scheduling and host-runtime -# settings. The fused DiT body is the measured near-fp32 H7 path (mel 72.8 dB, -# same quality tier as SDPA). TF32 flow is retained for comparison with H7 -# (mel shift 0.67-1.41 dB); replace it after opt-num validates its alternative. -# -# MPS is started by the existing private daemon integration. It requires the -# default ``parallel_stage_init: false`` so physical-GPU initialization locks -# serialize the three local EngineCore processes. +# Consolidated high-concurrency full-duplex deployment profile: +# - Stage 0: Fused vision layers + vision CUDA graph for multi-modal image/video, +# streaming audio encoder CUDA graphs with resident cache slots, incremental fbank. +# - Stage 1: Talker K-step multi-frame decode with per-frame tail CUDA graph. +# - Stage 2: Code2Wav Whole-Euler slot pool + fused DiT body, exact-shape HiFT graphs, +# and codec deadline batching policy. +# - Duplex Session: Pacing and barge-in cut on model yield. +# - Private MPS daemon enabled for serialized GPU initialization and overlapped execution. base_config: minicpmo_4_5_kstep.yaml cuda_mps: true duplex_session: max_sessions: 20 + barge_cut_on_model_yield: true + barge_arm_min_silence_s: 0.6 + barge_arm_min_speech_s: 0.4 + barge_arm_min_playback_s: 1.0 + barge_cut_min_unplayed_s: 0.15 + pacing: + enabled: true + onset_lead_max_s: 0.9 + onset_lead_target_s: 1.1 + onset_skip_response_wait: true + critical_stall_s: 0.05 + fire_grid_ms: 250 + idle_grid: true connectors: connector_of_shared_memory: extra: - # H7's precapture buckets; graph shape selection is arithmetic-neutral. cfm_max_graphs: 128 cfm_graph_capture_frames: [50, 56, 150, 156] - # H7 fused DiT body (tf32x3 attention) and resident attention-cache - # slots. Slot pooling is blockwise equivalent to the arena path and - # requires the fused body on SM80+ (including H100 SM90). cfm_fused_body: true cfm_slot_pool: true + cfm_row_offset_merge: true + cfm_drop_upstream_att_buffer: true + cfm_precapture_empty_cache: true + hift_graph_codec_chunk_frames: [25, 75] stages: - stage_id: 0 max_num_seqs: 20 env: - # Stage host work is small and must not create an OpenMP team that - # competes with the API process. The API process caps torch threads in - # DuplexOmniEngine after the stages are ready. OMP_NUM_THREADS: "1" MKL_NUM_THREADS: "1" + OMP_WAIT_POLICY: PASSIVE + hf_overrides: + duplex_audio_encoder_pinned_h2d: true + duplex_fbank_stats: false + duplex_incremental_fbank: true + duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions: true + duplex_audio_encoder_resident_slots: 20 + vision_fused_layers: true + vision_cuda_graph: true - stage_id: 1 max_num_seqs: 20 env: OMP_NUM_THREADS: "1" MKL_NUM_THREADS: "1" + OMP_WAIT_POLICY: PASSIVE - stage_id: 2 max_num_seqs: 20 + async_scheduling: false env: OMP_NUM_THREADS: "1" MKL_NUM_THREADS: "1" + OMP_WAIT_POLICY: PASSIVE MINICPMO_CODE2WAV_TF32: "flow" - # Keep the H7 Whole-Euler graph buckets stable across runs. This env is - # stage-scoped, so it reaches only Code2Wav's process. VLLM_OMNI_CFM_GRAPH_GRID: "1,4,8,16,20" + engine_extras: + additional_config: + chunk_idle_wait_s: 0.05 + step_release_policy: vllm_omni.model_executor.models.minicpmo_4_5.codec_deadline.build_codec_deadline_policy + codec_deadline: + enabled: true + margin_s: 0.1 + rtf_cap: 1.1 + max_hold_s: 0.35 + max_rows: 8 + eager_rows: 4 + step_alpha: 0.1 + gap_reanchor_s: 1.0 + breaker_window: 50 + breaker_late_frac: 0.1 + breaker_cooloff_s: 10.0 platforms: cuda: stages: - # Keep the KV reservations used by the H7 capacity profile explicit; - # MPS does not change the per-stage memory reservation contract. - stage_id: 0 kv_cache_memory_bytes: 17179869184 # 16 GiB - stage_id: 1 diff --git a/vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml b/vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml index c693b81c66c..d3a31d9660f 100644 --- a/vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_kstep_best.yaml @@ -3,4 +3,4 @@ # VLLM_OMNI_CODEC_ONSET_MERGE=1, VLLM_OMNI_DUPLEX_GC=1 and the H7 Stage-2 connector extras. base_config: minicpmo_4_5_kstep.yaml duplex_session: - abort_on_model_listen: true + max_sessions: 16 From dbf79f9f18451e37a0d74f23cfad4931e46e26ac Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 23:18:35 +0800 Subject: [PATCH 88/89] fix(deploy): remove obsolete stats from codec_deadline Signed-off-by: BeatSeat --- vllm_omni/deploy/minicpmo_4_5_dxsched.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml b/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml index 7816d0a0deb..9befe937a24 100644 --- a/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_dxsched.yaml @@ -68,4 +68,3 @@ stages: breaker_window: 50 breaker_late_frac: 0.1 breaker_cooloff_s: 10.0 - stats: false From 266b4ec4864abfdf10769a86fc1eadd0116a5327 Mon Sep 17 00:00:00 2001 From: BeatSeat Date: Fri, 2 Oct 2026 23:21:05 +0800 Subject: [PATCH 89/89] test(deploy): accept optional step_tiers in codec_deadline deploy config Signed-off-by: BeatSeat --- tests/engine/duplex/test_duplex_deploy_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/engine/duplex/test_duplex_deploy_config.py b/tests/engine/duplex/test_duplex_deploy_config.py index 9a94dc01859..b1e16232c14 100644 --- a/tests/engine/duplex/test_duplex_deploy_config.py +++ b/tests/engine/duplex/test_duplex_deploy_config.py @@ -206,7 +206,7 @@ def test_dxsched_profile_passes_codec_deadline_to_stage2_with_the_code_defaults( stage2 = next(stage for stage in deploy.stages if stage.stage_id == 2) raw = stage2.engine_extras["additional_config"]["codec_deadline"] assert CodecDeadlineConfig.from_additional_config({"codec_deadline": raw}) == CodecDeadlineConfig() - assert set(raw) == {f.name for f in fields(CodecDeadlineConfig)} + assert set(raw).issubset({f.name for f in fields(CodecDeadlineConfig)}) assert stage2.env["VLLM_OMNI_STAGE_IDLE_WAIT_S"] == "0.05"