From c11596c63711326ba1af4527f949f1558c63a8b4 Mon Sep 17 00:00:00 2001 From: shihongzhi Date: Tue, 22 Sep 2026 09:34:12 +0800 Subject: [PATCH 1/7] [Model] MammothModa2: enable Cache-DiT acceleration for the DiT stage (#7291) Signed-off-by: shihongzhi --- docs/design/feature/cache_dit.md | 48 ++++- .../diffusion/cache_acceleration/cache_dit.md | 33 ++++ docs/user_guide/diffusion_features.md | 3 +- .../text_to_image/text_to_image.py | 6 +- tests/diffusion/cache/test_cache_dit.py | 108 ++++++++++ .../test_pipeline_mammothmoda2_dit.py | 186 +++++++++++++++++- vllm_omni/deploy/mammoth_moda2.yaml | 10 + .../cache/cachedit/model_specific.py | 42 ++++ .../mammoth_moda2/mammothmoda2_dit_model.py | 10 +- .../pipeline_mammothmoda2_dit.py | 75 ++++++- 10 files changed, 507 insertions(+), 14 deletions(-) diff --git a/docs/design/feature/cache_dit.md b/docs/design/feature/cache_dit.md index 38b597ca3..69a8de613 100644 --- a/docs/design/feature/cache_dit.md +++ b/docs/design/feature/cache_dit.md @@ -180,17 +180,52 @@ cache_dit.enable_cache( > **Note:** For single transformer with multiple block lists, `refresh_context` works the same as standard models. +### Example 3: Sequential-CFG Diffusion Pipeline (MammothModa2) + +MammothModa2's DiT stage runs on the diffusion runner, which enables the +configured backend at startup and transfers ownership to the pipeline through +the request-scoped protocol (`adopt_cache_dit_backend` / `is_cache_dit_enabled`); +`forward()` then reconciles per-request state through +`RequestScopedCacheDiTRuntime` (MiniMax H3 precedent). + +**Key difference:** The enabler caches only the repeated main-layer stack and +marks the pipeline as requiring paired CFG forwards, because cache-dit's +separate-CFG accounting distinguishes cond/uncond passes purely by forward +parity: + +```python +pipeline._cache_dit_requires_paired_cfg = True +block_adapter = BlockAdapter( + transformer=pipeline.gen_transformer, + blocks=[pipeline.gen_transformer.layers], # Refiners stay outside + forward_pattern=[ForwardPattern.Pattern_3], + has_separate_cfg=True, + check_forward_pattern=True, +) +``` + +Per request, the pipeline calls `runtime.prepare(spec)` for CFG requests +(refreshing the step context) and `runtime.prepare(None)` for `guidance=1.0` +requests (disabling hooks — their single forward per step cannot be expressed +by the parity accounting). + ### Registering Custom Implementations -After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in +After writing your custom enabler, register it in +`CUSTOM_DIT_ENABLERS` via `register_custom_dit_enablers()` in `vllm_omni/diffusion/cache/cachedit/model_specific.py`: ```python -CUSTOM_DIT_ENABLERS = { - "Wan22Pipeline": enable_cache_for_wan22, - "LongCatImagePipeline": enable_cache_for_longcat_image, - "YourCustomPipeline": enable_cache_for_your_model, # Add here -} +CUSTOM_DIT_ENABLERS.update( + { + "Wan22Pipeline": enable_cache_for_wan22, + "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3, + "Krea2Pipeline": enable_cache_for_krea2, + "Magi2Pipeline": enable_cache_for_magi2, + "MammothModa2DiTPipeline": enable_cache_for_mammothmoda2, + "YourCustomPipeline": enable_cache_for_your_model, # Add here + } +) ``` --- @@ -274,6 +309,7 @@ Complete examples in the codebase: | **Wan2.2** | `cachedit.model_specific::enable_cache_for_wan22` | Single or dual-transformer | Auto-detects mode based on transformer_2 presence | | **LongCat** | `cachedit.config::CacheDiTAdapterConfig` | Declarative block adapter | Two block lists in one transformer | | **BAGEL** | `cachedit.model_specific::BagelCachedAdapter` | Custom cached adapter | Complex architecture | +| **MammothModa2** | `cachedit.model_specific::enable_cache_for_mammothmoda2` | Sub-stack + paired CFG | Generation-runner pipeline; lifecycle owned by `pipeline_mammothmoda2_dit.py` | --- diff --git a/docs/user_guide/diffusion/cache_acceleration/cache_dit.md b/docs/user_guide/diffusion/cache_acceleration/cache_dit.md index 4d1f69068..3b686a32e 100644 --- a/docs/user_guide/diffusion/cache_acceleration/cache_dit.md +++ b/docs/user_guide/diffusion/cache_acceleration/cache_dit.md @@ -72,6 +72,39 @@ omni = Omni( ## Example Script +### Multi-Stage Models (MammothModa2) + +MammothModa2 is a multi-stage AR→DiT pipeline; the diffusion runner owns the +Cache-DiT lifecycle of its DiT stage. Configure it on the **DiT stage entry** +of the deploy YAML so the acceleration only applies to the denoising stage: + +```yaml +# deploy YAML (see vllm_omni/deploy/mammoth_moda2.yaml) +stages: + - stage_id: 1 + # ... other DiT stage settings ... + cache_backend: cache_dit + cache_config: + Fn_compute_blocks: 1 + Bn_compute_blocks: 0 + max_warmup_steps: 4 + residual_diff_threshold: 0.24 + max_continuous_cached_steps: 3 + enable_cache_dit_summary: true # log skip-ratio stats per request +``` + +Behavior notes: + +- Only the repeated main-layer stack is cached; the Q-Former refiners always run. +- Sequential-CFG parity requires the unconditional pass on every denoise step, + so the `cfg_range` skip optimization is not used: outside the interval CFG is + neutralized with `scale=1.0` instead of skipping the uncond forward. +- Requests with `text_guidance_scale = 1.0` run with cache hooks disabled. +- `residual_diff_threshold` defaults to `0.24` (measured ~1.7-2.2x at 50 steps, + PSNR 27-34 dB vs uncached). For 20-step generation — especially with a + partial `cfg_range` — use `0.12` for noticeably closer parity at a modest + speed cost. + ### Offline Inference Use the example script under `examples/offline_inference/text_to_image`: diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index a8fb6d972..703a83f38 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -130,7 +130,7 @@ The following tables show which models support each feature: | **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **MammothModa2(T2I)** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | @@ -154,6 +154,7 @@ The following tables show which models support each feature: > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`. > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). > 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). +> 7. MammothModa2 runs its DiT stage on the diffusion runner (`StageExecutionType.DIFFUSION`); Cache-DiT is enabled through the standard diffusion-stage knobs on the stage entry of the deploy YAML (`cache_backend: cache_dit`, plus optional `cache_config` / `enable_cache_dit_summary`). The runner installs the backend at startup and the pipeline adopts it per request. Only the repeated main-layer stack is cached; requests with `text_guidance_scale = 1.0` bypass the cache hooks. ### VideoGen diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index c74236316..a8624b9ae 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -157,7 +157,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--enable-cache-dit-summary", action="store_true", - help="Enable cache-dit summary logging after diffusion forward passes.", + default=None, + help=( + "Enable cache-dit summary logging after diffusion forward passes. " + "Default: unset (defer to the deploy YAML's enable_cache_dit_summary)." + ), ) parser.add_argument( "--ulysses-degree", diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index c1d061ef2..7de2aaa38 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -60,6 +60,7 @@ def test_custom_cache_dit_enablers_are_registered_explicitly(): "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3, "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2, "Magi2Pipeline": cd_model_specific.enable_cache_for_magi2, + "MammothModa2DiTPipeline": cd_model_specific.enable_cache_for_mammothmoda2, } with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True): @@ -116,6 +117,33 @@ def test_magi2_cache_dit_targets_only_nested_repeated_layers(mock_block_adapter, assert get_transformer(pipeline) is transformer_block +@patch("vllm_omni.diffusion.cache.cachedit.model_specific.enable_cache_for_dit") +@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter") +def test_mammothmoda2_cache_dit_targets_only_main_layers(mock_block_adapter, mock_enable_cache): + pipeline = Mock() + transformer = pipeline.gen_transformer + layers = torch.nn.ModuleList([torch.nn.Identity()]) + transformer.layers = layers + adapter = mock_block_adapter.return_value + refresh = Mock() + mock_enable_cache.return_value = refresh + + result = cd_model_specific.enable_cache_for_mammothmoda2(pipeline, SAMPLE_CACHE_CONFIG) + + mock_block_adapter.assert_called_once() + adapter_kwargs = mock_block_adapter.call_args.kwargs + assert adapter_kwargs["transformer"] is transformer + assert adapter_kwargs["blocks"] == [layers] + assert adapter_kwargs["has_separate_cfg"] is True + assert adapter_kwargs["check_forward_pattern"] is True + assert result.refresh is refresh + assert result.targets == (adapter,) + get_transformer = mock_enable_cache.call_args.kwargs["get_pipeline_transformer"] + assert get_transformer(pipeline) is transformer + # Sequential-CFG parity is mandatory while cache-dit is installed. + assert pipeline._cache_dit_requires_paired_cfg is True + + @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit.summary") @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter.is_cached", return_value=True) def test_cache_summary_uses_custom_nested_targets(mock_is_cached, mock_summary): @@ -338,3 +366,83 @@ def test_summary_with_no_transformer_is_nonfatal(): pass cache_summary(pipeline=FakePipeline()) + + +# This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend +@pytest.mark.skipif( + current_omni_platform.is_rocm(), + reason="vLLM ROCm custom ops lack CPU fallback", +) +def test_mammothmoda2_cache_dit_runs_end_to_end_on_tiny_model(request: pytest.FixtureRequest): + """A tiny MammothModa2 DiT runs through Cache-DiT with paired CFG forwards. + + Exercises the Pattern_3 contract end to end: after enabling the custom + enabler, blocks must be callable as ``block(hidden_states, **kwargs)`` + and the pipeline's two-forwards-per-step CFG cadence must drive the + cache context without shape or parity errors. + """ + from cache_dit import BlockAdapter + + from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel + from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal + + # MammothModa2 attention resolves through the shared Omni attention layer, + # which picks the platform-default backend at construction time. Pin + # TORCH_SDPA so a CUDA-visible host does not select FA3 for CPU tensors. + with _force_torch_sdpa(): + model = Transformer2DModel( + patch_size=2, + in_channels=4, + hidden_size=96, + num_layers=4, + num_refiner_layers=1, + num_attention_heads=2, + num_kv_heads=2, + multiple_of=8, + axes_dim_rope=(16, 16, 16), + axes_lens=(300, 512, 512), + text_feat_dim=16, + ) + model.eval() + + MammothModa2Pipeline = type("MammothModa2DiTPipeline", (), {}) + pipeline = MammothModa2Pipeline() + pipeline.gen_transformer = model + backend = CacheDiTBackend(DiffusionCacheConfig()) + backend.enable(pipeline) + request.addfinalizer(lambda: backend.disable(pipeline)) + assert BlockAdapter.is_cached(model) + + num_inference_steps = 8 + backend.refresh(pipeline, num_inference_steps=num_inference_steps) + + freqs_cis = RotaryPosEmbedReal.get_freqs_real((16, 16, 16), (300, 512, 512), theta=10000) + latents = torch.randn(1, 4, 8, 8) + text = torch.randn(1, 5, 16) + negative = torch.randn(1, 3, 16) + text_mask = torch.ones(1, 5, dtype=torch.bool) + negative_mask = torch.ones(1, 3, dtype=torch.bool) + + with torch.no_grad(): + for step in range(num_inference_steps): + timestep = torch.full((1,), 1.0 - step / num_inference_steps) + # Conditional pass ... + cond = model( + hidden_states=latents, + timestep=timestep, + text_hidden_states=text, + freqs_cis=freqs_cis, + text_attention_mask=text_mask, + ) + # ... always followed by the unconditional pass (parity). + uncond = model( + hidden_states=latents, + timestep=timestep, + text_hidden_states=negative, + freqs_cis=freqs_cis, + text_attention_mask=negative_mask, + ) + latents = uncond + 4.0 * (cond - uncond) + + assert latents.shape == (1, 4, 8, 8) + assert torch.isfinite(latents).all() diff --git a/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py b/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py index 776ca5bc6..b49ada83e 100644 --- a/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py +++ b/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py @@ -1,13 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from contextlib import contextmanager from dataclasses import dataclass +from types import SimpleNamespace +import cache_dit import pytest import torch +from cache_dit import BlockAdapter from torch import nn -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, TransformerConfig +from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, RequestScopedCacheDiTRuntime +from vllm_omni.diffusion.cache.selector import get_cache_backend +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import ( + AttentionConfig, + DiffusionCacheConfig, + DiffusionOutput, + OmniDiffusionConfig, + TransformerConfig, +) +from vllm_omni.diffusion.models.interface import adopt_request_scoped_cache_dit from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import ( MammothModa2DiTPipeline, @@ -17,6 +31,7 @@ from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import ( from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -110,6 +125,10 @@ def _pipeline_shell() -> MammothModa2DiTPipeline: pipeline.device = torch.device("cpu") pipeline.config = _build_mammoth_config(_od_config()) pipeline._llm_hidden_size = 8 + # Cache-DiT runtime state that __init__ normally creates (disabled here, + # matching a pipeline built without a cache_backend). + pipeline._cache_dit_runtime = RequestScopedCacheDiTRuntime(pipeline) + pipeline._cache_dit_config = None return pipeline @@ -118,12 +137,14 @@ def _batch( request_id: str = "req-7", prompt: object | None = None, sampling: OmniDiffusionSamplingParams | None = None, + height: int = 32, + width: int = 48, ) -> DiffusionRequestBatch: if prompt is None: prompt = { "prompt": "", - "height": 32, - "width": 48, + "height": height, + "width": width, "additional_information": { "full_hidden_states": torch.arange(32, dtype=torch.float32).reshape(4, 8), "full_token_ids": [10, 11, 100, 101], @@ -390,3 +411,162 @@ def test_forward_rejects_missing_visual_tokens_before_model_access() -> None: } with pytest.raises(ValueError, match="no visual-token hidden states.*req-empty"): _pipeline_shell().forward(_batch(request_id="req-empty", prompt=prompt)) + + +@contextmanager +def _force_torch_sdpa(): + """Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3).""" + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"), + parallel_config=SimpleNamespace(ring_degree=1), + ) + with set_current_diffusion_config(od_config): + yield + + +def _constructible_raw_config() -> dict: + """Checkpoint config whose DiT/VAE are small enough to actually run on CPU.""" + return { + "model_type": "mammothmoda2", + "llm_config": { + "model_type": "mammothmoda2_qwen2_5_vl", + "text_config": { + "model_type": "mammothmoda2_qwen2_5_vl_text", + "hidden_size": 8, + "gen_vocab_start_index": 100, + }, + }, + "gen_vae_config": {"block_out_channels": [8, 8], "norm_num_groups": 8}, + "gen_dit_config": { + "patch_size": 2, + "in_channels": 4, + "hidden_size": 96, + "num_layers": 4, + "num_refiner_layers": 1, + "num_attention_heads": 2, + "num_kv_heads": 2, + "multiple_of": 8, + "axes_dim_rope": [16, 16, 16], + "axes_lens": [300, 512, 512], + "text_feat_dim": 16, + }, + # Top-level axes feed the pipeline-owned rotary table (see pipeline __init__). + "gen_axes_dim_rope": [16, 16, 16], + "gen_axes_lens": [300, 512, 512], + } + + +def _cache_dit_od_config() -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model="/models/MammothModa2-Preview", + model_class_name="MammothModa2DiTPipeline", + tf_model_config=TransformerConfig.from_dict(_constructible_raw_config()), + cache_backend="cache_dit", + cache_config=DiffusionCacheConfig(), + ) + + +@pytest.mark.skipif( + current_omni_platform.is_rocm(), + reason="vLLM ROCm custom ops lack CPU fallback", +) +def test_forward_transitions_request_scoped_cache_dit_across_requests(request: pytest.FixtureRequest) -> None: + """Cache-DiT request lifecycle on one production-built pipeline instance. + + Mirrors the diffusion runner startup (``get_cache_backend`` → ``enable`` → + ``adopt_request_scoped_cache_dit``) on a tiny pipeline built through the + real ``__init__``, then drives three requests through the real ``forward()``: + + 1. a partial-CFG request installs the request-scoped cache context: + cond/uncond forwards stay paired on every denoise step (parity + accounting), even outside ``cfg_range``, and blocks record cached steps; + 2. a guidance=1.0 request runs with hooks disabled (``prepare(None)``) and + its output is bit-identical to a reference captured before the backend + was ever enabled, proving a clean uninstall with no stale module state; + 3. a CFG request with a different step count and image size re-installs the + hooks through a fresh context, so no stale shape or step count survives. + """ + torch.manual_seed(0) # Deterministic random init → deterministic cache hits. + with _force_torch_sdpa(): + pipeline = MammothModa2DiTPipeline(od_config=_cache_dit_od_config()) + pipeline.eval() + assert pipeline._cache_dit_config is not None + request.addfinalizer(pipeline._cache_dit_runtime.disable) + + # The default prompt carries 2 text + 2 visual condition tokens, so the + # conditional pass sees 4 tokens and the uncond pass (empty prompt) sees 0. + cond_tokens, uncond_tokens = 4, 0 + seq_lens: list[int] = [] + + def record_call(_module, _args, kwargs): + text = kwargs.get("text_hidden_states") + seq_lens.append(int(text.shape[1]) if text is not None else -1) + + hook = pipeline.gen_transformer.register_forward_pre_hook(record_call, with_kwargs=True) + + unguided_sampling = OmniDiffusionSamplingParams( + height=32, width=48, seed=4321, guidance_scale=1.0, num_inference_steps=8 + ) + partial_cfg_sampling = OmniDiffusionSamplingParams( + height=32, + width=48, + seed=1234, + guidance_scale=4.0, + num_inference_steps=8, + extra_args={"cfg_range": [0.2, 0.8]}, + ) + changed_cfg_sampling = OmniDiffusionSamplingParams( + height=48, + width=64, + seed=777, + guidance_scale=6.0, + num_inference_steps=10, + extra_args={"cfg_range": [0.1, 0.9]}, + ) + + # Reference run for phase 2: identical request while hooks were never + # installed on this pipeline. + assert not pipeline.is_cache_dit_enabled() + reference = pipeline.forward(_batch(sampling=unguided_sampling)).output + assert seq_lens == [cond_tokens] * 8 + assert torch.isfinite(reference.float()).all() + + # Runner startup: enable once, then the pipeline owns all transitions. + backend = get_cache_backend("cache_dit", DiffusionCacheConfig()) + assert isinstance(backend, CacheDiTBackend) + backend.enable(pipeline) + assert adopt_request_scoped_cache_dit(pipeline, backend) + assert pipeline.is_cache_dit_enabled() + + # Phase 1: partial CFG keeps cond/uncond paired on every step, including + # the steps whose normalized index falls outside cfg_range [0.2, 0.8]. + seq_lens.clear() + pipeline.forward(_batch(sampling=partial_cfg_sampling)) + assert seq_lens == [cond_tokens, uncond_tokens] * 8 + assert pipeline.is_cache_dit_enabled() + assert BlockAdapter.is_cached(pipeline.gen_transformer) + hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0] + assert hits.cached_steps or hits.cfg_cached_steps + + # Phase 2: no-CFG request disables the hooks and must reproduce the + # never-enabled reference bit-for-bit. + seq_lens.clear() + unguided = pipeline.forward(_batch(sampling=unguided_sampling)).output + assert seq_lens == [cond_tokens] * 8 + assert not pipeline.is_cache_dit_enabled() + assert not BlockAdapter.is_cached(pipeline.gen_transformer) + assert torch.equal(unguided, reference) + + # Phase 3: CFG with a new step count and image size re-installs the hooks + # on a fresh context; output shape follows the new request geometry. + seq_lens.clear() + regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64)).output + assert seq_lens == [cond_tokens, uncond_tokens] * 10 + assert pipeline.is_cache_dit_enabled() + assert BlockAdapter.is_cached(pipeline.gen_transformer) + hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0] + assert hits.cached_steps or hits.cfg_cached_steps + assert regenerated.shape == (1, 3, 12, 16) + assert torch.isfinite(regenerated.float()).all() + + hook.remove() diff --git a/vllm_omni/deploy/mammoth_moda2.yaml b/vllm_omni/deploy/mammoth_moda2.yaml index 76027424c..e0b652915 100644 --- a/vllm_omni/deploy/mammoth_moda2.yaml +++ b/vllm_omni/deploy/mammoth_moda2.yaml @@ -20,3 +20,13 @@ stages: enforce_eager: true trust_remote_code: true enable_prefix_caching: false + # Optional Cache-DiT acceleration for the DiT stage (~2x faster denoise, + # training-free). Uncomment to enable; guidance=1.0 requests bypass it. + # cache_backend: cache_dit + # cache_config: + # Fn_compute_blocks: 1 + # Bn_compute_blocks: 0 + # max_warmup_steps: 4 + # residual_diff_threshold: 0.24 + # max_continuous_cached_steps: 3 + # enable_cache_dit_summary: true diff --git a/vllm_omni/diffusion/cache/cachedit/model_specific.py b/vllm_omni/diffusion/cache/cachedit/model_specific.py index 0a25abd3c..8bc36a8b0 100644 --- a/vllm_omni/diffusion/cache/cachedit/model_specific.py +++ b/vllm_omni/diffusion/cache/cachedit/model_specific.py @@ -867,6 +867,47 @@ def enable_cache_for_magi2(pipeline: Any, cache_config: Any) -> CacheDiTEnableRe return CacheDiTEnableResult(refresh=refresh, targets=(block_adapter,)) +def _get_mammothmoda2_transformer(pipeline: Any) -> torch.nn.Module: + return pipeline.gen_transformer + + +def enable_cache_for_mammothmoda2(pipeline: Any, cache_config: Any) -> CacheDiTEnableResult: + """Cache only MammothModa2's repeated main DiT stack. + + ``Transformer2DModel`` runs three Q-Former refiners (noise / ref-image / + context) whose inputs change every denoise step, plus the ``layers`` + stack that dominates per-step compute. Only ``layers`` is a repeated + residual stack, so the ``BlockAdapter`` wraps it and the refiners stay + outside the cached region. Blocks take ``hidden_states`` plus keyword + step context and return only hidden states (``Pattern_3``). + + The pipeline runs sequential CFG: each conditional forward is followed by + an unconditional forward. cache-dit tells cond/uncond apart purely by + transformer-forward parity (``has_separate_cfg=True``), so the + ``cfg_range`` optimization that skips the unconditional pass outside the + interval would desync that accounting. Like Cosmos3, we keep the passes + paired and neutralize CFG via scale=1.0 outside the interval; the + pipeline also disables hooks for no-CFG requests (single forward per + step), whose parity the accounting cannot represent. + """ + pipeline._cache_dit_requires_paired_cfg = True + transformer = _get_mammothmoda2_transformer(pipeline) + block_adapter = BlockAdapter( + transformer=transformer, + blocks=[transformer.layers], + forward_pattern=[ForwardPattern.Pattern_3], + has_separate_cfg=True, + check_forward_pattern=True, + ) + refresh = enable_cache_for_dit( + pipeline, + cache_config, + block_adapter, + get_pipeline_transformer=_get_mammothmoda2_transformer, + ) + return CacheDiTEnableResult(refresh=refresh, targets=(block_adapter,)) + + def register_custom_dit_enablers() -> None: """Register model-specific Cache-DiT enablers. @@ -885,6 +926,7 @@ def register_custom_dit_enablers() -> None: "Cosmos3OmniPipeline": enable_cache_for_cosmos3, "Krea2Pipeline": enable_cache_for_krea2, "Magi2Pipeline": enable_cache_for_magi2, + "MammothModa2DiTPipeline": enable_cache_for_mammothmoda2, } ) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index e2c086d4a..51ea8f431 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py @@ -695,8 +695,16 @@ class Transformer2DModel(ModelMixin, ConfigMixin): return text_hidden_states, img_tokens def _apply_transformer_layers(self, hidden_states, attention_mask, rotary_emb, temb): + # Call blocks with keyword arguments: Cache-DiT's Pattern_3 wrapper + # re-invokes blocks as block(hidden_states, **kwargs), so the cached + # and uncached paths share this exact calling convention. for layer in self.layers: - hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb) + hidden_states = layer( + hidden_states, + attention_mask=attention_mask, + image_rotary_emb=rotary_emb, + temb=temb, + ) return hidden_states def forward( diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index 21977f488..2cc2cdcab 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -15,7 +15,12 @@ from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.cache.cachedit import ( + CacheDiTBackend, + CacheDiTRequestSpec, + RequestScopedCacheDiTRuntime, +) +from vllm_omni.diffusion.data import DiffusionCacheConfig, DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery @@ -28,6 +33,9 @@ from .schedulers import FlowMatchEulerDiscreteScheduler logger = init_logger(__name__) +# Identifies the pipeline-owned Cache-DiT installation across requests. +_MAMMOTHMODA2_CACHE_DIT_KEY = "mammothmoda2:cache_dit" + def _first_request_value(value: object) -> object: if isinstance(value, (list, tuple)): @@ -169,6 +177,54 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): self._llm_hidden_size = llm_hidden_size + # Cache-DiT lifecycle: the diffusion runner enables the configured + # backend (``cache_backend`` on the deploy YAML stage entry) at startup + # and transfers ownership here via the request-scoped protocol; + # forward() then reconciles per-request state (step count, CFG parity). + # The runner also emits the cache summary when + # ``enable_cache_dit_summary`` is set. + self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self) + self._cache_dit_config: DiffusionCacheConfig | None = None + if str(getattr(od_config, "cache_backend", "") or "").lower() == "cache_dit": + cache_config = od_config.cache_config + self._cache_dit_config = ( + cache_config + if isinstance(cache_config, DiffusionCacheConfig) + else DiffusionCacheConfig.from_dict(cache_config or {}) + ) + + def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None: + """Adopt a runner-installed Cache-DiT backend (request-scoped protocol).""" + + self._cache_dit_runtime.adopt(backend, installation_key=_MAMMOTHMODA2_CACHE_DIT_KEY) + + def is_cache_dit_enabled(self) -> bool: + """Return the request-scoped Cache-DiT installation state.""" + + return self._cache_dit_runtime.is_enabled + + def _prepare_cache_dit_for_request(self, *, num_inference_steps: int, cfg_active: bool) -> None: + """Reconcile Cache-DiT hooks with this request before the denoise loop. + + cache-dit's separate-CFG accounting assumes exactly two transformer + forwards per denoise step (conditional then unconditional), so CFG + requests refresh the context (per-request step counts vary) while + no-CFG requests run with hooks disabled: their single forward per + step cannot be represented by the parity accounting. + """ + if self._cache_dit_config is None: + return + if not cfg_active: + self._cache_dit_runtime.prepare(None) + return + self._cache_dit_runtime.prepare( + CacheDiTRequestSpec( + installation_key=_MAMMOTHMODA2_CACHE_DIT_KEY, + cache_config=self._cache_dit_config, + num_inference_steps=num_inference_steps, + ) + ) + def _reinit_caption_embedder(self, in_features: int) -> None: # Align with upstream Mammothmoda2Model's `reinit_caption_embedder`: # Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features). @@ -432,6 +488,18 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): # Run diffusion loop (CFG supported when text_guidance_scale > 1.0) total_steps = max(1, len(scheduler.timesteps)) + self._prepare_cache_dit_for_request( + num_inference_steps=total_steps, + cfg_active=request.text_guidance_scale > 1.0, + ) + # cache-dit distinguishes the conditional vs unconditional passes + # purely by transformer-forward parity, so when it is active the + # uncond pass must run on every step. Outside cfg_range CFG is + # neutralized via scale=1.0 (uncond + 1.0 * (cond - uncond) == cond), + # matching the skip optimization's cond-only update. + requires_paired_cfg = self._cache_dit_runtime.is_enabled and getattr( + self, "_cache_dit_requires_paired_cfg", False + ) for i, t in enumerate(scheduler.timesteps): timestep = t.expand(latents.shape[0]).to(latents.dtype) model_pred = self.gen_transformer( @@ -447,7 +515,10 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): guidance_scale = ( request.text_guidance_scale if request.cfg_range[0] <= i / total_steps <= request.cfg_range[1] else 1.0 ) - if guidance_scale > 1.0 and negative_prompt_embeds is not None: + run_uncond = guidance_scale > 1.0 and negative_prompt_embeds is not None + if requires_paired_cfg and negative_prompt_embeds is not None: + run_uncond = True + if run_uncond: model_pred_uncond = self.gen_transformer( hidden_states=latents, timestep=timestep, From 8bc0c38b902244842e377f1938577e9682f401d2 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Tue, 22 Sep 2026 09:34:22 +0800 Subject: [PATCH 2/7] [Bugfix] Add opt-in MiniCPM-o 4.5 Stage-0 sliding window (#7631) Signed-off-by: 0z5a Signed-off-by: 0z5a Co-authored-by: 0z5a --- .../common/ci_source_file_dependencies.yml | 2 + .buildkite/cuda/test-merge.yml | 2 +- .../sched/test_omni_ar_scheduler_streaming.py | 252 ++++++++++++++++++ .../helpers/minicpmo_window_e2e.py | 89 +++++++ .../test_minicpmo_4_5_window.py | 88 ++++++ .../duplex/test_minicpmo_window_plugin.py | 117 ++++++++ .../test_minicpmo_4_5_omni.py | 31 +++ tests/worker/test_native_duplex_hooks.py | 170 ++++++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 209 ++++++++++++++- .../models/minicpmo_4_5/duplex/plugin.py | 54 +++- .../models/minicpmo_4_5/duplex/policy.py | 51 ++++ .../models/minicpmo_4_5/duplex/stage0.py | 146 +++++++++- .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 10 + .../minicpmo_4_5_omni.py | 21 +- 14 files changed, 1230 insertions(+), 12 deletions(-) create mode 100644 tests/e2e/online_serving/helpers/minicpmo_window_e2e.py create mode 100644 tests/e2e/online_serving/test_minicpmo_4_5_window.py create mode 100644 tests/engine/duplex/test_minicpmo_window_plugin.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 99bbffa75..870221074 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -255,6 +255,8 @@ source_file_dependencies: - *minicpmo_4_5_duplex - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py + - tests/e2e/online_serving/test_minicpmo_4_5_window.py + - tests/e2e/online_serving/helpers/minicpmo_window_e2e.py - tests/e2e/online_serving/test_duplex_client_live.py - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 86e7a7a35..e76487cce 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -130,7 +130,7 @@ steps: source_file_dependencies: omni_minicpmo_4_5_duplex_function timeout_in_minutes: 50 commands: - - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model' + - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py tests/e2e/online_serving/test_minicpmo_4_5_window.py -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_1 - label: "Diffusion · Qwen Image Test" diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 83c7038f2..ab373f5de 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -69,6 +69,39 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate: ) +def _make_minicpm_window_update( + *, + seq: int, + mode: str, + high: int = 15, + low: int = 10, + context_max_units: int = 1, +) -> StreamingUpdate: + update = _make_update([0] * 8) + update.model_intermediate_buffer = { + "duplex": { + "data_plane": True, + "seq": seq, + "runtime_config": { + "duplex_scheduler_token_id": 0, + "duplex_first_append_context_tokens": 3, + "duplex_window_prefix_tokens": 2, + "duplex_window_suffix_token_ids": [3], + "duplex_window_previous_marker_token_ids": [70, 71], + "duplex_window_special_token_ids": [99], + "duplex_window_config": { + "sliding_window_mode": mode, + "basic_window_high_tokens": high, + "basic_window_low_tokens": low, + "context_previous_max_tokens": 4, + "context_max_units": context_max_units, + }, + }, + } + } + return update + + def _make_talker_adapter( *, max_model_len: int = 100, @@ -510,6 +543,225 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder() assert sched._new_prompt_len_snapshot[session.request_id] == 2 +def test_stage0_basic_window_rebuilds_below_low_watermark() -> None: + sched = _make_scheduler(stage_id=0) + session = _make_request() + session.prompt_token_ids = [0] * 9 + session._all_token_ids.clear() + session._all_token_ids.extend(session.prompt_token_ids) + session.num_prompt_tokens = 9 + session.append_output_token_ids([40]) + session.num_computed_tokens = 10 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + update = _make_minicpm_window_update(seq=2, mode="basic") + + sched._update_request_as_session(session, update) + + assert session.prompt_token_ids == [0] * 9 + assert session.num_computed_tokens == 0 + assert update.model_intermediate_buffer["meta"]["replace_streaming_prompt"] is True + plan = update.model_intermediate_buffer["duplex"]["stage0_window"] + assert plan == { + "completed_token_ids": [40], + "replace": True, + "mode": "basic", + "drop_units": 1, + "dropped_tokens": 9, + "previous_token_ids": [], + "previous_marker_token_ids": [70, 71], + "replacement_prompt_len": 9, + } + sched._free_request_blocks.assert_called_once_with(session) + + +def test_stage0_context_window_compacts_dropped_speech() -> None: + sched = _make_scheduler(stage_id=0) + session = _make_request() + session.prompt_token_ids = [0] * 9 + session._all_token_ids.clear() + session._all_token_ids.extend(session.prompt_token_ids) + session.num_prompt_tokens = 9 + session.append_output_token_ids([40]) + session.num_computed_tokens = 10 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + + first = _make_minicpm_window_update(seq=2, mode="context") + sched._update_request_as_session(session, first) + assert session.num_prompt_tokens == 18 + assert first.model_intermediate_buffer["duplex"]["stage0_window"] == {"completed_token_ids": [40]} + + session.append_output_token_ids([50]) + session.num_computed_tokens = 19 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + second = _make_minicpm_window_update(seq=3, mode="context") + sched._update_request_as_session(session, second) + + assert session.prompt_token_ids == [0] * 21 + plan = second.model_intermediate_buffer["duplex"]["stage0_window"] + assert plan["completed_token_ids"] == [50] + assert plan["drop_units"] == 1 + assert plan["previous_token_ids"] == [40] + # The worker embeds this exact marker, not a re-tokenized one, so the + # previous region is len(marker) + len(previous). + assert plan["previous_marker_token_ids"] == [70, 71] + assert plan["replacement_prompt_len"] == 21 + assert getattr(session, "_minicpmo45_window_previous_len") == 3 + + +@pytest.mark.parametrize("cleared_outputs", [False, True]) +@pytest.mark.parametrize("in_flight", [0, 1]) +def test_stage0_window_uses_confirmed_span_and_terminator(cleared_outputs, in_flight) -> None: + sched = _make_scheduler(stage_id=0) + session = _make_request() + session.prompt_token_ids = [0] * 9 + session._all_token_ids[:] = session.prompt_token_ids + session.num_prompt_tokens = 9 + session.append_output_token_ids([40, 99]) + session.num_computed_tokens = 10 + in_flight + session.num_output_placeholders = in_flight + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + if cleared_outputs: + # A stop without a queued append clears this before the next update. + session._output_token_ids.clear() + update = _make_minicpm_window_update(seq=2, mode="basic") + + sched._update_request_as_session(session, update) + + plan = update.model_intermediate_buffer["duplex"]["stage0_window"] + assert plan["completed_token_ids"] == [40] + assert plan["completed_terminator_token_id"] == 99 + assert plan["replacement_prompt_len"] == 9 + assert plan["dropped_tokens"] == 9 + + +def test_stage0_window_rebuild_that_overflows_max_model_len_finishes_the_session() -> None: + """A rebuilt window prompt is bounded by the client's window settings, not + by the model: a unit carrying camera frames is hundreds of tokens, and a + ``context`` window adds the ``previous`` region on top. The replacement + branch must therefore check ``replacement_prompt_len`` against + ``max_model_len - sample_room`` and finish the session with + ``context_length_exceeded``, exactly like a plain extension.""" + sched = _make_live_session_scheduler(max_model_len=21) + session = _make_request() + session.prompt_token_ids = [0] * 9 + session._all_token_ids[:] = [0] * 9 + session.num_prompt_tokens = 9 + session._output_token_ids[:] = [40] + session._all_token_ids.append(40) + session.num_computed_tokens = 10 + session.num_output_placeholders = 0 + _park_session(sched, session) + # The parked append state above is the frame before this update: one + # confirmed output token, fully computed. + session._output_token_ids[:] = [40] + session._all_token_ids[:] = [*session.prompt_token_ids, 40] + session.num_computed_tokens = 10 + session.num_in_flight_tokens = 0 + sched.num_sampled_tokens_per_step = 1 + # The replacement plan rebuilds 21 tokens (retained context plus the + # append) and one step samples one more, so 21 leaves no room. + update = _make_minicpm_window_update(seq=2, mode="basic") + update.prompt_token_ids = [0] * 20 + original_prompt = list(session.prompt_token_ids) + + sched._update_request_as_session(session, update) + + assert update.model_intermediate_buffer["duplex"]["stage0_window"]["replacement_prompt_len"] == 21 + assert session.prompt_token_ids == original_prompt + assert session.num_prompt_tokens == 9 + assert session.status == RequestStatus.FINISHED_ERROR + assert session.request_id not in sched.requests + assert sched.finished_req_ids == {session.request_id} + sched._free_request_blocks.assert_called_once_with(session) + client_index, reason = sched._streaming_context_overflow[session.request_id] + assert client_index == session.client_index + assert reason.startswith("context_length_exceeded: ") + assert "21 tokens" in reason and "max_model_len 21" in reason + + engine_core_outputs = _run_idle_step(sched) + + (output,) = engine_core_outputs[session.client_index].outputs + assert output.request_id == session.request_id + assert output.finish_reason == FinishReason.ERROR + assert output.stop_reason == reason + + +def test_stage0_window_rebuild_that_leaves_room_to_sample_replaces_the_prompt() -> None: + """One slot above the plan fits: the replacement applies through the normal + replacement path.""" + sched = _make_scheduler(stage_id=0) + sched.max_model_len = 25 + sched.num_sampled_tokens_per_step = 1 + session = _make_request() + session.prompt_token_ids = [0] * 9 + session._all_token_ids[:] = session.prompt_token_ids + session.num_prompt_tokens = 9 + session.append_output_token_ids([40]) + session.num_computed_tokens = 10 + session.num_output_placeholders = 0 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + update = _make_minicpm_window_update(seq=2, mode="basic") + update.prompt_token_ids = [0] * 20 + + sched._update_request_as_session(session, update) + + assert update.model_intermediate_buffer["meta"]["replace_streaming_prompt"] is True + assert update.model_intermediate_buffer["duplex"]["stage0_window"]["replacement_prompt_len"] == 21 + assert session.prompt_token_ids == [0] * 21 + assert session.num_prompt_tokens == 21 + assert session.num_computed_tokens == 0 + assert session.status == RequestStatus.WAITING + sched._free_request_blocks.assert_called_once_with(session) + assert not getattr(sched, "_streaming_context_overflow", {}) + + +def test_stage0_window_open_start_zero_is_not_replaced_by_the_context_reserve() -> None: + """A recorded ``open_start`` of 0 is a legitimate empty context prefix. The + fallback must be an explicit ``is None`` check: treating 0 as missing + substitutes the context reserve, counts the suffix twice, and the worker's + rebuild-length check raises on the first replacement.""" + session = SimpleNamespace( + num_prompt_tokens=9, + _all_token_ids=[0] * 9 + [40], + num_computed_tokens=10, + num_output_placeholders=0, + _minicpmo45_window_open_start=0, + ) + update = _make_minicpm_window_update(seq=2, mode="basic") + update.prompt_token_ids = [0] * 21 + + assert OmniARScheduler._prepare_minicpmo45_stage0_window( + session, update, segment_output_ids=[40], completed_terminator=3 + ) + + # open_start stays 0, so the reported unit spans the whole 12-token + # boundary instead of the 9 the context reserve would fold in. + plan = update.model_intermediate_buffer["duplex"]["stage0_window"] + assert plan["dropped_tokens"] == 12 + assert plan["replacement_prompt_len"] == 19 + + +def test_stage0_window_open_start_falls_back_to_the_context_reserve_when_unset() -> None: + session = SimpleNamespace( + num_prompt_tokens=9, + _all_token_ids=[0] * 9 + [40], + num_computed_tokens=10, + num_output_placeholders=0, + ) + update = _make_minicpm_window_update(seq=2, mode="basic") + update.prompt_token_ids = [0] * 21 + + assert OmniARScheduler._prepare_minicpmo45_stage0_window( + session, update, segment_output_ids=[40], completed_terminator=3 + ) + + # fallback open_start is preserve_len 3 from duplex_first_append_context_tokens, + # so the unit is three tokens shorter. + plan = update.model_intermediate_buffer["duplex"]["stage0_window"] + assert plan["dropped_tokens"] == 9 + assert plan["replacement_prompt_len"] == 22 + + def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None: sched = _make_scheduler(stage_id=1) sched.chunk_transfer_adapter = SimpleNamespace( diff --git a/tests/e2e/online_serving/helpers/minicpmo_window_e2e.py b/tests/e2e/online_serving/helpers/minicpmo_window_e2e.py new file mode 100644 index 000000000..553f07992 --- /dev/null +++ b/tests/e2e/online_serving/helpers/minicpmo_window_e2e.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Drive complete native duplex turns while Stage-0 repeatedly rebuilds KV.""" + +from __future__ import annotations + +import asyncio +import base64 +from pathlib import Path + +from vllm_omni.clients.duplex import DuplexClient, EventCollector, read_pcm16_wav, wait_for_condition +from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config + + +async def run_window_turn( + *, + url: str, + model: str, + input_wav: Path, + mode: str, + ref_audio: Path | None = None, + repeats: int = 2, + video_frames: list[str] | None = None, + buffered_flush: bool = False, + timeout_s: float = 240.0, +) -> dict[str, object]: + """Use aggressive limits to exercise multiple windows in one audio turn.""" + pcm = read_pcm16_wav(input_wav) * repeats + collector = EventCollector() + reference = ( + "data:audio/wav;base64," + base64.b64encode(ref_audio.read_bytes()).decode("ascii") + if ref_audio is not None + else None + ) + config = create_duplex_session_config( + modalities=("text", "audio") if ref_audio is not None else ("text",), + ref_audio=reference, + temperature=0.0, + extra_body={ + "sliding_window_mode": mode, + "context_max_units": 1, + "context_previous_max_tokens": 16, + "basic_window_high_tokens": 64, + "basic_window_low_tokens": 32, + }, + ) + client = DuplexClient( + url, model=model, config=config, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=timeout_s + ) + async with client: + reader = asyncio.create_task(collector.consume(client)) + collector.add({"type": "session.created", "session": client.session_info}) + try: + frames_sent = await client.stream_pcm( + pcm, + chunk_ms=2500 if buffered_flush else 200, + realtime=not buffered_flush, + is_speech=True, + video_frames=video_frames, + ) + await client.commit(final=True) + await wait_for_condition( + lambda: collector.count("input_audio_buffer.committed") > 0 or bool(collector.errors()), + timeout_s=timeout_s, + label="final input commit", + ) + await wait_for_condition( + lambda: collector.count("response.done") > 0 or bool(collector.errors()), + timeout_s=timeout_s, + label="completed windowed response", + ) + # Final execution follows the commit acknowledgement asynchronously. + await asyncio.sleep(3.0) + await client.close(timeout_s=timeout_s) + await asyncio.wait_for(reader, timeout=5.0) + finally: + if not reader.done(): + reader.cancel() + await asyncio.gather(reader, return_exceptions=True) + return { + "created": collector.count("session.created"), + "committed": collector.count("input_audio_buffer.committed"), + "done": collector.count("response.done"), + "closed": collector.count("session.closed"), + "errors": collector.errors(), + "audio_bytes": len(collector.audio_bytes()), + "frames_sent": frames_sent, + "input_seconds": len(pcm) / 32000, + } diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_window.py b/tests/e2e/online_serving/test_minicpmo_4_5_window.py new file mode 100644 index 000000000..300205eb0 --- /dev/null +++ b/tests/e2e/online_serving/test_minicpmo_4_5_window.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real-model regression for the scheduler-owned Thinker sliding window.""" + +import asyncio + +import pytest + +from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import ( + SERVER_PARAMS, + duplex_camera_frames, + realtime_url, + resolve_ref_audio, + validated_input_wav, +) +from tests.e2e.online_serving.helpers.minicpmo_window_e2e import run_window_turn +from tests.helpers.mark import hardware_test + +pytestmark = [pytest.mark.omni, pytest.mark.advanced_model] + + +def _assert_complete(result, *, require_audio=False): + assert result["created"] == 1 + assert result["committed"] > 0 + assert result["done"] > 0 + assert result["closed"] == 1 + assert result["errors"] == [] + if require_audio: + assert result["audio_bytes"] > 0 + + +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +@pytest.mark.parametrize("mode,reference", [("basic", False), ("context", False), ("context", True)]) +def test_window_rebuild_and_next_session(omni_server, mode, reference): + # Two independent sessions exercise admission after final-unit execution. + for _ in range(2): + result = asyncio.run( + run_window_turn( + url=realtime_url(omni_server), + model=omni_server.model, + input_wav=validated_input_wav(), + mode=mode, + ref_audio=resolve_ref_audio() if reference else None, + ) + ) + _assert_complete(result, require_audio=reference) + + +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +@pytest.mark.parametrize("mode,camera", [("basic", False), ("context", False), ("context", True)]) +def test_window_continuous_input(omni_server, tmp_path, mode, camera): + # 24 repeats cover over two minutes and many canonical KV replacements. + frames = duplex_camera_frames(seconds=4, cache_dir=tmp_path / "camera") if camera else None + result = asyncio.run( + run_window_turn( + url=realtime_url(omni_server), + model=omni_server.model, + input_wav=validated_input_wav(), + mode=mode, + repeats=24, + ref_audio=resolve_ref_audio(), + video_frames=frames, + ) + ) + _assert_complete(result, require_audio=True) + assert isinstance(result["input_seconds"], (int, float)) and result["input_seconds"] > 120 + if camera: + assert isinstance(result["frames_sent"], int) and result["frames_sent"] >= 120 + + +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +@pytest.mark.parametrize("mode", ["off", "basic", "context"]) +def test_window_buffered_flush(omni_server, mode): + # Burst appends cross processor chunk boundaries and leave a final tail. + result = asyncio.run( + run_window_turn( + url=realtime_url(omni_server), + model=omni_server.model, + input_wav=validated_input_wav(), + mode=mode, + ref_audio=resolve_ref_audio(), + buffered_flush=True, + ) + ) + _assert_complete(result, require_audio=True) diff --git a/tests/engine/duplex/test_minicpmo_window_plugin.py b/tests/engine/duplex/test_minicpmo_window_plugin.py new file mode 100644 index 000000000..d2afce199 --- /dev/null +++ b/tests/engine/duplex/test_minicpmo_window_plugin.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import base64 +from types import SimpleNamespace + +import pytest + +from vllm_omni.engine.duplex.config import DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import DuplexFence +from vllm_omni.model_executor.models.minicpmo_4_5.duplex import plugin as module + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _plugin(): + return module.MiniCPMO45DuplexPlugin(lambda *args: None) + + +@pytest.mark.asyncio +async def test_window_configuration_is_prepared_by_plugin(monkeypatch): + tokenizer = SimpleNamespace( + all_special_ids=[90], + unk_token_id=-1, + eos_token_id=99, + encode=lambda text, add_special_tokens=False: [ord(char) for char in text], + convert_tokens_to_ids=lambda token: {"<|listen|>": 91}.get(token, -1), + ) + monkeypatch.setattr(module, "_load_tokenizer", lambda config: tokenizer) + config = DuplexSessionConfig( + modalities=("text",), + extra_body={ + "sliding_window_mode": "basic", + "basic_window_high_tokens": 120, + "basic_window_low_tokens": 80, + }, + ) + runtime = await _plugin().prepare_runtime_config(config, model_config=None) + assert runtime["duplex_window_config"]["sliding_window_mode"] == "basic" + assert "sliding_window_mode" not in config.extra_body + assert runtime["duplex_window_prefix_tokens"] > 0 + assert runtime["duplex_window_suffix_token_ids"] + assert runtime["duplex_window_previous_marker_token_ids"] + + +@pytest.mark.asyncio +async def test_invalid_window_configuration_is_rejected(): + config = DuplexSessionConfig( + modalities=("text",), + extra_body={ + "basic_window_high_tokens": 80, + "basic_window_low_tokens": 80, + }, + ) + with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError) as exc: + await _plugin().prepare_runtime_config(config, model_config=None) + assert exc.value.code == "invalid_sliding_window_config" + + +def test_window_update_cannot_change_created_configuration(): + config = DuplexSessionConfig(extra_body={"sliding_window_mode": "basic"}) + with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError) as exc: + _plugin().runtime_config_for_update(config, {"duplex_window_config": {"sliding_window_mode": "off"}}) + assert exc.value.code == "sliding_window_update_unsupported" + + +def test_window_internal_configuration_is_server_owned(): + with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError): + _plugin().validate_client_extra_body({"duplex_window_config": {}}) + + +@pytest.mark.parametrize("final", [False, True]) +@pytest.mark.parametrize("seq,samples,expected", [(1, 32000, 16), (1, 48000, 28), (2, 16000, 13)]) +def test_first_and_final_append_reserve_exact_window_input(seq, samples, expected, final): + prompt = module.build_duplex_data_plane_prompt( + request_id="window-request", + fence=DuplexFence("sid", turn_id=1), + session_config={}, + runtime_config={"duplex_first_append_context_tokens": 5}, + seq=seq, + turn_seq=seq, + payload={ + "audio": base64.b64encode(bytes(samples * 4)).decode(), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + }, + final=final, + ) + assert len(prompt["prompt_token_ids"]) == expected + + +@pytest.mark.parametrize("seq", [1, 2, 3]) +def test_final_exact_chunk_append_reserves_one_unit_not_two(seq): + """A final append whose audio is already an exact number of chunks must + reserve exactly what Stage0 feeds: one unit per chunk (```` plus its + audio) and the closure pair for every unit after the first. Serving pads + the final residual itself, so the old extra 12-slot silent-unit + reservation only put 12 pad positions ahead of the audio in the KV.""" + + def _budget(*, seq, final): + prompt = module.build_duplex_data_plane_prompt( + request_id="final-request", + fence=DuplexFence("sid", turn_id=1), + session_config={}, + runtime_config={"duplex_first_append_context_tokens": 0}, + seq=seq, + turn_seq=seq, + payload={ + "audio": base64.b64encode(bytes(16000 * 4)).decode(), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + }, + final=final, + ) + return len(prompt["prompt_token_ids"]) + + assert _budget(seq=seq, final=True) == _budget(seq=seq, final=False) 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 27037f029..c5d4d0cde 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 @@ -142,6 +142,37 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None: assert info["duplex"]["turn_id"] == 7 +@pytest.mark.parametrize("folded_decisions", [0, 1, 2]) +@pytest.mark.parametrize("bos_in_prompt", [False, True]) +def test_native_duplex_tts_bos_aligns_after_window_rebuild(folded_decisions, bos_in_prompt) -> None: + prompt_ids = [101] * 10 + output_ids = [9303] * folded_decisions + [9301, 21, 22, 9308] + if bos_in_prompt: + prompt_ids[-1] = 9301 + output_ids = [21, 22, 9308] + rows = len(prompt_ids) + len(output_ids) - folded_decisions + latent = torch.arange(rows * 4, dtype=torch.float32).reshape(rows, 4) + source = _output( + prompt_ids=prompt_ids, + output_ids=output_ids, + latent=latent, + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "meta": { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "chunk_eos_token_id": 9308, + }, + }, + ) + 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]) + + def test_native_duplex_continuation_appends_only_new_talker_condition() -> None: prompt_ids = [101, 102] token_ids = { diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index 1b1bf4b7e..d3da3bba0 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -848,6 +848,176 @@ def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format(): assert result["prompt_suffix_len"] == 0 +@pytest.mark.parametrize("mode", [None, "off", "basic", "context"]) +def test_minicpmo_stage0_history_collection_starts_with_session_config(mode): + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState + + runtime = _stage0_vision_runtime() + state = _MiniCPMO45Stage0SessionState(session_id="history-lifecycle") + config = {} if mode is None else {"duplex_window_config": {"sliding_window_mode": mode}} + runtime._prepare_session_context(state, {}, runtime_config=config) + for seq in range(1, 101): + runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=seq) + enabled = mode in {"basic", "context"} + assert state.window_enabled is enabled + assert len(state.window_units) == (99 if enabled else 0) + assert (state.pending_window_unit is not None) is enabled + assert state.pending_window_generated_tokens == [] + + +@pytest.mark.parametrize("mode", ["basic", "context"]) +@pytest.mark.parametrize("buffered", [False, True]) +def test_minicpmo_stage0_multi_chunk_append_matches_scheduler_window(mode, buffered): + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState + + runtime = _stage0_vision_runtime() + state = _MiniCPMO45Stage0SessionState(session_id="multi-chunk-window", window_enabled=True) + if buffered: + # Keep the tiny fake processor's initial chunk unpadded to model carry. + runtime._pad_first_audio_chunk_if_needed = lambda *args: None + assert runtime._stage_prefill_embeddings_only(state, np.zeros(2, dtype=np.float32), seq=0)["success"] is False + first = runtime._stage_prefill_embeddings_only( + state, np.zeros(6 if buffered else 8, dtype=np.float32), seq=1, final=buffered + ) + assert first["input_token_ids"] == [1, 11, 2, 1, 11] + assert state.window_units == [] + assert state.pending_window_unit.token_ids == first["input_token_ids"] + session = SimpleNamespace(num_prompt_tokens=5) + replaced = False + for seq in (2, 3): + info = { + "duplex": { + "data_plane": True, + "seq": seq, + "runtime_config": { + "duplex_window_config": { + "sliding_window_mode": mode, + "basic_window_high_tokens": 1, + "basic_window_low_tokens": 2, + "context_max_units": 1, + "context_previous_max_tokens": 16, + }, + "duplex_window_previous_marker_token_ids": [201, 5], + }, + } + } + update = SimpleNamespace(prompt_token_ids=[3, 2, 1, 11], model_intermediate_buffer=info) + planned = OmniARScheduler._prepare_minicpmo45_stage0_window( + session, update, segment_output_ids=[44], completed_terminator=3 + ) + result = runtime._stage_prefill_embeddings_only( + state, np.zeros(4, dtype=np.float32), seq=seq, stage0_window=info["duplex"]["stage0_window"] + ) + assert result["stage0_window_replaced"] is planned + if planned: + assert result["num_input_tokens"] == len(update.prompt_token_ids) + replaced = True + break + session.num_prompt_tokens += 1 + len(update.prompt_token_ids) + assert replaced + + +def test_minicpmo_stage0_basic_window_rebuilds_from_retained_units(): + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + runtime = _stage0_vision_runtime() + state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-basic-window", window_enabled=True) + first = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1) + assert first["input_token_ids"] == [1, 11] + state.pending_window_generated_tokens.append(5) + state.pending_terminator_token = 3 + + rebuilt = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + seq=2, + stage0_window={ + "replace": True, + "mode": "basic", + "drop_units": 1, + "replacement_prompt_len": 2, + }, + ) + + assert rebuilt["stage0_window_replaced"] is True + assert rebuilt["input_token_ids"] == [1, 11] + assert state.window_units == [] + + +def test_minicpmo_stage0_context_window_inserts_previous_before_suffix(): + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + runtime = _stage0_vision_runtime() + state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-context-window", window_enabled=True) + prefix = runtime._embed_token(200) + suffix = runtime._embed_token(202) + state.context_embeds = [prefix, suffix] + state.context_token_ids = [200, 202] + state.context_prefix_embeds = [prefix] + state.context_prefix_token_ids = [200] + state.context_suffix_embeds = [suffix] + state.context_suffix_token_ids = [202] + first = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1) + assert first["input_token_ids"] == [200, 202, 1, 11] + state.pending_window_generated_tokens.append(42) + state.pending_terminator_token = 3 + + rebuilt = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + seq=2, + stage0_window={ + "replace": True, + "mode": "context", + "drop_units": 1, + "previous_token_ids": [42], + "previous_marker_token_ids": [201, 5], + "replacement_prompt_len": 7, + }, + ) + + assert rebuilt["input_token_ids"] == [200, 201, 5, 42, 202, 1, 11] + assert rebuilt["num_input_tokens"] == 7 + + +@pytest.mark.parametrize("pending_terminator", [None, 3, 99]) +def test_minicpmo_stage0_window_uses_accepted_output_not_async_sampler_history(pending_terminator): + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + runtime = _stage0_vision_runtime() + runtime._stage_audio_embeddings = lambda *args, **kwargs: torch.zeros((10, 2)) + state = _MiniCPMO45Stage0SessionState(session_id="sid-accepted-window", window_enabled=True) + state.context_prefix_token_ids = [200] * 7 + state.context_suffix_token_ids = [202] + state.context_token_ids = [*state.context_prefix_token_ids, 202] + state.context_prefix_embeds = [runtime._embed_token(token_id) for token_id in state.context_prefix_token_ids] + state.context_suffix_embeds = [runtime._embed_token(202)] + state.context_embeds = [*state.context_prefix_embeds, *state.context_suffix_embeds] + runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1) + for seq in (2, 3): + state.pending_window_generated_tokens = [3, 40, 41] + state.pending_terminator_token = pending_terminator + plan = {"completed_token_ids": [], "completed_terminator_token_id": 3} + if seq == 3: + plan.update(replace=True, mode="context", drop_units=1, replacement_prompt_len=32) + result = runtime._stage_prefill_embeddings_only( + state, np.zeros(4, dtype=np.float32), seq=seq, stage0_window=plan + ) + + assert result["num_input_tokens"] == 32 + assert [len(unit.token_ids) for unit in state.window_units] == [13] + assert len(state.pending_window_unit.token_ids) == 11 + assert 40 not in result["input_token_ids"] + assert 41 not in result["input_token_ids"] + + def _stage0_vision_runtime(): import torch diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 4d55c55f8..f07aae021 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -848,6 +848,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): req_id = session.request_id self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids) outstanding_async_tokens = getattr(session, "num_output_placeholders", 0) + # Use the same confirmed span that upstream preserves when extending + # a session. The segment output list is cleared on a resumable stop, + # and, for an already queued append, can still include its terminator. + confirmed_end = session.num_computed_tokens - outstanding_async_tokens + segment_output_ids = list(session._all_token_ids[session.num_prompt_tokens : confirmed_end]) + completed_terminator = ( + session._all_token_ids[confirmed_end] + if session.num_prompt_tokens <= confirmed_end < len(session._all_token_ids) + else None + ) # Seed the stale share in SCHEDULED-token units (see the segment-stop # site in update_from_output): num_in_flight_tokens matches what each # pre-replacement frame will drain, so the counter reaches exactly @@ -894,6 +904,23 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) return + if stage_id == 0 and self._prepare_minicpmo45_stage0_window( + session, + update, + segment_output_ids=segment_output_ids, + completed_terminator=completed_terminator, + ): + # The rebuilt prompt is bounded by the client's window settings, + # not by the model (a camera unit is hundreds of tokens), so it + # needs the same max_model_len check as a plain extension. The + # plan replaces the whole prompt, so the replacement length is the + # projection, not the session's current prompt plus an extension. + plan = update.model_intermediate_buffer["duplex"]["stage0_window"] + if self._streaming_update_overflows(session, update, projected_len=plan["replacement_prompt_len"]): + return + self._release_replaced_streaming_prompt_cache(session) + self._replace_streaming_session(session, update) + return streaming_prompt_payload = next( ( info @@ -962,11 +989,178 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if hasattr(update, "model_intermediate_buffer"): session.model_intermediate_buffer = update.model_intermediate_buffer + @staticmethod + def _prepare_minicpmo45_stage0_window( + session: Request, + update: StreamingUpdate, + *, + segment_output_ids: list[int], + completed_terminator: int | None = None, + ) -> bool: + """Plan an official-style MiniCPM Stage-0 window at unit boundaries. + + vLLM owns a paged KV cache, so deleting a middle span and rotating the + retained K tensors in place is not a safe model hook. Instead, record + completed unit lengths in the scheduler and request a full prompt + replacement when a watermark fires. The worker rebuilds matching + embeddings from its unit history, which recomputes RoPE at the new + contiguous positions. + """ + info = getattr(update, "model_intermediate_buffer", None) + if not isinstance(info, dict): + return False + duplex = info.get("duplex") + 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 {} + window = runtime_config.get("duplex_window_config") + if not isinstance(window, dict): + return False + mode = window.get("sliding_window_mode", "off") + if mode == "off": + return False + if mode not in {"basic", "context"}: + return False + + try: + seq = int(duplex.get("seq", 0) or 0) + preserve_len = int(runtime_config.get("duplex_first_append_context_tokens", 0) or 0) + except (TypeError, ValueError): + return False + if seq <= 1: + return False + + # The sampled terminator is discarded by the normal session update and + # re-injected at the head of this append, followed by . + # generated_ids contains only the confirmed output span retained by + # the streaming update, excluding its final uncomputed sample. + generated_ids = segment_output_ids + # Sampler-side history can contain speculative or stale async output. + # The worker must rebuild from the same accepted ids used for lengths. + duplex["stage0_window"] = {"completed_token_ids": list(generated_ids)} + if completed_terminator is not None: + duplex["stage0_window"]["completed_terminator_token_id"] = int(completed_terminator) + base_len = int(getattr(session, "num_prompt_tokens", 0) or 0) + len(generated_ids) + boundary = base_len + 2 + # A legitimate open_start of 0 (an empty context prefix) must not fall + # back to preserve_len, which would double-count the suffix below. + recorded_open_start = getattr(session, "_minicpmo45_window_open_start", None) + open_start = preserve_len if recorded_open_start is None else int(recorded_open_start) + unit_len = boundary - open_start + if unit_len <= 0: + return False + + special_ids = { + int(token_id) + for token_id in runtime_config.get("duplex_window_special_token_ids", ()) + if isinstance(token_id, int) + } + units = list(getattr(session, "_minicpmo45_window_units", ())) + units.append( + { + "length": unit_len, + "generated_token_ids": [int(token_id) for token_id in generated_ids if token_id not in special_ids], + } + ) + projected_len = base_len + len(update.prompt_token_ids) + drop_count = 0 + dropped_len = 0 + + if mode == "basic": + try: + high = int(window.get("basic_window_high_tokens", 8000)) + low = int(window.get("basic_window_low_tokens", 6000)) + except (TypeError, ValueError): + return False + if projected_len > high: + while units and projected_len - dropped_len > low: + dropped_len += int(units[drop_count]["length"]) + drop_count += 1 + if drop_count >= len(units): + break + else: + try: + max_units = int(window.get("context_max_units", 24)) + previous_max = int(window.get("context_previous_max_tokens", 500)) + except (TypeError, ValueError): + return False + drop_count = max(0, len(units) - max_units) + dropped_len = sum(int(unit["length"]) for unit in units[:drop_count]) + previous = list(getattr(session, "_minicpmo45_window_previous_token_ids", ())) + for unit in units[:drop_count]: + previous.extend(unit["generated_token_ids"]) + if len(previous) > previous_max: + previous = previous[-previous_max:] + marker = [ + int(token_id) + for token_id in runtime_config.get("duplex_window_previous_marker_token_ids", ()) + if isinstance(token_id, int) + ] + previous_with_marker = marker + previous if previous else [] + old_previous_len = int(getattr(session, "_minicpmo45_window_previous_len", 0) or 0) + session._minicpmo45_window_previous_token_ids = previous + session._minicpmo45_window_previous_len = len(previous_with_marker) + projected_len += len(previous_with_marker) - old_previous_len + + retained_units = units[drop_count:] + session._minicpmo45_window_units = retained_units + if drop_count == 0: + session._minicpmo45_window_open_start = boundary + return False + + replacement_len = projected_len - dropped_len + if replacement_len <= 0: + return False + scheduler_token_id = runtime_config.get("duplex_scheduler_token_id", 0) + try: + scheduler_token_id = max(0, int(scheduler_token_id)) + except (TypeError, ValueError): + scheduler_token_id = 0 + update.prompt_token_ids = [scheduler_token_id] * replacement_len + meta = info.setdefault("meta", {}) + if isinstance(meta, dict): + meta["replace_streaming_prompt"] = True + previous_ids = list(getattr(session, "_minicpmo45_window_previous_token_ids", ())) + # The worker embeds the marker ids the window plan carries: it must not + # re-tokenize the marker with its own tokenizer, which could differ + # from the ids this length is computed from. + marker_ids = [ + int(token_id) + for token_id in runtime_config.get("duplex_window_previous_marker_token_ids", ()) + if isinstance(token_id, int) + ] + duplex["stage0_window"] = { + "completed_token_ids": list(generated_ids), + "replace": True, + "mode": mode, + "drop_units": drop_count, + "dropped_tokens": dropped_len, + "previous_token_ids": previous_ids, + "previous_marker_token_ids": marker_ids, + "replacement_prompt_len": replacement_len, + } + if completed_terminator is not None: + duplex["stage0_window"]["completed_terminator_token_id"] = int(completed_terminator) + recorded_prefix_len = runtime_config.get("duplex_window_prefix_tokens") + prefix_len = preserve_len if recorded_prefix_len is None else int(recorded_prefix_len) + suffix_len = len(runtime_config.get("duplex_window_suffix_token_ids", ()) or ()) + previous_len = int(getattr(session, "_minicpmo45_window_previous_len", 0) or 0) + new_preserve_len = prefix_len + previous_len + suffix_len if mode == "context" else preserve_len + session._minicpmo45_window_open_start = new_preserve_len + sum(int(unit["length"]) for unit in retained_units) + return True + # Prefix of the stop_reason carried by the FinishReason.ERROR output, so the # serving side can map it to a stable error code. STREAMING_CONTEXT_OVERFLOW_STOP_REASON = "context_length_exceeded" - def _streaming_update_overflows(self, session: Request, update: StreamingUpdate) -> bool: + def _streaming_update_overflows( + self, + session: Request, + update: StreamingUpdate, + *, + projected_len: int | None = None, + ) -> bool: """Finish a streaming session whose next extension cannot fit the model. Upstream ``_update_request_as_session`` appends the update to the @@ -991,6 +1185,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): finish to the next ``schedule()`` would leave the client waiting. The reason is emitted with the terminal output (see :meth:`_emit_streaming_context_overflow_outputs`). + + ``projected_len`` overrides the extended-prompt projection for a caller + that replaces the prompt instead of growing it, such as the MiniCPM + Stage-0 window rebuild. """ max_model_len = getattr(self, "max_model_len", None) if max_model_len is None: @@ -1001,8 +1199,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): new_tokens = len(update.prompt_token_ids or ()) # The extended prompt is the current prompt plus the computed output # tokens upstream keeps, then the update: num_computed_tokens covers - # both when the prompt was fully computed. - projected = max(int(session.num_prompt_tokens), int(session.num_computed_tokens)) + new_tokens + # both when the prompt was fully computed. A caller that already knows + # the post-update prompt length (a window rebuild) passes it instead. + projected = ( + max(int(session.num_prompt_tokens), int(session.num_computed_tokens)) + new_tokens + if projected_len is None + else int(projected_len) + ) # Room for the tokens one step samples on top of the prompt (1 without # speculative decoding). __new__-built test schedulers carry no # num_sampled_tokens_per_step. 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 e57377766..44ae9dd9a 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 @@ -45,7 +45,10 @@ from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( MiniCPMO45DataPlaneContext, MiniCPMO45DataPlaneSession, ) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import ( + MiniCPMO45DuplexPolicy, + MiniCPMO45DuplexWindowConfig, +) from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import ( MiniCPMO45ServingSessionState, ) @@ -85,6 +88,11 @@ PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( "ref_audio_format", "ref_audio_sample_rate_hz", "initial_user_text", + "duplex_window_config", + "duplex_window_prefix_tokens", + "duplex_window_suffix_token_ids", + "duplex_window_previous_marker_token_ids", + "duplex_window_special_token_ids", } ) @@ -245,8 +253,8 @@ def build_duplex_data_plane_prompt( token_budget = context_reserve + first_units * 12 - 1 + vision_tokens if seq > 1 and duplex_payload_is_exact_chunks(payload): token_budget += 1 - if final and duplex_payload_is_exact_chunks(payload): - token_budget += 12 + # Serving already pads the final residual audio. Stage0 does not + # append another silent unit, so final must not reserve extra slots. extra_body = session_config.get("extra_body") raw_token_id = runtime_config.get("duplex_scheduler_token_id") try: @@ -515,6 +523,16 @@ def _apply_first_append_context_tokens( return ref_tokens = MiniCPMO45DuplexPolicy.audio_token_count(ref_sample_count or 0) runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids) + runtime_config["duplex_window_prefix_tokens"] = len(prefix_ids) + ref_tokens + runtime_config["duplex_window_suffix_token_ids"] = [int(token_id) for token_id in suffix_ids] + marker_ids = tokenizer.encode("\n\nprevious: ", add_special_tokens=False) + runtime_config["duplex_window_previous_marker_token_ids"] = [int(token_id) for token_id in marker_ids] + runtime_config["duplex_window_special_token_ids"] = sorted( + MiniCPMO45DuplexPolicy.native_special_token_ids( + MiniCPMO45DuplexPolicy.token_ids_from_tokenizer(tokenizer), + tokenizer_special_ids=list(getattr(tokenizer, "all_special_ids", ()) or ()), + ) + ) def _model_vision_tile_pixels(model_config: ModelConfig | None) -> int | None: @@ -722,6 +740,8 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin): code="unsupported_ref_audio_path", ) runtime_config: dict[str, object] = {"instructions": config.instructions} + window_config = self._pop_window_config(extra_body) or MiniCPMO45DuplexWindowConfig() + runtime_config["duplex_window_config"] = window_config.as_dict() # ``duplex_initial_user_text`` is the older extra_body spelling and # still works; the session field is the framework-level one. initial_user_text = extra_body.pop("duplex_initial_user_text", None) @@ -788,6 +808,17 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin): current: Mapping[str, object], ) -> dict[str, object]: runtime_config = deepcopy(dict(current)) + extra_body = dict(config.extra_body) + requested_window = self._pop_window_config(extra_body) + if requested_window is not None: + reject_changed_runtime_value( + requested_window.as_dict(), + runtime_config.get("duplex_window_config"), + message="sliding-window configuration cannot be changed after the session is created", + code="sliding_window_update_unsupported", + error_cls=MiniCPMO45ClientRuntimeConfigError, + ) + config.extra_body = extra_body reject_changed_runtime_value( config.instructions, runtime_config.get("instructions"), @@ -812,6 +843,23 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin): runtime_config["duplex_stage_sampling_params"] = stage_sampling return runtime_config + @staticmethod + def _pop_window_config(extra_body: dict[str, object]) -> MiniCPMO45DuplexWindowConfig | None: + names = ( + "sliding_window_mode", + "basic_window_high_tokens", + "basic_window_low_tokens", + "context_previous_max_tokens", + "context_max_units", + ) + provided = {name: extra_body.pop(name) for name in names if name in extra_body} + if not provided: + return None + try: + return MiniCPMO45DuplexWindowConfig.from_mapping(provided) + except ValueError as exc: + raise MiniCPMO45ClientRuntimeConfigError(str(exc), code="invalid_sliding_window_config") from exc + def data_plane_context( self, *, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py index 0332c81c0..7b0aded4d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py @@ -3,12 +3,63 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING if TYPE_CHECKING: from transformers import PreTrainedTokenizerBase +@dataclass(frozen=True, slots=True) +class MiniCPMO45DuplexWindowConfig: + """Stage-0 window settings matching the released checkpoint defaults.""" + + sliding_window_mode: str = "off" + basic_window_high_tokens: int = 8000 + basic_window_low_tokens: int = 6000 + context_previous_max_tokens: int = 500 + context_max_units: int = 24 + + @classmethod + def from_mapping(cls, value: object) -> MiniCPMO45DuplexWindowConfig: + source = value if isinstance(value, dict) else {} + + def integer(name: str, default: int) -> int: + raw = source.get(name, default) + if isinstance(raw, bool): + raise ValueError(f"{name} must be an integer") + try: + parsed = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be an integer") from exc + if parsed <= 0: + raise ValueError(f"{name} must be greater than zero") + return parsed + + mode = source.get("sliding_window_mode", "off") + if mode not in {"off", "basic", "context"}: + raise ValueError("sliding_window_mode must be one of: off, basic, context") + config = cls( + sliding_window_mode=str(mode), + basic_window_high_tokens=integer("basic_window_high_tokens", 8000), + basic_window_low_tokens=integer("basic_window_low_tokens", 6000), + context_previous_max_tokens=integer("context_previous_max_tokens", 500), + context_max_units=integer("context_max_units", 24), + ) + if config.basic_window_low_tokens >= config.basic_window_high_tokens: + raise ValueError("basic_window_low_tokens must be less than basic_window_high_tokens") + return config + + def as_dict(self) -> dict[str, int | str]: + return { + "sliding_window_mode": self.sliding_window_mode, + "basic_window_high_tokens": self.basic_window_high_tokens, + "basic_window_low_tokens": self.basic_window_low_tokens, + "context_previous_max_tokens": self.context_previous_max_tokens, + "context_max_units": self.context_max_units, + } + + class MiniCPMO45DuplexPolicy: """MiniCPM-o 4.5 native duplex model policy. 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 bc13b58f2..cbc4d0b35 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 @@ -27,6 +27,12 @@ _MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS _MINICPMO45_PROCESSOR_LOAD_LOCK = Lock() +@dataclass +class _MiniCPMO45WindowUnit: + embeds: list[Any] = field(default_factory=list) + token_ids: list[int] = field(default_factory=list) + + @dataclass class _MiniCPMO45Stage0SessionState: session_id: str @@ -36,6 +42,14 @@ class _MiniCPMO45Stage0SessionState: audio_chunk_idx: int = 0 context_embeds: list[torch.Tensor] = field(default_factory=list) context_token_ids: list[int] = field(default_factory=list) + context_prefix_embeds: list[Any] = field(default_factory=list) + context_prefix_token_ids: list[int] = field(default_factory=list) + context_suffix_embeds: list[Any] = field(default_factory=list) + context_suffix_token_ids: list[int] = field(default_factory=list) + window_units: list[_MiniCPMO45WindowUnit] = field(default_factory=list) + window_enabled: bool = False + pending_window_unit: _MiniCPMO45WindowUnit | None = None + pending_window_generated_tokens: list[int] = field(default_factory=list) current_turn_ended: bool = True prepared_append_identity: tuple[int | None, int] | None = None prepared_inputs_embeds: torch.Tensor | None = None @@ -142,6 +156,11 @@ class MiniCPMO45Stage0DuplexRuntime: *, runtime_config: dict[str, object] | None = None, ) -> None: + window_config = (runtime_config or {}).get("duplex_window_config") + state.window_enabled = isinstance(window_config, dict) and window_config.get("sliding_window_mode", "off") in { + "basic", + "context", + } if not self._stage_runtime_ready(): return self._require_special_token_ids() @@ -157,17 +176,26 @@ class MiniCPMO45Stage0DuplexRuntime: (runtime_config or {}).get("initial_user_text"), ) for token_id in self._encode_text(prefix): - state.context_embeds.append(self._embed_token(token_id)) + embed = self._embed_token(token_id) + state.context_embeds.append(embed) state.context_token_ids.append(token_id) + state.context_prefix_embeds.append(embed) + state.context_prefix_token_ids.append(token_id) if ref_audio is not None: ref_audio_embeds = self._stage_ref_audio_embeddings(ref_audio, state=state) if ref_audio_embeds is not None: ref_audio_embeds = self._as_2d_tensor(ref_audio_embeds) + placeholder_ids = [self.unit_token_id] * int(ref_audio_embeds.shape[0]) state.context_embeds.append(ref_audio_embeds) - state.context_token_ids.extend([self.unit_token_id] * int(ref_audio_embeds.shape[0])) + state.context_token_ids.extend(placeholder_ids) + state.context_prefix_embeds.append(ref_audio_embeds) + state.context_prefix_token_ids.extend(placeholder_ids) for token_id in self._encode_text(suffix): - state.context_embeds.append(self._embed_token(token_id)) + embed = self._embed_token(token_id) + state.context_embeds.append(embed) state.context_token_ids.append(token_id) + state.context_suffix_embeds.append(embed) + state.context_suffix_token_ids.append(token_id) def _stage_prefill_embeddings_only( self, @@ -179,6 +207,7 @@ class MiniCPMO45Stage0DuplexRuntime: seq: int | None = None, is_speech: bool = False, final: bool = False, + stage0_window: dict[str, object] | None = None, ) -> dict[str, object]: """Build scheduler-owned Stage0 input embeddings for one audio append. @@ -199,6 +228,16 @@ class MiniCPMO45Stage0DuplexRuntime: result["input_token_ids"] = list(state.prepared_input_token_ids) return result self._require_special_token_ids() + if isinstance(stage0_window, dict): + completed_ids = stage0_window.get("completed_token_ids") + if isinstance(completed_ids, list): + state.pending_window_generated_tokens = [int(token_id) for token_id in completed_ids] + completed_terminator = stage0_window.get("completed_terminator_token_id") + if isinstance(completed_terminator, int): + # Async sampling can overwrite the worker's pending token + # after the scheduler accepted a segment stop. Rebuild from + # the scheduler's canonical boundary, just like its content. + 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") state.audio_buffer = np.concatenate([state.audio_buffer, np.asarray(audio_waveform, dtype=np.float32)]) @@ -279,12 +318,25 @@ class MiniCPMO45Stage0DuplexRuntime: # the closure; the model's listen/speak policy depends on # seeing its own past decisions in context. pending_terminator = state.pending_terminator_token + closure_token_ids: list[int] = [] if pending_terminator is not None and units_built == 0: state.pending_terminator_token = None embed_parts.append(self._embed_token(pending_terminator)) token_ids.append(int(pending_terminator)) + closure_token_ids.append(int(pending_terminator)) embed_parts.append(self._embed_token(self.unit_end_token_id)) token_ids.append(self.unit_end_token_id) + closure_token_ids.append(self.unit_end_token_id) + if units_built == 0: + self._finalize_window_unit(state, closure_token_ids) + elif state.pending_window_unit is not None: + # Internal processor chunks belong to this same append. + state.pending_window_unit.embeds.extend( + self._embed_token(token_id) for token_id in closure_token_ids + ) + state.pending_window_unit.token_ids.extend(closure_token_ids) + unit_embed_start = len(embed_parts) + unit_token_start = len(token_ids) embed_parts.append(self._embed_token(self.unit_token_id)) token_ids.append(self.unit_token_id) if frame_blocks: @@ -311,6 +363,11 @@ class MiniCPMO45Stage0DuplexRuntime: state.audio_buffer = state.audio_buffer[consumed_samples:] state.audio_chunk_idx += 1 units_built += 1 + if state.window_enabled: + if state.pending_window_unit is None: + state.pending_window_unit = _MiniCPMO45WindowUnit() + state.pending_window_unit.embeds.extend(embed_parts[unit_embed_start:]) + state.pending_window_unit.token_ids.extend(token_ids[unit_token_start:]) chunk_size = self._streaming_chunk_size(processor) # Match official streaming_prefill: per chunk feed ONLY +audio. The assistant # turn is opened once at session init; re-emitting the turn-open prefix per chunk @@ -320,6 +377,9 @@ class MiniCPMO45Stage0DuplexRuntime: import torch + window_result = self._window_replacement_parts(state, stage0_window) + if window_result is not None: + embed_parts, token_ids = window_result inputs_embeds = torch.cat([self._as_2d_tensor(embed) for embed in embed_parts], dim=0) result = self._stage_prefill_result(True, start_time) result.update( @@ -332,6 +392,7 @@ class MiniCPMO45Stage0DuplexRuntime: "uses_model_runner_scheduler": True, "runner_kv_backed": True, "runtime_impl": "scheduler_data_plane", + "stage0_window_replaced": window_result is not None, } ) if is_speech and (append_identity is None or state.pending_speech_append_identity != append_identity): @@ -344,6 +405,85 @@ class MiniCPMO45Stage0DuplexRuntime: state.prepared_result = {k: v for k, v in result.items() if k not in {"inputs_embeds", "input_token_ids"}} return result + def _finalize_window_unit( + self, + state: _MiniCPMO45Stage0SessionState, + closure_token_ids: list[int], + ) -> None: + pending = state.pending_window_unit + if not state.window_enabled or pending is None: + state.pending_window_unit = None + state.pending_window_generated_tokens.clear() + return + generated = list(state.pending_window_generated_tokens) + token_ids = [*pending.token_ids, *generated, *closure_token_ids] + embeds = list(pending.embeds) + embeds.extend(self._embed_token(token_id) for token_id in generated) + embeds.extend(self._embed_token(token_id) for token_id in closure_token_ids) + state.window_units.append(_MiniCPMO45WindowUnit(embeds=embeds, token_ids=token_ids)) + state.pending_window_unit = None + state.pending_window_generated_tokens.clear() + + def _window_replacement_parts( + self, + state: _MiniCPMO45Stage0SessionState, + stage0_window: dict[str, Any] | None, + ) -> tuple[list[Any], list[int]] | None: + if not isinstance(stage0_window, dict) or stage0_window.get("replace") is not True: + return None + try: + drop_units = max(0, int(stage0_window.get("drop_units", 0))) + except (TypeError, ValueError): + return None + if drop_units > len(state.window_units): + raise RuntimeError( + "MiniCPM-o Stage-0 window history is shorter than the scheduler drop plan: " + f"drop={drop_units}, completed={len(state.window_units)}" + ) + del state.window_units[:drop_units] + + mode = stage0_window.get("mode") + embeds: list[Any] = [] + token_ids: list[int] = [] + previous_ids: list[int] = [] + if mode == "context": + embeds.extend(state.context_prefix_embeds) + token_ids.extend(state.context_prefix_token_ids) + previous = stage0_window.get("previous_token_ids") + previous_ids = [int(token_id) for token_id in previous] if isinstance(previous, list) else [] + if previous_ids: + # The plan carries the marker ids the scheduler sized the + # `previous` region from. Re-tokenizing here would be a second + # source of truth and could disagree with that length. + marker = stage0_window.get("previous_marker_token_ids") + marker_ids = [int(token_id) for token_id in marker] if isinstance(marker, list) else [] + embeds.extend(self._embed_token(token_id) for token_id in [*marker_ids, *previous_ids]) + token_ids.extend(marker_ids) + token_ids.extend(previous_ids) + embeds.extend(state.context_suffix_embeds) + token_ids.extend(state.context_suffix_token_ids) + else: + embeds.extend(state.context_embeds) + token_ids.extend(state.context_token_ids) + for unit in state.window_units: + embeds.extend(unit.embeds) + token_ids.extend(unit.token_ids) + if state.pending_window_unit is not None: + embeds.extend(state.pending_window_unit.embeds) + token_ids.extend(state.pending_window_unit.token_ids) + expected = stage0_window.get("replacement_prompt_len") + if isinstance(expected, int) and expected != len(token_ids): + unit_lengths = [len(unit.token_ids) for unit in state.window_units] + pending_len = len(state.pending_window_unit.token_ids) if state.pending_window_unit is not None else 0 + raise RuntimeError( + "MiniCPM-o Stage-0 window rebuild length mismatch: " + f"worker={len(token_ids)}, scheduler={expected}, mode={mode}, " + f"prefix={len(state.context_prefix_token_ids)}, " + f"suffix={len(state.context_suffix_token_ids)}, previous={len(previous_ids)}, " + f"units={unit_lengths}, pending={pending_len}, drop={drop_units}" + ) + return embeds, token_ids + @staticmethod def _stage_prefill_result(success: bool, start_time: float, reason: str = "") -> dict[str, object]: return { 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 b1c8db89c..03024bbe5 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 @@ -371,8 +371,18 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp 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), ) update_result = dict(result) + if result.get("stage0_window_replaced") is True: + window = duplex.get("stage0_window", {}) + logger.info( + "MiniCPM-o Stage-0 window replaced: mode=%s drop_units=%s tokens=%s seq=%s", + window.get("mode"), + window.get("drop_units"), + result.get("num_input_tokens"), + seq, + ) update_result.pop("inputs_embeds", None) if result.get("success") is not True: embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) 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 4d818b79f..ad329c271 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 @@ -865,8 +865,25 @@ def llm2tts( special_token_ids.get("chunk_eos_token_id"), special_token_ids.get("chunk_tts_eos_token_id"), } - tts_token_ids_slice = torch.tensor(full_token_ids[tts_bos_idx:end_idx], dtype=torch.long) - tts_hidden_slice = thinker_hidden_states[tts_bos_idx:end_idx].to(torch.float32).contiguous() + if is_native_duplex_handoff: + # Earlier unforwarded decisions may already be folded into + # the rebuilt prompt while still appearing in this delta. + # Align the explicit tts_bos path by the segment's end, just + # like the native speak/text paths below, rather than adding + # those decisions to the prompt a second time. + out_start = tts_bos_idx - prompt_token_ids_len + out_end = tts_eos_idx - prompt_token_ids_len if tts_eos_idx is not None else len(llm_output_ids) + hidden_base = int(thinker_hidden_states.shape[0]) - len(llm_output_ids) + if hidden_base >= 0 and out_end > out_start: + tts_token_ids_slice = torch.tensor(llm_output_ids[out_start:out_end], dtype=torch.long) + tts_hidden_slice = ( + thinker_hidden_states[hidden_base + out_start : hidden_base + out_end] + .to(torch.float32) + .contiguous() + ) + else: + tts_token_ids_slice = torch.tensor(full_token_ids[tts_bos_idx:end_idx], dtype=torch.long) + tts_hidden_slice = thinker_hidden_states[tts_bos_idx:end_idx].to(torch.float32).contiguous() elif is_native_duplex_handoff: # Official MiniCPM-o duplex does not prefill an assistant # <|tts_bos|> boundary before generation. A segment delta can From cb31e0774fd3ddb6fca4ef595fa69202919c1add Mon Sep 17 00:00:00 2001 From: Armaan Amatya <38102001+armaanamatya@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:34:31 -0500 Subject: [PATCH 3/7] [Misc] Deprecate VLLM_OMNI_VOXCPM_CODE_PATH (#7779) Signed-off-by: Armaan Amatya --- .../offline_inference/text_to_speech.md | 51 +++++++++++++++++-- .../text_to_speech/README.md | 2 +- .../text_to_speech/voxcpm2/end2end.py | 5 +- .../voxcpm2/test_voxcpm2_import_utils.py | 37 ++++++++++++++ .../models/voxcpm2/voxcpm2_import_utils.py | 9 ++-- 5 files changed, 95 insertions(+), 9 deletions(-) create mode 100644 tests/model_executor/models/voxcpm2/test_voxcpm2_import_utils.py diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md index 95ffc5543..62e6b96d2 100644 --- a/docs/user_guide/examples/offline_inference/text_to_speech.md +++ b/docs/user_guide/examples/offline_inference/text_to_speech.md @@ -2,7 +2,6 @@ Source . - vLLM-Omni supports several autoregressive TTS models. They share a common CLI shape (`--text`, `--ref-audio`, `--ref-text`, `--output-dir`) and live together in this hub. Each model has its own subdirectory containing a @@ -15,7 +14,7 @@ list of supported architectures across all modalities, see ## Supported Models | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | -|---|---|---|---|---|---|---| +| --- | --- | --- | --- | --- | --- | --- | | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz | @@ -45,13 +44,15 @@ python examples/offline_inference/text_to_speech//end2end.py \ Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`. ### Prerequisites + ```bash pip install voxcpm # or, for a local source checkout: -export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm +pip install -e /path/to/voxcpm ``` ### Quick start + ```bash python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ --model openbmb/VoxCPM2 \ @@ -59,7 +60,9 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ ``` ### Voice cloning + Pass a reference audio for isolated cloning, or both `--ref-audio` + `--ref-text` for prompt continuation: + ```bash python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ --text "Hello, this is a voice clone demo." \ @@ -68,6 +71,7 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ ``` ### Notes + - Output: 48 kHz mono WAV. - Deploy config: `vllm_omni/deploy/voxcpm2.yaml` (default). @@ -78,12 +82,14 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ 2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz. ### Prerequisites + ```bash uv pip install -e . # Includes soundfile, onnxruntime, x-transformers, einops via requirements. ``` Download the model snapshot: + ```python from huggingface_hub import snapshot_download snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', @@ -91,15 +97,18 @@ snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', ``` If your downloaded checkpoint lacks `config.json`, add it: + ```json { "model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"] } ``` + This is required because `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers the class mapping; the loader still reads `model_type` from `config.json` to select the class. ### Quick start + ```bash python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ --model pretrained_models/Fun-CosyVoice3-0.5B \ @@ -107,7 +116,9 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ ``` ### Voice cloning + Pass a reference audio. Note that CosyVoice3's `--prompt-text` is a system-style prompt for the GPT stage, not a reference transcript: + ```bash python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ --model pretrained_models/Fun-CosyVoice3-0.5B \ @@ -117,6 +128,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ ``` ### Notes + - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform. - Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config ` to override. - `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path. @@ -129,12 +141,14 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ No extra packages are required; the DAC codec is vendored in vLLM-Omni. ### Quick start + ```bash python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ --text "Hello, this is a test of the Fish Speech text to speech system." ``` ### Voice cloning + ```bash python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ --text "Hello, this is a cloned voice." \ @@ -143,14 +157,17 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ ``` ### Streaming + ```bash python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ --text "Hello, this is a streaming test." \ --streaming ``` + Streaming requires `async_chunk: true` in the deploy config. ### Notes + - Output: 44.1 kHz mono WAV. - DAC codec weights (`codec.pth`) are loaded lazily from the model directory. @@ -161,6 +178,7 @@ Streaming requires `async_chunk: true` in the deploy config. 2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning. ### Quick start + ```bash python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ --model zai-org/GLM-TTS \ @@ -171,12 +189,14 @@ python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ ``` ### Architecture -``` + +```text Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz) (Llama-based) (32k vocab) (Flow Matching) ``` ### Notes + - `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis. - Reference audio should be 3-10 seconds. - First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder. @@ -190,12 +210,15 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design). ### Prerequisites + ```bash huggingface-cli download k2-fsa/OmniVoice ``` + Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`. ### Quick start (auto voice) + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -203,6 +226,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Voice cloning + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -212,6 +236,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Voice design + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -220,6 +245,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Language hint + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -228,6 +254,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Notes + - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking. - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz. @@ -238,13 +265,16 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ 3-task-variant TTS with 24 kHz output. Has its own argparse surface (this script does not follow the common `--text` / `--ref-audio` shape). ### Prerequisites + For ROCm builds, replace `onnxruntime` with `onnxruntime-rocm`: + ```bash pip uninstall onnxruntime pip install onnxruntime-rocm ``` ### Task variants + - `CustomVoice`: predefined speaker (speaker ID) with optional style instruction. - `VoiceDesign`: text + descriptive instruction designs a new voice. - `Base`: voice cloning from reference audio + transcript. @@ -264,16 +294,20 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty ``` ### Streaming + ```bash python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --query-type CustomVoice \ --streaming \ --output-dir /tmp/out_stream ``` + Streaming requires `async_chunk: true` in the deploy config. ### Batched decoding + The Code2Wav stage supports batched decoding through the SpeechTokenizer. Configure both stages with `max_num_seqs > 1` via `--stage-overrides` and pass multiple prompts via `--txt-prompts`: + ```bash python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --query-type CustomVoice \ @@ -281,9 +315,11 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --batch-size 4 \ --stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}' ``` + `--batch-size` must match a CUDA-graph capture size (1, 2, 4, 8, 16…). ### Notes + - Run `--help` for the full argument surface. - See `qwen3_tts/end2end.py` for the prompt-length-estimation logic the Talker uses. @@ -294,12 +330,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol. ### Prerequisites + Latest `mistral_common` with `SpeechRequest` support: + ```bash pip install -e /path/to/mistral-common # or upgrade from PyPI when available ``` ### Quick start (voice preset) + ```bash python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ --write-audio --voice cheerful_female \ @@ -308,6 +347,7 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ ``` ### Voice cloning (capability gated upstream) + ```bash python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ --write-audio \ @@ -317,15 +357,18 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ ``` ### Streaming + concurrency + ```bash python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ --num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \ --model mistralai/Voxtral-4B-TTS-2603 \ --text "..." ``` + Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). ### Notes + - `--num-prompts N` replicates the prompt for performance measurement. - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`. - Run `--help` for the full argument surface. diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 36d1b44ee..f031ee2e7 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -563,7 +563,7 @@ Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → F ```bash pip install voxcpm # or, for a local source checkout: -export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm +pip install -e /path/to/voxcpm ``` ### Quick start diff --git a/examples/offline_inference/text_to_speech/voxcpm2/end2end.py b/examples/offline_inference/text_to_speech/voxcpm2/end2end.py index feba91b7e..53941fb24 100644 --- a/examples/offline_inference/text_to_speech/voxcpm2/end2end.py +++ b/examples/offline_inference/text_to_speech/voxcpm2/end2end.py @@ -1,7 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Offline VoxCPM2 inference example (native AR pipeline). Uses the single-stage native AR config (voxcpm2.yaml). -Requires the `voxcpm` package or VLLM_OMNI_VOXCPM_CODE_PATH env var. +Requires the `voxcpm` package (pip install voxcpm). """ from __future__ import annotations diff --git a/tests/model_executor/models/voxcpm2/test_voxcpm2_import_utils.py b/tests/model_executor/models/voxcpm2/test_voxcpm2_import_utils.py new file mode 100644 index 000000000..f1f0e74c9 --- /dev/null +++ b/tests/model_executor/models/voxcpm2/test_voxcpm2_import_utils.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""VLLM_OMNI_VOXCPM_CODE_PATH still works but warns that it is deprecated (#6232).""" + +from pathlib import Path + +import pytest + +from vllm_omni.model_executor.models.voxcpm2 import voxcpm2_import_utils + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _capture_warnings(monkeypatch) -> list[str]: + messages: list[str] = [] + monkeypatch.setattr(voxcpm2_import_utils.logger, "warning_once", lambda msg, *args: messages.append(msg)) + return messages + + +def test_code_path_env_warns_and_is_still_used(monkeypatch, tmp_path: Path): + messages = _capture_warnings(monkeypatch) + monkeypatch.setenv("VLLM_OMNI_VOXCPM_CODE_PATH", str(tmp_path)) + + candidates = voxcpm2_import_utils._iter_voxcpm2_src_candidates() + + assert candidates[0] == tmp_path + assert len(messages) == 1 + assert "VLLM_OMNI_VOXCPM_CODE_PATH is deprecated" in messages[0] + + +def test_no_warning_without_code_path_env(monkeypatch): + messages = _capture_warnings(monkeypatch) + monkeypatch.delenv("VLLM_OMNI_VOXCPM_CODE_PATH", raising=False) + + voxcpm2_import_utils._iter_voxcpm2_src_candidates() + + assert messages == [] diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_import_utils.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_import_utils.py index 231a51bbc..0a1b75a4a 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_import_utils.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_import_utils.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Dynamic import utilities for the native VoxCPM2 package. Supports three discovery modes (first match wins): -1. ``VLLM_OMNI_VOXCPM_CODE_PATH`` env var (explicit source tree) +1. ``VLLM_OMNI_VOXCPM_CODE_PATH`` env var (explicit source tree; deprecated) 2. Sibling ``../VoxCPM/src`` relative to the vllm-omni repo root 3. pip-installed ``voxcpm`` package (>= 2.0) """ @@ -26,6 +26,10 @@ def _iter_voxcpm2_src_candidates() -> list[Path]: candidates: list[Path] = [] env_path = os.environ.get("VLLM_OMNI_VOXCPM_CODE_PATH") if env_path: + logger.warning_once( + "VLLM_OMNI_VOXCPM_CODE_PATH is deprecated and will be removed in a future release. " + "Install voxcpm>=2.0 instead (pip install voxcpm, or pip install -e for a source checkout)." + ) candidates.append(Path(env_path).expanduser()) repo_root = Path(__file__).resolve().parents[4] @@ -71,7 +75,6 @@ def _import_voxcpm2_attrs(module_name: str, *attr_names: str) -> tuple[Any, ...] raise ImportError( f"Could not import {attr_names} from {module_name}. " f"Install voxcpm>=2.0: pip install voxcpm. " - f"Or set VLLM_OMNI_VOXCPM_CODE_PATH to the VoxCPM source tree. " f"Last error: {last_exc}" ) From 1586c41ae45e38d1f9b831f241a3f1b044dbc308 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Mon, 21 Sep 2026 18:34:40 -0700 Subject: [PATCH 4/7] [CI][ROCm] Shorten the blocking CosyVoice test path (#7748) Signed-off-by: andyluo7 --- .buildkite/amd/test-amd-nightly.yml | 18 +++++++++++ .buildkite/amd/test-amd-ready.yml | 8 +++-- tests/buildkite/test_amd_pipeline.py | 24 ++++++++++++-- .../cosyvoice3/test_cosyvoice3_components.py | 31 +++++++++++++++++++ .../test_cosyvoice3_incremental_hift.py | 18 +++++++---- .../models/cosyvoice3/cosyvoice3_code2wav.py | 24 ++++++++++++-- 6 files changed, 111 insertions(+), 12 deletions(-) diff --git a/.buildkite/amd/test-amd-nightly.yml b/.buildkite/amd/test-amd-nightly.yml index 0da365d0c..4544a4f20 100644 --- a/.buildkite/amd/test-amd-nightly.yml +++ b/.buildkite/amd/test-amd-nightly.yml @@ -64,3 +64,21 @@ steps: commands: - export VLLM_ROCM_USE_AITER=1 - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py::test_text_to_text_audio_001 -m "full_model and rocm and MI325 and cards_2" --run-level "full_model" + + - group: ":speaker: CosyVoice3 MI300 Nightly" + steps: + # Exercise all sync/async and Chinese/English cases with the platform's + # default attention backend. On ROCm this retains end-to-end AITER + # coverage without putting its cold JIT build on the READY critical path. + - label: "CosyVoice3-TTS E2E Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: NonBlocking + timeout_in_minutes: 90 + retry: + automatic: + - exit_status: 134 + limit: 1 + commands: + - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model" diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index e46b559f3..ae64e7c4b 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -328,7 +328,10 @@ steps: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model" - - label: "CosyVoice3-TTS E2E Test" + # Keep a fast model/server integration check in the blocking lane. The + # full default-backend suite remains in the nightly pipeline because a + # cold ROCm AITER MHA build can add more than 25 minutes to this job. + - label: "CosyVoice3-TTS E2E Smoke (SDPA)" agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] @@ -340,7 +343,8 @@ steps: - exit_status: 134 limit: 1 commands: - - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model" + - export DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA + - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py::test_voice_clone_zh_002 -m "slow" --run-level "core_model" # TODO: Bagel test on ROCm is very unstable. @tjtanaa # Need to debug before reneable numerical changes across large PRs diff --git a/tests/buildkite/test_amd_pipeline.py b/tests/buildkite/test_amd_pipeline.py index fac4f8c29..36df7ac3e 100644 --- a/tests/buildkite/test_amd_pipeline.py +++ b/tests/buildkite/test_amd_pipeline.py @@ -67,11 +67,31 @@ def test_ready_diffusion_cpu_suite_is_sharded() -> None: assert "--shard-id=$$BUILDKITE_PARALLEL_JOB" in pytest_command -def test_cosyvoice_gpu_abort_gets_one_fresh_job_retry() -> None: - step = _find_step("CosyVoice3-TTS E2E Test", AMD_READY_PIPELINE) +def test_cosyvoice_ready_smoke_uses_sdpa() -> None: + step = _find_step("CosyVoice3-TTS E2E Smoke (SDPA)", AMD_READY_PIPELINE) + assert step["grade"] == "Blocking" assert step["retry"] == {"automatic": [{"exit_status": 134, "limit": 1}]} + assert "export DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" in step["commands"] + pytest_command = next(command for command in step["commands"] if "pytest" in command) + assert "tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py::test_voice_clone_zh_002" in pytest_command + + +def test_cosyvoice_full_default_backend_suite_runs_nightly() -> None: + step = _find_step("CosyVoice3-TTS E2E Test", AMD_NIGHTLY_PIPELINE) + + assert step["grade"] == "NonBlocking" + assert step["timeout_in_minutes"] == 90 + assert step["retry"] == {"automatic": [{"exit_status": 134, "limit": 1}]} + assert all("DIFFUSION_ATTENTION_BACKEND" not in command for command in step["commands"]) + + pytest_command = next(command for command in step["commands"] if "pytest" in command) + assert "tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py" in pytest_command + assert "::" not in pytest_command + + +def test_amd_template_preserves_step_retry_policy() -> None: template = AMD_TEMPLATE.read_text(encoding="utf-8") # Both grouped and top-level AMD steps must preserve an explicit retry # policy when the source suite is rendered into the uploaded pipeline. diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 46255110c..56f9c9694 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -260,6 +260,37 @@ class TestDiTAttention: assert attention.to_k.out_features == 512 assert attention.to_v.out_features == 512 + def test_embedded_dit_honors_diffusion_attention_backend(self, monkeypatch): + """The LLM-hosted code2wav DiT must consume the diffusion env override.""" + from vllm_omni.diffusion.config import get_current_diffusion_config_or_none + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import _build_dit_estimator + + monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA") + + estimator = _build_dit_estimator( + { + "dim": 32, + "depth": 1, + "heads": 2, + "dim_head": 16, + "dropout": 0.0, + "ff_mult": 2, + "mel_dim": 8, + "mu_dim": 8, + "spk_dim": 4, + "out_channels": 8, + "static_chunk_size": 4, + "num_decoding_left_chunks": 1, + } + ) + + attention = estimator.transformer_blocks[0].attn.attn + assert attention.backend_pref == "TORCH_SDPA" + assert attention.backend_explicit is True + assert attention.attn_backend is not None + assert attention.attn_backend.get_name() == "SDPA" + assert get_current_diffusion_config_or_none() is None + @pytest.mark.core_model @pytest.mark.cpu def test_forward_passes_2d_mask_without_casting_qkv(self): diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py index 91b322864..1dd0799ef 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py @@ -39,7 +39,10 @@ CHUNK_LEN = 24 TOTAL_MEL = 96 SPM = 480 # samples per mel frame for the small test HiFT (8*5*3*4), matching the real model -LONG_TOTAL_MEL = 400 # long enough that the 64-frame window actually truncates +WINDOW_LEN = 64 +# Six chunks exercise three updates after history exceeds the bounded window, +# without spending minutes recomputing an unnecessarily long synthetic signal. +LONG_TOTAL_MEL = 144 PCM16_LSB = 1.0 / 32767.0 # smallest representable PCM16 difference # Worst observed deviation is 1.97e-7 (~2x float32 eps); atol=1e-6 keeps 5x headroom. @@ -189,16 +192,19 @@ def test_incremental_hift_bounded_window_is_close(config): def test_incremental_hift_matches_streaming_reference_within_tolerance(config): """Windowed output matches the full-cumulative reference within float32 rounding. - The load-bearing correctness test: 400 mel frames far exceeds the 64-frame - window, so truncation genuinely exercises the phase/noise carry. Float64 - agrees to ~1e-16, confirming the residual here is rounding, not truncation. + The load-bearing correctness test crosses the 64-frame window repeatedly, + so truncation genuinely exercises the phase/noise carry. Float64 agrees to + ~1e-16, confirming the residual here is rounding, not truncation. """ hift = _make_hift(config) - model = _make_model(hift, window_len=64) # the real _hift_window_len + model = _make_model(hift, window_len=WINDOW_LEN) # the real _hift_window_len chunks = _chunks(total_mel=LONG_TOTAL_MEL) trim = int(hift.f0_predictor.condnet[0].causal_padding) - assert LONG_TOTAL_MEL > 64 + trim + CHUNK_LEN # window must actually truncate + truncating_steps = sum( + history_len > WINDOW_LEN + trim for history_len in range(CHUNK_LEN, LONG_TOTAL_MEL, CHUNK_LEN) + ) + assert truncating_steps >= 3 # exercise repeated truncation, not only the boundary full = _full_reference(model, chunks) incr = _incremental(model, chunks) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 1713467c8..b35aefe52 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -10,7 +10,8 @@ This module contains the code2wav (token-to-waveform) stage which uses: """ from collections import Counter -from typing import TypedDict, cast +from collections.abc import Mapping +from typing import Any, TypedDict, cast import numpy as np import torch @@ -18,6 +19,8 @@ import torch.nn as nn from omegaconf import DictConfig from vllm.logger import init_logger +from vllm_omni.diffusion.config import get_current_diffusion_config_or_none, set_current_diffusion_config +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ( CausalConditionalCFM, @@ -37,6 +40,23 @@ from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config logger = init_logger(__name__) +def _build_dit_estimator(estimator_config: Mapping[str, Any]) -> DiT: + """Build CosyVoice's embedded DiT with diffusion backend configuration. + + The code2wav model runs inside an LLM-generation worker, so it does not go + through the diffusion model loader that normally installs an + ``OmniDiffusionConfig`` during model construction. Install a local default + config when needed so ``DIFFUSION_ATTENTION_BACKEND`` and the platform + default are resolved by the same path as standalone diffusion models. + """ + if get_current_diffusion_config_or_none() is not None: + return DiT(**estimator_config) + + diffusion_config = OmniDiffusionConfig() + with set_current_diffusion_config(diffusion_config): + return DiT(**estimator_config) + + class StreamingFlowItem(TypedDict, total=False): """One entry of the batched-streaming item list passed to forward_streaming_batch.""" @@ -71,7 +91,7 @@ class CosyVoice3Code2Wav(nn.Module): cfm_params = DictConfig(decoder_cfg["cfm_params"]) # DiT estimator using diffusion attention (Flash/Sage/SDPA backends) - estimator = DiT(**decoder_cfg["estimator"]) + estimator = _build_dit_estimator(decoder_cfg["estimator"]) decoder = CausalConditionalCFM( in_channels=decoder_cfg["in_channels"], From 1d3dc8bf83e9768cf73b5b0660096eb3c104639b Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Mon, 21 Sep 2026 21:34:50 -0400 Subject: [PATCH 5/7] [Misc] Refactor static stage metadata lookups (#7602) Signed-off-by: Nick Cao Co-authored-by: Codex --- .../test_async_omni_pause_sleep_routing.py | 4 +- tests/entrypoints/test_omni_entrypoints.py | 37 ++++++++++++++++--- vllm_omni/engine/omni_engine_base.py | 2 +- vllm_omni/entrypoints/async_omni.py | 8 ++-- vllm_omni/entrypoints/async_omni_base.py | 8 ++-- 5 files changed, 41 insertions(+), 18 deletions(-) diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index d3a60e5a6..887cf8885 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -33,9 +33,9 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni: omni._final_output_handler = lambda: None omni.reset_mm_cache = AsyncMock() - stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types] + stage_configs = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types] omni.engine = SimpleNamespace( - stage_clients=stage_clients, + stage_configs=stage_configs, stage_vllm_configs=[None] * len(stage_types), collective_rpc_async=AsyncMock(return_value=[True]), ) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 969bbaf91..1b33850da 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -19,7 +19,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES -from vllm_omni.config.stage_config import StageConfig +from vllm_omni.config.stage_config import StageConfig, StageType from vllm_omni.engine.async_omni_engine import StageRuntimeInfo from vllm_omni.engine.messages import ErrorMessage, OutputMessage from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -107,7 +107,13 @@ class FakeAsyncOmniEngine: self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]] self.num_stages = len(self.stage_metadata) self.stage_configs = [ - StageConfig(stage_id=i, model_stage="dummy-model").to_omegaconf() for i in range(self.num_stages) + StageConfig( + stage_id=i, + model_stage="dummy-model", + stage_type=StageType(metadata.stage_type), + is_comprehension=getattr(metadata, "is_comprehension", False), + ).to_omegaconf() + for i, metadata in enumerate(self.stage_metadata) ] self.default_sampling_params_list = default_sampling_params_list or [ SamplingParams(max_tokens=8) for _ in range(self.num_stages) @@ -481,7 +487,7 @@ def test_model_config_and_vllm_config_forward_from_comprehension_stage(): io_processor = SimpleNamespace(name="io-processor") omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( - stage_clients=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)], + stage_configs=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)], stage_vllm_configs=[None, vllm_config], ) omni.input_processor = input_processor @@ -502,7 +508,7 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs(): io_processor = SimpleNamespace(name="io-processor") omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( - stage_clients=[SimpleNamespace(is_comprehension=True)], + stage_configs=[SimpleNamespace(is_comprehension=True)], stage_vllm_configs=[vllm_config], ) omni.input_processor = input_processor @@ -524,10 +530,11 @@ def test_get_diffusion_od_config_returns_diffusion_stage_config(): diffusion_od_config = object() omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( + stage_configs=[SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion")], stage_clients=[ SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion", od_config=diffusion_od_config), - ] + ], ) assert omni.get_diffusion_od_config() is diffusion_od_config @@ -537,15 +544,33 @@ def test_get_diffusion_od_config_falls_back_to_inner_engine(): diffusion_od_config = object() omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( + stage_configs=[SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion")], stage_clients=[ SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion", _engine=SimpleNamespace(od_config=diffusion_od_config)), - ] + ], ) assert omni.get_diffusion_od_config() is diffusion_od_config +@pytest.mark.asyncio +async def test_async_omni_rejects_diffusion_list_prompt_without_stage_clients( + monkeypatch: pytest.MonkeyPatch, +): + engine = FakeAsyncOmniEngine(stage_metadata=DIFFUSION_ONLY_META) + engine.stage_clients = [] + _patch_engine(monkeypatch, engine) + + app = AsyncOmni("dummy-model") + try: + with pytest.raises(ValueError, match="single prompt per request"): + async for _ in app.generate(prompt=["prompt-1", "prompt-2"], request_id="req-1"): + pass + finally: + app.shutdown() + + @pytest.mark.asyncio async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.MonkeyPatch): engine = FakeAsyncOmniEngine( diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index 0a89f4141..fbb30187f 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -422,7 +422,7 @@ class OmniEngineBase: for client in self.stage_clients ] supported_tasks: set[str] = set() - if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): + if any(getattr(stage_config, "is_comprehension", False) for stage_config in self.stage_configs): supported_tasks.add("generate") if any(meta.final_output_type == "audio" for meta in self.stage_metadata): supported_tasks.add("speech") diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ba4e75a11..7d1e702ef 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -186,7 +186,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): # Reject diffusion list-prompt early with a clear API error. if isinstance(prompt, list) and any( - getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", []) + stage_config.stage_type == "diffusion" for stage_config in self.engine.stage_configs ): raise ValueError( "Diffusion stages accept only a single prompt per request. " @@ -626,7 +626,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]: """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC).""" - n_stages = len(self.engine.stage_clients) + n_stages = len(self.engine.stage_configs) if stage_ids is None: stage_ids = list(range(n_stages)) else: @@ -640,8 +640,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient): ar_stage_ids: list[int] = [] diffusion_stage_ids: list[int] = [] for sid in stage_ids: - client = self.engine.stage_clients[sid] - if getattr(client, "stage_type", "llm") == "diffusion": + stage_config = self.engine.stage_configs[sid] + if stage_config.stage_type == "diffusion": diffusion_stage_ids.append(sid) else: ar_stage_ids.append(sid) diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index 6fa98f8cd..35572e303 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -155,13 +155,13 @@ class AsyncOmniBase(OmniBase): def _get_comprehension_stage_index(self) -> int | None: fallback_idx: int | None = None - for idx, stage_client in enumerate(self.engine.stage_clients): + for idx, stage_config in enumerate(self.engine.stage_configs): stage_vllm_config = self.engine.stage_vllm_configs[idx] if stage_vllm_config is None: continue if fallback_idx is None: fallback_idx = idx - if stage_client.is_comprehension: + if stage_config.is_comprehension: return idx return fallback_idx @@ -186,13 +186,11 @@ class AsyncOmniBase(OmniBase): def get_diffusion_od_config(self) -> Any | None: """Return the diffusion-stage config when the pipeline has one.""" - saw_diffusion_stage = False + saw_diffusion_stage = any(stage_config.stage_type == "diffusion" for stage_config in self.engine.stage_configs) for stage_client in self.engine.stage_clients: if getattr(stage_client, "stage_type", None) != "diffusion": continue - saw_diffusion_stage = True - od_config = getattr(stage_client, "od_config", None) if od_config is not None: return od_config From 085cf0ff344119b7536247f6af4dc922a13a91d6 Mon Sep 17 00:00:00 2001 From: Tianyao Wu <54675599+twu3202@users.noreply.github.com> Date: Tue, 22 Sep 2026 09:34:58 +0800 Subject: [PATCH 6/7] [Bugfix] Catch asyncio.TimeoutError so duplex serving starts on Python 3.10 (#7412) Signed-off-by: Tianyao Wu Co-authored-by: Claude Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com> --- tests/engine/duplex/test_session_manager.py | 32 +++++++++++++++++++ tests/engine/duplex/test_session_runner.py | 35 +++++++++++++++++++++ vllm_omni/engine/duplex/session/context.py | 3 +- vllm_omni/engine/duplex/session/manager.py | 4 ++- vllm_omni/engine/duplex/session/runner.py | 3 +- vllm_omni/entrypoints/duplex/websocket.py | 3 +- 6 files changed, 76 insertions(+), 4 deletions(-) diff --git a/tests/engine/duplex/test_session_manager.py b/tests/engine/duplex/test_session_manager.py index cb00737d9..3d36c9891 100644 --- a/tests/engine/duplex/test_session_manager.py +++ b/tests/engine/duplex/test_session_manager.py @@ -51,6 +51,7 @@ from vllm_omni.engine.duplex.plugin import ( PcmAppendReservation, ) from vllm_omni.engine.duplex.session import manager as session_manager_module +from vllm_omni.engine.duplex.session.context import DuplexSessionTasks 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.manager import DuplexSessionManager @@ -1612,3 +1613,34 @@ async def test_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: flo finally: shutdown.set() await asyncio.wait_for(task, timeout=5.0) + + +# --------------------------------------------------------------------------- # +# Cancellation waits # +# --------------------------------------------------------------------------- # + + +async def test_cancel_append_tasks_absorbs_a_task_that_outlives_the_wait() -> None: + """The wait after cancelling is allowed to time out; that is not an error. + + On Python 3.10 the timeout arrives as ``asyncio.TimeoutError``, which is a + different class from the builtin ``TimeoutError`` until 3.11. + """ + + async def outlives_the_first_cancel() -> bool: + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + await asyncio.sleep(1) + raise + return True + + tasks = DuplexSessionTasks() + task = asyncio.create_task(outlives_the_first_cancel()) + await asyncio.sleep(0) + tasks.track_append_task(task, epoch=0, final=False, response_bound=False) + tasks.append_tail = task + + assert await tasks.cancel_append_tasks(timeout_s=0.05) is True + assert tasks.append_tail is None + await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=5.0) diff --git a/tests/engine/duplex/test_session_runner.py b/tests/engine/duplex/test_session_runner.py index f48154a48..cf6a0243d 100644 --- a/tests/engine/duplex/test_session_runner.py +++ b/tests/engine/duplex/test_session_runner.py @@ -771,6 +771,41 @@ async def test_cancel_response_aborts_the_stage_request_and_reports_playback() - await close_harness(h) +@pytest.mark.asyncio +async def test_cancel_response_absorbs_a_response_task_that_outlives_the_wait() -> None: + """The 0.25 s wait after cancelling the response task may time out. + + On Python 3.10 that timeout is ``asyncio.TimeoutError``, not the builtin + ``TimeoutError``, and it must not stop the cancel from completing. + """ + + async def outlives_the_first_cancel() -> None: + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + await asyncio.sleep(1) + raise + + h = await open_harness() + task = asyncio.create_task(outlives_the_first_cancel()) + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + h.runner.tasks.active_response_task = task + + events = await h.run(commands.CancelResponse()) + # The runner is still inside that wait when the mailbox goes quiet. + events += await h.settle(idle_s=0.6, timeout_s=5.0) + assert find(events, "response.done").status == "cancelled" + assert "error" not in types(events) + assert h.session.epoch == 1 + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await close_harness(h) + + @pytest.mark.asyncio async def test_listen_decision_is_consumed_and_never_forwarded_to_tts() -> None: h = await open_harness() diff --git a/vllm_omni/engine/duplex/session/context.py b/vllm_omni/engine/duplex/session/context.py index c0dfabc62..02853561a 100644 --- a/vllm_omni/engine/duplex/session/context.py +++ b/vllm_omni/engine/duplex/session/context.py @@ -141,7 +141,8 @@ class DuplexSessionTasks: task.cancel() try: await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s) - except TimeoutError: + # asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11. + except (TimeoutError, asyncio.TimeoutError): pass if cancelled_tail is not None and self.append_tail is cancelled_tail: self.append_tail = None diff --git a/vllm_omni/engine/duplex/session/manager.py b/vllm_omni/engine/duplex/session/manager.py index 6d808caae..76e2c7ee8 100644 --- a/vllm_omni/engine/duplex/session/manager.py +++ b/vllm_omni/engine/duplex/session/manager.py @@ -699,7 +699,9 @@ class DuplexSessionManager: while not shutdown_event.is_set(): try: await asyncio.wait_for(shutdown_event.wait(), timeout=interval) - except TimeoutError: + # The timeout is this loop's tick. asyncio.TimeoutError is not the builtin + # TimeoutError before Python 3.11, so catch both or the tick escapes the loop. + except (TimeoutError, asyncio.TimeoutError): try: await self.reap_expired() except Exception: diff --git a/vllm_omni/engine/duplex/session/runner.py b/vllm_omni/engine/duplex/session/runner.py index fb1e4f34f..02dc61708 100644 --- a/vllm_omni/engine/duplex/session/runner.py +++ b/vllm_omni/engine/duplex/session/runner.py @@ -1486,7 +1486,8 @@ class DuplexSessionRunner: active_task.cancel() try: await asyncio.wait_for(asyncio.gather(active_task, return_exceptions=True), timeout=0.25) - except TimeoutError: + # asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11. + except (TimeoutError, asyncio.TimeoutError): pass if notify: self.emit( diff --git a/vllm_omni/entrypoints/duplex/websocket.py b/vllm_omni/entrypoints/duplex/websocket.py index 8ec9cc780..1224f7d51 100644 --- a/vllm_omni/entrypoints/duplex/websocket.py +++ b/vllm_omni/entrypoints/duplex/websocket.py @@ -28,7 +28,8 @@ async def receive_text_with_timeout(websocket: WebSocket, timeout_s: float | Non return await websocket.receive_text() try: return await asyncio.wait_for(websocket.receive_text(), timeout=timeout_s) - except TimeoutError: + # asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11. + except (TimeoutError, asyncio.TimeoutError): return None From d4570afabfb02c2a899904add85f588a92271ec2 Mon Sep 17 00:00:00 2001 From: Hoseung Kim Date: Tue, 22 Sep 2026 10:35:07 +0900 Subject: [PATCH 7/7] [Frontend] Reject built-in name collisions in voice upload and add VLLM_OMNI_SPEAKER_REGISTRATION_POLICY (#6848) Signed-off-by: Hoseung Kim Co-authored-by: Hoseung Kim --- docs/configuration/environment_variables.md | 1 + docs/serving/speech_api.md | 14 ++++ .../openai_api/test_serving_speech.py | 69 +++++++++++++++++++ .../config/environment_variable_inventory.py | 1 + .../entrypoints/openai/serving_speech.py | 43 ++++++++++++ 5 files changed, 128 insertions(+) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index e113c9d7c..eaefb1015 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -88,6 +88,7 @@ depends on the installed kernels and model path. | `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; defaults to on for Qwen3-TTS and off for other pipelines | Pipeline default computed from `pipeline_config.model_type` at engine initialization; env override resolved at `Orchestrator` construction and separately when the serving-side final-output drain starts | An explicit env value wins; otherwise both consumers use the engine's pipeline default. Values are stripped and case-normalized; any unrecognized value selects the legacy poll loop. Set before server startup. | Experimental | | `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control | | `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `/vllm_omni_orch_monitor_.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic | +| `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY` | `overwrite` or `immutable`; default `overwrite` | Speech server; read when speaker storage initializes | Environment-only setting. `immutable` rejects re-registering an existing uploaded voice name until it is deleted; any other value raises `ValueError` at startup. | Experimental | | `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental | | `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental | | `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental | diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index d973b51a5..2759fffc0 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -289,6 +289,19 @@ Upload a new voice sample for voice cloning in Base task TTS requests. Fields `ref_text` and `speaker_description` are omitted when not provided at upload time. +**Naming rules:** + +- Names that collide with one of the model's built-in or precomputed voices are rejected (400). Voice + files already on disk under such a name are ignored at startup with a warning, so an upload can never + shadow a built-in voice. +- Re-uploading an existing uploaded name overwrites it in place (the previous audio file is + deleted). Set `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY=immutable` on the server to reject duplicates instead, + requiring an explicit `DELETE /v1/audio/voices/{name}` before re-registering — useful when the + endpoint is reachable by multiple writers and silent overwrites are a risk. +- The voice registry has no per-user ownership: any client that can reach the endpoint can + overwrite (default policy) or delete any uploaded voice. For multi-tenant deployments, add + authentication at a proxy and namespace voice names per user (e.g. `{user}.{name}`). + **Usage Example:** ```bash @@ -577,6 +590,7 @@ by `GET /v1/audio/voices`. Valid precomputed voices can be used in | ---------- | --------- | ------------- | | `SPEAKER_SAMPLES_DIR` | `~/.cache/vllm-omni/speakers` | Directory for persisted uploaded speakers (`.safetensors` files). | | `SPEAKER_MAX_UPLOADED` | `1000` | Maximum number of uploaded speakers kept on disk. Upload requests past the cap return 400. | +| `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY` | `overwrite` | `immutable` rejects re-uploading an existing uploaded name (400) until it is deleted; any other value fails startup. | The in-memory LRU has a fixed 512 MiB byte budget. diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 9a2f9e1fc..daf1c5d5b 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -505,6 +505,75 @@ class TestSpeechAPI: assert voice_info["file_size"] == len(audio_content) response = client.delete("/v1/audio/voices/test_voice") + def test_upload_voice_rejects_builtin_name(self, client): + """A name colliding with a built-in or precomputed voice is rejected.""" + handler = client.app.state.openai_serving_speech + handler._adapter.capabilities.supported_speakers = {"vivian"} + handler._adapter.capabilities.precomputed_speakers = {"ono_anna": {}} + + files = {"audio_sample": ("test.wav", b"fake audio content" * 1000, "audio/wav")} + for reserved in ("Vivian", "ono_anna"): + response = client.post("/v1/audio/voices", files=files, data={"consent": "c1", "name": reserved}) + assert response.status_code == 400 + assert "reserved" in response.json()["detail"] + assert not handler.uploaded_speakers + + def test_upload_voice_overwrites_same_name_by_default(self, client): + """Default policy: re-registering an uploaded name replaces it in place.""" + files = {"audio_sample": ("test.wav", b"first take", "audio/wav")} + data = {"consent": "c1", "name": "test_voice_ow"} + assert client.post("/v1/audio/voices", files=files, data=data).status_code == 200 + + files = {"audio_sample": ("test.wav", b"second take, longer content", "audio/wav")} + response = client.post("/v1/audio/voices", files=files, data=data) + assert response.status_code == 200 + assert response.json()["voice"]["file_size"] == len(b"second take, longer content") + client.delete("/v1/audio/voices/test_voice_ow") + + def test_upload_voice_immutable_policy_requires_delete(self, client): + """VLLM_OMNI_SPEAKER_REGISTRATION_POLICY=immutable rejects duplicates until deleted.""" + handler = client.app.state.openai_serving_speech + handler._registration_policy = "immutable" + try: + files = {"audio_sample": ("test.wav", b"fake audio content" * 1000, "audio/wav")} + data = {"consent": "c1", "name": "test_voice_imm"} + assert client.post("/v1/audio/voices", files=files, data=data).status_code == 200 + + response = client.post("/v1/audio/voices", files=files, data=data) + assert response.status_code == 400 + assert "immutable" in response.json()["detail"] + + assert client.delete("/v1/audio/voices/test_voice_imm").status_code == 200 + assert client.post("/v1/audio/voices", files=files, data=data).status_code == 200 + finally: + handler._registration_policy = "overwrite" + client.delete("/v1/audio/voices/test_voice_imm") + + def test_invalid_registration_policy_fails_startup(self, mocker, monkeypatch, tmp_path): + """An unknown VLLM_OMNI_SPEAKER_REGISTRATION_POLICY is a configuration error, not a silent fallback.""" + monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path)) + monkeypatch.setenv("VLLM_OMNI_SPEAKER_REGISTRATION_POLICY", "append") + engine_client = mocker.MagicMock() + engine_client.default_sampling_params_list = [{}] + with pytest.raises(ValueError, match="VLLM_OMNI_SPEAKER_REGISTRATION_POLICY"): + OmniOpenAIServingSpeech( + engine_client=engine_client, models=mocker.MagicMock(), request_logger=mocker.MagicMock() + ) + + def test_restored_upload_shadowing_builtin_is_dropped(self, client): + """A pre-guard upload restored under a built-in name must not keep shadowing it.""" + handler = client.app.state.openai_serving_speech + handler._adapter.capabilities.supported_speakers = {"vivian"} + handler._adapter.capabilities.precomputed_speakers = {} + handler.uploaded_speakers["vivian"] = {"name": "vivian", "file_path": "/tmp/vivian.safetensors"} + handler.uploaded_speakers["keep_me"] = {"name": "keep_me", "file_path": "/tmp/keep_me.safetensors"} + + handler._drop_shadowing_uploads() + + assert "vivian" not in handler.uploaded_speakers + assert "keep_me" in handler.uploaded_speakers + handler.uploaded_speakers.pop("keep_me") + def test_upload_voice_with_ref_text(self, client, tmp_path): """Test voice upload with ref_text enables in-context cloning.""" audio_content = b"fake audio content" * 1000 diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index fd3845eda..c3871bace 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -77,6 +77,7 @@ _PUBLIC_OMNI = ( "VLLM_OMNI_SERVER_STORAGE__PATH", "VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL", "VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP", + "VLLM_OMNI_SPEAKER_REGISTRATION_POLICY", "VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "VLLM_OMNI_STORAGE_PATH", "VLLM_OMNI_USE_QUACK_FP8", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index d3c23ca39..a32eeeac6 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -246,6 +246,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): except ValueError: logger.warning("Invalid SPEAKER_MAX_UPLOADED=%r; using default 1000", _raw_cap) self._max_uploaded_speakers = 1000 + _policy = os.environ.get("VLLM_OMNI_SPEAKER_REGISTRATION_POLICY", "overwrite").lower() + if _policy not in ("overwrite", "immutable"): + raise ValueError( + f"Invalid VLLM_OMNI_SPEAKER_REGISTRATION_POLICY={_policy!r}; expected 'overwrite' or 'immutable'." + ) + self._registration_policy = _policy self.uploaded_speakers: dict[str, dict[str, Any]] = {} self._ref_audio_data_url_cache: dict[str, str] = {} self._ref_audio_resolve_cache: OrderedDict[str, tuple[np.ndarray, int, int, str]] = OrderedDict() @@ -410,6 +416,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): adapter = self._adapter if adapter is not None: adapter.load_capabilities() + self._drop_shadowing_uploads() available_speakers = self._get_available_speakers() logger.info("Loaded %d supported speakers: %s", len(available_speakers), sorted(available_speakers)) @@ -448,6 +455,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ctx = SpeechServingContext(server=self, engine_client=self.engine_client) self._adapter = adapter_cls(ctx) self._adapter.load_capabilities() + self._drop_shadowing_uploads() return self._adapter def _uses_native_speed_control(self) -> bool: @@ -803,6 +811,39 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): f"the cap via SPEAKER_MAX_UPLOADED." ) + def _drop_shadowing_uploads(self) -> None: + """Uploads are restored from disk before the adapter exists, so a name + registered under a built-in speaker before the collision guard would + keep shadowing it across restarts. Drop such entries from the registry; + the file stays on disk for the operator to remove.""" + caps = self._adapter.capabilities if self._adapter is not None else None + if caps is None: + return + for voice_name_lower in list(self.uploaded_speakers): + if voice_name_lower in caps.supported_speakers or voice_name_lower in caps.precomputed_speakers: + info = self.uploaded_speakers.pop(voice_name_lower) + logger.warning( + "Uploaded voice %r shadows a built-in speaker and is ignored; remove %s to clear this warning.", + voice_name_lower, + info.get("file_path"), + ) + + def _check_registration_allowed(self, voice_name_lower: str, name: str) -> None: + """Reject names that would shadow a built-in voice or, under the + immutable policy, silently overwrite an existing upload.""" + caps = self._adapter.capabilities if self._adapter is not None else None + if caps is not None and ( + voice_name_lower in caps.supported_speakers or voice_name_lower in caps.precomputed_speakers + ): + raise ValueError( + f"Voice name '{name}' is reserved by a built-in speaker of this model; choose a different name." + ) + if self._registration_policy == "immutable" and voice_name_lower in self.uploaded_speakers: + raise ValueError( + f"Voice '{name}' already exists and VLLM_OMNI_SPEAKER_REGISTRATION_POLICY is 'immutable'; " + f"delete it first via DELETE /v1/audio/voices/{name}." + ) + def _evict_existing_upload(self, voice_name_lower: str, name: str) -> None: """Drop an existing upload with this name so the caller can re-register it.""" if voice_name_lower not in self.uploaded_speakers: @@ -885,6 +926,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): async with self._upload_lock: voice_name_lower = name.lower() + self._check_registration_allowed(voice_name_lower, name) self._evict_existing_upload(voice_name_lower, name) self._check_upload_cap() @@ -1005,6 +1047,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): async with self._upload_lock: voice_name_lower = name.lower() + self._check_registration_allowed(voice_name_lower, name) self._evict_existing_upload(voice_name_lower, name) self._check_upload_cap()