diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index e54955c04..63cba58e5 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -14,7 +14,7 @@ The full set of backends and their platform defaults is in the **Backend Options | Value | Notes | |---|---| -| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. | +| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention and **SAGE** quantized attention — see below. Requires `flashinfer`. | | `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. | | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). | | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. | @@ -167,6 +167,48 @@ sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `T Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting `TRTLLM_ATTN` raises rather than silently degrading. +## TRTLLM_ATTN SAGE Quantization + +On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen +FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run +in low precision. Enable it through the typed `quant` block on the attention spec: + +| Key | Valid values | Meaning | +|---|---|---| +| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). | +| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). | +| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). | + +V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so +neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels. + +The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the +fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`), +`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`), +`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises. + +```bash +vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN", + "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}' +``` + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec + +AttentionConfig( + default=AttentionSpec( + backend="TRTLLM_ATTN", + quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16), + ), +) +``` + + +**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel +availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103** +(B300); `int8` QK kernels are compiled for **SM100 only**. + ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell) Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`. @@ -180,22 +222,6 @@ Same prompt and seed across runs. `Total generation time` from `text_to_video.py Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise. -## Known Limitations - -### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile - -LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121). - -**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0: - -```bash -DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model Lightricks/LTX-2 ... -``` - -### FA4 not yet integrated - -FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands. ## Choosing a Backend Manually @@ -353,64 +379,3 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_t --tensor-parallel-size 2 \ --output outputs/hv15_sage3.mp4 ``` - -### Mixed backends across roles - -Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention: - -```bash -python examples/offline_inference/text_to_video/text_to_video.py \ - --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ - --prompt "A dog running across a field of golden wheat." \ - --diffusion-attention-config.per_role.self.backend FLASH_ATTN \ - --diffusion-attention-config.per_role.cross.backend TORCH_SDPA \ - --tensor-parallel-size 2 \ - --output outputs/wan22_mixed.mp4 -``` - -### Compare against FlashAttention - -Unset the backend override, or explicitly use `FLASH_ATTN`: - -```bash -python examples/offline_inference/text_to_video/text_to_video.py \ - --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ - --prompt "A dog running across a field of golden wheat." \ - --height 704 --width 1280 --num-frames 49 \ - --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \ - --tensor-parallel-size 2 \ - --output outputs/wan22_fa3.mp4 -``` - -## Validation Guidance - -Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`. - -Always compare: - -- End-to-end runtime -- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown) -- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review) - -At minimum, keep the same: - -- model -- prompt -- seed -- resolution -- frame count / step count -- parallel config (TP / CFG-parallel / Ulysses degrees) - -## Reproducing the Benchmark Table - -The end-to-end numbers above were collected by running `text_to_video.py` / -`text_to_image.py` with the same prompt and seed while varying -`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the -backends without loading a model: - -```bash -python benchmarks/diffusion/bench_attention_backends.py --preset hv15 -``` - -It runs all three BF16 backends on representative DiT attention shapes and -prints a ranking table at the end. diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index 94fae5fa7..aea9a2b6a 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/examples/online_serving/minicpmo/README.md @@ -23,9 +23,11 @@ including `librosa`. The deploy config auto-loads via `--omni`. The default `vllm_omni/deploy/minicpmo_4_5.yaml` keeps all three stages on -logical device 0 with memory budgets of 55%, 22%, and 22%. The profile admits -at most four sequences per stage and bounds startup video profiling to 32 -frames per video. For throughput,`minicpmo_4_5_2gpu.yaml` gives the Thinker +logical device 0 with memory budgets of 55%, 15%, and 15%, leaving headroom +for runtime kernels such as the HiFi-GAN vocoder's cuDNN workspace. The +profile admits at most four sequences per stage and bounds startup video +profiling to 32 frames per video. For throughput, +`minicpmo_4_5_2gpu.yaml` gives the Thinker GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That profile admits at most four concurrent sequences per stage. diff --git a/pyproject.toml b/pyproject.toml index c442c3164..dd3080d84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ dev = [ ] demo = [ + "gradio>=6.7.0", "opencv-python>=4.12.0.88", # cli/stream_client frame reading "requests>=2.28.0", # cli/stream_client HTTP "cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次") diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index 340c89581..f87ee5b19 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/recipes/OpenBMB/MiniCPM-o-4_5.md @@ -77,7 +77,9 @@ a fallback because it would duplicate state machines and correctness paths. The default [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml) co-locates Thinker, codec-only Talker, and Code2Wav on GPU 0. Their -`gpu_memory_utilization` budgets are 0.55, 0.22, and 0.22. All stages admit +`gpu_memory_utilization` budgets are 0.55, 0.15, and 0.15. The remaining +device memory is left available for runtime workspaces such as the HiFi-GAN +vocoder's cuDNN kernels. All stages admit up to four sequences. Startup video profiling is bounded to 32 frames per video. Use the two-GPU profile for production throughput. @@ -202,9 +204,10 @@ speech output (TTS)"** checkbox on / off. #### Notes -- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.22, and - 0.22 of GPU 0. The larger Thinker share protects its multimodal KV cache; - all three model processes still share one CUDA device. +- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.15, and + 0.15 of GPU 0. The larger Thinker share protects its multimodal KV cache, + while the unreserved memory remains available to runtime kernels; all three + model processes still share one CUDA device. - `--trust-remote-code` is required — the HF repo ships a custom `MiniCPMO` config / model class. - Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 21167cb07..f37ac6ada 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -1,6 +1,8 @@ import json import os import threading +from collections.abc import Callable +from contextlib import AbstractContextManager, ExitStack, contextmanager from pathlib import Path from typing import Any @@ -63,35 +65,78 @@ paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_C _omni_server_lock = threading.Lock() +class _SingleActiveContext: + """Reuse one active context while its configuration key is unchanged.""" + + def __init__(self) -> None: + self._key: Any = None + self._stack: ExitStack | None = None + self._value: Any = None + + def acquire(self, key: Any, factory: Callable[[], AbstractContextManager[Any]]) -> Any: + if self._stack is not None and key == self._key: + return self._value + + self.close() + stack = ExitStack() + value = stack.enter_context(factory()) + self._key = key + self._stack = stack + self._value = value + return value + + def close(self) -> None: + stack = self._stack + self._key = None + self._stack = None + self._value = None + if stack is not None: + stack.close() + + +@contextmanager +def _start_omni_server(server_param): + test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = server_param + + print(f"Starting OmniServer with test: {test_name}, model: {model}") + + server_args: list[str] = [] + if use_omni: + server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"] + # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py): + # deploy-config sets the base; stage-overrides are applied on top. Both can be set. + if stage_config_path: + server_args = ["--deploy-config", stage_config_path] + server_args + if stage_overrides: + server_args = ["--stage-overrides", stage_overrides] + server_args + if extra_cli_args: + server_args = list(extra_cli_args) + server_args + with OmniServer(model, server_args, use_omni=use_omni) as server: + server.test_name = test_name + print("OmniServer started successfully") + yield server + print("OmniServer stopping...") + + print("OmniServer stopped") + + @pytest.fixture(scope="module") -def omni_server(request): +def omni_server_context(): """Start vLLM-Omni server as a subprocess with actual model weights. - Uses session scope so the server starts only once for the entire test session. + Reuse it for adjacent benchmark cases with the same server configuration. Multi-stage initialization can take 10-20+ minutes. """ with _omni_server_lock: - test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = request.param + active_context = _SingleActiveContext() + try: + yield active_context + finally: + active_context.close() - print(f"Starting OmniServer with test: {test_name}, model: {model}") - server_args: list[str] = [] - if use_omni: - server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"] - # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py): - # deploy-config sets the base; stage-overrides are applied on top. Both can be set. - if stage_config_path: - server_args = ["--deploy-config", stage_config_path] + server_args - if stage_overrides: - server_args = ["--stage-overrides", stage_overrides] + server_args - if extra_cli_args: - server_args = list(extra_cli_args) + server_args - with OmniServer(model, server_args, use_omni=use_omni) as server: - server.test_name = test_name - print("OmniServer started successfully") - yield server - print("OmniServer stopping...") - - print("OmniServer stopped") +@pytest.fixture +def omni_server(request, omni_server_context): + return omni_server_context.acquire(request.param, lambda: _start_omni_server(request.param)) @pytest.fixture diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py index 71d4aa95b..d6d1e0b30 100644 --- a/tests/dfx/perf/tests/test_runner_metadata.py +++ b/tests/dfx/perf/tests/test_runner_metadata.py @@ -1,6 +1,8 @@ """Tests for DFX runner metadata field exclusion.""" import json +from contextlib import contextmanager +from types import SimpleNamespace import pytest @@ -254,8 +256,9 @@ def test_benchmark_param_id_suffix_from_task_eval_phase(): ] -def test_create_paired_omni_benchmark_pytest_params(tmp_path): +def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_path, monkeypatch): from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params + from tests.dfx.perf.scripts import run_benchmark configs = [ { @@ -285,3 +288,51 @@ def test_create_paired_omni_benchmark_pytest_params(tmp_path): assert bench_row == ("test_tts", 0) assert any(m.name == "tts" for m in by_id["test_tts-p0"].marks) assert not any(m.name == "tts" for m in by_id["test_omni-p0"].marks) + + omni_p0_server = by_id["test_omni-p0"].values[0] + omni_p1_server = by_id["test_omni-p1"].values[0] + tts_server = by_id["test_tts-p0"].values[0] + assert omni_p0_server == omni_p1_server + assert omni_p0_server != tts_server + + events = [] + + @contextmanager + def fake_start(server_param): + events.append(("start", server_param)) + yield object() + events.append(("stop", server_param)) + + monkeypatch.setattr(run_benchmark, "_start_omni_server", fake_start) + active_context = run_benchmark._SingleActiveContext() + try: + first = run_benchmark.omni_server.__wrapped__( + SimpleNamespace(param=omni_p0_server), + active_context, + ) + second = run_benchmark.omni_server.__wrapped__( + SimpleNamespace(param=omni_p1_server), + active_context, + ) + assert first is second + assert events == [("start", omni_p0_server)] + + third = run_benchmark.omni_server.__wrapped__( + SimpleNamespace(param=tts_server), + active_context, + ) + assert third is not first + assert events == [ + ("start", omni_p0_server), + ("stop", omni_p0_server), + ("start", tts_server), + ] + finally: + active_context.close() + + assert events == [ + ("start", omni_p0_server), + ("stop", omni_p0_server), + ("start", tts_server), + ("stop", tts_server), + ] diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index adf13d51c..72b409de5 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -52,6 +52,44 @@ class TestAttentionSpec: with pytest.raises(ValueError, match="only supported by the TRTLLM_ATTN"): AttentionSpec(backend="TORCH_SDPA", skip_softmax={"target_sparsity": 0.5}) + def test_quant_serialized_with_defaults_and_overrides(self): + assert AttentionSpec(backend="TRTLLM_ATTN", quant={"dtype_qk": "fp8_e4m3"}).backend_kwargs()["quant"] == { + "dtype_qk": "fp8_e4m3", + "q_block_size": 1, + "k_block_size": 16, + } + assert AttentionSpec( + backend="TRTLLM_ATTN", quant={"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16} + ).backend_kwargs()["quant"] == {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16} + + def test_quant_superset_fields_passed_through(self): + # FLASHINFER_ATTN-style config (dtype_vo / flashinfer_backend) round-trips on the shared spec. + bk = AttentionSpec( + backend="FLASHINFER_ATTN", + quant={"dtype_qk": "bfloat16", "dtype_vo": "fp8_e4m3", "flashinfer_backend": "trtllm-gen"}, + ).backend_kwargs()["quant"] + assert bk["dtype_qk"] == "bfloat16" + assert bk["dtype_vo"] == "fp8_e4m3" + assert bk["flashinfer_backend"] == "trtllm-gen" + + def test_quant_and_skip_softmax_coexist(self): + bk = AttentionSpec( + backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5} + ).backend_kwargs() + assert bk["target_sparsity"] == 0.5 and bk["quant"]["dtype_qk"] == "int8" + + @pytest.mark.parametrize( + "spec, match", + [ + ({"backend": "TORCH_SDPA", "quant": {"dtype_qk": "int8"}}, "only supported by the TRTLLM_ATTN"), + ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int4"}}, "quant.dtype_qk"), + ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int8", "k_block_size": 8}}, "quant.k_block_size"), + ], + ) + def test_quant_validation_rejects(self, spec, match): + with pytest.raises(ValueError, match=match): + AttentionSpec(**spec) + class TestAttentionConfig: def test_empty_config(self): diff --git a/tests/diffusion/attention/test_trtllm_attn.py b/tests/diffusion/attention/test_trtllm_attn.py index 75d04c2b8..a2f855596 100644 --- a/tests/diffusion/attention/test_trtllm_attn.py +++ b/tests/diffusion/attention/test_trtllm_attn.py @@ -47,6 +47,44 @@ def test_skip_config_pure_resolution(): assert cfg.resolve_factor(4096, timestep=None) == pytest.approx(0.01 * 4096) +def test_quant_config_from_backend_kwargs(): + from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig + + assert not QuantConfig.from_backend_kwargs(None).enabled + assert not QuantConfig.from_backend_kwargs({}).enabled + assert QuantConfig.from_backend_kwargs({"quant": {"dtype_qk": "fp8_e4m3"}}) == QuantConfig("fp8_e4m3", 1, 16) + qc = QuantConfig.from_backend_kwargs({"quant": {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}}) + assert qc.enabled and (qc.dtype_qk, qc.q_block_size, qc.k_block_size) == ("int8", 4, 16) + + +def test_quant_rejects_non_sage_dtype(): + with pytest.raises(RuntimeError, match="supports dtype_qk"): + _impl(quant={"dtype_qk": "bfloat16"}) + + +def test_quant_quantize_requires_flashinfer_routine(monkeypatch): + import vllm_omni.diffusion.attention.backends.trtllm_attn as mod + + monkeypatch.setattr(mod, "_sage_kernel_available", lambda: True) + monkeypatch.setattr(mod, "_sage_quantize_fn", lambda: None) + with pytest.raises(RuntimeError, match="trtllm_sage_attention_quantize"): + _impl(quant={"dtype_qk": "int8"}) + + +def test_quant_quantize_calls_routine_and_shapes_sfs(): + from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig + + captured = {} + + def fake_quantize(q, k, v, q_block_size, k_block_size, qk_quant_dtype): + captured.update(q_block_size=q_block_size, k_block_size=k_block_size, qk_quant_dtype=qk_quant_dtype) + return "qq", "kq", "vq", "qsfs", "ksfs", "vsfs" + + q_q, k_q, v_q, sfs, blk = QuantConfig(dtype_qk="int8").quantize(object(), object(), object(), fake_quantize) + assert (q_q, k_q, v_q, sfs, blk) == ("qq", "kq", "vq", ("qsfs", "ksfs", None, "vsfs"), (1, 16, 0, 1)) + assert captured == {"q_block_size": 1, "k_block_size": 16, "qk_quant_dtype": torch.int8} + + def test_skip_factor_none_without_curve(): assert _impl(target_sparsity=0.5)._resolve_skip_factor(4096) is None @@ -257,6 +295,32 @@ def test_bf16_dense_matches_sdpa(): assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high" +requires_sage = pytest.mark.skipif( + not (_has_trtllm_attn() and tg._sage_quantize_fn() is not None), + reason="requires Blackwell SM100+ GPU with flashinfer >= 0.6.16rc1 (trtllm_sage_attention_quantize)", +) + + +@requires_sage +@pytest.mark.parametrize("dtype_qk", ["fp8_e4m3", "int8"]) +def test_sage_quant_matches_sdpa(dtype_qk): + torch.manual_seed(0) + B, S, H, D = 2, 256, 8, 128 + scale = 1.0 / math.sqrt(D) + q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + impl = _impl(quant={"dtype_qk": dtype_qk, "q_block_size": 1, "k_block_size": 16}) + try: + out = impl.forward_cuda(q, k, v, None) + except RuntimeError as e: # pragma: no cover - arch-dependent (int8 is SM100-only) + if "Missing TRTLLM-GEN kernel" in str(e): + pytest.skip(f"no trtllm-gen SAGE kernel for dtype_qk={dtype_qk} on this GPU arch") + raise + assert out.shape == (B, S, H, D) and torch.isfinite(out).all() + ref = _sdpa_ref(q, k, v, scale) + rel = (out.float() - ref).abs().mean() / ref.abs().mean() + assert rel < 0.1, f"SAGE({dtype_qk}) rel err {rel:.4f} too high" + + @requires_trtllm_attn def test_skip_tiny_threshold_matches_dense(): torch.manual_seed(0) diff --git a/tests/e2e/features/fullduplex/openai/test_runtime_control_redaction.py b/tests/e2e/features/fullduplex/openai/test_runtime_control_redaction.py new file mode 100644 index 000000000..16bea686d --- /dev/null +++ b/tests/e2e/features/fullduplex/openai/test_runtime_control_redaction.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Regression tests for duplex ``runtime_control`` payload redaction. + +The redacted result is emitted on the WebSocket (for example +``session.created.runtime_control``), so it must always survive +``json.dumps``. Control results carry typed values straight from the engine -- +``fence`` and ``accepted_fence`` are ``DuplexFence`` instances added by +``DuplexControlClient.execute`` -- which previously reached +``WebSocket.send_json`` unconverted and killed the session at handshake. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from enum import Enum + +import pytest + +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.openai.runtime_bridge import NativeRuntimeBridgeMixin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_redact = NativeRuntimeBridgeMixin._redact_runtime_control_result + + +def _control_result() -> dict[str, object]: + """The shape ``DuplexControlClient.execute`` returns for an open.""" + return { + "ok": True, + "operation": "open", + "session_id": "sess-1", + "fence": DuplexFence("sess-1", epoch=0, turn_id=0, response_seq=0, incarnation=0), + "accepted_fence": DuplexFence("sess-1", epoch=1, turn_id=2, response_seq=3, incarnation=4), + "stage_results": [{"stage_id": -1, "replica_id": -1, "result": {"supported": True}}], + } + + +def test_redacted_control_result_is_json_serializable() -> None: + """The whole payload must serialize the way send_json does.""" + redacted = _redact(_control_result()) + # Must not raise. + json.dumps({"type": "session.created", "runtime_control": redacted}) + + +def test_fences_are_converted_to_plain_dicts() -> None: + redacted = _redact(_control_result()) + assert redacted["fence"] == { + "session_id": "sess-1", + "epoch": 0, + "turn_id": 0, + "response_seq": 0, + "incarnation": 0, + } + # Field values are preserved, not just the type. + assert redacted["accepted_fence"]["epoch"] == 1 + assert redacted["accepted_fence"]["turn_id"] == 2 + assert redacted["accepted_fence"]["response_seq"] == 3 + assert redacted["accepted_fence"]["incarnation"] == 4 + + +def test_nested_dataclasses_are_converted() -> None: + """Fences can appear inside stage results, not only at the top level.""" + redacted = _redact( + { + "stage_results": [ + {"stage_id": 0, "result": {"fence": DuplexFence("sess-2", epoch=7)}}, + ] + } + ) + json.dumps(redacted) + assert redacted["stage_results"][0]["result"]["fence"]["epoch"] == 7 + + +def test_enums_are_reduced_to_values() -> None: + class _Mode(str, Enum): + APPEND = "append_audio_chunk" + + redacted = _redact({"mode": _Mode.APPEND}) + json.dumps(redacted) + assert redacted["mode"] == "append_audio_chunk" + + +def test_redaction_still_drops_heavy_payload_keys() -> None: + """The pre-existing contract must be preserved.""" + redacted = _redact( + { + "ok": True, + "stage_handoff": object(), + "tts_handoff": object(), + "omni_payload": object(), + "tts_hidden_states": object(), + "tts_token_ids": object(), + "traceback": "long trace", + } + ) + assert redacted == {"ok": True} + + +def test_plain_values_pass_through_unchanged() -> None: + payload = {"ok": True, "count": 3, "name": "x", "missing": None, "items": [1, "two"]} + assert _redact(payload) == payload + + +def test_dataclass_types_are_not_treated_as_instances() -> None: + """A class object is not a dataclass instance and must not be asdict'd.""" + + @dataclass + class _Shape: + x: int = 1 + + assert _redact(_Shape) is _Shape diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py index 089a304a7..6e76904b6 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py @@ -150,7 +150,16 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None: request_config = { "model": omni_server.model, "messages": messages, + "modalities": ["text", "audio"], "stream": True, + # Delimit the assistant answer with the TTS template so the Talker + # does not spend its codec-token budget speaking hidden reasoning. + "extra_body": { + "chat_template_kwargs": { + "use_tts_template": True, + "enable_thinking": False, + } + }, "key_words": {"audio": ["Beijing"]}, } diff --git a/tests/engine/test_stage_metadata.py b/tests/engine/test_stage_metadata.py index f56bbc2a9..0b21f7573 100644 --- a/tests/engine/test_stage_metadata.py +++ b/tests/engine/test_stage_metadata.py @@ -158,3 +158,16 @@ def test_extract_stage_metadata_preserves_legacy_one_argument_api(): assert metadata.stage_id == 0 assert metadata.model_stage == "thinker" assert metadata.custom_process_input_func is operator.add + + +def test_extract_stage_metadata_defaults_missing_engine_input_source(): + pipeline, deploy = _metadata_inputs() + legacy_config = merge_pipeline_deploy( + pipeline, + deploy, + )[0].to_omegaconf() + del legacy_config["engine_input_source"] + + metadata = extract_stage_metadata(legacy_config) + + assert metadata.engine_input_source == [] diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 4c8522ae0..c6d8e0354 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -3,6 +3,8 @@ import os from collections import Counter from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace import pytest import torch @@ -13,6 +15,10 @@ 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.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.stage_init_utils import ( + get_stage_connector_spec, + load_omni_transfer_config_for_model, +) from vllm_omni.entrypoints.utils import ( _convert_dataclasses_to_dict, _filter_dict_like_object, @@ -313,6 +319,58 @@ class TestResolveModelConfigPath: assert result is not None assert "glm_image.yaml" in result + def test_uses_registry_default_deploy_when_hf_model_type_differs( + self, + mocker: MockerFixture, + tmp_path, + ): + from vllm_omni.config.stage_config import _DEPLOY_DIR as _REAL_DEPLOY_DIR + from vllm_omni.entrypoints import utils as utils_mod + + hf_config = SimpleNamespace( + model_type="minicpmo", + architectures=["MiniCPMO"], + version="4.5", + ) + model = "local/minicpmo-registry-default-test" + mocker.patch( + "vllm_omni.entrypoints.utils.get_config", + return_value=hf_config, + ) + mocker.patch( + "vllm_omni.config.config_factory.get_config", + return_value=hf_config, + ) + # A bare deploy/.yaml must not win over the registered + # pipeline default when HF model_type is not an OMNI_PIPELINES key. + (tmp_path / "minicpmo.yaml").write_text("stages: []\n", encoding="utf-8") + (tmp_path / "minicpmo_4_5.yaml").write_bytes((_REAL_DEPLOY_DIR / "minicpmo_4_5.yaml").read_bytes()) + mocker.patch.object(utils_mod, "_DEPLOY_DIR", tmp_path) + + result = resolve_model_config_path(model) + + assert result is not None + assert Path(result).as_posix().endswith("minicpmo_4_5.yaml") + assert Path(result).name == "minicpmo_4_5.yaml" + transfer_config = load_omni_transfer_config_for_model(model, result) + assert transfer_config is not None + + sender = get_stage_connector_spec( + omni_transfer_config=transfer_config, + stage_id=1, + async_chunk=True, + ) + receiver = get_stage_connector_spec( + omni_transfer_config=transfer_config, + stage_id=2, + async_chunk=True, + ) + assert sender["extra"]["role"] == "sender" + assert receiver["extra"]["role"] == "receiver" + assert sender["extra"]["connector_get_sleep_s"] == 0.01 + assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000 + assert sender["extra"]["connector_get_max_wait"] == 300 + class TestLoadAndResolveStageConfigs: def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture): diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index b6f11ff5d..84b5ca11c 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -163,11 +163,13 @@ class TestDeployTopology: assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor) if filename == "minicpmo_4_5.yaml": assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4] - assert [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] == [ + memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] + assert memory_utilizations == [ 0.55, - 0.22, - 0.22, + 0.15, + 0.15, ] + assert sum(memory_utilizations) <= 0.9 + 1e-6 assert stages[0].yaml_engine_args["limit_mm_per_prompt"] == {"video": {"count": 1, "num_frames": 32}} elif filename in {"minicpmo_4_5_2gpu.yaml"}: assert [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] == [ diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py index 3938051d3..e0b52d1aa 100644 --- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py +++ b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py @@ -235,6 +235,9 @@ class _NoopPerf: def stop(self, name: str) -> None: pass + def reset(self) -> None: + pass + class _FakeTTS: feat_decoder = SimpleNamespace(estimator=SimpleNamespace(_compiled=False)) @@ -253,6 +256,164 @@ class _FakeTTS: return x.squeeze(1).mean(dim=1).repeat(1, 4) +class TestChunkedPrefillContract: + @staticmethod + def _make_talker(prompt_len: int = 30): + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import _PrefillInputs + + class FakeTTS: + def feat_encoder(self, x: torch.Tensor) -> torch.Tensor: + return x + + def enc_to_lm_proj(self, x: torch.Tensor) -> torch.Tensor: + return x.squeeze(2) + + def stop_proj(self, x: torch.Tensor) -> torch.Tensor: + return x + + def stop_actn(self, x: torch.Tensor) -> torch.Tensor: + return x + + def stop_head(self, x: torch.Tensor) -> torch.Tensor: + return x[:, :2] + + def lm_to_dit_proj(self, x: torch.Tensor) -> torch.Tensor: + return x + + def res_to_dit_proj(self, x: torch.Tensor) -> torch.Tensor: + return x + + class FakeScaffold: + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(torch.float32).unsqueeze(-1).repeat(1, 1, 4) + + talker = _make_bare_talker() + talker._tts = FakeTTS() + talker.model = FakeScaffold() + talker.config = SimpleNamespace(bos_token_id=-1, hidden_size=4) + talker._side_dtype = torch.float32 + talker._get_multichar_zh_split = lambda: {} + talker._perf = _NoopPerf() + talker._enable_torch_compile = False + talker._setup_cfm_buffers = lambda: None + talker._run_cfm_for_state = lambda *_args: torch.ones(1, 1, 4) + + text_token = torch.arange(prompt_len, dtype=torch.int32).unsqueeze(0) + audio_feat = torch.zeros(1, prompt_len, 1, 4) + text_mask = torch.ones(1, prompt_len, dtype=torch.int32) + audio_mask = torch.zeros(1, prompt_len, dtype=torch.int32) + talker._build_prefill_inputs = lambda *_args, **_kwargs: _PrefillInputs( + text_token=text_token, + audio_feat=audio_feat, + text_mask=text_mask, + audio_mask=audio_mask, + ) + return talker + + def test_preprocess_slices_full_embeddings_by_request_position(self) -> None: + talker = self._make_talker() + + _, first, _ = talker.preprocess( + input_ids=torch.ones(15, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + _omni_is_prefill=True, + _omni_num_computed_tokens=0, + ) + state = talker._active_states["req"] + state.decode_step_count = 7 + + _, second, _ = talker.preprocess( + input_ids=torch.ones(15, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + _omni_is_prefill=True, + _omni_num_computed_tokens=15, + ) + + torch.testing.assert_close(first[:, 0], torch.arange(15, dtype=torch.float32)) + torch.testing.assert_close(second[:, 0], torch.arange(15, 30, dtype=torch.float32)) + assert state.decode_step_count == 7 + assert state.is_last_prefill_chunk is True + + def test_prefill_cache_lifecycle_across_chunks(self) -> None: + talker = self._make_talker() + prompt_cache = {"mode": "reference", "ref_audio_feat": torch.ones(1)} + build_calls = [] + + def _build_prompt_cache(**_kwargs): + build_calls.append(1) + return prompt_cache + + talker._build_prompt_cache = _build_prompt_cache + + talker.preprocess( + input_ids=torch.ones(15, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + reference_audio="ref.wav", + _omni_is_prefill=True, + _omni_num_computed_tokens=0, + ) + state = talker._active_states["req"] + prefill_embeds = state.prefill_embeds + prefill_masks = state.prefill_masks + + talker.preprocess( + input_ids=torch.ones(15, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + reference_audio="ref.wav", + _omni_is_prefill=True, + _omni_num_computed_tokens=15, + ) + + assert build_calls == [1] + assert state.prompt_cache is prompt_cache + assert state.prefill_embeds is prefill_embeds + assert state.prefill_masks is prefill_masks + + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import _PrefillResidualMeta + + talker._finish_prefill( + state, + _PrefillResidualMeta( + lm_hidden=torch.ones(1, 4), + prefix_feat_cond=torch.ones(1, 1, 4), + ), + res_out=torch.ones(1, 4), + dev=torch.device("cpu"), + ) + + assert state.prefill_completed is True + assert state.prompt_cache is None + assert state.prefill_embeds is None + assert state.prefill_masks is None + assert talker._active_states["req"] is state + + def test_single_token_tail_is_still_prefill(self) -> None: + talker = self._make_talker() + + _, embeds, _ = talker.preprocess( + input_ids=torch.ones(1, dtype=torch.int32), + input_embeds=None, + request_id="req", + text_token_ids=[list(range(29))], + _omni_is_prefill=True, + _omni_num_computed_tokens=29, + ) + is_last_prefill_chunk = talker._active_states["req"].is_last_prefill_chunk + is_prefill = talker._pending_requests[-1][1] + assert embeds.shape == (1, 4) + assert embeds[0, 0] == 29 + assert is_prefill is True + assert is_last_prefill_chunk is True + + class TestDecodeBatchContract: def test_forward_skip_pending_prefill_keeps_audio_queue_aligned(self) -> None: class FakeScaffold: diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index ea1fdaf2f..abe8b5ad2 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -38,7 +38,8 @@ stages: - stage_id: 1 max_num_seqs: 4 - gpu_memory_utilization: 0.22 + # Leave device-level headroom for the colocated HiFi-GAN cuDNN workspace. + gpu_memory_utilization: 0.15 enforce_eager: false max_num_batched_tokens: 8192 devices: "0" @@ -58,7 +59,7 @@ stages: - stage_id: 2 max_num_seqs: 4 - gpu_memory_utilization: 0.22 + gpu_memory_utilization: 0.15 enforce_eager: true enable_chunked_prefill: false max_num_batched_tokens: 65536 diff --git a/vllm_omni/diffusion/attention/backends/trtllm_attn.py b/vllm_omni/diffusion/attention/backends/trtllm_attn.py index 9e3b305f1..718484a59 100644 --- a/vllm_omni/diffusion/attention/backends/trtllm_attn.py +++ b/vllm_omni/diffusion/attention/backends/trtllm_attn.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools +import inspect import math from dataclasses import dataclass, replace @@ -83,6 +85,66 @@ except Exception as e: # pragma: no cover - import guard ) +@functools.lru_cache(maxsize=1) +def _sage_kernel_available() -> bool: + if not HAS_FLASHINFER: + return False + try: + return "sage_attn_sfs" in inspect.signature(trtllm_ragged_attention_deepseek).parameters + except (TypeError, ValueError): + return False + + +@functools.lru_cache(maxsize=1) +def _sage_quantize_fn(): + try: + from flashinfer import trtllm_sage_attention_quantize + + return trtllm_sage_attention_quantize + except Exception: # pragma: no cover + return None + + +_QK_QUANT_DTYPES = { + "int8": torch.int8, + "fp8_e4m3": torch.float8_e4m3fn, +} + + +@dataclass(frozen=True) +class QuantConfig: + dtype_qk: str | None = None + q_block_size: int = 1 + k_block_size: int = 16 + + @classmethod + def from_backend_kwargs(cls, backend_kwargs: dict | None) -> "QuantConfig": + q = (backend_kwargs or {}).get("quant") or {} + return cls( + dtype_qk=q.get("dtype_qk"), + q_block_size=int(q.get("q_block_size", 1) or 1), + k_block_size=int(q.get("k_block_size", 16) or 16), + ) + + @property + def enabled(self) -> bool: + return self.dtype_qk is not None + + def quantize(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, quantize_fn): + qk_quant_dtype = _QK_QUANT_DTYPES[self.dtype_qk] + q_q, k_q, v_q, q_sfs, k_sfs, v_sfs = quantize_fn( + q, + k, + v, + q_block_size=self.q_block_size, + k_block_size=self.k_block_size, + qk_quant_dtype=qk_quant_dtype, + ) + sage_attn_sfs = (q_sfs, k_sfs, None, v_sfs) + num_elts_per_sage_attn_blk = (self.q_block_size, self.k_block_size, 0, 1) + return q_q, k_q, v_q, sage_attn_sfs, num_elts_per_sage_attn_blk + + def _workspace_bytes() -> int: import vllm.envs as envs @@ -129,6 +191,30 @@ class TrtllmAttentionImpl(AttentionImpl): self.skip = SkipSoftmaxConfig.from_backend_kwargs(backend_kwargs) self._warned_missing_timestep = False + self.quant = QuantConfig.from_backend_kwargs(backend_kwargs) + # Resolve the SAGE quantize fn once at init so the compiled forward path never calls the + # lru_cache-wrapped getter (which triggers a Dynamo graph break every step). + self._sage_quantize_fn = None + if self.quant.enabled: + if self.quant.dtype_qk not in _QK_QUANT_DTYPES: + raise RuntimeError( + f"TRTLLM_ATTN quant (SAGE) supports dtype_qk in {sorted(_QK_QUANT_DTYPES)}, got " + f"{self.quant.dtype_qk!r}. FLASHINFER_ATTN dtypes (float16/bfloat16) are not SAGE." + ) + if not _sage_kernel_available(): + raise RuntimeError( + "TRTLLM_ATTN quant (SAGE) was requested but this FlashInfer build does not " + "expose the trtllm-gen sage_attn_sfs kernel path. Install a FlashInfer build " + "that provides it, or remove the quant config." + ) + self._sage_quantize_fn = _sage_quantize_fn() + if self._sage_quantize_fn is None: + raise RuntimeError( + "TRTLLM_ATTN quant (SAGE) was requested but this FlashInfer build lacks " + "trtllm_sage_attention_quantize (added in flashinfer >= 0.6.16rc1). Upgrade " + "FlashInfer, or remove the quant config." + ) + def set_layer_calibration(self, a: float, b: float) -> None: self.skip = replace(self.skip, a=a, b=b) @@ -203,6 +289,15 @@ class TrtllmAttentionImpl(AttentionImpl): _skip_factor = self._resolve_skip_factor(kv_len) + # SAGE kwargs are only understood by newer FlashInfer builds; pass them exclusively when + # SAGE quant is active (which already requires the kernel, checked at init) so the dense + # path stays compatible with older builds that lack these parameters. + sage_kwargs: dict = {} + if self.quant.enabled: + q, k, v, sage_attn_sfs, sage_block_sizes = self.quant.quantize(q, k, v, self._sage_quantize_fn) + sage_kwargs["sage_attn_sfs"] = sage_attn_sfs + sage_kwargs["num_elts_per_sage_attn_blk"] = sage_block_sizes + out = trtllm_ragged_attention_deepseek( query=q, key=k, @@ -222,5 +317,6 @@ class TrtllmAttentionImpl(AttentionImpl): is_causal=self.causal, return_lse=False, skip_softmax_threshold_scale_factor=_skip_factor, + **sage_kwargs, ) return out.reshape(batch, q_len, num_q_heads, head_dim) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 86f8d8980..622d3ce7d 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1419,44 +1419,88 @@ class SkipSoftmaxSpec: raise ValueError("skip_softmax: set either target_sparsity or threshold, not both.") +@dataclass +class AttnQuantSpec: + dtype_qk: str | None = None + dtype_vo: str | None = None + q_block_size: int = 1 + k_block_size: int = 16 + flashinfer_backend: str | None = None + + _VALID_DTYPES = frozenset({"float16", "bfloat16", "int8", "fp8_e4m3"}) + _VALID_BLOCK_SIZES = frozenset({1, 4, 16}) + + def __post_init__(self) -> None: + for name, v in (("dtype_qk", self.dtype_qk), ("dtype_vo", self.dtype_vo)): + if v is not None and v not in self._VALID_DTYPES: + raise ValueError(f"quant.{name}={v!r} unsupported; use one of {sorted(self._VALID_DTYPES)}.") + for name, v in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)): + if v not in self._VALID_BLOCK_SIZES: + raise ValueError( + f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}." + ) + + @property + def enabled(self) -> bool: + return self.dtype_qk is not None or self.dtype_vo is not None + + @dataclass class AttentionSpec: """Specifies a backend and its typed backend-specific config for one attention role.""" backend: str skip_softmax: SkipSoftmaxSpec | None = None + quant: AttnQuantSpec | None = None skip_calibration: dict | None = field(default=None, repr=False) def __post_init__(self) -> None: if not isinstance(self.backend, str): raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}") - self.skip_softmax = self._coerce_skip_softmax(self.skip_softmax) + self.skip_softmax = self._coerce(self.skip_softmax, SkipSoftmaxSpec, "skip_softmax") + self.quant = self._coerce(self.quant, AttnQuantSpec, "quant") if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN": raise ValueError( f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. " - f"Remove skip_softmax or set backend to TRTLLM_ATTN." + "Remove skip_softmax or set backend to TRTLLM_ATTN." + ) + if self.quant is not None and self.backend.upper() not in ("TRTLLM_ATTN", "FLASHINFER_ATTN"): + raise ValueError( + f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but " + f"backend={self.backend!r}. Remove quant or set a supported backend." ) @staticmethod - def _coerce_skip_softmax(value: Any) -> SkipSoftmaxSpec | None: - if value is None or isinstance(value, SkipSoftmaxSpec): + def _coerce(value: Any, cls: type, field_name: str) -> Any: + if value is None or isinstance(value, cls): return value if isinstance(value, Mapping): - return SkipSoftmaxSpec(**dict(value)) - raise TypeError(f"Expected dict or SkipSoftmaxSpec for skip_softmax, got {type(value)!r}") + return cls(**dict(value)) + raise TypeError(f"Expected dict or {cls.__name__} for {field_name}, got {type(value)!r}") def backend_kwargs(self) -> dict[str, Any] | None: """Serialize typed backend config into the kwargs dict the backend impl consumes.""" - if self.skip_softmax is None: - return None - ss = self.skip_softmax kw: dict[str, Any] = {} - if ss.threshold is not None: - kw["skip_softmax_threshold"] = ss.threshold - if ss.target_sparsity is not None: - kw["target_sparsity"] = ss.target_sparsity - if ss.disabled_until_timestep: - kw["disabled_until_timestep"] = ss.disabled_until_timestep + if self.skip_softmax is not None: + ss = self.skip_softmax + if ss.threshold is not None: + kw["skip_softmax_threshold"] = ss.threshold + if ss.target_sparsity is not None: + kw["target_sparsity"] = ss.target_sparsity + if ss.disabled_until_timestep: + kw["disabled_until_timestep"] = ss.disabled_until_timestep + if self.quant is not None and self.quant.enabled: + q = self.quant + quant_kw: dict[str, Any] = { + "dtype_qk": q.dtype_qk, + "q_block_size": q.q_block_size, + "k_block_size": q.k_block_size, + } + if q.dtype_vo is not None: + quant_kw["dtype_vo"] = q.dtype_vo + if q.flashinfer_backend is not None: + quant_kw["flashinfer_backend"] = q.flashinfer_backend + kw["quant"] = quant_kw return kw or None @@ -1526,7 +1570,7 @@ class AttentionConfig: normalized[role] = node return - spec_keys = {"backend", "skip_softmax"} + spec_keys = {"backend", "skip_softmax", "quant"} node_dict = dict(node) node_keys = set(node_dict) if node_keys & spec_keys: diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 5ce22d9c8..3638b791e 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -372,7 +372,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: engine_args["attention_backend"] = "TRITON_ATTN" runtime_cfg = stage_config.runtime - engine_input_source: list[int] = stage_config.engine_input_source + engine_input_source: list[int] = _get_attr_or_item(stage_config, "engine_input_source", []) final_output: bool = stage_config.final_output final_output_type: str | None = stage_config.final_output_type diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 20e9747f9..dab544d62 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -13,6 +13,8 @@ from vllm.transformers_utils.config import get_config, get_hf_file_to_dict from vllm.transformers_utils.repo_utils import file_or_path_exists from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.config.stage_config import _DEPLOY_DIR from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero from vllm_omni.entrypoints.stage_utils import _to_dict @@ -226,22 +228,49 @@ def _try_resolve_omni_model_type(model: str) -> str | None: return best_match -def resolve_model_config_path(model: str) -> str: - """Resolve the stage config file path from the model name. +def _registry_default_deploy_path(model: str) -> str | None: + """Return the registered pipeline's default deploy YAML path, if any. - Resolves stage configuration path based on the model type and device type. - First tries to find a device-specific YAML file from stage_configs/{device_type}/ - directory. If not found, falls back to the default config file. + A checkpoint's HF ``model_type`` is not always the registered Omni pipeline + key. MiniCPM-o 4.5, for example, reports ``minicpmo`` while its + predicate-selected pipeline is ``minicpmo_4_5``. Stage construction already + loads that pipeline's default deploy config, so connector discovery must + resolve the same path. + """ + pipeline_config = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=True, + ) + if pipeline_config is None or pipeline_config.default_deploy_config_name is None: + return None + default_deploy_path = _DEPLOY_DIR / pipeline_config.default_deploy_config_name + if default_deploy_path.is_file(): + return str(default_deploy_path) + return None + + +def resolve_model_config_path(model: str) -> str | None: + """Resolve the stage/deploy config file path from the model name. + + Resolves configuration path based on the model type and device type. + Order: + 1. Device-specific stage config under ``stage_configs/{device_type}/`` + 2. If HF ``model_type`` is not an ``OMNI_PIPELINES`` key, the registered + pipeline's ``default_deploy_config_name`` (keeps connectors aligned with + stage construction when HF type and pipeline key differ) + 3. ``deploy/{model_type}.yaml`` + 4. Legacy ``stage_configs/{model_type}.yaml`` + 5. Registered pipeline default deploy config (final fallthrough) Args: model: Model name or path (used to determine model_type) Returns: - String path to the stage configuration file + String path to the stage/deploy configuration file, or ``None`` if no + matching config file exists. Raises: ValueError: If model_type cannot be determined - FileNotFoundError: If no stage config file exists for the model type """ # Try to get config from standard transformers format first try: @@ -296,15 +325,24 @@ def resolve_model_config_path(model: str) -> str: if os.path.exists(complete_config_path): return str(complete_config_path) - deploy_config_path = PROJECT_ROOT / "vllm_omni" / "deploy" / model_type_str + # Prefer the registry default before deploy/.yaml when the + # HF type is not itself a pipeline key. Otherwise a future + # deploy/minicpmo.yaml would desync connectors from stages for MiniCPM-o 4.5. + if normalized_model_type not in OMNI_PIPELINES: + registry_path = _registry_default_deploy_path(model) + if registry_path is not None: + return registry_path + + deploy_config_path = _DEPLOY_DIR / model_type_str if os.path.exists(deploy_config_path): return str(deploy_config_path) stage_config_file = f"vllm_omni/model_executor/stage_configs/{normalized_model_type}.yaml" stage_config_path = PROJECT_ROOT / stage_config_file - if not os.path.exists(stage_config_path): - return None - return str(stage_config_path) + if os.path.exists(stage_config_path): + return str(stage_config_path) + + return _registry_default_deploy_path(model) def load_stage_configs_from_model( diff --git a/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py index 3d6674621..258f512b6 100644 --- a/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py +++ b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py @@ -4,6 +4,8 @@ import asyncio import base64 import inspect from collections.abc import Mapping +from dataclasses import asdict, is_dataclass +from enum import Enum import numpy as np from vllm.logger import init_logger @@ -620,6 +622,17 @@ class NativeRuntimeBridgeMixin: return redacted if isinstance(value, list | tuple): return [cls._redact_runtime_control_result(item) for item in value] + if is_dataclass(value) and not isinstance(value, type): + # Control results carry typed values straight from the engine: + # `fence` and `accepted_fence` are `DuplexFence` instances added by + # `DuplexControlClient.execute`. This redacted value is emitted on + # the wire (for example `session.created.runtime_control` in + # `session_runner`), where `WebSocket.send_json` would raise + # `TypeError: Object of type DuplexFence is not JSON serializable` + # and kill the session during the handshake. + return {key: cls._redact_runtime_control_result(child) for key, child in asdict(value).items()} + if isinstance(value, Enum): + return value.value return value async def _send_native_duplex_events( diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 7ac2ba123..1bea2a124 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -89,7 +89,7 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: class _RotaryEmbedding(CustomOp): - """RoPE with HuggingFace-compatible CPU/CUDA math and cached NPU tables.""" + """RoPE with cached cos/sin tables; ``forward_native`` keeps the HF on-the-fly math.""" def __init__(self, config) -> None: super().__init__() @@ -101,25 +101,34 @@ class _RotaryEmbedding(CustomOp): rope_theta = getattr(config, "rope_theta", 10000.0) inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) - if current_omni_platform.is_npu(): - max_seq = int(getattr(config, "num_code_groups", 0) or 0) + 1 - positions = torch.arange(max_seq, dtype=torch.float32) - freqs = torch.outer(positions, inv_freq) - emb = torch.cat((freqs, freqs), dim=-1) - self.register_buffer("cos_cached", emb.cos(), persistent=False) - self.register_buffer("sin_cached", emb.sin(), persistent=False) + + # Build the cos/sin lookup tables once. ``num_code_groups + 1`` + # positions cover every re-prefill step of the code predictor. Compute + # in float32 (matching HF) and cast per-call in forward. + max_seq = int(getattr(config, "num_code_groups", 0) or 0) + 1 + positions = torch.arange(max_seq, dtype=torch.float32) + freqs = torch.outer(positions, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos(), persistent=False) + self.register_buffer("sin_cached", emb.sin(), persistent=False) + + def _lookup(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # position_ids: [batch, seq_len] -> cos/sin: [batch, seq_len, head_dim] + cos = self.cos_cached[position_ids] + sin = self.sin_cached[position_ids] + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) def forward_npu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return self.cos_cached[position_ids].to(dtype=x.dtype), self.sin_cached[position_ids].to(dtype=x.dtype) + return self._lookup(x, position_ids) def forward_cuda(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return self.forward_native(x, position_ids) + return self._lookup(x, position_ids) def forward_xpu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return self.forward_native(x, position_ids) + return self._lookup(x, position_ids) def forward_native(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - # position_ids: [batch, seq_len] + # HuggingFace on-the-fly computation, kept as a numeric reference. inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 94ac454d8..3a1677d49 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -264,10 +264,12 @@ class _RequestState: decode_step_count: int = 0 request_start_time: float = 0.0 prefill_completed: bool = False + is_last_prefill_chunk: bool = False prompt_cache: dict | None = None prefill_masks: tuple | None = None is_stopping: bool = False precomputed_is_stopping: bool | None = None + prefill_embeds: torch.Tensor | None = None @dataclasses.dataclass @@ -2054,7 +2056,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if is_prefill: flush_decode_fsq_batch() - res_input, meta = self._prepare_residual_prefill(state, req_hidden, dev) + res_input, meta = self._prepare_residual_prefill(state, req_hidden, dev, req_pos) elif state.prefill_completed: if ( self._enable_batched_fsq_fusion @@ -2141,7 +2143,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): offset += n if is_prefill: - prefill_batch.append((state, meta, res_out)) + if state.is_last_prefill_chunk: + prefill_batch.append((state, meta, res_out)) else: self._finish_decode(state, meta, res_out) decoded_states.append(state) @@ -2152,20 +2155,28 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): for state, meta, res_out in prefill_batch: self._finish_prefill(state, meta, res_out, dev) - collect_states = [state for state, _, _ in req_metas] - self._precompute_stop_flags_for_audio_collect(collect_states) - ready_audio_by_req = self._drain_ready_audio_copies_for_states(collect_states) - audio_by_req = self._collect_audio_batch( - collect_states, - initial_delayed_chunks_by_req={ - state.request_id: ready_audio_by_req.get(state.request_id) - for state, is_prefill, _ in req_metas - if not is_prefill - }, - ) + collect_states = [ + state for state, is_prefill, _ in req_metas if not is_prefill or state.is_last_prefill_chunk + ] + audio_by_req: dict[str, torch.Tensor | None] = {} + if collect_states: + self._precompute_stop_flags_for_audio_collect(collect_states) + ready_audio_by_req = self._drain_ready_audio_copies_for_states(collect_states) + audio_by_req = self._collect_audio_batch( + collect_states, + initial_delayed_chunks_by_req={ + state.request_id: ready_audio_by_req.get(state.request_id) + for state, is_prefill, _ in req_metas + if not is_prefill + }, + ) for state, is_prefill, _ in req_metas: - self._results_queue.append((state.request_id, state.precomputed_stop_logits)) - self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id))) + if is_prefill and not state.is_last_prefill_chunk: + self._results_queue.append((state.request_id, None)) + self._audio_queue.append((state.request_id, None)) + else: + self._results_queue.append((state.request_id, state.precomputed_stop_logits)) + self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id))) self._pending_requests.clear() self._flush_deferred_cleanup() @@ -2179,11 +2190,14 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state: _RequestState, base_lm_out: torch.Tensor, dev: torch.device, + req_pos: torch.Tensor, ) -> tuple[torch.Tensor, _PrefillResidualMeta]: tts = self.tts text_mask, feat_mask, feat, feat_embed = state.prefill_masks - state.prefill_masks = None - + text_mask = text_mask[:, req_pos] + feat_mask = feat_mask[:, req_pos] + feat = feat[:, req_pos] + feat_embed = feat_embed[:, req_pos] tts_len = text_mask.shape[1] scaffold_len = base_lm_out.shape[0] assert scaffold_len == tts_len, ( @@ -2424,6 +2438,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state.decode_step_count = 0 state.request_start_time = time.perf_counter() state.prefill_completed = True + state.prefill_masks = None + state.prefill_embeds = None + state.prompt_cache = None self._perf.stop("prefill_tail_batch") if self._enable_profiling: @@ -2465,6 +2482,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state.decode_step_count = 0 state.request_start_time = time.perf_counter() state.prefill_completed = True + state.prefill_embeds = None + state.prefill_masks = None + state.prompt_cache = None self._perf.stop("prefill_tail") if logger.isEnabledFor(logging.DEBUG): @@ -3033,7 +3053,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): span_len = int(input_ids.shape[0]) dev = input_ids.device req_id = info_dict.get("request_id", "default") - is_prefill = span_len > 1 + is_prefill = bool(info_dict.get("_omni_is_prefill", span_len > 1)) if is_prefill: # Do not evict state here: _pending_requests is a per-step prefix, @@ -3056,22 +3076,29 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): token_ids = token_ids[1:] state = self._get_or_create_state(req_id) - state.decode_pad = None - state.prefill_completed = False - state.decode_step_count = 0 - state.precomputed_stop_logits = None - state.precomputed_is_stopping = None - state.last_audio_patch_gpu = None - if not hasattr(state, "pending_audio_chunks_gpu"): - state.pending_audio_chunks_gpu = [] - if not hasattr(state, "pending_audio_copies"): - state.pending_audio_copies = [] - state.pending_audio_chunks_gpu.clear() - state.pending_audio_copies.clear() - state.curr_embed_for_next = None - state.prev_feat_embed = None - state.curr_prefix_feat_cond = None - state.is_stopping = False + num_computed_tokens = int(info_dict.get("_omni_num_computed_tokens", 0)) + is_first_prefill_chunk = num_computed_tokens == 0 + if is_first_prefill_chunk: + state.decode_pad = None + state.prefill_completed = False + state.is_last_prefill_chunk = False + state.decode_step_count = 0 + state.precomputed_stop_logits = None + state.precomputed_is_stopping = None + state.last_audio_patch_gpu = None + if not hasattr(state, "pending_audio_chunks_gpu"): + state.pending_audio_chunks_gpu = [] + if not hasattr(state, "pending_audio_copies"): + state.pending_audio_copies = [] + state.pending_audio_chunks_gpu.clear() + state.pending_audio_copies.clear() + state.curr_embed_for_next = None + state.prev_feat_embed = None + state.curr_prefix_feat_cond = None + state.is_stopping = False + state.prompt_cache = None + state.prefill_embeds = None + state.prefill_masks = None # Voice clone / continuation ref_audio = info_dict.get("reference_audio") or info_dict.get("ref_audio") @@ -3088,13 +3115,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): voice_profile = voice_profile[0] if voice_profile else None requires_precomputed_cache = isinstance(voice_profile, dict) - state.prompt_cache = None voice_name = info_dict.get("voice_name") if isinstance(voice_name, list): voice_name = voice_name[0] if voice_name else None _created_at = int(info_dict.get("voice_created_at") or 0) - if voice_name: + if state.prompt_cache is None and voice_name: _cache_key = self._speaker_cache.make_cache_key( voice_name, model_type="voxcpm2", created_at=_created_at ) @@ -3131,14 +3157,24 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): _key = self._speaker_cache.make_cache_key(voice_name, model_type="voxcpm2", created_at=_created_at) self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()}) logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name) - - inputs = self._build_prefill_inputs(token_ids, dev, req_id) - tts = self.tts - feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs.audio_feat)) - text_embed = self.model.embed_input_ids(inputs.text_token.to(dev)) - text_mask, feat_mask = inputs.text_mask, inputs.audio_mask - embeds = (text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed).squeeze(0) - state.prefill_masks = (text_mask, feat_mask, inputs.audio_feat, feat_embed) + if state.prefill_embeds is None or state.prefill_masks is None: + inputs = self._build_prefill_inputs(token_ids, dev, req_id) + tts = self.tts + feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs.audio_feat)) + text_embed = self.model.embed_input_ids(inputs.text_token.to(dev)) + text_mask, feat_mask = inputs.text_mask, inputs.audio_mask + embeds = (text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed).squeeze(0) + state.prefill_masks = (text_mask, feat_mask, inputs.audio_feat, feat_embed) + state.prefill_embeds = embeds + chunk_end = num_computed_tokens + span_len + state.is_last_prefill_chunk = chunk_end == state.prefill_embeds.shape[0] + if chunk_end > state.prefill_embeds.shape[0]: + raise ValueError( + "VoxCPM2 chunked prefill exceeds the constructed prompt: " + f"request_id={req_id} start={num_computed_tokens} end={chunk_end} " + f"prompt_len={state.prefill_embeds.shape[0]}" + ) + embeds = state.prefill_embeds[num_computed_tokens:chunk_end] else: state = self._active_states.get(req_id) curr = state.curr_embed_for_next if state else None