anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge commit '6fe1ff7709c72c08bbfc33a7e14b73ec94135549' into codex-issue-7035-h100-determinism

Browse Source 
Signed-off-by: david6666666 <530634352@qq.com>
This commit is contained in:
david6666666 committed 2026-09-17 09:22:22 +00:00
2 parents 5eb51516db 6fe1ff7709
commit f9253ae1ad
96 files changed
+8480 -822
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/config/composable_parallel/test_factory_integration.py
+30 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Integration tests: strategy overlay on the real merge_pipeline_deploy seam."""
from __future__ import annotations
@@ -20,6 +20,7 @@ from vllm_omni.config.composable_parallel import (
apply_strategy_specs,
)
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
@@ -137,3 +138,31 @@ def test_cli_overrides_strategy_with_warning():
assert _resolved(stages, "talker").runtime.num_replicas == 3
# ...and the override was warned about, naming the conflicting field.
assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages)
def test_typed_cli_overrides_strategy():
pipeline = OMNI_PIPELINES["qwen2_5_omni"]
config = VllmOmniConfig.from_pipeline_config(
pipeline,
user_deploy_config=load_deploy_config(_DEPLOY),
cli_overrides={"stage_1_num_replicas": 3},
strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
)
talker = next(stage for stage in config.stage_configs if stage.model_stage == "talker")
assert talker.runtime_config.num_replicas == 3
assert config.orchestrator_config.omni_lb_policy == "round-robin"
assert config.strategy_omni_lb_policy == "round-robin"
def test_typed_tp_only_strategy_has_no_derived_lb_policy():
pipeline = OMNI_PIPELINES["qwen2_5_omni"]
config = VllmOmniConfig.from_pipeline_config(
pipeline,
user_deploy_config=load_deploy_config(_DEPLOY),
cli_overrides={"omni_lb_policy": "round-robin"},
strategy_specs={"thinker": [_tp(1)]},
)
assert config.orchestrator_config.omni_lb_policy == "round-robin"
assert config.strategy_omni_lb_policy is None
tests/config/test_forced_aligner_injection.py
+36
View File
@@ -2,13 +2,19 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
from vllm.pooling_params import PoolingParams
import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry)
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY
from vllm_omni.config.stage_config import (
DeployConfig,
StageExecutionType,
)
from vllm_omni.engine.stage_init_utils import (
build_engine_args_dict_from_omni_stage_config,
extract_stage_metadata_from_omni_stage_config,
)
from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
POOLING_OUTPUT_DECODER_PATH,
TIMESTAMPS_MODALITY,
@@ -94,3 +100,33 @@ def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_ch
assert stages[2].engine_args.runner == "pooling"
assert len(pipeline.stages) == 2
assert deploy.stages == []
def test_injected_aligner_survives_typed_config_and_runtime_projections():
pipeline, deploy = inject_forced_aligner_stage(
_PIPELINE_REGISTRY["qwen3_tts"],
DeployConfig(),
{"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"},
)
config = VllmOmniConfig.from_pipeline_config(
pipeline,
user_deploy_config=deploy,
cli_overrides={"model": "/models/Qwen3-TTS"},
)
stage = config.stage_configs[-1]
assert stage.pooling_config.runner == "pooling"
assert stage.pooling_config.pooling_output_decoder == POOLING_OUTPUT_DECODER_PATH
assert isinstance(stage.pooling_config.default_pooling_params, PoolingParams)
assert stage.pooling_config.default_pooling_params.task == deploy.stages[-1].default_pooling_params["task"]
metadata = extract_stage_metadata_from_omni_stage_config(stage)
assert isinstance(metadata.default_sampling_params, PoolingParams)
assert metadata.default_sampling_params.task == deploy.stages[-1].default_pooling_params["task"]
engine_args = build_engine_args_dict_from_omni_stage_config(
stage,
stage.model_config.model or "/models/Qwen3-ForcedAligner-0.6B",
)
assert engine_args["runner"] == "pooling"
assert engine_args["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH
tests/config/test_omni_config.py
+54 -1
View File
@@ -11,6 +11,7 @@ from pathlib import Path
import msgspec
import pytest
import torch
from pydantic import ValidationError
from pydantic.fields import FieldInfo
from transformers import Qwen3OmniMoeConfig
@@ -123,6 +124,18 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
def test_nested_stage_override_deep_merges_structured_model_config() -> None:
config = _from_pipeline_key(
"cosmos3_policy",
deploy_config_path=get_deploy_config_path("cosmos3_policy_droid.yaml"),
cli_overrides={"stage_0_model_config": {"guardrails": False}},
)
model_config = config.stage_by_id(0).diffusion_config.model_config
assert model_config["guardrails"] is False
assert model_config["policy_server_config"]["action_space"] == "joint_position"
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
pipeline = _resolve_pipeline_or_skip(model_type)
@@ -654,6 +667,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
"load_config",
"cache_config",
"scheduler_config",
"pooling_config",
"connector_config",
"runtime_config",
"parallel_config",
@@ -889,6 +903,17 @@ def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_saf
msgspec.msgpack.encode(payload)
def test_structured_diffusion_torch_dtype_is_msgpack_transport_safe():
stage_config = _from_pipeline_key(
"hunyuan_image3_dit",
cli_overrides={"dtype": torch.bfloat16},
).stage_by_id(0)
serialized = _serialize_stage_config(stage_config)
assert serialized["diffusion_config"]["dtype"] == "bfloat16"
msgspec.msgpack.encode(serialized)
def test_diffusion_parallel_config_fields_cover_legacy_surface():
from vllm_omni.diffusion.data import DiffusionParallelConfig
@@ -1480,7 +1505,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model=str(tmp_path))
od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
@@ -1739,3 +1764,31 @@ def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder):
builder(pipeline, deploy)
else:
builder(pipeline, user_deploy_config=deploy)
@pytest.mark.parametrize("explicit", [False, True])
def test_diffusion_quantization_origin_survives_projection_and_transport(monkeypatch, explicit):
from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
from vllm_omni.quantization import build_quant_config
checkpoint = TransformerConfig.from_dict(
{
"quantization_config": {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"w4a8_fallback_steps": [37],
}
}
)
requested = build_quant_config("mxfp4", w4a8_fallback_steps=[]) if explicit else None
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
tf_model_config=checkpoint,
quantization_config=requested,
)
assert cfg.quantization_config_is_auto_detected is not explicit
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda self: self.set_tf_model_config(checkpoint))
cfg.enrich_config()
restored = ForkingPickler.loads(ForkingPickler.dumps(cfg))
assert restored.quantization_config_is_auto_detected is not explicit
assert restored.quantization_config.w4a8_fallback_steps == ([] if explicit else [37])
tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py
+679
View File
@@ -0,0 +1,679 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import json
from collections.abc import Callable, Iterable
import pytest
import torch
import torch.nn as nn
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig
from vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint import (
FASTH3_BASE_MODEL,
FASTH3_V2_BASE_SCHEDULE,
FASTH3_V2_MODEL_ID,
FastH3CheckpointSpec,
)
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3DiTArchConfig,
MiniMaxH3DiTModel,
)
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
resolve_minimax_h3_diffusion_model_path,
)
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
from vllm_omni.errors import OmniClientError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
is_minimax_h3_modular,
resolve_minimax_h3_model_root,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _fastvideo_metadata() -> dict[str, object]:
"""Return the release-owned metadata read from fastvideo_inference.json."""
return {
"model_id": FASTH3_V2_MODEL_ID,
"schema_version": "fasth3-inference-contract-v1",
"dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
"transformer_forwards": 8,
"video_scheduler_shift": 10.0,
"audio_scheduler_shift": 3.0,
"guidance_scale": 1.0,
"attention_backend": "VIDEO_SPARSE_ATTN_H3",
"vsa_sparsity": 0.8,
"vsa_tile_size": 64,
"task": "t2av",
}
def _sampling(
*,
num_inference_steps: int | str | None = 8,
guidance_scale: float | None = 1.0,
extra_args: dict[str, object] | None = None,
lora_request: LoRARequest | None = None,
timesteps: torch.Tensor | None = None,
sigmas: list[float] | None = None,
) -> OmniDiffusionSamplingParams:
return OmniDiffusionSamplingParams(
# Keep the string form to exercise malformed client input validation.
num_inference_steps=num_inference_steps, # type: ignore[arg-type]
guidance_scale=guidance_scale,
extra_args=extra_args or {},
lora_request=lora_request,
timesteps=timesteps,
sigmas=sigmas,
)
def _od_config(
*,
backend: str = "FASTVIDEO_VSA",
topk: int | None = None,
per_role: dict[str, AttentionSpec] | None = None,
lora_path: str | None = None,
ring_degree: int = 1,
allgather_degree: int = 1,
) -> OmniDiffusionConfig:
attention_config = AttentionConfig(
default=AttentionSpec(backend=backend, fastvideo_vsa_topk=topk),
per_role=per_role or {},
)
return OmniDiffusionConfig(
diffusion_attention_config=attention_config,
lora_path=lora_path,
parallel_config=DiffusionParallelConfig(ring_degree=ring_degree, allgather_degree=allgather_degree),
)
def test_v2_release_has_nine_sigma_nodes_and_eight_intervals():
assert FASTH3_V2_BASE_SCHEDULE.base_schedule == (
0.999,
0.874,
0.749,
0.624,
0.5,
0.375,
0.25,
0.125,
0.0,
)
assert len(FASTH3_V2_BASE_SCHEDULE.base_schedule) == 9
assert FASTH3_V2_BASE_SCHEDULE.num_inference_steps == 8
assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(10.0) == pytest.approx(
[
0.9998999099189271,
0.985788405143244,
0.9675752486758817,
0.94316807738815,
0.9090909090909091,
0.8571428571428571,
0.7692307692307693,
0.5882352941176471,
0.0,
],
abs=1e-12,
)
assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(3.0) == pytest.approx(
[
0.9996664442961973,
0.9541484716157204,
0.8995196156925539,
0.8327402135231315,
0.75,
0.6428571428571429,
0.5,
0.3,
0.0,
],
abs=1e-12,
)
def test_fastvideo_metadata_builds_release_sampling_metadata():
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
assert spec.vsa_sparsity == 0.8
assert spec.release_metadata() == {
"partition": "fl2va",
"tasks": ["t2va"],
"sigma_shift_scales": {"video": 10.0, "audio": 3.0},
"base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
}
@pytest.mark.parametrize(
("field", "value"),
[
("model_id", "FastVideo/FastH3-other"),
("schema_version", "fasth3-inference-contract-v0"),
("dmd_denoising_steps", [999, 749, 500, 250]),
("transformer_forwards", 4),
("video_scheduler_shift", 12.0),
("audio_scheduler_shift", 4.0),
("guidance_scale", 2.0),
("attention_backend", "FLASH_ATTN"),
("vsa_sparsity", 0.9),
("vsa_tile_size", 32),
("task", "t2va"),
],
)
def test_fastvideo_metadata_rejects_contract_drift(field, value):
metadata = _fastvideo_metadata()
metadata[field] = value
with pytest.raises(ValueError, match=field):
FastH3CheckpointSpec.from_metadata(metadata)
def test_fastvideo_metadata_rejects_missing_contract_field():
metadata = _fastvideo_metadata()
metadata.pop("task")
with pytest.raises(ValueError, match="task"):
FastH3CheckpointSpec.from_metadata(metadata)
def test_legacy_h3_metadata_is_not_treated_as_fastvideo_release():
with pytest.raises(ValueError, match="model_id"):
FastH3CheckpointSpec.from_metadata(
{
"schema_version": 1,
"partition": "fl2va",
"tasks": ["t2va", "fl2va"],
"sigma_shift_scales": {"video": 12.0, "audio": 3.0},
}
)
def test_v2_serving_contract_accepts_backend_selected_for_h3_self_role():
config = _od_config(
backend="TORCH_SDPA",
per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA")},
)
FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).check_serving_contract(
partition="fl2va", od_config=config
)
@pytest.mark.parametrize(
("kwargs", "partition", "match"),
[
({}, "ref2va", "task-type fl2va"),
({"lora_path": "/tmp/adapter"}, "fl2va", "LoRA"),
({"backend": "TORCH_SDPA"}, "fl2va", "FASTVIDEO_VSA"),
({"topk": 2}, "fl2va", "fixed fastvideo_vsa_topk"),
({"ring_degree": 2}, "fl2va", "local attention"),
({"allgather_degree": 2}, "fl2va", "local attention"),
],
)
def test_v2_serving_contract_rejects_incompatible_runtime(kwargs, partition, match):
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
with pytest.raises(ValueError, match=match):
spec.check_serving_contract(partition=partition, od_config=_od_config(**kwargs))
def test_v2_serving_contract_reads_a_per_role_fixed_topk_override():
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
config = _od_config(
per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=1)},
)
with pytest.raises(ValueError, match="fixed fastvideo_vsa_topk"):
spec.check_serving_contract(partition="fl2va", od_config=config)
def test_v2_request_contract_accepts_omitted_steps_and_default_shifts():
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
spec.check_request(_sampling(num_inference_steps=None))
spec.check_request(_sampling(extra_args={"flow_shift": 10, "audio_flow_shift": 3}, guidance_scale=1))
@pytest.mark.parametrize(
"sampling",
[
_sampling(num_inference_steps=None),
_sampling(timesteps=torch.tensor([0.9, 0.5, 0.0])),
_sampling(sigmas=[1.0, 0.5, 0.0]),
],
)
def test_v2_step_execution_requires_pinned_count_and_schedule(sampling):
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
with pytest.raises(OmniClientError, match="step execution"):
spec.check_request(sampling, step_execution=True)
def test_v2_step_execution_accepts_only_explicit_pinned_count():
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
spec.check_request(_sampling(num_inference_steps=8), step_execution=True)
@pytest.mark.parametrize(
("kwargs", "match"),
[
({"num_inference_steps": 7}, "num_inference_steps=8"),
({"num_inference_steps": 9}, "num_inference_steps=8"),
({"num_inference_steps": "8"}, "num_inference_steps=8"),
({"extra_args": {"flow_shift": 9}}, "flow_shift=10"),
({"extra_args": {"audio_flow_shift": 4}}, "audio_flow_shift=3"),
({"extra_args": {"flow_shift": "bad"}}, "flow_shift=10"),
({"guidance_scale": 2.0}, "guidance_scale=1"),
(
{"lora_request": LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/test.safetensors")},
"per-request LoRA",
),
],
)
def test_v2_request_contract_rejects_sampling_drift(kwargs, match):
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
with pytest.raises(OmniClientError, match=match):
spec.check_request(_sampling(**kwargs))
def test_resolve_native_vaes_uses_pinned_revision_and_only_vae_patterns(tmp_path, monkeypatch):
model_root = tmp_path / "FastH3"
model_root.mkdir()
revision = "base-revision-123"
(model_root / "provenance.json").write_text(
json.dumps({"base_model": f"hf://{FASTH3_BASE_MODEL}@{revision}"}),
encoding="utf-8",
)
snapshot = tmp_path / "hf-snapshot"
calls: dict[str, object] = {}
def fake_download(**kwargs: object) -> str:
calls.update(kwargs)
snapshot.mkdir()
return str(snapshot)
monkeypatch.setattr(
"vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint.download_weights_from_hf_specific",
fake_download,
)
result = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
assert result == snapshot / "FL2VA"
assert calls == {
"model_name_or_path": FASTH3_BASE_MODEL,
"cache_dir": None,
"allow_patterns": ["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
"revision": revision,
"require_all": True,
}
assert sorted(path.name for path in model_root.iterdir()) == ["provenance.json"]
assert not (model_root / "model_index.json").exists()
assert not (model_root / "conversion.json").exists()
@pytest.mark.parametrize(
"base_model",
[
"MiniMaxAI/MiniMax-H3@main",
"hf://MiniMaxAI/MiniMax-H3@",
"hf://Other/MiniMax-H3@revision",
"",
None,
],
)
def test_resolve_native_vaes_rejects_unpinned_provenance(tmp_path, base_model):
model_root = tmp_path / "FastH3"
model_root.mkdir()
(model_root / "provenance.json").write_text(json.dumps({"base_model": base_model}), encoding="utf-8")
with pytest.raises(ValueError, match="provenance"):
FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
def test_modular_local_model_root_uses_text_encoder_without_partition_files(tmp_path):
root = tmp_path / "FastH3"
(root / "text_encoder").mkdir(parents=True)
(root / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
(root / "modular_model_index.json").write_text(
json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
encoding="utf-8",
)
assert is_minimax_h3_modular(str(root)) is True
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "text_encoder")
assert not (root / "FL2VA").exists()
def test_modular_hub_model_root_downloads_only_text_encoder(monkeypatch, tmp_path):
from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
snapshot = tmp_path / "modular-snapshot"
calls: dict[str, object] = {}
def fake_index(model: str, *, revision: str | None = None) -> dict[str, str]:
assert model == "FastVideo/FastVideo-FastH3-8-Step-V2"
assert revision == "release-revision"
return {"_class_name": "MiniMaxH3ModularPipeline"}
def fake_download(**kwargs: object) -> str:
calls.update(kwargs)
snapshot.mkdir()
return str(snapshot)
monkeypatch.setattr(checkpoint_module, "get_diffusion_model_index", fake_index)
monkeypatch.setattr(checkpoint_module, "download_weights_from_hf_specific", fake_download)
result = resolve_minimax_h3_model_root(
"FastVideo/FastVideo-FastH3-8-Step-V2",
"release-revision",
"fl2va",
)
assert result == str(snapshot / "text_encoder")
assert calls == {
"model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
"cache_dir": None,
"allow_patterns": ["text_encoder/**"],
"revision": "release-revision",
"require_all": True,
}
def test_native_hub_model_with_modular_class_keeps_partition_layout(monkeypatch):
from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
monkeypatch.setattr(
checkpoint_module,
"get_diffusion_model_index",
lambda _model, *, revision=None: {"_class_name": "MiniMaxH3ModularPipeline"},
)
assert checkpoint_module.is_minimax_h3_modular("MiniMaxAI/MiniMax-H3") is False
def test_native_local_model_root_keeps_partition_layout(tmp_path):
root = tmp_path / "MiniMax-H3"
(root / "FL2VA" / "text_encoder").mkdir(parents=True)
(root / "FL2VA" / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
assert is_minimax_h3_modular(str(root)) is False
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
def test_diffusion_resolver_returns_local_modular_root_without_partition_files(tmp_path):
root = tmp_path / "FastH3"
root.mkdir()
(root / "modular_model_index.json").write_text(
json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
encoding="utf-8",
)
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(root)
assert not (root / "FL2VA").exists()
assert not (root / "Ref2VA").exists()
def test_diffusion_resolver_hub_modular_downloads_transformer_contract_only(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
snapshot = tmp_path / "modular-snapshot"
calls: dict[str, object] = {}
def fake_modular(_model: str, _revision: str | None = None) -> bool:
return True
def fake_download(**kwargs: object) -> str:
calls.update(kwargs)
snapshot.mkdir()
return str(snapshot)
monkeypatch.setattr(pipeline_module, "is_minimax_h3_modular", fake_modular)
monkeypatch.setattr(pipeline_module, "download_weights_from_hf_specific", fake_download)
result = resolve_minimax_h3_diffusion_model_path(
"FastVideo/FastVideo-FastH3-8-Step-V2", "release-revision", "fl2va"
)
assert result == str(snapshot)
assert calls == {
"model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
"cache_dir": None,
"allow_patterns": [
"modular_model_index.json",
"fastvideo_inference.json",
"provenance.json",
"transformer/**",
],
"revision": "release-revision",
"require_all": True,
}
patterns = calls["allow_patterns"]
assert isinstance(patterns, list)
assert all("video_vae" not in pattern and "audio_vae" not in pattern for pattern in patterns)
def test_diffusers_arch_config_normalizes_all_aliases():
arch = MiniMaxH3DiTArchConfig.from_mapping(
{
"num_refiner_layers": 3,
"ffn_dim": 12,
"in_channels": 4,
"audio_in_channels": 6,
"freq_dim": 8,
"time_embed_hidden_dim": 10,
"rope_freq_dim": 5,
"patch_size": [1, 2, 2],
}
)
assert arch.token_refiner_num_layers == 3
assert arch.ffn_hidden_size == 12
assert arch.latents_dim == 4
assert arch.audio_latents_dim == 6
assert arch.timestep_input_dim == 8
assert arch.time_embed_hidden_size == 10
assert arch.rope_inv_freq_len == 5
assert arch.patch_size == (1, 2, 2)
class _WeightTarget(nn.Module):
def __init__(self, shape: tuple[int, ...], loader: Callable[..., None]) -> None:
super().__init__()
self.weight = nn.Parameter(torch.zeros(shape))
self.weight.weight_loader = loader
class _PipelineLoadTransformer(nn.Module):
def __init__(self, loaded: set[str]) -> None:
super().__init__()
self.arch = MiniMaxH3DiTArchConfig(num_layers=1)
self._loaded = loaded
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
list(weights)
return self._loaded
def post_load_weights(self) -> None:
pass
def test_pipeline_load_weights_rejects_missing_fast_h3_compression_gate():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
nn.Module.__init__(pipeline)
pipeline.transformer = _PipelineLoadTransformer({"blocks.0.attn.qkv_proj.weight"})
pipeline.video_vae = None
pipeline.audio_vae = None
pipeline._fasth3 = None
pipeline._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
with pytest.raises(ValueError, match="missing compression gates"):
pipeline.load_weights([("transformer.blocks.0.attn.qkv_proj.weight", torch.zeros(1))])
def _load_fixture(
*, diffusers_weights: bool
) -> tuple[MiniMaxH3DiTModel, list[tuple[str | None, torch.Tensor]], list[tuple[int, torch.Tensor]]]:
model = object.__new__(MiniMaxH3DiTModel)
nn.Module.__init__(model)
model.arch = MiniMaxH3DiTArchConfig(
num_layers=1,
hidden_size=3,
num_attention_heads=2,
attention_head_dim=2,
ffn_hidden_size=6,
rope_inv_freq_len=2,
)
model._diffusers_weights = diffusers_weights
model._rope_theta = 10000.0
qkv_calls: list[tuple[str | None, torch.Tensor]] = []
fc1_calls: list[tuple[int, torch.Tensor]] = []
def qkv_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str | None = None) -> None:
if shard_id is None:
assert loaded_weight.shape == (12, 3)
param.data.copy_(loaded_weight)
qkv_calls.append((None, loaded_weight.clone()))
return
assert loaded_weight.shape == (4, 3)
assert shard_id in {"q", "k", "v"}
offset = {"q": 0, "k": 4, "v": 8}[shard_id]
param.data[offset : offset + 4].copy_(loaded_weight)
qkv_calls.append((shard_id, loaded_weight.clone()))
def fc1_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: int) -> None:
assert loaded_weight.shape == (3, 3)
param.data[shard_id * 3 : (shard_id + 1) * 3].copy_(loaded_weight)
fc1_calls.append((shard_id, loaded_weight.clone()))
model.blocks = nn.ModuleList([nn.Module()])
model.blocks[0].attn = nn.Module()
model.blocks[0].attn.qkv_proj = _WeightTarget((12, 3), qkv_loader)
model.blocks[0].mlp = nn.Module()
model.blocks[0].mlp.fc1 = _WeightTarget((6, 3), fc1_loader)
model.rope = nn.Module()
model.rope.register_buffer("inv_freq", torch.zeros(2))
return model, qkv_calls, fc1_calls
def test_diffusers_load_weights_uses_shape_aware_qkv_and_mlp_slices():
model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=True)
q = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 10
k = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100
v = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 200
mlp = torch.arange(18, dtype=torch.float32).reshape(6, 3) + 300
loaded = model.load_weights(
[
("transformer_blocks.0.attn.to_q.weight", q),
("transformer_blocks.0.attn.to_k.weight", k),
("transformer_blocks.0.attn.to_v.weight", v),
("transformer_blocks.0.ff.net.0.proj.weight", mlp),
]
)
assert loaded == {
"blocks.0.attn.qkv_proj.weight",
"blocks.0.mlp.fc1.weight",
"rope.inv_freq",
}
assert [shard_id for shard_id, _ in qkv_calls] == ["q", "k", "v"]
assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, torch.cat((q, k, v)))
torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, torch.cat((mlp[3:], mlp[:3])))
torch.testing.assert_close(
model.rope.inv_freq,
1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)),
rtol=0,
atol=0,
)
def test_native_load_weights_keeps_grouped_qkv_and_gate_first_layouts():
model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=False)
grouped_qkv = torch.arange(36, dtype=torch.float32).reshape(12, 3)
fc1 = torch.arange(18, dtype=torch.float32).reshape(6, 3)
loaded = model.load_weights(
[
("blocks.0.attn.qkv_proj.weight", grouped_qkv),
("blocks.0.mlp.fc1.weight", fc1),
]
)
expected_qkv = torch.cat(
(
grouped_qkv[:2],
grouped_qkv[6:8],
grouped_qkv[2:4],
grouped_qkv[8:10],
grouped_qkv[4:6],
grouped_qkv[10:],
)
)
assert loaded == {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight"}
assert [shard_id for shard_id, _ in qkv_calls] == [None]
assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, expected_qkv)
torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, fc1)
assert torch.equal(model.rope.inv_freq, torch.zeros(2))
def test_diffusers_load_weights_rejects_unknown_name():
model, _, _ = _load_fixture(diffusers_weights=True)
with pytest.raises(ValueError, match="unsupported Diffusers H3 weight"):
model.load_weights([("transformer_blocks.0.attn.not_a_projection.weight", torch.zeros(4, 3))])
def test_diffusers_load_weights_rejects_duplicate_source_name():
model, _, _ = _load_fixture(diffusers_weights=True)
weight = torch.zeros(4, 3)
with pytest.raises(ValueError, match="duplicate Diffusers H3 weight"):
model.load_weights(
[
("transformer_blocks.0.attn.to_q.weight", weight),
("transformer_blocks.0.attn.to_q.weight", weight),
]
)
def test_diffusers_load_weights_rejects_incomplete_qkv_group():
model, _, _ = _load_fixture(diffusers_weights=True)
with pytest.raises(ValueError, match="incomplete Diffusers H3 QKV"):
model.load_weights(
[
("transformer_blocks.0.attn.to_q.weight", torch.zeros(4, 3)),
("transformer_blocks.0.attn.to_k.weight", torch.zeros(4, 3)),
]
)
def test_diffusers_load_weights_rejects_missing_compression_gate_parameter():
model, _, _ = _load_fixture(diffusers_weights=True)
with pytest.raises(ValueError, match="has no model parameter"):
model.load_weights([("transformer_blocks.0.attn.to_gate_compress.weight", torch.zeros(4, 3))])
def test_resolve_sigma_positions_uses_release_metadata_schedule():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
nn.Module.__init__(pipeline)
pipeline._fasth3 = None
pipeline._lora_sigma_schedules = {}
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
pipeline._base_schedule_by_partition = {
"fl2va": DMD2SigmaSchedule.from_metadata(spec.release_metadata()),
}
positions, steps = pipeline._resolve_sigma_positions("t2va", _sampling())
assert positions == FASTH3_V2_BASE_SCHEDULE.base_schedule
assert steps == 8 == len(positions) - 1
with pytest.raises(OmniClientError, match="must be 8"):
pipeline._resolve_sigma_positions("t2va", _sampling(num_inference_steps=9))
tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+1
View File
@@ -423,6 +423,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
assert pipeline.partition == expected_partition
assert pipeline.supported_tasks == expected_tasks
assert len(created["dit"]) == expected_dits
assert all(kwargs["diffusers_weights"] is False for _, kwargs in created["dit"])
component_path = tmp_path / component_partition
assert created["video_vae"] == [str(component_path / "video_vae")]
assert created["audio_vae"] == [str(component_path / "audio_vae")]
tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py
+4 -1
View File
@@ -262,7 +262,10 @@ def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path):
path = tmp_path / "fasth3" / "adapter_model.safetensors"
_write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))})
_write_adapter(
path,
tensors={"transformer_blocks.0.attn.not_a_projection.lora_A.weight": torch.ones((_RANK, _HIDDEN))},
)
# Dropping it silently would load a model that is not the distilled student.
with pytest.raises(FastH3AdapterError, match="name no known"):
_load(path.parent)
tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py
+22 -1
View File
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from unittest.mock import Mock
@@ -41,7 +42,7 @@ class _FakeAttention(nn.Module):
super().__init__()
def _small_od_config():
def _small_od_config(*, class_name=None):
arch = {
"num_layers": 1,
"token_refiner_num_layers": 1,
@@ -60,6 +61,8 @@ def _small_od_config():
"final_adaln_out_features": 2 * 8,
"rope_inv_freq_len": 2,
}
if class_name is not None:
arch["_class_name"] = class_name
return SimpleNamespace(
tf_model_config=arch,
parallel_config=SimpleNamespace(ulysses_degree=1),
@@ -133,6 +136,24 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
)
def test_explicit_native_weight_format_overrides_diffusers_class_name(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "Attention", _FakeAttention)
monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
model = h3.MiniMaxH3DiTModel(
_small_od_config(class_name="MiniMaxH3Transformer3DModel"),
diffusers_weights=False,
)
assert model._diffusers_weights is False
class _WeightTarget(nn.Module):
def __init__(self, loader):
super().__init__()
tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
+7 -1
View File
@@ -187,8 +187,14 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
"model": model_ref,
"enforce_eager": True,
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
"stage_0_tensor_parallel_size": 2,
# This test loads the direct FL2VA diffusion pipeline (not the
# disaggregated two-stage MiniMax-H3 pipeline), so configure both the
# DiT world and the Qwen3-VL encoder TP explicitly.
"tensor_parallel_size": 2,
"text_encoder_tp_size": 2,
"vae_use_tiling": True,
# MiniMax H3's VAE is implemented by code shipped with the checkpoint.
"trust_remote_code": True,
}
omni_bl = Omni(**common_kwargs)
tests/diffusion/models/wan2_2/test_create_transformer_quant.py
+212 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for transformer quant-config auto-detection.
The loader path at pipeline_wan2_2.py carries two quantization contracts:
@@ -27,6 +27,7 @@ from types import SimpleNamespace
import pytest
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_
# od_config must remain unchanged — _create_transformer does not modify it.
assert od_config.quantization_config is None
def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch):
import json
import torch
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
for component in ("transformer", "transformer_2"):
directory = tmp_path / component
directory.mkdir()
config = dict(_MIN_CFG)
expert_quant_config = (
disk_quant_config[0 if component == "transformer" else 1]
if isinstance(disk_quant_config, tuple)
else disk_quant_config
)
if expert_quant_config is not None:
config["quantization_config"] = expert_quant_config
(directory / "config.json").write_text(json.dumps(config))
(tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]}))
class FakeComponent(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
class FakeTransformer(FakeComponent):
def __init__(self, quant_config=None, **kwargs):
super().__init__()
self.received_quant_config = quant_config
for module in (wan22_module, i2v_module):
monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent())
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer)
pipeline_cls = {
"t2v": wan22_module.Wan22Pipeline,
"i2v": i2v_module.Wan22I2VPipeline,
"vace": vace_module.Wan22VACEPipeline,
}[pipeline_kind]
monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None)
od_config = OmniDiffusionConfig(
model=str(tmp_path),
model_class_name=pipeline_cls.__name__,
dtype=torch.bfloat16,
quantization_config=active,
boundary_ratio=0.875,
flow_shift=12.0,
enable_diffusion_pipeline_profiler=False,
)
od_config.enrich_config()
pipeline = pipeline_cls(od_config=od_config)
return pipeline
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
@pytest.mark.parametrize("offline", [False, True])
def test_pipeline_constructor_routes_each_expert_quantization(
pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker
):
"""Run the actual pipeline constructors, expert factories and quant-method dispatch."""
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.quantization import build_quant_config
active = build_quant_config(
{
"transformer": {
"method": method,
"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
"w4a8_fallback_steps": [0, 1],
},
"transformer_2": {
"method": method,
"w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
"w4a8_fallback_steps": [],
},
}
)
disk_quant_config = None
if offline:
disk_quant_config = {
"quant_method": method,
flag: True,
"ignored_layers": ["blocks.12.ffn.net_2"],
"w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"],
"w4a8_fallback_steps": [9],
}
pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch)
for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)):
received = transformer.received_quant_config
assert received.get_name() == method
assert getattr(received, flag) == offline
assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else [])
selected = f"blocks.{10 + index}.attn1.to_qkv"
other = f"blocks.{11 - index}.attn1.to_qkv"
assert received.w4a8_fallback_layers == [selected]
assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer
assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer
if offline:
assert isinstance(
received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod
)
assert pipeline.od_config.quantization_config is active
@pytest.mark.parametrize("explicit_none", [False, True])
def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch):
from vllm_omni.quantization import build_quant_config
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None}
active = build_quant_config(components)
pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
pipeline._create_transformer(_MIN_CFG, component="transformer_2")
assert captured[-1].get("quant_config") is None
with pytest.raises(ValueError, match="disabled for component 'transformer_2'"):
pipeline._create_transformer(
{**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}},
component="transformer_2",
)
def test_pipeline_component_default_reaches_expert_factory(monkeypatch):
from vllm_omni.quantization import build_quant_config
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}})
pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
pipeline._create_transformer(_MIN_CFG, component="transformer_2")
assert captured[-1]["quant_config"] is active.default_config
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
@pytest.mark.parametrize(
"configured_expert, with_default",
[
("transformer", False),
("transformer", True),
("transformer_2", False),
("transformer_2", True),
],
)
def test_pipeline_partial_expert_config_does_not_match_sibling(
pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch
):
from vllm_omni.quantization import build_quant_config
component_spec = {
configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
}
if with_default:
component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]}
active = build_quant_config(component_spec)
pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch)
sibling = "transformer_2" if configured_expert == "transformer" else "transformer"
selected = getattr(pipeline, configured_expert).received_quant_config
other = getattr(pipeline, sibling).received_quant_config
assert selected is active.component_configs[configured_expert]
assert other is active.default_config
if with_default:
assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"]
else:
assert other is None
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
@pytest.mark.parametrize("different_ignored_layers", [False, True])
def test_enriched_auto_quantization_preserves_each_expert_policy(
pipeline_kind, different_ignored_layers, tmp_path, monkeypatch
):
"""Startup auto-detection must not turn the first expert into a user override."""
high = {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"ignored_layers": ["blocks.0.ffn.net_2"],
"w4a8_fallback_steps": [],
"w4a8_fallback_layers": [],
"mxfp4_scale_alg": 0,
"require_smooth_scale": False,
}
low = {
**high,
"ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"],
"w4a8_fallback_steps": [37],
"w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
"mxfp4_scale_alg": 2,
"require_smooth_scale": True,
}
pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch)
assert pipeline.od_config.quantization_config is not None # enrich_config really detected the first expert.
for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]:
config = transformer.received_quant_config
for policy in (
"ignored_layers",
"w4a8_fallback_steps",
"w4a8_fallback_layers",
"mxfp4_scale_alg",
"require_smooth_scale",
):
assert getattr(config, policy) == expected[policy], policy
tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
+97 -1
View File
@@ -20,9 +20,11 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
def _cpu_platform(monkeypatch):
def _cpu_pipeline_runtime(monkeypatch):
# These lightweight pipelines have no accelerator or distributed groups.
module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1))
class _StubTransformer(nn.Module):
@@ -477,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
assert torch.equal(result, torch.full_like(latents, 10.0))
@pytest.mark.parametrize("fail_first_request", [False, True])
def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch):
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.forward_context import (
ForwardContext,
get_forward_context,
is_forward_context_available,
override_forward_context,
)
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
contexts: list[ForwardContext] = []
events = []
class InjectedPredictionError(Exception):
pass
class RecordingTransformer(_StubTransformer):
def __init__(self, expert):
super().__init__()
self.expert = expert
def forward(self, hidden_states, encoder_hidden_states, **kwargs):
context = get_forward_context()
assert context is contexts[-1]
request_index = len(contexts) - 1
step = context.denoise_step_idx
positive = bool(encoder_hidden_states[0, 0, 0] > 0)
events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2])))
if fail_first_request and request_index == 0 and step == 1 and not positive:
raise InjectedPredictionError
return (torch.ones_like(hidden_states) * (1 if positive else -1),)
def encode_prompt(**kwargs):
context = get_forward_context()
assert context.denoise_step_idx is None
contexts.append(context)
embeds = torch.ones(1, 8, 8)
return embeds, -embeds
pipeline = _make_pipeline()
pipeline.transformer = RecordingTransformer("high")
pipeline.transformer_2 = RecordingTransformer("low")
pipeline.scheduler = _StubScheduler([900, 500, 100])
monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt)
monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents)
# Use the real runner context manager, Wan forward/diffuse, and CFG
# dispatch. Only model computation, scheduler math and unrelated KV I/O
# are replaced; no test code publishes a denoise step.
runner = object.__new__(DiffusionModelRunner)
runner.pipeline = pipeline
runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32)
runner.vllm_config = None
runner.cache_backend = None
monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None)
with override_forward_context(None):
for request_index in range(2):
request = OmniDiffusionRequest(
prompt="a fox walks",
request_id=f"request-{request_index}",
sampling_params=OmniDiffusionSamplingParams(
num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent"
),
)
def execute():
return runner._execute_request_list(
[request],
od_config=runner.od_config,
allow_single_output=True,
require_request_batch_support=False,
record_name="test_wan_step_context",
record_output_peak_memory=False,
)
if fail_first_request and request_index == 0:
with pytest.raises(InjectedPredictionError):
execute()
else:
execute()
assert not is_forward_context_available()
assert len(contexts) == 2 and contexts[0] is not contexts[1]
expected = [
(request, step, "high" if step == 0 else "low", positive, step in (0, 2))
for request in range(2)
for step in range(2 if fail_first_request and request == 0 else 3)
for positive in (True, False)
]
assert events == expected
class _StubDMDScheduler:
def __init__(self) -> None:
self.predict_clean_calls: list[tuple[float, float, float]] = []
tests/diffusion/quantization/test_mxfp4_checkpoint.py
+294
View File
@@ -0,0 +1,294 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU coverage for MXFP4 checkpoint loading and the W4A8 dtype contract."""
import sys
from types import ModuleType
import pytest
import torch
from vllm.config.load import LoadConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.forward_context import (
ForwardContext,
override_forward_context,
set_forward_context_denoise_step_idx,
)
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.quantization.mxfp4_config import (
DiffusionMXFP4Config,
DiffusionMXFP4DualScaleMixedConfig,
NPUMxfp4DualScaleLinearMethod,
NPUMxfp4LinearMethod,
NPUMxfp4OnlineLinearMethod,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _patch_tp_state(monkeypatch):
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
class _CheckpointLayer(torch.nn.Module):
def __init__(self, method, dtype=torch.bfloat16):
super().__init__()
method.create_weights(
self,
input_size_per_partition=512,
output_partition_sizes=[2],
input_size=512,
output_size=2,
params_dtype=dtype,
weight_loader=default_weight_loader,
)
def load_weights(self, weights):
params = dict(self.named_parameters())
loaded = set()
for name, value in weights:
if name in params:
params[name].weight_loader(params[name], value)
loaded.add(name)
return loaded
def _load_checkpoint(layer, config, weights, monkeypatch):
od_config = OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
loader.counter_before_loading_weights = 0.0
monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
loader.load_weights(layer)
@pytest.mark.parametrize("serialized", [False, True])
def test_single_scale_requires_checkpoint_scale_only_when_serialized(serialized, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized)
# Exercise the shared implementation directly as well as the online MRO
# below: checkpoint requirements must follow the configured storage mode.
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {"weight": torch.ones((2, 512), dtype=torch.bfloat16)}
if serialized:
with pytest.raises(ValueError, match="Required weights.*weight_scale"):
_load_checkpoint(layer, config, weights, monkeypatch)
else:
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_online_mxfp4_loads_bf16_checkpoint_without_scale(scale_alg, monkeypatch):
config = DiffusionMXFP4Config(mxfp4_scale_alg=scale_alg)
method = NPUMxfp4OnlineLinearMethod(config)
layer = _CheckpointLayer(method)
assert set(dict(layer.named_parameters())) == {"weight"}
assert layer.weight.is_meta
weight = torch.arange(1024, dtype=torch.bfloat16).reshape(2, 512)
packed_weight = torch.zeros((2, 256), dtype=torch.uint8)
generated_scale = torch.full((2, 16), 127, dtype=torch.uint8)
quantized_inputs = []
def quantize(value, **kwargs):
quantized_inputs.append(value.clone())
expected = dict(dst_type="fp4", axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
if scale_alg == 2:
expected["dst_type_max"] = 7.25
assert kwargs == expected
return packed_weight, generated_scale
# Keep real creation, lazy loading, processing, and strict loader checks;
# substitute only the device quantization operator for this CPU test.
npu = ModuleType("torch_npu")
npu.__dict__.update(float4_e2m1fn_x2="fp4", npu_dynamic_mx_quant=quantize)
monkeypatch.setitem(sys.modules, "torch_npu", npu)
_load_checkpoint(layer, config, {"weight": weight}, monkeypatch)
assert len(quantized_inputs) == 1
torch.testing.assert_close(quantized_inputs[0], weight, rtol=0, atol=0)
torch.testing.assert_close(layer.weight, packed_weight, rtol=0, atol=0)
torch.testing.assert_close(layer.weight_scale, generated_scale.reshape(2, 8, 2), rtol=0, atol=0)
assert not getattr(layer.weight_scale, "is_checkpoint_required", False)
method.process_weights_after_loading(layer)
assert len(quantized_inputs) == 1
@pytest.mark.parametrize("require_smooth", [False, True])
@pytest.mark.parametrize("has_smooth", [False, True])
def test_single_scale_checkpoint_smooth_default_and_required(require_smooth, has_smooth, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=require_smooth)
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
scale = torch.linspace(0.5, 3.0, 512)
if has_smooth:
weights["mul_scale"] = scale
else:
weights.pop("mul_scale")
if require_smooth and not has_smooth:
with pytest.raises(ValueError, match="mul_scale"):
_load_checkpoint(layer, config, weights, monkeypatch)
else:
_load_checkpoint(layer, config, weights, monkeypatch)
torch.testing.assert_close(layer.mul_scale, scale if has_smooth else torch.ones(512), rtol=0, atol=0)
@pytest.mark.parametrize("bad_scale", ["shape", "name"])
def test_declared_single_scale_smooth_rejects_malformed_checkpoint(bad_scale, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
if bad_scale == "shape":
weights["mul_scale"] = torch.ones(511)
with pytest.raises(AssertionError):
_load_checkpoint(layer, config, weights, monkeypatch)
else:
weights["multiply_scale"] = weights.pop("mul_scale")
with pytest.raises(ValueError, match="mul_scale"):
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("required", [False, True])
@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), 0.0, -1.0])
def test_single_scale_rejects_invalid_smooth_values(required, bad_value, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=required)
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
weights["mul_scale"][7] = bad_value
with pytest.raises(ValueError, match="finite, strictly positive"):
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("bad_scale", [torch.ones(512, dtype=torch.int32), torch.ones(1, 512)])
def test_single_scale_rejects_invalid_smooth_type_or_rank(bad_scale, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
weights["mul_scale"] = bad_scale
with pytest.raises(ValueError, match="one-dimensional floating-point"):
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("bad_value", [1e10, 1e-10])
def test_single_scale_rejects_smooth_overflow_or_underflow_after_cast(bad_value, monkeypatch):
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
method = NPUMxfp4LinearMethod(config)
layer = _CheckpointLayer(method, torch.float16)
layer.mul_scale.data.fill_(bad_value)
# Validation must fail before any device packing or GEMM.
monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu"))
with pytest.raises(ValueError, match="finite, strictly positive"):
method.process_weights_after_loading(layer)
@pytest.mark.parametrize(
("mode", "missing"),
[
("single", "weight"),
("single", "weight_scale"),
("dual", "weight"),
("dual", "weight_scale"),
("dual", "weight_dual_scale"),
("dual", "mul_scale"),
],
)
def test_offline_mxfp4_required_tensors_do_not_use_legacy_scale_tolerance(mode, missing, monkeypatch):
if mode == "single":
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
method = NPUMxfp4LinearMethod(config)
else:
config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
method = NPUMxfp4DualScaleLinearMethod(config)
layer = _CheckpointLayer(method)
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters() if name != missing}
with pytest.raises(ValueError, match=missing):
_load_checkpoint(layer, config, weights, monkeypatch)
@pytest.mark.parametrize("rank", [0, 1])
def test_single_scale_smooth_uses_real_row_parallel_loader(rank):
from vllm.model_executor.layers.linear import RowParallelLinear
layer = _CheckpointLayer(NPUMxfp4LinearMethod(DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)))
layer.tp_rank = rank
checkpoint_scale = torch.arange(1024, dtype=torch.float32)
RowParallelLinear.weight_loader(layer, layer.mul_scale, checkpoint_scale)
torch.testing.assert_close(layer.mul_scale, checkpoint_scale[rank * 512 : (rank + 1) * 512], rtol=0, atol=0)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("fallback", [False, True])
@pytest.mark.parametrize("with_bias", [False, True])
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_single_scale_smooth_precedes_quantization_and_preserves_output(
dtype, fallback, with_bias, scale_alg, monkeypatch
):
config = DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=True, w4a8_fallback_steps=[0], mxfp4_scale_alg=scale_alg
)
method = NPUMxfp4LinearMethod(config)
layer = _CheckpointLayer(method, dtype)
weights = {
"weight": torch.linspace(-0.5, 1.5, 1024, dtype=dtype).reshape(2, 512),
"weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
"mul_scale": torch.linspace(0.5, 3.0, 512),
}
_load_checkpoint(layer, config, weights, monkeypatch)
x = torch.linspace(-1.0, 2.0, 3072, dtype=dtype).reshape(2, 3, 512)
bias = torch.tensor([0.5, -1.0], dtype=dtype) if with_bias else None
quantized_inputs = []
# Only the two device operators are substituted. Exercise real apply(),
# step selection, Smooth multiplication, bias handling, and reshape logic.
npu = ModuleType("torch_npu")
def quantize(value, **kwargs):
quantized_inputs.append(value.clone())
expected = dict(
dst_type=torch.float8_e4m3fn if fallback else "fp4",
axis=-1,
block_size=32,
round_mode="rint",
scale_alg=0 if fallback else scale_alg,
)
if not fallback and scale_alg == 2:
expected["dst_type_max"] = 7.25
assert kwargs == expected
return value, torch.ones((6, 16), dtype=torch.uint8)
def matmul(x_q, weight, weight_scale, *, bias, output_dtype, **kwargs):
if fallback:
assert output_dtype == torch.bfloat16
if bias is not None:
assert bias.dtype == output_dtype
assert bias.shape == (1, 2)
result = x_q.float() @ weight.float()
if bias is not None:
result += bias.float()
return result.to(output_dtype)
npu.__dict__.update(
float8_e4m3fn=torch.float8_e4m3fn,
float8_e8m0fnu=torch.uint8,
float4_e2m1fn_x2="fp4",
npu_dynamic_mx_quant=quantize,
npu_quant_matmul=matmul,
)
monkeypatch.setitem(sys.modules, "torch_npu", npu)
with override_forward_context(ForwardContext()):
set_forward_context_denoise_step_idx(0 if fallback else 1)
result = method.apply(layer, x, bias)
expected_input = x.reshape(-1, 512) * weights["mul_scale"].to(dtype)
assert len(quantized_inputs) == 1
torch.testing.assert_close(quantized_inputs[0], expected_input, rtol=0, atol=0)
expected = expected_input.float() @ weights["weight"].float().T
if bias is not None:
expected += bias.float()
if fallback:
expected = expected.bfloat16()
torch.testing.assert_close(result, expected.to(dtype).reshape(2, 3, 2), rtol=0, atol=0)
assert result.dtype == dtype
tests/diffusion/quantization/test_mxfp4_config.py
+552 -3
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config."""
import pytest
@@ -27,12 +27,42 @@ def test_mxfp4_config_get_name():
assert DiffusionMXFP4Config.get_name() == "mxfp4"
@pytest.mark.parametrize("value", [True, False, None, "2", 2.0, 1, -1, 3])
def test_mxfp4_scale_alg_rejects_invalid_values(value):
from vllm_omni.quantization import build_quant_config
with pytest.raises(ValueError, match="mxfp4_scale_alg"):
build_quant_config({"method": "mxfp4", "mxfp4_scale_alg": value})
@pytest.mark.parametrize("runtime_alg", [0, 2])
@pytest.mark.parametrize("serialized", [False, True])
def test_runtime_algorithm_survives_offline_storage_rebuild(runtime_alg, serialized):
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, mxfp4_scale_alg=runtime_alg)
resolved = resolve_quant_config_from_disk(
active,
{
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"mxfp4_scale_alg": 2,
"ignored_layers": ["proj_out"],
},
)
assert resolved.is_checkpoint_mxfp4_serialized
assert resolved.mxfp4_scale_alg == runtime_alg
assert resolved.w4a8_fallback_layers == resolved.w4a8_fallback_steps == []
def test_mxfp4_config_from_config_defaults():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = DiffusionMXFP4Config.from_config({})
assert cfg.is_checkpoint_mxfp4_serialized is False
assert cfg.ignored_layers == []
assert cfg.w4a8_fallback_layers == []
def test_mxfp4_config_from_config_serialized():
@@ -42,6 +72,98 @@ def test_mxfp4_config_from_config_serialized():
assert cfg.is_checkpoint_mxfp4_serialized is True
def test_mxfp4_config_requires_offline_smooth():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
values = {"is_checkpoint_mxfp4_serialized": True, "require_smooth_scale": True}
assert DiffusionMXFP4Config.from_config(values).require_smooth_scale
assert build_quant_config({"method": "mxfp4", **values}).require_smooth_scale
with pytest.raises(ValueError, match="offline"):
build_quant_config({"method": "mxfp4", "require_smooth_scale": True})
with pytest.raises(ValueError, match="boolean"):
DiffusionMXFP4Config.from_config({**values, "require_smooth_scale": "false"})
@pytest.mark.parametrize("serialized", [False, True])
def test_mxfp4_disk_smooth_requirement_preserves_step_policy(serialized):
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, w4a8_fallback_steps=[1])
resolved = resolve_quant_config_from_disk(
active,
{
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"require_smooth_scale": True,
},
)
assert resolved.require_smooth_scale
assert resolved.is_checkpoint_mxfp4_serialized
assert resolved.w4a8_fallback_steps == [1]
@pytest.mark.parametrize("disk_requirement", [None, False, True])
def test_mxfp4_expert_rebuild_cannot_weaken_explicit_smooth_requirement(disk_requirement):
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
disk = {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"ignored_layers": ["blocks.0.attn2.to_q"],
}
if disk_requirement is not None:
disk["require_smooth_scale"] = disk_requirement
resolved = resolve_quant_config_from_disk(active, disk)
assert resolved.require_smooth_scale
assert resolved.ignored_layers == disk["ignored_layers"]
@pytest.mark.parametrize("disk_qc", ["mxfp4", {"quant_method": "mxfp4"}])
def test_mxfp4_method_only_disk_config_preserves_caller_layer_policy(disk_qc):
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
active = DiffusionMXFP4Config(
ignored_layers=["blocks.0.attn2.to_q"],
w4a8_fallback_steps=[1],
w4a8_fallback_layers=["blocks.1.attn1.to_qkv"],
)
resolved = resolve_quant_config_from_disk(active, disk_qc)
assert resolved is active
assert resolved.ignored_layers == ["blocks.0.attn2.to_q"]
assert resolved.w4a8_fallback_steps == [1]
assert resolved.w4a8_fallback_layers == ["blocks.1.attn1.to_qkv"]
def test_quantization_disk_string_rejects_mismatched_active_method():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
active = build_quant_config("mxfp4")
with pytest.raises(ValueError, match="mxfp4_dualscale.*active quantization config is 'mxfp4'"):
resolve_quant_config_from_disk(active, "mxfp4_dualscale")
def test_quantization_disk_string_accepts_equivalent_method_alias():
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
class ActiveAliasConfig:
@staticmethod
def get_name() -> str:
return "inc"
active = ActiveAliasConfig()
assert resolve_quant_config_from_disk(active, "auto-round") is active
def test_mxfp4_config_from_config_ignored_layers():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
@@ -57,6 +179,202 @@ def test_mxfp4_config_from_config_modules_to_not_convert_fallback():
assert cfg.ignored_layers == ["proj_out"]
@pytest.mark.parametrize("serialized", [False, True])
@pytest.mark.parametrize(
"policy",
[
{"w4a8_fallback_steps": [0]},
{"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
{"w4a8_fallback_steps": [1], "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
],
)
def test_dualscale_rejects_w4a8_configuration(serialized, policy):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
config = {"quant_method": "mxfp4_dualscale", "is_checkpoint_serialized": serialized, **policy}
with pytest.raises(ValueError, match="does not support W4A8"):
DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=serialized, **policy)
with pytest.raises(ValueError, match="does not support W4A8"):
DiffusionMXFP4DualScaleMixedConfig.from_config(config)
with pytest.raises(ValueError, match="does not support W4A8"):
build_quant_config({"transformer": config})
with pytest.raises(ValueError, match="does not support W4A8"):
resolve_quant_config_from_disk(None, config)
@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
@pytest.mark.parametrize("steps", [[-1], [True], [1.5], ["2"], "0,2", (0, 2)])
def test_mxfp4_rejects_invalid_fallback_steps(method, steps):
from vllm_omni.quantization import build_quant_config
with pytest.raises(ValueError, match="non-negative integer"):
build_quant_config({"method": method, "w4a8_fallback_steps": steps})
def test_mxfp4_fallback_steps_follow_request_context():
from vllm_omni.diffusion.forward_context import (
ForwardContext,
override_forward_context,
set_forward_context_denoise_step_idx,
)
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
steps = [2, 0, 2]
cfg = build_quant_config({"method": "mxfp4", "w4a8_fallback_steps": steps})
steps.append(1)
assert cfg.w4a8_fallback_steps == [0, 2]
with override_forward_context(None):
assert not _is_w4a8_fallback_step([])
with pytest.raises(RuntimeError, match="publishes its denoise step"):
_is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
for _request in range(2):
with override_forward_context(ForwardContext()):
for step, expected in [(0, True), (1, False), (2, True), (3, False)]:
set_forward_context_denoise_step_idx(step)
# Both CFG passes and both experts read the same index.
assert [_is_w4a8_fallback_step(cfg.w4a8_fallback_steps) for _cfg_pass in range(2)] == [
expected,
expected,
]
with pytest.raises(RuntimeError, match="publishes its denoise step"):
_is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
@pytest.mark.parametrize("method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
@pytest.mark.parametrize("steps", [[], [0, 2]])
@pytest.mark.parametrize("already_serialized", [False, True])
def test_mxfp4_disk_rebuild_preserves_runtime_fallback(method, flag, steps, already_serialized):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
layers = ["blocks.10.attn1.to_qkv"] if steps else []
active = build_quant_config(
{
"method": method,
flag: already_serialized,
"w4a8_fallback_steps": steps,
"w4a8_fallback_layers": layers,
}
)
# Separate experts may have distinct high precision layers. Both must
# retain the requested step policy during online/offline reconciliation.
for ignored in [["blocks.0.attn1.to_q"], ["blocks.1.attn1.to_q"]]:
resolved = resolve_quant_config_from_disk(
active,
{
"quant_method": method,
flag: True,
"ignored_layers": ignored,
"w4a8_fallback_steps": [9],
"w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
},
)
assert getattr(resolved, flag)
assert resolved.ignored_layers == ignored
assert resolved.w4a8_fallback_steps == steps
assert resolved.w4a8_fallback_layers == layers
assert active.ignored_layers == []
@pytest.mark.parametrize(
"runtime_policy",
[{}, {"w4a8_fallback_steps": [1]}, {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}],
)
def test_mxfp4_explicit_omission_disables_saved_policies_for_both_experts(runtime_policy):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
active = build_quant_config({"method": "mxfp4", **runtime_policy})
resolved_experts = []
for ignored in (["blocks.0.attn2.to_q"], ["blocks.1.ffn.net_2"]):
disk = {
"quant_method": "mxfp4",
"is_checkpoint_mxfp4_serialized": True,
"ignored_layers": ignored,
"w4a8_fallback_steps": [9],
"w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
}
resolved = resolve_quant_config_from_disk(active, disk)
assert resolved.w4a8_fallback_steps == runtime_policy.get("w4a8_fallback_steps", [])
assert resolved.w4a8_fallback_layers == runtime_policy.get("w4a8_fallback_layers", [])
assert resolved.ignored_layers == ignored
detected = resolve_quant_config_from_disk(None, disk)
assert detected.w4a8_fallback_steps == [9]
assert detected.w4a8_fallback_layers == ["blocks.9.ffn.net_2"]
resolved_experts.append(resolved)
assert resolved_experts[0] is not resolved_experts[1]
assert active.is_checkpoint_mxfp4_serialized is False
assert active.ignored_layers == []
@pytest.mark.parametrize(
"method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
)
def test_serialized_mxfp4_disk_omission_clears_active_ignored_layers(method, flag):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
active_spec = {
"method": method,
flag: True,
"ignored_layers": ["blocks.0.attn2.to_q"],
"w4a8_fallback_steps": [1] if method == "mxfp4" else [],
"w4a8_fallback_layers": ["blocks.1.attn1.to_qkv"] if method == "mxfp4" else [],
}
if method == "mxfp4":
active_spec["require_smooth_scale"] = True
active = build_quant_config(active_spec)
resolved = resolve_quant_config_from_disk(active, {"quant_method": method, flag: True})
assert resolved is not active
assert resolved.ignored_layers == []
assert resolved.w4a8_fallback_steps == active.w4a8_fallback_steps
assert resolved.w4a8_fallback_layers == active.w4a8_fallback_layers
if method == "mxfp4":
assert resolved.require_smooth_scale
@pytest.mark.parametrize(
"method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
)
@pytest.mark.parametrize("with_active", [False, True])
def test_serialized_mxfp4_disk_uses_modules_to_not_convert(method, flag, with_active):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
active = None
disk_steps = [9] if method == "mxfp4" else []
active_steps = [1] if method == "mxfp4" else []
expected_steps = disk_steps
if with_active:
active = build_quant_config(
{
"method": method,
flag: True,
"ignored_layers": ["blocks.0.attn2.to_q"],
"w4a8_fallback_steps": active_steps,
}
)
expected_steps = active_steps
resolved = resolve_quant_config_from_disk(
active,
{
"quant_method": method,
flag: True,
"modules_to_not_convert": ["blocks.2.ffn.net_2"],
"w4a8_fallback_steps": disk_steps,
},
)
assert resolved.ignored_layers == ["blocks.2.ffn.net_2"]
assert resolved.w4a8_fallback_steps == expected_steps
# ---------------------------------------------------------------------------
# build_quant_config integration
# ---------------------------------------------------------------------------
@@ -594,18 +912,31 @@ def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch):
assert isinstance(method, ROCmMxfp4OnlineLinearMethod)
def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch):
@pytest.mark.parametrize("steps", [[], [0, 2]])
def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch, steps):
"""A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe."""
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"])
cfg = DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"], w4a8_fallback_steps=steps
)
assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod)
def test_rocm_rejects_w4a8_step_fallback(_rocm_platform, monkeypatch):
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
cfg = DiffusionMXFP4Config(w4a8_fallback_steps=[0, 2])
with pytest.raises(NotImplementedError, match="only supported on NPU"):
cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch):
"""MXFP4 on ROCm requires gfx950; any other arch must raise."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
@@ -652,3 +983,221 @@ def test_rocm_create_weights_row_parallel_tp2(_rocm_platform):
assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
assert layer.input_size_per_partition == _TP2_K // _TP2
assert layer.output_size_per_partition == _TP2_N
@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
@pytest.mark.parametrize("layers", ["blocks.1.attn1.to_qkv", [1], [True], [""], [" a"], ["blocks.*"], ["a..b"]])
def test_mxfp4_rejects_invalid_fallback_layers(method, layers):
from vllm_omni.quantization import build_quant_config
with pytest.raises(ValueError, match="exact runtime Linear paths"):
build_quant_config({"method": method, "w4a8_fallback_layers": layers})
@pytest.mark.parametrize("method", ["mxfp4"])
def test_mxfp4_layer_policy_from_checkpoint_and_component_config(method):
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.component_config import resolve_component_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
names = ["blocks.10.attn1.to_qkv", "blocks.10.attn1.to_qkv"]
components = build_quant_config(
{"transformer": {"method": method, "w4a8_fallback_layers": names}, "transformer_2": {"method": method}}
)
names.append("blocks.11.attn1.to_qkv")
assert resolve_component_quant_config(components, "transformer").w4a8_fallback_layers == [names[0]]
assert resolve_component_quant_config(components, "transformer_2").w4a8_fallback_layers == []
detected = resolve_quant_config_from_disk(None, {"quant_method": method, "w4a8_fallback_layers": [names[0]]})
assert detected.w4a8_fallback_layers == [names[0]]
assert type(detected).from_config({"w4a8_fallback_layers": [names[0]]}).w4a8_fallback_layers == [names[0]]
@pytest.mark.parametrize(
"kind, serialized_flag",
[("mxfp4", "is_checkpoint_mxfp4_serialized")],
)
@pytest.mark.parametrize("offline", [False, True])
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_mxfp4_layer_and_step_union_routes_actual_methods(
kind, serialized_flag, offline, scale_alg, monkeypatch, mocker
):
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
from vllm_omni.quantization import build_quant_config, mxfp4_config
monkeypatch.setattr(mxfp4_config.current_omni_platform, "is_npu", lambda: True)
selected = "blocks.10.attn1.to_qkv"
other = "blocks.11.attn2.to_q"
preserved = "blocks.12.ffn.net_2"
config = build_quant_config(
{
"method": kind,
"mxfp4_scale_alg": scale_alg,
serialized_flag: offline,
"w4a8_fallback_layers": [selected, preserved],
"ignored_layers": [preserved],
"w4a8_fallback_steps": [1],
}
)
# Real config dispatch creates separate methods for selected and other layers.
linear = mocker.Mock(spec=LinearBase)
assert isinstance(config.get_quant_method(linear, preserved), UnquantizedLinearMethod)
x = torch.ones(2, 64, dtype=torch.bfloat16)
layer = torch.nn.Module()
for name in ("weight", "weight_scale", "mul_scale"):
layer.register_buffer(name, torch.ones(64, dtype=x.dtype))
calls = []
def record_a8(*args):
calls.append("A8")
return x
def record_a4(*args):
calls.append("A4")
return x
monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", record_a8)
for name in (selected, other):
method = config.get_quant_method(linear, name)
assert method.is_w4a8_fallback_layer == (name == selected)
scales = (x, x)
monkeypatch.setattr(method, "_quantize_activation", lambda *args: scales)
monkeypatch.setattr(method, "_quant_matmul", record_a4)
for step in (0, 1, 2):
with override_forward_context(ForwardContext(denoise_step_idx=step)):
method.apply(layer, x)
assert calls == ["A8", "A8", "A8", "A4", "A8", "A4"]
# A permanently selected layer works without a request context, even if steps are also configured.
with override_forward_context(None):
config.get_quant_method(linear, selected).apply(layer, x)
assert calls[-1] == "A8"
@pytest.mark.parametrize("online", [False, True])
@pytest.mark.parametrize("selected", [False, True])
@pytest.mark.parametrize("scale_alg", [0, 2])
def test_single_scale_prepares_one_weight_for_both_activation_precisions(online, selected, scale_alg, monkeypatch):
import sys
from types import SimpleNamespace
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
from vllm_omni.quantization import mxfp4_config
monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
prepared = []
packed = []
def quantize(x, **kwargs):
prepared.append(kwargs)
return x.to(torch.int8), torch.full((x.shape[0], 2), 127, dtype=torch.uint8)
def pack(x, dtype):
packed.append(x)
return x.to(dtype)
fake_npu = SimpleNamespace(
float4_e2m1fn_x2=torch.int8,
float8_e8m0fnu=torch.uint8,
float8_e4m3fn=torch.float8_e4m3fn,
npu_dtype_cast=pack,
npu_dynamic_mx_quant=quantize,
)
monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
config = mxfp4_config.DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=not online,
w4a8_fallback_layers=["blocks.0.attn1.to_qkv"] if selected else [],
w4a8_fallback_steps=[1],
mxfp4_scale_alg=scale_alg,
)
cls = mxfp4_config.NPUMxfp4OnlineLinearMethod if online else mxfp4_config.NPUMxfp4LinearMethod
method = cls(config, prefix="blocks.0.attn1.to_qkv")
layer = torch.nn.Linear(64, 2, bias=False, dtype=torch.bfloat16)
layer.orig_dtype = torch.bfloat16
if not online:
layer.register_parameter("weight_scale", torch.nn.Parameter(torch.full((2, 2), 127, dtype=torch.uint8), False))
layer.register_parameter("mul_scale", torch.nn.Parameter(torch.linspace(0.5, 2, 64), False))
method.process_weights_after_loading(layer)
pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
method.process_weights_after_loading(layer)
assert len(prepared) == int(online)
expected = dict(dst_type=torch.int8, axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
if scale_alg == 2:
expected["dst_type_max"] = 7.25
assert prepared == ([expected] if online else [])
assert len(packed) == int(not online)
branches = []
def a4(x_q, x_scale, target, bias, dtype):
assert (target.weight.data_ptr(), target.weight_scale.data_ptr()) == pointers
branches.append("A4")
return torch.zeros(x_q.shape[0], 2, dtype=dtype)
def a8(x, weight, scale, bias, dtype):
assert (weight.data_ptr(), scale.data_ptr()) == pointers
branches.append("A8")
return torch.zeros(x.shape[0], 2, dtype=dtype)
monkeypatch.setattr(method, "_quant_matmul", a4)
monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", a8)
monkeypatch.setattr(method, "_quantize_activation", lambda x: (x, x))
monkeypatch.setattr(fake_npu, "npu_dynamic_mx_quant", lambda *a, **kw: pytest.fail("weight requantization"))
monkeypatch.setattr(fake_npu, "npu_dtype_cast", lambda *a, **kw: pytest.fail("weight repacking"))
for step in (0, 1, 2):
with override_forward_context(ForwardContext(denoise_step_idx=step)):
method.apply(layer, torch.ones(2, 64, dtype=torch.bfloat16))
assert branches == (["A8"] * 3 if selected else ["A4", "A8", "A4"])
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
assert set(layer.state_dict()) == (
{"weight", "weight_scale"} if online else {"weight", "weight_scale", "mul_scale"}
)
assert not hasattr(layer, "w4a8_weight")
assert not hasattr(layer, "w4a8_weight_scale")
@pytest.mark.parametrize("online", [False, True])
def test_dualscale_without_fallback_keeps_original_a4_path(online, monkeypatch):
import sys
from types import SimpleNamespace
from vllm_omni.quantization import mxfp4_config
monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
original_ones = torch.ones
monkeypatch.setattr(torch, "ones", lambda *args, **kw: original_ones(*args, **{**kw, "device": "cpu"}))
calls = []
def quantize_dual(x, **kw):
calls.append(kw.get("smooth_scale"))
return x.to(torch.int8), torch.ones(x.shape[0], 1), torch.ones(x.shape[0], 16, dtype=torch.uint8)
fake_npu = SimpleNamespace(
float4_e2m1fn_x2=torch.int8,
npu_dtype_cast=lambda x, dtype: x.to(dtype),
npu_format_cast=lambda x, *args, **kw: x,
npu_dynamic_dual_level_mx_quant=quantize_dual,
)
monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
config = mxfp4_config.DiffusionMXFP4DualScaleMixedConfig(
is_checkpoint_serialized=not online, w4a8_fallback_steps=[], w4a8_fallback_layers=[]
)
cls = mxfp4_config.NPUMxfp4DualScaleOnlineLinearMethod if online else mxfp4_config.NPUMxfp4DualScaleLinearMethod
method = cls(config)
layer = torch.nn.Module()
layer.input_size_per_partition = 512
for name, value in {
"weight": torch.ones(2, 512, dtype=torch.bfloat16),
"weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
"weight_dual_scale": torch.ones(2, 1, 1),
"mul_scale": torch.ones(512),
}.items():
layer.register_parameter(name, torch.nn.Parameter(value, requires_grad=False))
method.process_weights_after_loading(layer)
method.process_weights_after_loading(layer)
assert len(calls) == int(online)
assert set(layer.state_dict()) == {"weight", "weight_scale", "weight_dual_scale", "mul_scale"}
assert not hasattr(layer, "w4a8_weight")
monkeypatch.setattr(method, "_quant_matmul", lambda *args: torch.zeros(2, 2, dtype=torch.bfloat16))
method.apply(layer, torch.ones(2, 512, dtype=torch.bfloat16))
assert len(calls) == 1 + int(online)
torch.testing.assert_close(calls[-1], layer.mul_scale)
tests/diffusion/quantization/test_wan_mxfp4_loading.py
+237
View File
@@ -0,0 +1,237 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Exercise Wan's real name mapping and vLLM's QKV parameter loaders on CPU."""
import pytest
import torch
from vllm.config.load import LoadConfig
from vllm.model_executor.layers.linear import QKVParallelLinear
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, DiffusionMXFP4DualScaleMixedConfig
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _patch_tp_state(monkeypatch):
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
prefix = "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer"
monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_world_size", lambda: 1)
class _WanCheckpoint(torch.nn.Module):
load_weights = WanTransformer3DModel.load_weights
_log_w4a8_fallback_load_summary = WanTransformer3DModel._log_w4a8_fallback_load_summary
def __init__(self, config, *, num_blocks=1):
super().__init__()
blocks = []
for block_idx in range(num_blocks):
block = torch.nn.Module()
block.attn1 = torch.nn.Module()
block.attn1.to_qkv = QKVParallelLinear(
hidden_size=512,
head_size=32,
total_num_heads=4,
bias=False,
params_dtype=torch.bfloat16,
quant_config=config,
prefix=f"blocks.{block_idx}.attn1.to_qkv",
disable_tp=True,
)
blocks.append(block)
self.blocks = torch.nn.ModuleList(blocks)
def _checkpoint(mode, fused):
if mode == "single":
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
elif mode == "dual":
config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
else:
config = None
model = _WanCheckpoint(config)
weights = {}
expected = {}
for name, param in model.named_parameters():
suffix = name.rsplit(".", 1)[1]
output_dim = getattr(param, "output_dim", None)
if output_dim is None:
value = torch.full_like(param, 2.5)
shards = [value] * 3
else:
parts = param.chunk(3, dim=output_dim)
shards = [
torch.full_like(part, (121 if suffix == "weight_scale" else 1) + i) for i, part in enumerate(parts)
]
value = torch.cat(shards, dim=output_dim)
expected[name] = value
if fused:
weights[name] = value
else:
for shard_id, shard in zip(("q", "k", "v"), shards):
weights[f"blocks.0.attn1.to_{shard_id}.{suffix}"] = shard
return model, config, weights, expected
def _load(model, config, weights, monkeypatch):
loader = DiffusersPipelineLoader(
LoadConfig(), OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
)
loader.counter_before_loading_weights = 0.0
monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
loader.load_weights(model)
@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
@pytest.mark.parametrize("fused", [False, True])
def test_wan_qkv_loads_actual_weights_and_scales(mode, fused, monkeypatch):
model, config, weights, expected = _checkpoint(mode, fused)
_load(model, config, weights, monkeypatch)
for name, param in model.named_parameters():
torch.testing.assert_close(param, expected[name], rtol=0, atol=0)
@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
def test_wan_incomplete_qkv_weight_is_not_marked_fully_loaded(mode, monkeypatch):
model, config, weights, _ = _checkpoint(mode, fused=False)
weights.pop("blocks.0.attn1.to_k.weight")
weights.pop("blocks.0.attn1.to_v.weight")
with pytest.raises(ValueError, match="to_qkv.weight"):
_load(model, config, weights, monkeypatch)
@pytest.mark.parametrize("mode", ["single", "dual"])
def test_wan_incomplete_qkv_scale_is_not_marked_fully_loaded(mode, monkeypatch):
model, config, weights, _ = _checkpoint(mode, fused=False)
weights.pop("blocks.0.attn1.to_k.weight_scale")
weights.pop("blocks.0.attn1.to_v.weight_scale")
with pytest.raises(ValueError, match="to_qkv.weight_scale"):
_load(model, config, weights, monkeypatch)
@pytest.mark.parametrize("mode", ["single", "dual"])
@pytest.mark.parametrize("reverse", [False, True])
def test_wan_split_qkv_rejects_different_smooth_scales(mode, reverse, monkeypatch):
model, config, weights, _ = _checkpoint(mode, fused=False)
weights["blocks.0.attn1.to_k.mul_scale"] = weights["blocks.0.attn1.to_k.mul_scale"] * 2
if reverse:
weights = dict(reversed(weights.items()))
with pytest.raises(ValueError, match="must share the same Smooth tensor"):
_load(model, config, weights, monkeypatch)
@pytest.mark.parametrize("name", ["to_q.missing", "to_qkv.missing", "to_qkv_extra.weight", "to_query.weight"])
def test_wan_unknown_qkv_key_is_not_reported_as_loaded(name):
model, _, _, _ = _checkpoint("single", fused=True)
loaded = model.load_weights([(f"blocks.0.attn1.{name}", torch.ones(2))])
assert loaded == set()
@pytest.mark.parametrize("fused", [False, True])
def test_wan_single_scale_rejects_relabelled_dualscale_checkpoint(fused, monkeypatch):
model, config, weights, _ = _checkpoint("single", fused=fused)
projection = "to_qkv" if fused else "to_q"
weights[f"blocks.0.attn1.{projection}.weight_dual_scale"] = torch.ones(128, 1, 1)
with pytest.raises(ValueError, match="cannot load DualScale tensor"):
_load(model, config, weights, monkeypatch)
def test_wan_layer_policy_validates_runtime_qkv_and_bf16_override():
config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv"])
model = _WanCheckpoint(config)
# Uses the same validation called at the end of WanTransformer3DModel construction.
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
for bad in ("blocks.0.attn1.to_q", "blocks.0.attn1", "blocks.99.attn1.to_qkv", "transformer.blocks.0.attn1.to_qkv"):
config.w4a8_fallback_layers = [bad]
with pytest.raises(ValueError, match="exact runtime Linear paths"):
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
def test_wan_layer_policy_defers_only_other_pp_stage_paths():
from vllm.model_executor.models.utils import PPMissingLayer
config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv", "blocks.1.attn1.to_qkv"])
model = _WanCheckpoint(config)
model.blocks.append(PPMissingLayer())
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
config.w4a8_fallback_layers = ["blocks.1"]
with pytest.raises(ValueError, match="exact runtime Linear paths"):
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
@pytest.mark.parametrize("is_checkpoint_serialized", [False, True], ids=["online", "offline"])
def test_loader_logs_wan_expert_summaries_after_weight_processing(is_checkpoint_serialized, mocker):
selected = "blocks.0.attn1.to_qkv"
bf16 = "blocks.1.attn1.to_qkv"
transformer_config = DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
w4a8_fallback_layers=[selected],
)
transformer_2_config = DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
ignored_layers=[bf16],
w4a8_fallback_layers=[selected, bf16],
)
pipeline = torch.nn.Module()
pipeline.transformer = _WanCheckpoint(transformer_config)
pipeline.transformer_2 = _WanCheckpoint(transformer_2_config, num_blocks=2)
WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer, transformer_config)
WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer_2, transformer_2_config)
loader = DiffusersPipelineLoader(
LoadConfig(),
OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=transformer_config),
)
events = []
summaries = []
def process_weights(model, _target_device):
events.append("process")
for module in model.modules():
method = getattr(module, "quant_method", None)
if getattr(method, "is_w4a8_fallback_layer", False):
module._already_called_process_weights_after_loading = True
def capture_summary(message, *args):
rendered = message % args
summaries.append(rendered)
events.append(f"summary:{args[0]}")
loader._init_from_load_format = lambda *_args, **_kwargs: pipeline # type: ignore[method-assign]
loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
loader._process_weights_after_loading = process_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
mocker.patch(
"vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info",
side_effect=capture_summary,
)
assert loader.load_model(load_device="cpu") is pipeline
assert events == ["load", "process", "summary:transformer", "summary:transformer_2"]
assert "component=transformer; selected=1 ['transformer.blocks.0.attn1.to_qkv']" in summaries[0]
assert "BF16 overrides=0 []" in summaries[0]
assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[0]
assert "component=transformer_2; selected=1 ['transformer_2.blocks.0.attn1.to_qkv']" in summaries[1]
assert "BF16 overrides=1 ['transformer_2.blocks.1.attn1.to_qkv']" in summaries[1]
assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[1]
def test_wan_load_summary_does_not_claim_unprocessed_layer_is_ready(mocker):
selected = "blocks.0.attn1.to_qkv"
config = DiffusionMXFP4Config(w4a8_fallback_layers=[selected])
model = _WanCheckpoint(config)
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
logger_info = mocker.patch("vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info")
WanTransformer3DModel._log_w4a8_fallback_load_summary(model, "transformer")
message, *args = logger_info.call_args.args
rendered = message % tuple(args)
assert "weight/A8 processing state=not-ready (0/1 selected layers processed)" in rendered
tests/diffusion/test_diffusion_config_propagation.py
+6
View File
@@ -25,6 +25,12 @@ from vllm_omni.diffusion.model_metadata import (
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _local_model_paths(monkeypatch):
# These tests exercise config transport, not model repository resolution.
monkeypatch.setattr("vllm_omni.diffusion.data.get_model_path", lambda model, revision: model)
def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
"""Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
tests/e2e/features/comfyui/test_comfyui_integration.py
+1 -1
View File
@@ -534,7 +534,7 @@ def mock_async_omni(
# 1. The API layer uses its stage_list and stage_configs attributes
# 2. Its __init__ method has slow side effects (model & config loading).
# These cases use the non-duplex mock engine without loading model configs.
mocker.patch("vllm_omni.entrypoints.openai.api_server._is_duplex_model", return_value=False)
mocker.patch("vllm_omni.entrypoints.openai.api_server._should_serve_duplex", return_value=False)
mock_async_omni_cls = mocker.patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni")
monkeypatch.setattr(
"vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat",
tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
+10
View File
@@ -21,6 +21,7 @@ import torch.distributed as dist
from transformers import AutoTokenizer
from tests.helpers.mark import hardware_test
from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
from vllm_omni.entrypoints.async_omni import AsyncOmni
@@ -136,6 +137,15 @@ async def _generate_once(
def _assert_live_uniproc_worker(engine: AsyncOmni) -> None:
"""Default single-GPU path must use the in-process executor, not mp IPC."""
stage_configs = engine.engine.stage_configs
assert len(stage_configs) == 1
stage_config = stage_configs[0]
assert isinstance(stage_config, VllmOmniDiffusionStageConfig)
assert stage_config.model_config.model == MODEL
assert stage_config.model_config.enforce_eager is True
assert stage_config.scheduler_config.max_num_seqs == 1
assert stage_config.diffusion_config.custom_pipeline_args == {"pipeline_class": CUSTOM_PIPELINE_CLASS}
assert stage_config.diffusion_config.worker_extension_cls == WORKER_EXTENSION_CLASS
clients = engine.engine.stage_clients
assert len(clients) == 1
client = clients[0]
tests/e2e/offline_inference/test_hunyuanimage3.py
+138 -24
View File
@@ -1,5 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# ruff: noqa: E501
import copy
import importlib.util
import os
import sys
from collections.abc import Generator
from pathlib import Path
@@ -10,20 +15,37 @@ import torch.nn.functional as F
from PIL import Image
from transformers import CLIPModel, CLIPProcessor
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni import Omni
from vllm_omni.config.omni_config import (
VllmOmniARStageConfig,
VllmOmniDiffusionStageConfig,
)
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator, get_model_class_name
from vllm_omni.platforms import current_omni_platform
PROMPT = "A brown and white dog is running on the grass"
MODEL_NAME = "tencent/HunyuanImage-3.0"
MODEL_NAME = os.environ.get("HUNYUAN_IMAGE3_MODEL", "tencent/HunyuanImage-3.0-Instruct")
LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32"
DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml")
_PRODUCTION_DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml")
# Keep the production Mooncake connector, but use its TCP transport for this
# single-host test. The shared server's RDMA device is not usable reliably;
# TCP still exercises the Mooncake AR→DiT data plane without changing deploy
# semantics in production.
DEPLOY_CONFIG_PATH = modify_stage_config(
_PRODUCTION_DEPLOY_CONFIG_PATH,
updates={"connectors.rdma_connector.extra.protocol": "tcp"},
)
pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion]
pytestmark = [
pytest.mark.advanced_model,
pytest.mark.diffusion,
]
# System prompt type. Options: None, dynamic, en_vanilla, en_recaption, en_think_recaption, en_unified
# Below are the CLIP embedding tensors from the official HunyuanImage model (seed=1234, prompt: "A brown and white dog is running on the grass").
@@ -285,25 +307,97 @@ def omni() -> Generator[Omni, None, None]:
yield runner.omni
@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
@hardware_test(res={"cuda": "H100"}, num_cards=4)
def test_structured_mixed_pipeline_config_reaches_runtime(omni: Omni) -> None:
"""Mixed-pipeline deploy settings reach both terminal engine configs."""
engine = omni.engine
# Topology, connector selection, and placement remain Omni-owned.
resolved_stages = engine.stage_configs
assert len(resolved_stages) == 2
ar_stage, diffusion_stage = resolved_stages
assert isinstance(ar_stage, VllmOmniARStageConfig)
assert ar_stage.stage_id == 0
assert ar_stage.model_stage == "AR"
assert ar_stage.final_output_type == "text"
assert ar_stage.runtime_config.devices == "0,1"
assert ar_stage.connector_config.output_connectors == {"to_stage_1": "rdma_connector"}
assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig)
assert diffusion_stage.stage_id == 1
assert diffusion_stage.model_stage == "dit"
assert diffusion_stage.final_output_type == "image"
assert diffusion_stage.input_sources == [0]
assert diffusion_stage.runtime_config.devices == "2,3"
assert diffusion_stage.connector_config.input_connectors == {"from_stage_0": "rdma_connector"}
# Stage 0 is materialized as the VllmConfig actually passed to its engine.
ar_vllm_config, diffusion_vllm_config = engine.stage_vllm_configs
assert ar_vllm_config is not None
assert diffusion_vllm_config is None
assert ar_vllm_config.model_config.enforce_eager is True
assert ar_vllm_config.model_config.trust_remote_code is True
assert ar_vllm_config.parallel_config.tensor_parallel_size == 2
assert ar_vllm_config.scheduler_config.max_num_seqs == 1
assert ar_vllm_config.scheduler_config.max_num_batched_tokens == 32768
assert ar_vllm_config.cache_config.gpu_memory_utilization == 0.95
ar_omni_config = ar_vllm_config.model_config.omni_config
assert ar_omni_config.omni_kv_config["need_send_cache"] is True
expected_ar_sampling = {
"temperature": 0.0,
"top_p": 1,
"top_k": -1,
"max_tokens": 8192,
"detokenize": True,
"skip_special_tokens": False,
"include_stop_str_in_output": True,
}
ar_sampling, diffusion_sampling = engine.default_sampling_params_list
assert all(getattr(ar_sampling, name) == value for name, value in expected_ar_sampling.items())
assert diffusion_sampling.num_inference_steps == 50
assert diffusion_sampling.guidance_scale == 0
# Stage 1 is a diffusion engine, so its terminal owner is the launched
# OmniDiffusionConfig rather than a VllmConfig.
od_config = omni.get_diffusion_od_config()
assert od_config is not None
assert od_config.model == MODEL_NAME
assert od_config.enforce_eager is True
assert od_config.distributed_executor_backend == "mp"
assert od_config.parallel_config.tensor_parallel_size == 2
assert od_config.parallel_config.enable_expert_parallel is True
assert od_config.max_num_seqs == 1
diffusion_kv = od_config.omni_kv_config
assert diffusion_kv["need_recv_cache"] is True
assert diffusion_kv["enable_kv_async_prefetch"] is True
assert diffusion_kv["kv_prefetch_min_free_mem_ratio"] == 0.2
def _extract_generated_image(outputs: list[object]) -> Image.Image:
if not outputs:
raise AssertionError("No outputs were returned from Omni.generate()")
first_output = outputs[0]
if images := getattr(first_output, "images", None):
return images[0]
request_output = first_output
if request_output is not None and (images := getattr(request_output, "images", None)):
return images[0]
for output in outputs:
if images := getattr(output, "images", None):
return images[0]
multimodal_output = getattr(output, "multimodal_output", None)
if isinstance(multimodal_output, dict) and (images := multimodal_output.get("images")):
return images[0]
raise AssertionError("No generated image found in Omni output")
def extract_embedding(image: Image.Image, clip_model: CLIPModel, clip_processor: CLIPProcessor) -> torch.Tensor:
inputs = clip_processor(images=image.convert("RGB"), return_tensors="pt")
# The full CLIP forward API is stable across supported Transformers
# versions, but it expects text inputs as well as pixels. An empty text
# prompt is sufficient when only the projected image embedding is needed.
inputs = clip_processor(text=[""], images=[image.convert("RGB")], return_tensors="pt", padding=True)
with torch.inference_mode():
features = clip_model.get_image_features(**inputs)
# Use the public CLIP forward output, matching the shared CLIPScorer
# helper. ``get_image_features`` changed from Tensor to
# BaseModelOutputWithPooling in Transformers 5.x.
features = clip_model(**inputs).image_embeds
features = F.normalize(features, p=2, dim=-1)
return features.squeeze(0)
@@ -320,16 +414,36 @@ def compare_semantic(
def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
generator_device = current_omni_platform.device_type or "cuda"
sampling_params = OmniDiffusionSamplingParams(
seed=1234,
generator=torch.Generator(device=generator_device).manual_seed(1234),
num_outputs_per_prompt=1,
)
if use_system_prompt is not None:
sampling_params.extra_args = {"use_system_prompt": use_system_prompt}
from transformers import AutoTokenizer
outputs = omni.generate({"prompt": PROMPT}, sampling_params)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
prompt_tokens = build_prompt_tokens(
PROMPT,
tokenizer,
task="t2i",
bot_task="think",
sys_type=use_system_prompt or "None",
)
stop_token_ids = resolve_stop_token_ids(task="t2i", bot_task="think", tokenizer=tokenizer, image_size="1024x1024")
params_list = copy.deepcopy(omni.default_sampling_params_list)
generator_device = current_omni_platform.device_type or "cuda"
for params in params_list:
if isinstance(params, OmniDiffusionSamplingParams):
params.seed = 1234
params.generator = torch.Generator(device=generator_device).manual_seed(1234)
params.num_outputs_per_prompt = 1
elif hasattr(params, "stop_token_ids"):
params.stop_token_ids = stop_token_ids
prompt = {
"prompt": PROMPT,
"prompt_token_ids": prompt_tokens.token_ids,
"use_system_prompt": prompt_tokens.system_prompt_type,
"modalities": ["image"],
"height": 1024,
"width": 1024,
}
outputs = omni.generate(prompt, params_list)
return _extract_generated_image(outputs)
@@ -342,7 +456,7 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
MIN_SEMANTIC_SIMILARITY = 0.5
@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
@pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES)
def test_system_prompt_scores(
omni: Omni,
@@ -378,7 +492,7 @@ def _load_text_to_image_module():
return module
@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.")
def test_shared_script_ar_path_reaches_generation(
omni: Omni,
clip_bundle: tuple[CLIPModel, CLIPProcessor],
tests/e2e/offline_inference/test_qwen3_omni.py
+85
View File
@@ -17,6 +17,10 @@ import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.media import generate_synthetic_video
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni.config.omni_config import (
VllmOmniARStageConfig,
VllmOmniGenerationStageConfig,
)
from vllm_omni.platforms import current_omni_platform
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
@@ -25,6 +29,7 @@ thinker_only_models = ["Qwen/Qwen3-Omni-30B-A3B-Captioner"]
# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
# platforms: section. Only CUDA needs an extra enforce_eager tweak.
_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
_PRODUCTION_DEPLOY = get_deploy_config_path("qwen3_omni_moe.yaml")
def get_cuda_graph_config():
@@ -46,6 +51,7 @@ else:
# Create parameter combinations for model and stage config
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
production_test_params = [(model, _PRODUCTION_DEPLOY) for model in models]
# we can use the same config for a model that only has thinker (i.e., does not
# enable audio output) because the resolver should figure out that it doesn't
# need the full pipeline based on the HF config.
@@ -59,6 +65,85 @@ def get_question(prompt_type="video"):
return prompts.get(prompt_type, prompts["video"])
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", production_test_params, indirect=True)
def test_structured_multistage_config_reaches_runtime(omni_runner, offline_client) -> None:
"""Deploy settings reach the materialized configs used by live stages."""
engine = omni_runner.omni.engine
# The resolved typed stages remain the Omni-owned source for topology and
# placement, which are not fields of the terminal vLLM config.
resolved_stages = engine.stage_configs
assert len(resolved_stages) == 3
thinker, talker, code2wav = resolved_stages
assert isinstance(thinker, VllmOmniARStageConfig)
assert isinstance(talker, VllmOmniARStageConfig)
assert isinstance(code2wav, VllmOmniGenerationStageConfig)
assert [stage.stage_id for stage in resolved_stages] == [0, 1, 2]
assert [stage.model_stage for stage in resolved_stages] == ["thinker", "talker", "code2wav"]
assert code2wav.final_output_type == "audio"
assert thinker.runtime_config.devices == "0"
assert talker.runtime_config.devices == code2wav.runtime_config.devices == "1"
assert talker.connector_config.input_connectors == {"from_stage_0": "connector_of_shared_memory"}
assert code2wav.connector_config.input_connectors == {"from_stage_1": "connector_of_shared_memory"}
# Engine-owned deploy settings must survive the startup projection into
# the actual VllmConfig objects retained by the stage pools.
assert len(engine.stage_vllm_configs) == 3
thinker_vllm, talker_vllm, code2wav_vllm = engine.stage_vllm_configs
assert thinker_vllm is not None
assert talker_vllm is not None
assert code2wav_vllm is not None
assert thinker_vllm.scheduler_config.max_num_seqs == 64
assert talker_vllm.scheduler_config.max_num_seqs == 64
assert code2wav_vllm.scheduler_config.max_num_seqs == 64
assert thinker_vllm.scheduler_config.max_num_batched_tokens == 32768
assert talker_vllm.scheduler_config.max_num_batched_tokens == 32768
assert code2wav_vllm.scheduler_config.max_num_batched_tokens == 65536
assert thinker_vllm.cache_config.gpu_memory_utilization == 0.9
assert talker_vllm.cache_config.gpu_memory_utilization == 0.6
assert code2wav_vllm.cache_config.gpu_memory_utilization == 0.1
assert all(config.cache_config.enable_prefix_caching is False for config in engine.stage_vllm_configs)
assert all(config.model_config.trust_remote_code is True for config in engine.stage_vllm_configs)
assert code2wav_vllm.scheduler_config.enable_chunked_prefill is False
assert code2wav_vllm.scheduler_config.async_scheduling is False
# Sampling defaults are consumed from StageClient metadata, rather than
# from either the resolver output or VllmConfig.
expected_sampling = (
{"temperature": 0.0, "max_tokens": 2048},
{"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "repetition_penalty": 1.05},
{
"temperature": 0.0,
"top_p": 1.0,
"top_k": -1,
"max_tokens": 65536,
"repetition_penalty": 1.1,
},
)
assert len(engine.default_sampling_params_list) == len(expected_sampling)
for runtime_params, expected in zip(engine.default_sampling_params_list, expected_sampling, strict=True):
# SamplingParams normalization may add backend defaults (for example
# ``detokenize=True``). Verify deploy-owned values without rejecting
# those non-lossy normalized fields.
assert all(getattr(runtime_params, name) == value for name, value in expected.items())
if current_omni_platform.is_cuda():
assert thinker_vllm.model_config.enforce_eager is False
assert talker_vllm.model_config.enforce_eager is False
# Exercise the materialized stages and their connector path with a real
# request; the assertions above are not merely resolver-shape checks.
offline_client.send_omni_request(
{
"prompts": "Answer with one short sentence: what is the capital of China?",
"modalities": ["text"],
}
)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
tests/e2e/offline_inference/test_wan22_t2v.py
+29
View File
@@ -5,6 +5,8 @@ import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OfflineOmniClient
from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
@@ -12,6 +14,33 @@ PROMPT = "Two anthropomorphic cats in boxing gear on a spotlighted stage."
NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1})
@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
def test_structured_diffusion_config_reaches_runtime(omni_runner) -> None:
"""The default WAN resolver launches the typed diffusion config itself."""
omni = omni_runner.omni
resolved_stages = omni.engine.stage_configs
stage_clients = omni.engine.stage_clients
assert len(resolved_stages) == 1
assert len(stage_clients) == 1
stage_config = resolved_stages[0]
od_config = stage_clients[0].od_config
assert isinstance(stage_config, VllmOmniDiffusionStageConfig)
assert isinstance(od_config, OmniDiffusionConfig)
assert stage_config.stage_id == 0
assert stage_config.model_stage == "dit"
assert stage_config.final_output_type == "video"
assert omni.engine.stage_vllm_configs == [None]
assert od_config.stage_id == 0
assert od_config.model == MODEL
assert od_config.model_class_name == "WanPipeline"
assert od_config.output_type == "pil"
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1})
tests/e2e/online_serving/personaplex_realtime_duplex.py
+479 -88
View File
@@ -10,24 +10,24 @@ import base64
import hashlib
import json
import math
import time
import uuid
import wave
from collections.abc import Sequence
from collections.abc import Awaitable, Callable, Sequence
from pathlib import Path
from typing import Protocol
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import numpy as np
try:
import websockets
from websockets.exceptions import ConnectionClosed
from websockets.exceptions import ConnectionClosed, WebSocketException
except ImportError as exc: # pragma: no cover - driver dependency.
raise SystemExit("Install websockets first: pip install websockets") from exc
from vllm_omni.clients.duplex import (
EventCollector,
write_pcm16_wav,
)
from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats
from vllm_omni.clients.duplex import EventCollector, write_pcm16_wav
from vllm_omni.clients.duplex import (
wait_for_condition as wait_for,
)
@@ -35,6 +35,24 @@ from vllm_omni.clients.duplex import (
SAMPLE_RATE_HZ = 24_000
FRAME_SAMPLES = 1_920
FRAME_PERIOD_S = FRAME_SAMPLES / SAMPLE_RATE_HZ
AUDIO_DELTA_EVENT_TYPES = frozenset({"response.audio.delta", "response.output_audio.delta"})
class _AudioAcceptanceError(AssertionError):
"""A local acceptance failure with numeric diagnostics safe to report."""
def __init__(self, check: str, message: str, stats: dict[str, object]) -> None:
super().__init__(message)
self.check = check
self.stats = dict(stats)
class _ProbeTransport(Protocol):
async def recv(self) -> str | bytes: ...
async def send(self, message: str) -> None: ...
async def close(self) -> None: ...
class RawRealtimeProbe:
@@ -46,29 +64,33 @@ class RawRealtimeProbe:
keeps a transport-only client and reads everything from ``.events``.
"""
def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None:
def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024, close_timeout_s: float = 10.0) -> None:
self.url = url
self.max_size = max_size
self.close_timeout_s = close_timeout_s
self.events = EventCollector()
self._ws = None
self._ws: _ProbeTransport | None = None
self._reader_task: asyncio.Task[None] | None = None
async def __aenter__(self) -> RawRealtimeProbe:
self._ws = await websockets.connect(self.url, max_size=self.max_size)
self._ws = await websockets.connect(self.url, max_size=self.max_size, close_timeout=self.close_timeout_s)
self._reader_task = asyncio.create_task(self._read_events())
return self
async def __aexit__(self, exc_type, exc, traceback) -> None:
if self._ws is not None:
await self._ws.close()
if self._reader_task is not None:
self._reader_task.cancel()
try:
await self._reader_task
except asyncio.CancelledError:
pass
try:
if self._ws is not None:
await self._ws.close()
finally:
if self._reader_task is not None:
self._reader_task.cancel()
try:
await self._reader_task
except asyncio.CancelledError:
pass
async def _read_events(self) -> None:
assert self._ws is not None
try:
while True:
raw = await self._ws.recv()
@@ -81,6 +103,7 @@ class RawRealtimeProbe:
return
async def send(self, event: dict[str, object]) -> None:
assert self._ws is not None
await self._ws.send(json.dumps(event))
@@ -142,6 +165,38 @@ def _events(client: RawRealtimeProbe, event_type: str) -> list[dict[str, object]
return [event for event in client.events.events if event.get("type") == event_type]
def _audio_events(client: RawRealtimeProbe) -> list[dict[str, object]]:
"""Return audio deltas from both current and legacy Realtime event names."""
return [event for event in client.events.events if event.get("type") in AUDIO_DELTA_EVENT_TYPES]
def _validated_audio_event(event: dict[str, object]) -> bytes:
"""Decode one wire audio delta without depending on client event aliases."""
encoded = event.get("delta")
if not isinstance(encoded, str):
raise ValueError("missing_audio_delta")
try:
chunk = base64.b64decode(encoded, validate=True)
except ValueError:
raise ValueError("invalid_audio_base64") from None
if chunk:
if len(chunk) % 2:
raise ValueError("unaligned_pcm16_packet")
if event.get("sample_rate_hz") != SAMPLE_RATE_HZ:
raise ValueError("invalid_audio_sample_rate")
if not isinstance(event.get("response_id"), str) or not event.get("response_id"):
raise ValueError("missing_audio_response_id")
return chunk
def _validated_audio_chunks(client: RawRealtimeProbe) -> list[tuple[dict[str, object], bytes]]:
return [(event, _validated_audio_event(event)) for event in _audio_events(client)]
def _audio_bytes(client: RawRealtimeProbe) -> bytes:
return b"".join(chunk for _, chunk in _validated_audio_chunks(client) if chunk)
async def _open_session(
args: argparse.Namespace,
*,
@@ -151,22 +206,7 @@ async def _open_session(
) -> tuple[RawRealtimeProbe, dict[str, object]]:
client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id))
await client.__aenter__()
await client.send(
{
"type": "session.update",
"session": {
"session_id": session_id,
"model": args.model,
"modalities": ["audio", "text"],
"input_audio_format": "pcm_f32le",
"output_audio_format": "pcm16",
"voice": args.voice,
"instructions": persona,
"turn_detection": None,
"extra_body": {"auto_response": True},
},
}
)
await client.send(_session_update(args, session_id=session_id, persona=persona))
event_type = "error" if expect_error else "session.created"
await wait_for(
lambda: client.events.count(event_type) > 0,
@@ -176,6 +216,23 @@ async def _open_session(
return client, _events(client, event_type)[-1]
def _session_update(args: argparse.Namespace, *, session_id: str, persona: str) -> dict[str, object]:
return {
"type": "session.update",
"session": {
"session_id": session_id,
"model": args.model,
"modalities": ["audio", "text"],
"input_audio_format": "pcm_f32le",
"output_audio_format": "pcm16",
"voice": args.voice,
"instructions": persona,
"turn_detection": None,
"extra_body": {"auto_response": True},
},
}
async def _close_session(client: RawRealtimeProbe, *, timeout_s: float) -> None:
await client.send({"type": "session.close"})
await wait_for(
@@ -200,23 +257,24 @@ async def _stream_frames(
frame_count = min(frame_count, max_frames)
for seq in range(frame_count):
frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size))
frame = np.ascontiguousarray(frame, dtype="<f4")
await client.send(
{
"type": "input_audio_buffer.append",
"audio": base64.b64encode(frame).decode("ascii"),
"format": "pcm_f32le",
"sample_rate_hz": SAMPLE_RATE_HZ,
"duration_ms": 80,
"audio_end_ms": (seq + 1) * 80,
"is_speech": bool(np.sqrt(np.mean(np.square(frame))) >= 1e-4),
}
)
await client.send(_frame_event(frame, seq))
await asyncio.sleep(FRAME_PERIOD_S)
return frame_count
def _frame_event(frame: np.ndarray, seq: int) -> dict[str, object]:
frame = np.ascontiguousarray(np.pad(frame, (0, FRAME_SAMPLES - frame.size)), dtype="<f4")
return {
"type": "input_audio_buffer.append",
"audio": base64.b64encode(frame).decode("ascii"),
"format": "pcm_f32le",
"sample_rate_hz": SAMPLE_RATE_HZ,
"duration_ms": 80,
"audio_end_ms": (seq + 1) * 80,
"is_speech": bool(np.sqrt(np.mean(np.square(frame))) >= 1e-4),
}
def _audio_frame_stats(
raw: bytes,
*,
@@ -242,11 +300,7 @@ def _audio_frame_stats(
def _response_ids(client: RawRealtimeProbe) -> set[str]:
return {
str(event["response_id"])
for event in _events(client, "response.output_audio.delta")
if isinstance(event.get("response_id"), str)
}
return {str(event["response_id"]) for event in _audio_events(client) if isinstance(event.get("response_id"), str)}
def _session_result(
@@ -256,40 +310,66 @@ def _session_result(
args: argparse.Namespace,
minimum_chunks: int,
) -> tuple[bytes, set[str], dict[str, object]]:
chunks = sum(len(items) for items in client.events.response_audio.values())
raw = client.events.audio_bytes()
if chunks < minimum_chunks or not raw or len(raw) % 2:
raise AssertionError(f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}")
if isinstance(client, RawRealtimeProbe):
audio_chunks = _validated_audio_chunks(client)
chunks = sum(bool(chunk) for _, chunk in audio_chunks)
raw = b"".join(chunk for _, chunk in audio_chunks if chunk)
else:
# Unit fixtures may provide an already-decoded EventCollector without
# retaining the base64 wire field. Real websocket probes always take
# the strict wire-validation path above.
chunks = sum(len(items) for items in client.events.response_audio.values())
raw = client.events.audio_bytes()
stats: dict[str, object] = {"audio_chunks": chunks, "audio_samples": len(raw) // 2}
if not raw or len(raw) % 2:
raise _AudioAcceptanceError("pcm16_output", f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}", stats)
pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
rms = float(np.sqrt(np.mean(np.square(pcm))))
stats["audio_rms"] = rms
# This low floor only catches corrupt, non-finite, or effectively zero PCM.
# Audible speech is checked per frame below with the higher 1e-3 default.
if not np.isfinite(pcm).all() or rms <= 1e-5:
raise AssertionError(f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}")
rates = {
int(event["sample_rate_hz"])
for event in _events(client, "response.output_audio.delta")
if isinstance(event.get("sample_rate_hz"), int)
}
raise _AudioAcceptanceError(
"audio_rms", f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}", stats
)
rates = {rate for event in _audio_events(client) if isinstance(rate := event.get("sample_rate_hz"), int)}
if rates != {SAMPLE_RATE_HZ}:
raise AssertionError(f"unexpected output sample rates: {sorted(rates)}")
stats = _audio_frame_stats(
raw,
input_frames=input_frames,
voiced_frame_rms_threshold=args.voiced_frame_rms_threshold,
)
if not 0 <= int(stats["frame_deficit"]) <= args.max_frame_deficit:
raise AssertionError(f"invalid synchronous frame accounting: input={input_frames}, stats={stats}")
raise _AudioAcceptanceError("sample_rate", f"unexpected output sample rates: {sorted(rates)}", stats)
if input_frames < 1:
raise _AudioAcceptanceError("input_frames", f"audio received without input frames: input={input_frames}", stats)
try:
frame_stats = _audio_frame_stats(
raw,
input_frames=input_frames,
voiced_frame_rms_threshold=args.voiced_frame_rms_threshold,
)
except AssertionError as exc:
# _audio_frame_stats only formats local sample counts, not wire data.
raise _AudioAcceptanceError("whole_codec_frames", str(exc), stats) from exc
stats.update(frame_stats)
if chunks < minimum_chunks:
raise _AudioAcceptanceError(
"audio_chunks", f"insufficient audio chunks: chunks={chunks}, minimum={minimum_chunks}", stats
)
if not 0 <= int(frame_stats["frame_deficit"]) <= args.max_frame_deficit:
raise _AudioAcceptanceError(
"frame_deficit",
f"invalid synchronous frame accounting: input={input_frames}, maximum_deficit={args.max_frame_deficit}, "
f"stats={stats}",
stats,
)
min_voiced_frames = int(args.min_voiced_frames)
if min_voiced_frames < 1:
raise ValueError(f"min_voiced_frames must be positive, got {min_voiced_frames}")
if int(stats["voiced_frames"]) < min_voiced_frames:
raise AssertionError(
f"output contained insufficient audible speech: minimum_voiced_frames={min_voiced_frames}, stats={stats}"
if int(frame_stats["voiced_frames"]) < min_voiced_frames:
raise _AudioAcceptanceError(
"voiced_frames",
f"output contained insufficient audible speech: minimum_voiced_frames={min_voiced_frames}, stats={stats}",
stats,
)
runtime_metadata = [
metadata["vllm_omni"]
for event in _events(client, "response.output_audio.delta")
for event in _audio_events(client)
if isinstance((metadata := event.get("metadata")), dict) and isinstance(metadata.get("vllm_omni"), dict)
]
if not any(
@@ -298,18 +378,19 @@ def _session_result(
and item.get("runner_kv_backed") is True
for item in runtime_metadata
):
raise AssertionError(f"audio events did not prove the scheduler data plane: {runtime_metadata}")
raise _AudioAcceptanceError(
"scheduler_data_plane", "audio events did not prove the scheduler data plane", stats
)
response_ids = _response_ids(client)
if len(response_ids) != 1:
raise AssertionError(f"continuous stream used multiple responses: {sorted(response_ids)}")
raise _AudioAcceptanceError(
"response_identity", f"continuous stream requires one response: count={len(response_ids)}", stats
)
return (
raw,
response_ids,
{
**stats,
"audio_chunks": chunks,
"audio_samples": len(raw) // 2,
"audio_rms": rms,
"response_ids": sorted(response_ids),
},
)
@@ -332,6 +413,8 @@ def _save(output_dir: Path, name: str, client: RawRealtimeProbe, raw: bytes) ->
async def run(args: argparse.Namespace) -> dict[str, object]:
if getattr(args, "sessions", None) is not None:
return await run_load(args)
input_path = Path(args.input_wav)
input_identity = _input_identity(
input_path,
@@ -381,14 +464,12 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
)
await asyncio.gather(
wait_for(
lambda: len(primary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
lambda: len(_audio_bytes(primary)) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
timeout_s=args.timeout_s,
label="primary frame coverage",
),
wait_for(
lambda: (
len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit
),
lambda: (len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit),
timeout_s=args.timeout_s,
label="secondary frame coverage",
),
@@ -426,16 +507,14 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
await asyncio.gather(
wait_for(
lambda: (
len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES)
>= secondary_total_frames - args.max_frame_deficit
len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_total_frames - args.max_frame_deficit
),
timeout_s=args.timeout_s,
label="survivor frame coverage after slot recycle",
),
wait_for(
lambda: (
len(replacement.events.audio_bytes()) // (2 * FRAME_SAMPLES)
>= replacement_frames - args.max_frame_deficit
len(_audio_bytes(replacement)) // (2 * FRAME_SAMPLES) >= replacement_frames - args.max_frame_deficit
),
timeout_s=args.timeout_s,
label="replacement frame coverage",
@@ -493,6 +572,284 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
return result
def _distribution(values: list[float]) -> dict[str, int | float | None]:
if not values:
return {"count": 0, "median": None, "p99": None, "max": None, "argmax_index": None}
return {
"count": len(values),
"median": float(np.median(values)),
"p99": float(np.quantile(values, 0.99)),
"max": max(values),
"argmax_index": int(np.argmax(values)),
}
def _load_metrics(client: RawRealtimeProbe, sends: list[tuple[float, float, float]]) -> dict[str, object]:
# Use the collector's local clock, never a timestamp supplied by the server.
packets: list[dict[str, object]] = []
response_ids: set[str] = set()
packet_times: list[float] = []
packet_bytes: list[int] = []
invalid_packets: list[dict[str, object]] = []
samples = 0
for index, (raw, received) in enumerate(zip(client.events.events, client.events.event_received_at_s, strict=True)):
if raw.get("type") not in AUDIO_DELTA_EVENT_TYPES:
continue
try:
chunk = _validated_audio_event(raw)
except ValueError as exc:
# Only fixed diagnostic codes from the validator are exported,
# never audio payloads or details from the remote endpoint.
invalid_packets.append({"event_index": index, "received_at_s": received, "code": str(exc)})
continue
if chunk:
samples += len(chunk) // 2
packet_times.append(received)
packet_bytes.append(len(chunk))
response_id = raw.get("response_id")
packets.append({"received_at_s": received, "samples": len(chunk) // 2, "response_id": response_id})
if isinstance(response_id, str) and response_id:
response_ids.add(response_id)
intervals = [(b - a) * 1000 for a, b in zip(packet_times, packet_times[1:])]
first_send = sends[0][1] if sends else None
audio_duration = samples / SAMPLE_RATE_HZ
first_audio_ms = (packet_times[0] - first_send) * 1000 if packet_times and first_send is not None else None
client_rtf = (packet_times[-1] - first_send) / audio_duration if audio_duration and first_send is not None else None
continuity = compute_continuity_stats(packet_times, packet_bytes, SAMPLE_RATE_HZ) if packet_times else None
lateness = _distribution([max(0.0, sent - planned) * 1000 for planned, sent, _ in sends])
p99_lateness = lateness["p99"]
return {
"input_frames": len(sends),
"output_samples": samples,
# Populate frame accounting from _session_result, including on failure.
"frame_deficit": None,
"response_ids": sorted(response_ids),
"client_first_audio_after_stream_start_ms": first_audio_ms,
"client_stream_rtf": client_rtf,
"client_audio_packet_interval_ms": _distribution(intervals),
"client_send_lateness_ms": lateness,
"client_pacing_warning": isinstance(p99_lateness, float) and p99_lateness > FRAME_PERIOD_S * 1000,
"audio_underrun_s": continuity.max_underrun_s if continuity else None,
"audio_underrun_event_count": continuity.underrun_event_count if continuity else None,
"audio_continuity_ok": continuity.is_continuous if continuity else None,
"input_timeline": [
{"frame_index": i, "planned_at_s": planned, "send_started_at_s": sent, "send_completed_at_s": completed}
for i, (planned, sent, completed) in enumerate(sends)
],
"audio_packet_timeline": packets,
"invalid_audio_packets": invalid_packets,
}
def _check_load_connection(client: RawRealtimeProbe) -> None:
if client.events.errors():
raise RuntimeError("server emitted a Realtime error (see the session error count)")
if client._reader_task is not None and client._reader_task.done():
client._reader_task.result() # Propagate malformed JSON/reader failures.
raise RuntimeError("connection closed before session completion")
async def _wait_load_event(client: RawRealtimeProbe, event_type: str, timeout_s: float) -> None:
def received() -> bool:
if client.events.count(event_type):
return True
_check_load_connection(client)
return False
await wait_for(received, timeout_s=timeout_s, label=event_type)
async def _paced_load_frames(
client: RawRealtimeProbe,
pcm: np.ndarray,
*,
epoch: float,
timeout_s: float,
sends: list[tuple[float, float, float]],
clock: Callable[[], float] = time.monotonic,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
) -> None:
next_send = epoch
for seq in range(math.ceil(pcm.size / FRAME_SAMPLES)):
await sleep(max(0.0, next_send - clock()))
_check_load_connection(client)
event = _frame_event(pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES], seq)
started = clock()
await asyncio.wait_for(client.send(event), timeout=timeout_s)
completed = clock()
sends.append((epoch + seq * FRAME_PERIOD_S, started, completed))
# Do not burst to catch up after a stalled send; report accumulated lag.
next_send = max(epoch + (seq + 1) * FRAME_PERIOD_S, completed + FRAME_PERIOD_S)
async def _request_load_close(client: RawRealtimeProbe, timeout_s: float) -> None:
await client.send({"type": "session.close"})
await _wait_load_event(client, "session.closed", timeout_s)
async def _run_load_session(
args: argparse.Namespace,
pcm: np.ndarray,
index: int,
ready: asyncio.Future[None],
start: asyncio.Future[float],
) -> dict[str, object]:
session_id = f"personaplex-load-{index}-{uuid.uuid4().hex}"
client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id), close_timeout_s=args.cleanup_timeout_s)
sends: list[tuple[float, float, float]] = []
error: str | None = None
cleanup_error: str | None = None
metrics: dict[str, object] | None = None
session_stats: dict[str, object] = {}
acceptance_check: str | None = None
phase = "admission"
try:
async with client:
try:
await asyncio.wait_for(
client.send(_session_update(args, session_id=session_id, persona=args.persona)), args.timeout_s
)
await _wait_load_event(client, "session.created", args.timeout_s)
capabilities = _capabilities(_events(client, "session.created")[-1])
if (
capabilities.get("chunk_period_ms") != 80
or capabilities.get("supports_multi_session_same_replica") is not True
):
raise AssertionError("server did not advertise PersonaPlex multi-session capabilities")
ready.set_result(None)
phase = "stream"
epoch = await asyncio.shield(start)
await _paced_load_frames(client, pcm, epoch=epoch, timeout_s=args.timeout_s, sends=sends)
phase = "drain"
# Fixed, reported observation window. Closing must not extend
# it and let a slow session pass by flushing after the cutoff.
await asyncio.sleep(args.drain_s)
_check_load_connection(client)
_, _, session_stats = _session_result(
client, input_frames=len(sends), args=args, minimum_chunks=args.minimum_audio_chunks
)
finally:
metrics = _load_metrics(client, sends)
if client.events.count("session.created") and not client.events.count("session.closed"):
try:
# Send and acknowledgement share one short cleanup budget,
# independent of the potentially long inference timeout.
await asyncio.wait_for(
_request_load_close(client, args.cleanup_timeout_s), args.cleanup_timeout_s
)
except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, WebSocketException) as exc:
cleanup_error = f"cleanup: {type(exc).__name__}"
except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, AssertionError, WebSocketException) as exc:
if isinstance(exc, _AudioAcceptanceError):
acceptance_check = exc.check
session_stats = exc.stats
error = f"{phase}: AssertionError: {exc}"
else:
# Other assertions/errors may embed handshake credentials or wire
# metadata. Only our local numeric acceptance diagnostics are safe.
error = f"{phase}: {type(exc).__name__}"
finally:
# Admission failure must not leave the other sessions at the barrier.
if not ready.done():
ready.set_result(None)
if metrics is None:
metrics = _load_metrics(client, sends)
if cleanup_error:
error = f"{error}; {cleanup_error}" if error else cleanup_error
if client.events.errors():
error = error or "server emitted a Realtime error"
if metrics["invalid_audio_packets"]:
error = error or "invalid audio packets"
first_audio = metrics["client_first_audio_after_stream_start_ms"]
if isinstance(first_audio, float) and first_audio < 0:
error = error or "audio arrived before the first input frame"
rtf = metrics["client_stream_rtf"]
if args.max_client_rtf is not None and isinstance(rtf, float) and rtf > args.max_client_rtf:
error = error or f"client_stream_rtf exceeds {args.max_client_rtf}"
return {
"index": index,
"session_id": session_id,
"ok": error is None,
"error": error,
"acceptance_check": acceptance_check,
"server_error_count": len(client.events.errors()),
"server_error_codes": [
str(body.get("code", "unknown")) if isinstance(body := event.get("error"), dict) else "unknown"
for event in client.events.errors()
],
**metrics,
**session_stats,
}
async def run_load(args: argparse.Namespace) -> dict[str, object]:
input_path = Path(args.input_wav)
identity = _input_identity(input_path, expected_sha256=args.expected_input_sha256)
pcm = _read_wav_as_float32(input_path)
if args.tail_s:
pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)])
pcm = pcm[: args.load_frames * FRAME_SAMPLES]
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
# Refuse to silently overwrite evidence from a previous sweep.
result_path = output_dir / "load-result.json"
if result_path.exists():
raise FileExistsError(result_path)
loop = asyncio.get_running_loop()
start: asyncio.Future[float] = loop.create_future()
ready: list[asyncio.Future[None]] = [loop.create_future() for _ in range(args.sessions)]
tasks = [asyncio.create_task(_run_load_session(args, pcm, i, r, start)) for i, r in enumerate(ready)]
try:
await asyncio.gather(*ready)
start.set_result(time.monotonic())
rows = await asyncio.gather(*tasks)
finally:
for task in tasks:
if not task.done():
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
seen: set[str] = set()
shared: set[str] = set()
for row in rows:
assert isinstance(row["response_ids"], list)
response_ids = set(row["response_ids"])
shared.update(seen & response_ids)
seen.update(response_ids)
for row in rows:
assert isinstance(row["response_ids"], list)
if shared.intersection(row["response_ids"]):
row.update(ok=False, error="concurrent sessions shared a response id")
result = {
"schema_version": 1,
"mode": "paced_multi_session",
"model": args.model,
"requested_sessions": args.sessions,
"passed_sessions": sum(row["ok"] is True for row in rows),
"ok": all(row["ok"] is True for row in rows),
"input": {**identity, "tail_s": args.tail_s, "load_frames": args.load_frames},
"frame_period_s": FRAME_PERIOD_S,
"drain_s": args.drain_s,
"cleanup_timeout_s": args.cleanup_timeout_s,
"client_pacing_warning_threshold_ms": FRAME_PERIOD_S * 1000,
"client_pacing_warning_sessions": [row["index"] for row in rows if row["client_pacing_warning"]],
"max_client_rtf": args.max_client_rtf,
"acceptance": {
"max_frame_deficit": args.max_frame_deficit,
"min_voiced_frames": args.min_voiced_frames,
"minimum_audio_chunks": args.minimum_audio_chunks,
"voiced_frame_rms_threshold": args.voiced_frame_rms_threshold,
},
"shared_response_ids": sorted(shared),
"sessions": rows,
"reported_server": {"revision": args.server_revision, "hardware": args.server_hardware},
"measurement_origin": "client monotonic clock; admission excluded; streaming and drain retained",
"interval_definition": "nonempty audio packet arrivals, not server ticks or per-frame inference latency",
}
with result_path.open("x", encoding="utf-8") as stream:
json.dump(result, stream, ensure_ascii=False, indent=2, allow_nan=False)
return result
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
@@ -510,6 +867,12 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser.add_argument("--tail-s", type=float, default=2.0)
parser.add_argument("--drain-s", type=float, default=2.0)
parser.add_argument("--timeout-s", type=float, default=180.0)
parser.add_argument(
"--cleanup-timeout-s",
type=float,
default=5.0,
help="Load mode: shared session-close/ack budget; also bounds the websocket close handshake.",
)
parser.add_argument(
"--max-frame-deficit",
type=int,
@@ -536,11 +899,39 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
default=0,
help="Frames sent after slot reuse; 0 replays the full input so the replacement must produce audible output.",
)
return parser.parse_args(argv)
parser.add_argument(
"--sessions", type=int, help="Opt into paced load mode with N sessions; omit for the lifecycle test."
)
parser.add_argument(
"--load-frames", type=int, default=1000, help="Maximum input frames per load session (no repetition)."
)
parser.add_argument("--max-client-rtf", type=float, help="Optional client stream-window RTF acceptance ceiling.")
parser.add_argument("--server-revision", help="User-reported server commit/version; not verified by this client.")
parser.add_argument(
"--server-hardware", help="User-reported server hardware; not inferred from the client machine."
)
args = parser.parse_args(argv)
if args.sessions is not None:
if not 1 <= args.sessions <= 64 or not 1 <= args.load_frames <= 10000:
parser.error("load mode requires sessions in [1,64] and load-frames in [1,10000]")
for name in ("timeout_s", "cleanup_timeout_s", "drain_s", "tail_s"):
value = getattr(args, name)
if not math.isfinite(value) or value < 0 or (name in ("timeout_s", "cleanup_timeout_s") and value == 0):
parser.error(f"{name} must be finite and non-negative (timeout must be positive)")
if args.max_client_rtf is not None and (not math.isfinite(args.max_client_rtf) or args.max_client_rtf <= 0):
parser.error("max-client-rtf must be finite and positive")
if args.max_frame_deficit < 0 or args.min_voiced_frames < 1 or args.minimum_audio_chunks < 1:
parser.error("frame-deficit must be non-negative; voiced frames and audio chunks must be positive")
if not math.isfinite(args.voiced_frame_rms_threshold) or args.voiced_frame_rms_threshold < 0:
parser.error("voiced-frame-rms-threshold must be finite and non-negative")
return args
def main() -> None:
print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
result = asyncio.run(run(parse_args()))
print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False))
if result.get("mode") == "paced_multi_session" and not result["ok"]:
raise SystemExit(1)
if __name__ == "__main__":
tests/e2e/online_serving/test_personaplex_load_driver.py
+568
View File
@@ -0,0 +1,568 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import base64
import json
import wave
from contextlib import asynccontextmanager
from typing import TypedDict
import numpy as np
import pytest
import websockets
from tests.e2e.online_serving import personaplex_realtime_duplex as driver
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _ServerState(TypedDict):
index: int
closed: bool
frames: int
close_requested: bool
def _args(tmp_path, *extra):
wav_path = tmp_path / "input.wav"
with wave.open(str(wav_path), "wb") as stream:
stream.setnchannels(1)
stream.setsampwidth(2)
stream.setframerate(driver.SAMPLE_RATE_HZ)
stream.writeframes(np.full(driver.FRAME_SAMPLES * 2, 2000, dtype="<i2").tobytes())
return driver.parse_args(
[
"--model",
"fixture-only",
"--input-wav",
str(wav_path),
"--output-dir",
str(tmp_path / "out"),
"--sessions",
"2",
"--load-frames",
"2",
"--tail-s",
"0",
"--drain-s",
"0.04",
"--timeout-s",
"3",
"--cleanup-timeout-s",
"1",
"--min-voiced-frames",
"1",
"--minimum-audio-chunks",
"1",
"--max-frame-deficit",
"0",
*extra,
]
)
def _audio(response_id="r", samples=1920, event_type="response.output_audio.delta"):
return {
"type": event_type,
"response_id": response_id,
"delta": base64.b64encode(np.full(samples, 2000, dtype="<i2").tobytes()).decode("ascii"),
"sample_rate_hz": driver.SAMPLE_RATE_HZ,
"metadata": {
"vllm_omni": {
"runtime_impl": "scheduler_data_plane",
"uses_model_runner_scheduler": True,
"runner_kv_backed": True,
}
},
}
@asynccontextmanager
async def _server(mode="success"):
states: list[_ServerState] = []
async def handle(ws):
index = len(states)
state: _ServerState = {"index": index, "closed": False, "frames": 0, "close_requested": False}
states.append(state)
response_id = "shared" if mode == "shared_response" else f"r-{index}"
try:
async for raw in ws:
message = json.loads(raw)
if message["type"] == "session.update":
if mode == "reject_second" and index == 1:
await ws.send(json.dumps({"type": "error", "error": {"code": "resource_exhausted"}}))
continue
if mode == "missing_capabilities":
await ws.send(
json.dumps(
{
"type": "session.created",
"session": {"resume_token": "must-not-be-exported"},
}
)
)
continue
await ws.send(
json.dumps(
{
"type": "session.created",
"session": {
"resume_token": "must-not-be-exported",
"capabilities": {
"chunk_period_ms": 80,
"supports_multi_session_same_replica": True,
},
},
}
)
)
await ws.send(json.dumps({"type": "response.created", "response": {"id": response_id}}))
elif message["type"] == "input_audio_buffer.append":
state["frames"] += 1
if mode == "short_audio" and state["frames"] == 2:
continue
if mode == "disconnect":
await ws.close()
return
if mode == "bad_json":
await ws.send("{")
elif mode not in ("silent", "late_on_close"):
event_type = (
"response.audio.delta" if mode == "legacy_audio_event" else "response.output_audio.delta"
)
packet = _audio(response_id, event_type=event_type)
if mode == "private_metadata":
packet["metadata"] = {"vllm_omni": {"private": "remote-detail-must-not-be-exported"}}
if mode == "odd_packet":
raw_audio = base64.b64decode(packet["delta"])
raw_audio = raw_audio + b"\x00" if state["frames"] == 1 else raw_audio[:-1]
packet["delta"] = base64.b64encode(raw_audio).decode("ascii")
if state["frames"] == 2:
if mode == "missing_packet_rate":
packet.pop("sample_rate_hz")
elif mode == "string_packet_rate":
packet["sample_rate_hz"] = "24000"
elif mode == "invalid_base64":
packet["delta"] = "!" + packet["delta"]
elif mode == "orphan_audio":
orphan = _audio(response_id)
orphan.pop("response_id")
await ws.send(json.dumps(orphan))
await ws.send(json.dumps(packet))
elif message["type"] == "session.close":
state["close_requested"] = True
if mode == "late_on_close":
await ws.send(json.dumps(_audio(response_id, samples=driver.FRAME_SAMPLES * state["frames"])))
if mode == "error_on_close":
await ws.send(json.dumps({"type": "error", "error": {"code": "cleanup_failed"}}))
if mode != "no_close_ack":
await ws.send(json.dumps({"type": "session.closed"}))
finally:
state["closed"] = True
async with websockets.serve(handle, "127.0.0.1", 0) as server:
port = server.sockets[0].getsockname()[1]
yield f"ws://127.0.0.1:{port}/v1/realtime", states
def test_omitted_sessions_preserves_lifecycle_mode():
args = driver.parse_args(["--model", "unused", "--input-wav", "unused.wav"])
assert args.sessions is None
@pytest.mark.parametrize(
"options",
[
["--sessions", "0"],
["--sessions", "65"],
["--load-frames", "0"],
["--load-frames", "10001"],
["--drain-s", "nan"],
["--tail-s", "-1"],
["--timeout-s", "0"],
["--timeout-s", "inf"],
["--cleanup-timeout-s", "0"],
["--cleanup-timeout-s", "nan"],
["--max-client-rtf", "-1"],
["--max-frame-deficit", "-1"],
["--minimum-audio-chunks", "0"],
["--min-voiced-frames", "0"],
["--voiced-frame-rms-threshold", "nan"],
],
)
def test_load_rejects_invalid_configuration(tmp_path, options):
with pytest.raises(SystemExit):
_args(tmp_path, *options)
def test_metric_clock_and_packet_frame_distinction():
client = driver.RawRealtimeProbe("ws://unused")
client.events.add({"type": "session.created", "session": {"resume_token": "secret"}}, received_at_s=7)
for when in (10.4, 10.8, 11.6):
event = _audio(samples=driver.FRAME_SAMPLES * 5)
event["_client_received_at_s"] = -999.0
client.events.add(event, received_at_s=when)
report = driver._load_metrics(client, [(10, 10, 10.01)])
intervals = report["client_audio_packet_interval_ms"]
assert isinstance(intervals, dict)
assert intervals["count"] == 2
assert intervals["median"] == pytest.approx(600)
assert intervals["p99"] == pytest.approx(796)
assert intervals["max"] == pytest.approx(800)
assert intervals["argmax_index"] == 1
assert report["client_first_audio_after_stream_start_ms"] == pytest.approx(400)
assert report["output_samples"] == 15 * driver.FRAME_SAMPLES
assert report["audio_underrun_s"] == pytest.approx(0.4)
assert report["audio_continuity_ok"] is False
underruns = report["audio_underrun_event_count"]
assert isinstance(underruns, int) and underruns >= 1
assert report["client_pacing_warning"] is False
assert "secret" not in json.dumps(report)
assert "-999" not in json.dumps(report)
def test_empty_output_has_missing_not_zero_latency():
report = driver._load_metrics(driver.RawRealtimeProbe("ws://unused"), [])
assert report["client_stream_rtf"] is None
assert report["audio_continuity_ok"] is None
assert report["audio_underrun_s"] is None
intervals = report["client_audio_packet_interval_ms"]
assert isinstance(intervals, dict)
assert intervals["p99"] is None
def test_failed_acceptance_retains_frame_counts_without_remote_data():
client = driver.RawRealtimeProbe("ws://unused")
client.events.add(_audio(), received_at_s=1.0)
args = driver.parse_args(
["--model", "unused", "--input-wav", "unused.wav", "--max-frame-deficit", "0", "--min-voiced-frames", "1"]
)
with pytest.raises(AssertionError) as raised:
driver._session_result(client, input_frames=2, args=args, minimum_chunks=1)
assert getattr(raised.value, "check", None) == "frame_deficit"
stats = getattr(raised.value, "stats", {})
assert stats["frame_deficit"] == 1 and type(stats["frame_deficit"]) is int
assert stats["output_frames"] == stats["voiced_frames"] == 1
assert stats["frame_coverage_ratio"] == 0.5
@pytest.mark.parametrize("event_type", sorted(driver.AUDIO_DELTA_EVENT_TYPES))
def test_session_result_accepts_current_and_legacy_audio_event_names(tmp_path, event_type):
client = driver.RawRealtimeProbe("ws://unused")
client.events.add(_audio(samples=driver.FRAME_SAMPLES, event_type=event_type), received_at_s=1.0)
raw, response_ids, stats = driver._session_result(
client,
input_frames=1,
args=_args(tmp_path, "--sessions", "1"),
minimum_chunks=1,
)
assert len(raw) == driver.FRAME_SAMPLES * 2
assert response_ids == {"r"}
assert stats["frame_deficit"] == 0
def test_legacy_audio_event_name_contributes_load_metrics():
client = driver.RawRealtimeProbe("ws://unused")
client.events.add(_audio(samples=driver.FRAME_SAMPLES * 5, event_type="response.audio.delta"), received_at_s=10.4)
report = driver._load_metrics(client, [(10.0, 10.0, 10.01)])
assert report["output_samples"] == driver.FRAME_SAMPLES * 5
assert report["response_ids"] == ["r"]
assert report["invalid_audio_packets"] == []
@pytest.mark.asyncio
@pytest.mark.parametrize("sessions", [1, 2, 4])
async def test_real_websocket_load_success(tmp_path, sessions):
args = _args(tmp_path, "--sessions", str(sessions))
async with _server() as (url, states):
args.url = url
result = await driver.run(args)
assert result["ok"] is True
assert result["passed_sessions"] == sessions
rows = result["sessions"]
assert isinstance(rows, list)
assert len(rows) == sessions
assert all(row["frame_deficit"] == 0 for row in rows)
for row in rows:
assert type(row["frame_deficit"]) is int
assert row["output_frames"] == row["voiced_frames"] == row["audio_chunks"] == 2
assert row["silent_frames"] == 0
assert row["frame_coverage_ratio"] == 1.0
assert row["audio_rms"] == pytest.approx(2000 / 32768)
assert row["acceptance_check"] is None
assert isinstance(row["audio_continuity_ok"], bool)
assert all(s["closed"] and s["close_requested"] for s in states)
assert "must-not-be-exported" not in json.dumps(result)
saved = json.loads((tmp_path / "out/load-result.json").read_text())
assert saved == result
origins = [row["input_timeline"][0]["planned_at_s"] for row in rows]
assert len(set(origins)) == 1
@pytest.mark.asyncio
async def test_real_websocket_load_accepts_legacy_audio_event_name(tmp_path):
args = _args(tmp_path, "--sessions", "2")
async with _server("legacy_audio_event") as (url, states):
args.url = url
result = await driver.run(args)
assert result["ok"] is True
assert result["passed_sessions"] == 2
assert result["shared_response_ids"] == []
assert all(s["closed"] and s["close_requested"] for s in states)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"mode",
[
"reject_second",
"silent",
"disconnect",
"no_close_ack",
"shared_response",
"late_on_close",
"error_on_close",
"bad_json",
"missing_capabilities",
],
)
async def test_real_websocket_failures_remain_in_results(tmp_path, mode):
args = _args(tmp_path)
async with _server(mode) as (url, states):
args.url = url
result = await asyncio.wait_for(driver.run(args), timeout=4)
assert result["ok"] is False
assert "must-not-be-exported" not in json.dumps(result)
assert result["requested_sessions"] == 2
rows = result["sessions"]
assert isinstance(rows, list)
assert len(rows) == 2
assert all(s["closed"] for s in states)
if mode == "reject_second":
assert result["passed_sessions"] == 1
rejected = next(row for row in rows if not row["ok"])
assert rejected["input_frames"] == 0
assert rejected["client_stream_rtf"] is None
assert rejected["server_error_codes"] == ["resource_exhausted"]
if mode == "silent":
# A failed measurement must still release its admitted session.
assert all(s["close_requested"] for s in states)
if mode == "late_on_close":
assert all(row["output_samples"] == 0 for row in rows)
if mode == "shared_response":
assert result["passed_sessions"] == 0
@pytest.mark.asyncio
async def test_existing_report_is_not_overwritten(tmp_path):
args = _args(tmp_path)
path = tmp_path / "out/load-result.json"
path.parent.mkdir()
path.write_text("previous measurement")
with pytest.raises(FileExistsError):
await driver.run(args)
assert path.read_text() == "previous measurement"
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["success", "no_close_ack"])
async def test_cancellation_closes_all_connections(tmp_path, mode):
args = _args(tmp_path, "--drain-s", "10", "--timeout-s", "30")
async with _server(mode) as (url, states):
args.url = url
task = asyncio.create_task(driver.run(args))
async def started():
while len(states) < 2 or not all(state["frames"] for state in states):
await asyncio.sleep(0.01)
await asyncio.wait_for(started(), timeout=3)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=3)
assert all(s["closed"] for s in states)
@pytest.mark.asyncio
async def test_client_rtf_ceiling_marks_slow_sessions_failed(tmp_path):
args = _args(tmp_path, "--max-client-rtf", "0.001")
async with _server() as (url, _):
args.url = url
result = await driver.run(args)
assert result["passed_sessions"] == 0
rows = result["sessions"]
assert isinstance(rows, list)
assert all("client_stream_rtf exceeds" in row["error"] for row in rows)
@pytest.mark.asyncio
async def test_stalled_send_does_not_cause_catchup_burst():
class Clock:
now = 10.0
def monotonic(self):
return self.now
async def sleep(self, delay):
self.now += delay
class Client(driver.RawRealtimeProbe):
async def send(self, event):
if event["audio_end_ms"] == 80:
clock.now += 0.25
clock = Clock()
sends: list[tuple[float, float, float]] = []
await driver._paced_load_frames(
Client("ws://unused"),
np.zeros(driver.FRAME_SAMPLES * 3, dtype="<f4"),
epoch=10.0,
timeout_s=1.0,
sends=sends,
clock=clock.monotonic,
sleep=clock.sleep,
)
assert len(sends) == 3
assert sends[1][1] >= sends[0][2] + driver.FRAME_PERIOD_S
assert sends[1][1] - sends[1][0] == pytest.approx(0.25)
assert sends[2][1] >= sends[1][2] + driver.FRAME_PERIOD_S
metrics = driver._load_metrics(Client("ws://unused"), sends)
assert metrics["client_pacing_warning"] is True
assert asyncio.sleep != clock.sleep
@pytest.mark.asyncio
async def test_failed_frame_acceptance_preserves_numeric_diagnostics(tmp_path):
args = _args(tmp_path, "--sessions", "1")
async with _server("short_audio") as (url, states):
args.url = url
result = await driver.run(args)
rows = result["sessions"]
assert isinstance(rows, list)
row = rows[0]
assert row["ok"] is False
assert row["acceptance_check"] == "frame_deficit"
assert row["frame_deficit"] == 1 and type(row["frame_deficit"]) is int
assert row["output_frames"] == row["voiced_frames"] == row["audio_chunks"] == 1
assert row["frame_coverage_ratio"] == 0.5
assert row["audio_rms"] == pytest.approx(2000 / 32768)
assert "input=2" in row["error"] and "maximum_deficit=0" in row["error"]
assert all(state["closed"] and state["close_requested"] for state in states)
assert json.loads((tmp_path / "out/load-result.json").read_text()) == result
@pytest.mark.asyncio
async def test_acceptance_failure_does_not_export_remote_metadata(tmp_path):
args = _args(tmp_path, "--sessions", "1")
async with _server("private_metadata") as (url, _):
args.url = url
result = await driver.run(args)
rows = result["sessions"]
assert isinstance(rows, list)
row = rows[0]
assert row["ok"] is False
assert row["acceptance_check"] == "scheduler_data_plane"
assert "remote-detail" not in json.dumps(result)
assert "must-not-be-exported" not in json.dumps(result)
@pytest.mark.parametrize("ok", [True, False])
def test_load_main_exit_status(tmp_path, monkeypatch, ok):
args = _args(tmp_path)
async def completed(_args):
return {"mode": "paced_multi_session", "ok": ok}
monkeypatch.setattr(driver, "parse_args", lambda: args)
monkeypatch.setattr(driver, "run", completed)
if ok:
driver.main()
else:
with pytest.raises(SystemExit) as exc:
driver.main()
assert exc.value.code == 1
@pytest.mark.asyncio
async def test_handshake_rejection_is_retained_without_server_details(tmp_path, monkeypatch):
from websockets.exceptions import InvalidHandshake
async def rejected(_self):
raise InvalidHandshake("credentials-in-remote-handshake-detail")
monkeypatch.setattr(driver.RawRealtimeProbe, "__aenter__", rejected)
result = await asyncio.wait_for(driver.run(_args(tmp_path)), timeout=3)
assert result["ok"] is False
rows = result["sessions"]
assert isinstance(rows, list) and len(rows) == 2
assert all(row["error"] == "admission: InvalidHandshake" for row in rows)
assert "credentials-in-remote-handshake-detail" not in json.dumps(result)
@pytest.mark.asyncio
@pytest.mark.parametrize("sessions", [1, 2])
@pytest.mark.parametrize(
"mode",
[
"missing_packet_rate",
"string_packet_rate",
"invalid_base64",
"orphan_audio",
"odd_packet",
],
)
async def test_load_rejects_invalid_individual_audio_packets(tmp_path, sessions, mode):
# The concatenated PCM can still meet coverage/voicing thresholds, so the
# validator must also check each packet before accepting the whole run.
args = _args(tmp_path, "--sessions", str(sessions))
async with _server(mode) as (url, states):
args.url = url
result = await asyncio.wait_for(driver.run(args), timeout=4)
assert result["ok"] is False
assert result["passed_sessions"] == 0
rows = result["sessions"]
assert isinstance(rows, list) and len(rows) == sessions
assert all(row["error"] is not None for row in rows)
expected_code = {
"missing_packet_rate": "invalid_audio_sample_rate",
"string_packet_rate": "invalid_audio_sample_rate",
"invalid_base64": "invalid_audio_base64",
"orphan_audio": "missing_audio_response_id",
"odd_packet": "unaligned_pcm16_packet",
}[mode]
assert all({packet["code"] for packet in row["invalid_audio_packets"]} == {expected_code} for row in rows)
assert all(
sum(packet["samples"] for packet in row["audio_packet_timeline"]) == row["output_samples"] for row in rows
)
assert all(state["closed"] and state["close_requested"] for state in states)
assert json.loads((tmp_path / "out/load-result.json").read_text()) == result
def test_load_allows_empty_and_multi_frame_audio_packets():
client = driver.RawRealtimeProbe("ws://unused")
client.events.add(_audio(samples=0), received_at_s=1.0)
client.events.add(_audio(samples=driver.FRAME_SAMPLES * 5), received_at_s=1.4)
metrics = driver._load_metrics(client, [(1.0, 1.0, 1.001)])
assert metrics["invalid_audio_packets"] == []
assert metrics["output_samples"] == driver.FRAME_SAMPLES * 5
packets = metrics["audio_packet_timeline"]
assert isinstance(packets, list) and len(packets) == 1
def test_load_packet_diagnostics_do_not_copy_remote_payload():
client = driver.RawRealtimeProbe("ws://unused")
event = _audio()
event["delta"] = "private-remote-payload!"
client.events.add(event, received_at_s=1.4)
metrics = driver._load_metrics(client, [])
assert metrics["output_samples"] == 0
assert metrics["invalid_audio_packets"] == [
{
"event_index": 0,
"received_at_s": 1.4,
"code": "invalid_audio_base64",
}
]
assert "private-remote-payload" not in json.dumps(metrics)
tests/engine/test_async_omni_engine_abort_ack.py
+3 -2
View File
@@ -66,9 +66,10 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture):
assert msg.rpc_id == "abort-rpc-1"
assert not task.done()
rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
partial_output = SimpleNamespace(request_id="req-a", finished=True)
rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True, abort_outputs=[partial_output]))
result = await task
assert result == []
assert result == [partial_output]
try:
asyncio.run(_run())
tests/engine/test_async_omni_engine_stage_init.py
+590 -2
View File
@@ -8,25 +8,66 @@ import json
import os
import time
import types
from dataclasses import dataclass, field
import pytest
from omegaconf import OmegaConf
from vllm.v1.engine.utils import EngineZmqAddresses
from vllm_omni.config.omni_config import OmniStageRuntimeConfig
from vllm_omni.diffusion.data import AttentionConfig
from vllm_omni.engine import omni_engine_base as async_omni_engine_module
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
from vllm_omni.engine.stage_init_utils import (
LogicalStageInitPlan,
ReplicaInitPlan,
build_stage0_input_processor,
compute_replica_layout,
split_devices_for_replicas,
stage_runtime_env,
)
from vllm_omni.engine.stage_runtime import StageRuntime
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@dataclass
class _FakeParallelConfig:
enable_fault_tolerance: bool = False
enable_elastic_ep: bool = False
data_parallel_size: int = 1
use_ray: bool = False
@dataclass
class _FakeVllmConfig:
parallel_config: _FakeParallelConfig = field(default_factory=_FakeParallelConfig)
@dataclass
class _FakeRuntimeConfig:
devices: str
@dataclass
class _FakeStageConfig:
stage_id: int
stage_type: str
engine_args: dict[str, object]
runtime: _FakeRuntimeConfig
def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch):
env_key = "VLLM_OMNI_TEST_TYPED_STAGE_ENV"
monkeypatch.delenv(env_key, raising=False)
with stage_runtime_env(0, OmniStageRuntimeConfig(env={env_key: "typed-value"})):
assert os.environ[env_key] == "typed-value"
assert env_key not in os.environ
def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch):
engine = object.__new__(AsyncOmniEngine)
engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True)
@@ -188,7 +229,7 @@ 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):
def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):
import vllm_omni.engine.omni_engine_base as engine_mod
engine = object.__new__(AsyncOmniEngine)
@@ -206,6 +247,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
engine._omni_lb_policy = "random"
engine.request_queue = types.SimpleNamespace()
engine._log_stats = True
engine._client_config = engine_mod.OmniClientConfig(client_count=2, client_index=1, stage_addresses={})
engine._parallel_stage_init = False
captured: dict[str, object] = {}
@@ -221,6 +263,51 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
assert captured["stage_init_timeout"] == 7
assert captured["log_stats"] is True
assert captured["client_config"] is engine._client_config
def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch):
import vllm_omni.engine.omni_engine_base 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.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 = False
engine._parallel_stage_init = False
prompt_transform = object()
client = types.SimpleNamespace(
prompt_transform_func=prompt_transform,
prompt_expand_func=None,
default_sampling_params=types.SimpleNamespace(),
final_output=True,
final_output_type="text",
stage_type="llm",
model_stage="text_encoder",
is_comprehension=True,
)
pool = types.SimpleNamespace(
stage_client=client,
stage_vllm_config=None,
output_processor=None,
)
runtime = types.SimpleNamespace(stage_pools=[pool], initialize=lambda: None)
monkeypatch.setattr(engine_mod, "create_stage_runtime", lambda **_kwargs: runtime)
engine._initialize_stages(stage_init_timeout=7)
assert engine.prompt_transform_func is prompt_transform
def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
@@ -777,6 +864,289 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
assert captured["client_log_stats"] is True
def test_stage_runtime_attaches_external_llm_client_without_launching_engine(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
client_addresses = {
"input_address": "ipc://stage0-input-1",
"output_address": "ipc://stage0-output-1",
}
client_config = {
"client_count": 2,
"client_index": 1,
"stage_addresses": {0: {0: client_addresses}},
}
runtime = StageRuntime(
stage_configs=[],
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
async_chunk=False,
client_config=client_config,
)
parallel_config = _FakeParallelConfig()
vllm_config = _FakeVllmConfig(parallel_config=parallel_config)
plan = _make_llm_plan(0, stage_id=0, vllm_config=vllm_config).replicas[0]
captured: dict[str, object] = {}
stage_client = object()
def _capture_client(**kwargs):
captured.update(kwargs)
return stage_client
monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", _capture_client)
monkeypatch.setattr(
runtime_mod,
"launch_stage_replica",
lambda **_kwargs: pytest.fail("external client must not launch a stage engine"),
)
assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
assert captured["client_addresses"] == client_addresses
assert captured["client_count"] == 2
assert captured["client_index"] == 1
assert not hasattr(parallel_config, "_api_process_count")
assert not hasattr(parallel_config, "_api_process_rank")
@pytest.mark.parametrize("stage_ids", [(0,), (0, 1)], ids=["single-stage", "multi-stage"])
@pytest.mark.parametrize("client_count", [1, 2])
def test_stage_runtime_launches_shared_engines_with_per_client_addresses(monkeypatch, stage_ids, client_count):
import vllm_omni.engine.stage_runtime as runtime_mod
runtime = _make_stage_runtime()
stage_plans = []
for stage_id in stage_ids:
parallel_config = _FakeParallelConfig()
plan = _make_llm_plan(
stage_id,
stage_id=stage_id,
vllm_config=_FakeVllmConfig(parallel_config=parallel_config),
)
plan.replicas[0].engine_args_dict = {}
stage_plans.append(plan)
stage_plans[0].replicas[0].metadata.runtime_cfg = {"env": {"VLLM_OMNI_TEST_STAGE_RUNTIME_ENV": "enabled"}}
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda stage_id, _cfg: str(stage_id))
monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
events: list[tuple[str, int]] = []
captured_launch_kwargs: list[dict[str, object]] = []
captured_launch_env: list[str | None] = []
@contextlib.contextmanager
def _fake_launch_stage_replica(**kwargs):
stage_id = kwargs["stage_id"]
events.append(("enter", stage_id))
captured_launch_kwargs.append(kwargs)
captured_launch_env.append(os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV"))
yield StageReplicaResources(
addresses=EngineZmqAddresses(
inputs=[f"ipc://stage{stage_id}-input-{idx}" for idx in range(client_count)],
outputs=[f"ipc://stage{stage_id}-output-{idx}" for idx in range(client_count)],
),
)
events.append(("exit", stage_id))
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
with runtime.launch_stage_engines(client_count) as launch:
expected = (
[("enter", stage_id) for stage_id in stage_ids]
if client_count > 1
else [(event, stage_id) for stage_id in stage_ids for event in ("enter", "exit")]
)
assert events == expected
for client_index, config in enumerate(launch.client_configs):
assert config["client_count"] == client_count
assert config["client_index"] == client_index
for stage_id in stage_ids:
assert config["stage_addresses"][stage_id][0] == {
"input_address": f"ipc://stage{stage_id}-input-{client_index}",
"output_address": f"ipc://stage{stage_id}-output-{client_index}",
}
assert events == (
[(event, stage_id) for event in ("enter", "exit") for stage_id in stage_ids] if client_count > 1 else expected
)
assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_count")
assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_rank")
assert all(
kwargs["watched_frontend_processes"] is (launch.watched_frontend_processes if client_count > 1 else None)
for kwargs in captured_launch_kwargs
)
assert captured_launch_env == ["enabled" if stage_id == 0 else None for stage_id in stage_ids]
assert os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV") is None
@pytest.mark.parametrize("client_count", [1, 2])
def test_stage_runtime_overlapping_devices_acquire_real_locks_once(monkeypatch, tmp_path, client_count):
import fcntl
import vllm_omni.engine.stage_init_utils as init_utils
import vllm_omni.engine.stage_runtime as runtime_mod
runtime = _make_stage_runtime()
plans = [_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig()) for stage_id in (0, 1)]
for stage_id, plan in enumerate(plans):
plan.replicas[0].engine_args_dict = {"tensor_parallel_size": stage_id + 1}
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: plans)
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda sid, _cfg: "0" if sid == 0 else "0,1")
monkeypatch.setattr(
init_utils, "device_init_lock_path", lambda device_id: str(tmp_path / f"device-{device_id}.lock")
)
monkeypatch.setattr(
init_utils.time, "sleep", lambda _: pytest.fail("Overlapping stages must not wait on their own lock")
)
@contextlib.contextmanager
def launch(**kwargs):
# Real flock on a separate descriptor proves the runtime holds each lock.
for device_id in range(kwargs["stage_id"] + 1):
with open(tmp_path / f"device-{device_id}.lock", "a") as probe:
with pytest.raises(BlockingIOError):
fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
yield StageReplicaResources(
addresses=EngineZmqAddresses(
inputs=[f"ipc://input-{i}" for i in range(client_count)],
outputs=[f"ipc://output-{i}" for i in range(client_count)],
)
)
monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
with runtime.launch_stage_engines(client_count):
pass
for device_id in (0, 1):
with open(tmp_path / f"device-{device_id}.lock", "a") as probe:
fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
fcntl.flock(probe, fcntl.LOCK_UN)
@pytest.mark.parametrize("diffusion_stage_id", [0, 1], ids=["diffusion-only", "enginecore-to-diffusion"])
def test_stage_runtime_multi_api_rejects_diffusion_before_launch(monkeypatch, diffusion_stage_id):
import vllm_omni.engine.stage_runtime as runtime_mod
runtime = _make_stage_runtime()
stage_plans = [
_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig())
for stage_id in range(diffusion_stage_id)
]
stage_plans.append(_make_diffusion_plan(diffusion_stage_id, stage_id=diffusion_stage_id))
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
monkeypatch.setattr(
runtime_mod,
"acquire_device_locks",
lambda *_args: pytest.fail("unsupported topology must fail before acquiring devices"),
)
monkeypatch.setattr(
runtime_mod,
"launch_stage_replica",
lambda **_kwargs: pytest.fail("unsupported topology must not partially launch EngineCore stages"),
)
with pytest.raises(ValueError, match="diffusion stage\\(s\\) are not supported"):
with runtime.launch_stage_engines(2):
pytest.fail("unsupported topology must not yield a launch handle")
def test_stage_runtime_multi_api_maps_stage_devices_from_launcher_visibility(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
runtime = _make_stage_runtime()
stage_plan = _make_llm_plan(
0,
stage_id=0,
vllm_config=_FakeVllmConfig(),
)
replica = stage_plan.replicas[0]
replica.engine_args_dict = {}
replica.metadata.runtime_cfg = {"devices": "0"}
device_env = runtime_mod.current_omni_platform.device_control_env_var
monkeypatch.setenv(device_env, "5")
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
captured_devices: list[str | None] = []
@contextlib.contextmanager
def _fake_launch_stage_replica(**kwargs):
captured_devices.append(kwargs["stage_visible_devices"])
yield StageReplicaResources(
addresses=EngineZmqAddresses(
inputs=["ipc://input-0", "ipc://input-1"],
outputs=["ipc://output-0", "ipc://output-1"],
),
)
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
with runtime.launch_stage_engines(2):
pass
assert captured_devices == ["5"]
assert os.environ[device_env] == "5"
def test_stage_runtime_multi_api_rejects_elastic_ep(monkeypatch):
runtime = _make_stage_runtime()
stage_plan = _make_llm_plan(
0,
stage_id=0,
vllm_config=_FakeVllmConfig(parallel_config=_FakeParallelConfig(enable_elastic_ep=True)),
)
stage_plan.replicas[0].engine_args_dict = {}
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
with pytest.raises(ValueError, match="enable-elastic-ep"):
with runtime.launch_stage_engines(2):
pass
def test_stage_runtime_multi_api_failure_shuts_down_before_exceptional_context_exit(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
runtime = _make_stage_runtime()
stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig())
stage_plan.replicas[0].engine_args_dict = {}
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_args: None)
monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
events: list[str] = []
class _Manager:
def shutdown(self) -> None:
events.append("shutdown")
@contextlib.contextmanager
def _fake_launch_stage_replica(**_kwargs):
try:
yield StageReplicaResources(
manager=_Manager(),
addresses=EngineZmqAddresses(
inputs=["tcp://127.0.0.1:0", "tcp://127.0.0.1:1"],
outputs=["tcp://127.0.0.1:2", "tcp://127.0.0.1:3"],
),
)
except RuntimeError as exc:
events.append(f"exceptional-exit:{exc}")
raise
else:
events.append("normal-exit")
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
with pytest.raises(RuntimeError, match="deferred TCP addresses"):
with runtime.launch_stage_engines(2):
pass
assert events == [
"shutdown",
"exceptional-exit:Stage 0 returned deferred TCP addresses; multi-API launch requires fixed ports or IPC addresses",
]
def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
@@ -1023,7 +1393,7 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
num_replicas=1,
launch_mode="local",
stage_cfg=types.SimpleNamespace(engine_args={}, runtime=types.SimpleNamespace(devices="0")),
metadata=types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}),
metadata=types.SimpleNamespace(stage_id=0, stage_type="llm", runtime_cfg={"devices": "0"}),
stage_connector_spec={},
omni_kv_connector=(None, None, None),
stage_vllm_config=fake_vllm_config,
@@ -1715,6 +2085,51 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology():
assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
def test_inject_kv_stage_info_updates_typed_connector_config():
from vllm_omni.config.omni_config import (
OmniStageConnectorConfig,
OmniStageDiffusionParallelConfig,
VllmOmniDiffusionStageConfig,
)
from vllm_omni.config.stage_config import StageExecutionType, StagePipelineConfig
from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
from vllm_omni.entrypoints.utils import inject_omni_kv_config
stage0 = VllmOmniDiffusionStageConfig(
stage_pipeline_config=StagePipelineConfig(
stage_id=0,
model_stage="diffusion",
execution_type=StageExecutionType.DIFFUSION,
),
connector_config=OmniStageConnectorConfig(
omni_kv_config={
"need_send_cache": True,
"omni_from_stage": "0",
"omni_to_stage": "1",
}
),
parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=4),
)
stage1 = VllmOmniDiffusionStageConfig(
stage_pipeline_config=StagePipelineConfig(
stage_id=1,
model_stage="diffusion",
execution_type=StageExecutionType.DIFFUSION,
input_sources=(0,),
),
connector_config=OmniStageConnectorConfig(omni_kv_config={"need_recv_cache": True}),
parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=2),
)
inject_omni_kv_config(stage0, {"kv_connector": "P2pNcclConnector"}, "0", "1")
inject_kv_stage_info(stage0, 0, [stage0, stage1])
assert stage0.connector_config.omni_kv_config["stage_id"] == 0
assert stage0.connector_config.omni_kv_config["connector_config"] == {"kv_connector": "P2pNcclConnector"}
assert stage0.connector_config.omni_kv_config["engine_input_source"] == []
assert stage0.connector_config.omni_kv_config["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
"""ROCm default attention logic only applies to LLM stages, not diffusion."""
from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata
@@ -1832,6 +2247,102 @@ def test_port_from_zmq_address_parsing():
assert _port_from_zmq_address("tcp://host:not-a-port") is None
@pytest.mark.parametrize("parallel", [False, True])
@pytest.mark.parametrize("failure", ["ready", "attach", None])
def test_single_api_common_launch_ownership_and_rollback(monkeypatch, parallel, failure):
import vllm_omni.engine.stage_runtime as runtime_mod
runtime = _make_stage_runtime()
runtime._parallel_stage_init = parallel
plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig()).replicas[0]
plan.engine_args_dict = {}
events = []
class Manager:
def shutdown(self):
events.append("shutdown")
manager = Manager()
addresses = EngineZmqAddresses(inputs=["ipc://input"], outputs=["ipc://output"])
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
def acquire(*_):
events.append("lock")
return [42]
monkeypatch.setattr(runtime_mod, "acquire_device_locks", acquire)
monkeypatch.setattr(runtime_mod, "release_device_locks", lambda _: events.append("unlock"))
@contextlib.contextmanager
def launch(**kwargs):
assert kwargs["num_api_servers"] == 1
assert kwargs["omni_parallel_stage_init"] is parallel
events.append("spawn")
yield StageReplicaResources(manager=manager, addresses=addresses)
assert runtime._replica_launch_lock.locked() is False
events.append("ready")
if failure == "ready":
raise RuntimeError("ready")
def attach(**kwargs):
assert "ready" in events
assert kwargs["engine_manager"] is manager
events.append("attach")
if failure == "attach":
raise RuntimeError("attach")
return manager
monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", attach)
if failure:
with pytest.raises(RuntimeError, match=failure):
runtime._initialize_local_llm_replica(plan, 1)
assert events.count("shutdown") == 1
else:
assert runtime._initialize_local_llm_replica(plan, 1) is manager
assert "shutdown" not in events
assert events.count("lock") == events.count("unlock") == (0 if parallel else 1)
@pytest.mark.parametrize("client_count", [1, 2])
def test_common_launch_parallel_admission_before_spawn(monkeypatch, client_count):
import vllm_omni.engine.stage_runtime as runtime_mod
runtime = _make_stage_runtime()
runtime._parallel_stage_init = True
plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig())
plan.replicas[0].engine_args_dict = {}
events = []
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [plan])
monkeypatch.setattr(runtime, "_reject_unguardable_executors", lambda _: events.append("guard"))
monkeypatch.setattr(runtime, "_run_stage_admission", lambda _: events.append("admit"))
monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
def forbidden_lock(*_):
pytest.fail("parallel initialization must use child phase locks")
monkeypatch.setattr(runtime_mod, "acquire_device_locks", forbidden_lock)
@contextlib.contextmanager
def launch(**kwargs):
assert events == ["guard", "admit"]
assert kwargs["omni_parallel_stage_init"] is True
events.append("spawn")
yield StageReplicaResources(
addresses=EngineZmqAddresses(
inputs=[f"ipc://in{i}" for i in range(client_count)],
outputs=[f"ipc://out{i}" for i in range(client_count)],
)
)
assert not runtime._replica_launch_lock.locked()
events.append("ready")
monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
with runtime.launch_stage_engines(client_count):
pass
assert events == ["guard", "admit", "spawn", "ready"]
@pytest.mark.parametrize(
"stage_modes,expected", [([False, True, True], True), ([False, False], False), ([True, False], True)]
)
@@ -1852,3 +2363,80 @@ def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes,
with pytest.raises(ConfigResolvedError):
engine.__init__("test-model")
assert engine.async_chunk is expected
@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
def test_dist_stage_runtime_applies_local_dp_to_stage_config(typed):
from vllm_omni.config.omni_config import VllmOmniARStageConfig
from vllm_omni.config.stage_config import StagePipelineConfig
from vllm_omni.engine.stage_runtime import DistStageRuntime
if typed:
stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="ar"))
runtime_cfg = stage.runtime_config
else:
runtime_cfg = types.SimpleNamespace(num_replicas=1)
stage = types.SimpleNamespace(stage_id=0, runtime=runtime_cfg)
runtime = DistStageRuntime(
stage_configs=[stage],
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
async_chunk=False,
single_stage_id_filter=0,
omni_master_address="127.0.0.1",
omni_master_port=12345,
omni_dp_size_local=2,
)
runtime._validate_single_stage_mode_replica_constraints()
assert runtime_cfg.num_replicas == 2
@pytest.mark.parametrize("num_replicas", [1, 2, 3])
def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(mocker, num_replicas):
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.engine import stage_runtime as runtime_module
stage = StageConfigFactory.create_typed_default_diffusion(
"generic-diffusion", {"model_class_name": "QwenImagePipeline"}
).stage_configs[0]
devices = ",".join(str(i) for i in range(num_replicas))
stage.runtime_config.devices = devices
stage.runtime_config.num_replicas = num_replicas
runtime = StageRuntime(
stage_configs=[stage],
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
async_chunk=False,
)
mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
mocker.patch.object(runtime, "_stage_device_scope", side_effect=lambda *_: contextlib.nullcontext())
client = mocker.Mock()
launch = mocker.patch.object(runtime_module, "launch_diffusion_stage_replica", return_value=(client, None))
# Replanning must not inherit the first replica's narrowed device slice.
for _ in range(2):
counts, device_map = compute_replica_layout([stage])
plans = runtime._build_logical_stage_init_plans(
omni_transfer_config=None,
replicas_per_stage=counts,
replica_devices_map=device_map,
)
replicas = plans[0].replicas
assert stage.runtime_config.devices == devices
assert len({id(plan.stage_cfg) for plan in replicas}) == num_replicas
for i, plan in enumerate(replicas):
assert plan.metadata.runtime_cfg is plan.stage_cfg.runtime_config
assert plan.stage_cfg.runtime_config.devices == str(i)
if num_replicas > 1:
assert plan.stage_cfg is not stage
assert runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) is client
assert launch.call_args.kwargs["stage_config"] is plan.stage_cfg
assert launch.call_args.kwargs["metadata"] is plan.metadata
assert launch.call_args.kwargs["use_inline"] is (num_replicas == 1)
assert launch.call_count == 2 * num_replicas
tests/engine/test_parallel_stage_init.py
+17
View File
@@ -306,6 +306,23 @@ def test_check_admission_admits_diffusion_with_explicit_util():
assert ledgers[0].kv_budget_bytes == int(cap * 0.3)
def test_check_admission_reads_typed_diffusion_cache_config():
cap = 80 * 1024**3
diff = _local_diffusion_replica(0, 0, "0", util=None)
diff.stage_cfg.cache_config = types.SimpleNamespace(gpu_memory_utilization=0.35)
plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[diff])
ledgers = check_admission(
[plan],
resolve_physical_devices=lambda _r: [0],
device_total_memory=lambda _d: cap,
external_reserve_bytes=0,
safety_margin_bytes=0,
)
assert ledgers[0].kv_budget_bytes == int(cap * 0.35)
def test_check_admission_unresolved_local_raises():
"""A local replica invisible to the ledger must fail admission, not skip it:
it would otherwise initialize concurrently with unbounded demand."""
tests/engine/test_single_stage_mode.py
+18 -3
View File
@@ -16,7 +16,8 @@ from pytest_mock import MockerFixture
from vllm.v1.engine.utils import EngineZmqAddresses
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.config.omni_config import VllmOmniARStageConfig, VllmOmniDiffusionStageConfig
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, StageExecutionType, StagePipelineConfig
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
from vllm_omni.engine.stage_engine_startup import (
@@ -57,6 +58,10 @@ def _make_llm_plan(
vllm_config: Any | None = None,
) -> LogicalStageInitPlan:
stage_cfg = _make_stage_cfg(stage_id)
if launch_mode == "remote":
stage_cfg = VllmOmniARStageConfig(
stage_pipeline_config=StagePipelineConfig(stage_id=stage_id, model_stage="thinker"),
)
metadata = SimpleNamespace(
stage_id=stage_id,
stage_type="llm",
@@ -98,6 +103,16 @@ def _make_diffusion_plan(
launch_mode: str,
) -> LogicalStageInitPlan:
stage_cfg = _make_stage_cfg(stage_id, stage_type="diffusion")
if launch_mode == "remote":
stage_cfg = VllmOmniDiffusionStageConfig(
stage_pipeline_config=StagePipelineConfig(
stage_id=stage_id,
model_stage="diffusion",
execution_type=StageExecutionType.DIFFUSION,
final_output=True,
final_output_type="image",
),
)
metadata = SimpleNamespace(
stage_id=stage_id,
stage_type="diffusion",
@@ -960,6 +975,7 @@ class TestSingleStageReplicaInitialization:
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 client_kwargs["metadata"].model_stage == "thinker"
assert events == ["enter", "exit", "attach"]
def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture):
@@ -1125,11 +1141,9 @@ class TestSingleStageReplicaInitialization:
return _ctx()
remote_metadata = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0].metadata
plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0]
sentinel_client = SimpleNamespace()
mocker.patch.object(runtime_mod, "extract_legacy_stage_metadata", return_value=remote_metadata)
mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect)
mock_from_addresses = mocker.patch(
"vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses",
@@ -1146,6 +1160,7 @@ class TestSingleStageReplicaInitialization:
replica_id=0,
)
mock_from_addresses.assert_called_once()
assert mock_from_addresses.call_args.args[0].final_output_type == "image"
def test_initialize_local_diffusion_replica_registers_with_master(self, mocker: MockerFixture):
import vllm_omni.engine.stage_runtime as runtime_mod
tests/engine/test_stage_device_layout.py
+63 -13
View File
@@ -18,6 +18,7 @@ import types
from unittest import mock
import pytest
from pytest_mock import MockerFixture
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.stage_init_utils import (
@@ -156,6 +157,43 @@ def test_replica_guard_and_splitter_share_local_world_size(devices, expected):
assert replica_devices_map == {0: expected}
@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
@pytest.mark.parametrize("local_dp", [None, 0, 1, 2])
@pytest.mark.parametrize("pool_devices", [False, True], ids=["template", "pool"])
def test_replica_layout_preserves_local_dp_semantics(typed, local_dp, pool_devices):
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig
from vllm_omni.engine.stage_engine_startup import get_headless_replica_devices
engine_args = {
"tensor_parallel_size": 2,
"data_parallel_size": 4,
"data_parallel_size_local": local_dp,
"pipeline_parallel_size": 2,
}
width = 2 * max(1, local_dp if local_dp is not None else 4) * 2
devices = ",".join(str(i) for i in range(width * (2 if pool_devices else 1)))
if typed:
pipeline = PipelineConfig(
model_type="test",
model_arch="TestModel",
stages=(StagePipelineConfig(stage_id=0, model_stage="ar", final_output=True),),
)
stage = VllmOmniConfig.from_pipeline_config(
pipeline,
cli_overrides={
f"stage_0_{key}": value for key, value in {**engine_args, "devices": devices, "num_replicas": 2}.items()
},
).stage_configs[0]
else:
stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args)
expected = [",".join(str(i) for i in range(r * width, (r + 1) * width)) for r in range(2)]
assert get_stage_devices_per_replica(stage) == width
assert compute_replica_layout([stage]) == ([2], {0: expected})
assert get_headless_replica_devices(stage, stage_id=0, omni_dp_size_local=2) == expected
def test_non_tp_layout_error_uses_generic_guidance():
"""A PP mismatch must not be reported as a top-level TP broadcast."""
stage = _stage(0, devices="0")
@@ -203,7 +241,7 @@ def test_build_vllm_config_fails_before_engine_config_on_mismatch():
get_class.assert_not_called()
def test_build_vllm_config_proceeds_on_consistent_layout():
def test_build_vllm_config_proceeds_on_consistent_layout(mocker: MockerFixture) -> None:
"""The guard must not false-positive: a consistent single-GPU stage flows past
it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual."""
stage = _stage(1, devices="1")
@@ -212,18 +250,30 @@ def test_build_vllm_config_proceeds_on_consistent_layout():
model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()),
)
sentinel_executor = object()
with (
mock.patch.object(
stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config
) as create_engine_config,
mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor),
mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant),
):
vllm_config, executor_class = build_vllm_config(
stage,
model="dummy-model",
engine_args_dict={"tensor_parallel_size": 1},
)
captured_api_process_config: dict[str, int] = {}
def capture_create_engine_config(engine_args, *args, **kwargs):
captured_api_process_config["count"] = engine_args._api_process_count
captured_api_process_config["rank"] = engine_args._api_process_rank
return fake_config
create_engine_config = mocker.patch.object(
stage_init_utils.OmniEngineArgs,
"create_engine_config",
autospec=True,
side_effect=capture_create_engine_config,
)
mocker.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor)
mocker.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant)
vllm_config, executor_class = build_vllm_config(
stage,
model="dummy-model",
engine_args_dict={"tensor_parallel_size": 1},
api_process_count=3,
api_process_rank=2,
)
create_engine_config.assert_called_once()
assert vllm_config is fake_config
assert executor_class is sentinel_executor
assert captured_api_process_config == {"count": 3, "rank": 2}
tests/engine/test_stage_engine_args.py
+16
View File
@@ -95,6 +95,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
"num_weight_load_threads",
"omni_kv_config",
"parallel_config",
"pooling_output_decoder",
"silence_ban_frames",
"subtalker_sampling_params",
"task_type",
@@ -366,6 +367,21 @@ def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults():
assert inherited_defaults.isdisjoint(engine_args)
def test_typed_llm_projection_omits_diffusion_only_and_process_only_defaults():
stage_config = VllmOmniARStageConfig(
stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
)
engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
assert {
"enable_multithread_weight_load",
"num_weight_load_threads",
"disable_autocast",
"log_level",
}.isdisjoint(engine_args)
def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch):
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
tests/entrypoints/openai_api/test_api_server_guards.py
+198
View File
@@ -42,7 +42,10 @@ from __future__ import annotations
import asyncio
import json
import signal
from argparse import Namespace
from contextlib import asynccontextmanager
from dataclasses import dataclass
from types import SimpleNamespace
import pytest
@@ -59,11 +62,78 @@ from vllm_omni.entrypoints.serve.utils import errors as serve_errors
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@dataclass
class _FakeParallelConfig:
_api_process_rank: int
@dataclass
class _FakeVllmConfig:
parallel_config: _FakeParallelConfig
def test_omni_api_worker_sets_parent_death_signal_before_validation(mocker):
args = _minimal_args(_omni_stage_client_configs=[{"stage_addresses": {}}])
death_signal = mocker.patch.object(api_server, "set_death_signal")
with pytest.raises(RuntimeError, match="client configuration"):
api_server.run_omni_api_server_worker_proc(
"127.0.0.1:8000",
mocker.Mock(),
args,
{"client_count": 2, "client_index": -1},
)
death_signal.assert_called_once_with(signal.SIGTERM)
@pytest.mark.asyncio
async def test_omni_run_server_worker_forwards_client_config(monkeypatch):
captured: dict[str, object] = {}
@asynccontextmanager
async def fake_build_async_omni(*_args, **kwargs):
captured.update(kwargs)
yield _FakeEngineClient(vllm_config=_FakeVllmConfig(_FakeParallelConfig(_api_process_rank=1)))
monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
monkeypatch.setattr(api_server, "build_openai_app", lambda *_args: FastAPI())
monkeypatch.setattr(api_server, "omni_init_app_state", lambda *_args: asyncio.sleep(0))
monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_args: None)
async def fake_storage_start() -> None:
return None
monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start)
async def fake_serve_http(app, **_kwargs):
assert app.state.api_server_count == 2
task = asyncio.create_task(asyncio.sleep(0))
await asyncio.sleep(0)
return task
monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
args = _minimal_args(model="dummy")
config = {"client_count": 2, "client_index": 1, "stage_addresses": {}}
await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), args, config)
assert captured["client_config"] is config
# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
# would invent ("HEAD", path) entries for every GET and fail the current census.
# TODO: extend this set and the expected census if @router.head / @router.options appear.
_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
def test_profiler_endpoints_detect_typed_stage_config():
stage = SimpleNamespace(profiler_config=SimpleNamespace(profiler="torch"))
assert api_server._should_enable_profiler_endpoints([stage])
# Exact Omni router census on main. Update intentionally when routes change.
_EXPECTED_ROUTER_ROUTES = {
("POST", "/v1/chat/completions"),
@@ -619,6 +689,48 @@ async def test_realtime_route_defaults_to_configured_duplex_handler(
assert calls == [expected_handler]
@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime", "/v1/realtime?duplex=1"])
def test_multi_api_duplex_reconnect_rejected_on_each_worker(path: str) -> None:
class SessionHandler:
async def handle_realtime_session(self, _websocket) -> None:
pytest.fail("Multi-API requests must not access process-local sessions")
class PendingWarmup:
def is_set(self) -> bool:
raise AssertionError("Multi-API rejection must not wait for duplex warmup")
# Separate apps represent workers with independent frontend registries.
# Opening on A then reconnecting to B must reject at the same boundary.
for worker_index in (0, 1):
app = FastAPI()
app.state.api_server_count = 2
app.state.api_server_index = worker_index
app.state.openai_serving_duplex = SessionHandler()
app.state.duplex_warmup_done = PendingWarmup()
app.include_router(api_server.router)
with TestClient(app) as client, client.websocket_connect(path) as websocket:
assert websocket.receive_json()["code"] == "multi_api_duplex_unsupported"
with pytest.raises(WebSocketDisconnect) as exc:
websocket.receive_text()
assert exc.value.code == 1008
@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime?duplex=1"])
def test_single_api_duplex_session_still_reaches_handler(path: str) -> None:
class SessionHandler:
async def handle_realtime_session(self, websocket) -> None:
await websocket.accept()
await websocket.send_json({"type": "session.created"})
await websocket.close()
app = FastAPI()
app.state.api_server_count = 1
app.state.openai_serving_duplex = SessionHandler()
app.include_router(api_server.router)
with TestClient(app) as client, client.websocket_connect(path) as websocket:
assert websocket.receive_json() == {"type": "session.created"}
def test_health_without_engine_returns_stable_unhealthy_response() -> None:
"""Lock ``/health`` with no engine initialized.
@@ -708,6 +820,38 @@ def test_speech_without_handler_preserves_not_found_http_error() -> None:
assert exc_info.value.detail == "The model does not support Speech API"
@pytest.mark.asyncio
async def test_multi_api_rejects_runtime_voice_upload() -> None:
app = FastAPI()
app.state.api_server_count = 2
raw_request = _request_for(app, method="POST", path="/v1/audio/voices")
with pytest.raises(HTTPException) as exc_info:
await api_server.upload_voice(
raw_request,
audio_sample=None,
speaker_embedding=None,
consent="consent-id",
name="probe",
)
assert exc_info.value.status_code == 409
assert "process-local frontend state" in exc_info.value.detail
@pytest.mark.asyncio
async def test_multi_api_rejects_sleep_route() -> None:
app = FastAPI()
app.state.api_server_count = 2
raw_request = _request_for(app, method="POST", path="/v1/omni/sleep")
with pytest.raises(HTTPException) as exc_info:
await api_server.omni_sleep(api_server.OmniSleepRequest(stage_ids=[0]), raw_request)
assert exc_info.value.status_code == 409
assert "process-local frontend state" in exc_info.value.detail
def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
"""Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
@@ -930,3 +1074,57 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
must_be_none=_MULTISTAGE_MUST_BE_NONE,
)
@pytest.mark.parametrize("count", [None, 0, -1, "2", True])
def test_process_local_guard_rejects_unknown_topology(count):
app = FastAPI()
if count is not None:
app.state.api_server_count = count
# CLI args cannot substitute for missing/invalid actual runtime state.
app.state.args = Namespace(api_server_count=1)
with pytest.raises(HTTPException) as exc:
api_server._reject_process_local_state_with_multiple_api_workers(_request_for(app), "Sleep")
assert exc.value.status_code == 503
@pytest.mark.parametrize(
"requested,config,expected",
[(None, None, 1), (1, None, 1), (2, {"client_count": 2}, 2), (None, {"client_count": 2}, 2)],
)
def test_api_topology_resolved_from_launch(requested, config, expected):
assert api_server._resolve_api_server_count(Namespace(api_server_count=requested), config) == expected
@pytest.mark.parametrize("requested,config", [(2, None), (2, {}), (2, {"client_count": 1}), (2, {"client_count": "2"})])
def test_api_topology_rejects_inconsistent_launch(requested, config):
with pytest.raises(ValueError):
api_server._resolve_api_server_count(Namespace(api_server_count=requested), config)
@pytest.mark.parametrize("count,status", [(2, 409), (None, 503)])
@pytest.mark.parametrize(
"method,path",
[
("POST", "/v1/videos"),
("GET", "/v1/videos"),
("GET", "/v1/videos/job"),
("GET", "/v1/videos/job/content"),
("DELETE", "/v1/videos/job"),
],
)
def test_video_lifecycle_guard_precedes_processing(monkeypatch, count, status, method, path):
from fastapi.testclient import TestClient
app = FastAPI()
if count is not None:
app.state.api_server_count = count
app.include_router(api_server.router)
async def forbidden_form():
pytest.fail("video form processing must not run before topology rejection")
app.dependency_overrides[api_server._parse_video_form] = forbidden_form
with TestClient(app) as client:
response = client.request(method, path)
assert response.status_code == status
tests/entrypoints/openai_api/test_audex_serving_guards.py
+45
View File
@@ -15,6 +15,7 @@ from types import SimpleNamespace
import pytest
import torch
from vllm_omni.config.stage_config import StagePipelineConfig
from vllm_omni.entrypoints.openai.serving_speech import (
_AUDEX_NO_AUDIO_GUARD_MODEL_TYPES,
OmniOpenAIServingSpeech,
@@ -251,9 +252,28 @@ class TestAudexStageDetection:
stage = serving._find_tts_stage()
assert stage is not None and stage.engine_args.model_stage == "audex_thinker"
def test_typed_pipeline_detected_without_legacy_engine_args(self):
serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
stages = [
SimpleNamespace(
stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage),
model_config=SimpleNamespace(model_arch=None),
worker_type="generation",
)
for index, model_stage in enumerate(("audex_thinker", "audex_code2wav"))
]
serving.engine_client = SimpleNamespace(stage_configs=stages)
stage = serving._find_tts_stage()
serving._tts_stage = stage
assert stage is stages[0]
assert serving._detect_tts_model_type() == "audex"
def test_guided_injection_sets_guided_temperature():
"""Guided requests decode at the measured CFG temperature (0.05)."""
serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
cond = TestAudexCfgInjection()._cond_prompt()
params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
@@ -265,6 +285,31 @@ def test_guided_injection_sets_guided_temperature():
assert out.temperature == 0.1
def test_typed_full_pipeline_uses_full_checkpoint_tokenizer(mocker):
from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
serving._tts_stage = SimpleNamespace(
stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="audex_omni"),
model_config=SimpleNamespace(model_arch=None),
worker_type="ar",
)
serving.engine_client = SimpleNamespace(
model_config=SimpleNamespace(model="/models/audex"),
)
ensure_snapshot = mocker.patch(
"vllm_omni.model_executor.models.audex.checkpoint.ensure_audex_snapshot",
return_value="/snapshots/audex",
)
from_pretrained = mocker.patch("transformers.AutoTokenizer.from_pretrained", return_value=object())
adapter = AudexAdapter(SpeechServingContext(server=serving, engine_client=serving.engine_client))
adapter._get_tokenizer()
ensure_snapshot.assert_called_once_with("/models/audex", profile="full")
from_pretrained.assert_called_once_with("/snapshots/audex/checkpoint_folder_full")
# ---------------------------------------------------------------- TTA serving surface
from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter # noqa: E402
tests/entrypoints/openai_api/test_omni_sleep_wakeup.py
+2
View File
@@ -20,6 +20,7 @@ def _make_app(engine_client):
from vllm_omni.entrypoints.openai.api_server import router
app = FastAPI()
app.state.api_server_count = 1
app.include_router(router)
app.state.engine_client = engine_client
app.state.sleeping_stages = set()
@@ -253,6 +254,7 @@ def pure_diffusion_app(pure_diffusion_engine, mocker):
)
app = FastAPI()
app.state.api_server_count = 1
app.include_router(router)
mocker.patch(
tests/entrypoints/openai_api/test_openpi_serving.py
+30
View File
@@ -138,6 +138,36 @@ def test_policy_server_config_reads_stage_config_model_config():
assert serving.policy_server_config.to_dict() == policy_config
@pytest.mark.parametrize("inline", [False, True], ids=["out_of_process", "inline"])
def test_policy_server_config_reads_typed_runtime(inline):
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE
from vllm_omni.entrypoints.async_omni import AsyncOmni
typed_stage = VllmOmniConfig.from_pipeline_config(
PI0_PIPELINE, deploy_config_path=get_deploy_config_path("pi0.yaml")
).stage_configs[0]
policy_config = typed_stage.diffusion_config.model_config["policy_server_config"]
stage_client = SimpleNamespace(stage_type="diffusion")
if inline:
# Worker-enriched inline metadata remains authoritative over the
# startup stage config; only the process client needs the fallback.
policy_config = {**policy_config, "worker_metadata": "loaded"}
stage_client.od_config = SimpleNamespace(model_config={"policy_server_config": policy_config})
engine_client = AsyncOmni.__new__(AsyncOmni)
engine_client.engine = SimpleNamespace(
stage_clients=[stage_client],
stage_configs=[typed_stage],
stage_vllm_configs=[None],
get_diffusion_od_config=lambda: SimpleNamespace(model_class_name="Pi0Pipeline"),
)
serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client)
assert serving.policy_server_config.to_dict() == policy_config
def test_policy_server_config_reads_omegaconf_stage_config():
engine_client = SimpleNamespace(
get_diffusion_od_config=lambda: None,
tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py
+32 -1
View File
@@ -27,7 +27,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@dataclass
class _EngineClientStub:
stage_configs: list[dict[str, object]]
stage_configs: list[object]
@dataclass
@@ -149,6 +149,37 @@ async def test_chat_reference_audio_reaches_code2wav_transport(
assert code2wav_payload.meta.ref_audio_sr == 16000
@pytest.mark.asyncio
async def test_typed_stage_enables_reference_audio(
serving_chat: OmniOpenAIServingChat,
mocker: MockerFixture,
) -> None:
reference_audio_url = "data:audio/wav;base64,AAAA"
reference_waveform = np.array([0.25, -0.5], dtype=np.float32)
serving_chat.engine_client = _EngineClientStub(
[
SimpleNamespace(
model_config=SimpleNamespace(model_arch="MiniCPMO45OmniForConditionalGeneration"),
model_stage="llm",
)
]
)
serving_chat.model_config = _ModelConfigStub()
engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
media_connector = mocker.patch(
"vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
).return_value
media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000))
await serving_chat._attach_minicpmo45_reference_audio(
engine_prompt,
_chat_request(reference_audio_url),
)
media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url)
assert MINICPMO45_REFERENCE_AUDIO_KEY in cast(dict[str, object], engine_prompt)
@pytest.mark.parametrize("request_shape", ["extra_body", "model_extra"])
@pytest.mark.asyncio
async def test_reference_audio_accepts_request_extra_mappings(
tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py
+22
View File
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for MiniCPM-o 4.5 audio-streaming output-kind coercion.
Chat ``stream=true`` coerces every AR stage to DELTA. For MiniCPM-o 4.5 that
@@ -33,6 +34,13 @@ def _stage(model_arch: str, model_stage: str) -> SimpleNamespace:
return SimpleNamespace(engine_args=SimpleNamespace(model_arch=model_arch, model_stage=model_stage))
def _typed_stage(model_arch: str, model_stage: str) -> SimpleNamespace:
return SimpleNamespace(
model_config=SimpleNamespace(model_arch=model_arch),
model_stage=model_stage,
)
def _minicpmo45_stages():
return [_stage(MINICPMO45_ARCH, "llm"), _stage(MINICPMO45_ARCH, "tts")]
@@ -50,6 +58,20 @@ def test_thinker_forced_final_only_talker_stays_delta(serving_chat):
assert out[1].output_kind is RequestOutputKind.DELTA
def test_typed_stages_force_thinker_final_only(serving_chat):
serving_chat.engine_client = SimpleNamespace(
stage_configs=[
_typed_stage(MINICPMO45_ARCH, "llm"),
_typed_stage(MINICPMO45_ARCH, "tts"),
]
)
out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(), ["text", "audio"])
assert out[0].output_kind is RequestOutputKind.FINAL_ONLY
assert out[1].output_kind is RequestOutputKind.DELTA
@pytest.mark.parametrize("modalities", [None, [], ["text"]])
def test_non_audio_request_is_untouched(serving_chat, modalities):
serving_chat.engine_client = SimpleNamespace(stage_configs=_minicpmo45_stages())
tests/entrypoints/openai_api/test_serving_speech.py
+34
View File
@@ -27,6 +27,7 @@ from pydantic import ValidationError
from pytest_mock import MockerFixture
from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse
from vllm_omni.config.stage_config import StagePipelineConfig
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
@@ -1039,6 +1040,38 @@ class TestTTSMethods:
asyncio.run(server._prepare_speech_generation(request))
server.shutdown()
def test_prepare_speech_rejects_typed_non_tts_omni_model(self, mocker: MockerFixture):
"""Typed Qwen3-Omni stages must hit the same talker guard as legacy stages."""
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.tts_max_instructions_length = None
mock_engine_client.stage_configs = [
SimpleNamespace(
stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage),
model_config=SimpleNamespace(model_arch=None),
worker_type=worker_type,
)
for index, (model_stage, worker_type) in enumerate(
(
("thinker", "ar"),
("talker", "ar"),
("code2wav", "generation"),
)
)
]
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=mocker.MagicMock(),
)
request = OpenAICreateSpeechRequest(input="Hello world")
with pytest.raises(ValueError, match="only supported for dedicated TTS models"):
asyncio.run(server._prepare_speech_generation(request))
server.shutdown()
def test_estimate_prompt_len_fallback(self, speech_server):
"""Test prompt length estimation falls back to 2048 when model is unavailable."""
tts_params = {"text": ["Hello"], "task_type": ["CustomVoice"]}
@@ -3793,6 +3826,7 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt
def _make_api_server_request(handler, *, method: str = "POST", path: str = "/v1/audio/voices") -> Request:
app = FastAPI()
app.state.openai_serving_speech = handler
app.state.api_server_count = 1
scope = {
"type": "http",
"app": app,
tests/entrypoints/openai_api/test_tts_adapter.py
+26
View File
@@ -22,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import (
detect_tts_model_type,
resolve_adapter,
)
from vllm_omni.entrypoints.openai.tts_adapters.base import resolve_stage_model_path
from vllm_omni.entrypoints.openai.tts_adapters.covo_audio import CovoAudioAdapter
from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter
from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
@@ -99,6 +100,31 @@ def test_resolve_unknown_returns_none():
assert resolve_adapter(None) is None
def test_stage_model_path_prefers_typed_override():
engine_client = SimpleNamespace(
stage_configs=[
SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model")),
SimpleNamespace(
model_config=SimpleNamespace(model="typed-stage-model"),
),
],
model="served-model",
)
assert resolve_stage_model_path(engine_client) == "typed-stage-model"
def test_stage_model_path_falls_back_to_legacy_then_served_model():
legacy_client = SimpleNamespace(
stage_configs=[SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model"))],
model="served-model",
)
served_client = SimpleNamespace(stage_configs=[SimpleNamespace()], model="served-model")
assert resolve_stage_model_path(legacy_client) == "legacy-stage-model"
assert resolve_stage_model_path(served_client) == "served-model"
def test_voxcpm2_resolves():
"""VoxCPM2 (the served ``latent_generator`` model) resolves cleanly.
tests/entrypoints/openai_api/test_video_server.py
+58
View File
@@ -463,6 +463,7 @@ async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch)
@pytest.fixture
def test_client():
app = FastAPI()
app.state.api_server_count = 1
app.include_router(router)
app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion(
diffusion_engine=FakeAsyncOmni(),
@@ -1134,6 +1135,32 @@ def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_fal
assert handler.supports_mixed_reference_inputs
def test_typed_stage_drives_video_capability_checks():
from vllm_omni.config.config_factory import StageConfigFactory
minimax_stage = StageConfigFactory.create_typed_default_diffusion(
"minimax-h3",
{"model_class_name": "MiniMaxH3Pipeline"},
).stage_configs[0]
minimax_handler = OmniOpenAIServingVideo.for_diffusion(
SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)),
model_name="minimax-h3",
stage_configs=[minimax_stage],
)
assert minimax_handler.supports_mixed_reference_inputs
cosmos_stage = StageConfigFactory.create_typed_default_diffusion(
"cosmos3",
{"model_class_name": "Cosmos3OmniDiffusersPipeline"},
).stage_configs[0]
cosmos_handler = OmniOpenAIServingVideo.for_diffusion(
SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)),
model_name="cosmos3",
stage_configs=[cosmos_stage],
)
assert cosmos_handler.supported_control_upload_types == frozenset({"edge", "blur", "depth", "seg", "wsm"})
@pytest.mark.parametrize("model_class_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
def test_control_upload_capability_is_declared_only_by_cosmos3(test_client, model_class_name):
handler = test_client.app.state.openai_serving_video
@@ -1191,6 +1218,35 @@ def test_cosmos3_reference_video_limit_uses_v2v_condition_frames():
assert spec.keep == "first"
@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
@pytest.mark.parametrize(
("num_frames", "extra_params", "expected"),
[
(189, {"condition_frame_indexes_vision": [0, 2], "condition_video_keep": "last"}, (9, "last")),
(189, {"condition_frame_indexes_vision": [0, 2]}, (9, "first")),
(5, {"condition_frame_indexes_vision": [0, 20]}, (5, "first")),
(None, {"action_mode": "inverse_dynamics", "action_chunk_size": 16}, (17, "first")),
],
)
def test_cosmos3_reference_video_decode_policy_with_runtime_configs(typed, num_frames, extra_params, expected):
from vllm_omni.config.config_factory import StageConfigFactory
if typed:
stages = list(
StageConfigFactory.create_typed_default_diffusion(
"cosmos3",
{"model_class_name": "Cosmos3OmniDiffusersPipeline"},
).stage_configs
)
else:
stages = _cosmos3_stage_configs()
request = VideoGenerationRequest(prompt="Continue this motion.", num_frames=num_frames, extra_params=extra_params)
spec = _reference_video_decode_spec(request, stages)
assert (spec.max_frames, spec.keep) == expected
def test_cosmos3_reference_video_limit_preserves_action_frames():
request = VideoGenerationRequest(
prompt="Predict the action.",
@@ -1749,6 +1805,7 @@ def test_action_extraction_accepts_multimodal_actions_payload():
def test_missing_handler_returns_503():
app = FastAPI()
app.state.api_server_count = 1
app.include_router(router)
app.state.openai_serving_video = None
client = TestClient(app)
@@ -2735,6 +2792,7 @@ def test_cosmos3_control_upload_rejects_invalid_size(control_bytes, message, tes
def test_sync_missing_handler_returns_503():
app = FastAPI()
app.state.api_server_count = 1
app.include_router(router)
app.state.openai_serving_video = None
client = TestClient(app)
tests/entrypoints/test_async_omni_diffusion_config.py
+134 -2
View File
@@ -2,13 +2,15 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
from typing import Any
import pytest
import torch
from pydantic import ValidationError
from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config.resolver import OmniConfigResolution
from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -117,7 +119,7 @@ def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocke
trust_remote_code=False,
)
assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
assert stage_configs[0].diffusion_config.extras["ltx2_use_conv_vae"] is True
def test_default_stage_rejects_unknown_nested_parallel_config_key():
@@ -239,6 +241,33 @@ def test_default_stage_config_includes_default_sampling_params():
}
@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
@pytest.mark.parametrize("sampling_defaults", [{"0": {"guidance_scale": 7.5}}, '{"0":{"guidance_scale":7.5}}'])
def test_generic_diffusion_sampling_defaults_remain_overridable(typed, sampling_defaults):
from vllm_omni.config.yaml_util import create_config
from vllm_omni.entrypoints.omni_base import OmniBase
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
kwargs = {"model_class_name": "QwenImagePipeline", "default_sampling_params": sampling_defaults}
if typed:
stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0]
metadata = stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage)
else:
stage = create_config(StageConfigFactory.create_default_diffusion(kwargs))[0]
metadata = stage_init_utils.extract_legacy_stage_metadata(stage)
base = OmniBase.__new__(OmniBase)
base.engine = SimpleNamespace(num_stages=1, stage_configs=[stage])
base.default_sampling_params_list = [metadata.default_sampling_params]
base.sampling_constraints_list = base._get_sampling_constraints_list([stage])
assert base.resolve_sampling_params_list(None)[0].guidance_scale == 7.5
requested = OmniDiffusionSamplingParams(guidance_scale=2.0)
assert base.resolve_sampling_params_list(requested)[0].guidance_scale == 2.0
assert requested.guidance_scale == 2.0
assert base.default_sampling_params_list[0].guidance_scale == 7.5
assert base.sampling_constraints_list == [{}]
def test_default_stage_config_includes_diffusion_attention_backend():
"""Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config."""
stage_cfg = StageConfigFactory.create_default_diffusion(
@@ -726,3 +755,106 @@ def test_default_stage_config_includes_quantization_config():
stage_cfg = StageConfigFactory.create_default_diffusion({"quantization_config": quantization_config})[0]
assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
@pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HunyuanVideo15Pipeline"])
def test_generic_diffusion_uses_canonical_video_output_type(model_class_name):
config = StageConfigFactory.create_typed_default_diffusion(
"generic-video",
{"model_class_name": model_class_name},
)
assert config.stage_configs[0].final_output_type == "video"
def test_generic_diffusion_resolves_structured_stage_without_legacy_conversion(mocker):
"""Generic diffusion reaches runtime as the structured stage itself."""
mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
mocker.patch(
"vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
return_value=(True, "FakeDiffusionPipeline"),
)
resolved = resolve_omni_config(
"generic-diffusion",
trust_remote_code=False,
deploy_config_path=None,
cli_overrides={
"num_gpus": 4,
"tensor_parallel_size": 2,
"default_sampling_params": '{"0": {"guidance_scale": 7.5}}',
},
stage_overrides=None,
strategy_config_path=None,
)
assert resolved.pipeline_config is not None
assert resolved.pipeline_config.model_type == "generic_diffusion"
assert len(resolved.stage_configs) == 1
stage = resolved.stage_configs[0]
assert isinstance(stage, VllmOmniDiffusionStageConfig)
assert stage.model_config.model == "generic-diffusion"
assert stage.model_config.default_sampling_params == {"guidance_scale": 7.5}
assert stage.parallel_config.tensor_parallel_size == 2
assert stage.parallel_config.data_parallel_size == 2
assert stage.parallel_config.world_size == 4
assert stage.runtime_config.devices == "0,1,2,3"
def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker):
"""Standard runtime resolves and starts the typed stage through the real launcher."""
from vllm_omni.engine import stage_engine_startup as startup_module
from vllm_omni.engine import stage_runtime as runtime_module
from vllm_omni.engine.stage_runtime import StageRuntime
mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
mocker.patch(
"vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
return_value=(True, "FakeDiffusionPipeline"),
)
resolved = resolve_omni_config(
"generic-diffusion",
trust_remote_code=False,
deploy_config_path=None,
cli_overrides={"num_gpus": 1},
stage_overrides=None,
strategy_config_path=None,
)
stage = resolved.stage_configs[0]
launched: dict[str, Any] = {}
client = SimpleNamespace(input_address=None, shutdown=mocker.Mock())
mocker.patch.object(runtime_module, "prepare_engine_environment")
mocker.patch.object(runtime_module, "load_omni_transfer_config_for_model", return_value=None)
mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
def _initialize_typed_stage(stage_id, model, stage_config, metadata, **kwargs):
launched.update(
stage_id=stage_id,
model=model,
stage_config=stage_config,
metadata=metadata,
**kwargs,
)
return client
mocker.patch.object(startup_module, "initialize_diffusion_stage", side_effect=_initialize_typed_stage)
runtime = StageRuntime(
stage_configs=list(resolved.stage_configs),
model="generic-diffusion",
config_path="",
stage_init_timeout=10,
async_chunk=False,
)
runtime.initialize()
assert launched["stage_config"] is stage
assert isinstance(launched["stage_config"], VllmOmniDiffusionStageConfig)
assert launched["metadata"].stage_type == "diffusion"
assert launched["metadata"].model_stage == "diffusion"
assert launched["use_inline"] is True
assert launched["model"] == "generic-diffusion"
assert launched["stage_id"] == 0
assert runtime.stage_pools[0].clients == [client]
tests/entrypoints/test_omni_entrypoints.py
+49 -13
View File
@@ -17,6 +17,8 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_config import StageConfig
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
@@ -202,6 +204,49 @@ def _make_base():
return obj
@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
def test_resolve_sampling_params_list_preserves_stage_constraints(typed):
from vllm_omni.config.omni_config import VllmOmniARStageConfig
from vllm_omni.config.stage_config import StagePipelineConfig
base = _make_base()
base.engine.num_stages = 1
base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])]
if typed:
stage = VllmOmniARStageConfig(
stage_pipeline_config=StagePipelineConfig(
stage_id=0,
model_stage="dummy-model",
sampling_constraints={"detokenize": False, "stop_token_ids": [42]},
)
)
else:
stage = StageConfig(
stage_id=0,
model_stage="dummy-model",
sampling_constraints={"detokenize": False, "stop_token_ids": [42]},
).to_omegaconf()
base.engine.stage_configs = [stage]
base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs)
assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}]
resolved_defaults = base.resolve_sampling_params_list(None)
assert resolved_defaults[0].stop_token_ids == [42]
caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7])
resolved = base.resolve_sampling_params_list(caller_params)
assert resolved[0] is not caller_params
assert resolved[0].seed == 1234
assert resolved[0].max_tokens == 7
assert resolved[0].detokenize is False
assert resolved[0].stop_token_ids == [7, 42]
assert 7 in resolved[0]._all_stop_token_ids
assert 42 in resolved[0]._all_stop_token_ids
assert caller_params.detokenize is True
assert caller_params.stop_token_ids == [7]
def test_resolve_sampling_params_list_merges_required_stop_tokens():
base = _make_base()
base.engine.num_stages = 1
@@ -219,22 +264,13 @@ def test_resolve_sampling_params_list_merges_required_stop_tokens():
base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs)
assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": required_stop_ids}]
resolved_defaults = base.resolve_sampling_params_list(None)
assert resolved_defaults[0].stop_token_ids == [99, 151704, 151645]
def test_sampling_constraints_are_forwarded_by_typed_stage_configs():
config = VllmOmniConfig.from_pipeline_config(OMNI_PIPELINES["qwen3_tts"])
caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[100, 151704])
constraints = OmniBase._get_sampling_constraints_list(config.stage_configs)
resolved = base.resolve_sampling_params_list(caller_params)
assert resolved[0] is not caller_params
assert resolved[0].seed == 1234
assert resolved[0].max_tokens == 7
assert resolved[0].detokenize is False
assert resolved[0].stop_token_ids == [100, 151704, 151645]
assert {100, 151704, 151645}.issubset(resolved[0]._all_stop_token_ids)
assert caller_params.detokenize is True
assert caller_params.stop_token_ids == [100, 151704]
assert constraints == [dict(stage.stage_pipeline_config.sampling_constraints) for stage in config.stage_configs]
@pytest.mark.parametrize("use_defaults", [False, True])
tests/entrypoints/test_resolve_dreamzero_config.py
+6 -3
View File
@@ -32,15 +32,18 @@ def test_dreamzero_resolves_through_registry_with_model_defaults(monkeypatch):
stage_overrides=None,
strategy_config_path=None,
)
engine_args = resolved.stage_configs[0].engine_args
stage = resolved.stage_configs[0]
assert resolved.config_path is not None
assert resolved.config_path.endswith("vllm_omni/deploy/dreamzero.yaml")
assert engine_args.model_class_name == "DreamZeroPipeline"
assert engine_args.model_config.policy_server_config.action_space == "joint_position"
assert stage.diffusion_config.model_class_name == "DreamZeroPipeline"
assert stage.diffusion_config.model_config["policy_server_config"]["action_space"] == "joint_position"
def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
def _get_hf_config(path, _model):
return None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}
monkeypatch.setattr(
"vllm.transformers_utils.config.get_hf_file_to_dict",
lambda path, _model, **_kwargs: (
tests/entrypoints/test_serve.py
+382
View File
@@ -6,12 +6,19 @@
from __future__ import annotations
import argparse
from dataclasses import dataclass
from multiprocessing.reduction import ForkingPickler
from types import SimpleNamespace
from typing import Any
import pytest
from pytest_mock import MockerFixture
from vllm.v1.engine.utils import EngineZmqAddresses
from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig
from vllm_omni.config.resolver import OmniConfigResolution
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
from vllm_omni.engine.stage_runtime import StageEngineLaunch
from vllm_omni.entrypoints.cli.serve import (
OmniServeCommand,
_parse_stage_overrides,
@@ -23,6 +30,31 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingName
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@dataclass(frozen=True)
class _FakeProcess:
sentinel: str
exitcode: int | None
name: str
pid: int
@dataclass
class _FakeProcessOwner:
processes: list[_FakeProcess]
@dataclass
class _FakeStageEngineArgs:
async_chunk: bool
enable_sleep_mode: bool
@dataclass
class _FakeStageConfig:
stage_id: int
engine_args: _FakeStageEngineArgs
def _resolved(*stages: SimpleNamespace) -> OmniConfigResolution:
return OmniConfigResolution(config_path="/fake/stages.yaml", stage_configs=tuple(stages))
@@ -42,6 +74,8 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
explicit = args.get_explicit_kwargs_dict()
assert args.get_explicit_kwargs_dict()
assert not explicit["async_chunk"]
assert (args.api_server_count or 1) == 1
assert "api_server_count" not in explicit
def _parse_serve_args(argv: list[str]) -> TrackingNamespace:
@@ -167,6 +201,244 @@ def test_serve_parser_rejects_invalid_video_output_transport(value: str) -> None
parser.parse_args(["serve", "fake-model", "--omni", "--video-output-transport", value])
def test_tracking_namespace_is_picklable_for_spawned_api_workers() -> None:
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="subcommand")
OmniServeCommand().subparser_init(subparsers)
args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2"])
args._omni_stage_client_configs = [{"stage_addresses": {0: {0: {"input_address": "ipc://input"}}}}]
restored = ForkingPickler.loads(ForkingPickler.dumps(args))
assert restored.api_server_count == 2
assert restored.get_explicit_kwargs_dict()["api_server_count"] == 2
assert restored._omni_stage_client_configs == args._omni_stage_client_configs
def test_serve_validate_rejects_multiple_api_servers_for_diffusion(mocker: MockerFixture) -> None:
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="subcommand")
cmd = OmniServeCommand()
cmd.subparser_init(subparsers)
args = parser.parse_args(["serve", "fake-diffusion-model", "--omni", "--api-server-count", "2"])
mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=True)
validate = mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
with pytest.raises(ValueError, match="not supported for diffusion"):
cmd.validate(args)
validate.assert_not_called()
def test_serve_validate_rejects_sleep_mode_with_multiple_api_servers(mocker: MockerFixture) -> None:
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="subcommand")
cmd = OmniServeCommand()
cmd.subparser_init(subparsers)
args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2", "--enable-sleep-mode"])
mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=False)
mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
with pytest.raises(ValueError, match="enable-sleep-mode"):
cmd.validate(args)
@pytest.mark.parametrize("typed", [False, True])
@pytest.mark.parametrize("downstream_async_chunk", [False, True])
def test_build_multi_api_stage_runtime_matches_current_constructor(
mocker: MockerFixture, typed: bool, downstream_async_chunk: bool
) -> None:
from vllm_omni.entrypoints.cli import serve as serve_module
args = TrackingNamespace(
argparse.Namespace(
model="dummy-model",
model_tag=None,
stage_init_timeout=1,
tokenizer=None,
disable_log_stats=False,
),
frozenset({"model"}),
)
if typed:
from vllm_omni.config.omni_config import VllmOmniARStageConfig
from vllm_omni.config.stage_config import StagePipelineConfig
stages = [
VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=i, model_stage="ar"))
for i in range(2)
]
stages[1].connector_config.async_chunk = downstream_async_chunk
else:
stages = [
_FakeStageConfig(
stage_id=i,
engine_args=_FakeStageEngineArgs(
async_chunk=bool(i and downstream_async_chunk), enable_sleep_mode=False
),
)
for i in range(2)
]
mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
mocker.patch(
"vllm_omni.config.config_factory.with_trust_remote_code_override",
side_effect=lambda kwargs, _trust: kwargs,
)
mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
mocker.patch(
"vllm_omni.config.resolver.resolve_omni_config",
return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=tuple(stages)),
)
runtime = serve_module._build_multi_api_stage_runtime(args, 2)
assert runtime._client_count == 1
assert runtime._client_index == 0
assert runtime._async_chunk is downstream_async_chunk
@pytest.mark.parametrize("typed", [False, True])
def test_build_multi_api_stage_runtime_rejects_sleep_enabled_in_stage_config(
mocker: MockerFixture, typed: bool
) -> None:
from vllm_omni.entrypoints.cli import serve as serve_module
args = TrackingNamespace(
argparse.Namespace(
model="dummy-model",
model_tag=None,
stage_init_timeout=1,
tokenizer=None,
disable_log_stats=False,
),
frozenset({"model"}),
)
stage_config = _FakeStageConfig(
stage_id=3,
engine_args=_FakeStageEngineArgs(async_chunk=False, enable_sleep_mode=True),
)
if typed:
from vllm_omni.config.omni_config import VllmOmniARStageConfig
from vllm_omni.config.stage_config import StagePipelineConfig
typed_stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=3, model_stage="ar"))
typed_stage.model_config.enable_sleep_mode = True
stage_config = typed_stage
mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
mocker.patch(
"vllm_omni.config.config_factory.with_trust_remote_code_override",
side_effect=lambda kwargs, _trust: kwargs,
)
mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
mocker.patch(
"vllm_omni.config.resolver.resolve_omni_config",
return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=(stage_config,)),
)
with pytest.raises(ValueError, match=r"sleep mode.*stage\(s\) \[3\]"):
serve_module._build_multi_api_stage_runtime(args, 2)
def test_run_multi_api_server_omni_starts_workers_after_shared_engine_launch(mocker: MockerFixture) -> None:
from vllm_omni.entrypoints.cli import serve as serve_module
args = TrackingNamespace(
argparse.Namespace(api_server_count=2, shutdown_timeout=1),
frozenset({"api_server_count", "shutdown_timeout"}),
)
socket = mocker.Mock()
primary_addresses = EngineZmqAddresses(
inputs=["ipc://input-0", "ipc://input-1"],
outputs=["ipc://output-0", "ipc://output-1"],
)
engine_launch = StageEngineLaunch(
client_configs=[{"client_count": 2, "client_index": index, "stage_addresses": {}} for index in range(2)],
resources=[StageReplicaResources(addresses=primary_addresses)],
)
runtime = mocker.MagicMock()
runtime.launch_stage_engines.return_value.__enter__.return_value = engine_launch
runtime.launch_stage_engines.return_value.__exit__.return_value = False
manager = mocker.Mock()
manager.processes = [mocker.Mock()]
manager.gather_actual_addresses.return_value = (primary_addresses.inputs, primary_addresses.outputs)
mocker.patch.object(serve_module, "_build_multi_api_stage_runtime", return_value=runtime)
mocker.patch.object(serve_module, "_wait_for_multi_api_server_completion")
mocker.patch("signal.signal")
mocker.patch("vllm.entrypoints.openai.api_server.setup_server", return_value=("127.0.0.1:8000", socket))
mocker.patch("vllm.v1.metrics.prometheus.setup_multiprocess_prometheus")
start_manager = mocker.patch.object(serve_module, "_start_api_server_process_manager", return_value=manager)
serve_module.run_multi_api_server_omni(args)
assert args._omni_stage_client_configs is engine_launch.client_configs
start_manager.assert_called_once()
manager_kwargs = start_manager.call_args.kwargs
assert manager_kwargs["cleanup_timeout"] == 1
assert manager_kwargs["num_servers"] == 2
assert manager_kwargs["input_addresses"] == ["ipc://input-0", "ipc://input-1"]
assert manager_kwargs["target_server_fn"] is serve_module.run_omni_api_server_worker_proc
assert engine_launch.watched_frontend_processes == manager.processes
runtime.shutdown.assert_called_once_with()
socket.close.assert_called_once_with()
def test_start_api_server_process_manager_cleans_up_partial_start(mocker: MockerFixture) -> None:
from vllm_omni.entrypoints.cli import serve as serve_module
first_process = mocker.Mock(name="first_process")
second_process = mocker.Mock(name="second_process")
first_process.pid = 1234
second_process.pid = None
second_process.start.side_effect = RuntimeError("spawn failed")
parent_pipes = [mocker.Mock(name="parent_pipe_0"), mocker.Mock(name="parent_pipe_1")]
child_pipes = [mocker.Mock(name="child_pipe_0"), mocker.Mock(name="child_pipe_1")]
spawn_context = mocker.Mock()
spawn_context.Process.side_effect = [first_process, second_process]
spawn_context.Pipe.side_effect = list(zip(parent_pipes, child_pipes))
mocker.patch("vllm.v1.utils.multiprocessing.get_context", return_value=spawn_context)
shutdown = mocker.patch("vllm.v1.utils.shutdown")
with pytest.raises(RuntimeError, match="spawn failed"):
serve_module._start_api_server_process_manager(
cleanup_timeout=1.5,
listen_address="127.0.0.1:8000",
sock=mocker.Mock(),
args=argparse.Namespace(),
num_servers=2,
input_addresses=["ipc://input-0", "ipc://input-1"],
output_addresses=["ipc://output-0", "ipc://output-1"],
target_server_fn=mocker.Mock(),
)
first_process.start.assert_called_once_with()
second_process.start.assert_called_once_with()
for pipe in parent_pipes:
pipe.close.assert_called_once_with()
shutdown.assert_called_once_with([first_process], timeout=1.5)
def test_wait_for_multi_api_server_completion_rejects_engine_failure(
mocker: MockerFixture,
) -> None:
from vllm_omni.entrypoints.cli import serve as serve_module
api_process = _FakeProcess(sentinel="api-0", exitcode=None, name="api-0", pid=10)
engine_process = _FakeProcess(sentinel="engine-0", exitcode=4, name="engine-0", pid=20)
manager = _FakeProcessOwner(processes=[api_process])
engine_launch = StageEngineLaunch(
client_configs=[],
resources=[StageReplicaResources(manager=_FakeProcessOwner(processes=[engine_process]))],
)
mocker.patch("multiprocessing.connection.wait", return_value=["engine-0"])
with pytest.raises(RuntimeError, match=r"Shared stage engine process engine-0 .* exited with code 4"):
serve_module._wait_for_multi_api_server_completion(manager, engine_launch)
def test_serve_parser_rejects_stage_configs_path() -> None:
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="subcommand")
@@ -334,6 +606,26 @@ def test_parse_stage_overrides_invalid_json_raises() -> None:
assert f"Got: {bad!r}" in message
@pytest.mark.parametrize(
"payload",
['{"abc": {}}', '{"-1": {}}', '{"1.5": {}}', '{"\uff10": {}}'],
)
def test_parse_stage_overrides_rejects_invalid_stage_ids(payload: str) -> None:
with pytest.raises(argparse.ArgumentTypeError, match="non-negative integer stage ids"):
_parse_stage_overrides(payload)
def test_parse_stage_overrides_preserves_arbitrary_override_fields() -> None:
parsed = _parse_stage_overrides(
'{"0": {"extras": {"ltx2_use_conv_vae": true}, "kv_cache_dtype": "fp8", "typo_field_xyz": 1}}'
)
assert parsed["0"] == {
"extras": {"ltx2_use_conv_vae": True},
"kv_cache_dtype": "fp8",
"typo_field_xyz": 1,
}
def test_parse_stage_overrides_rejects_non_dict_top_level() -> None:
"""Top-level must be a JSON object (dict). A list, scalar, or non-dict
mapping is rejected with a ValueError naming ``--stage-overrides`` and
@@ -706,6 +998,77 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
assert manager_kwargs["omni_replica_id"] == 0
def test_run_headless_generic_diffusion_launches_structured_stage(mocker: MockerFixture) -> None:
"""Headless resolution starts the typed stage through the real group launcher."""
from vllm_omni.engine import stage_engine_startup as startup_module
mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
mocker.patch(
"vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
return_value=(True, "FakeDiffusionPipeline"),
)
mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
mocker.patch.object(startup_module.stage_init_utils, "load_omni_transfer_config_for_model", return_value=None)
mocker.patch.object(
startup_module.initialization,
"resolve_omni_kv_config_for_stage",
return_value=(None, None, None),
)
captured: dict[str, Any] = {}
od_config = SimpleNamespace()
def _build_diffusion_config(model, stage_config, metadata):
captured.update(model=model, stage_config=stage_config, metadata=metadata)
return od_config
def _launch_replica_group(**kwargs):
captured.update(group_kwargs=kwargs)
captured["manager"] = kwargs["launch_one"](0)
def _launch_replica(**kwargs):
captured.update(replica_kwargs=kwargs)
return SimpleNamespace(exitcode=None)
mocker.patch.object(startup_module.stage_init_utils, "build_diffusion_config", side_effect=_build_diffusion_config)
mocker.patch.object(startup_module, "launch_headless_replica_group", side_effect=_launch_replica_group)
mocker.patch.object(startup_module, "launch_headless_diffusion_replica", side_effect=_launch_replica)
explicit_keys = frozenset(
{
"model",
"stage_id",
"omni_master_address",
"omni_master_port",
"worker_backend",
"model_class_name",
"num_gpus",
}
)
args = _make_headless_args(
explicit_keys=explicit_keys,
model="generic-diffusion",
model_class_name="FakeDiffusionPipeline",
num_gpus=1,
)
run_headless(args)
stage = captured["stage_config"]
assert isinstance(stage, VllmOmniDiffusionStageConfig)
assert captured["metadata"].stage_type == "diffusion"
assert captured["metadata"].model_stage == "diffusion"
assert captured["model"] == "generic-diffusion"
assert captured["group_kwargs"]["stage_id"] == 0
assert captured["group_kwargs"]["omni_dp_size_local"] == 1
assert captured["group_kwargs"]["per_replica_devices"] == ["0"]
assert captured["replica_kwargs"]["stage_config"] is stage
assert captured["replica_kwargs"]["stage_id"] == 0
assert captured["replica_kwargs"]["omni_master_address"] == "127.0.0.1"
assert captured["replica_kwargs"]["omni_master_port"] == 26000
assert captured["replica_kwargs"]["od_config"] is od_config
assert stage.runtime_config.devices == "0"
def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None:
"""A diffusion replica that exits with a non-zero code must surface as a
RuntimeError from ``run_headless`` (the head needs the signal to roll
@@ -754,3 +1117,22 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
with pytest.raises(RuntimeError, match=r"exited with code 137"):
run_headless(_make_headless_args(stage_id=1))
@pytest.mark.parametrize("name", ["unfiltered_ns", "explicit_keys", "missing", "__dict__"])
def test_tracking_namespace_uninitialized_access_raises_attribute_error(name):
namespace = object.__new__(TrackingNamespace)
with pytest.raises(AttributeError):
getattr(namespace, name)
@pytest.mark.parametrize("deep", [False, True])
def test_tracking_namespace_copy_preserves_tracking(deep):
import copy
namespace = TrackingNamespace(
argparse.Namespace(model="example", api_server_count=2), frozenset({"api_server_count"})
)
restored = copy.deepcopy(namespace) if deep else copy.copy(namespace)
assert restored.get_explicit_kwargs_dict() == {"api_server_count": 2}
assert restored.model == "example"
tests/entrypoints/test_utils.py
+142 -50
View File
@@ -13,6 +13,14 @@ import torch
from pytest_mock import MockerFixture
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm_omni.config.composable_parallel import (
Broadcast,
FanInByStage,
MeshAxisSpec,
RouteByStage,
StrategySpec,
TakeRank,
)
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.resolver import (
OmniConfigResolution,
@@ -21,7 +29,6 @@ from vllm_omni.config.resolver import (
resolve_omni_config,
)
from vllm_omni.config.stage_config import PipelineConfig
from vllm_omni.config.yaml_util import create_config
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.entrypoints.utils import (
@@ -356,16 +363,12 @@ class TestResolveOmniConfig:
)
assert resolved.config_path is None
assert len(resolved.stage_configs) == 1
engine_args = resolved.stage_configs[0]["engine_args"]
assert "dtype" in engine_args
assert engine_args["engine_backend"] == engine_backend
assert engine_args["revision"] == "pinned-revision"
assert resolved.stage_configs[0].diffusion_config.dtype is torch.float32
assert resolved.stage_configs[0].diffusion_config.engine_backend == engine_backend
assert resolved.stage_configs[0].diffusion_config.revision == "pinned-revision"
def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture):
mocker.patch(
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
return_value=None,
)
mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
resolve_model_class = mocker.patch(
"vllm_omni.config.resolver.resolve_model_class_name",
return_value="WanImageToVideoPipeline",
@@ -374,59 +377,98 @@ class TestResolveOmniConfig:
"vllm_omni.config.resolver.DiffusionModelRegistry.get_supported_archs",
return_value={"WanImageToVideoPipeline"},
)
load_model_class = mocker.patch("vllm_omni.config.resolver.DiffusionModelRegistry._try_load_model_cls")
resolved = resolve_omni_config(
"/models/Wan2.2-I2V",
trust_remote_code=False,
deploy_config_path=None,
cli_overrides={
"diffusion_load_format": "diffusers",
"revision": "pinned-revision",
},
cli_overrides={"diffusion_load_format": "diffusers", "revision": "pinned-revision"},
stage_overrides=None,
strategy_config_path=None,
)
resolve_model_class.assert_called_once_with("/models/Wan2.2-I2V", "diffusers", "pinned-revision")
load_model_class.assert_not_called()
stage = resolved.stage_configs[0]
resolve_model_class.assert_called_once_with(
"/models/Wan2.2-I2V",
"diffusers",
"pinned-revision",
)
assert stage.engine_args.model_class_name == "WanImageToVideoPipeline"
assert stage.diffusion_config.model_class_name == "WanImageToVideoPipeline"
assert stage.diffusion_config.revision == "pinned-revision"
assert stage.final_output_type == "video"
def test_generic_diffusion_stage_overrides_reach_typed_backend(self, mocker: MockerFixture):
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
mocker.patch(
"vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
return_value=(True, "LTX2Pipeline"),
)
resolved = resolve_omni_config(
"/models/LTX-2.5-Diffusers",
trust_remote_code=False,
deploy_config_path=None,
cli_overrides={
"extras": {"ltx2_use_conv_vae": False, "keep": "global"},
"tensor_parallel_size": 1,
"diffusion_streaming_output": True,
},
stage_overrides={
"0": {
"extras": {"ltx2_use_conv_vae": True},
"tensor_parallel_size": 2,
"devices": "2,3",
"engine_backend": "custom.backend",
},
},
strategy_config_path=None,
)
stage = resolved.stage_by_id(0)
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="/models/LTX-2.5-Diffusers")
assert stage.runtime_config.devices == "2,3"
assert stage.parallel_config.world_size == 2
assert engine_args["engine_backend"] == "custom.backend"
assert engine_args["extras"]["ltx2_use_conv_vae"] is True
assert engine_args["extras"]["keep"] == "global"
assert engine_args["streaming_output"] is True
def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture):
endpoint_restriction = SimpleNamespace(name="chat")
typed_stage = SimpleNamespace(stage_id=1)
structured_config = SimpleNamespace(
orchestrator_config=SimpleNamespace(deploy_config_path="/resolved/deploy.yaml"),
orchestrator_config=SimpleNamespace(
deploy_config_path="/resolved/deploy.yaml",
omni_lb_policy="round_robin",
),
pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)),
stage_configs=(typed_stage,),
strategy_omni_lb_policy="round-robin",
)
runtime_stage = create_config(
{
"stage_id": 1,
"runtime": {"devices": "1,2,3", "num_replicas": 3},
"engine_args": {"model": "dummy-model"},
}
)
legacy_stage = SimpleNamespace(to_omegaconf=lambda: runtime_stage)
create_structured = mocker.patch(
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
return_value=structured_config,
)
create_legacy = mocker.patch(
"vllm_omni.config.resolver.StageConfigFactory._resolve_legacy_from_registry",
return_value=SimpleNamespace(
stage_configs=[legacy_stage],
pipeline_config=structured_config.pipeline_config,
omni_lb_policy="round_robin",
),
"vllm_omni.config.resolver.StageConfigFactory.create_legacy_stage_configs_from_model",
)
strategy_specs = {"stage_1": {"dp": 3}}
strategy_specs = {
"stage_1": [
StrategySpec(
"stage_replica",
MeshAxisSpec("stage_replica", 2),
RouteByStage("round_robin"),
FanInByStage(),
)
]
}
load_strategy = mocker.patch(
"vllm_omni.config.resolver._load_strategy_specs",
return_value=strategy_specs,
)
apply_strategy = mocker.patch(
"vllm_omni.config.composable_parallel.apply_strategy_specs",
return_value=SimpleNamespace(omni_lb_policy="round_robin"),
)
resolved = resolve_omni_config(
"dummy-model",
@@ -447,31 +489,63 @@ class TestResolveOmniConfig:
trust_remote_code=None,
cli_overrides=expected_overrides,
deploy_config_path="deploy.yaml",
)
create_legacy.assert_called_once_with(
structured_config.pipeline_config,
expected_overrides,
"/resolved/deploy.yaml",
strategy_specs=strategy_specs,
)
create_legacy.assert_not_called()
load_strategy.assert_called_once_with("strategy.yaml")
apply_strategy.assert_not_called()
assert resolved.config_path == "/resolved/deploy.yaml"
assert resolved.pipeline_config is structured_config.pipeline_config
assert resolved.omni_lb_policy == "round_robin"
assert resolved.omni_lb_policy == "round-robin"
assert resolved.endpoint_restrictions == (endpoint_restriction,)
assert resolved.stage_configs == (runtime_stage,)
assert resolved.stage_configs == (typed_stage,)
def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
pipeline = OMNI_PIPELINES["qwen3_tts"]
assert isinstance(pipeline, PipelineConfig)
def test_tp_only_strategy_does_not_report_default_lb_policy_as_derived(self, mocker: MockerFixture):
structured_config = SimpleNamespace(
orchestrator_config=SimpleNamespace(deploy_config_path=None),
pipeline_config=pipeline,
orchestrator_config=SimpleNamespace(
deploy_config_path="/resolved/deploy.yaml",
omni_lb_policy="round-robin",
),
pipeline_config=SimpleNamespace(endpoint_restrictions=()),
stage_configs=(SimpleNamespace(stage_id=0),),
strategy_omni_lb_policy=None,
)
mocker.patch(
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
return_value=structured_config,
)
mocker.patch(
"vllm_omni.config.resolver._load_strategy_specs",
return_value={
"thinker": [
StrategySpec(
"tp",
MeshAxisSpec("tp", 2),
Broadcast(),
TakeRank(),
)
]
},
)
resolved = resolve_omni_config(
"dummy-model",
trust_remote_code=True,
deploy_config_path="deploy.yaml",
cli_overrides={"omni_lb_policy": "round-robin"},
stage_overrides=None,
strategy_config_path="strategy.yaml",
)
assert resolved.omni_lb_policy is None
def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
pipeline = OMNI_PIPELINES["qwen3_tts"]
assert isinstance(pipeline, PipelineConfig)
mocker.patch(
"vllm_omni.config.config_factory.StageConfigFactory.get_pipeline_config",
return_value=pipeline,
)
resolved = resolve_omni_config(
"Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
@@ -483,9 +557,9 @@ class TestResolveOmniConfig:
)
assert resolved.pipeline_config is not None
pipeline_stage_ids = [stage.stage_id for stage in resolved.pipeline_config.stages]
runtime_stage_ids = [stage.stage_id for stage in resolved.stage_configs]
assert pipeline_stage_ids == runtime_stage_ids
assert [stage.stage_id for stage in resolved.pipeline_config.stages] == [
stage.stage_id for stage in resolved.stage_configs
]
assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner"
assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1
@@ -529,3 +603,21 @@ class TestCumulativeStreamingCoercion:
assert all([isinstance(r, SamplingParams) for r in result])
assert result[0].output_kind == RequestOutputKind.DELTA
assert result[1].output_kind == RequestOutputKind.DELTA
@pytest.mark.parametrize("api_server_count", [1, 2])
def test_prepare_stage_config_inputs_filters_frontend_options(api_server_count):
from vllm_omni.entrypoints.utils import prepare_stage_config_inputs
kwargs = {
"api_server_count": api_server_count,
"disable_log_stats": True,
"model_tag": "example",
"max_num_seqs": 8,
"stage_overrides": '{"0": {"max_num_seqs": 4}}',
}
inputs = prepare_stage_config_inputs("example", kwargs, trust_remote_code=None)
assert inputs.kwargs == {"max_num_seqs": 8}
assert inputs.stage_overrides == {"0": {"max_num_seqs": 4}}
assert kwargs["api_server_count"] == api_server_count
tests/examples/offline_inference/test_text_to_video_prompts.py
+48 -1
View File
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import functools
import importlib.util
import sys
from types import SimpleNamespace
import pytest
import torch
@@ -24,6 +26,51 @@ def _load_text_to_video():
return module
@pytest.mark.parametrize(
("arguments", "expected"),
[
(["--quantization", "mxfp4"], {"quantization": "mxfp4"}),
(
["--quantization-config", '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}'],
{"quantization_config": {"method": "mxfp4", "w4a8_fallback_steps": [0, 2]}},
),
],
)
def test_quantization_cli_reaches_omni(arguments, expected, monkeypatch):
mod = _load_text_to_video()
captured = {}
class StopBeforeModelLoadError(Exception):
pass
def capture_omni(**kwargs):
captured.update(kwargs)
raise StopBeforeModelLoadError
monkeypatch.setattr(sys, "argv", ["text_to_video.py", "--model", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", *arguments])
monkeypatch.setattr(mod, "current_omni_platform", SimpleNamespace(device_type="cpu"))
monkeypatch.setattr(mod, "Omni", capture_omni)
with pytest.raises(StopBeforeModelLoadError):
mod.main()
assert {key: captured[key] for key in ("quantization", "quantization_config") if key in captured} == expected
def test_quantization_cli_options_are_mutually_exclusive(monkeypatch):
mod = _load_text_to_video()
monkeypatch.setattr(
sys,
"argv",
["text_to_video.py", "--quantization", "mxfp4", "--quantization-config", '{"method":"mxfp4"}'],
)
with pytest.raises(SystemExit) as exc:
mod.parse_args()
assert exc.value.code == 2
@pytest.mark.parametrize(
("negative_prompt", "expected_negative_prompt"),
[(None, None), ("", ""), ("blurry", "blurry")],
tests/platforms/npu/quant/test_mxfp4_step_fallback.py
+168
View File
@@ -0,0 +1,168 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Real A5 kernels: step selection, bias, shape, and reuse of prepared W4 weights."""
import pytest
import torch
pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.npu]
torch_npu = pytest.importorskip("torch_npu")
@torch.inference_mode()
def _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, *, policy, scale_alg=2):
from vllm_omni.diffusion.forward_context import (
ForwardContext,
override_forward_context,
set_forward_context_denoise_step_idx,
)
from vllm_omni.quantization.mxfp4_config import (
DiffusionMXFP4Config,
NPUMxfp4LinearMethod,
NPUMxfp4OnlineLinearMethod,
)
prefix = "blocks.10.attn1.to_qkv"
steps = [1] if policy in ("step", "layer_and_step") else []
layers = [prefix] if policy in ("layer", "layer_and_step") else []
fallback = {"w4a8_fallback_steps": steps, "w4a8_fallback_layers": layers, "mxfp4_scale_alg": scale_alg}
torch.manual_seed(42)
layer = torch.nn.Linear(1024, 128, bias=with_bias, device="npu", dtype=dtype)
if mode == "single":
method = NPUMxfp4OnlineLinearMethod(DiffusionMXFP4Config(**fallback), prefix=prefix)
elif mode == "single_offline":
method = NPUMxfp4LinearMethod(
DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, **fallback), prefix=prefix
)
layer.orig_dtype = dtype
values = torch.tensor([-4.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0])
layer.weight.copy_(values[torch.randint(8, (128, 1024))].to(layer.weight))
layer.register_parameter(
"weight_scale", torch.nn.Parameter(torch.full((128, 32), 124, dtype=torch.uint8, device="npu"), False)
)
layer.register_parameter("mul_scale", torch.nn.Parameter((torch.rand(1024) * 1.5 + 0.5).npu(), False))
else:
raise ValueError(f"Unsupported single-scale test mode: {mode}")
method.process_weights_after_loading(layer)
x = torch.randn(input_shape, device="npu", dtype=dtype)
pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
method.process_weights_after_loading(layer)
quantize = torch_npu.npu_dynamic_mx_quant
activation_types = []
def record_quantize(x, **kwargs):
dst = kwargs["dst_type"]
expected_alg = 0 if dst == torch_npu.float8_e4m3fn else scale_alg
assert kwargs["scale_alg"] == expected_alg
assert kwargs["axis"] == -1 and kwargs["block_size"] == 32 and kwargs["round_mode"] == "rint"
if expected_alg == 2:
assert kwargs["dst_type_max"] == 7.25
else:
assert "dst_type_max" not in kwargs
activation_types.append(dst)
return quantize(x, **kwargs)
monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", record_quantize)
results = []
for _request in range(2):
with override_forward_context(ForwardContext()):
for step in range(3):
set_forward_context_denoise_step_idx(step)
output = method.apply(layer, x, layer.bias)
torch.npu.synchronize()
assert output.shape == (*input_shape[:-1], 128)
assert output.dtype == dtype
assert torch.isfinite(output).all()
results.append(output.cpu())
expected_types = [
torch_npu.float8_e4m3fn if layers or step in steps else torch_npu.float4_e2m1fn_x2 for step in range(3)
]
assert activation_types == expected_types * 2
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
assert not hasattr(layer, "w4a8_weight")
assert not hasattr(layer, "w4a8_weight_scale")
for first, second in zip(results[:3], results[3:]):
torch.testing.assert_close(first, second, rtol=0, atol=0)
@pytest.mark.parametrize("with_bias", [False, True])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("input_shape", [(16, 1024), (2, 8, 1024)])
@pytest.mark.parametrize("mode", ["single", "single_offline"])
def test_mxfp4_step_fallback(with_bias, dtype, input_shape, mode, monkeypatch):
_exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, policy="step")
@pytest.mark.cards_1
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@torch.inference_mode()
def test_single_scale_step_fallback_matches_independent_dense_reference(dtype):
if not torch.npu.is_available() or "Ascend950" not in torch.npu.get_device_name():
pytest.skip("MXFP4 kernels require an Ascend950 NPU")
from vllm_omni.diffusion.forward_context import (
ForwardContext,
override_forward_context,
set_forward_context_denoise_step_idx,
)
from vllm_omni.quantization.mxfp4_config import (
DiffusionMXFP4Config,
NPUMxfp4LinearMethod,
)
# Synthetic serialized single-scale values and E8M0 exponent bytes.
# Powers of two make both activation precisions exact: this independent
# CPU dense oracle needs no production packing, dequantization or apply.
# Real UOS-quantized checkpoint validation remains a separate checkpoint/E2E gate.
values = torch.tensor([1.0, -1.0, 2.0, -2.0]).repeat(32).unsqueeze(1).expand(128, 512)
weight = torch.cat((values, values), dim=1)
fine = torch.cat((torch.full((128, 16), 127), torch.full((128, 16), 128)), dim=1).to(torch.uint8)
smooth = torch.cat((torch.full((512,), 0.5), torch.full((512,), 2.0)))
bias = torch.tensor([0.25, -0.25]).repeat(64)
x = torch.tensor([1.0, -1.0, 0.5, -0.5]).repeat(4).unsqueeze(1).expand(16, 1024) / 64
expected = torch.nn.functional.linear(x * smooth, torch.cat((values, 2 * values), dim=1), bias)
layer = torch.nn.Linear(1024, 128, bias=True, device="npu", dtype=dtype)
layer.weight.copy_(weight)
layer.bias.copy_(bias)
layer.orig_dtype = dtype
for name, value in {"weight_scale": fine, "mul_scale": smooth}.items():
layer.register_parameter(name, torch.nn.Parameter(value.npu(), requires_grad=False))
method = NPUMxfp4LinearMethod(
DiffusionMXFP4Config(
is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True, w4a8_fallback_steps=[1], mxfp4_scale_alg=2
)
)
method.process_weights_after_loading(layer)
prepared_pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
assert set(layer.state_dict()) == {"weight", "bias", "weight_scale", "mul_scale"}
assert not hasattr(layer, "w4a8_weight")
assert not hasattr(layer, "w4a8_weight_scale")
x = x.to(device="npu", dtype=dtype)
with override_forward_context(ForwardContext()):
for step in range(3):
set_forward_context_denoise_step_idx(step)
actual = method.apply(layer, x, layer.bias).cpu()
# A8's documented GEMM output is BF16 before restoring the caller
# dtype; the reference does not reuse quantization or unpacking.
reference = expected.bfloat16().to(dtype) if step == 1 else expected.to(dtype)
torch.testing.assert_close(actual, reference, rtol=0, atol=0)
assert actual.dtype == dtype
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == prepared_pointers
@pytest.mark.cards_1
@pytest.mark.parametrize("mode", ["single", "single_offline"])
@pytest.mark.parametrize("policy", ["layer", "layer_and_step"])
def test_mxfp4_layer_fallback(mode, policy, monkeypatch):
"""Reuse the kernel/shape/cache assertions without multiplying the full dtype/bias matrix."""
_exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy=policy)
@pytest.mark.cards_1
@pytest.mark.parametrize("mode", ["single", "single_offline"])
def test_mxfp4_ocp_compatibility(mode, monkeypatch):
"""OCP compatibility: A4 -> A8 -> A4 with the same W4 across requests."""
_exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy="step", scale_alg=0)