Merge branch 'main' into fix/tp-local-kv-recv-consensus

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.

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.
 

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次" -> "一次")

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

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

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),
+    ]

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):

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)

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

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"]},
     }
 

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 == []

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/<hf_model_type>.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):

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] == [

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:

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

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)

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:

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
 

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/<hf_model_type>.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(

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(

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()
 

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