From fe02e8b81e92ef2a45b6d29f307ad5df6599b6b9 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Sat, 1 Aug 2026 10:57:12 +0800 Subject: [PATCH 01/10] [CI][NPU] Add MiniCPM-o 4.5 dependencies (#5660) Signed-off-by: Weiming Liao --- docker/Dockerfile.npu | 4 ++++ docker/Dockerfile.npu.a3 | 4 ++++ docker/Dockerfile.npu.ci | 1 + docker/Dockerfile.npu.ci.a3 | 1 + 4 files changed, 10 insertions(+) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 2dc94c2b4..05a780542 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -26,6 +26,10 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation +RUN python3 -m pip install -v --no-deps \ + "stepaudio2-minicpmo>=0.1.1" \ + "step-audio2>=1.0.0" + ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index 246467fbe..d226a505f 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -26,6 +26,10 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation +RUN python3 -m pip install -v --no-deps \ + "stepaudio2-minicpmo>=0.1.1" \ + "step-audio2>=1.0.0" + ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index c56589e9e..fa2c0d45f 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -35,6 +35,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ "psutil>=7.2.0" \ "soundfile>=0.13.1" \ "pyttsx3>=2.99" \ + "opencc==1.4.1" \ "opencv-python>=4.12.0.88" \ "imageio[ffmpeg]>=0.6.0" \ "jiwer>=3.0.0" \ diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 index 6a8dd4146..b605567c5 100644 --- a/docker/Dockerfile.npu.ci.a3 +++ b/docker/Dockerfile.npu.ci.a3 @@ -35,6 +35,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ "psutil>=7.2.0" \ "soundfile>=0.13.1" \ "pyttsx3>=2.99" \ + "opencc==1.4.1" \ "opencv-python>=4.12.0.88" \ "imageio[ffmpeg]>=0.6.0" \ "jiwer>=3.0.0" \ From 0a033592245706f6f7cb6b763c6262ca6400e8d5 Mon Sep 17 00:00:00 2001 From: NATURE Date: Sat, 1 Aug 2026 11:01:06 +0800 Subject: [PATCH 02/10] [Bugfix][MiniCPM-o] Stabilize single-GPU audio generation (#5637) Signed-off-by: natureofnature --- tests/config/test_config_factory.py | 26 +++++++++++++++++++ tests/helpers/minicpmo_4_5_duplex.py | 13 ++++++++++ vllm_omni/deploy/minicpmo_4_5.yaml | 7 +++++ .../minicpmo_4_5_3gpu_stage1_replicas.yaml | 7 +++++ .../minicpmo_4_5_4gpu_stage1_replicas.yaml | 7 +++++ .../models/minicpmo_4_5/batched_token2wav.py | 21 +++++++++++++++ 6 files changed, 81 insertions(+) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index a811868c1..33dcd0551 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2214,6 +2214,32 @@ class TestPlatformOverrides: assert rocm.stages[0].enforce_eager is None assert rocm.stages[1].enforce_eager is True + def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self): + pipeline = resolve_pipeline_config("minicpmo_4_5") + assert isinstance(pipeline, PipelineConfig) + deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml")) + + cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda") + cuda_stages = merge_pipeline_deploy(pipeline, cuda) + assert cuda_stages[1].yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3 + + # The CUDA memory budget must not leak into the existing NPU profile. + npu = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu") + npu_stages = merge_pipeline_deploy(pipeline, npu) + assert "kv_cache_memory_bytes" not in npu_stages[1].yaml_engine_args + + # Multi-GPU replica profiles place Talker on dedicated devices and + # preserve their previous automatic KV-cache sizing. + for filename in ( + "minicpmo_4_5_3gpu_stage1_replicas.yaml", + "minicpmo_4_5_4gpu_stage1_replicas.yaml", + ): + replica = _apply_platform_overrides( + load_deploy_config(Path(get_deploy_config_path(filename))), platform="cuda" + ) + replica_stages = merge_pipeline_deploy(pipeline, replica) + assert replica_stages[1].yaml_engine_args["kv_cache_memory_bytes"] is None + def test_npu_overrides(self): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): diff --git a/tests/helpers/minicpmo_4_5_duplex.py b/tests/helpers/minicpmo_4_5_duplex.py index dd32c0dcf..e19903e84 100644 --- a/tests/helpers/minicpmo_4_5_duplex.py +++ b/tests/helpers/minicpmo_4_5_duplex.py @@ -24,6 +24,19 @@ DEPLOY_CONFIG = modify_stage_config( 1: {"kv_cache_memory_bytes": 512 * 1024 * 1024}, 2: {"kv_cache_memory_bytes": 256 * 1024 * 1024}, }, + # Platform overrides are applied after ordinary stage settings. Keep + # the constrained duplex test's Talker budget ahead of the base CUDA + # profile's single-GPU default. + "platforms": { + "cuda": { + "stages": [ + { + "stage_id": 1, + "kv_cache_memory_bytes": 512 * 1024 * 1024, + } + ] + } + }, }, ) CORE_DEPLOY_CONFIG = modify_stage_config( diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index abe8b5ad2..b3e211175 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -75,6 +75,13 @@ stages: detokenize: true platforms: + cuda: + stages: + - stage_id: 1 + # Four Talker streams need much less KV than automatic sizing. + # Leave single-GPU headroom for concurrent Stage 2 diffusion batches. + kv_cache_memory_bytes: 2147483648 # 2 GiB + npu: stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml index 17cd014e9..7a0513167 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml @@ -15,3 +15,10 @@ stages: devices: "1,2" num_replicas: 2 async_scheduling: false + +platforms: + cuda: + stages: + - stage_id: 1 + # Stage 1 has dedicated GPUs in this profile; retain automatic sizing. + kv_cache_memory_bytes: null diff --git a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml index e4fe45922..da28456fe 100644 --- a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml @@ -15,3 +15,10 @@ stages: devices: "1,2,3" num_replicas: 3 async_scheduling: false + +platforms: + cuda: + stages: + - stage_id: 1 + # Stage 1 has dedicated GPUs in this profile; retain automatic sizing. + kv_cache_memory_bytes: null diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index f0e07e50b..f2d96be0d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -65,6 +65,27 @@ class BatchedToken2Wav(nn.Module): self._token2wav = token2wav self.flow = token2wav.flow self.hift = token2wav.hift + hift_parameter = next(self.hift.parameters(), None) + if hift_parameter is not None and hift_parameter.device.type == "cuda": + # Prime the CUDA state used by HiFT during backend construction. + # Otherwise, the first live audio chunk can fail when async stages + # share one GPU. + device = hift_parameter.device + dtype = hift_parameter.dtype + mel_channels = int(self.hift.conv_pre.in_channels) + with ( + torch.inference_mode(), + torch.random.fork_rng(devices=[device]), + _autocast_disabled(device), + ): + # 50 mel frames match the default first streamed vocoder chunk. + speech, source = self.hift( + torch.zeros((1, mel_channels, 50), device=device, dtype=dtype), + torch.zeros((1, 1, 0), device=device, dtype=dtype), + ) + torch.accelerator.synchronize(device) + del speech, source + torch.accelerator.empty_cache() self.float16 = bool(token2wav.float16) self.n_timesteps = int(token2wav.n_timesteps) self.mel_cache_len = int(token2wav.mel_cache_len) From 8cb79308c68d8ada6ea07855811ff533c6900965 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Harenome=20=E5=AE=89=E4=BB=81=20Ranaivoarivony-Razanajato?= <676776+harenome@users.noreply.github.com> Date: Sat, 1 Aug 2026 05:19:21 +0200 Subject: [PATCH 03/10] diffusion: Fix silent override of explicit `guidance_scale=0` (#4999) Signed-off-by: Harenome Ranaivoarivony-Razanajato --- tests/diffusion/test_diffusion_request.py | 37 +++++++++++++++++++ .../models/flux/pipeline_flux_kontext.py | 4 +- .../models/glm_image/pipeline_glm_image.py | 2 +- .../pipeline_soulx_singer_base.py | 2 +- vllm_omni/diffusion/request.py | 10 ++--- vllm_omni/inputs/data.py | 2 +- 6 files changed, 48 insertions(+), 9 deletions(-) diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py index ffaab3c53..259492c1e 100644 --- a/tests/diffusion/test_diffusion_request.py +++ b/tests/diffusion/test_diffusion_request.py @@ -72,3 +72,40 @@ def test_tp_seed_same_across_ranks_and_varies_across_requests(): # Seeds must vary across requests (non-determinism preserved). assert len(set(seeds)) == n_requests, f"Expected {n_requests} unique seeds but got {len(set(seeds))}: {seeds}" + + +def _request_with_guidance(**sampling_kwargs) -> OmniDiffusionRequest: + return OmniDiffusionRequest( + prompt={"prompt": "a cup of coffee on a table"}, + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1, **sampling_kwargs), + request_id="request-test", + ) + + +def test_omitted_guidance_scale_is_unset_and_resolved_to_one(): + # No guidance_scale passed: the caller did not set it, so it must NOT count as + # "provided" (each pipeline applies its own default), and __post_init__ resolves the + # None "unset" sentinel to a concrete float so no consumer ever observes None. + req = _request_with_guidance() + + assert req.sampling_params.guidance_scale_provided is False + assert req.sampling_params.guidance_scale == 1.0 + + +def test_explicit_zero_guidance_scale_is_honored(): + # Regression (#4998): an explicit guidance_scale=0 (disable CFG) must be honored, + # not mistaken for "unset" and overridden by the pipeline's model default. This is + # the sentinel collision that the None default + ``is not None`` check fixes. + req = _request_with_guidance(guidance_scale=0.0) + + assert req.sampling_params.guidance_scale_provided is True + assert req.sampling_params.guidance_scale == 0.0 + + +@pytest.mark.parametrize("value", [0.999, 1.0, 5.0]) +def test_explicit_nonzero_guidance_scale_is_honored(value): + # Any explicitly-set value (truthy or not) is honored as provided and preserved. + req = _request_with_guidance(guidance_scale=value) + + assert req.sampling_params.guidance_scale_provided is True + assert req.sampling_params.guidance_scale == value diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py index 5aff911c6..7edc4751a 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py @@ -640,7 +640,9 @@ class FluxKontextPipeline( width = req.sampling_params.width or width num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps sigmas = req.sampling_params.sigmas or sigmas - guidance_scale = req.sampling_params.guidance_scale or guidance_scale + guidance_scale = ( + req.sampling_params.guidance_scale if req.sampling_params.guidance_scale_provided else guidance_scale + ) num_images_per_prompt = ( req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py index 9f3d29400..6ceb5959b 100644 --- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py +++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py @@ -733,7 +733,7 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsCompon ) width = ar_width or req.sampling_params.width or img_width or self.default_sample_size * self.vae_scale_factor num_inference_steps = req.sampling_params.num_inference_steps or 50 - guidance_scale = req.sampling_params.guidance_scale or 1.5 + guidance_scale = req.sampling_params.guidance_scale if req.sampling_params.guidance_scale_provided else 1.5 # Ensure dimensions are multiples of vae_scale_factor * patch_size multiple_of = self.vae_scale_factor * self._patch_size diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py index a48b5236d..0961e6916 100644 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py @@ -353,7 +353,7 @@ class FlowMatchingAudioPipeline( sampling_params.extra_args = extra_args num_inference_steps = sampling_params.num_inference_steps or 32 - guidance_scale = sampling_params.guidance_scale or 3.0 + guidance_scale = sampling_params.guidance_scale if sampling_params.guidance_scale_provided else 3.0 generator = self._resolve_diffusion_generator(sampling_params) with self._stage_timer("consume_payload"): diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 0edb30d23..9ea879212 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -42,11 +42,11 @@ class OmniDiffusionRequest: self.sampling_params.seed = random.randint(0, 2**31 - 1) # Detect whether user explicitly provided guidance_scale. - # The sentinel default is 0.0 (false-like); any truthy value means - # the caller set it intentionally. We must resolve this BEFORE - # auto-filling guidance_scale_2, otherwise the sentinel leaks into - # guidance_scale_2. - if self.sampling_params.guidance_scale: + # The unset sentinel is None (see OmniDiffusionSamplingParams); any value + # the caller set -- including an explicit 0.0 (disable CFG) -- must be + # honored. Resolve before auto-filling guidance_scale_2, else the sentinel + # leaks into guidance_scale_2. + if self.sampling_params.guidance_scale is not None: self.sampling_params.guidance_scale_provided = True else: self.sampling_params.guidance_scale = 1.0 diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index f0344c994..5e6b5275e 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -263,7 +263,7 @@ class OmniDiffusionSamplingParams: # Scheduler parameters – ``None`` means "not explicitly set by the caller"; # each pipeline's ``forward()`` decides its own model-specific default. num_inference_steps: int | None = None - guidance_scale: float = 0.0 + guidance_scale: float | None = None guidance_scale_provided: bool = False guidance_scale_2: float | None = None guidance_scale_2_provided: bool = False From cd5daaea6a8e3a856ea57b24ae7c46ec70f9f5fb Mon Sep 17 00:00:00 2001 From: akshatvishu Date: Sat, 1 Aug 2026 09:24:17 +0530 Subject: [PATCH 04/10] [BugFix][Metrics] Fix per-stage time and token stats in metrics API (#4974) Signed-off-by: akshatvishu --- docs/design/metrics.md | 8 +- ...mni_generation_scheduler_update_session.py | 12 ++ .../test_chunk_transfer_adapter.py | 46 ++++ .../test_async_omni_engine_stage_init.py | 131 +++++++++++- tests/engine/test_orchestrator.py | 200 +++++++++++++++--- tests/engine/test_output_processor.py | 49 ++++- tests/engine/test_single_stage_mode.py | 11 +- .../entrypoints/test_stream_finish_reason.py | 100 ++++++++- .../core/sched/omni_generation_scheduler.py | 22 ++ .../chunk_transfer_adapter.py | 19 +- vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/engine/orchestrator.py | 20 +- vllm_omni/engine/stage_pool.py | 7 +- vllm_omni/engine/stage_runtime.py | 19 +- vllm_omni/entrypoints/omni_base.py | 1 + vllm_omni/entrypoints/openai/serving_chat.py | 22 +- vllm_omni/outputs/__init__.py | 4 + vllm_omni/outputs/output_processor.py | 24 ++- 18 files changed, 618 insertions(+), 78 deletions(-) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 09f0907df..6e4946f4e 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -148,7 +148,13 @@ Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, retur `OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead. -The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second. +The orchestrator side does not add its own throttle on top. Scheduler gauges are +recorded when `raw_outputs.scheduler_stats is not None` (i.e. this replica's +scheduler passed its own 1 Hz gate), while output-batch request/token metrics +are recorded from `IterationStats` whenever a non-empty output batch is +processed. A previous global `_last_stats_ts` on the orchestrator was removed +because it starved every replica other than the first to emit within each +second. ## Metric Definitions diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index 5cd829925..922cc5eb6 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py @@ -21,6 +21,7 @@ import vllm_omni # noqa: F401 - import for side effects (patch vLLM) from vllm.sampling_params import SamplingParams from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.engine import EngineCoreEventType +from vllm.v1.metrics.stats import PrefillStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler @@ -93,6 +94,17 @@ def _make_update(**kwargs) -> StreamingUpdate: return StreamingUpdate(**defaults) +def test_generation_scheduler_records_prefill_stats_for_metrics() -> None: + request = _make_request(prompt_token_ids=[1, 2, 3, 4]) + request.prefill_stats = PrefillStats() + + OmniGenerationScheduler._record_prefill_stats(request) + + assert request.prefill_stats.num_prompt_tokens == 4 + assert request.prefill_stats.num_computed_tokens == 4 + assert request.prefill_stats.num_cached_tokens == 0 + + def test_resumable_generation_stop_marks_segment_boundary() -> None: session = _make_request(request_id="req-generation-segment") session.status = RequestStatus.RUNNING diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 97b1618b5..e1a1e36f9 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -10,8 +10,10 @@ import pytest import torch from pytest_mock import MockerFixture from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler +from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats from vllm.v1.request import RequestStatus +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase @@ -37,8 +39,10 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None) external_req_id=external_req_id or req_id, status=status, prompt_token_ids=[], + num_prompt_tokens=0, num_computed_tokens=0, num_output_placeholders=0, + prefill_stats=None, additional_information=None, is_finished=lambda: status == RequestStatus.FINISHED_STOPPED, ) @@ -485,6 +489,48 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt assert request.additional_information["meta"]["finished"].item() is True +def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="generation") + request = _req("req-later-chunk", RequestStatus.WAITING, external_req_id="ext-later-chunk") + request.prefill_stats = PrefillStats() + adapter.request_ids_mapping[request.request_id] = request.external_req_id + + connector.get.return_value = ( + { + "codes": {"audio": torch.tensor([7, 8], dtype=torch.long)}, + "meta": {"finished": torch.tensor(False, dtype=torch.bool)}, + }, + 8, + ) + assert adapter._poll_single_request(request) is True + OmniGenerationScheduler._record_prefill_stats(request) + first_chunk_stats = request.prefill_stats + request.prefill_stats = None + + connector.get.return_value = ( + { + "codes": {"audio": torch.tensor([9, 10, 11], dtype=torch.long)}, + "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, + }, + 8, + ) + assert adapter._poll_single_request(request) is True + assert isinstance(request.prefill_stats, PrefillStats) + OmniGenerationScheduler._record_prefill_stats(request) + second_chunk_stats = request.prefill_stats + + assert first_chunk_stats is not None + assert first_chunk_stats.num_prompt_tokens == 2 + assert second_chunk_stats is not None + assert second_chunk_stats.num_prompt_tokens == 3 + + prompt_token_stats = PromptTokenStats() + prompt_token_stats.update_from_output(first_chunk_stats) + prompt_token_stats.update_from_output(second_chunk_stats) + assert prompt_token_stats.total == 5 + assert prompt_token_stats.computed == 5 + + def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter): adapter, _ = build_adapter(stage_id=1, connector_extra={"role": "sender"}) request = _req("req-1", RequestStatus.WAITING) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index c366fb313..67765aba6 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -172,6 +172,41 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): ) +def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch): + import vllm_omni.engine.async_omni_engine as engine_mod + + engine = object.__new__(AsyncOmniEngine) + engine.stage_configs = [types.SimpleNamespace()] + engine.model = "dummy-model" + engine.config_path = "dummy-config" + engine.single_stage_mode = False + engine.diffusion_batch_size = 1 + engine.async_chunk = False + engine.tokenizer = None + engine._single_stage_id_filter = None + engine._omni_master_address = None + engine._omni_master_port = None + engine._omni_dp_size_local = 1 + engine._omni_heartbeat_timeout = 30.0 + engine._omni_lb_policy = "random" + engine.request_queue = types.SimpleNamespace() + engine._log_stats = True + + captured: dict[str, object] = {} + runtime = types.SimpleNamespace(stage_pools=[], initialize=lambda: None) + + def _capture_create_stage_runtime(**kwargs): + captured.update(kwargs) + return runtime + + monkeypatch.setattr(engine_mod, "create_stage_runtime", _capture_create_stage_runtime) + + engine._initialize_stages(stage_init_timeout=7) + + assert captured["stage_init_timeout"] == 7 + assert captured["log_stats"] is True + + def test_compute_replica_layout_splits_diffusion_devices_by_world_size(): stage_cfg = types.SimpleNamespace( stage_id=0, @@ -342,6 +377,95 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch): assert runtime.stage_pools[1].output_processor is stage1_output_processor +def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + log_stats=True, + ) + cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) + plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg).replicas[0] + plan.engine_args_dict = {} + + captured: dict[str, object] = {} + addresses = types.SimpleNamespace( + inputs=["tcp://127.0.0.1:1"], + outputs=["tcp://127.0.0.1:2"], + frontend_stats_publish_address=None, + ) + resources = types.SimpleNamespace(manager=object(), coordinator=None, addresses=addresses) + stage_client = types.SimpleNamespace() + + @contextlib.contextmanager + def _capture_launch_stage_replica(**kwargs): + captured["log_stats"] = kwargs["log_stats"] + yield resources + + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: []) + monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica) + monkeypatch.setattr( + runtime_mod.StageEngineCoreClientBase, + "make_async_mp_client", + lambda **kwargs: (captured.__setitem__("client_log_stats", kwargs["log_stats"]) or stage_client), + ) + + assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client + assert captured["log_stats"] is True + assert captured["client_log_stats"] is True + + +def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + log_stats=True, + ) + cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) + stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg) + stage_client = types.SimpleNamespace( + stage_type="llm", + is_comprehension=False, + final_output=True, + final_output_type=None, + default_sampling_params=types.SimpleNamespace(), + ) + captured: dict[str, object] = {} + output_processor = object() + + def _capture_build_llm_stage_output_processor(plan, stage_vllm_config, *, log_stats=False): + captured["plan"] = plan + captured["stage_vllm_config"] = stage_vllm_config + captured["log_stats"] = log_stats + return output_processor + + monkeypatch.setattr( + runtime_mod, + "build_llm_stage_output_processor", + _capture_build_llm_stage_output_processor, + ) + + pools = runtime._assemble_stage_pools([stage_plan], {0: [stage_client]}) + + assert pools[0].output_processor is output_processor + assert captured == { + "plan": stage_plan, + "stage_vllm_config": cfg, + "log_stats": True, + } + + def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch): import vllm_omni.engine.stage_runtime as runtime_mod @@ -537,10 +661,11 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag stage_configs=[], model="dummy-model", config_path="dummy-config", - stage_init_timeout=302, + stage_init_timeout=1, diffusion_batch_size=1, async_chunk=False, ) + stage_init_timeout = 302 fake_vllm_config = types.SimpleNamespace() fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None) @@ -587,14 +712,14 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag ) try: - runtime._initialize_local_llm_replica(plan, 302) + runtime._initialize_local_llm_replica(plan, stage_init_timeout) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) else: os.environ[device_env_var] = prev_device_env - assert captured_timeout == 302 + assert captured_timeout == stage_init_timeout def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path): diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index c058ed907..b6c4f8509 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -15,6 +15,7 @@ import pytest from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams from vllm.v1.engine.exceptions import EngineDeadError +from vllm.v1.metrics.stats import IterationStats from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -145,7 +146,7 @@ class FakeStageClient: try: return self._engine_core_outputs.get_nowait() except queue.Empty: - return SimpleNamespace(outputs=[]) + return SimpleNamespace(outputs=[], scheduler_stats=None, finished_requests=None) def get_diffusion_output_nowait(self): try: @@ -253,12 +254,40 @@ class FakeOutputProcessor: return None +class IterationStatsOutputProcessor(FakeOutputProcessor): + def process_outputs(self, *_args, **kwargs): + iteration_stats = kwargs.get("iteration_stats") + if iteration_stats is None and len(_args) >= 3: + iteration_stats = _args[2] + assert iteration_stats is not None + iteration_stats.num_generation_tokens += 3 + return super().process_outputs(*_args, **kwargs) + + +class RecordingOutputProcessor(FakeOutputProcessor): + def __init__(self, *, request_outputs: list[object] | None = None) -> None: + super().__init__(request_outputs=request_outputs) + self.process_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + + def process_outputs(self, *args, **kwargs): + self.process_calls.append((args, kwargs)) + return super().process_outputs(*args, **kwargs) + + +class RecordingStatLogger: + def __init__(self) -> None: + self.records: list[tuple[Any, Any, int]] = [] + + def record(self, scheduler_stats, iteration_stats, *, engine_idx: int = 0) -> None: + self.records.append((scheduler_stats, iteration_stats, engine_idx)) + + def _sampling_params(max_tokens: int = 4) -> SamplingParams: return SamplingParams(max_tokens=max_tokens) def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace: - return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None) + return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None) def _build_request_output( @@ -329,6 +358,7 @@ def _build_harness( output_processors: list[object] | None = None, stage_vllm_configs: list[object] | None = None, async_chunk: bool = False, + log_stats: bool = False, stage_pools: list[StagePool] | None = None, ) -> OrchestratorFixture: """Build an Orchestrator test harness. @@ -364,6 +394,7 @@ def _build_harness( rpc_async_queue=rpc_queue.async_q, stage_pools=stage_pools, async_chunk=async_chunk, + log_stats=log_stats, ) ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue)) await orchestrator.run() @@ -2137,38 +2168,143 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator await _shutdown_orchestrator(orchestrator_fixture) -def test_orchestrator_does_not_re_introduce_global_stats_throttle() -> None: - """Regression: each (stage, replica) must independently publish its wrapped - vllm:* stats when its scheduler emits non-None scheduler_stats. - - A previous version of Orchestrator carried a global self._last_stats_ts / - _stats_interval_s gate around _stat_logger.record(). Because - OmniSchedulerMixin.make_stats() already throttles at 1 Hz per scheduler - (one per (stage, replica)), the extra global gate starved every replica - other than the first to emit within each second — their {stage, replica} - gauges/counters went stale. - - The fix removed the global gate entirely; the only signal needed is - 'this replica's scheduler emitted non-None scheduler_stats'. This test - fails loudly if someone reintroduces the global throttle. - """ - import inspect - - from vllm_omni.engine.orchestrator import Orchestrator - - source = inspect.getsource(Orchestrator) - assert "_last_stats_ts" not in source, ( - "Orchestrator must not gate stat recording on a global timestamp. " - "OmniSchedulerMixin.make_stats() already throttles per scheduler " - "(per (stage, replica)); an outer global gate starves all but the " - "first replica to emit within each 1s window." +@pytest.mark.asyncio +async def test_orchestrator_records_iteration_stats_without_scheduler_stats(orchestrator_factory) -> None: + stage0 = FakeStageClient(stage_type="llm", final_output=True) + processor = IterationStatsOutputProcessor( + request_outputs=[_build_request_output("req-stats", token_ids=[1], finished=True)] ) - assert "_stats_interval_s" not in source - assert "raw_outputs.scheduler_stats is not None" in source, ( - "Orchestrator must gate stat recording solely on " - "raw_outputs.scheduler_stats being non-None — the per-scheduler 1Hz " - "throttle in OmniSchedulerMixin.make_stats() is the only gate needed." + orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True) + stat_logger = RecordingStatLogger() + orchestrator_fixture.orchestrator._stat_logger = stat_logger + request = SimpleNamespace(request_id="req-stats", prompt_token_ids=[1, 2]) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-stats", + prompt=request, + original_prompt={"prompt": "hello"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + + stage0.push_engine_core_outputs(_engine_core_outputs("raw-without-scheduler-stats", 1.0)) + + await _wait_for(lambda: len(stat_logger.records) == 1) + scheduler_stats, iteration_stats, engine_idx = stat_logger.records[0] + assert scheduler_stats is None + assert iteration_stats is not None + assert iteration_stats.num_generation_tokens == 3 + assert engine_idx == 0 + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_orchestrator_records_scheduler_stats_without_outputs(orchestrator_factory) -> None: + stage0 = FakeStageClient(stage_type="llm", final_output=True) + processor = RecordingOutputProcessor() + orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True) + stat_logger = RecordingStatLogger() + orchestrator_fixture.orchestrator._stat_logger = stat_logger + scheduler_stats = SimpleNamespace(num_running_reqs=0, kv_cache_usage=0.5) + + try: + stage0.push_engine_core_outputs( + SimpleNamespace( + outputs=[], + timestamp=1.0, + scheduler_stats=scheduler_stats, + finished_requests=None, + ) + ) + + await _wait_for(lambda: len(stat_logger.records) == 1) + recorded_scheduler_stats, iteration_stats, engine_idx = stat_logger.records[0] + assert recorded_scheduler_stats is scheduler_stats + assert iteration_stats is None + assert engine_idx == 0 + assert processor.process_calls[0][0][2] is None + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_orchestrator_does_not_build_iteration_stats_for_finished_only_batch(orchestrator_factory) -> None: + stage0 = FakeStageClient(stage_type="llm", final_output=True) + processor = RecordingOutputProcessor() + orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True) + stat_logger = RecordingStatLogger() + orchestrator_fixture.orchestrator._stat_logger = stat_logger + + try: + stage0.push_engine_core_outputs( + SimpleNamespace( + outputs=[], + timestamp=1.0, + scheduler_stats=None, + finished_requests={"req-finished"}, + ) + ) + + await _wait_for(lambda: bool(processor.process_calls)) + assert processor.process_calls[0][0][2] is None + assert stat_logger.records == [] + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_orchestrator_does_not_build_iteration_stats_without_stat_logger(orchestrator_factory) -> None: + stage0 = FakeStageClient(stage_type="llm", final_output=True) + processor = RecordingOutputProcessor() + orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True) + orchestrator_fixture.orchestrator._stat_logger = None + + try: + stage0.push_engine_core_outputs(_engine_core_outputs("raw-output", 1.0)) + + await _wait_for(lambda: bool(processor.process_calls)) + assert processor.process_calls[0][0][2] is None + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_stage_pool_preserves_none_iteration_stats() -> None: + client = FakeStageClient(stage_type="llm", final_output=True) + processor = RecordingOutputProcessor() + pool = StagePool( + 0, + [client], + output_processor=processor, + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), ) + raw_outputs = SimpleNamespace(outputs=["raw"], timestamp=1.0, scheduler_stats=None) + + await pool.process_llm_raw_outputs(0, raw_outputs, iteration_stats=None) + + assert processor.process_calls[0][0][2] is None + + +@pytest.mark.asyncio +async def test_stage_pool_process_llm_raw_outputs_mutates_iteration_stats() -> None: + client = FakeStageClient(stage_type="llm", final_output=True) + processor = IterationStatsOutputProcessor() + pool = StagePool( + 0, + [client], + output_processor=processor, + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + raw_outputs = SimpleNamespace(outputs=["raw"], timestamp=1.0, scheduler_stats=None) + iteration_stats = IterationStats() + + await pool.process_llm_raw_outputs(0, raw_outputs, iteration_stats=iteration_stats) + + assert iteration_stats.num_generation_tokens == 3 @pytest.mark.asyncio diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 81f0c8303..23ac601b3 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -10,9 +10,11 @@ import pytest import torch from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind -from vllm.v1.engine import FinishReason +from vllm.v1.engine import EngineCoreEvent, EngineCoreEventType, FinishReason from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor +from vllm.v1.metrics.stats import IterationStats, PrefillStats +from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.outputs import output_processor from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState @@ -616,3 +618,48 @@ def test_no_detokenizer_final_only(): result = s.make_request_output([], None, FinishReason.STOP, None) assert result is not None assert AUDIO in result.outputs[0].multimodal_output + + +def test_mm_only_outputs_update_iteration_stats(): + processor = MultimodalOutputProcessor( + tokenizer=None, + log_stats=True, + engine_core_output_type=AUDIO, + ) + state = _make_no_detok_state(RequestOutputKind.CUMULATIVE) + processor.request_states[state.request_id] = state + processor.external_req_ids[state.external_req_id].append(state.request_id) + + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=4, + num_local_cached_tokens=0, + num_external_cached_tokens=0, + ) + output = OmniEngineCoreOutput( + request_id=state.request_id, + new_token_ids=[10, 11], + finish_reason=FinishReason.STOP, + events=[ + EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0), + EngineCoreEvent(EngineCoreEventType.SCHEDULED, 2.0), + ], + prefill_stats=prefill_stats, + multimodal_output={AUDIO: torch.ones(1, 8)}, + ) + iteration_stats = IterationStats() + + processor.process_outputs( + [output], + engine_core_timestamp=3.0, + iteration_stats=iteration_stats, + ) + + assert iteration_stats.num_prompt_tokens == 4 + assert iteration_stats.num_generation_tokens == 2 + assert len(iteration_stats.time_to_first_tokens_iter) == 1 + assert len(iteration_stats.finished_requests) == 1 + finished = iteration_stats.finished_requests[0] + assert finished.finish_reason == FinishReason.STOP + assert finished.num_prompt_tokens == state.prompt_len + assert finished.num_generation_tokens == 2 diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index f90a290ca..8d70f0ed1 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -778,6 +778,7 @@ class TestSingleStageReplicaInitialization: single_stage_id_filter=None, omni_master_address="127.0.0.1", omni_master_port=26000, + log_stats=True, ) runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "llm"} @@ -811,10 +812,17 @@ class TestSingleStageReplicaInitialization: sentinel_client = SimpleNamespace() mock_connect = mocker.patch.object(runtime_mod, "connect_remote_engine_cores", side_effect=_fake_connect) + client_kwargs: dict[str, Any] = {} + + def _capture_make_async_mp_client(**kwargs): + client_kwargs.update(kwargs) + events.append("attach") + return sentinel_client + mocker.patch.object( StageEngineCoreClientBase, "make_async_mp_client", - side_effect=lambda **_: (events.append("attach"), sentinel_client)[1], + side_effect=_capture_make_async_mp_client, ) result = runtime._initialize_remote_replica(plan, stage_init_timeout=60) @@ -824,6 +832,7 @@ class TestSingleStageReplicaInitialization: assert mock_connect.call_args.kwargs["vllm_config"].parallel_config.data_parallel_size_local == 0 assert mock_connect.call_args.kwargs["stage_id"] == 1 assert mock_connect.call_args.kwargs["replica_id"] == 0 + assert client_kwargs["log_stats"] is True assert events == ["enter", "exit", "attach"] def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture): diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 5ba3fea68..3d6052045 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -17,6 +17,7 @@ Key invariants tested: """ import json +import time from unittest.mock import MagicMock import pytest @@ -80,28 +81,37 @@ def _make_audio_omni_output( request_id: str = "test-req", index: int = 0, num_prompt_tokens: int = 3, + stage_id: int | None = None, + replica_id: int | None = None, + audio_samples: int = 0, ) -> OmniRequestOutput: """Build an OmniRequestOutput for audio (no torch dependency).""" + completion = CompletionOutput( + index=index, + text="", + token_ids=[], + cumulative_logprob=0.0, + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + if audio_samples > 0: + completion.multimodal_output = { + "audio": [MagicMock(numel=MagicMock(return_value=audio_samples))], + "sr": 24000, + } res = RequestOutput( request_id=request_id, prompt="test", prompt_token_ids=list(range(num_prompt_tokens)), prompt_logprobs=None, - outputs=[ - CompletionOutput( - index=index, - text="", - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="stop", - stop_reason=None, - ) - ], + outputs=[completion], finished=True, ) return OmniRequestOutput( request_id=request_id, + stage_id=stage_id, + replica_id=replica_id, final_output_type="audio", request_output=res, finished=True, @@ -321,6 +331,74 @@ async def test_single_modality_audio_only_one_stop(): assert finish_reasons[-1] == "stop" +@pytest.mark.parametrize( + ("output_replica_id", "pool_replica_id", "expected_replica_id", "pool_lookup"), + [ + (7, None, 7, False), + (None, 5, 5, True), + ], +) +@pytest.mark.asyncio +async def test_streaming_audio_metrics_resolve_replica_id( + monkeypatch, + output_replica_id, + pool_replica_id, + expected_replica_id, + pool_lookup, +): + import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod + from vllm_omni.entrypoints.client_request_state import ClientRequestState + + serving_chat = _build_serving_chat() + request = _make_request(modalities=["audio"]) + req_state = ClientRequestState( + request_id="internal-req", + external_request_id="test-req", + ) + req_state.request_arrival_ts = time.time() - 1.0 + serving_chat.engine_client.request_states = {"internal-req": req_state} + serving_chat.engine_client.mod_metrics = object() + serving_chat.engine_client.engine = MagicMock() + serving_chat.engine_client.engine.stage_pools = [MagicMock(), MagicMock(), MagicMock()] + serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.return_value = pool_replica_id + + first_packet_calls = [] + finalize_calls = [] + + def fake_observe_audio_first_packet(*args, **kwargs): + first_packet_calls.append(kwargs) + + def fake_observe_audio_streaming_finalize(*args, **kwargs): + finalize_calls.append(kwargs) + + monkeypatch.setattr(serving_chat_mod, "observe_audio_first_packet", fake_observe_audio_first_packet) + monkeypatch.setattr(serving_chat_mod, "observe_audio_streaming_finalize", fake_observe_audio_streaming_finalize) + + async def result_generator(): + yield _make_audio_omni_output(stage_id=2, replica_id=output_replica_id, audio_samples=2400) + + await _collect_stream( + serving_chat.chat_completion_stream_generator( + request=request, + result_generator=result_generator(), + request_id="test-req", + model_name="test-model", + conversation=[], + tokenizer=MagicMock(), + request_metadata=MagicMock(), + ) + ) + + if pool_lookup: + serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.assert_called_once_with("internal-req") + else: + serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.assert_not_called() + assert first_packet_calls[0]["stage_id"] == 2 + assert first_packet_calls[0]["replica_id"] == expected_replica_id + assert finalize_calls[0]["stage_id"] == 2 + assert finalize_calls[0]["replica_id"] == expected_replica_id + + # --------------------------------------------------------------------------- # Tests: fallback stop chunk when declared modality is not produced # --------------------------------------------------------------------------- diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 69c6c3c5c..47e6bb4d2 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -56,6 +56,24 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): ) self._latest_omni_connector_output: OmniConnectorOutput | None = None + @staticmethod + def _record_prefill_stats(request: Request) -> None: + """Mirror upstream first-prefill prompt accounting for generation stages. + + Async-chunk generation stages can grow ``prompt_token_ids`` after the + first schedule; until this path tracks final chunk lengths, this records + the first scheduled prompt snapshot. Keep this in sync with vLLM + Scheduler.schedule() if this fast path gains prefix-cache or external-KV + accounting. + """ + if request.prefill_stats is None: + return + request.prefill_stats.set( + num_prompt_tokens=request.num_prompt_tokens, + num_local_cached_tokens=0, + num_external_cached_tokens=0, + ) + def _handle_stopped_request(self, request: Request) -> bool: if ( request.resumable @@ -147,6 +165,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): break if self.log_stats: request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) + if num_computed_tokens == 0: + self._record_prefill_stats(request) req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens cached_prompt_token_ids[request.request_id] = request.prompt_token_ids @@ -213,6 +233,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self.running.append(request) if self.log_stats: request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) + if request.num_computed_tokens == 0: + self._record_prefill_stats(request) req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index cce5db472..d722adad2 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping from typing import Any import torch +from vllm.v1.metrics.stats import PrefillStats from vllm.v1.request import Request, RequestStatus from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload @@ -109,6 +110,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0) return max(0, num_computed - num_placeholders) + @staticmethod + def _refresh_generation_chunk_prefill_state(request: Request) -> None: + request.num_prompt_tokens = len(request.prompt_token_ids) + if getattr(request, "prefill_stats", None) is None: + request.prefill_stats = PrefillStats() + @classmethod def create_connector(cls, model_config: Any): connector_config = getattr(model_config, "stage_connector_config", None) @@ -249,15 +256,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): new_ids = payload_data.get("codes", {}).get("audio") has_tensor_codes = isinstance(new_ids, torch.Tensor) use_tensor_codes = has_tensor_codes and new_ids.ndim >= 2 + prompt_token_ids: list[int] if use_tensor_codes: - request.prompt_token_ids = [0] if new_ids.numel() > 0 else [] + prompt_token_ids = [0] if new_ids.numel() > 0 else [] elif has_tensor_codes: new_ids = new_ids.tolist() + prompt_token_ids = new_ids elif new_ids is None: new_ids = [] - request.prompt_token_ids = new_ids - if not use_tensor_codes: - request.prompt_token_ids = new_ids + prompt_token_ids = new_ids + else: + prompt_token_ids = new_ids + request.prompt_token_ids = prompt_token_ids prev_info = getattr(request, "additional_information", None) info = dict(prev_info) if isinstance(prev_info, dict) else {} for key, value in payload_data.items(): @@ -293,6 +303,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # chunk as ready, otherwise Stage1 can consume before the # first DAC frame arrives. return False + self._refresh_generation_chunk_prefill_state(request) # Mark as finished for consumption self._finished_load_reqs.add(req_id) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index e72fd01c9..454266b7a 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -331,6 +331,7 @@ class AsyncOmniEngine: omni_heartbeat_timeout=self._omni_heartbeat_timeout, omni_lb_policy=self._omni_lb_policy, request_queue=self.request_queue, + log_stats=self._log_stats, ) self._runtime.initialize() diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 963054688..0b4f17aae 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -362,6 +362,7 @@ class Orchestrator: self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools + self.log_stats = log_stats self._orch_monitor = create_orch_monitor( enabled=enable_orch_monitor, replica_sampler=self._sample_replica_metrics, @@ -446,7 +447,10 @@ class Orchestrator: already no-ops on ``scheduler_stats is None`` (which is what the upstream scheduler returns when its own log_stats is False), so this gate is mainly to keep the ``/metrics`` surface clean - when the user did not request stats. + when the user did not request stats. When stats are enabled, record() + must still receive IterationStats even on ticks where scheduler_stats + is None; upstream PrometheusStatLogger records token/request metrics + from iteration_stats independently of scheduler gauges. """ self._running_counter = running_counter self._transfer_emitter = transfer_emitter @@ -937,19 +941,17 @@ class Orchestrator: ) if req_state.streaming.enabled: await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state) - # OmniSchedulerMixin.make_stats() already throttles - # per-scheduler at 1 Hz, so raw_outputs.scheduler_stats - # being non-None means this replica passed its own gate. - # A second global throttle here would drop stats for - # other (stage, replica) pairs in the same 1s window. - record_stats = self._stat_logger is not None and raw_outputs.scheduler_stats is not None - iteration_stats = IterationStats() if record_stats else None + iteration_stats = ( + IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None + ) raw_output = await pool.process_llm_raw_outputs( replica_id, raw_outputs, iteration_stats=iteration_stats, ) - if record_stats: + if self._stat_logger is not None and ( + raw_outputs.scheduler_stats is not None or iteration_stats is not None + ): self._stat_logger.record( raw_outputs.scheduler_stats, iteration_stats, diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 820b28510..354032a7d 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -1081,7 +1081,11 @@ class StagePool: async def _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None: """Pull raw EngineCoreOutputs from a stage replica without processing.""" outputs = await client.get_output_async() - if not outputs.outputs: + # Keep scheduler-only / finished-only batches. Omni schedulers mirror + # upstream by emitting SchedulerStats on throttled ticks even when no + # request output is produced, and dropping those batches loses KV/queue + # gauges for that interval. + if not outputs.outputs and outputs.scheduler_stats is None and not outputs.finished_requests: return None return outputs @@ -1097,7 +1101,6 @@ class StagePool: return [] client = cast(StagePoolLLMClient, raw_client) processor = self.output_processor - iteration_stats = IterationStats() processed = processor.process_outputs( raw_outputs.outputs, raw_outputs.timestamp, diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 1f2024fd2..8f13a3643 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -126,6 +126,7 @@ class StageRuntime: diffusion_batch_size: int, async_chunk: bool, tokenizer: str | None = None, + log_stats: bool = False, ) -> None: self._stage_configs = stage_configs self._model = model @@ -134,6 +135,7 @@ class StageRuntime: self._diffusion_batch_size = diffusion_batch_size self._async_chunk = async_chunk self._tokenizer = tokenizer + self._log_stats = log_stats self._num_stages = len(stage_configs) # Populated by initialize() @@ -575,7 +577,7 @@ class StageRuntime: with launch_stage_replica( vllm_config=vllm_config, executor_class=executor_class, - log_stats=False, + log_stats=self._log_stats, stage_id=plan.metadata.stage_id, replica_id=plan.replica_id, stage_config=plan.stage_cfg, @@ -594,6 +596,7 @@ class StageRuntime: stage_client = StageEngineCoreClientBase.make_async_mp_client( vllm_config=vllm_config, executor_class=executor_class, + log_stats=self._log_stats, metadata=plan.metadata, client_addresses=self._client_addresses_from_zmq(resources.addresses), engine_manager=resources.manager, @@ -704,7 +707,11 @@ class StageRuntime: stage_vllm_config = plan.replicas[0].stage_vllm_config if stage_vllm_config is None: raise RuntimeError(f"Stage {plan.stage_id} is missing vllm_config") - output_processor = build_llm_stage_output_processor(plan, stage_vllm_config) + output_processor = build_llm_stage_output_processor( + plan, + stage_vllm_config, + log_stats=self._log_stats, + ) stage_pools.append( StagePool( @@ -742,10 +749,11 @@ class DistStageRuntime(StageRuntime): stage_init_timeout: int, diffusion_batch_size: int, async_chunk: bool, - tokenizer: str | None = None, single_stage_id_filter: int | None, omni_master_address: str, omni_master_port: int, + tokenizer: str | None = None, + log_stats: bool = False, omni_dp_size_local: int = 1, omni_heartbeat_timeout: float = 30.0, omni_lb_policy: str = "random", @@ -759,6 +767,7 @@ class DistStageRuntime(StageRuntime): diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, + log_stats=log_stats, ) self._single_stage_id_filter = single_stage_id_filter self._omni_master_address = omni_master_address @@ -1051,6 +1060,7 @@ class DistStageRuntime(StageRuntime): client = StageEngineCoreClientBase.make_async_mp_client( vllm_config=vllm_config, executor_class=ctx.executor_class, + log_stats=self._log_stats, metadata=metadata, client_addresses=client_addresses, engine_manager=resources.manager, @@ -1089,6 +1099,7 @@ def create_stage_runtime( omni_heartbeat_timeout: float = 30.0, omni_lb_policy: str = "random", request_queue: janus.Queue[EngineQueueMessage] | None = None, + log_stats: bool = False, ) -> StageRuntime: """Factory: select StageRuntime or DistStageRuntime.""" if single_stage_mode: @@ -1102,6 +1113,7 @@ def create_stage_runtime( diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, + log_stats=log_stats, single_stage_id_filter=single_stage_id_filter, omni_master_address=omni_master_address, omni_master_port=omni_master_port, @@ -1118,4 +1130,5 @@ def create_stage_runtime( diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, + log_stats=log_stats, ) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 423aa65b9..55ebffa02 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -612,6 +612,7 @@ class OmniBase(PDDisaggregationMixin): request_id=req_id or "", finished=finished, stage_id=stage_id, + replica_id=result.replica_id, final_output_type=output_type, request_output=engine_outputs, images=images, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 7c807cfc9..5dc3002ca 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2062,15 +2062,19 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): now_ts = time.time() if req_state is not None and req_state.first_audio_ts is None: req_state.first_audio_ts = now_ts - stage_pools = getattr(self.engine_client.engine, "stage_pools", None) - # The orchestrator binds requests by their internal id, - # not the user-visible external id, so look up the - # replica with req_state.request_id (internal). - replica_id = ( - stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id) - if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools) - else None - ) + replica_id = getattr(omni_res, "replica_id", None) + if replica_id is None: + stage_pools = getattr(self.engine_client.engine, "stage_pools", None) + # Fallback for older outputs. The orchestrator binds + # requests by internal id, but can release that + # binding before the API sees the final audio chunk, + # so prefer the OutputMessage replica copied onto + # OmniRequestOutput above. + replica_id = ( + stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id) + if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools) + else None + ) req_state.audio_emit_stage_id = omni_res.stage_id req_state.audio_emit_replica_id = replica_id observe_audio_first_packet( diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index 34731eeca..d0e41dbf0 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -83,6 +83,7 @@ class OmniRequestOutput: request_id: Unique identifier for this request finished: Whether generation is complete stage_id: Identifier of the stage that produced this output (pipeline mode) + replica_id: Identifier of the stage replica that produced this output final_output_type: Type of output ("text", "image", "audio", "latents") request_output: The underlying RequestOutput from the stage (pipeline mode) images: List of generated PIL images (diffusion mode) @@ -96,6 +97,7 @@ class OmniRequestOutput: # Pipeline stage fields stage_id: int | None = None + replica_id: int | None = None final_output_type: str = "text" request_output: RequestOutput | None = None @@ -154,6 +156,7 @@ class OmniRequestOutput: stage_id: int, final_output_type: str, request_output: RequestOutput, + replica_id: int | None = None, ) -> "OmniRequestOutput": """Create output from pipeline stage. @@ -168,6 +171,7 @@ class OmniRequestOutput: return cls( request_id=getattr(request_output, "request_id", ""), stage_id=stage_id, + replica_id=replica_id, final_output_type=final_output_type, request_output=request_output, finished=True, diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index 4d06f4124..9558d66e5 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -523,7 +523,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): upstream_outputs.append(eco) # Handle multimodal-only outputs (generation stages) locally. - mm_request_outputs = self._process_mm_only_outputs(mm_only_outputs) + mm_request_outputs = self._process_mm_only_outputs( + mm_only_outputs, + engine_core_timestamp=engine_core_timestamp, + iteration_stats=iteration_stats, + ) # Delegate text/pooling outputs to upstream. processed = super().process_outputs( @@ -537,6 +541,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): def _process_mm_only_outputs( self, engine_core_outputs: list[EngineCoreOutput], + *, + engine_core_timestamp: float | None = None, + iteration_stats: IterationStats | None = None, ) -> list[OmniRequestOutput | PoolingRequestOutput]: """Handle outputs from generation stages that have no detokenizer. @@ -554,7 +561,15 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): stop_reason = eco.stop_reason kv_transfer_params = eco.kv_transfer_params routed_experts = eco.routed_experts - req_state.num_cached_tokens = eco.num_cached_tokens + self._update_stats_from_output( + req_state, + eco, + engine_core_timestamp, + iteration_stats, + ) + prefill_stats = getattr(eco, "prefill_stats", None) + if prefill_stats is not None: + req_state.num_cached_tokens = prefill_stats.num_cached_tokens req_state.is_prefilling = False is_non_final_audio_chunk = ( @@ -578,6 +593,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): is_segment_finished = bool(getattr(eco, "is_segment_finished", False)) if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk: self._finish_request(req_state) + self._update_stats_from_finished( + req_state, + finish_reason, + iteration_stats, + ) return request_outputs def _update_stats_from_output( From 9d67f8e02fd4b386a62b8bdac2bed95039ef6753 Mon Sep 17 00:00:00 2001 From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Date: Fri, 31 Jul 2026 21:44:28 -0700 Subject: [PATCH 05/10] [CI/Build] Let audio assertions pin the transcript language (#5646) Signed-off-by: Shenglei Fu --- .../test_qwen3_omni_expansion.py | 4 + tests/helpers/assertions.py | 9 ++- tests/helpers/media.py | 15 ++-- tests/helpers/tests/test_assertions.py | 65 ++++++++++++++- tests/helpers/tests/test_media.py | 81 +++++++++++++++++++ 5 files changed, 165 insertions(+), 9 deletions(-) create mode 100644 tests/helpers/tests/test_media.py diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 3ded285ce..1ba94c011 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -468,6 +468,10 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None: "messages": messages, "stream": True, "key_words": {"text": ["london"]}, + # A one-word English clip is too short for whisper's language detection to + # be stable; unpinned it renders "London" as 런던 / ランデ / 梁敦 and the + # containment check then fails on correct audio. + "transcript_language": "en", } # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/helpers/assertions.py). diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 11e86ed8c..9fe34f9b5 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -527,7 +527,7 @@ def _resolve_audio_transcript( audio_bytes = getattr(response, "audio_bytes", None) if not audio_bytes: return None - return convert_audio_bytes_to_text(audio_bytes) + return convert_audio_bytes_to_text(audio_bytes, language=request_config.get("transcript_language")) def assert_omni_response(response: Any, request_config: dict[str, Any], run_level): @@ -632,6 +632,7 @@ def _assert_transcript_matches( *, threshold: float, escalation_model: str | None = None, + language: str | None = None, ) -> None: """Assert spoken audio matches ``expected_text``. @@ -647,7 +648,8 @@ def _assert_transcript_matches( clip is re-transcribed with that stronger ASR and the assertion is decided on its verdict -- so a weak whisper-``small`` mishear on a short clip does not flake the gate, while a genuine model artifact still fails (the strong ASR - mismatches too). + mismatches too). ``language`` applies to that escalated pass too, so a request + that pins a language does not silently fall back to auto-detection on retry. """ expected = str(expected_text).strip().lower() similarity = cosine_similarity_text(transcript.strip().lower(), expected) @@ -664,7 +666,7 @@ def _assert_transcript_matches( f"whisper-small below threshold ({similarity:.2f} <= {threshold}); " f"escalating to whisper-{escalation_model} to rule out an ASR mishear" ) - strong_transcript = convert_audio_bytes_to_text(audio_bytes, model_size=escalation_model) + strong_transcript = convert_audio_bytes_to_text(audio_bytes, model_size=escalation_model, language=language) strong_similarity = cosine_similarity_text(strong_transcript.strip().lower(), expected) print( f"audio content (whisper-{escalation_model}): {strong_transcript}\n" @@ -741,6 +743,7 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], expected_text, threshold=0.9, escalation_model=request_config.get("transcript_escalation_model"), + language=request_config.get("transcript_language"), ) _assert_preset_voice_gender_from_audio( response.audio_bytes, diff --git a/tests/helpers/media.py b/tests/helpers/media.py index be444051d..93e136f57 100644 --- a/tests/helpers/media.py +++ b/tests/helpers/media.py @@ -652,7 +652,9 @@ def _serialize_whisper_model_download(model_size: str = "small"): f.close() -def _whisper_transcribe_in_current_process(output_path: str, model_size: str = "small") -> str: +def _whisper_transcribe_in_current_process( + output_path: str, model_size: str = "small", language: str | None = None +) -> str: import whisper device_index = None @@ -684,6 +686,9 @@ def _whisper_transcribe_in_current_process(output_path: str, model_size: str = " temperature=0.0, word_timestamps=True, condition_on_previous_text=False, + # None keeps whisper's auto-detection. Do not default this to a + # language: callers include non-English audio tests. + language=language, )["text"] finally: del model @@ -694,15 +699,15 @@ def _whisper_transcribe_in_current_process(output_path: str, model_size: str = " return text or "" -def convert_audio_file_to_text(output_path: str, model_size: str = "small") -> str: +def convert_audio_file_to_text(output_path: str, model_size: str = "small", language: str | None = None) -> str: """Convert an audio file to text in an isolated subprocess.""" ctx = multiprocessing.get_context("spawn") with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor: - future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size) + future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size, language) return future.result() -def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small") -> str: +def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small", language: str | None = None) -> str: output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav") os.close(output_fd) if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"): @@ -710,7 +715,7 @@ def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small") -> data, samplerate = sf.read(io.BytesIO(raw_bytes)) sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16") print(f"audio data is saved: {output_path}") - return convert_audio_file_to_text(output_path, model_size) + return convert_audio_file_to_text(output_path, model_size, language) __all__ = [ diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py index 8293349fa..9cb1560b9 100644 --- a/tests/helpers/tests/test_assertions.py +++ b/tests/helpers/tests/test_assertions.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest -from tests.helpers.assertions import _assert_transcript_matches +from tests.helpers import assertions +from tests.helpers.assertions import _assert_transcript_matches, _resolve_audio_transcript pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -25,3 +28,63 @@ def test_short_transcript_unrelated_text_still_fails(): expected_text="how are you", threshold=0.9, ) + + +def _capture_transcribe(monkeypatch) -> dict: + captured: dict = {} + + def fake_convert(raw_bytes, model_size="small", language=None): + captured["model_size"] = model_size + captured["language"] = language + return "London" + + monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", fake_convert) + return captured + + +def test_resolve_transcript_honors_declared_language(monkeypatch): + captured = _capture_transcribe(monkeypatch) + response = SimpleNamespace(audio_content=None, audio_bytes=b"fake-wav") + + _resolve_audio_transcript( + response, + {"modalities": ["text", "audio"], "transcript_language": "en"}, + "full_model", + speech_api=False, + ) + + assert captured["language"] == "en" + + +def test_resolve_transcript_leaves_language_unset_by_default(monkeypatch): + captured = _capture_transcribe(monkeypatch) + response = SimpleNamespace(audio_content=None, audio_bytes=b"fake-wav") + + _resolve_audio_transcript( + response, + {"modalities": ["text", "audio"]}, + "full_model", + speech_api=False, + ) + + assert captured["language"] is None + assert captured["model_size"] == "small" + + +def test_escalated_transcript_keeps_declared_language(monkeypatch): + # The escalated pass must honour the same language, otherwise a request that + # pins one silently falls back to auto-detection on retry. + captured = _capture_transcribe(monkeypatch) + + with pytest.raises(AssertionError, match="after ASR escalation"): + _assert_transcript_matches( + "totally different words", + audio_bytes=b"fake-wav", + expected_text="how are you", + threshold=0.9, + escalation_model="large-v3", + language="en", + ) + + assert captured["model_size"] == "large-v3" + assert captured["language"] == "en" diff --git a/tests/helpers/tests/test_media.py b/tests/helpers/tests/test_media.py new file mode 100644 index 000000000..4a8e5393d --- /dev/null +++ b/tests/helpers/tests/test_media.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +from types import SimpleNamespace + +import numpy as np +import pytest +import soundfile as sf + +from tests.helpers import media + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _stub_transcribe(monkeypatch) -> dict: + """Capture the kwargs the helper hands to whisper, with no model and no device.""" + captured: dict = {} + + class FakeModel: + def transcribe(self, path, **kwargs): + captured.update(kwargs) + return {"text": "London"} + + monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=lambda size, device=None: FakeModel())) + # Keep the unit test off the accelerator probe so it stays hermetic. + monkeypatch.setitem( + sys.modules, + "vllm_omni.platforms", + SimpleNamespace(current_omni_platform=SimpleNamespace(is_available=lambda: False)), + ) + return captured + + +def test_transcribe_forwards_requested_language(monkeypatch): + captured = _stub_transcribe(monkeypatch) + + media._whisper_transcribe_in_current_process("/tmp/does-not-matter.wav", "small", language="en") + + assert captured["language"] == "en" + + +def test_transcribe_defaults_to_auto_language(monkeypatch): + # Auto-detect must stay the default: forcing a language globally would break + # the non-English audio tests (e.g. the Chinese Qwen3-Omni prompts). + captured = _stub_transcribe(monkeypatch) + + media._whisper_transcribe_in_current_process("/tmp/does-not-matter.wav", "small") + + assert captured.get("language") is None + + +def test_bytes_entrypoint_forwards_language_to_subprocess(monkeypatch, tmp_path): + """Cover the two hops the tests above skip: bytes -> file -> executor.submit. + + The real call crosses a spawn ProcessPoolExecutor, so capture what gets + submitted rather than what the worker eventually does. + """ + submitted: dict = {} + + class FakeExecutor: + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def submit(self, fn, *args): + submitted["fn"] = fn + submitted["args"] = args + return SimpleNamespace(result=lambda: "London") + + monkeypatch.setattr(media.concurrent.futures, "ProcessPoolExecutor", lambda *a, **kw: FakeExecutor()) + + wav = tmp_path / "clip.wav" + sf.write(wav, np.zeros(2400, dtype=np.float32), 24000) + + assert media.convert_audio_bytes_to_text(wav.read_bytes(), "small", "en") == "London" + + assert submitted["fn"] is media._whisper_transcribe_in_current_process + assert submitted["args"][1:] == ("small", "en") From 67aa6c55a4408f96cd195f76e367f24dbb888238 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Sat, 1 Aug 2026 13:41:00 +0800 Subject: [PATCH 06/10] =?UTF-8?q?[Bugfix][MiniCPM-o]=20Align=20Daily-Omni?= =?UTF-8?q?=20bench=20with=20MiniCPM=20interleaved=20A=E2=80=A6=20(#5625)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: amy-why-3459 Co-authored-by: Cursor --- .../benchmarks/test_daily_omni_pack_modes.py | 249 ++++++++++++++ .../models/minicpmo_4_5/test_pipeline.py | 7 +- .../data_modules/daily_omni_dataset.py | 314 +++++++++++++++++- vllm_omni/benchmarks/patch/patch.py | 2 + vllm_omni/deploy/minicpmo_4_5.yaml | 10 +- .../entrypoints/cli/benchmark/cli_args.py | 17 + .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 25 +- 7 files changed, 606 insertions(+), 18 deletions(-) create mode 100644 tests/benchmarks/test_daily_omni_pack_modes.py diff --git a/tests/benchmarks/test_daily_omni_pack_modes.py b/tests/benchmarks/test_daily_omni_pack_modes.py new file mode 100644 index 000000000..46ca369e5 --- /dev/null +++ b/tests/benchmarks/test_daily_omni_pack_modes.py @@ -0,0 +1,249 @@ +"""Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes. + +The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches +OpenBMB's reported Daily-Omni score when frames and 1s audio segments arrive strictly +interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and an empty MiniCPM +system message (not the Qwen system prompt). These tests pin that contract so a refactor +cannot silently regress the score. + +vllm stubs are installed by tests/benchmarks/conftest.py before collection. +""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import numpy as np +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Load the data module directly (bypasses vllm_omni.__init__ heavy imports). +_REPO_ROOT = Path(__file__).resolve().parents[2] +_MODULE_PATH = _REPO_ROOT / "vllm_omni" / "benchmarks" / "data_modules" / "daily_omni_dataset.py" +_MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset" + +if _MODULE_NAME not in sys.modules: + _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH) + _mod = importlib.util.module_from_spec(_spec) + sys.modules[_MODULE_NAME] = _mod + _spec.loader.exec_module(_mod) + +from vllm_omni.benchmarks.data_modules.daily_omni_dataset import ( # noqa: E402 + DailyOmniDataset, + _uniform_sample_indices, +) + +_VIDEO_ID = "vid001" +_NUM_FRAMES = 5 + + +@pytest.fixture +def video_dir(tmp_path: Path) -> Path: + """Official Daily-Omni ``Videos/`` layout with placeholder media files.""" + per_video = tmp_path / "Videos" / _VIDEO_ID + per_video.mkdir(parents=True) + (per_video / f"{_VIDEO_ID}_video.mp4").write_bytes(b"fake-mp4") + (per_video / f"{_VIDEO_ID}_audio.wav").write_bytes(b"fake-wav") + return tmp_path / "Videos" + + +@pytest.fixture +def qa_json(tmp_path: Path) -> Path: + path = tmp_path / "qa.json" + path.write_text( + json.dumps( + [ + { + "video_id": _VIDEO_ID, + "question": "What happens first?", + "choice": {"A": "a", "B": "b", "C": "c", "D": "d"}, + "answer": "A", + "task_type": "AV Event Alignment", + "video_duration": "30s", + } + ] + ), + encoding="utf-8", + ) + return path + + +def _make_dataset(qa_json: Path, video_dir: Path, **kwargs) -> DailyOmniDataset: + return DailyOmniDataset( + qa_json_path=str(qa_json), + video_dir=str(video_dir), + disable_shuffle=True, + **kwargs, + ) + + +def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) -> None: + """Replace decord/ffmpeg extraction with deterministic in-memory frames + segments.""" + from PIL import Image + + def fake_extract(video_path, *, audio_path, include_audio, **_kw): + counter.append(1) + frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)] + segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else [] + return frames, segments + + monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract) + + +# --------------------------------------------------------------------------- +# _uniform_sample_indices +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (1000, 64), (7, 7), (3, 5), (1, 1)]) +def test_uniform_sample_indices_matches_openbmb_midpoint(n: int, k: int) -> None: + """Pin OpenBMB midpoint-bin sampling (not np.linspace endpoints).""" + got = _uniform_sample_indices(n, k) + if n <= k: + assert got == list(range(n)) + else: + gap = n / k + expected = [int(i * gap + gap / 2) for i in range(k)] + assert got == expected + # Endpoints differ from linspace for typical Daily-Omni downsample sizes. + linspace = [int(i) for i in np.linspace(0, n - 1, k, dtype=int)] + if n == 1000 and k == 64: + assert got[0] == 7 and got[-1] == 992 + assert linspace[0] == 0 and linspace[-1] == 999 + assert got != linspace + + +def test_uniform_sample_indices_non_positive_k() -> None: + assert _uniform_sample_indices(10, 0) == [] + assert _uniform_sample_indices(10, -1) == [] + + +# --------------------------------------------------------------------------- +# minicpm-interleave packing +# --------------------------------------------------------------------------- + + +def test_minicpm_interleave_alternates_image_and_audio(monkeypatch, qa_json, video_dir) -> None: + ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all") + _stub_extract(monkeypatch, ds, []) + + parts, extra, position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None) + + assert position == "first" + assert [p["type"] for p in parts] == ["image_url", "audio_url"] * _NUM_FRAMES + assert extra["mm_processor_kwargs"] == { + "use_audio_in_video": False, + "max_slice_nums": 1, + "use_image_id": False, + } + + +def test_minicpm_interleave_visual_mode_emits_frames_only(monkeypatch, qa_json, video_dir) -> None: + ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="visual") + _stub_extract(monkeypatch, ds, []) + + parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None) + + assert [p["type"] for p in parts] == ["image_url"] * _NUM_FRAMES + + +def test_minicpm_interleave_writes_file_urls_under_video_dir(monkeypatch, qa_json, video_dir) -> None: + ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all") + _stub_extract(monkeypatch, ds, []) + + parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None) + + segment_dir = video_dir / ".minicpm_daily_omni_interleave" / _VIDEO_ID + for part in parts: + url = part[part["type"]]["url"] + assert url.startswith("file://") + assert Path(url[len("file://") :]).parent == segment_dir.resolve() + assert sorted(p.name for p in segment_dir.glob("frame_*.jpg")) == [f"frame_{i:04d}.jpg" for i in range(_NUM_FRAMES)] + + +def test_minicpm_interleave_inline_mode_emits_data_urls(monkeypatch, qa_json, video_dir) -> None: + ds = _make_dataset( + qa_json, + video_dir, + pack_mode="minicpm-interleave", + input_mode="all", + inline_local_video=True, + ) + _stub_extract(monkeypatch, ds, []) + + parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None) + + assert parts[0]["image_url"]["url"].startswith("data:image/jpeg;base64,") + assert parts[1]["audio_url"]["url"].startswith("data:audio/wav;base64,") + assert not (video_dir / ".minicpm_daily_omni_interleave").exists() + + +def test_minicpm_interleave_caches_file_parts_but_not_inline_payloads(monkeypatch, qa_json, video_dir) -> None: + """Inline base64 must not be retained: ~700 Daily-Omni videos would pin GBs of RSS.""" + file_calls: list[int] = [] + ds_file = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all") + _stub_extract(monkeypatch, ds_file, file_calls) + ds_file._compose_daily_omni_multimodal(_VIDEO_ID, None) + ds_file._compose_daily_omni_multimodal(_VIDEO_ID, None) + assert len(file_calls) == 1 + assert ds_file._minicpm_interleave_cache + + inline_calls: list[int] = [] + ds_inline = _make_dataset( + qa_json, + video_dir, + pack_mode="minicpm-interleave", + input_mode="all", + inline_local_video=True, + ) + _stub_extract(monkeypatch, ds_inline, inline_calls) + ds_inline._compose_daily_omni_multimodal(_VIDEO_ID, None) + ds_inline._compose_daily_omni_multimodal(_VIDEO_ID, None) + assert len(inline_calls) == 2 + assert not ds_inline._minicpm_interleave_cache + + +def test_minicpm_interleave_missing_local_video_returns_none(qa_json, tmp_path) -> None: + empty = tmp_path / "empty_videos" + empty.mkdir() + ds = _make_dataset(qa_json, empty, pack_mode="minicpm-interleave", input_mode="all") + + parts, extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None) + + assert parts is None + assert extra is None + + +# --------------------------------------------------------------------------- +# System prompt selection +# --------------------------------------------------------------------------- + + +def test_minicpm_interleave_keeps_empty_system_message(monkeypatch, qa_json, video_dir) -> None: + """Official MiniCPM omni mode still sends role=system with empty content.""" + ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all") + _stub_extract(monkeypatch, ds, []) + parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None) + + messages = ds._build_daily_omni_openai_messages(parts, "q?", {"A": "a"}) + + assert [m["role"] for m in messages] == ["system", "user"] + assert messages[0]["content"][0]["text"] == "" + + +def test_qwen_pack_mode_keeps_system_message(qa_json, video_dir) -> None: + ds = _make_dataset(qa_json, video_dir, pack_mode="qwen", input_mode="all") + + messages = ds._build_daily_omni_openai_messages([], "q?", {"A": "a"}) + + assert [m["role"] for m in messages] == ["system", "user"] + assert "Qwen" in messages[0]["content"][0]["text"] + + +def test_invalid_pack_mode_rejected(qa_json, video_dir) -> None: + with pytest.raises(ValueError, match="pack_mode must be"): + _make_dataset(qa_json, video_dir, pack_mode="bogus") diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 84b5ca11c..5236ea883 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -170,7 +170,12 @@ class TestDeployTopology: 0.15, ] assert sum(memory_utilizations) <= 0.9 + 1e-6 - assert stages[0].yaml_engine_args["limit_mm_per_prompt"] == {"video": {"count": 1, "num_frames": 32}} + # Daily-Omni minicpm-interleave: up to 64 image/audio items (+ optional video). + assert stages[0].yaml_engine_args["limit_mm_per_prompt"] == { + "image": 64, + "audio": 64, + "video": 1, + } elif filename in {"minicpmo_4_5_2gpu.yaml"}: assert [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] == [ 0.9, diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py index 65918414f..95d58e4b3 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py @@ -52,11 +52,14 @@ Usage: """ import base64 +import io import json import logging +import math import os import shutil import tarfile +import wave from dataclasses import dataclass from pathlib import Path from typing import Any, Literal @@ -77,6 +80,10 @@ except ImportError: logger = logging.getLogger(__name__) +# Official MiniCPM ``get_video_frame_audio_segments`` default (env ``MAX_NUM_FRAMES``). +_MINICPM_MAX_NUM_FRAMES = int(os.environ.get("MAX_NUM_FRAMES", "64")) +_MINICPM_AUDIO_SR = 16000 + def _daily_omni_hf_cache_root() -> Path: return Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")).expanduser().resolve() @@ -203,12 +210,62 @@ class _ListDatasetIterator: # Aligns with Lliar-liar/Daily-Omni CLI ``--input_mode`` (test_model/*/testmodel.py). DailyOmniInputMode = Literal["all", "visual", "audio"] +# How multimodal parts are packed into OpenAI chat messages: +# - ``qwen``: one ``video_url`` + one ``audio_url`` (upstream Daily-Omni / Qwen protocol). +# - ``minicpm-interleave``: 1fps ``[image_i, audio_i, ...]`` matching OpenBMB MiniCPM-o +# ``get_video_frame_audio_segments`` + interleaved omni contents (paper ~80% recipe). +DailyOmniPackMode = Literal["qwen", "minicpm-interleave"] + # ``build_conversation()`` in Daily-Omni ``test_model/Qwen2.5-Omni/testmodel.py`` (verbatim). DAILY_OMNI_SYSTEM_TEXT = ( "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, " "capable of perceiving auditory and visual inputs, as well as generating text and speech." ) +# MiniCPM-o ``get_sys_prompt(mode="omni")`` without ref-audio is an empty system string +# (still sent as role=system; see ``_build_daily_omni_openai_messages``). +MINICPM_OMNI_SYSTEM_TEXT = "" + + +def _uniform_sample_indices(n: int, k: int) -> list[int]: + """Evenly pick ``k`` indices from ``0..n-1`` (OpenBMB ``uniform_sample`` midpoint bins). + + Matches OmniEvalKit / MiniCPM ``uniform_sample``: midpoint of each equal-width bin, + not ``np.linspace`` endpoints. See + https://github.com/OpenBMB/OmniEvalKit/blob/1adac0577258d539efc03f16d8a0f4b3f7df6c19/o_e_Kit/utils/utils.py#L50-L54 + """ + if k <= 0: + return [] + if n <= k: + return list(range(n)) + gap = n / k + return [int(i * gap + gap / 2) for i in range(k)] + + +def _numpy_to_wav_bytes(audio_np: Any, sample_rate: int = _MINICPM_AUDIO_SR) -> bytes: + """Encode a mono float/int numpy waveform as 16-bit PCM WAV bytes.""" + import numpy as np + + arr = np.asarray(audio_np) + if arr.ndim > 1: + arr = arr.reshape(-1) + if arr.dtype != np.int16: + arr = np.clip(arr.astype(np.float64), -1.0, 1.0) + arr = (arr * 32767.0).astype(np.int16) + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(arr.tobytes()) + return buf.getvalue() + + +def _pil_to_jpeg_bytes(image: Any, quality: int = 85) -> bytes: + buf = io.BytesIO() + image.save(buf, format="JPEG", quality=quality) + return buf.getvalue() + @dataclass class DailyOmniSampleRequest(SampleRequest): @@ -262,6 +319,11 @@ class DailyOmniDataset(BenchmarkDataset): ``all`` — video + WAV (default; official audio-visual protocol); ``visual`` — video only; ``audio`` — extracted WAV only (requires ``{video_id}/{video_id}_audio.wav`` under ``video_dir``). + pack_mode: Multimodal packing recipe: + ``qwen`` — one ``video_url`` + optional ``audio_url`` (default; Daily-Omni/Qwen); + ``minicpm-interleave`` — MiniCPM-o paper recipe: 1fps frames interleaved with + matching 1s audio segments as ``image_url``/``audio_url`` pairs (required to approach + OpenBMB ~80% Daily-Omni; needs local video+WAV under ``video_dir``). max_duration_seconds: Reserved for future ffprobe-based filtering; currently **not applied** when building requests (metadata ``video_duration`` is still passed through for eval). dataset_subset: Optional HuggingFace subset name (``load_dataset(..., name=...)``); used by bench @@ -270,7 +332,8 @@ class DailyOmniDataset(BenchmarkDataset): inline_local_video: If True, embed local MP4 as ``data:video/mp4;base64,...`` in requests so the API server does not need ``--allowed-local-media-path`` (large JSON; use for small runs). When ``input_mode`` is ``audio`` or ``all``, local WAV is embedded the same way - (``data:audio/wav;base64,...``). + (``data:audio/wav;base64,...``). For ``minicpm-interleave``, embeds JPEG/WAV segment + data URLs instead of the raw MP4. trust_remote_code: Whether to trust remote code when loading HuggingFace dataset (online mode only). """ @@ -291,6 +354,7 @@ class DailyOmniDataset(BenchmarkDataset): random_seed: int = 0, video_dir: str | None = None, input_mode: DailyOmniInputMode = "all", + pack_mode: DailyOmniPackMode = "qwen", inline_local_video: bool = False, trust_remote_code: bool = False, max_duration_seconds: float | None = None, @@ -300,6 +364,8 @@ class DailyOmniDataset(BenchmarkDataset): ) -> None: if input_mode not in ("all", "visual", "audio"): raise ValueError(f"input_mode must be 'all', 'visual', or 'audio', got {input_mode!r}") + if pack_mode not in ("qwen", "minicpm-interleave"): + raise ValueError(f"pack_mode must be 'qwen' or 'minicpm-interleave', got {pack_mode!r}") # Validate arguments: need either local JSON or HF path if qa_json_path is None and dataset_path is None: @@ -318,11 +384,14 @@ class DailyOmniDataset(BenchmarkDataset): self.video_dir = Path(video_dir) if video_dir else None self.inline_local_video = inline_local_video self.input_mode: DailyOmniInputMode = input_mode + self.pack_mode: DailyOmniPackMode = pack_mode self.max_duration_seconds = max_duration_seconds self.trust_remote_code = trust_remote_code #: In-process cache of ffprobe durations only (no disk persistence). self._video_durations: dict[str, float] = {} + #: Cache of MiniCPM 1fps interleaved expansions: video_id -> content parts. + self._minicpm_interleave_cache: dict[str, list[dict[str, Any]]] = {} # Initialize parent BenchmarkDataset super().__init__( @@ -336,11 +405,13 @@ class DailyOmniDataset(BenchmarkDataset): # Verify dataset info logger.info( - "Loaded Daily-Omni dataset: mode=%s, source=%s, random_seed=%d, input_mode=%s, max_duration=%s", + "Loaded Daily-Omni dataset: mode=%s, source=%s, random_seed=%d, " + "input_mode=%s, pack_mode=%s, max_duration=%s", "local_json" if self.qa_json_path else "huggingface", str(self.qa_json_path) if self.qa_json_path else f"{dataset_path}/{dataset_split}", random_seed, input_mode, + pack_mode, f"{max_duration_seconds}s" if max_duration_seconds else "unlimited", ) @@ -575,7 +646,8 @@ class DailyOmniDataset(BenchmarkDataset): messages = self._build_daily_omni_openai_messages(mm_payload, question, choice) user_text = self._official_daily_omni_user_prompt(question, choice) # Text-only length estimate (same as before: no MM token count in bench). - prompt_len = len(tokenizer.encode(f"{DAILY_OMNI_SYSTEM_TEXT}\n{user_text}")) + system_text = self._system_text_for_pack_mode() + prompt_len = len(tokenizer.encode(f"{system_text}\n{user_text}" if system_text else user_text)) return DailyOmniSampleRequest( prompt=user_text, @@ -766,10 +838,11 @@ class DailyOmniDataset(BenchmarkDataset): video_id: str, video_url: str | None, ) -> tuple[dict[str, Any] | list[dict[str, Any]] | None, dict[str, Any] | None, Literal["first", "last"]]: - """Build ``multi_modal_data`` and request extras for the active ``input_mode``. + """Build multimodal OpenAI content parts and request extras for the active modes.""" + if self.pack_mode == "minicpm-interleave": + return self._compose_minicpm_interleaved_multimodal(video_id, video_url) - Mirrors upstream Daily-Omni: separate video + WAV with ``use_audio_in_video=False``. - """ + # Qwen / upstream Daily-Omni: separate video + WAV with ``use_audio_in_video=False``. extra: dict[str, Any] = {"mm_processor_kwargs": {"use_audio_in_video": False}} mode = self.input_mode @@ -787,6 +860,219 @@ class DailyOmniDataset(BenchmarkDataset): return None, None, "first" return [v, a], extra, "first" + def _compose_minicpm_interleaved_multimodal( + self, + video_id: str, + video_url: str | None, + ) -> tuple[list[dict[str, Any]] | None, dict[str, Any] | None, Literal["first", "last"]]: + """Pack media like MiniCPM-o official omni chat: 1fps ``[frame_i, audio_i, ...]``. + + OpenBMB's HF / utils path expands ``video_url`` with ``use_audio=True`` into interleaved + PIL frames and 1s audio ndarrays. vLLM-Omni does not implement that expansion, so the + bench client builds ordered ``image_url`` / ``audio_url`` parts instead. MiniCPM-o 5.0's + OpenAI chat template preserves content-part order, so placeholders stay time-aligned. + + ``max_slice_nums=1`` and ``use_image_id=False`` mirror MiniCPM-o's omni path + (``modeling_minicpmo.py`` passes both for interleaved frames): slicing every frame with + the config default of 9 would blow past ``max_model_len``, and ``N`` + prefixes would label each frame as a separate picture instead of a video timeline. + """ + del video_url # Hub HTTP URLs are not used; local extract is required. + extra: dict[str, Any] = { + "mm_processor_kwargs": { + "use_audio_in_video": False, + "max_slice_nums": 1, + "use_image_id": False, + } + } + mode = self.input_mode + if mode == "audio": + # Fall back to whole-WAV for audio-only ablations. + a = self._get_audio_content(video_id) + return ([a] if a else None), extra, "first" + + parts = self._get_minicpm_interleaved_parts(video_id, include_audio=(mode == "all")) + if not parts: + return None, None, "first" + return parts, extra, "first" + + def _get_minicpm_interleaved_parts( + self, + video_id: str, + *, + include_audio: bool, + ) -> list[dict[str, Any]] | None: + """Return cached or freshly extracted MiniCPM-style interleaved OpenAI content parts.""" + cache_key = f"{video_id}|audio={int(include_audio)}|inline={int(self.inline_local_video)}" + cached = self._minicpm_interleave_cache.get(cache_key) + if cached is not None: + return cached + + self._lazy_ensure_hub_media_dir() + video_path = self._resolve_local_video_path(video_id) + if video_path is None: + logger.warning( + "MiniCPM interleave pack requires local video under video_dir=%s for video_id=%r", + self.video_dir, + video_id, + ) + return None + audio_path = self._resolve_local_audio_path(video_id) if include_audio else None + if include_audio and audio_path is None: + logger.warning( + "MiniCPM interleave pack (input_mode=all) missing WAV for video_id=%r under %s", + video_id, + self.video_dir, + ) + return None + + try: + frames, audio_segments = self._extract_minicpm_frame_audio_segments( + video_path, + audio_path=audio_path, + include_audio=include_audio, + ) + except Exception: + logger.exception( + "Failed MiniCPM 1fps extract for video_id=%r path=%s", + video_id, + video_path, + ) + return None + + if not frames: + logger.warning("No frames extracted for video_id=%r", video_id) + return None + if include_audio and len(audio_segments) != len(frames): + logger.warning( + "Frame/audio length mismatch for video_id=%r: frames=%d audios=%d", + video_id, + len(frames), + len(audio_segments), + ) + n = min(len(frames), len(audio_segments)) + frames = frames[:n] + audio_segments = audio_segments[:n] + + cache_root: Path | None = None + if not self.inline_local_video: + assert self.video_dir is not None + cache_root = self.video_dir / ".minicpm_daily_omni_interleave" / video_id + cache_root.mkdir(parents=True, exist_ok=True) + + parts: list[dict[str, Any]] = [] + for i, frame in enumerate(frames): + jpeg = _pil_to_jpeg_bytes(frame) + if self.inline_local_video: + b64 = base64.b64encode(jpeg).decode("ascii") + parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{b64}"}, + } + ) + else: + assert cache_root is not None + frame_path = cache_root / f"frame_{i:04d}.jpg" + if not frame_path.is_file() or frame_path.stat().st_size == 0: + frame_path.write_bytes(jpeg) + parts.append( + { + "type": "image_url", + "image_url": {"url": frame_path.expanduser().resolve().as_uri()}, + } + ) + + if include_audio: + wav = _numpy_to_wav_bytes(audio_segments[i]) + if self.inline_local_video: + b64 = base64.b64encode(wav).decode("ascii") + parts.append( + { + "type": "audio_url", + "audio_url": {"url": f"data:audio/wav;base64,{b64}"}, + } + ) + else: + assert cache_root is not None + audio_seg_path = cache_root / f"audio_{i:04d}.wav" + if not audio_seg_path.is_file() or audio_seg_path.stat().st_size == 0: + audio_seg_path.write_bytes(wav) + parts.append( + { + "type": "audio_url", + "audio_url": {"url": audio_seg_path.expanduser().resolve().as_uri()}, + } + ) + + # File URLs are tiny; cache them so repeated QA on the same video skips re-extract. + # Inline base64 must not be retained: ~700 Daily-Omni videos would pin GBs of RSS. + if not self.inline_local_video: + self._minicpm_interleave_cache[cache_key] = parts + logger.debug( + "MiniCPM interleave packed video_id=%r frames=%d include_audio=%s parts=%d", + video_id, + len(frames), + include_audio, + len(parts), + ) + return parts + + @staticmethod + def _extract_minicpm_frame_audio_segments( + video_path: Path, + *, + audio_path: Path | None, + include_audio: bool, + max_num_frames: int = _MINICPM_MAX_NUM_FRAMES, + ) -> tuple[list[Any], list[Any]]: + """Port of MiniCPM ``get_video_frame_audio_segments`` (stack_frames=1, 1fps).""" + import numpy as np + from decord import VideoReader, cpu + from PIL import Image + from vllm.multimodal.media.audio import load_audio + + vr = VideoReader(str(video_path), ctx=cpu(0)) + avg_fps = float(vr.get_avg_fps()) or 1.0 + duration = len(vr) / avg_fps + num_seconds = max(1, int(math.ceil(duration))) + second_timestamps = list(range(num_seconds)) + + if duration > max_num_frames: + timestamps = [round(i * 0.1, 1) for i in range(int(duration / 0.1))] + frame_idx = [min(int(ts * avg_fps), len(vr) - 1) for ts in timestamps] + pick = _uniform_sample_indices(len(frame_idx), max_num_frames) + frame_idx = [frame_idx[i] for i in pick] + timestamps = [timestamps[i] for i in pick] + else: + frame_idx = [min(int(i * avg_fps), len(vr) - 1) for i in range(num_seconds)] + timestamps = second_timestamps + + video = vr.get_batch(frame_idx).asnumpy() + frames = [Image.fromarray(v.astype("uint8")).convert("RGB") for v in video] + + audio_segments: list[Any] = [] + if include_audio: + load_path = str(audio_path) if audio_path is not None else str(video_path) + audio_np, sr = load_audio(load_path, sr=_MINICPM_AUDIO_SR, mono=True) + for i, start_time in enumerate(timestamps): + end_time = timestamps[i + 1] if i < len(timestamps) - 1 else duration + start_sample = int(start_time * sr) + end_sample = int(end_time * sr) + segment = audio_np[start_sample:end_sample] + if i == len(timestamps) - 1 and len(segment) < 1600: + segment = np.concatenate([segment, np.zeros(1600 - len(segment), dtype=segment.dtype)]) + if len(segment) == 0: + segment = np.zeros(1600, dtype=np.float32) + audio_segments.append(segment) + + return frames, audio_segments + + def _system_text_for_pack_mode(self) -> str: + if self.pack_mode == "minicpm-interleave": + return MINICPM_OMNI_SYSTEM_TEXT + return DAILY_OMNI_SYSTEM_TEXT + @staticmethod def _media_desc_for_official_prompt(mode: DailyOmniInputMode) -> str: """``media_desc`` in upstream ``build_conversation``.""" @@ -835,14 +1121,19 @@ class DailyOmniDataset(BenchmarkDataset): question: str, choice: Any, ) -> list[dict[str, Any]]: - """Map upstream conversation to OpenAI Chat Completions ``messages`` (video_url / audio_url parts).""" + """Map upstream conversation to OpenAI Chat Completions ``messages``.""" user_text = self._official_daily_omni_user_prompt(question, choice) mm_list: list[dict[str, Any]] = mm_payload if isinstance(mm_payload, list) else [mm_payload] user_content: list[dict[str, Any]] = [*mm_list, {"type": "text", "text": user_text}] - return [ - {"role": "system", "content": [{"type": "text", "text": DAILY_OMNI_SYSTEM_TEXT}]}, - {"role": "user", "content": user_content}, - ] + system_text = self._system_text_for_pack_mode() + messages: list[dict[str, Any]] = [] + # MiniCPM-o ``get_sys_prompt(mode="omni")`` still emits role=system with empty + # content (no ref-audio). Truthiness would drop that role and change chat-template + # control tokens; keep the empty system message for minicpm-interleave. + if self.pack_mode == "minicpm-interleave" or system_text: + messages.append({"role": "system", "content": [{"type": "text", "text": system_text}]}) + messages.append({"role": "user", "content": user_content}) + return messages def sample_by_task_type( self, @@ -900,6 +1191,7 @@ class DailyOmniDataset(BenchmarkDataset): f"dataset_split={self.dataset_split!r}, " f"video_dir={self.video_dir!r}, " f"input_mode={self.input_mode!r}, " + f"pack_mode={self.pack_mode!r}, " f"inline_local_video={self.inline_local_video!r}, " f"max_duration_seconds={self.max_duration_seconds}, " f"random_seed={self.random_seed}" diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 00af6ee23..2d816d718 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -280,6 +280,7 @@ def get_samples(args, tokenizer): random_seed=args.seed, video_dir=video_dir, input_mode=getattr(args, "daily_omni_input_mode", "all"), + pack_mode=getattr(args, "daily_omni_pack_mode", "qwen"), inline_local_video=getattr(args, "daily_omni_inline_local_video", False), trust_remote_code=getattr(args, "trust_remote_code", False), disable_shuffle=getattr(args, "disable_shuffle", False), @@ -310,6 +311,7 @@ def get_samples(args, tokenizer): random_seed=args.seed, video_dir=video_dir, input_mode=getattr(args, "daily_omni_input_mode", "all"), + pack_mode=getattr(args, "daily_omni_pack_mode", "qwen"), inline_local_video=getattr(args, "daily_omni_inline_local_video", False), trust_remote_code=getattr(args, "trust_remote_code", False), no_stream=getattr(args, "no_stream", False), diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index b3e211175..66a5f3ac9 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -24,9 +24,10 @@ stages: max_num_batched_tokens: 16384 devices: "0" limit_mm_per_prompt: - video: - count: 1 - num_frames: 32 + # Daily-Omni MiniCPM pack: up to 64 interleaved image/audio pairs (1fps). + image: 64 + audio: 64 + video: 1 default_sampling_params: temperature: 0.0 top_p: 1.0 @@ -34,7 +35,8 @@ stages: max_tokens: 2048 seed: 42 detokenize: true - repetition_penalty: 1.1 + # 1.0 for MCQ accuracy; 1.1 can bias short A–D answers. + repetition_penalty: 1.0 - stage_id: 1 max_num_seqs: 4 diff --git a/vllm_omni/entrypoints/cli/benchmark/cli_args.py b/vllm_omni/entrypoints/cli/benchmark/cli_args.py index 7f7671496..97b549fe1 100644 --- a/vllm_omni/entrypoints/cli/benchmark/cli_args.py +++ b/vllm_omni/entrypoints/cli/benchmark/cli_args.py @@ -95,6 +95,23 @@ def add_daily_omni_cli_args(parser: argparse.ArgumentParser) -> None: "'all': video + WAV together. Sets mm_processor_kwargs.use_audio_in_video=false " "and matches official separate video/audio streams.", ) + group.add_argument( + "--daily-omni-pack-mode", + type=str, + choices=["qwen", "minicpm-interleave"], + default="qwen", + help="How to pack multimodal parts into OpenAI chat messages. " + "'qwen' (default): one video_url + one audio_url (Daily-Omni/Qwen protocol). " + "'minicpm-interleave': MiniCPM-o official recipe — 1fps frames interleaved with " + "matching 1s audio segments as image_url/audio_url pairs (needed to approach " + "OpenBMB ~80%% Daily-Omni; requires local Videos extract). By default this mode " + "writes JPEG/WAV segments under " + "/.minicpm_daily_omni_interleave and sends file:// URLs, " + "so start the server with --allowed-local-media-path " + "(or a parent). Alternatively pass --daily-omni-inline-local-video to embed " + "segments as data URLs. For MiniCPM-o 4.5 string chat templates also start the " + "server with --interleave-mm-strings.", + ) group.add_argument( "--daily-omni-save-eval-items", action="store_true", diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 577ce5aa7..4671e3b96 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -3412,10 +3412,23 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn ) -> bool: return False - def get_image_prompt_texts(self, image_size: ImageSize, image_idx: int = 0) -> str: + def get_image_prompt_texts( + self, + image_size: ImageSize, + image_idx: int = 0, + max_slice_nums: int | None = None, + use_image_id: bool | None = None, + ) -> str: + # ``process_images`` forwards ``mm_processor_kwargs`` to the pixel path, so the + # placeholder grid must honor the same ``max_slice_nums``; otherwise a request that + # sets it (e.g. interleaved omni frames with max_slice_nums=1) produces fewer image + # features than placeholder ```` slots. ``use_image_id`` is likewise per-request + # so callers can match MiniCPM-o's omni path, which emits no ```` prefix. return self.info.get_slice_image_placeholder( image_size, image_idx=image_idx, + max_slice_nums=max_slice_nums, + **({} if use_image_id is None else {"use_image_id": use_image_id}), ) def get_video_prompt_texts(self, image_size: ImageSize, num_frames: int) -> str: @@ -3605,13 +3618,21 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn additional_placeholders.append((modality, sub_pattern)) placeholders += additional_placeholders + image_max_slice_nums = hf_processor_mm_kwargs.get("max_slice_nums") + image_use_image_id = hf_processor_mm_kwargs.get("use_image_id") + def get_image_replacement(item_idx: int): images = mm_items.get_items("image", (MiniCPMVImageEmbeddingItems, ImageProcessorItems)) image_size = images.get_image_size(item_idx) return PromptUpdateDetails.select_text( - self.get_image_prompt_texts(image_size, item_idx), + self.get_image_prompt_texts( + image_size, + item_idx, + max_slice_nums=None if image_max_slice_nums is None else int(image_max_slice_nums), # type: ignore[arg-type] + use_image_id=None if image_use_image_id is None else bool(image_use_image_id), + ), "", ) From 185e3d6df60276166349e69fdc1171fd0e50c79f Mon Sep 17 00:00:00 2001 From: Librium <129784133+pxljs@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:05:57 +0800 Subject: [PATCH 07/10] [Bugfix] Preserve channel-last video tensor layout (#5418) Signed-off-by: pxljs <1621352782@qq.com> --- .../openai_api/test_video_api_utils.py | 19 +++++++++++++++++++ .../entrypoints/openai/video_api_utils.py | 5 ++--- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 722819682..652adef4d 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -48,6 +48,25 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): assert mux_calls[0]["audio"] is None +@pytest.mark.parametrize("frame_count", [3, 4]) +def test_channel_last_video_tensor_preserves_channel_sized_frame_count(frame_count): + video = torch.arange(frame_count * 2 * 5 * 3, dtype=torch.uint8).reshape(frame_count, 2, 5, 3) + + frames = video_api_utils._coerce_video_to_uint8_frames(video) + + assert frames.shape == (frame_count, 2, 5, 3) + np.testing.assert_array_equal(frames, video.numpy()) + + +def test_channel_first_video_tensor_is_converted_to_channel_last(): + video = torch.arange(3 * 5 * 2 * 6, dtype=torch.uint8).reshape(3, 5, 2, 6) + + frames = video_api_utils._coerce_video_to_uint8_frames(video) + + assert frames.shape == (5, 2, 6, 3) + np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy()) + + def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch): muxers = [] diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index aa6479189..f0d2a59f2 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -308,9 +308,6 @@ def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray: video_tensor = video_tensor.detach().cpu() if video_tensor.dim() == 5: raise ValueError("Batched video tensors are not supported for single-video encoding.") - elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4): - # [C, F, H, W] -> [F, H, W, C] - video_tensor = video_tensor.permute(1, 2, 3, 0) if video_tensor.is_floating_point(): # Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no @@ -329,6 +326,8 @@ def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray: if video_array.ndim == 4: # Convert channel-first layouts to channel-last + # Prefer an explicit channel-last dimension for ambiguous 3/4-frame + # videos before interpreting a leading dimension as channels. if video_array.shape[0] in (3, 4) and video_array.shape[-1] not in (3, 4): video_array = np.transpose(video_array, (1, 2, 3, 0)) elif video_array.shape[1] in (3, 4) and video_array.shape[-1] not in (3, 4): From 6f98c807d29ea9d7ca347bee3bf573a4e69fe028 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:06:05 +0800 Subject: [PATCH 08/10] [Bugfix] Restore BAGEL CFG KV transfers (#5620) Signed-off-by: lishunyang12 Co-authored-by: lishunyang12 --- .../test_omni_ar_scheduler_kv_transfer.py | 57 +++++++++++++++++++ tests/engine/test_async_omni_engine_input.py | 30 ++++++++++ tests/worker/test_omni_connector_mixin.py | 35 ++++++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 4 +- .../omni_connectors/kv_transfer_manager.py | 6 ++ vllm_omni/engine/async_engine_utils.py | 7 ++- vllm_omni/engine/async_omni_engine.py | 8 +-- .../omni_connector_model_runner_mixin.py | 21 +++++++ 8 files changed, 160 insertions(+), 8 deletions(-) create mode 100644 tests/core/sched/test_omni_ar_scheduler_kv_transfer.py diff --git a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py new file mode 100644 index 000000000..cffe3f289 --- /dev/null +++ b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py @@ -0,0 +1,57 @@ +from types import SimpleNamespace + +import pytest +from vllm import SamplingParams +from vllm.v1.engine import EngineCoreRequest + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata +from vllm_omni.engine.serialization import deserialize_additional_information + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_engine_request() -> EngineCoreRequest: + return EngineCoreRequest( + request_id="req", + prompt_token_ids=[1], + mm_features=None, + sampling_params=SamplingParams(max_tokens=1), + pooling_params=None, + arrival_time=0.0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + ) + + +def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]: + tagged = apply_omni_final_stage_metadata( + _make_engine_request(), + final_stage_id=0, + force_kv_transfer=force_kv_transfer, + ) + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler._omits_kv_transfer_cache = {} + request = SimpleNamespace( + request_id="req", + additional_information=tagged.additional_information, + ) + result = scheduler._request_omits_kv_transfer_to_next_stage(request) + metadata = deserialize_additional_information(tagged.additional_information) + return result, metadata + + +def test_stage_zero_request_omits_kv_transfer(): + omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False) + + assert omits_transfer + assert "omni_force_kv_transfer" not in metadata + + +def test_cfg_companion_forces_kv_transfer_without_downstream_payload(): + omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True) + + assert not omits_transfer + assert metadata["omni_final_stage_id"] == 0 + assert metadata["omni_force_kv_transfer"] is True diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 91fdc752d..740064173 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -7,7 +7,9 @@ from vllm.v1.engine import EngineCoreRequest from vllm_omni.distributed.omni_coordinator import ReplicaInfo, ReplicaStatus from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo +from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -107,6 +109,34 @@ def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: M assert torch.equal(info["hidden_states"]["tts"], hidden) +def test_cfg_companion_suppresses_payload_but_forces_kv_transfer(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = lambda *_args: [ + ExpandedPrompt( + prompt={"prompt": "negative"}, + role="cfg_text", + request_id_suffix="__cfg_text", + ) + ] + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text") + engine.request_queue = mocker.Mock() + + engine._enqueue_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + + message = engine.request_queue.sync_q.put.call_args.args[0] + metadata = deserialize_additional_information(message.prompt.additional_information) + assert metadata["omni_final_stage_id"] == 0 + assert metadata["omni_force_kv_transfer"] is True + + def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) params = SamplingParams(max_tokens=8) diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index cf4e09cb9..7a2a0053b 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -17,6 +17,9 @@ from unittest.mock import MagicMock, patch import pytest import torch +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVTransferManager, +) from vllm_omni.outputs import OmniConnectorOutput from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, @@ -1198,6 +1201,38 @@ class TestRankAwareKVRouting(unittest.TestCase): ) host.shutdown_omni_connectors() + def test_single_rank_hooks_use_default_manager_key(self): + host = self._make_host(from_tp=1, to_tp=1, local_rank=0) + expected = ["omni_0_to_1_kv_cache_req"] + + self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected) + self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected) + host.shutdown_omni_connectors() + + def test_init_uses_effective_manager_tp2_topology(self): + model_config = _make_model_config(stage_id=0) + model_config.omni_kv_config = { + "need_send_cache": True, + "omni_from_stage": "0", + "omni_to_stage": "1", + "rank_mapping": {"from_tp": 2, "to_tp": 2}, + } + with patch( + "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", + return_value=1, + ): + manager = OmniKVTransferManager.from_model_config(model_config) + + host = MixinHost() + host.init_omni_connectors(model_config=model_config, kv_transfer_manager=manager) + + expected = ["req_0_0_1_1"] + self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected) + self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected) + self.assertEqual(host.get_kv_rank_mapping()["from_tp"], 2) + self.assertEqual(host.get_kv_rank_mapping()["to_tp"], 2) + host.shutdown_omni_connectors() + def test_send_keys_route_from_rank_gt_to_rank(self): host = self._make_host(from_tp=4, to_tp=2, local_rank=3) self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0), ["req_0_0_3_1"]) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index b8a2fbd9b..514cc9b94 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -147,7 +147,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return None def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool: - """True when orchestrator will not run stage 1+ for this request (e.g. text-only). + """True when this stage-zero-final request does not need downstream KV. The result is cached per request to avoid repeated deserialization of additional_information on every scheduler tick. @@ -162,7 +162,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): result = False else: info = deserialize_additional_information(payload) - result = info.get("omni_final_stage_id") == 0 + result = info.get("omni_final_stage_id") == 0 and not bool(info.get("omni_force_kv_transfer", False)) self._omits_kv_transfer_cache[rid] = result return result diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index efed75241..8616cf9a6 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -403,6 +403,7 @@ class OmniKVTransferManager: self._prefetch_executor: ThreadPoolExecutor | None = ( ThreadPoolExecutor(max_workers=1, thread_name_prefix="kv-prefetch") if async_prefetch else None ) + self._prefetch_futures: dict[str, Any] = {} self._bg_copy_stream: current_omni_platform.Stream | None = None @@ -415,6 +416,11 @@ class OmniKVTransferManager: except Exception as e: logger.warning("Failed to eagerly initialize sender connector: %s", e) + @property + def tp_topology(self) -> KVTPTopology: + """Return the effective topology after config and runtime detection.""" + return self._tp_topo + @property def topo_config(self) -> _TransferTopoConfig: """Lazily build the parallel topology snapshot on first access.""" diff --git a/vllm_omni/engine/async_engine_utils.py b/vllm_omni/engine/async_engine_utils.py index 58f8351b0..6f544239b 100644 --- a/vllm_omni/engine/async_engine_utils.py +++ b/vllm_omni/engine/async_engine_utils.py @@ -77,12 +77,17 @@ def upgrade_to_omni_request( def apply_omni_final_stage_metadata( request: EngineCoreRequest, final_stage_id: int, + *, + force_kv_transfer: bool = False, ) -> EngineCoreRequest: - """Tag a request so the AR scheduler can skip DiT KV for stage zero.""" + """Tag a request with its final stage and optional KV-transfer override.""" merged: dict[str, Any] = {} if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None: merged = deserialize_additional_information(request.additional_information) merged["omni_final_stage_id"] = final_stage_id + merged.pop("omni_force_kv_transfer", None) + if force_kv_transfer: + merged["omni_force_kv_transfer"] = True payload = serialize_additional_information(merged) return OmniEngineCoreRequest.from_request( request, diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 454266b7a..730018f9e 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -798,11 +798,9 @@ class AsyncOmniEngine: supported_tasks=self.supported_tasks, ) request.external_req_id = cid - # Companions are stage-0-final: without this worker-visible tag - # the AR runner would still emit downstream connector payloads - # for them (the orchestrator-side final_stage_id=0 registration - # alone does not reach the worker). - request = apply_omni_final_stage_metadata(request, 0) + # Companions are stage-0-final for ordinary downstream payloads, + # but diffusion still needs their CFG KV caches. + request = apply_omni_final_stage_metadata(request, 0, force_kv_transfer=True) # Registration of this companion on stage-0's output processor is # deferred to Orchestrator._handle_add_companion, which routes diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 363b850d4..ae7a0c0fd 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -122,6 +122,19 @@ class OmniConnectorModelRunnerMixin: # -- heterogeneous TP rank support -- rank_cfg = self._parse_rank_mapping(model_config) + if self._kv_transfer_manager is not None: + topology = getattr(self._kv_transfer_manager, "tp_topology", None) + effective_mapping = ( + getattr(topology, "source_tp_size", None), + getattr(topology, "target_tp_size", None), + getattr(topology, "local_rank", None), + ) + if all(isinstance(value, int) for value in effective_mapping): + rank_cfg = { + "from_tp": effective_mapping[0], + "to_tp": effective_mapping[1], + "local_rank": effective_mapping[2], + } self._from_tp: int = rank_cfg["from_tp"] self._to_tp: int = rank_cfg["to_tp"] self._local_rank: int = rank_cfg["local_rank"] @@ -1275,6 +1288,10 @@ class OmniConnectorModelRunnerMixin: For heterogeneous TP receive, the local rank is the target rank and must fetch one or more source-rank shards keyed as ``from_rank -> to_rank``. """ + if self._from_tp <= 1 and self._to_tp <= 1: + resolved_to_stage = self._next_stage_id if to_stage is None else to_stage + return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] + remote_ranks = self.get_kv_remote_ranks() return [ self.get_kv_connector_key( @@ -1307,6 +1324,10 @@ class OmniConnectorModelRunnerMixin: chunk_id: int = 0, ) -> list[str]: """Build send-side connector keys for this rank's KV shard(s).""" + if self._from_tp <= 1 and self._to_tp <= 1: + resolved_to_stage = self._next_stage_id if to_stage is None else to_stage + return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] + target_ranks = self.get_kv_target_ranks_for_send() return [ self.get_kv_connector_key( From 6e8c279a68c5dda298d418c369d10af2a2a8035c Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:09:27 +0800 Subject: [PATCH 09/10] [Bugfix] Match image file metadata to output format (#5619) Signed-off-by: lishunyang12 Co-authored-by: lishunyang12 --- .../openai_api/test_image_server.py | 77 +++++++++++++++--- vllm_omni/entrypoints/openai/api_server.py | 78 +++++++++++-------- .../entrypoints/openai/image_api_utils.py | 3 +- .../entrypoints/openai/protocol/images.py | 19 ++++- 4 files changed, 129 insertions(+), 48 deletions(-) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 34ae5ba20..2d6c0299e 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -1301,7 +1301,11 @@ def test_generate_images_rejects_model_mismatch(test_client): assert "model mismatch" in response.json()["detail"].lower() -def test_image_file_response_format_multiple(test_client): +@pytest.mark.parametrize( + ("output_format", "expected_pil_format"), + [("png", "PNG"), ("jpeg", "JPEG"), ("webp", "WEBP")], +) +def test_image_file_response_format_multiple(test_client, output_format, expected_pil_format): """Test response_format=file with n>1 returns ZIP archive""" response = test_client.post( "/v1/images/generations", @@ -1309,6 +1313,7 @@ def test_image_file_response_format_multiple(test_client): "prompt": "a dog", "n": 3, "response_format": "file", + "output_format": output_format, }, ) @@ -1317,23 +1322,32 @@ def test_image_file_response_format_multiple(test_client): assert "attachment" in response.headers.get("content-disposition", "") assert ".zip" in response.headers.get("content-disposition", "") - # Verify it's a valid ZIP with 3 PNG files + # Verify it's a valid ZIP with 3 correctly labelled image files import zipfile zip_buffer = io.BytesIO(response.content) with zipfile.ZipFile(zip_buffer, "r") as zf: files = zf.namelist() assert len(files) == 3 - assert all(f.endswith(".png") for f in files) + assert all(f.endswith(f".{output_format}") for f in files) - # Verify each file is a valid PNG + # Verify each file's bytes match its advertised format for filename in files: img_bytes = zf.read(filename) img = Image.open(io.BytesIO(img_bytes)) - assert img.format == "PNG" + assert img.format == expected_pil_format -def test_image_file_response_format_single(test_client): +@pytest.mark.parametrize( + ("output_format", "expected_media_type", "expected_pil_format"), + [ + ("png", "image/png", "PNG"), + ("jpg", "image/jpeg", "JPEG"), + ("jpeg", "image/jpeg", "JPEG"), + ("webp", "image/webp", "WEBP"), + ], +) +def test_image_file_response_format_single(test_client, output_format, expected_media_type, expected_pil_format): """Test response_format=file with n=1 returns a single image file.""" response = test_client.post( "/v1/images/generations", @@ -1341,16 +1355,61 @@ def test_image_file_response_format_single(test_client): "prompt": "a dog", "n": 1, "response_format": "file", + "output_format": output_format, }, ) assert response.status_code == 200 - assert response.headers["content-type"] == "image/png" + assert response.headers["content-type"] == expected_media_type assert "attachment" in response.headers.get("content-disposition", "") - assert ".png" in response.headers.get("content-disposition", "") + assert f".{output_format}" in response.headers.get("content-disposition", "") img = Image.open(io.BytesIO(response.content)) - assert img.format == "PNG" + assert img.format == expected_pil_format + + +@pytest.mark.parametrize( + ("output_format", "image_count", "expected_media_type", "expected_pil_format"), + [ + ("jpeg", 1, "image/jpeg", "JPEG"), + ("webp", 2, "application/zip", "WEBP"), + ], +) +def test_multistage_image_file_response_uses_requested_format( + async_omni_test_client, + output_format, + image_count, + expected_media_type, + expected_pil_format, +): + engine = async_omni_test_client.app.state.engine_client + engine._images = [Image.new("RGB", (16, 16), color="green") for _ in range(image_count)] + + response = async_omni_test_client.post( + "/v1/images/generations", + json={ + "prompt": "a dog", + "n": image_count, + "response_format": "file", + "output_format": output_format, + }, + ) + + assert response.status_code == 200 + assert response.headers["content-type"] == expected_media_type + if image_count == 1: + assert f".{output_format}" in response.headers["content-disposition"] + assert Image.open(io.BytesIO(response.content)).format == expected_pil_format + return + + import zipfile + + with zipfile.ZipFile(io.BytesIO(response.content), "r") as archive: + filenames = archive.namelist() + assert len(filenames) == image_count + assert all(filename.endswith(f".{output_format}") for filename in filenames) + for filename in filenames: + assert Image.open(io.BytesIO(archive.read(filename))).format == expected_pil_format def make_test_image_bytes(size=(64, 64)) -> bytes: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 6d088f65e..2e92c864b 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -98,7 +98,6 @@ from vllm_omni.entrypoints.openai.duplex_capability import should_enable_duplex_ from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError from vllm_omni.entrypoints.openai.image_api_utils import ( SUPPORTED_LAYERED_RESOLUTIONS, - encode_image_base64, encode_image_base64_with_compression, parse_size, validate_layered_layers, @@ -1731,6 +1730,39 @@ async def show_available_models(raw_request: Request) -> JSONResponse: # Image generation API endpoints +def _build_image_generation_response( + *, + images: list[Image.Image], + request: ImageGenerationRequest, + stage_durations: Any, + peak_memory_mb: Any, +) -> ImageGenerationResponse | StreamingResponse: + """Encode generated images and apply the requested response format.""" + output_format = _choose_output_format(request.output_format or "png", None) + image_data = [ + ImageData( + b64_json=encode_image_base64_with_compression(image, format=output_format), + revised_prompt=None, + ) + for image in images + ] + response_kwargs: dict[str, Any] = { + "created": int(time.time()), + "data": image_data, + "output_format": output_format, + "metrics": { + "stage_durations": stage_durations or None, + "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, + }, + } + if request.size is not None: + response_kwargs["size"] = request.size + response = ImageGenerationResponse(**response_kwargs) + if request.response_format == ResponseFormat.FILE: + return response.stream_response() + return response + + @router.post( "/v1/images/generations", dependencies=[Depends(validate_json_request)], @@ -1832,15 +1864,11 @@ async def generate_images( content=generation_result.model_dump(), ) flat_images, stage_durations, peak_memory_mb, _ = generation_result - image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in flat_images] - - return ImageGenerationResponse( - created=int(time.time()), - data=image_data, - metrics={ - "stage_durations": stage_durations or None, - "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, - }, + return _build_image_generation_response( + images=flat_images, + request=request, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, ) # Build params - pass through user values directly @@ -1928,32 +1956,14 @@ async def generate_images( logger.debug(f"Successfully generated {len(images)} image(s)") - # Determine output format (default to png) - output_format = _choose_output_format(request.output_format or "png", None) - - # Encode images to base64 with the specified format - image_data = [ - ImageData(b64_json=encode_image_base64_with_compression(img, format=output_format), revised_prompt=None) - for img in images - ] - stage_durations = getattr(result, "stage_durations", None) peak_memory_mb = getattr(result, "peak_memory_mb", None) - response_kwargs = { - "created": int(time.time()), - "data": image_data, - "output_format": output_format, - "metrics": { - "stage_durations": stage_durations or None, - "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, - }, - } - if request.size: - response_kwargs["size"] = size_str - response = ImageGenerationResponse(**response_kwargs) - if request.response_format != ResponseFormat.FILE: - return response - return response.stream_response() + return _build_image_generation_response( + images=images, + request=request, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) except (EngineGenerateError, EngineDeadError) as exc: return _create_engine_error_json_response(raw_request, exc) diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py index 5bd79e936..ed759a080 100644 --- a/vllm_omni/entrypoints/openai/image_api_utils.py +++ b/vllm_omni/entrypoints/openai/image_api_utils.py @@ -101,13 +101,14 @@ def encode_image_base64_with_compression( """ buffer = io.BytesIO() image = prepare_image_for_output_format(image, format) + pil_format = "jpeg" if format == "jpg" else format save_kwargs = {} if format in ("jpg", "jpeg", "webp"): save_kwargs["quality"] = output_compression elif format == "png": save_kwargs["compress_level"] = max(0, min(9, 9 - output_compression // 11)) - image.save(buffer, format=format, **save_kwargs) + image.save(buffer, format=pil_format, **save_kwargs) buffer.seek(0) return base64.b64encode(buffer.read()).decode("utf-8") diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index e83c7b5c0..0a20ba1c8 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -21,6 +21,13 @@ from pydantic import BaseModel, Field, field_validator from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers +_IMAGE_FILE_METADATA = { + "jpg": ("jpg", "image/jpeg"), + "jpeg": ("jpeg", "image/jpeg"), + "png": ("png", "image/png"), + "webp": ("webp", "image/webp"), +} + class ResponseFormat(str, Enum): """Image response format""" @@ -183,7 +190,7 @@ class ImageGenerationResponse(BaseModel): created: int = Field(..., description="Unix timestamp of when the generation completed") data: list[ImageData] = Field(..., description="Array of generated images") - output_format: str = Field(None, description="The output format of the image generation") + output_format: str | None = Field(None, description="The output format of the image generation") size: str = Field(None, description="The size of the image generated") cot_output: str | None = Field( None, @@ -201,12 +208,16 @@ class ImageGenerationResponse(BaseModel): status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail="No image data available for file response.", ) + extension, media_type = _IMAGE_FILE_METADATA.get( + (self.output_format or "png").lower(), + _IMAGE_FILE_METADATA["png"], + ) if len(self.data) == 1: image_bytes = base64.b64decode(self.data[0].b64_json) - filename = f"image_{uuid.uuid4().hex[:8]}.png" + filename = f"image_{uuid.uuid4().hex[:8]}.{extension}" return StreamingResponse( io.BytesIO(image_bytes), - media_type="image/png", + media_type=media_type, headers={ "Content-Disposition": f'attachment; filename="{filename}"', "Content-Length": str(len(image_bytes)), @@ -217,7 +228,7 @@ class ImageGenerationResponse(BaseModel): with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf: for idx, item in enumerate(self.data): if item.b64_json: - zf.writestr(f"image_{idx}.png", base64.b64decode(item.b64_json)) + zf.writestr(f"image_{idx}.{extension}", base64.b64decode(item.b64_json)) zip_bytes = zip_buffer.getvalue() filename = f"images_{uuid.uuid4().hex[:8]}.zip" return StreamingResponse( From 4b1547b4d83cbbbbfc472bed6559240fcf26a594 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:17:33 +0800 Subject: [PATCH 10/10] [Bugfix] Fix Wan2.2 omitted guidance scale resolution (#5615) Signed-off-by: lishunyang12 Signed-off-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> Co-authored-by: lishunyang12 --- .../wan2_2/test_wan22_pipeline_diffuse.py | 37 +++++++++++-------- .../models/wan2_2/pipeline_wan2_2.py | 37 ++++++++++--------- .../models/wan2_2/pipeline_wan2_2_i2v.py | 19 +--------- vllm_omni/diffusion/request.py | 9 ++--- 4 files changed, 47 insertions(+), 55 deletions(-) diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 198c31908..8c06df034 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -9,7 +9,9 @@ import torch from torch import nn from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -91,7 +93,20 @@ def _make_pipeline() -> Wan22Pipeline: return pipeline -def test_forward_delegates_denoising_to_diffuse() -> None: +@pytest.mark.parametrize( + ("sampling_params_kwargs", "expected_low", "expected_high"), + [ + ({}, 4.0, 4.0), + ({"guidance_scale": 0.0}, 0.0, 0.0), + ({"guidance_scale": 1.0}, 1.0, 1.0), + ({"guidance_scale": 3.0, "guidance_scale_2": 5.0}, 3.0, 5.0), + ], +) +def test_forward_delegates_denoising_to_diffuse( + sampling_params_kwargs: dict[str, float], + expected_low: float, + expected_high: float, +) -> None: pipeline = _make_pipeline() captured: dict[str, object] = {} @@ -101,25 +116,15 @@ def test_forward_delegates_denoising_to_diffuse() -> None: pipeline.diffuse = _fake_diffuse # type: ignore[method-assign] - mock_req = SimpleNamespace( + mock_req = OmniDiffusionRequest( prompt="prompt", request_id="test-req", - sampling_params=SimpleNamespace( - height=None, - width=None, + sampling_params=OmniDiffusionSamplingParams( num_frames=1, num_inference_steps=2, - guidance_scale_provided=True, - guidance_scale=1.0, - guidance_scale_2=None, - boundary_ratio=None, - generator=None, - seed=None, - num_outputs_per_prompt=1, max_sequence_length=32, - latents=None, output_type="latent", - extra_args={}, + **sampling_params_kwargs, ), ) batch = DiffusionRequestBatch(requests=[mock_req]) @@ -129,8 +134,8 @@ def test_forward_delegates_denoising_to_diffuse() -> None: assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8))) assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8)) assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps) - assert captured["guidance_low"] == 1.0 - assert captured["guidance_high"] == 1.0 + assert captured["guidance_low"] == expected_low + assert captured["guidance_high"] == expected_high assert captured["boundary_timestep"] == pytest.approx(875.0) assert captured["latent_condition"] is None assert captured["first_frame_mask"] is None diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index e3617f84d..0fc9aca65 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -37,7 +37,7 @@ from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from vllm_omni.inputs.data import OmniTextPrompt +from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -84,6 +84,24 @@ def resolve_wan_flow_shift(req: OmniDiffusionRequest, od_config: OmniDiffusionCo raise ValueError(f"Invalid flow_shift={raw_flow_shift!r}. flow_shift must be a float.") from exc +def resolve_wan_guidance_scales( + sampling_params: OmniDiffusionSamplingParams, + default_guidance_scale: float, +) -> tuple[float, float]: + guidance_scale = ( + sampling_params.guidance_scale if sampling_params.guidance_scale_provided else default_guidance_scale + ) + guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0] + guidance_high = ( + sampling_params.guidance_scale_2 + if sampling_params.guidance_scale_2_provided and sampling_params.guidance_scale_2 is not None + else ( + guidance_scale[1] if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1 else guidance_low + ) + ) + return guidance_low, guidance_high + + def retrieve_latents( encoder_output: torch.Tensor, generator: torch.Generator | None = None, @@ -571,25 +589,10 @@ class Wan22Pipeline( width = (width // mod_value) * mod_value num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps - # Respect per-request guidance_scale when explicitly provided. - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale - else: - guidance_scale = 4.0 - output_type = req.sampling_params.output_type or "np" attention_kwargs: dict | None = None - guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0] - guidance_high = ( - req.sampling_params.guidance_scale_2 - if req.sampling_params.guidance_scale_2 is not None - else ( - guidance_scale[1] - if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1 - else guidance_low - ) - ) + guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=4.0) # record guidance for properties self._guidance_scale = guidance_low diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index d8c1d4a8c..180f8bb7a 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -36,6 +36,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( create_transformer_from_config, load_transformer_config, resolve_wan_flow_shift, + resolve_wan_guidance_scales, resolve_wan_sample_solver, retrieve_latents, ) @@ -475,26 +476,10 @@ class Wan22I2VPipeline( num_frames = req.sampling_params.num_frames or 81 num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps - # Respect per-request guidance_scale when explicitly provided. - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale - else: - guidance_scale = 5.0 - output_type = req.sampling_params.output_type or "np" attention_kwargs: dict | None = None - # Handle guidance scales - guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0] - guidance_high = ( - req.sampling_params.guidance_scale_2 - if req.sampling_params.guidance_scale_2 is not None - else ( - guidance_scale[1] - if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1 - else guidance_low - ) - ) + guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=5.0) self._guidance_scale = guidance_low self._guidance_scale_2 = guidance_high diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 9ea879212..e1147c3d6 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -41,11 +41,10 @@ class OmniDiffusionRequest: if self.sampling_params.generator is None and self.sampling_params.seed is None: self.sampling_params.seed = random.randint(0, 2**31 - 1) - # Detect whether user explicitly provided guidance_scale. - # The unset sentinel is None (see OmniDiffusionSamplingParams); any value - # the caller set -- including an explicit 0.0 (disable CFG) -- must be - # honored. Resolve before auto-filling guidance_scale_2, else the sentinel - # leaks into guidance_scale_2. + # Detect whether the caller explicitly provided guidance_scale before + # resolving the omitted value. ``0.0`` is API-valid and must not be + # treated as omission because that would unexpectedly enable CFG in + # pipelines with a model-specific default above zero. if self.sampling_params.guidance_scale is not None: self.sampling_params.guidance_scale_provided = True else: