[Test] Align CosyVoice3 hardware mocks across platform layers

Signed-off-by: Sy03 <1370724210@qq.com>

.buildkite/common/ci_source_file_dependencies.yml

@@ -16,6 +16,9 @@ x-qwen3-tts: &qwen3_tts
   - vllm_omni/worker/gpu_model_runner.py
   - vllm_omni/model_executor/models/qwen3_tts/
   - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+  - vllm_omni/model_executor/models/common/fused_sampling.py
+  - vllm_omni/model_executor/models/common/short_kv_attention.py
+  - vllm_omni/model_executor/models/common/talker_first_audio.py
   - vllm_omni/model_executor/models/common/snake_activation.py
   - vllm_omni/model_executor/models/common/alias_free_activation.py
   - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
@@ -85,11 +88,16 @@ x-qwen3-omni: &qwen3_omni
   - vllm_omni/entrypoints/openai/serving_chat.py
   - vllm_omni/model_executor/models/qwen3_omni/
   - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+  - vllm_omni/model_executor/models/common/fused_sampling.py
+  - vllm_omni/model_executor/models/common/short_kv_attention.py
+  - vllm_omni/model_executor/models/common/talker_first_audio.py
   - vllm_omni/model_executor/models/common/snake_activation.py
   - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
   - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
   - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
   - vllm_omni/deploy/qwen3_omni_moe.yaml
+  - vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml
+  - vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml
   - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
 
 x-wan22: &wan22

.buildkite/cuda/test-merge.yml

@@ -124,6 +124,30 @@ steps:
           - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         mirror_hardwares: h100_2
 
+      - label: "Omni · Qwen3-Omni Async-Chunk SHM Leak"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_omni_shm_leak.py
+          - tests/helpers/runtime.py
+          - tests/helpers/stage_config.py
+          - vllm_omni/core/sched/omni_ar_scheduler.py
+          - vllm_omni/distributed/omni_connectors/connectors/shm_connector.py
+          - vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
+          - vllm_omni/engine/orchestrator.py
+          - vllm_omni/engine/stage_engine_core_proc.py
+          - vllm_omni/engine/stage_pool.py
+          - vllm_omni/entrypoints/openai/serving_chat.py
+          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+          - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
+          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+          - vllm_omni/deploy/qwen3_omni_moe.yaml
+        timeout_in_minutes: 40
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_shm_leak.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_2
+
       - label: "Omni · MiniCPM-o 4.5 Test"
         source_file_dependencies: omni_minicpmo_4_5_function
         timeout_in_minutes: 50

.buildkite/cuda/test-ready.yml

@@ -25,6 +25,16 @@ steps:
           - "pytest -sv tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py -k packed -m 'core_model and cuda' --run-level=core_model"
         mirror_hardwares: h100_1
 
+      - label: "Simple · MRV2 FA3 Graph Bounds Test"
+        timeout_in_minutes: 15
+        source_file_dependencies:
+          - tests/worker_v2/test_capture_query_bound_gpu.py
+          - tests/worker_v2/test_decode_query_bound_gpu.py
+          - vllm_omni/worker_v2/model_states/omni_model_state.py
+        commands:
+          - "pytest -sv tests/worker_v2/test_capture_query_bound_gpu.py tests/worker_v2/test_decode_query_bound_gpu.py -m 'core_model and cuda and H100' --run-level=core_model"
+        mirror_hardwares: h100_1
+
       - label: "Simple · Diffusion Test"
         timeout_in_minutes: 30
         commands:

docs/configuration/hsdp_pre_sharded_loading.md

@@ -0,0 +1,22 @@
+# Pre-sharded HSDP loading
+
+Use `--use-hsdp --hsdp-shard-size 8 --hsdp-weight-load-strategy pre_sharded`
+to shard transformer parameters before loading checkpoint values. Set the shard
+size to the available GPU count (or configure replication explicitly).
+The default strategy, `full`, retains the existing load-then-shard behavior.
+The same `hsdp_weight_load_strategy` option is available through Python and stage
+engine configuration.
+
+The pre-sharded strategy reads rank-local slices from ordinary Hugging Face
+safetensors checkpoints into FSDP-owned storage. No converted checkpoint is
+required. Cosmos 3 and Cosmos 3 Edge transformer blocks release their initialized
+parameter storage to meta as they are constructed; nonpersistent buffers retain
+their values.
+Encoders and VAEs follow their existing loading paths.
+
+This strategy requires the default diffusion pipeline loader, complete dedicated
+transformer checkpoint sources, and supported checkpoint key mappings. Tensor
+layouts must match the runtime layout; tensor transforms are unsupported.
+Checkpoint tensor dtypes must match runtime tensor dtypes at preflight; dtype conversion is unsupported.
+Quantization and LoRA are unsupported. Every worker needs checkpoint access.
+`num_weight_load_threads` controls the safetensors reader's threads per rank.

docs/configuration/stage_configs.md

@@ -319,7 +319,9 @@ ROCm and MUSA. Select one of these deployment profiles:
 | `qwen3_tts_mrv2.yaml` | V2 | B1 | Explicit MRV2 profile (same runner selection as the default) |
 | `qwen3_tts_high_concurrency_mrv2.yaml` | V2 | B1, B2 | Opt-in throughput tuning |
 | `qwen3_tts_high_concurrency_mrv2_b4.yaml` | V2 | B1, B2, B3, B4 | Experimental throughput / buffered playback |
+| `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` | V2 | B1–B8 | Experimental two-stage deployment on one GPU, with MPS |
 | `qwen3_tts_high_concurrency.yaml` | V1 | Existing defaults | V1 high-concurrency control |
+| `qwen3_tts_fused_single_gpu.yaml` | V2 | N/A (in-Talker decoder) | Opt-in single-stage CUDA pipeline on one GPU |
 
 ```bash
 # MRV2 is the default; pass a copy with `model_runner: v1` to force V1.
@@ -350,12 +352,73 @@ supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local
 (with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains
 its existing `talker_mtp` entry point for V1.
 
+### Selecting one or two Qwen3-TTS stages
+
+Choose the deployment profile explicitly with `--deploy-config`. The default
+`qwen3_tts` pipeline retains two stages: the Talker publishes codec chunks to
+Code2Wav, which produces audio. Its stage devices can be placed on the same
+GPU or on separate GPUs. `qwen3_tts_fused` is a separate, opt-in pipeline for
+one CUDA GPU: the Talker owns the stateful codec decoder and emits PCM directly,
+without a Code2Wav engine or stage connector.
+
+```bash
+# Two stages on one GPU; the two-stage pipeline remains independently selectable.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml
+
+# One stage on one GPU; no inter-stage transport or Code2Wav scheduling.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml
+```
+
+The single-stage profile requires CUDA, V2, asynchronous chunks, TP=PP=1,
+an in-process worker and disabled prefix caching. Select the full profile;
+adding `talker_stream_decode: true` to a two-stage Talker is rejected because
+its latent output belongs to the Code2Wav input contract. Single-stage options
+live in `stages[].additional_config`; two-stage transport and first-frame
+options live in `connectors.*.extra`.
+
+The single-stage profile defaults `talker_stream_first_audio` to `false` for
+throughput. To opt into a separate first-frame delivery queue, save this
+overlay beside the bundled deployment files and pass its path to
+`--deploy-config`:
+
+```yaml
+base_config: qwen3_tts_fused_single_gpu.yaml
+stages:
+  - stage_id: 0
+    additional_config:
+      talker_stream_decode: true
+      talker_stream_first_audio: true
+      ref_code_context_frames: 72
+      code_predictor_kv_cache: true
+      code_predictor_fused_sampling: true
+      code_predictor_fused: true
+```
+
+This queue uses PCM from the single-stage decoder. The two-stage
+`talker_first_audio` option below uses its own first-frame decoder and does
+not require single-stage execution. The two options cannot be enabled
+together. Single-stage prefix-cache reset with
+`reset_running_requests=True` returns `False` while requests are active;
+wait for completion or abort them before retrying.
+
+Reference-cloning context is selected independently by each profile. The
+single-stage profile explicitly uses 72 reference frames, matching the
+`qwen3_tts.yaml` context boundary. A shorter
+`additional_config.ref_code_context_frames` (for example 25) can change
+timbre continuity even when WER is similar;
+the throughput measurements for CustomVoice do not qualify Base voice
+cloning quality. Reference layout, grouping and priming are model hooks;
+the shared worker does not select Qwen's context length or codebook layout.
+
 ### Qwen3-TTS first-frame delivery and rollback
 
 The standard `qwen3_tts.yaml`, `qwen3_tts_high_concurrency.yaml`,
 `qwen3_tts_mrv2.yaml` and `qwen3_tts_high_concurrency_mrv2.yaml` profiles
 enable the `talker_first_audio` connector option, as does the experimental
-single-GPU profile. This changes the default CUDA MRV2 streaming path:
+two-stage `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile.
+This changes the default CUDA MRV2 streaming path:
 residual prediction runs eagerly after the Talker sample, the Talker loads
 an additional first-frame decoder with its weights and CUDA graphs, and the
 orchestrator orders its audio before subsequent Code2Wav chunks.

docs/design/feature/attention_backend_selection.md

@@ -4,6 +4,9 @@ This document defines the selection and extension contract for diffusion
 attention backends. User-facing backend choices, installation, and tuning are
 in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).
 
+The [execution contract PoC](attention_execution_contract_poc.md) describes the
+experimental post-selection capability and lifecycle contract from RFC #7226.
+
 ## Scope
 
 The contract applies to DiT and other diffusion attention layers. It is

docs/design/feature/attention_execution_contract_poc.md

@@ -0,0 +1,133 @@
+# Diffusion attention execution contract PoC
+
+This prototype accompanies [RFC #7226](https://github.com/vllm-project/vllm-omni/issues/7226)
+and extends the [attention selection design](attention_backend_selection.md).
+It demonstrates path-specific capabilities with dense BF16 FA4, NPU/ROCm
+routing examples, and tensor-state lifetime with a test-only attention module.
+
+## Execution contract
+
+`ExecutionContext` describes the requested execution path. `ExecutionPathResult`
+reports its identity, support status and reason, and compilation mode.
+`AttentionBackend.resolve_capabilities()` provides conservative pre-construction
+results. After initialization, `Attention.resolve_execution_path()` supplies the
+active parallel, paged-KV, and HSDP context; the backend combines it with the
+selected kernel and normalized metadata.
+
+Resolution runs outside compiled execution. FA4 reports `SUPPORTED` and
+`CUSTOM_OP` for dense, noncausal BF16 without parallel or HSDP boundaries when
+its kernel accepts the head dimensions. Dimension validation delegates to FA4's
+architecture-specific rules through `backends/utils/fa.py`. Kernel rejections
+report `UNSUPPORTED` with an actionable reason. Missing private validators and
+other unmigrated paths report `UNMIGRATED` with advisory `EAGER_ONLY` defaults.
+`requested_support()` checks a fullgraph request; it does not enforce selection
+or change existing execution.
+
+FA4 execution uses an opaque custom op. Its fake output preserves Q's batch,
+sequence, and head count and uses V's head dimension. Output is contiguous even
+when Q is noncontiguous. Metadata normalization is
+shared with dispatch. Producers must update published mask semantics when masks
+change; unpublished masks remain runtime-dependent to avoid synchronization.
+Callers resolve again when execution metadata changes.
+
+## NPU and ROCm worked examples
+
+The same `FLASH_ATTN` backend resolves MindIE paths on NPU and AITER paths on
+ROCm. Both examples report `UNMIGRATED`: their route is inspected, while hardware
+correctness and compilation remain unvalidated. `EAGER_ONLY` is an advisory
+default, not a measured compiler limitation. Existing execution is unchanged.
+
+| Input | NPU / MindIE | ROCm / AITER |
+| --- | --- | --- |
+| Dense | `npu_dense` | `rocm_dense` |
+| Padding mask | `npu_masked`: expand the key mask to `[B, 1, Q, K]` | `rocm_masked_varlen`: unpad, call varlen attention, then restore rows |
+| Packed inputs | Opt-in `[real, pad]` contract; choose varlen or prefix-KV slicing from the environment | Forward complete packed-document boundaries to varlen attention |
+| Incomplete packing | Rebuild a mask when possible; otherwise reject before the kernel | Shared CUDA-like normalization rejects incomplete metadata |
+| Unpublished mask semantics | Route to the masked wrapper | Keep resolution runtime-dependent without reading mask values |
+
+NPU resolution reuses `_resolve_packed_seq_npu`; ROCm reuses the metadata
+normalization used by dispatch. AITER identity comes from initialization rather
+than caller-provided context. Quantized, paged, piecewise, and parallel paths
+remain outside these worked examples.
+
+The new L1 tests use CPU tensors and replace only vendor kernel calls. They
+check path identity, forwarded masks/boundaries, fallback rejection, and
+conservative fullgraph-request handling. They do not test vendor numerics or
+compiler support. Existing L1 CI already collects this test file.
+
+```bash
+python -m pytest tests/diffusion/attention/test_flash_attn.py \
+  -k 'npu_contract or rocm_contract' -m 'core_model and cpu' -q
+```
+
+### Hardware validation required before migration
+
+On an Ascend/MindIE or ROCm/AITER runner, record device, driver, PyTorch, vendor
+library, and compiler versions, then validate the corresponding rows above:
+
+1. Compare eager outputs with FP32 SDPA for dense and padding-mask inputs,
+   including unequal Q/K lengths. Compare valid query rows; do not assume both
+   platforms define padded-query outputs identically.
+2. Compare packed outputs with separate per-document references. On NPU, cover
+   both varlen and laser prefix-slicing modes, plus mask reconstruction and the
+   missing-mask error. A padding fallback does not establish multi-document
+   isolation. On ROCm, cover multiple real documents.
+3. Check output shape, dtype, device, finite values, and unchanged inputs.
+   Start BF16 comparisons at `atol=rtol=1e-2` and have the platform maintainer
+   confirm the tolerance against its reference tests.
+4. Run `torch.compile(fullgraph=True, dynamic=True)` with the platform's supported
+   compiler over multiple sequence lengths. Record graph breaks and recompiles;
+   choose `TRACEABLE`, `CUSTOM_OP`, or verified `EAGER_ONLY` from that evidence.
+   A custom-op path additionally needs schema/fake checks and repeated replay.
+
+Promote only the paths validated on that platform. These hardware checks have
+not been run locally; ROCm/NPU fullgraph support is not claimed.
+
+## State lifetime example
+
+`test_attention_state_lifetime.py` prepares an owned query-scale tensor eagerly
+and passes it explicitly to an opaque attention op. The compiled module retains
+its state after the caller drops its reference and releases ownership when the
+compiled callable is deleted. Equivalent instances reuse a graph with their own
+state values.
+
+This example tests ownership and release, dynamic replay across instances,
+custom-op schema/fake behavior, and Inductor execution. It adds no production
+backend or state registry. Planning keys, caching, fallback-policy declarations,
+opaque handles, CUDA graph capture, and TRTLLM integration are deferred.
+
+## Validation
+
+```bash
+python -m pytest tests/diffusion/attention/test_flash_attn_compile.py \
+  tests/diffusion/attention/test_flash_attn.py \
+  tests/diffusion/attention/test_attention_capabilities.py \
+  tests/diffusion/attention/test_attention_state_lifetime.py -q -rs
+```
+
+The focused suite covers capability decisions and metadata changes, existing
+padding/mask regressions, tensor-state lifetime, and real FA4 fullgraph execution
+against FP32 SDPA and eager FA4. Representative Q/K and V dimensions are
+(32, 32), (64, 64), (80, 48), (192, 128), and (256, 256), with batches 1 and 2
+and sequence lengths up to 1,024. Invalid dimensions are checked against actual
+kernel errors. These samples exercise the contract rather than define support.
+BF16 comparisons use `atol=rtol=1e-2`; `torch.library.opcheck` checks schema and
+fake-output correctness.
+
+Single-graph reuse is asserted in the tensor-state test. The FA4 numerical tests
+retain compiler caches across shapes within each case but do not require one
+graph across different batch sizes, Q/K length equality, or head dimensions.
+
+Tested environment: GB300, PyTorch `2.13.0+cu130`, FA4 `4.0.0b18`, CUTLASS DSL
+`4.6.2` with CUDA 13 libraries, Quack `0.6.4`, and TVM FFI `0.1.11`.
+Real-kernel tests require CUDA and CuTe FA4; their availability is checked at runtime.
+
+Before the rebase, the focused suite passed all 97 tests without skips, including
+11 CPU NPU/ROCm routing tests. After rebasing onto upstream `e3be42e05`, the
+expanded suite reports 99 passed and 11 failed in the local vLLM 0.28 environment.
+All 11 failures are import errors: upstream now requires `compute_layout_strides`
+from vLLM, which this environment lacks. Both noncontiguous FA4 regression cases
+and real-kernel schema/fake checks pass. The full suite must be rerun with the
+vLLM version required by upstream. NPU/ROCm device numerics and compilation remain
+unvalidated. Pre-commit passes with the CI hook skips. Upstream PyTorch/CUTLASS
+warnings remain.

docs/design/index.md

@@ -68,6 +68,7 @@ guides list selectable backends, platform defaults, installation, and tuning.
 The design contracts separate selection mechanics from backend algorithms:
 
 - [Attention Backend Selection](feature/attention_backend_selection.md)
+- [Attention Execution Contract PoC](feature/attention_execution_contract_poc.md)
 - [Skip-Softmax](feature/skip_softmax.md)
 
 #### CPU offloading
@@ -116,4 +117,5 @@ The pre-#5137 pages are preserved in the
 [legacy module archive](module/archive/README.md) for historical reference and
 are not active design contracts.
 
+- [Qwen3-Omni MRv2 performance](qwen3_omni_mrv2_performance.md)
 - [MiniCPM-o 4.5 turn-mode MRv2 performance](minicpm_o45_mrv2_performance.md)

docs/design/qwen3_omni_mrv2_performance.md

@@ -0,0 +1,22 @@
+# Qwen3-Omni MRv2 performance
+
+This model adaptation builds on #8184's existing stage runtime and output
+contracts. It adds Talker CPU metadata parsing, sampled-token embedding
+handoff, an optional first-frame decoder, length-grouped Code2Wav graph
+buckets, and opt-in incremental short-KV code prediction. The H200 profile
+also enables conservative fused sampling and bounded cuDNN algorithm search.
+The shared predictor keeps the upstream re-prefill helper and its default
+sampling behavior; short KV is explicitly selected by the deployment profile.
+
+Use `vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml` for the experimental H200
+profile. Preserve a control with `code_predictor_fused_sampling=false` and
+`codec_cudnn_benchmark=false` when C64 P99 is the primary requirement.
+cuDNN search restores process settings on exit and has startup cost. Kernel
+arithmetic is not bitwise-equivalent; the first-frame decoder uses additional
+model memory. Full-duplex performance is not established by turn-mode tests.
+
+First-frame delivery requires `talker_first_audio: true` in connector extra,
+a one-frame initial codec chunk, CUDA, a local single-rank Talker and prefix caching disabled.
+Other deployments retain the regular Code2Wav path.
+`codec_fused_snake` explicitly opts CUDA profiles into fused decoder blocks;
+the base V1 profile retains its original blocks and indexed text cache.

docs/serving/video_stream_api.md

@@ -16,6 +16,10 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
     --trust-remote-code
 ```
 
+To enable [incremental prefill](#incremental-prefill), set
+`enable_prefix_caching: true` on stage 0 (the thinker) in the deploy config. Otherwise
+the server falls back to rebuilding the prompt per query.
+
 ### Run the Example Client
 
 ```bash
@@ -54,7 +58,7 @@ WebSocket /v1/video/chat/stream
 
 When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt.
 
-Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path.
+Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. With incremental prefill active, every buffered frame is included instead of stride-sampling, with the same decode-failure exclusion. Excluded frames are not replaced with other buffered frames. The incremental path waits for the whole window to finish decoding. The legacy path waits only for its sampled positions, so an unselected frame still decoding does not delay the query.
 
 **Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`.
 
@@ -64,10 +68,10 @@ Queries stride-sample the buffered frames, keeping the last frame, then exclude
 | ------- | ------ | --------- | ------------- |
 | `model` | string or null | null | Optional model name. Usually omitted because the server hosts one model. |
 | `modalities` | list[string] | `["text", "audio"]` | Output modalities. Use `["text"]`, `["audio"]`, or both. |
-| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. |
-| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. |
+| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. Legacy path only: with incremental prefill active, every buffered frame is submitted and frame density is the client's responsibility (clients push discrete frames at their own rate). |
+| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Below the cap both paths keep every frame. Legacy evicts the oldest frame once full. Incremental prefill drops the oldest half at once, then appends until full again. |
 | `system_prompt` | string or null | null | Optional custom system prompt. |
-| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. |
+| `use_audio_in_video` | bool | `true` | Incremental prefill: forwarded on warmup and every query when this is true, so multimodal hashes stay aligned. Legacy: forwarded only when this query has input audio. |
 | `sampling_params_list` | list or null | null | Optional per-stage parameter dictionaries. Each provided entry replaces that stage's deployment sampling settings. |
 | `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. |
 | `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. |
@@ -107,9 +111,42 @@ The server accepts these legacy field names and rewrites them before validation.
 
 EVS compares downsampled frames and drops near-duplicate frames before they enter the session frame buffer. `frame_filter_threshold` controls retention: higher values are more permissive and keep more frames; lower values are more aggressive and drop more similar frames.
 
+## Incremental Prefill
+
+Server-side only (stage-0 `enable_prefix_caching: true` and `modalities: ["text"]`).
+There is no `session.config` flag. Audio-output sessions stay on the legacy path.
+
+Input is unchanged: `video.frame` and `audio.chunk` are separate buffers. Warmup
+prefills `history + frames` as they arrive (`max_tokens=1`). When
+`use_audio_in_video` is true, warmup and query both pass that processor kwarg.
+`video.query` is the same prompt plus optional input audio and the question text,
+so the warmed vision prefix cache is reused and the query pays for the audio/text
+suffix. All buffered frames are
+submitted in arrival order (`num_frames` subsample is legacy-only). Once the buffer
+hits `max_frames`, the oldest half is dropped in one step and a warning is logged;
+until then the buffer matches the legacy path. After a turn,
+history compresses to the last two text messages and the next warmup follows the
+new prefix.
+
+### Legacy path (incremental prefill inactive)
+
+Stage-0 prefix caching is off, or `modalities` includes `"audio"`: no warmup,
+frames are re-sampled up to `num_frames` each query, and the query waits only
+for those sampled frames to decode. Prompt shape is the same
+(`image_pil` frames, optional trailing `input_audio`, text). Both paths keep the
+frame buffer across queries (`max_frames`) and compress history the same way.
+
 ## Known Limitations
 
-- Session KV reuse and incremental prefill are not implemented in this PR. Each `video.query` rebuilds the model prompt from the retained frame and audio buffers.
+- Past `max_frames`, incremental prefill keeps a suffix of the newest frames, not a resample of the whole stream. Rebuilding that longer window is left to a follow-up.
+- Incremental prefill applies only to text-only sessions with stage-0 prefix caching enabled
+  on `WS /v1/video/chat/stream`. It does not accelerate `/v1/realtime` camera input, which
+  uses duplex conversation items rather than `video.frame` / `video.query`. Audio-output
+  sessions on this endpoint rebuild the prompt each query.
+- Stage-0 prefix caching is safe for audio-output requests: on a hit, the omni prefix cache
+  restores the thinker hidden states the talker needs. A preempted request that resumes on a
+  hit does not get them back, so size the deployment to avoid preemption; see
+  [Prefix Caching](../design/feature/prefix_caching.md).
 - Back-to-back short replies can still expose an engine-layer scheduler race. The PR notes an observed workaround of at least 200 ms idle between turns when clients repeatedly see idle timeouts.
 - If the audio buffer exceeds the server limit, the server emits `Audio buffer overflow` and clears the currently buffered audio for the session.
 - The API is intended for Qwen3-Omni streaming video understanding; other models may not support the same multimodal processor arguments.

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -130,7 +130,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- H200 automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when the CUDA MPS control tool is available; other devices and environments use `cosyvoice3.yaml`. See the [serving defaults and tradeoffs](../online_serving/text_to_speech.md#cosyvoice3). Pass `--deploy-config cosyvoice3.yaml` for the previous profile or `--deploy-config <path>` to override. The example uses the resolved profile's sampling parameters.
 - `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
 
 ---

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -155,16 +155,22 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
 ./cosyvoice3/run_server.sh
 ```
 
-Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
+Streaming is on by default. On CUDA Hopper GPUs with at least 140 GiB of device memory (H200), serving automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when `nvidia-cuda-mps-control` is available on `PATH`: AR32 / codec32, standard sampling, packed Flow, cached ISTFT, HiFT decode graphs, and local CUDA MPS. No `--hf-overrides` is needed. Other devices and environments without the MPS control tool use `cosyvoice3.yaml` with RAS sampling and AR8 / codec8.
 
-Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call:
+The optimized profile is not a universal hardware configuration. Its packed kernels are Hopper-only; the talker's Model Runner V2 is unsupported on NPU/XPU, and CUDA MPS requires a local NVIDIA GPU and its control tool. Larger batches, 25-frame chunks, and graph allocations also require per-device validation on smaller or slower GPUs. Automatic selection retains the generic profile outside the validated device class; an explicit deploy configuration takes precedence.
+
+The H200 default targets streaming throughput at C32/C64. It changes sampling from RAS to standard (temperature 0.7, top-p 0.8, top-k 20, repetition penalty 1.21); output lengths and quality can differ. HiFT captures a recurring exact shape on its third use, which adds latency to that request, and its graph pool needs memory headroom beyond the engine reservations. This profile is not a claim of optimal first-audio latency at low concurrency. Pass `--deploy-config cosyvoice3.yaml` to retain the previous configuration, or `--deploy-config <path>` to select another profile.
+
+Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path; use `--deploy-config cosyvoice3.yaml` alongside it when selecting that path on H200.
+
+Cross-request Stage-1 flow batching is enabled in the H200 profile. For the generic configuration, enable it in the server environment when concurrent requests should share a flow-estimator call:
 
 ```bash
 export COSYVOICE3_BATCH_FLOW=1
 vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
 ```
 
-Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
+Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Use `--deploy-config cosyvoice3.yaml` and leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
 
 ### CLI client
 
@@ -194,7 +200,7 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- Deploy config auto-loads the device-specific CosyVoice3 default described above. Pass `--deploy-config <path>` to override.
 - For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
 
 ---

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import argparse
+import copy
 import functools
 import json
 import os
@@ -10,13 +11,11 @@ from typing import Any
 
 import numpy as np
 import soundfile as sf
-from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 # Upstream zero-shot reference clip
 ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
@@ -59,8 +58,7 @@ def run_e2e():
         "--deploy-config",
         type=str,
         default=None,
-        help="Override the deploy config path. If unset, auto-loads "
-        "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
+        help="Override the deploy config path. If unset, selects the bundled CosyVoice3 profile for the device.",
     )
     parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
     parser.add_argument(
@@ -105,8 +103,6 @@ def run_e2e():
         profiler_config=args.profiler_config,
     )
 
-    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
-
     print("Model initialized. Preparing inputs...")
     ref_audio_path = args.ref_audio or _default_ref_audio()
     if not os.path.exists(ref_audio_path):
@@ -138,7 +134,7 @@ def run_e2e():
 
     print(f"Generating for prompt: {args.text}")
 
-    config = CosyVoice3Config()
+    config = omni.engine.stage_vllm_configs[0].model_config.hf_config
     tokenizer = get_qwen_tokenizer(
         token_path=args.tokenizer,
         skip_special_tokens=config.skip_special_tokens,
@@ -149,29 +145,10 @@ def run_e2e():
     min_len = int(base_len * config.min_token_text_ratio)
     max_len = int(base_len * config.max_token_text_ratio)
 
-    # Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
-    gpt_sampling = SamplingParams(
-        temperature=1.0,
-        top_p=sampling_cfg["top_p"],
-        top_k=sampling_cfg["top_k"],
-        repetition_penalty=2.0,
-        min_tokens=min_len,
-        max_tokens=max_len,
-        stop_token_ids=[sampling_cfg["eos_token_id"]],
-        # allowed_token_ids=list(range(6561+3)),
-        detokenize=False,
-    )
-    # Not used
-    s2mel_sampling = SamplingParams(
-        temperature=1.0,
-        top_p=1.0,
-        top_k=-1,
-        repetition_penalty=2.0,
-        max_tokens=256,
-        detokenize=False,
-    )
-
-    sampling_params_list = [gpt_sampling, s2mel_sampling]
+    # Keep the deploy profile's sampling policy, penalties and required stops.
+    sampling_params_list = copy.deepcopy(omni.default_sampling_params_list)
+    sampling_params_list[0].max_tokens = max(1, min(2048, max_len))
+    sampling_params_list[0].min_tokens = min(max(1, min_len), sampling_params_list[0].max_tokens)
 
     profiler_enabled = args.profiler_config is not None
     if profiler_enabled:

examples/offline_inference/yue2/end2end.py

@@ -28,7 +28,9 @@ import torch
 from vllm_omni import Omni
 from vllm_omni.model_executor.models.yue2.yue2 import (
     ABC_SAMPLING,
+    CONTEXT,
     KEY_MAX_AUDIO_FRAMES,
+    KEY_MAX_HOLD_STEPS,
     KEY_MIN_TOKENS,
     KEY_PENALTY_WINDOW,
     KEY_PHASE,
@@ -41,10 +43,12 @@ from vllm_omni.model_executor.models.yue2.yue2 import (
     KEY_TOP_P,
     SEMANTIC_SAMPLING,
     STOP_TOKEN_IDS,
+    SYNTHESIS_HOLD_STEPS,
 )
 from vllm_omni.tokenizers.yue2_prompt import (
     abc_ids_from_generated,
     abc_prefix_ids,
+    semantic_frames,
     semantic_prefix_ids,
 )
 from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
@@ -71,7 +75,13 @@ def sampling_params(engine, *, phase, seed, max_frames, prompt_ids):
         # conditioning prefix from the scheduled tokens alone.
         KEY_PREFIX_IDS: list(prompt_ids),
     }
-    params.max_tokens = preset["max_tokens"] if phase == "abc" else max_frames + 1
+    if phase == "abc":
+        params.max_tokens = preset["max_tokens"]
+    else:
+        # The song is synthesized on a side stream while the request is held;
+        # leave max_tokens headroom for the hold tokens.
+        params.max_tokens = min(max_frames + 1 + SYNTHESIS_HOLD_STEPS, CONTEXT - len(prompt_ids))
+        params.extra_args[KEY_MAX_HOLD_STEPS] = max(0, params.max_tokens - max_frames - 1)
     params.stop_token_ids = list(STOP_TOKEN_IDS)
     params.detokenize = False
     return params
@@ -144,6 +154,9 @@ def main() -> None:
         payload = {
             "prompt_token_ids": prompt_ids,
             "generated_token_ids": generated_ids,
+            # Codec frames only: without the end token and the hold tokens
+            # the request emits while its song is synthesized.
+            "semantic_frames": semantic_frames(generated_ids),
             "seed": args.seed,
             "cot": args.cot,
             "max_frames": args.max_frames,

examples/online_serving/text_to_speech/cosyvoice3/run_server.sh

@@ -1,11 +1,15 @@
 #!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Launch vLLM-Omni server for CosyVoice3 TTS
 #
 # Usage:
 #   ./run_server.sh
 #   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
 #
-# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
+# Streaming (async-chunk) is on by default. H200 selects the optimized profile;
+# other devices use vllm_omni/deploy/cosyvoice3.yaml.
 # Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
 
 set -e

recipes/m-a-p/YuE2-3B.md

@@ -40,9 +40,11 @@ user-downloaded weights and bundles none.
 | Output | 48 kHz stereo, whole song delivered when the semantic request finishes |
 
 Generation is one or two requests on a single AR stage: the abc phase writes
-the score, the semantic phase generates codec frames and, on its last step,
-solves the ODE and decodes the whole song in-engine (weights load once; the
-VAE loads at startup from `$YUE2_VAE` or the hub id).
+the score, the semantic phase generates codec frames, then queues the ODE and VAE on a
+high-priority CUDA stream. Other requests keep decoding while the finishing
+request emits HOLD tokens (up to 4096 steps, within its token/context budget).
+The scheduler is unchanged. Weights load once; the VAE loads at startup from
+`$YUE2_VAE` or the hub id.
 
 ## Running
 
@@ -67,15 +69,50 @@ ending.
 
 ## Notes
 
-- **Memory (RTX 4090, 24 GB, `gpu_memory_utilization: 0.70`):** the engine
-  reserves ~15.8 GiB after startup; with 4 concurrent requests pinned to the
-  9000-frame cap (worst case), peak usage during the terminal NAR/VAE
-  finishing pass reaches ~20.9 GiB, leaving ~3 GiB of headroom. At 0.85 the
-  same workload OOMs inside the finishing pass, so the deploy yaml pins 0.70.
-- **Known limitations:** 4 concurrent full-length requests are the verified
-  shape (`max_num_seqs: 4`); CUDA graph capture is a
-  follow-up; the abc phase runs eagerly
-  after prefill (its tokens are the product, not audio).
+The lifecycle e2e regression forces KV-cache preemption with chunked and
+unchunked recompute, mixes ABC and semantic sampling, and cancels a request
+after NAR work has started. It checks history alignment, released synthesis
+buffers, and successful generation after cancellation. Run with one H100/H200
+and cached YuE2-3B and YuE2-Vae weights:
+
+```bash
+export YUE2_MODEL_DIR=/path/to/YuE2-3B
+export YUE2_VAE=/path/to/YuE2-Vae
+CUDA_VISIBLE_DEVICES=0 python -m pytest tests/e2e/online_serving/test_yue2.py \
+    -k preemption_and_synthesis_abort -m 'slow and tts' --run-level full_model -q -s
+```
+
+The test retains the module's weekly TTS routing; it is not a per-PR CI gate.
+
+- **Memory:** `gpu_memory_utilization` budgets the vLLM engine; NAR K/V,
+  acoustic graph buffers and VAE activations also need room during synthesis.
+  The 0.70 / 4-slot defaults require validation on the target card. The earlier
+  RTX 4090 measurement (15.8 GiB after startup, 20.9 GiB peak at 9000 frames)
+  predates the async/compiled implementation; it does not validate this
+  version or long ABC prefixes on a 4090. Reduce the fraction if the target
+  workload needs more synthesis memory, especially for long ABC prefixes and
+  the 9000-frame cap. Releasing completed chunks bounds live NAR buffers,
+  while the shared graph allocator retains reserved storage for reuse.
+- **CUDA graphs:** the default deploy keeps the AR backbone eager. Both ABC
+  and semantic sampling use prewarmed power-of-two graph buckets. Custom
+  sampler captures have a bounded cache and fall back to eager sampling when
+  it fills. The NAR velocity pass uses a chunk graph on the serialized synthesis
+  stream. Chunk graphs share one allocator pool; completed chunks release
+  their engines and buffers after their own events. FA3 runs on Hopper; other
+  CUDA cards use SDPA, which is also warmed at startup.
+- **Serving:** run `vllm serve m-a-p/YuE2-3B --omni` with the default
+  `yue2.yaml`, which uses eager AR with 4 slots. Validate memory capacity
+  for the target card and workload.
+  For the tested single-H200 graph settings, reuse this deploy with
+  `--stage-overrides '{"0":{"max_num_seqs":32,"gpu_memory_utilization":0.5,"enforce_eager":false,"compilation_config":{"cudagraph_mode":"FULL_AND_PIECEWISE"}}}'`.
+  Decode uses FULL and mixed prefill uses PIECEWISE; VAE and the Python
+  synthesis queue remain outside these graphs.
+  Audio is delivered as a whole song; time to first audio equals completion
+  latency. Aborting a running song stops future work units and waits only for
+  the at-most-two submitted units before releasing its buffers.
+- **Preemption:** before drawing for a resumed row, model-owned sampling
+  reconciles the codec history and terminal flags with the engine's retained
+  sequence. Discarded draws cannot advance its penalty window or frame budget.
 - Audio is not expected to match the upstream torch reference bit-for-bit
   (fused vs eager kernels). Measured against the upstream torch reference
   (bd90e4c, same style/lyrics/seed 831001, cot=off, 200 frames): prompt token

tests/config/test_config_factory.py

@@ -571,6 +571,56 @@ class TestPipelineDiscovery:
         assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
         assert "qwen3_tts" in OMNI_PIPELINES
 
+    @pytest.mark.parametrize(
+        "cuda,major,memory_gib,mps,optimized",
+        [
+            (False, None, 0, True, False),
+            (True, None, 141, True, False),
+            (True, 8, 141, True, False),
+            (True, 9, 80, True, False),
+            (True, 10, 180, True, False),
+            (True, 9, 141, False, False),
+            (True, 9, 141, True, True),
+        ],
+    )
+    def test_cosyvoice3_device_default_and_explicit_overrides(
+        self, monkeypatch, cuda, major, memory_gib, mps, optimized
+    ):
+        from vllm.platforms import current_platform
+        from vllm.platforms.interface import DeviceCapability
+
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
+        monkeypatch.setattr(current_omni_platform, "device_name", "cuda" if cuda else "cpu")
+        capability = DeviceCapability(major, 0) if major is not None else None
+        monkeypatch.setattr(current_platform, "get_device_capability", lambda: capability)
+        monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: memory_gib * 1024**3)
+        monkeypatch.setattr(
+            "vllm_omni.model_executor.models.cosyvoice3.pipeline.shutil.which",
+            lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps else None,
+        )
+        pipeline = resolve_pipeline_config("cosyvoice3")
+        assert pipeline is not None
+        config = VllmOmniConfig.from_pipeline_config(pipeline)
+        talker, codec = config.stage_configs
+        assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == (32 if optimized else 8)
+        assert talker.runtime_config.cuda_mps == codec.runtime_config.cuda_mps == optimized
+        assert talker.model_config.hf_overrides == ({"cosyvoice3_sampling_mode": "standard"} if optimized else None)
+        assert pipeline.stages[0].sampling_constraints["stop_token_ids"] == list(range(6561, 6761))
+        if optimized:
+            assert codec.runtime_config.env["COSYVOICE3_CACHED_ISTFT"] == "1"
+            assert codec.runtime_config.env["COSYVOICE3_HIFT_GRAPH"] == "1"
+            override = VllmOmniConfig.from_pipeline_config(
+                pipeline, cli_overrides={"hf_overrides": {"cosyvoice3_sampling_mode": "ras", "seed": 7}}
+            )
+            assert override.stage_by_id(0).model_config.hf_overrides == {"cosyvoice3_sampling_mode": "ras", "seed": 7}
+        explicit = VllmOmniConfig.from_pipeline_config(
+            pipeline, deploy_config_path=get_deploy_config_path("cosyvoice3.yaml")
+        )
+        assert explicit.stage_by_id(0).scheduler_config.max_num_seqs == 8
+        assert explicit.stage_by_id(0).model_config.hf_overrides is None
+
     def test_registry_resolver_qwen3_omni_all_stages(self):
         """Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
         pipeline = resolve_pipeline_config(

tests/config/test_config_merge.py

@@ -35,3 +35,161 @@ def test_recursive_deploy_inheritance(tmp_path):
     merged = resolve_deploy_yaml(leaf)
     assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}}
     assert merged["async_chunk"] is False
+
+
+def test_single_stage_streaming_capability_uses_central_async_resolver():
+    from vllm_omni.config.stage_config import DeployConfig, _resolve_pipeline_async_chunk_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    assert _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig())
+    with pytest.raises(ValueError, match="requires async_chunk=True"):
+        _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig(async_chunk=False))
+
+
+@pytest.mark.parametrize(
+    "pipeline_key,profile,stage_count",
+    [
+        ("qwen3_tts", "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml", 2),
+        ("qwen3_tts_fused", "qwen3_tts_fused_single_gpu.yaml", 1),
+    ],
+)
+def test_qwen3_tts_deployment_selection_preserves_output_and_transport_contract(pipeline_key, profile, stage_count):
+    from tests.helpers.stage_config import get_deploy_config_path
+    from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+    from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+    pipeline = resolve_pipeline_config(pipeline_key)
+    assert pipeline is not None
+    deploy = load_deploy_config(get_deploy_config_path(profile))
+    stages = merge_pipeline_deploy(pipeline, deploy)
+    assert len(stages) == stage_count
+    talker = stages[0]
+    if stage_count == 1:
+        assert talker.final_output and talker.final_output_type == "audio"
+        assert talker.yaml_engine_args["engine_output_type"] == "audio"
+        assert "custom_process_next_stage_input_func" not in talker.yaml_engine_args
+        assert not talker.yaml_extras.get("output_connectors")
+        assert talker.yaml_engine_args["supports_running_prefix_cache_reset"] is False
+        assert talker.yaml_engine_args["additional_config"]["ref_code_context_frames"] == 72
+    else:
+        assert not talker.final_output
+        assert talker.yaml_engine_args["engine_output_type"] == "latent"
+        assert talker.yaml_engine_args["custom_process_next_stage_input_func"].endswith("talker2code2wav_async_chunk")
+        assert talker.yaml_extras["output_connectors"]
+        assert stages[1].input_sources == [0] and stages[1].final_output_type == "audio"
+        assert talker.yaml_engine_args.get("supports_running_prefix_cache_reset", True)
+
+
+def test_single_stage_running_reset_capability_cannot_be_overridden():
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    deploy = DeployConfig(
+        stages=[StageDeployConfig(stage_id=0, engine_extras={"supports_running_prefix_cache_reset": True})]
+    )
+    legacy = merge_pipeline_deploy(QWEN3_TTS_FUSED_PIPELINE, deploy)[0]
+    legacy.runtime_overrides["supports_running_prefix_cache_reset"] = True
+    assert legacy.to_omegaconf().engine_args.supports_running_prefix_cache_reset is False
+    with pytest.raises(ValueError, match="no structured config owner: supports_running_prefix_cache_reset"):
+        VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE, user_deploy_config=deploy)
+    structured = VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE)
+    assert structured.stage_by_id(0).model_config.supports_running_prefix_cache_reset is False
+
+
+@pytest.fixture
+def single_stage_tts_config(mocker):
+    from vllm.config import VllmConfig
+
+    config = mocker.Mock(spec=VllmConfig)
+    config.additional_config = {"talker_stream_decode": True}
+    config.model_config = mocker.Mock(
+        stage_connector_config={},
+        engine_output_type="audio",
+        final_output=True,
+        supports_running_prefix_cache_reset=False,
+        use_v2_model_runner=True,
+        async_chunk=True,
+    )
+    config.parallel_config = mocker.Mock(
+        tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend="uni"
+    )
+    config.cache_config = mocker.Mock(enable_prefix_caching=False)
+    config.device_config = mocker.Mock(device="cuda")
+    mocker.patch(
+        "vllm_omni.model_executor.models.qwen3_tts.stream_decode.current_omni_platform.is_cuda", return_value=True
+    )
+    return config
+
+
+def test_single_stage_tts_rejects_unsupported_decoder_configuration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.cache_config.enable_prefix_caching = True
+    with pytest.raises(ValueError, match="enable_prefix_caching=False"):
+        talker_stream_decode_enabled(config)
+    config.cache_config.enable_prefix_caching = False
+    assert talker_stream_decode_enabled(config)
+
+
+@pytest.mark.parametrize("option_location", ["additional_config", "connector"])
+@pytest.mark.parametrize("output_type", ["latent", "audio"])
+def test_two_stage_talker_rejects_single_stage_pcm_option(single_stage_tts_config, option_location, output_type):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = output_type
+    config.model_config.final_output = False
+    if option_location == "connector":
+        config.model_config.stage_connector_config = {"extra": config.additional_config}
+        config.additional_config = {}
+    with pytest.raises(ValueError, match="final audio-output Talker"):
+        talker_stream_decode_enabled(config)
+
+
+def test_two_stage_first_audio_remains_independent(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.first_audio import talker_first_audio_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = "latent"
+    config.additional_config = {}
+    config.model_config.stage_connector_config = {"extra": {"talker_first_audio": True}}
+    assert not talker_stream_decode_enabled(config)
+    assert talker_first_audio_enabled(config)
+    config.model_config.stage_connector_config["extra"]["talker_stream_first_audio"] = True
+    with pytest.raises(ValueError, match="requires the single-stage"):
+        talker_stream_decode_enabled(config)
+
+
+def test_single_stage_pcm_requires_cuda_device_and_own_first_audio_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.device_config.device = "cpu"
+    with pytest.raises(ValueError, match="requires CUDA"):
+        talker_stream_decode_enabled(config)
+    config.device_config.device = "cuda"
+    config.additional_config["talker_first_audio"] = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        talker_stream_decode_enabled(config)
+    config.additional_config.pop("talker_first_audio")
+    config.additional_config["talker_stream_first_audio"] = True
+    assert talker_stream_decode_enabled(config)
+
+
+def test_single_stage_audio_talker_requires_pcm_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.additional_config = {}
+    with pytest.raises(ValueError, match="requires talker_stream_decode=True"):
+        talker_stream_decode_enabled(single_stage_tts_config)
+
+
+def test_single_stage_pcm_requires_reset_capability_declaration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.model_config.supports_running_prefix_cache_reset = True
+    with pytest.raises(ValueError, match="must disable running prefix-cache reset"):
+        talker_stream_decode_enabled(single_stage_tts_config)

tests/config/test_environment_variables.py

@@ -167,9 +167,9 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 32,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 75,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 77,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
-        EnvironmentVariableCategory.INTERNAL: 6,
+        EnvironmentVariableCategory.INTERNAL: 7,
     }
 
     disposition_counts = Counter(
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
     )
     assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 46,
+        ModelEnvironmentVariableDisposition.PROMOTE: 48,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

tests/config/test_omni_config.py

@@ -1046,6 +1046,7 @@ def test_sub_config_fields_match_structured_scopes():
         "interleave_mm_strings",
         "media_io_kwargs",
         "final_output",
+        "supports_running_prefix_cache_reset",
         "active_stream_window",
         "session_mode",
         "duplex_max_sessions",

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -1,10 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 
 import pytest
 from vllm import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.request import RequestStatus
 
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
 from vllm_omni.engine.serialization import deserialize_additional_information
 
@@ -25,7 +31,7 @@ def _make_engine_request() -> EngineCoreRequest:
     )
 
 
-def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
+def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
     tagged = apply_omni_final_stage_metadata(
         _make_engine_request(),
         final_stage_id=0,
@@ -37,21 +43,182 @@ def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
         request_id="req",
         additional_information=tagged.additional_information,
     )
-    result = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request)
     metadata = deserialize_additional_information(tagged.additional_information)
-    return result, metadata
+    return omits_kv, omits_chunk, metadata
 
 
 def test_stage_zero_request_omits_kv_transfer():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False)
 
-    assert omits_transfer
+    assert omits_kv
+    assert omits_chunk
     assert "omni_force_kv_transfer" not in metadata
 
 
 def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True)
 
-    assert not omits_transfer
+    assert not omits_kv
+    assert omits_chunk
     assert metadata["omni_final_stage_id"] == 0
     assert metadata["omni_force_kv_transfer"] is True
+
+
+class _ChunkRequest(SimpleNamespace):
+    def __hash__(self):
+        return hash(self.request_id)
+
+    def __eq__(self, other):
+        return isinstance(other, _ChunkRequest) and other.request_id == self.request_id
+
+
+def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False):
+    tagged = apply_omni_final_stage_metadata(
+        _make_engine_request(),
+        final_stage_id=final_stage_id,
+        force_kv_transfer=force_kv_transfer,
+    )
+    request = _ChunkRequest(
+        request_id="req-shm-leak",
+        additional_information=tagged.additional_information,
+        num_in_flight_tokens=1,
+        num_stale_output_tokens=0,
+        sampling_params=SimpleNamespace(num_logprobs=None),
+        has_encoder_inputs=False,
+        pooling_params=None,
+        status=RequestStatus.RUNNING,
+        resumable=False,
+        num_output_placeholders=0,
+        spec_token_ids=[],
+        _output_token_ids=[],
+        num_computed_tokens=1,
+        stop_reason=None,
+        num_nans_in_logits=None,
+        client_index=0,
+        trace_headers=None,
+    )
+    request.is_finished = lambda: False
+    request.get_finished_reason = lambda: None
+    request.take_prefill_stats = lambda: None
+    request.take_events = lambda: None
+    return request
+
+
+def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None):
+    adapter = mocker.MagicMock()
+    adapter._confirmed_num_computed_tokens.return_value = 1
+
+    sched = mocker.MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._update_request_with_output.return_value = ([42], True)
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._handle_stopped_request.return_value = True
+
+    def _free_request(req, delay_free_blocks=False):
+        sched._omits_kv_transfer_cache.pop(req.request_id, None)
+        return None, None
+
+    sched._free_request.side_effect = _free_request
+    sched.chunk_transfer_adapter = adapter
+    sched.running = [request]
+    sched.waiting_for_transfer_free = set()
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.pending_stop_after_extraction = set()
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
+
+    scheduler_output = mocker.MagicMock()
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = mocker.MagicMock()
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    model_runner_output.inter_stage_outputs = inter_stage_outputs
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+    return adapter
+
+
+@pytest.mark.parametrize(
+    ("final_stage_id", "force_kv_transfer", "expect_save"),
+    [
+        pytest.param(0, False, False, id="stage-zero-final"),
+        pytest.param(1, False, True, id="downstream-final"),
+    ],
+)
+def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save):
+    request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer)
+    adapter = _run_finished_save_step(mocker, request)
+
+    assert adapter.save_async.called is expect_save
+
+
+def test_stage_zero_final_skips_inter_stage_output(mocker):
+    warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning")
+    request = _make_chunk_request(0)
+    adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}])
+
+    adapter.save_async.assert_not_called()
+    assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list)
+
+
+def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    scheduler.skipped_waiting = []
+    scheduler.log_stats = False
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    session = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+        status=RequestStatus.WAITING,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert "req" in scheduler._omits_kv_transfer_cache
+    update = SimpleNamespace(
+        additional_information={"omni_final_stage_id": 1},
+        model_intermediate_buffer=None,
+        arrival_time=0.0,
+        sampling_params=None,
+    )
+
+    OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update)
+
+    assert scheduler._omits_kv_transfer_cache == {}
+    assert session.additional_information == {"omni_final_stage_id": 1}
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert not scheduler._request_omits_kv_transfer_to_next_stage(session)
+
+
+def test_omits_cache_misses_when_payload_object_is_replaced():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    request = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(request)
+    request.additional_information = {"omni_final_stage_id": 1}
+
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(request)

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -40,6 +40,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -103,10 +104,11 @@ def _make_drain_sched(session: Request) -> MagicMock:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
     return sched
 
 
-def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
+def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool:
     """Feed one model-runner frame through update_from_output.
 
     Returns True when the frame's tokens were delivered (the append ran),
@@ -114,7 +116,7 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     """
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled}
-    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {}
     scheduler_output.num_invalid_spec_tokens = 0
 
     model_runner_output = MagicMock(spec=ModelRunnerOutput)
@@ -134,6 +136,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     return sched._update_request_with_output.called
 
 
+def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None:
+    from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+    session = _make_session()
+    session.num_computed_tokens = 6
+    session.num_in_flight_tokens = 1
+    session.num_output_placeholders = 1
+    sched = _make_drain_sched(session)
+    sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id}
+    sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id}
+    sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2}
+
+    OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched)
+    assert session.num_stale_output_tokens == 1 and session.drop_stale_output
+    assert session.num_output_placeholders == 0
+    sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session)
+    assert _run_step(sched, session, num_scheduled=1, token=42) is False
+    assert session.num_stale_output_tokens == 0
+    assert _run_step(sched, session, num_scheduled=1, token=43) is True
+
+
 def test_exact_drain_delivers_new_segment_frame() -> None:
     """One in-flight decode frame at replacement: its late output drains the
     counter exactly and the new segment's first frame is delivered."""
@@ -225,3 +248,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None:
     assert _run_step(sched, session, num_scheduled=5, token=42) is False  # late prefill chunk dropped
     assert session.num_stale_output_tokens == 0
     assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives
+
+
+@pytest.mark.parametrize("drop", [False, True])
+def test_preemption_stale_output_follows_upstream_delivery_policy(drop):
+    request = _make_session()
+    request.status = RequestStatus.PREEMPTED
+    request.num_stale_output_tokens = 1
+    request.num_in_flight_tokens = 1
+    request.drop_stale_output = drop
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop)
+    assert request.num_stale_output_tokens == 0
+    if not drop:
+        sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True)
+
+
+def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters():
+    request = _make_session()
+    request.num_stale_output_tokens = 2
+    request.num_in_flight_tokens = 2
+    request.num_computed_tokens = 3
+    request.num_output_placeholders = 2
+    request.drop_stale_output = False
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41])
+    assert request.num_computed_tokens == 3
+    assert request.num_output_placeholders == 2
+
+
+@pytest.mark.parametrize(
+    "arch,supports_reset,running,reset_running,blocked",
+    [
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False),
+        ("AnotherStatefulAudioModel", False, True, True, True),
+        ("AnotherStatefulAudioModel", True, True, True, False),
+    ],
+)
+def test_reset_running_streaming_codec_is_rejected_before_preemption(
+    mocker, arch, supports_reset, running, reset_running, blocked
+):
+    from vllm.config import VllmConfig
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    from vllm_omni.config.model import OmniModelConfig
+
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.vllm_config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(
+            spec=OmniModelConfig,
+            model_arch=arch,
+            supports_running_prefix_cache_reset=supports_reset,
+        ),
+    )
+    request = _make_session()
+    request.num_in_flight_tokens = 1
+    request.num_output_placeholders = 1
+    scheduler.running = [request] if running else []
+
+    def reset_cache(_self, reset_running_requests=False, reset_connector=False):
+        assert reset_running_requests is reset_running and reset_connector is True
+        return True
+
+    reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache)
+
+    assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked)
+    if blocked:
+        reset.assert_not_called()
+        assert scheduler.running == [request]
+        assert request.status == RequestStatus.RUNNING
+        assert request.num_in_flight_tokens == request.num_output_placeholders == 1
+        assert not request.drop_stale_output
+    else:
+        reset.assert_called_once()

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -26,6 +26,7 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARS
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -175,6 +176,7 @@ def _run_resumable_segment_stop(
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -312,6 +314,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -400,6 +403,7 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from types import SimpleNamespace
@@ -76,7 +79,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.input_coordinator = coordinator
     scheduler.chunk_transfer_adapter = None
-    scheduler._omits_kv_transfer_cache = {"req-free": True}
+    scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)}
     scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
     scheduler.finished_req_ids = set()
     scheduler.finished_req_ids_dict = None

tests/diffusion/ar_diffusion/test_capability_runner.py

@@ -376,10 +376,17 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
 
 def test_budget_reduced_capacity_drives_runner_lru():
     pipeline = CapablePipeline(tiny_spec(capacity=2))
-    # One tiny LingBot-like session requires 128 bytes; two require 192.
+    # One tiny LingBot-like session requires 1,808 bytes; two require 2,592.
+    #
+    # These read 128 and 192 while the paging unit was the frame. tiny_spec
+    # declares one token per frame, so it used to page a single token at a
+    # time -- a block size no attention kernel accepts. Paging at 16 makes
+    # each page sixteen times larger. The block *counts* are unchanged, and
+    # so is what this test is about: a budget that fits fewer sessions than
+    # the pipeline asked for must drive the runner's LRU.
     runner = make_runner(
         pipeline,
-        available_bytes=128,
+        available_bytes=2048,
         gpu_memory_fraction=1.0,
     )
     kv = runner.kv_cache
@@ -665,3 +672,106 @@ def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch):
     monkeypatch.setattr(runner, "execute_model", fail_if_called)
     runner._warmup_ar_rollout()
     assert execute.called is False
+
+
+# ── paging unit vs eviction unit ────────────────────────────────────────────
+
+
+def test_a_frame_that_is_already_a_legal_block_keeps_its_paging():
+    """Every resolution that runs today must page exactly as it does today."""
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for tokens_per_frame in (16, 320, 1440, 4096):
+        assert tokens_per_frame % 16 == 0
+        assert paging_block_size(tokens_per_frame) == tokens_per_frame
+
+
+def test_a_frame_the_kernel_would_reject_pages_at_the_finest_legal_unit():
+    """832x480 is the checkpoint's own default and the kernel rejects it.
+
+    (480/16) x (832/16) = 1560 tokens per frame, and 1560 % 16 == 8, so a
+    frame-sized block is not something FlashAttention's paged kernel accepts.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 16 == 8
+    assert paging_block_size(1560) == 16
+    # And whatever it returns is always something the kernel takes.
+    for tokens_per_frame in range(1, 200):
+        assert paging_block_size(tokens_per_frame) % 16 == 0
+
+
+# ── the kernel's constraint is data, not a constant ─────────────────────────
+
+
+def _multiple_of(base):
+    from vllm.v1.attention.backend import MultipleOf
+
+    return MultipleOf(base)
+
+
+def test_the_dispatching_module_answers_in_vllms_own_vocabulary():
+    """AR-Diffusion calls flash_attn itself, so no backend can be asked.
+
+    The answer has to come from the module that picks the kernel, and it has to
+    have the shape vLLM uses, so the policy that consumes it does not care
+    where it came from.
+    """
+    from vllm.v1.attention.backend import MultipleOf
+
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes
+
+    advertised = supported_kernel_block_sizes()
+    assert advertised, "a kernel that accepts nothing cannot be paged for"
+    for entry in advertised:
+        assert isinstance(entry, int | MultipleOf)
+
+
+def test_a_kernel_that_takes_one_fixed_size_gets_that_size():
+    """hpc_attn advertises [64] -- not a multiple, a single legal value.
+
+    A frame of 1560 tokens is not 64, and 1560 % 64 == 24, so the frame cannot
+    be the block and the only thing left to page at is 64. Assuming multiples
+    of 16 here would hand the kernel a block size it rejects.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 64 == 24
+    assert paging_block_size(1560, [64]) == 64
+    assert paging_block_size(64, [64]) == 64
+    assert paging_block_size(128, [64]) == 64  # a legal multiple is still not the advertised size
+
+
+def test_a_kernel_with_large_pages_still_keeps_a_whole_frame_when_it_fits():
+    """FlashInfer offers 128-and-up pages on Blackwell and not on Hopper.
+
+    The same checkpoint on the two cards must therefore page differently, which
+    is exactly what a constant cannot express.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    hopper = [_multiple_of(16)]
+    blackwell = [_multiple_of(16), 128, 256]
+    assert paging_block_size(1440, hopper) == 1440  # a legal multiple of 16
+    assert paging_block_size(1440, blackwell) == 1440
+    assert paging_block_size(1560, hopper) == 16
+    assert paging_block_size(1560, blackwell) == 16  # finest legal unit, not the largest
+
+
+def test_whatever_comes_back_is_always_something_the_kernel_accepts():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for advertised, check in (
+        ([_multiple_of(16)], lambda n: n % 16 == 0),
+        ([64], lambda n: n == 64),
+        ([_multiple_of(32), 48], lambda n: n % 32 == 0 or n == 48),
+    ):
+        for tokens_per_frame in range(1, 300):
+            assert check(paging_block_size(tokens_per_frame, advertised))
+
+
+def test_a_kernel_that_advertises_nothing_is_an_error_not_a_guess():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    with pytest.raises(ValueError, match="no legal block size"):
+        paging_block_size(1560, [])

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -625,6 +625,74 @@ def test_non_contiguous_branch_indices_rejected():
         )
 
 
+# --- eviction against the block table ---------------------------------------
+
+
+class _RecordingManager(ChunkWindowManager):
+    """Capture the block range eviction asks for, without a real block pool.
+
+    ``remove_skipped_blocks`` reads three attributes and calls one method. A
+    pool would add nothing: the bug under test is the conversion from the
+    spec's units into the block table's, and that is entirely visible here.
+    """
+
+    def __init__(self, spec):
+        self.kv_cache_spec = spec
+        self.sliding_window = spec.sliding_window
+        self.block_size = spec.block_size
+        self.freed: tuple[int, int] | None = None
+
+    def _remove_blocks_in_range(self, request_id, first_block, last_block):
+        self.freed = (first_block, last_block)
+
+
+def test_sink_is_preserved_when_a_frame_spans_many_blocks():
+    # The shipped geometry: 832x480 is 1560 tokens per frame, 1560 % 16 == 8 is
+    # not a legal block size, so paging falls back to 16 and a frame spans 98
+    # blocks. Every other eviction test uses chunk_size == BLOCK, where a frame
+    # is a block and a sink counted in frames indexes the block table correctly
+    # by accident -- which is why this went unnoticed.
+    chunk_size, sink_chunks, window_chunks = 1560, 9, 9
+    spec = ChunkWindowSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+    )
+    manager = _RecordingManager(spec)
+
+    # Far enough past sink + window that eviction has definitely started.
+    manager.remove_skipped_blocks("req", total_computed_tokens=30 * chunk_size)
+    assert manager.freed is not None
+    first_freed, last_freed = manager.freed
+
+    # The assertion is on the token boundary rather than a block count, so it
+    # stays meaningful if the paging unit changes again.
+    sink_tokens = sink_chunks * chunk_size
+    assert first_freed * spec.block_size >= sink_tokens, (
+        f"eviction starts at token {first_freed * spec.block_size}, inside the {sink_tokens}-token sink"
+    )
+    # And it must not over-protect: at most one straddling block beyond the sink.
+    assert (first_freed - 1) * spec.block_size < sink_tokens
+    # Only wholly-skipped blocks are released.
+    skipped = manager.get_num_skipped_tokens(30 * chunk_size)
+    assert last_freed * spec.block_size <= sink_tokens + skipped
+
+
+def test_sink_conversion_is_identity_when_a_frame_is_a_block():
+    # Every geometry that paged one frame per block must be untouched.
+    spec = make_spec(chunk_size=BLOCK, window_chunks=2, sink_chunks=1)
+    manager = _RecordingManager(spec)
+    manager.remove_skipped_blocks("req", total_computed_tokens=8 * BLOCK)
+    assert manager.freed is not None
+    first_freed, _ = manager.freed
+    assert first_freed == spec.sink_chunks
+
+
 def test_staged_window_reuse_matches_a_full_gather():
     """Reusing the staged history must be byte-identical to re-gathering the whole window.
 

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -23,6 +23,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache import paged_attention as paged_attention_module
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ChunkWindowManager
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -36,10 +37,26 @@ NEG = "negative"
 
 
 def make_state(
-    *, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False
+    *,
+    num_layers=1,
+    window_chunks=2,
+    chunk_size=BLOCK,
+    sink_chunks=0,
+    reset_at_boundary=False,
+    dtype=torch.float32,
+    device=torch.device("cpu"),
+    reuse_history_staging=False,
 ):
+    """Build a cache. ``chunk_size`` defaults to the block size, but the two are
+    independent -- the shipped 832x480 gives 1560 tokens per frame against
+    16-token blocks, so a frame is 97.5 blocks."""
     cfg = ARDiffusionKVConfig(
-        enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
+        enable=True,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+        reuse_history_staging=reuse_history_staging,
     )
     kv = ARDiffusionKVCache(
         cfg,
@@ -109,11 +126,13 @@ def _commit_video_span(
     n_chunks: int,
     dtype: torch.dtype,
     device: torch.device,
+    chunk_size: int = BLOCK,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    span = n_chunks * chunk_size
+    ctx = st.get_kv_caches(kv_branch, seq_len=span, commit_current=True)[0].forward_ctx
     ctx.ensure_video_slots(device)
-    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    k = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
     kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
     kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
     st.commit_paged_context(kv_branch)
@@ -547,6 +566,644 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypa
         torch._dynamo.reset()
 
 
+# ── a frame that is not a whole number of blocks ────────────────────────────
+
+# 24 tokens per chunk against 16-token blocks leaves 8 over, the same remainder
+# the shipped 832x480 produces (1560 % 16 == 8). Every fixture above holds
+# chunk_size == BLOCK, so nothing there can reach this path.
+RAGGED_CHUNK = 24
+
+
+@pytest.mark.parametrize("commit_current", [False, True])
+@pytest.mark.parametrize("action_len", [0, 3])
+def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action_len):
+    """Attention reads a sequence, not a set of blocks.
+
+    With one committed chunk of 24 tokens the history stops 8 slots into its
+    last block. The committing path writes straight after it and is fine. The
+    scratch path used to restart at slot zero of a fresh region, which left
+    those 8 slots unwritten -- and since the kernel consumes the block table as
+    one contiguous run, it read them as if they were tokens and shifted every
+    token after them by 8. Shapes stayed correct throughout, so only the values
+    showed it.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    # window_chunks=2 with one chunk of history keeps the whole sequence
+    # resident, so the window does not also round up here; that case is
+    # test_a_ragged_window_keeps_the_block_table_a_fixed_shape below.
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
+
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=RAGGED_CHUNK
+    )
+
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    assert ctx.start_offset == RAGGED_CHUNK % BLOCK == 8
+
+    current_k = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, RAGGED_CHUNK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len, query_len=query.shape[1], device=device
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    new_k = torch.cat([history_k, current_k], dim=1)
+    new_v = torch.cat([history_v, current_v], dim=1)
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+
+def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
+    """The write targets themselves must leave no hole.
+
+    This is the property the values above depend on, asserted directly so a
+    regression names the cause rather than showing a numeric mismatch.
+    """
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
+    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=torch.float32, device=device, chunk_size=RAGGED_CHUNK)
+
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+
+    history_tail_block = ctx.history_block_ids[-1]
+    # The chunk starts inside the history's own last block, not at a fresh one.
+    assert ctx.current_video_block_ids[0] == history_tail_block
+    first_slot = int(ctx.current_video_slot_mapping[0])
+    assert first_slot == history_tail_block * BLOCK + 8
+
+    # And the action region must not be charged for that managed block.
+    ctx.ensure_action_slots(3, device)
+    assert ctx._scratch_blocks_used == len(ctx.current_video_block_ids) - 1
+    assert ctx.action_scratch_block_ids[0] not in ctx.current_video_block_ids
+
+
+def test_a_ragged_window_keeps_the_block_table_a_fixed_shape():
+    """The table's shape must not track how full the window is.
+
+    With a sink the visible window is 2*24 + 24 = 72 tokens, which is 4.5
+    blocks. Deriving the width by flooring that would understate the capacity,
+    the max() against the live block count would take over, and the shape would
+    change as the window filled -- recompiling the graph. The same floor would
+    also let max_seq_len come out under the kv_len actually passed.
+    """
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
+    # The window is deliberately not a whole number of blocks.
+    max_video_tokens = kv.spec.sliding_window + kv.spec.sink_chunks * kv.spec.chunk_size
+    assert max_video_tokens % BLOCK != 0
+
+    widths: set[int] = set()
+    kv_lens: list[int] = []
+    for _ in range(6):
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(device)
+        block_table, _, seq_lens, _, max_seq_len = ctx.build_block_table(
+            action_len=0, query_len=RAGGED_CHUNK, device=device
+        )
+        widths.add(int(block_table.shape[1]))
+        # A bound handed to the kernel has to actually bound the sequence.
+        assert int(seq_lens[0]) <= max_seq_len
+        kv_lens.append(int(seq_lens[0]))
+        st.commit_paged_context(POS)
+
+    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
+    # The window has to have actually filled, or none of the above was tested.
+    assert max(kv_lens) > min(kv_lens)
+    assert max(kv_lens) <= ctx.max_video_blocks * BLOCK
+    # Whole blocks are read, so each of the two boundaries -- the sink's end and
+    # the recent window's start -- can contribute up to BLOCK - 1 tokens the
+    # window no longer needs.
+    assert max(kv_lens) <= max_video_tokens + 2 * (BLOCK - 1)
+
+
+def _poison_pools(kv: ARDiffusionKVCache) -> None:
+    for pool in (*kv._k_pools, *kv._v_pools):
+        pool.fill_(float("nan"))
+
+
+def _write_prepared_chunk(
+    kv: ARDiffusionKVCache, ctx, *, dtype: torch.dtype, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Write fresh K/V for a prepared chunk, poisoning the rest of its blocks first.
+
+    Blocks are handed out whole and a reused block keeps whatever it held, so a
+    slot read without having been written would not look unusual. Every slot of
+    the chunk's blocks that this chunk does not write is set to NaN first --
+    except the leading slots of a block the history still occupies -- so reading
+    one turns the attention output NaN.
+    """
+    mapping = ctx.current_video_slot_mapping
+    first_slot = int(mapping[0])
+    for index, block in enumerate(ctx.current_video_block_ids):
+        start = first_slot if index == 0 else block * BLOCK
+        kv._k_pools[0][start : (block + 1) * BLOCK] = float("nan")
+        kv._v_pools[0][start : (block + 1) * BLOCK] = float("nan")
+    k = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][mapping] = k
+    kv._v_pools[0][mapping] = v
+    return k, v
+
+
+def _window_positions(end: int, *, sink: int, window: int, written, history: int, resident: set[int]) -> list[int]:
+    """Token positions attention must read at ``end``.
+
+    Every token of the current chunk, plus every written history token in a
+    block that is still resident and holds a sink token or one of the window's
+    most recent tokens.
+    """
+    start = end - window
+    blocks = {p // BLOCK for p in range(min(sink, end))} | {p // BLOCK for p in range(max(start, 0), end)}
+    return [p for p in sorted(written) if p >= history or (p // BLOCK in blocks and p // BLOCK in resident)]
+
+
+def _model_block_indices(kv, adapter) -> dict[int, int]:
+    """Each resident block's index in model positions, by block id.
+
+    Compaction removes the evicted gap after the sink's blocks and shifts the
+    rest of the table down by it, so a table index past the sink is
+    ``compacted_tokens`` behind the position it holds.
+    """
+    block_size = kv.block_size
+    sink_blocks = -(-(kv.spec.sink_chunks * kv.spec.chunk_size) // block_size)
+    shift = adapter.compacted_tokens // block_size
+    return {
+        int(block): index if index < sink_blocks else index + shift
+        for index, block in enumerate(kv.block_table(adapter))
+        if block != kv.null_block_id
+    }
+
+
+@pytest.mark.parametrize(
+    ("sink_chunks", "window_chunks", "reset_at_boundary"),
+    [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
+)
+def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten(
+    sink_chunks, window_chunks, reset_at_boundary
+):
+    """Once the window slides, the table must still cover the window exactly.
+
+    The sink and the recent window are separate token ranges, and each can
+    straddle a block edge on its own: 24-token chunks against 16-token blocks
+    put a boundary 8 tokens past an edge on alternating chunks. Counting blocks
+    off the ends of the resident list dropped the block holding the window's
+    first tokens whenever both boundaries straddled, and the returned length
+    ran past the last written slot. Neither changes a shape, so this compares
+    values against dense attention over the tokens the window keeps -- after
+    eviction, with every unwritten slot poisoned, on both the scratch and the
+    committing forward of every chunk.
+
+    The reset case keeps only the sink across a boundary, so the recent window
+    reaches back over blocks that are no longer resident. Those must be skipped
+    by position rather than read.
+
+    Eight chunks are enough for compaction to drop the evicted gap from the
+    table, after which storage positions run behind model positions. Removing
+    a gap that is not a whole number of chunks moves where eviction snaps, and
+    in the reset case the window then keeps tokens a boundary should have
+    dropped.
+    """
+    torch.manual_seed(0)
+    device, dtype = torch.device("cpu"), torch.float32
+    kv, st = make_state(
+        device=device,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+        chunk_size=RAGGED_CHUNK,
+    )
+    _poison_pools(kv)
+    sink, window = sink_chunks * RAGGED_CHUNK, int(kv.spec.sliding_window)
+    committed: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+    evicted = compacted = False
+
+    for chunk in range(8):
+        history = int(st.adapter(POS).absolute_num_computed_tokens)
+        end = history + RAGGED_CHUNK
+        for commit_current in (False, True):
+            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+            ctx.ensure_video_slots(device)
+            k, v = _write_prepared_chunk(kv, ctx, dtype=dtype, device=device)
+            written = dict(committed)
+            written.update({history + i: (k[i], v[i]) for i in range(RAGGED_CHUNK)})
+
+            query = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+            block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+                action_len=0, query_len=RAGGED_CHUNK, device=device
+            )
+            resident = set(_model_block_indices(kv, st.adapter(POS)).values())
+            positions = _window_positions(
+                end, sink=sink, window=window, written=written, history=history, resident=resident
+            )
+            label = f"chunk {chunk}, commit_current={commit_current}"
+            assert int(seq_lens[0]) == len(positions), f"{label}: kv_len does not match the written slots it covers"
+            paged = ar_diffusion_paged_attention(
+                query,
+                kv.key_cache(0),
+                kv.value_cache(0),
+                block_table=block_table,
+                query_start_loc=query_start_loc,
+                seq_lens=seq_lens,
+                max_query_len=max_query_len,
+                max_seq_len=max_seq_len,
+                softmax_scale=HEAD_DIM**-0.5,
+                causal=False,
+            )
+            ref = _dense_attention(
+                query,
+                torch.stack([written[p][0] for p in positions]).unsqueeze(0),
+                torch.stack([written[p][1] for p in positions]).unsqueeze(0),
+            )
+            torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5, msg=lambda m, label=label: f"{label}: {m}")
+        committed = written
+        st.commit_paged_context(POS)
+        compacted = compacted or st.adapter(POS).compacted_tokens > 0
+        evicted = evicted or compacted or kv.null_block_id in kv.block_table(st.adapter(POS))
+
+    assert evicted, "the window never slid, so nothing after eviction was tested"
+    assert compacted, "the table was never compacted, so storage and model positions never diverged"
+
+
+@pytest.mark.parametrize(
+    ("sink_chunks", "window_chunks", "reset_at_boundary"),
+    [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
+)
+def test_compaction_does_not_change_which_tokens_stay_resident(
+    monkeypatch, sink_chunks, window_chunks, reset_at_boundary
+):
+    """Compaction renumbers the block table; it must not change what eviction keeps.
+
+    Storage positions run ``compacted_tokens`` behind model positions, and
+    eviction snaps to chunk boundaries counted in storage positions. Removing a
+    gap that is not a whole number of chunks moves every later snap, so eviction
+    keeps a different set of tokens -- in the reset case, tokens a boundary
+    should have dropped, which the window then reads. Attention checked against
+    what is resident cannot see that, so the oracle here is the same rollout
+    with compaction switched off.
+    """
+
+    def resident_blocks_per_tick(compact: bool) -> tuple[list[set[int]], int]:
+        with monkeypatch.context() as patch:
+            if not compact:
+                patch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
+            kv, st = make_state(
+                window_chunks=window_chunks,
+                sink_chunks=sink_chunks,
+                reset_at_boundary=reset_at_boundary,
+                chunk_size=RAGGED_CHUNK,
+            )
+            ticks = []
+            for _ in range(12):
+                ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+                ctx.ensure_video_slots(torch.device("cpu"))
+                ticks.append(set(_model_block_indices(kv, st.adapter(POS)).values()))
+                st.commit_paged_context(POS)
+            return ticks, st.adapter(POS).compacted_tokens
+
+    expected, _ = resident_blocks_per_tick(compact=False)
+    actual, compacted_tokens = resident_blocks_per_tick(compact=True)
+    assert compacted_tokens > 0, "the table was never compacted, so nothing was tested"
+    for tick, (want, got) in enumerate(zip(expected, actual)):
+        assert got == want, f"tick {tick}: compaction kept blocks {sorted(got - want)} and dropped {sorted(want - got)}"
+
+
+def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch):
+    """The staging buffers must fit the table width this module pads to.
+
+    A sink or window that is not a whole number of blocks spans one more block
+    than its token count suggests, so sizing staging from tokens left it short
+    and the first prepared forward refused to run. And because such a chunk
+    shares its first block with the history's tail, the history a probe sees is
+    not left untouched by the next probe, so it is restaged rather than reused.
+    """
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device = torch.device("cpu")
+    kv, st = make_state(
+        device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK, reuse_history_staging=True
+    )
+    capacity = int(kv.history_staging[0][0].shape[0])
+    for chunk in range(6):
+        for probe in range(2):
+            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
+            ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
+            assert int(ctx.max_seq_len) <= capacity
+            assert not ctx.reuse_history, f"chunk {chunk}, probe {probe}: a ragged chunk reused the staged history"
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
+        st.commit_paged_context(POS)
+
+
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_current):
+    """Action K/V follows the video blocks in the table, and the kernel reads it as one run.
+
+    A 24-token chunk from an empty history ends 8 slots into its second 16-token
+    block, so the run would read those 8 unwritten slots as the first action
+    tokens and never reach the last ones. Every shape stays right, which is why
+    this is refused rather than computed.
+    """
+    _, st = make_state(window_chunks=2, chunk_size=RAGGED_CHUNK)
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+    with pytest.raises(ValueError, match="partly written video block"):
+        ctx.build_block_table(action_len=3, query_len=RAGGED_CHUNK + 3, device=torch.device("cpu"))
+
+
+def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick():
+    """832x480 with the checkpoint's own sink and window, slot by slot.
+
+    A 9-frame sink and a 9-frame window are each 8 tokens past a 16-token edge
+    but 28080 tokens together -- a whole number of blocks -- so a width taken
+    from their sum was one block short on every tick once the window had
+    filled. Twelve ticks of three frames, checking every slot the kernel reads
+    against the token that slot is supposed to hold. No attention is computed:
+    at this size the bookkeeping is the whole question.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    frame, frames_per_tick = (480 // 16) * (832 // 16), 3
+    block = paging_block_size(frame)
+    seq_len = frames_per_tick * frame
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=frame, window_chunks=9, sink_chunks=9)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=1,
+        num_kv_heads=1,
+        head_size=8,
+        dtype=torch.float16,
+        block_size=block,
+        max_model_len=1 << 20,
+        available_bytes=1 << 27,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
+        session_capacity=1,
+        frames_per_block=frames_per_tick,
+        max_scratch_tokens_per_branch=seq_len,
+        device=torch.device("cpu"),
+    )
+    st = ARDiffusionKVState(kv, "s1", {POS: kv.begin_request("r-pos")}, num_layers=1)
+    sink, window = 9 * frame, int(kv.spec.sliding_window)
+    holds: dict[int, int] = {}  # slot -> token position last written there
+    widths: set[int] = set()
+    filled = compacted = False
+
+    for tick in range(12):
+        history = int(st.adapter(POS).absolute_num_computed_tokens)
+        end = history + seq_len
+        ctx = st.get_kv_caches(POS, seq_len=seq_len, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(torch.device("cpu"))
+        for offset, slot in enumerate(ctx.current_video_slot_mapping.tolist()):
+            holds[slot] = history + offset
+        table, _, seq_lens, _, _ = ctx.build_block_table(action_len=0, query_len=seq_len, device=torch.device("cpu"))
+        widths.add(int(table.shape[1]))
+        kv_len = int(seq_lens[0])
+
+        position_of_block = {b: i * block for b, i in _model_block_indices(kv, st.adapter(POS)).items()}
+        read = []
+        for b in table[0].tolist()[: -(-kv_len // block)]:
+            assert b in position_of_block, f"tick {tick}: read block {b}, which holds no token of this session"
+            read.extend((b * block + o, position_of_block[b] + o) for o in range(block))
+        read = read[:kv_len]
+        wrong = [(slot, want) for slot, want in read if holds.get(slot) != want]
+        assert not wrong, f"tick {tick}: {len(wrong)} slots read hold a token other than the one their position implies"
+
+        start = end - window
+        kept = set(range(min(sink, end))) | set(range(max(start, 0), end))
+        missed = kept - {want for _, want in read}
+        assert not missed, f"tick {tick}: {len(missed)} tokens the window keeps were not read"
+
+        filled = filled or end > sink + window
+        st.commit_paged_context(POS)
+        compacted = compacted or st.adapter(POS).compacted_tokens > 0
+
+    assert filled, "the window never filled, so the case that fails was never reached"
+    assert compacted, "the table was never compacted, so storage and model positions never diverged"
+    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
+
+
+class _CountedBlockIds:
+    """Wraps a vLLM block manager's ``get_blocks`` and counts the block ids read.
+
+    Every path to a request's block ids goes through ``get_blocks``, so a copy of
+    the whole table is counted in full even if only a few entries are indexed.
+    """
+
+    def __init__(self, manager):
+        self.reads = 0
+        self._get_blocks = manager.get_blocks
+
+    def __call__(self, request_id):
+        from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+
+        blocks = self._get_blocks(request_id)
+        return KVCacheBlocks(tuple([_CountedBlock(block, self) for block in group] for group in blocks.blocks))
+
+
+class _CountedBlock:
+    def __init__(self, block, counter: _CountedBlockIds):
+        self._block = block
+        self._counter = counter
+
+    @property
+    def block_id(self) -> int:
+        self._counter.reads += 1
+        return self._block.block_id
+
+
+def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch):
+    """However long the block table is, reading the window must not cost more.
+
+    Evicted positions stay in the table as null entries until compaction drops
+    them a whole number of chunks at a time -- up to two frames' worth of blocks
+    at 832x480 -- so the table can still be far longer than the window. Only the
+    sink's blocks and the recent window's blocks can hold a kept token, so those
+    are all the read may touch. Compaction is switched off here so the table
+    grows to many times the window, and reads are counted where the ids leave
+    vLLM's block manager.
+    """
+    monkeypatch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
+    reads: list[int] = []
+    table_lengths: list[int] = []
+    for _ in range(60):
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(device)
+        counted = _CountedBlockIds(kv.manager)
+        kv.manager.get_blocks = counted
+        try:
+            ctx.build_block_table(action_len=0, query_len=RAGGED_CHUNK, device=device)
+        finally:
+            del kv.manager.get_blocks
+        reads.append(counted.reads)
+        table_lengths.append(len(kv.block_table(st.adapter(POS))))
+        st.commit_paged_context(POS)
+
+    assert table_lengths[-1] > 10 * ctx.max_video_blocks, "the table never outgrew the window, so nothing was tested"
+    assert max(reads) <= ctx.max_video_blocks, (
+        f"read {max(reads)} table entries for a {ctx.max_video_blocks}-block window"
+    )
+
+
+def test_the_checkpoints_own_default_resolution_can_build_a_cache():
+    """832x480 is what LingBot World v2 ships as its default, and it could not run.
+
+    A frame-sized block made the resolution a kernel-compatibility question:
+    1560 tokens per frame is not a multiple of 16, so FlashAttention's paged
+    kernel rejected it and the default resolution had no realtime path at all.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+    assert tokens_per_frame % 16 == 8, "the whole point is that a frame is not a legal block"
+
+    block_size = paging_block_size(tokens_per_frame)
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=tokens_per_frame, window_chunks=2)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=1,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=block_size,
+        max_model_len=1 << 16,
+        available_bytes=1 << 28,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
+        session_capacity=1,
+        frames_per_block=2,
+        max_scratch_tokens_per_branch=block_size,
+        device=torch.device("cpu"),
+    )
+
+    assert kv.block_size % 16 == 0
+    # The eviction unit is still the frame; only the paging unit changed.
+    assert kv.spec.chunk_size == tokens_per_frame
+    assert kv.blocks_per_frame == -(-tokens_per_frame // block_size) == 98
+
+
+def test_the_shipped_resolution_geometry_matches_dense_attention():
+    """The real number, not a scaled-down stand-in: 1560 tokens per frame.
+
+    This is the case the issue reported failing at ~4e-02 on the scratch path.
+    It runs on CPU because the reference is dense attention, not a kernel.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
+    )
+
+    # The non-committing path is the one that was wrong: four of the five
+    # forwards per generated block take it.
+    ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
+
+    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=0, query_len=tokens_per_frame, device=device
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    ref = _dense_attention(
+        query,
+        torch.cat([history_k, current_k], dim=1),
+        torch.cat([history_v, current_v], dim=1),
+    )
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
+    """The shipped 832x480 geometry through FlashAttention's paged kernel.
+
+    The CPU test above proves the addressing is right against a dense
+    reference. This one proves the real kernel agrees, at the real 1560
+    tokens per frame, through the production path -- host prep followed by
+    the fused write+attend op, which is what a forward actually calls.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
+    )
+
+    layer_ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    ctx.prepare(device=device, action_len=0, query_len=tokens_per_frame)
+    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
+    assert kv.block_size == 16
+
+    inputs = layer_ctx.to_layer_inputs()
+    paged = paged_write_attn(inputs, query[0], current_k[0], current_v[0], None, None, HEAD_DIM**-0.5).unsqueeze(0)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
 @pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
 @pytest.mark.parametrize("history_chunks", [0, 1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])

tests/diffusion/attention/test_attention_capabilities.py

@@ -0,0 +1,359 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import sys
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_validate_fa4_head_dims = fa.validate_fa4_head_dims
+
+
+@pytest.fixture(autouse=True)
+def kernel_validator(monkeypatch):
+    # Capability unit tests do not require FA4 or a CUDA device. Real-kernel
+    # validation is exercised by test_flash_attn_compile.py.
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+
+def _impl(*, kernel_variant: str | None = "fa4", causal: bool = False):
+    impl = FlashAttentionImpl.__new__(FlashAttentionImpl)
+    impl._kernel_variant = kernel_variant
+    impl.causal = causal
+    return impl
+
+
+def _resolve(
+    impl,
+    *,
+    context: ExecutionContext | None = None,
+    attn_metadata: AttentionMetadata | None = None,
+    dtype=torch.bfloat16,
+):
+    tensor = torch.empty((1, 16, 8, 64), dtype=dtype)
+    return impl.resolve_execution_path(
+        context or ExecutionContext(platform="cuda"),
+        tensor,
+        tensor,
+        tensor,
+        attn_metadata,
+    )
+
+
+class _UnmigratedBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "UNMIGRATED_TEST"
+
+    @staticmethod
+    def get_impl_cls():
+        return None
+
+    @staticmethod
+    def get_metadata_cls():
+        return None
+
+    @staticmethod
+    def get_builder_cls():
+        return None
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return []
+
+
+def test_capability_result_is_tri_state():
+    assert CapabilityResult.supported().status is SupportStatus.SUPPORTED
+    assert CapabilityResult.unsupported("not available").status is SupportStatus.UNSUPPORTED
+    assert CapabilityResult.unmigrated().status is SupportStatus.UNMIGRATED
+    with pytest.raises(ValueError, match="actionable reason"):
+        CapabilityResult(SupportStatus.UNSUPPORTED)
+
+
+def test_unmigrated_backend_default_is_conservative():
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    result = _UnmigratedBackend.resolve_capabilities(context)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    requested = result.requested_support(context)
+    assert requested.status is SupportStatus.UNSUPPORTED
+    assert "no verified fullgraph declaration" in requested.reason
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize("mode", list(CompilationMode))
+def test_fullgraph_request_checks_compilation_mode(mode):
+    result = replace(_resolve(_impl()), compilation_mode=mode)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    expected = SupportStatus.UNSUPPORTED if mode is CompilationMode.EAGER_ONLY else SupportStatus.SUPPORTED
+    assert result.requested_support(context).status is expected
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.SUPPORTED
+
+
+def test_request_preserves_verified_rejection_reason():
+    result = replace(_resolve(_impl()), support=CapabilityResult.unsupported("unsupported head size"))
+    assert result.requested_support(ExecutionContext(platform="cuda", require_fullgraph=True)) == result.support
+
+
+def test_exact_cuda_fa4_dense_path_is_supported():
+    context = ExecutionContext(
+        platform="cuda",
+        require_fullgraph=True,
+    )
+    result = _resolve(_impl(), context=context)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+    assert result.compilation_mode is CompilationMode.CUSTOM_OP
+    assert result.requested_support(context).status is SupportStatus.SUPPORTED
+
+
+def test_causal_fa4_path_remains_unmigrated():
+    result = _resolve(_impl(causal=True))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_fa4_delegates_dimensions_to_kernel(monkeypatch):
+    head_dims = (80, 48)
+    validator = Mock(return_value=True)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", validator)
+    query = torch.empty((1, 16, 8, head_dims[0]), dtype=torch.bfloat16)
+    value = torch.empty((1, 16, 8, head_dims[1]), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, query, value, None)
+    assert result.support.status is SupportStatus.SUPPORTED
+    validator.assert_called_once_with(*head_dims, 8)
+
+
+def test_fa4_reports_kernel_rejection(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", Mock(side_effect=AssertionError("kernel dimension constraint")))
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "kernel dimension constraint" in result.support.reason
+    assert "another backend" in result.support.reason
+
+
+def test_fa4_missing_validator_remains_unmigrated(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: False)
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_fa4_rejects_mismatched_query_key_dimensions():
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    key = torch.empty((1, 16, 8, 32), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, key, key, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "Q and K head dimensions must match" in result.support.reason
+
+
+@pytest.mark.parametrize("arch", [90, 100, 120])
+def test_fa4_validator_adapter_uses_selected_architecture(monkeypatch, arch):
+    validator = Mock()
+    monkeypatch.setitem(
+        sys.modules,
+        "flash_attn.cute.interface",
+        SimpleNamespace(_get_device_arch=lambda: arch, _validate_head_dims=validator),
+    )
+    # Call the real adapter, not the fixture's stand-in.
+    verified = _validate_fa4_head_dims(80, 48, 8)
+    if arch == 120:
+        assert not verified
+        validator.assert_not_called()
+    else:
+        assert verified
+        validator.assert_called_once_with(80, 48, arch // 10, 8)
+
+
+def test_fa4_validator_adapter_handles_missing_private_api(monkeypatch):
+    monkeypatch.setitem(sys.modules, "flash_attn.cute.interface", SimpleNamespace())
+    assert not _validate_fa4_head_dims(64, 64, 8)
+
+
+def test_initialized_kernel_identity_overrides_caller_claim():
+    result = _resolve(
+        _impl(kernel_variant=None),
+        context=ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+        ),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_preconstruction_without_resolved_cuda_kernel_is_unmigrated():
+    result = FlashAttentionBackend.resolve_capabilities(ExecutionContext(platform="cuda"))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.path == "unverified"
+
+
+@pytest.mark.parametrize(
+    "changes",
+    [
+        {"platform": "rocm"},
+        {"paged_kv": True},
+        {"parallel_strategy": ParallelStrategy.ULYSSES},
+        {"parallel_strategy": ParallelStrategy.RING},
+        {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING},
+        {"outer_boundaries": frozenset({OuterBoundary.HSDP})},
+    ],
+)
+def test_unverified_flash_attention_variants_remain_unmigrated(changes):
+    result = _resolve(
+        _impl(),
+        context=replace(ExecutionContext(platform="cuda"), **changes),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_unverified_dtype_remains_unmigrated():
+    result = _resolve(_impl(), dtype=torch.float32)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize(
+    "attn_metadata",
+    [
+        AttentionMetadata(full_attn_spans=[[(0, 8)]]),
+        AttentionMetadata(
+            extra={
+                "cu_seqlens_q": torch.tensor([0, 16], dtype=torch.int32),
+                "cu_seqlens_k": torch.tensor([0, 16], dtype=torch.int32),
+                "max_seqlen_q": 16,
+                "max_seqlen_k": 16,
+            }
+        ),
+        AttentionMetadata(extra={"kv_cache_dtype": "fp8"}),
+    ],
+)
+def test_metadata_normalization_prevents_dense_support(attn_metadata):
+    result = _resolve(_impl(), attn_metadata=attn_metadata)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_unpublished_all_true_mask_remains_runtime_dependent():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "runtime_mask_dependent"
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_producer_published_noop_mask_resolves_dense_fa4():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+
+
+def test_fixed_shape_mask_semantic_change_invalidates_dense_resolution():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    impl = _impl()
+    dense = _resolve(impl, attn_metadata=metadata)
+    assert dense.support.status is SupportStatus.SUPPORTED
+
+    metadata.attn_mask[:, 8:] = False
+    metadata.extra["attention_mask_mode"] = "padding"
+    masked = _resolve(impl, attn_metadata=metadata)
+    assert masked.support.status is SupportStatus.UNMIGRATED
+    request = ExecutionContext(platform="cuda", require_fullgraph=True)
+    assert masked.requested_support(request).status is SupportStatus.UNSUPPORTED
+
+
+def test_piecewise_dispatch_ignores_unused_incomplete_packed_metadata(monkeypatch):
+    from vllm_omni.diffusion.attention.backends import flash_attn
+    from vllm_omni.diffusion.attention.backends.utils import fa
+
+    metadata = AttentionMetadata(full_attn_spans=[[(0, 16)]], extra={"max_seqlen_q": 16})
+    impl = _impl()
+    impl.fa_deterministic = False
+    impl.softmax_scale = 0.125
+    assert _resolve(impl, attn_metadata=metadata).support.status is SupportStatus.UNMIGRATED
+
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", lambda *args, **kwargs: query)
+    monkeypatch.setattr(flash_attn, "piecewise_attn", lambda *args, **kwargs: query)
+    assert impl.forward_cuda(query, query, query, metadata) is query
+
+
+@pytest.mark.parametrize("input_index", [1, 2])
+@pytest.mark.parametrize("mismatch", ["dtype", "device"])
+def test_fa4_rejects_incompatible_key_value(input_index, mismatch):
+    tensors = [torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) for _ in range(3)]
+    tensors[input_index] = (
+        tensors[input_index].to(dtype=torch.float16) if mismatch == "dtype" else tensors[input_index].to("meta")
+    )
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *tensors, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert f"{mismatch}s must match" in result.support.reason
+
+
+@pytest.mark.parametrize("skip_quant", [False, True])
+def test_layer_resolution_uses_effective_kv_quantization(skip_quant):
+    from vllm_omni.diffusion.attention.layer import Attention
+    from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+
+    layer = Attention.__new__(Attention)
+    torch.nn.Module.__init__(layer)
+    layer.attention = _impl(kernel_variant=None)
+    layer._hsdp_compile_boundary_enabled = False
+    layer.skip_sequence_parallel = False
+    layer._no_parallel_strategy = NoParallelAttention()
+    layer.parallel_strategy = layer._no_parallel_strategy
+    layer.use_ring = False
+    layer.paged_kv_cache_role = None
+    layer._kv_cache_dtype = "fp8"
+    layer._disable_kv_quant = False
+    layer._kv_cache_skip_steps = None
+    layer._kv_cache_skip_layers = {0} if skip_quant else None
+    layer.layer_idx = 0
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": "fp8"})
+    result = layer.resolve_execution_path(ExecutionContext(platform="npu"), query, query, query, metadata)
+    assert result.path == ("npu_dense" if skip_quant else "npu_unverified")
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert metadata.extra == {"kv_cache_dtype": "fp8"}

tests/diffusion/attention/test_attention_state_lifetime.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Stateful contract example, independent of an optional attention kernel.
+
+The prepared module owns tensor state and passes it explicitly to an opaque op.
+This deliberately does not model opaque Python/C++ planning handles.
+"""
+
+import gc
+import weakref
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@torch.library.custom_op("vllm_omni_test::planned_attention", mutates_args=())
+def _planned_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_scale: torch.Tensor,
+) -> torch.Tensor:
+    return F.scaled_dot_product_attention(query * query_scale, key, value).contiguous()
+
+
+@_planned_attention.register_fake
+def _planned_attention_fake(query, key, value, query_scale):
+    return query.new_empty((*query.shape[:-1], value.shape[-1]))
+
+
+class _PreparedAttention(torch.nn.Module):
+    """Initialization/planning runs eagerly; forward only executes tensors."""
+
+    def __init__(self, query_scale: torch.Tensor):
+        super().__init__()
+        self.register_buffer("query_scale", query_scale.detach().clone())
+
+    def forward(self, query, key, value):
+        return _planned_attention(query, key, value, self.query_scale)
+
+
+def _inputs(length):
+    generator = torch.Generator().manual_seed(length)
+    return tuple(torch.randn(2, 3, length, 8, generator=generator) for _ in range(3))
+
+
+def _reference(plan, inputs):
+    query, key, value = inputs
+    return F.scaled_dot_product_attention(query * plan.query_scale, key, value)
+
+
+def test_dynamic_replay_uses_each_instances_state_without_recompilation():
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    # Same tensor metadata, different values: state must be a graph input,
+    # never a Python instance identifier or a captured first-instance constant.
+    plans = [_PreparedAttention(torch.full((1, 3, 1, 1), scale)) for scale in (0.5, 2.0)]
+    compiled = torch.compile(
+        lambda plan, q, k, v: plan(q, k, v),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    for length in (5, 9, 13):
+        inputs = _inputs(length)
+        for plan in plans:
+            torch.testing.assert_close(compiled(plan, *inputs), _reference(plan, inputs))
+        assert not torch.allclose(_reference(plans[0], inputs), _reference(plans[1], inputs))
+    assert compile_count == 1
+
+
+def test_compiled_callable_retains_state_until_released():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    plan_ref = weakref.ref(plan)
+    state_ref = weakref.ref(plan.query_scale)
+    inputs = _inputs(5)
+    expected = _reference(plan, inputs)
+    compiled = torch.compile(plan, backend="eager", fullgraph=True, dynamic=True)
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    # Deleting the caller's reference must not invalidate subsequent replay.
+    del plan
+    gc.collect()
+    assert plan_ref() is not None
+    assert state_ref() is not None
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    del compiled
+    gc.collect()
+    assert plan_ref() is None
+    assert state_ref() is None
+
+
+def test_stateful_custom_op_schema_and_fake_implementation():
+    inputs = _inputs(5)
+    # Unequal Q/K lengths and V head dimensions exercise the fake output shape.
+    query, _, _ = inputs
+    key = torch.randn(2, 3, 7, 8)
+    value = torch.randn(2, 3, 7, 4)
+    scale = torch.full((1, 3, 1, 1), 0.5)
+    torch.library.opcheck(_planned_attention, (query, key, value, scale))
+
+
+def test_stateful_attention_compiles_with_inductor():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    compiled = torch.compile(plan, fullgraph=True, dynamic=True)
+    for length in (5, 9):
+        inputs = _inputs(length)
+        torch.testing.assert_close(compiled(*inputs), _reference(plan, inputs))

tests/diffusion/attention/test_flash_attn.py

@@ -21,6 +21,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
+from vllm_omni.diffusion.attention.capabilities import CompilationMode, ExecutionContext, SupportStatus
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -432,6 +433,35 @@ def test_varlen_masked_routing_by_role(monkeypatch):
     assert cu_seqlens_k.tolist() == [0, 2, 5]
 
 
+@pytest.mark.cpu
+def test_published_noop_mask_uses_dense_path(monkeypatch):
+    calls = []
+
+    def fake_dense_func(q, k, v, **_kwargs):
+        calls.append("dense")
+        return q
+
+    def fail_varlen(*_args, **_kwargs):
+        raise AssertionError("published no-op mask must not use varlen attention")
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_dense_func)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fail_varlen)
+
+    impl = FlashAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 4, 2, 8)
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 4), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert torch.equal(output, query)
+    assert calls == ["dense"]
+
+
 @pytest.mark.cpu
 def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     calls = []
@@ -550,6 +580,199 @@ def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None):
     return fake
 
 
+@pytest.fixture
+def npu_contract(monkeypatch, mocker):
+    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("NPU must not validate FA4"))
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(return_value=torch.zeros_like(query.squeeze(0)))
+    _fake_mindiesd(monkeypatch, attention_forward=dense, attention_forward_varlen=packed)
+    return _npu_impl(), query, dense, packed
+
+
+def _assert_platform_path(impl, platform, query, metadata, path):
+    context = ExecutionContext(platform=platform, require_fullgraph=True)
+    result = impl.resolve_execution_path(context, query, query, query, metadata)
+    assert result.path == path
+    assert result.kernel_variant == {"npu": "mindiesd", "rocm": "aiter"}[platform]
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+def _real_pad_metadata():
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    return AttentionMetadata(
+        extra={
+            "npu_attn_varlen": True,
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 5,
+            "max_seqlen_k": 5,
+        }
+    )
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("kv_cache_dtype", [None, "float", "auto"])
+def test_npu_contract_dense_dispatch(npu_contract, kv_cache_dtype):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": kv_cache_dtype})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_dense")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_mask_layout(npu_contract):
+    impl, query, dense, packed = npu_contract
+    mask = torch.arange(8)[None] < 5
+    metadata = AttentionMetadata(attn_mask=mask)
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask[:, None, None, :].expand(1, 1, 8, 8))
+
+
+@pytest.mark.cpu
+def test_npu_contract_packed_dispatch(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[3:] == ([0, 5, 8], [0, 5, 8])
+
+
+@pytest.mark.cpu
+def test_npu_contract_prefix_dispatch(npu_contract, monkeypatch):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    # The same metadata resolves differently after the kernel choice changes.
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    _assert_platform_path(impl, "npu", query, metadata, "npu_prefix_kv_slice")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.args[1], query[:, :5])
+    torch.testing.assert_close(dense.call_args.args[2], query[:, :5])
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_incomplete_packing_rebuilds_mask(npu_contract):
+    impl, query, dense, packed = npu_contract
+    # CUDA-like normalization rejects this metadata, but NPU can rebuild a mask.
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True, "max_seqlen_q": 5, "valid_kv_length": 5})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    mask = (torch.arange(8) < 5)[None, None, None, :].expand(1, 1, 8, 8)
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask)
+
+
+@pytest.mark.cpu
+def test_npu_contract_unusable_packing_does_not_run_unmasked(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    with pytest.raises(ValueError, match="refusing to run unmasked attention"):
+        impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_not_called()
+
+
+@pytest.fixture
+def rocm_contract(monkeypatch, mocker):
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(side_effect=lambda *args, **kwargs: torch.zeros_like(args[0] if args else kwargs["q"]))
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", True)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", dense)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", packed)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("ROCm must not validate FA4"))
+    return _npu_impl(), query, dense, packed
+
+
+@pytest.mark.cpu
+def test_rocm_contract_dense_dispatch(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    _assert_platform_path(impl, "rocm", query, None, "rocm_dense")
+    impl.forward_hip(query, query, query)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_masked_unpadding(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.arange(8)[None] < 5, extra={"attention_mask_mode": "padding"})
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_masked_varlen")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[0].shape == (5, 2, 4)
+    torch.testing.assert_close(packed.call_args.kwargs["cu_seqlens_k"], torch.tensor([0, 5], dtype=torch.int32))
+
+
+@pytest.mark.cpu
+def test_rocm_contract_multiple_packed_documents(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    cu = torch.tensor([0, 2, 4, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 4,
+            "max_seqlen_k": 4,
+        }
+    )
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_packed_varlen")
+    # NPU's [real, pad] path cannot consume this packing contract.
+    assert impl._resolve_packed_seq_npu(query, query, metadata.extra) is None
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.kwargs["cu_seqlens_q"] is cu
+    assert packed.call_args.kwargs["cu_seqlens_k"] is cu
+
+
+@pytest.mark.cpu
+def test_rocm_contract_unknown_mask_stays_runtime_dependent(rocm_contract, monkeypatch):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.ones(1, 8, dtype=torch.bool))
+    # Resolution may inspect shape/metadata, but must not read mask values.
+    with monkeypatch.context() as patch:
+        patch.setattr(torch, "any", lambda *_args, **_kwargs: pytest.fail("resolution read mask values"))
+        _assert_platform_path(impl, "rocm", query, metadata, "rocm_runtime_mask_dependent")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_requires_initialized_aiter(monkeypatch):
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", False)
+    impl = _npu_impl()
+    query = torch.zeros(1, 8, 2, 4)
+    result = impl.resolve_execution_path(
+        ExecutionContext(platform="rocm", kernel_variant="aiter"), query, query, query, None
+    )
+    assert result.path == "rocm_unverified"
+    assert result.kernel_variant is None
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 
 

tests/diffusion/attention/test_flash_attn_compile.py

@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import ForwardContext
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+
+@pytest.fixture(autouse=True)
+def isolated_compiler_cache():
+    # Each case changes kernel/configuration on the same Attention.forward code
+    # object. Keep unrelated cases from exhausting its Dynamo recompile limit;
+    # retain the cache across all shapes and replays within an individual test.
+    torch.compiler.reset()
+    yield
+    torch.compiler.reset()
+
+
+@dataclass
+class _ParallelStrategyStub:
+    name: str
+    enabled: bool = True
+
+
+def _make_attention_layer(monkeypatch, *, head_size=64):
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    monkeypatch.setattr(
+        attention_layer,
+        "get_attn_backend_for_role",
+        lambda **_kwargs: (FlashAttentionBackend, None),
+    )
+    monkeypatch.setattr(
+        attention_layer,
+        "build_parallel_attention_strategy",
+        lambda **_kwargs: NoParallelAttention(),
+    )
+    return attention_layer.Attention(num_heads=8, head_size=head_size, softmax_scale=head_size**-0.5, causal=False)
+
+
+@pytest.fixture
+def fake_fa4(monkeypatch, tmp_path):
+    marker = tmp_path / "kernel"
+    marker.write_text("loaded", encoding="utf-8")
+
+    @functools.cache
+    def cached_kernel_loader():
+        with open(marker, encoding="utf-8") as handle:
+            handle.read()
+
+    def fake_attention(query, key, value, **kwargs):
+        cached_kernel_loader()
+        return (
+            torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2),
+                key.transpose(1, 2),
+                value.transpose(1, 2),
+                scale=kwargs["softmax_scale"],
+                is_causal=kwargs["causal"],
+            )
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", True)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_attention)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+    return FlashAttentionImpl(
+        num_heads=8,
+        head_size=64,
+        softmax_scale=0.125,
+        causal=False,
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_fa4_dense_dispatch_is_opaque_to_dynamic_torch_compile(fake_fa4, noncontiguous):
+    """Fast Dynamo/custom-op regression; Inductor coverage is separate."""
+    impl = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q = q.transpose(1, 2).contiguous().transpose(1, 2)
+        assert not q.is_contiguous()
+    path = impl.resolve_execution_path(
+        ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+            dtype="bfloat16",
+            require_fullgraph=True,
+        ),
+        q,
+        q,
+        q,
+        None,
+    )
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    compiled = torch.compile(
+        lambda query, key, value: impl.forward_cuda(query, key, value),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    out = compiled(q, q, q)
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q2 = q2.transpose(1, 2).contiguous().transpose(1, 2)
+    out2 = compiled(q2, q2, q2)
+    assert out.is_contiguous()
+    assert out2.is_contiguous()
+
+    assert out.shape == q.shape
+    assert out2.shape == q2.shape
+    torch.testing.assert_close(out, impl.forward_cuda(q, q, q))
+    torch.testing.assert_close(out2, impl.forward_cuda(q2, q2, q2))
+    assert compile_count == 1
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_fa4_production_attention_entry_compiles_with_inductor(
+    monkeypatch,
+    fake_fa4,
+):
+    """Compile through Attention.forward with the default Inductor backend."""
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    layer = _make_attention_layer(monkeypatch)
+    layer.attention = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+
+    resolved_contexts = []
+    resolve_execution_path = fake_fa4.resolve_execution_path
+
+    def record_execution_path(context, query, key, value, attn_metadata):
+        resolved_contexts.append(context)
+        return resolve_execution_path(
+            context,
+            query,
+            key,
+            value,
+            attn_metadata,
+        )
+
+    monkeypatch.setattr(fake_fa4, "resolve_execution_path", record_execution_path)
+    context = ExecutionContext(
+        platform="cuda",
+        kernel_variant="fa4",
+        dtype="bfloat16",
+        paged_kv=True,
+        parallel_strategy=ParallelStrategy.ULYSSES,
+        require_fullgraph=True,
+    )
+    path = layer.resolve_execution_path(context, q, q, q, None)
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+    assert resolved_contexts[-1].paged_kv is False
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    out = compiled(q, q, q)
+
+    assert out.shape == q.shape
+    torch.testing.assert_close(out, layer(q, q, q))
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    torch.testing.assert_close(compiled(q2, q2, q2), layer(q2, q2, q2))
+
+    layer.paged_kv_cache_role = "self"
+    monkeypatch.setattr(layer, "_active_paged_kv_adapter", lambda: object())
+    paged_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].paged_kv is True
+    assert paged_path.support.status is SupportStatus.UNMIGRATED
+    layer.paged_kv_cache_role = None
+
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ulysses")
+    hybrid_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.HYBRID_ULYSSES_RING
+    assert hybrid_path.support.status is SupportStatus.UNMIGRATED
+    layer.use_ring = False
+    layer.parallel_strategy = NoParallelAttention()
+
+    layer._hsdp_compile_boundary_enabled = True
+    hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert hsdp_path.support.status is SupportStatus.UNMIGRATED
+    assert hsdp_path.compilation_mode is CompilationMode.EAGER_ONLY
+
+    layer._hsdp_compile_boundary_enabled = False
+    runtime_context = ForwardContext(
+        omni_diffusion_config=OmniDiffusionConfig(
+            parallel_config=DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=1),
+        ),
+    )
+    monkeypatch.setattr(attention_layer, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(attention_layer, "get_forward_context", lambda: runtime_context)
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ring")
+    runtime_hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+    assert runtime_hsdp_path.support.status is SupportStatus.UNMIGRATED
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(32, 32), (64, 64), (80, 48), (192, 128), (256, 256)])
+def test_real_fa4_fullgraph_matches_sdpa_across_shapes(monkeypatch, head_dims):
+    """Validate real FA4 numerics through the complete compiled attention layer."""
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major < 10:
+        pytest.skip("Requires Blackwell")
+    if not fa.IS_FLASH_ATTN_4 or fa.flash_attn_func is None:
+        pytest.skip("Requires CuTe FlashAttention-4")
+
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    # Tile boundaries, cross-attention lengths, batching, and a longer sequence.
+    for batch, q_length, kv_length in ((1, 16, 16), (1, 129, 257), (2, 257, 129), (1, 1024, 1024)):
+        query, key, value = (
+            torch.randn(batch, length, 8, dim, device="cuda", dtype=torch.bfloat16, generator=generator)
+            for length, dim in ((q_length, head_size), (kv_length, head_size), (kv_length, value_head_size))
+        )
+        path = layer.resolve_execution_path(context, query, key, value, None)
+        assert path.requested_support(context).status is SupportStatus.SUPPORTED
+        with sdpa_kernel(SDPBackend.MATH):
+            reference = torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2).float(),
+                key.transpose(1, 2).float(),
+                value.transpose(1, 2).float(),
+                scale=head_size**-0.5,
+            ).transpose(1, 2)
+        actual = compiled(query, key, value)
+        assert actual.dtype == query.dtype
+        assert actual.device == query.device
+        torch.testing.assert_close(actual.float(), reference, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(actual, layer(query, key, value))
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(192, 192), (65, 64), (64, 65)])
+def test_real_fa4_dimension_rejection_matches_kernel(monkeypatch, head_dims):
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major not in (10, 11) or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires SM100/SM110 CuTe FlashAttention-4")
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    query = torch.randn(1, 16, 8, head_size, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 16, 8, value_head_size, device="cuda", dtype=torch.bfloat16)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    path = layer.resolve_execution_path(context, query, query, value, None)
+    assert path.requested_support(context).status is SupportStatus.UNSUPPORTED
+    with pytest.raises(AssertionError) as error:
+        layer(query, query, value)
+    assert str(error.value) in path.support.reason
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_real_fa4_custom_op_unequal_value_dimension_schema(noncontiguous):
+    if not current_omni_platform.is_cuda() or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires CuTe FlashAttention-4")
+    query = torch.randn(1, 17, 8, 80, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(1, 25, 8, 80, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 25, 8, 48, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        query, key, value = (t.transpose(1, 2).contiguous().transpose(1, 2) for t in (query, key, value))
+        assert not query.is_contiguous()
+    torch.library.opcheck(
+        torch.ops.vllm_omni.fa4_dense_attention.default,
+        (query, key, value, 80**-0.5, False, False),
+    )

tests/diffusion/distributed/test_cosmos3_pre_sharded.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare full and pre-sharded Cosmos loading with real two-rank FSDP."""
+
+from contextlib import ExitStack
+from datetime import timedelta
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.gpu]
+
+
+class _SDPAAttention(nn.Module):
+    """Use deterministic native attention while testing loading and FSDP."""
+
+    def __init__(self, *, causal, softmax_scale, **kwargs):
+        super().__init__()
+        self.causal = causal
+        self.scale = softmax_scale
+
+    def forward(self, query, key, value):
+        return F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            is_causal=self.causal,
+            scale=self.scale,
+        ).transpose(1, 2)
+
+
+def _load_cosmos_worker(rank, rendezvous, checkpoint_dir, edge):
+    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
+
+    bootstrap_vllm_layer_custom_op_modules()
+
+    from safetensors.torch import save_file
+    from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+    from vllm.config.load import LoadConfig
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+    from vllm.utils.torch_utils import set_default_torch_dtype
+
+    from vllm_omni.diffusion.distributed import hsdp
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.platforms import current_omni_platform
+
+    device = torch.device("cuda", rank)
+    current_omni_platform.set_device(device)
+    dist.init_process_group("nccl", init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=120))
+    try:
+        with ExitStack() as stack:
+            stack.enter_context(set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda"))))
+            stack.enter_context(set_default_torch_dtype(torch.bfloat16))
+            # HSDP spans both workers, while every vLLM linear uses TP=1.
+            for module in (linear, parameter):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_rank", return_value=0))
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+            for module in (transformer_cosmos3, transformer_cosmos3_edge):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+                stack.enter_context(patch.object(module, "FrameworkAttention", _SDPAAttention))
+            stack.enter_context(patch.object(transformer_cosmos3, "_get_ulysses_state", return_value=(1, 0, None)))
+            stack.enter_context(
+                patch.object(hsdp, "get_world_group", return_value=SimpleNamespace(world_size=2, rank_in_group=rank))
+            )
+
+            model_cls = (
+                transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer
+                if edge
+                else transformer_cosmos3.Cosmos3VFMTransformer
+            )
+            tf_config = dict(
+                hidden_size=16,
+                num_hidden_layers=1,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=8,
+                intermediate_size=32,
+                vocab_size=32,
+                latent_patch_size=2,
+                latent_channel=48,
+                temporal_compression_factor=4,
+                rope_scaling={"mrope_section": [2, 1, 1]},
+            )
+            if edge:
+                tf_config.update(
+                    backbone_type=transformer_cosmos3_edge.COSMOS3_EDGE_BACKBONE_TYPE,
+                    qk_norm_for_text=False,
+                    use_und_k_norm_for_gen=True,
+                )
+            config = SimpleNamespace(
+                tf_model_config=tf_config,
+                dtype=torch.bfloat16,
+                quantization_config=None,
+                lora_path=None,
+                num_weight_load_threads=1,
+                hsdp_weight_load_strategy="full",
+                parallel_config=SimpleNamespace(
+                    use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=2, tensor_parallel_size=1
+                ),
+            )
+
+            class Pipeline(nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.transformer = model_cls(config)
+                    self.transformer.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+                    self.vae = nn.Linear(2, 2)
+                    self.weights_sources = [
+                        DiffusersPipelineLoader.ComponentSource(checkpoint_dir, None, None, "transformer.", False)
+                    ]
+
+                def load_weights(self, weights):
+                    state = dict(weights)
+                    self.transformer.load_state_dict(
+                        {name.removeprefix("transformer."): t for name, t in state.items()}
+                    )
+                    self.transformer.post_load_weights()
+                    self.transformer.eval()
+                    self.transformer.validate_loaded_weights(set(state))
+                    return set(state)
+
+            full = Pipeline()
+            if rank == 0:
+                torch.manual_seed(42)
+                with torch.no_grad():
+                    for name, tensor in full.transformer.named_parameters():
+                        tensor.copy_(torch.randn_like(tensor) * 0.02)
+                        if "norm" in name:
+                            tensor.fill_(1)
+                save_file(full.transformer.state_dict(), str(Path(checkpoint_dir) / "model.safetensors"))
+            dist.barrier()
+
+            def load(model):
+                loader = DiffusersPipelineLoader(LoadConfig(), config)
+                with patch.object(loader, "_init_from_load_format", return_value=model):
+                    return loader._load_model_with_hsdp(device)
+
+            full = load(full)
+            config.hsdp_weight_load_strategy = "pre_sharded"
+            pre_sharded = load(Pipeline())
+            full_parameters = dict(full.transformer.named_parameters())
+            from torch.distributed.tensor import DTensor
+
+            for name, actual in pre_sharded.transformer.named_parameters():
+                expected = full_parameters[name]
+                if name.startswith("time_embedder."):
+                    assert not isinstance(actual, DTensor)
+                    assert actual.dtype == torch.float32
+                else:
+                    assert isinstance(actual, DTensor)
+                    actual = actual.full_tensor()
+                    expected = expected.full_tensor()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            assert pre_sharded.vae.weight.device == device
+            torch.testing.assert_close(pre_sharded.transformer.scratch, full.transformer.scratch)
+            assert all(t.device.type != "meta" for t in pre_sharded.transformer.buffers())
+
+            inputs = dict(
+                hidden_states=torch.linspace(-1, 1, 48 * 4, device=device).reshape(1, 48, 1, 2, 2),
+                timestep=torch.tensor([1.0], device=device),
+                text_ids=torch.tensor([[1, 2]], device=device),
+                text_mask=torch.ones(1, 2, dtype=torch.long, device=device),
+                video_shape=(1, 2, 2),
+                fps=24.0,
+            )
+            with torch.no_grad():
+                # Also exercise the cached UND path on the second denoising step.
+                for step in (1.0, 0.5):
+                    inputs["timestep"].fill_(step)
+                    expected = full.transformer(**inputs)
+                    actual = pre_sharded.transformer(**inputs)
+                    assert torch.isfinite(actual).all()
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("edge", [False, True], ids=["cosmos3", "cosmos3_edge"])
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 2 or not dist.is_nccl_available(),
+    reason="requires two CUDA GPUs and NCCL",
+)
+def test_cosmos3_pre_sharded_matches_full_hsdp(tmp_path, edge):
+    torch.multiprocessing.spawn(
+        _load_cosmos_worker,
+        args=(f"file://{tmp_path / 'rendezvous'}", str(tmp_path), edge),
+        nprocs=2,
+        join=True,
+    )

tests/diffusion/distributed/test_hsdp.py

@@ -3,11 +3,13 @@
 """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
 import gc
+from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.fsdp import fully_shard
 from torch.distributed.tensor import DeviceMesh, DTensor
 from vllm.utils.network_utils import get_file_store_init_method
 
@@ -18,6 +20,7 @@ from vllm_omni.diffusion.distributed.hsdp import (
     _unshardable_parameters,
     shard_model,
 )
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
@@ -338,3 +341,92 @@ class TestHSDPShardConditions:
                 matched.append(name)
         assert "blocks.0" in matched
         assert "blocks.1" in matched
+
+
+def test_pre_sharded_to_empty_materializes_fsdp_internal_storage(cpu_process_group):
+    with torch.device("meta"):
+        model = nn.Module()
+        model.transformer = nn.Sequential(nn.Linear(2, 2, bias=False))
+    model.transformer.register_buffer("scratch", torch.tensor([3.0]), persistent=False)
+    shard_model(
+        model.transformer,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "0"],
+    )
+
+    load_plan = SimpleNamespace(
+        roots=[SimpleNamespace(name="transformer", module=model.transformer)],
+        groups=[SimpleNamespace(module=model.transformer[0])],
+        binding_names=frozenset({"transformer.0.weight"}),
+    )
+    sharded_weight = model.transformer[0].weight
+
+    DiffusersPipelineLoader._materialize_pre_sharded_hsdp_state(model, load_plan, torch.device("cpu"))
+
+    fsdp_param = fully_shard.state(model.transformer[0])._fsdp_param_group.fsdp_params[0]
+    assert model.transformer[0].weight is sharded_weight
+    assert model.transformer[0].weight.to_local().device.type == "cpu"
+    assert fsdp_param._sharded_param_data.device.type == "cpu"
+    assert (
+        fsdp_param._sharded_param_data.untyped_storage().data_ptr()
+        == model.transformer[0].weight.to_local().untyped_storage().data_ptr()
+    )
+    assert torch.equal(model.transformer.scratch, torch.tensor([3.0]))
+
+
+def test_pre_sharded_load_restores_checkpoint_values(cpu_process_group, tmp_path, monkeypatch):
+    from safetensors.torch import save_file
+    from vllm.config.load import LoadConfig
+
+    from vllm_omni.diffusion.distributed import hsdp as hsdp_module
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    class Transformer(nn.Module):
+        _hsdp_shard_conditions = [lambda name, module: name == "block"]
+        _hsdp_ignored_modules = ["time_embedder"]
+
+        def __init__(self):
+            super().__init__()
+            self.block = nn.Linear(3, 4, bias=False)
+            self.time_embedder = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16)
+            self.root_weight = nn.Parameter(torch.zeros(2))
+            self.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+            release_module_parameters_to_meta(self.block)
+
+        def post_load_weights(self):
+            self.time_embedder.to(torch.float32)
+
+    model = nn.Module()
+    model.transformer = Transformer()
+    model.vae = nn.Linear(2, 2)
+    tensors = {
+        "block.weight": torch.arange(12, dtype=torch.float32).reshape(4, 3),
+        "root_weight": torch.tensor([2.0, 3.0]),
+        "time_embedder.weight": torch.arange(4, dtype=torch.bfloat16).reshape(2, 2),
+    }
+    save_file(tensors, str(tmp_path / "model.safetensors"))
+    model.weights_sources = [DiffusersPipelineLoader.ComponentSource(str(tmp_path), None, None, "transformer.", False)]
+    config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization_config=None,
+        lora_path=None,
+        hsdp_weight_load_strategy="pre_sharded",
+        num_weight_load_threads=1,
+        parallel_config=SimpleNamespace(use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=1),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: model)
+    monkeypatch.setattr(hsdp_module, "get_world_group", lambda: SimpleNamespace(world_size=1, rank_in_group=0))
+    monkeypatch.setattr(hsdp_module, "current_omni_platform", SimpleNamespace(device_type="cpu"))
+    loaded = loader._load_model_with_hsdp(torch.device("cpu"))
+    assert loaded is model
+    for name, expected in tensors.items():
+        actual = dict(model.transformer.named_parameters())[name]
+        if name.startswith("time_embedder."):
+            assert not isinstance(actual, DTensor)
+            assert actual.dtype == torch.float32
+            torch.testing.assert_close(actual, expected.float())
+        else:
+            assert isinstance(actual, DTensor)
+            torch.testing.assert_close(actual.to_local(), expected)
+    torch.testing.assert_close(model.transformer.scratch, torch.tensor([7.0]))

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -27,6 +27,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     HostWeightPlanResult,
     TensorBinding,
+    build_checkpoint_binding_plan,
 )
 from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.models.helios import HeliosPipeline
@@ -2089,3 +2090,206 @@ def test_hsdp_broadcast_weight_load_online_quant_multiprocess():
             )
         assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag"))
         assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag"))
+
+
+def test_pre_sharded_hsdp_rejects_quantization_before_loading():
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=object(),
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    with pytest.raises(ValueError, match="does not support quantization"):
+        loader._load_model_with_hsdp(torch.device("cpu"))
+
+
+def test_pre_sharded_hsdp_strategy_dispatches_without_using_full_loader(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=None,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    discovered = PipelineModules(
+        dits=[model.transformer],
+        dit_names=["transformer"],
+        vaes=[],
+        encoders=[],
+        encoder_names=[],
+        resident_modules=[],
+        resident_names=[],
+    )
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=AssertionError("full loader must not run"))  # type: ignore[method-assign]
+    mocker.patch.object(loader_mod.ModuleDiscovery, "discover", return_value=discovered)
+    pre_sharded_load = mocker.patch.object(
+        loader,
+        "_load_model_with_pre_sharded_hsdp",
+        return_value=model,
+    )
+
+    assert loader._load_model_with_hsdp(torch.device("cpu")) is model
+    pre_sharded_load.assert_called_once()
+
+
+def test_pre_sharded_hsdp_rejects_tensor_transform():
+    def transform(tensor):
+        return tensor.transpose(-2, -1)
+
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+                transform=transform,
+            )
+        }
+    )
+
+    with pytest.raises(ValueError, match="tensor transforms are unsupported"):
+        DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_accepts_direct_safetensors_binding():
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+            )
+        }
+    )
+
+    DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_uses_fsdp_destination_directly():
+    target = torch.empty(8, 4, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+
+    checkpoint_states, local_bytes = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": target},
+    )
+    checkpoint_target = checkpoint_states[Path("transformer")]["source_weight"]
+
+    assert checkpoint_target is target
+    assert checkpoint_target.shape == (8, 4, 3)
+    assert checkpoint_target.device.type == "cpu"
+    assert checkpoint_target.untyped_storage().data_ptr() == target.untyped_storage().data_ptr()
+    assert local_bytes == target.numel() * target.element_size()
+
+
+def test_pre_sharded_hsdp_checkpoint_writes_runtime_storage_directly():
+    runtime_target = torch.empty(2, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+    checkpoint_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    checkpoint_states, _ = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": runtime_target},
+    )
+    checkpoint_states[Path("transformer")]["source_weight"].copy_(checkpoint_tensor)
+
+    assert torch.equal(runtime_target, checkpoint_tensor)
+
+
+def test_pre_sharded_hsdp_rejects_duplicate_checkpoint_destinations():
+    binding = TensorBinding(checkpoint_key="weight", file_path="transformer/model.safetensors")
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.first": binding,
+            "transformer.second": binding,
+        }
+    )
+
+    with pytest.raises(ValueError, match="multiple runtime tensors"):
+        DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+            load_plan,
+            {"transformer.first": torch.empty(2, 3), "transformer.second": torch.empty(2, 3)},
+        )
+
+
+def test_pre_sharded_hsdp_managed_module_traversal_is_child_first_and_skips_nested_roots():
+    root = nn.Module()
+    root.left = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+    root.right = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+
+    names = {module: name or "root" for name, module in root.named_modules()}
+    modules = DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+        root,
+        nested_hsdp_roots={root.left},
+    )
+
+    assert [names[module] for module in modules] == ["right.0", "right.1", "right", "root"]
+
+
+@pytest.mark.parametrize("checkpoint_key", ["renamed", "unexpected"])
+def test_hsdp_checkpoint_plan_honors_remap_and_rejects_missing(tmp_path, checkpoint_key, monkeypatch):
+    import vllm_omni.diffusion.model_loader.host_weight_plan as plan_mod
+
+    monkeypatch.setattr(plan_mod, "get_direct_mmap_adapter", lambda _model: None)
+
+    class Pipeline(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.Linear(2, 2, bias=False)
+
+        @staticmethod
+        def remap_checkpoint_key(name):
+            return "transformer.weight" if name == "transformer.renamed" else name
+
+    checkpoint = tmp_path / "model.safetensors"
+    save_file({checkpoint_key: torch.ones(2, 2)}, str(checkpoint))
+    source = SimpleNamespace(
+        model_or_path=str(tmp_path),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    model = Pipeline()
+    result = build_checkpoint_binding_plan(
+        model,
+        dit_modules=(("transformer", model.transformer),),
+        sources=(source,),
+        model_path=None,
+        tensor_parallel_size=1,
+        online_quantization=False,
+    )
+
+    if checkpoint_key == "renamed":
+        assert result.plan is not None
+        assert result.plan.bindings["transformer.weight"].checkpoint_key == "renamed"
+    else:
+        assert result.plan is None
+        assert "no checkpoint binding" in result.fallback_reason

tests/diffusion/models/auk/test_auk_cudagraph.py

@@ -61,7 +61,10 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa
 
     # The conditioning is prepared once per request and reused by every step.
     assert prepare_spy.call_count == 1
-    assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
+    # _prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
+    assert prepare_spy.call_args.args[6] is (cfg_strength >= 1e-5)
+    # The time grid reaches prepare, so the adaLN modulations are computed once.
+    assert prepare_spy.call_args.args[7].numel() == 2
     assert step_spy.call_count == 2
     capture_spy.assert_not_called()
     assert not wrapper._cache
@@ -84,7 +87,8 @@ def test_graph_inputs_use_bounded_length_buckets() -> None:
     assert padded[4].shape == (1, 100, 4)
     assert padded[5].shape == (1, 100)
     assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51]
-    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
+    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False, 0)
+    assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32)
     assert wrapper.max_graphs == 32
 
 
@@ -153,18 +157,21 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
             inputs["ref"],
             inputs["ref_mask"],
         )
-        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
+        # sample_latents passes the grid, so the key carries its step count (two steps here).
+        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5, 2)
         assert key in wrapper._cache
 
         # The static context was refreshed for this request's conditioning.
         entry = wrapper._cache[key]
-        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
+        grid = torch.tensor([0.0, 0.4], device="cuda")
+        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid)
         for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True):
             if want is None:
                 assert static is None
             else:
                 torch.testing.assert_close(static, want)
         torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
+        assert int(entry.static_step) == 1
 
     assert len(wrapper._cache) == 1
 

tests/diffusion/models/auk/test_auk_dit_fastpaths.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The AuK DiT fast paths keep the reference numerics.
+
+Two paths are checked: the grouped conv position embedding as one batched
+GEMM, and the adaLN modulations precomputed for a whole time grid, eager and
+under the per-step CUDA graph.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.auk import auk_transformer as T
+from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, build_time_grid
+from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_dit(device: str = "cpu") -> AuKTransformer:
+    torch.manual_seed(5)
+    return (
+        AuKTransformer(
+            dim=32,
+            heads=2,
+            dim_head=16,
+            ff_mult=2,
+            latent_dim=4,
+            text_hidden_dim=8,
+            num_layers=2,
+            num_single_layers=2,
+        )
+        .eval()
+        .to(device)
+    )
+
+
+def _reference_conv_pos_embed(module: T.ConvPosEmbedding, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+    """The channels-first conv stack as the reference implementation writes it."""
+    keep = None if mask is None else mask.unsqueeze(1)
+    x = x.transpose(1, 2)
+    if keep is not None:
+        x = x.masked_fill(~keep, 0.0)
+    for layer in module.conv1d:
+        x = layer(x)
+        if keep is not None and isinstance(layer, torch.nn.Conv1d):
+            x = x.masked_fill(~keep, 0.0)
+    return x.transpose(1, 2)
+
+
+@torch.inference_mode()
+def test_grouped_conv_as_batched_gemm_matches_conv1d() -> None:
+    torch.manual_seed(0)
+    conv = torch.nn.Conv1d(64, 64, 31, groups=16, padding=15)
+    x = torch.randn(2, 37, 64)
+    expected = conv(x.transpose(1, 2)).transpose(1, 2)
+    torch.testing.assert_close(T._grouped_conv1d_btc(x, conv), expected, atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("masked", [False, True])
+def test_conv_pos_embedding_matches_the_channels_first_stack(masked: bool) -> None:
+    torch.manual_seed(1)
+    module = T.ConvPosEmbedding(64).eval()
+    x = torch.randn(2, 37, 64)
+    mask = None
+    if masked:
+        mask = torch.ones(2, 37, dtype=torch.bool)
+        mask[1, 30:] = False
+    torch.testing.assert_close(module(x, mask), _reference_conv_pos_embed(module, x, mask), atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_infer", [False, True])
+def test_modulation_table_matches_per_step_modulation(cfg_infer: bool) -> None:
+    dit = _make_dit()
+    text = torch.randn(1, 5, 8)
+    ref = torch.randn(1, 3, 4)
+    x = torch.randn(1, 6, 4)
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cpu")
+    plain = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer)
+    tabled = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer, timesteps=grid[:-1])
+    assert tabled.modulation is not None and tabled.modulation.shape[0] == 4
+    for i in range(4):
+        expected = dit.step(x, grid[i], plain)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=i), expected, atol=1e-5, rtol=1e-5)
+        index = torch.tensor(i)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=index), expected, atol=1e-5, rtol=1e-5)
+
+
+def _padded_inputs(device: str) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(3)
+    text = torch.randn(1, 7, 8, generator=generator)
+    c_mask = torch.ones(1, 7, dtype=torch.bool)
+    c_mask[:, 5:] = False
+    text[:, 5:] = 0.0
+    ref = torch.randn(1, 6, 4, generator=generator)
+    ref_mask = torch.ones(1, 6, dtype=torch.bool)
+    ref_mask[:, 4:] = False
+    mask = torch.ones(1, 9, dtype=torch.bool)
+    mask[:, 7:] = False
+    inputs = {"text": text, "c_mask": c_mask, "ref": ref, "ref_mask": ref_mask, "mask": mask}
+    return {name: value.to(device) for name, value in inputs.items()}
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_graph_with_time_grid_matches_eager_with_time_grid(cfg_strength: float) -> None:
+    dit = _make_dit("cuda")
+    inputs = _padded_inputs("cuda")
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cuda")
+    x = torch.randn(1, 9, 4, device="cuda")
+    wrapper = AuKCUDAGraphWrapper(dit)
+    common = dict(
+        text=inputs["text"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        ref_mask=inputs["ref_mask"],
+        cfg_strength=cfg_strength,
+        timesteps=grid[:-1],
+    )
+    for i in range(4):
+        graph = wrapper(x=x, timestep=grid[i], new_request=i == 0, step_index=i, **common)
+        ctx = dit.prepare(
+            inputs["text"],
+            target_len=9,
+            c_mask=inputs["c_mask"],
+            ref=inputs["ref"],
+            ref_mask=inputs["ref_mask"],
+            cfg_infer=cfg_strength >= 1e-5,
+            timesteps=grid[:-1],
+        )
+        eager = wrapper._step(x, grid[i], ctx, cfg_strength, i)
+        torch.testing.assert_close(graph, eager, atol=1e-4, rtol=1e-4)
+    # One graph serves every step of the grid.
+    assert len(wrapper._cache) == 1
+    assert next(iter(wrapper._cache))[-1] == 4

tests/diffusion/models/auk/test_auk_step_context.py

@@ -127,7 +127,9 @@ def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> Non
             v = dit(x, time=timesteps[i], **inputs)
         x = x + (timesteps[i + 1] - timesteps[i]) * v
 
-    torch.testing.assert_close(got, x, rtol=0, atol=0)
+    # sample_latents precomputes the adaLN modulations of the whole grid in one
+    # GEMM per layer, whose rows round differently from one-row GEMMs in fp32.
+    torch.testing.assert_close(got, x, rtol=0, atol=1e-6)
     assert dit.text_cond is None and dit.text_uncond is None
 
 

tests/diffusion/models/auk/test_auk_vae_fastpath.py

@@ -12,7 +12,8 @@ import pytest
 import torch
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
+from vllm_omni.diffusion.models.auk import auk_vae
+from vllm_omni.diffusion.models.auk.auk_vae import AliasFreeActivation, AuKVAE, LowPass, SnakeBeta, Upsample
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles
 from vllm_omni.platforms import current_omni_platform
 
@@ -246,6 +247,72 @@ def test_tiled_decode_matches_the_whole_decode() -> None:
         AuKVAEDecodeGraph(vae, tile_frames=60)
 
 
+@pytest.mark.cpu
+@torch.inference_mode()
+def test_fused_activation_is_skipped_off_cuda() -> None:
+    activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True))
+    x = torch.randn(1, 4, 16)
+    assert not activation._fused_eligible(x)
+    assert activation(x).shape == (1, 4, 16)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
+@torch.inference_mode()
+def test_fused_activation_is_skipped_without_triton(monkeypatch: pytest.MonkeyPatch) -> None:
+    activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)).to("cuda")
+    x = torch.randn(1, 4, 16, device="cuda")
+    assert activation._fused_eligible(x)
+    monkeypatch.setattr(auk_vae, "HAS_TRITON", False)
+    assert not activation._fused_eligible(x)
+    assert activation(x).shape == (1, 4, 16)
+
+
+def _activation_input(layout: str, length: int) -> torch.Tensor:
+    """A ``[2, 6, length]`` input, contiguous or as a strided view of a larger buffer."""
+    if layout == "contiguous":
+        x = torch.randn(2, 6, length, device="cuda")
+    elif layout == "transposed":
+        x = torch.randn(2, length, 6, device="cuda").transpose(1, 2)
+    elif layout == "sliced":
+        x = torch.randn(2, 8, 2 * length, device="cuda")[:, 1:7, ::2]
+    else:
+        raise ValueError(layout)
+    assert x.shape == (2, 6, length)
+    if length > 1:  # contiguity of a single-sample view depends on the layout
+        assert x.is_contiguous() == (layout == "contiguous")
+    return x * 3
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
+@pytest.mark.parametrize("layout", ["contiguous", "transposed", "sliced"])
+@pytest.mark.parametrize("causal", [False, True])
+# The kernels tile the output in blocks of 1024 samples: the 2x oversampled
+# signal crosses a block boundary at 512 input samples and the decimated one at
+# 1024, so both sides of each boundary are covered, plus short and odd lengths.
+@pytest.mark.parametrize("length", [1, 7, 511, 512, 513, 1000, 1023, 1024, 1025])
+@torch.inference_mode()
+def test_fused_activation_matches_the_eager_modules(layout: str, causal: bool, length: int) -> None:
+    torch.manual_seed(0)
+    activation = AliasFreeActivation(SnakeBeta(6, alpha_logscale=True), causal=causal).to("cuda")
+    with torch.no_grad():
+        activation.act.alpha.normal_(0.0, 0.3)
+        activation.act.beta.normal_(0.0, 0.3)
+    x = _activation_input(layout, length)
+    original = x.clone()
+
+    eager = activation.downsample(activation.act(activation.upsample(x)))
+    assert activation._fused_eligible(x)
+    fused = activation(x)
+
+    # Same fp32 arithmetic, only the order in which the taps are summed differs.
+    assert fused.shape == eager.shape
+    torch.testing.assert_close(fused, eager, atol=1e-5, rtol=1e-5)
+    # Neither path writes into its (possibly strided) input.
+    torch.testing.assert_close(x, original, atol=0, rtol=0)
+
+
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
 @torch.inference_mode()

tests/diffusion/models/auk/test_pipeline_auk.py

@@ -291,6 +291,18 @@ class TestRequestParsing:
         with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"):
             build_pipeline(model_config={"max_dit_graphs": max_dit_graphs})
 
+    @pytest.mark.parametrize(
+        "model_config, expected", [({}, 16), ({"auk_ref_cache_size": 0}, 0), ({"auk_ref_cache_size": 4}, 4)]
+    )
+    def test_ref_cache_size_is_taken_from_model_config(self, build_pipeline, model_config, expected):
+        pipeline, _ = build_pipeline(model_config=model_config)
+        assert pipeline._ref_cache_size == expected
+
+    @pytest.mark.parametrize("size", [-1, 1.9, True, False, "4", None])
+    def test_invalid_ref_cache_size_is_rejected(self, build_pipeline, size):
+        with pytest.raises(ValueError, match="AuK auk_ref_cache_size must be a non-negative integer"):
+            build_pipeline(model_config={"auk_ref_cache_size": size})
+
     def test_pipeline_declares_audio_output(self, build_pipeline):
         pipeline, _ = build_pipeline()
 
@@ -492,6 +504,40 @@ class TestRequestParsing:
         assert call["sample"] is True
         assert isinstance(call["generator"], torch.Generator)
 
+    def test_reference_latents_are_cached_by_content(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+        clip = _silence(1.0)
+
+        for _ in range(3):
+            pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+        # A different clip is a new entry, the first one still hits.
+        other = (np.full_like(clip[0], 0.25), clip[1])
+        pipeline.forward(_batch(_prompt(audio=other, knobs={"gen_seconds": 1.0}), seed=1))
+        pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+
+        assert len(pipeline.vae.encode_calls) == 2
+        assert calls[0]["ref"] is calls[1]["ref"] is calls[4]["ref"]
+
+    def test_posterior_draws_and_disabled_cache_always_encode(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        for _ in range(2):
+            pipeline.forward(
+                _batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1)
+            )
+        assert len(pipeline.vae.encode_calls) == 2
+
+        pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 0})
+        for _ in range(2):
+            pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0}), seed=1))
+        assert len(pipeline.vae.encode_calls) == 2
+
+    def test_reference_cache_is_bounded(self, build_pipeline):
+        pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 2})
+        for level in (0.1, 0.2, 0.3):
+            clip = (np.full(SAMPLE_RATE, level, dtype=np.float32), SAMPLE_RATE)
+            pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
+        assert len(pipeline._ref_cache) == 2
+
     def test_one_request_per_forward(self, build_pipeline):
         pipeline, _ = build_pipeline()
         batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1)

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -868,6 +868,93 @@ def test_qwen_gen_mlp_remains_gated() -> None:
     assert hasattr(model.gen_layers[0].mlp, "gate_proj")
 
 
+@pytest.mark.parametrize("edge", [False, True])
+@pytest.mark.parametrize(
+    "use_hsdp,strategy", [(False, "full"), (False, "pre_sharded"), (True, "full"), (True, "pre_sharded")]
+)
+def test_cosmos3_releases_blocks_only_for_pre_sharded_hsdp(edge, use_hsdp, strategy, monkeypatch):
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    released = []
+
+    def release_and_check(block):
+        parameters = dict(block.named_parameters())
+        metadata = {name: dict(parameter.__dict__) for name, parameter in parameters.items()}
+        release_module_parameters_to_meta(block)
+        for name, parameter in block.named_parameters():
+            assert parameter is parameters[name]
+            assert parameter.device.type == "meta"
+            assert parameter.__dict__ == metadata[name]
+        released.append(block)
+
+    monkeypatch.setattr(transformer_cosmos3, "release_module_parameters_to_meta", release_and_check)
+    monkeypatch.setattr(transformer_cosmos3_edge, "release_module_parameters_to_meta", release_and_check)
+    model_cls = (
+        transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer if edge else transformer_cosmos3.Cosmos3VFMTransformer
+    )
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+            hsdp_weight_load_strategy=strategy,
+        )
+    )
+    blocks = [*model.language_model.layers, *model.gen_layers]
+    expected_meta = use_hsdp and strategy == "pre_sharded"
+    assert len(released) == (len(blocks) if expected_meta else 0)
+    assert all((p.device.type == "meta") == expected_meta for block in blocks for p in block.parameters())
+    assert model.time_embedder.linear_1.weight.device.type != "meta"
+    assert all(buffer.device.type != "meta" for buffer in model.buffers())
+
+
+@pytest.mark.parametrize("edge", [False, True])
+def test_cosmos3_pre_sharded_meta_processing_accepts_real_unquantized_blocks(edge, monkeypatch):
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import Cosmos3EdgeVFMTransformer, Cosmos3VFMTransformer
+
+    model_cls = Cosmos3EdgeVFMTransformer if edge else Cosmos3VFMTransformer
+    config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
+    model = model_cls(
+        SimpleNamespace(
+            tf_model_config=config(num_hidden_layers=1),
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(use_hsdp=True),
+            hsdp_weight_load_strategy="pre_sharded",
+        )
+    )
+    groups = [SimpleNamespace(module=block) for block in (*model.language_model.layers, *model.gen_layers)]
+    plan = SimpleNamespace(roots=[SimpleNamespace(module=model)], groups=groups)
+    processed = []
+    # Exercise model selection with real vLLM linears, without invoking the
+    # CPU-specific weight packing hook on meta tensors in CPU-only CI.
+    monkeypatch.setattr(
+        UnquantizedLinearMethod, "process_weights_after_loading", lambda self, layer: processed.append(layer)
+    )
+    DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    expected = [module for group in groups for module in group.module.modules() if hasattr(module, "quant_method")]
+    assert expected
+    assert processed == expected
+
+    model._hsdp_pre_sharded_meta_post_load = False
+    with pytest.raises(ValueError, match="has not declared"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    model._hsdp_pre_sharded_meta_post_load = True
+
+    class UnsupportedMethod(UnquantizedLinearMethod):
+        pass
+
+    expected[-1].quant_method = UnsupportedMethod()
+    processed.clear()
+    with pytest.raises(ValueError, match="UnsupportedMethod"):
+        DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
+    assert not processed
+
+
 def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 

tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py

@@ -148,7 +148,8 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
     assert array.shape == (2, 4, 4, 3)
     video = post({"video": frames}, SimpleNamespace(output_type="pt"))
     assert set(video) == {"video"}
-    assert video["video"] is frames
+    assert isinstance(video["video"], np.ndarray)
+    np.testing.assert_array_equal(video["video"], frames.numpy())
     image = post(
         {"image": torch.ones(4, 4, 3)},
         SimpleNamespace(output_type="pt"),
@@ -162,6 +163,40 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
     assert latent_output["image"] is latent
 
 
+@pytest.mark.parametrize("output_type", [None, "pt", "np"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_postprocess_preserves_video_range_for_serving(output_type: str | None, dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+    from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_frames
+
+    # Include black and midtones: all-white frames cannot expose double normalization.
+    frames = torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0], dtype=dtype)
+    frames = frames.reshape(5, 1, 1, 1).expand(5, 2, 4, 3).clone()
+    expected = frames.float().numpy().copy()
+    sampling = None if output_type is None else OmniDiffusionSamplingParams(output_type=output_type)
+    post = get_lingbot_video_post_process_func(None)
+
+    processed = post({"video": frames}, sampling)
+    encoded_frames = np.stack(_coerce_video_to_frames(processed["video"]))
+
+    assert set(processed) == {"video"}
+    assert encoded_frames.dtype == np.float32
+    np.testing.assert_array_equal(encoded_frames, expected)
+    np.testing.assert_array_equal(frames.float().numpy(), expected)
+
+
+@pytest.mark.parametrize("output_key", ["image", "video"])
+def test_postprocess_preserves_latent_payload(output_key: str) -> None:
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+
+    latents = torch.tensor([-2.0, 0.0, 2.0]).reshape(1, 3, 1, 1, 1)
+    sampling = OmniDiffusionSamplingParams(output_type="latent")
+    processed = get_lingbot_video_post_process_func(None)({output_key: latents}, sampling)
+
+    assert set(processed) == {output_key}
+    assert processed[output_key] is latents
+
+
 def test_postprocess_preserves_image_contract_for_serving():
     from vllm_omni.diffusion.data import DiffusionOutput
     from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
@@ -181,7 +216,9 @@ def test_postprocess_preserves_image_contract_for_serving():
         sampling_params=sampling,
         request_id="lingbot-image-contract-test",
     )
-    raw_output = {"image": torch.ones(8, 12, 3)}
+    image = torch.zeros(8, 12, 3)
+    image[4:] = 1.0
+    raw_output = {"image": image}
     processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling)
     normalized = normalize_diffusion_postprocess_output(processed)
 
@@ -199,6 +236,8 @@ def test_postprocess_preserves_image_contract_for_serving():
     endpoint_images = _extract_images_from_result(result)
     assert len(endpoint_images) == 1
     assert endpoint_images[0].size == (12, 8)
+    assert endpoint_images[0].getpixel((0, 0)) == (0, 0, 0)
+    assert endpoint_images[0].getpixel((0, 7)) == (255, 255, 255)
 
 
 def test_forward_resolves_t2v_sampling_and_flow_shift_alias():

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -13,7 +13,6 @@ import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
 from tests.helpers.mark import hardware_marks
-from tests.helpers.monitor import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -55,6 +54,7 @@ IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = {
 }
 
 _OFFLOAD_STATE_PROBE = "vllm_omni_layerwise_offload_test"
+_OFFLOAD_MEMORY_PROBE = "vllm_omni_layerwise_offload_memory_test"
 # A second CI build can reverse the same exact-head measurement without a
 # source edit; the default keeps the existing baseline-first workload.
 _MEASUREMENT_ORDER_ENV = "VLLM_OMNI_OFFLOAD_TEST_ORDER"
@@ -65,7 +65,7 @@ _MEASUREMENT_ORDERS: dict[str, tuple[bool, bool]] = {
 
 
 class OffloadStateProbe:
-    """Test-only worker extension exposing the configured offloader state."""
+    """Test-only worker extension exposing offloader and allocator state."""
 
     model_runner: Any
 
@@ -86,6 +86,21 @@ class OffloadStateProbe:
             "block_count": sum(group_sizes),
         }
 
+    def reset_peak_memory_for_test(self) -> dict[str, str]:
+        """Start a process-local inference peak after model initialization."""
+        current_omni_platform.synchronize()
+        torch.accelerator.reset_peak_memory_stats()
+        return {"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"}
+
+    def get_peak_memory_for_test(self) -> dict[str, str | float]:
+        """Return this worker's allocator peak without counting other jobs."""
+        current_omni_platform.synchronize()
+        return {
+            "probe": _OFFLOAD_MEMORY_PROBE,
+            "operation": "snapshot",
+            "peak_allocated_mb": torch.accelerator.max_memory_allocated() / (1024**2),
+        }
+
 
 def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float = 1e-2) -> bool:
     if not np.allclose(audio1, audio2, atol=atol):
@@ -96,13 +111,6 @@ def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float
     return True
 
 
-def _device_used_mb(device_index: int) -> float:
-    current_omni_platform.synchronize()
-    with current_omni_platform.device(device_index):
-        free_bytes, total_bytes = current_omni_platform.mem_get_info()
-    return (total_bytes - free_bytes) / (1024**2)
-
-
 def _extract_audio(output: Any) -> np.ndarray | None:
     if not output:
         return None
@@ -117,17 +125,29 @@ def _extract_audio(output: Any) -> np.ndarray | None:
     return np.asarray(audio)
 
 
-def _collect_offload_states(value: Any) -> list[dict[str, Any]]:
+def _collect_probe_results(value: Any, probe: str) -> list[dict[str, Any]]:
     if isinstance(value, dict):
-        return [value] if value.get("probe") == _OFFLOAD_STATE_PROBE else []
+        return [value] if value.get("probe") == probe else []
     if isinstance(value, (list, tuple)):
-        states: list[dict[str, Any]] = []
+        results: list[dict[str, Any]] = []
         for item in value:
-            states.extend(_collect_offload_states(item))
-        return states
+            results.extend(_collect_probe_results(item, probe))
+        return results
     return []
 
 
+def _collect_offload_states(value: Any) -> list[dict[str, Any]]:
+    return _collect_probe_results(value, _OFFLOAD_STATE_PROBE)
+
+
+def _collect_memory_results(value: Any, operation: str) -> list[dict[str, Any]]:
+    return [
+        result
+        for result in _collect_probe_results(value, _OFFLOAD_MEMORY_PROBE)
+        if result.get("operation") == operation
+    ]
+
+
 def _measurement_order() -> tuple[bool, bool]:
     """Return the requested order without changing the default CI workload."""
     configured = os.environ.get(_MEASUREMENT_ORDER_ENV, "baseline-first")
@@ -144,8 +164,6 @@ def run_inference(
     num_inference_steps: int = 3,
 ) -> dict[str, Any]:
     current_omni_platform.empty_cache()
-    device_index = current_omni_platform.current_device()
-    initial_used_mb = _device_used_mb(device_index)
 
     if model_name in AUDIO_MODEL:
         params = AUDIO_MODEL_PARAMS
@@ -166,52 +184,45 @@ def run_inference(
         if not offload_states:
             raise AssertionError("The offload-state worker probe returned no diffusion worker results")
 
-        # Measure steady-state inference memory, not model construction. Enabling
-        # layerwise offload first loads the model and then replaces each block's
-        # device storage with CPU-backed weights.  Monitoring that transition
-        # captures both the original model and temporary staging allocations,
-        # which is not representative of layerwise-offloaded inference.
-        monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
-        current_omni_platform.reset_peak_memory_stats()
-        monitor.start()
+        # Reset and read allocator peaks in the diffusion worker process. A
+        # device-wide mem_get_info() sample includes unrelated sibling jobs on
+        # the same GPU and produced negative "savings" in shared ROCm CI.
+        reset_results = _collect_memory_results(
+            runner.omni.engine.collective_rpc(method="reset_peak_memory_for_test", timeout=60),
+            "reset",
+        )
+        if not reset_results:
+            raise AssertionError("The worker peak-memory probe returned no reset acknowledgements")
 
-        try:
-            # Refer to tests/e2e/offline_inference/test_wan22.py
-            # Use minimal settings for testing
-            output = runner.omni.generate(
-                "A cat sitting on a table",
-                OmniDiffusionSamplingParams(
-                    generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
-                    guidance_scale=1.0,
-                    num_inference_steps=num_inference_steps,
-                    **params["sampler_params"],
-                ),
-            )
-        finally:
-            monitor.stop()
+        # Refer to tests/e2e/offline_inference/test_wan22.py.
+        # Use minimal settings for testing.
+        output = runner.omni.generate(
+            "A cat sitting on a table",
+            OmniDiffusionSamplingParams(
+                generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+                guidance_scale=1.0,
+                num_inference_steps=num_inference_steps,
+                **params["sampler_params"],
+            ),
+        )
+        memory_results = _collect_memory_results(
+            runner.omni.engine.collective_rpc(method="get_peak_memory_for_test", timeout=60),
+            "snapshot",
+        )
+        if not memory_results:
+            raise AssertionError("The worker peak-memory probe returned no snapshots")
 
         audio = _extract_audio(output)
         del output
 
-    # DeviceMemoryMonitor reports absolute device usage. Subtract this run's
-    # starting usage. Each mode runs in its own spawned process below, so model
-    # objects and process-local allocator/compiler/backend workspace state cannot
-    # carry into the other measurement. Shared filesystem caches can persist;
-    # _MEASUREMENT_ORDER_ENV enables a reversed-order validation run for that.
-    peak_used_mb = monitor.peak_used_mb
-    incremental_peak_mb = max(0.0, peak_used_mb - initial_used_mb)
-
     del runner
     gc.collect()
     cleanup_dist_env_and_memory()
     current_omni_platform.empty_cache()
-    post_cleanup_used_mb = _device_used_mb(device_index)
 
     return {
-        "initial_used_mb": initial_used_mb,
-        "peak_used_mb": peak_used_mb,
-        "incremental_peak_mb": incremental_peak_mb,
-        "post_cleanup_used_mb": post_cleanup_used_mb,
+        "peak_allocated_mb": max(float(result["peak_allocated_mb"]) for result in memory_results),
+        "memory_results": memory_results,
         "audio": audio,
         "offload_states": offload_states,
     }
@@ -259,12 +270,8 @@ def _assert_audio_outputs(
 
 
 def _print_measurement(label: str, measurement: dict[str, Any]) -> None:
-    print(
-        f"{label}: initial={measurement['initial_used_mb']:.1f} MB, "
-        f"peak={measurement['peak_used_mb']:.1f} MB, "
-        f"incremental_peak={measurement['incremental_peak_mb']:.1f} MB, "
-        f"post_cleanup={measurement['post_cleanup_used_mb']:.1f} MB"
-    )
+    print(f"{label}: worker_peak_allocated={measurement['peak_allocated_mb']:.1f} MB")
+    print(f"{label} worker memory: {measurement['memory_results']}")
     print(f"{label} offload state: {measurement['offload_states']}")
 
 
@@ -298,6 +305,33 @@ def test_measurement_order_rejects_unknown_value(monkeypatch) -> None:
         _measurement_order()
 
 
+@pytest.mark.diffusion
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_collect_memory_results_filters_nested_worker_responses() -> None:
+    value = [
+        None,
+        [
+            {"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"},
+            {
+                "probe": _OFFLOAD_MEMORY_PROBE,
+                "operation": "snapshot",
+                "peak_allocated_mb": 123.0,
+            },
+        ],
+        {"probe": _OFFLOAD_STATE_PROBE, "operation": "snapshot"},
+    ]
+
+    assert _collect_memory_results(value, "reset") == [{"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"}]
+    assert _collect_memory_results(value, "snapshot") == [
+        {
+            "probe": _OFFLOAD_MEMORY_PROBE,
+            "operation": "snapshot",
+            "peak_allocated_mb": 123.0,
+        }
+    ]
+
+
 @pytest.mark.diffusion
 @pytest.mark.core_model
 @pytest.mark.cpu
@@ -366,8 +400,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):
 
     # Verify that layerwise offloading significantly reduces memory usage
     # Passes only if the actual savings meets the expected savings
-    no_offload_peak_memory = no_offload["incremental_peak_mb"]
-    layerwise_offload_peak_memory = layerwise_offload["incremental_peak_mb"]
+    no_offload_peak_memory = no_offload["peak_allocated_mb"]
+    layerwise_offload_peak_memory = layerwise_offload["peak_allocated_mb"]
     actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory
     assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, (
         f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "

tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py

@@ -77,6 +77,7 @@ class _PipelineBackedEngine:
     def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None:
         self.pipeline = pipeline
         self.executor = SimpleNamespace(
+            is_dead=False,
             register_failure_callback=MagicMock(),
             check_health=MagicMock(),
         )
@@ -127,9 +128,9 @@ def _streaming_diffusion_output(
 
 def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]:
     envelope = output.output if isinstance(output.output, dict) else {}
-    payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {}
-    metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {}
-    return payload, metadata
+    payload = envelope.get("payload")
+    metadata = envelope.get("metadata")
+    return payload if isinstance(payload, dict) else {}, metadata if isinstance(metadata, dict) else {}
 
 
 class TestPipelineStreamingOutputToStageDiffusionClient:
@@ -230,7 +231,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
         client._response_socket.bind(response_address)
         client._encoder = OmniMsgpackEncoder()
         client._decoder = OmniMsgpackDecoder()
-        client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
+        client._output_queue = asyncio.Queue[OmniRequestOutput]()
         client._rpc_results = {}
         client._pending_rpcs = set()
         client._tasks = {}

tests/diffusion/test_attention_compile.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
@@ -14,6 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = False
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -38,6 +39,7 @@ def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
 
 def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = False
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -60,8 +62,10 @@ def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
     assert calls == ["impl"]
 
 
-def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
+@pytest.mark.parametrize("initialized_with_hsdp", [False, True])
+def test_attention_compile_boundary_without_diffusion_config(monkeypatch, initialized_with_hsdp):
     attention = object.__new__(Attention)
+    attention._hsdp_compile_boundary_enabled = initialized_with_hsdp
     calls = []
 
     def _boundary(query, key, value, attn_metadata=None):
@@ -80,4 +84,4 @@ def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
     with set_forward_context(vllm_config=None):
         assert Attention.forward(attention, query, query, query) is query
 
-    assert calls == ["impl"]
+    assert calls == (["boundary"] if initialized_with_hsdp else ["impl"])

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -15,6 +15,7 @@ from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import Request, RequestStatus
 
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
@@ -95,6 +96,30 @@ def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatc
     assert not adapter.code_prompt_token_ids
 
 
+def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
+    adapter, connector = shm_sender
+    ext_id = f"release_{uuid.uuid4().hex}"
+    request = _req("done", RequestStatus.RUNNING, external_req_id=ext_id)
+    sibling_key = f"{ext_id}_1_0_0"
+    assert connector.put("0", "1", sibling_key, "sibling")[0]
+    for _ in range(3):
+        adapter.save_async(None, request)
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+    # The consumer read only the first chunk before its own request ended.
+    assert connector.get("0", "1", f"{ext_id}_0_0") is not None
+
+    adapter.release_shm_resources(ext_id)
+    # The release is queued behind sends; nothing is unlinked on the caller.
+    assert f"{ext_id}_0_1" in connector._pending_keys
+    while adapter._pending_save_reqs:
+        adapter._send_single_request(adapter._pending_save_reqs.popleft())
+
+    assert adapter._reclaimed_shm_total == 2
+    assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
+    assert connector.get("0", "1", f"{ext_id}_0_1") is None
+    assert connector.get("0", "1", sibling_key)[0] == "sibling"
+
+
 @pytest.mark.parametrize("fail_terminal", [False, True])
 def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal):
     adapter, connector = shm_sender
@@ -2620,10 +2645,9 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
     adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
 
-    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
-
     scheduler = mocker.MagicMock()
     scheduler.chunk_transfer_adapter = adapter_mock
+    bind_omits_transfer_helpers(scheduler)
     scheduler.connector = None
     scheduler.perf_metrics = None
     scheduler.log_stats = False

tests/distributed/omni_connectors/test_shm_connector.py

@@ -213,7 +213,7 @@ class TestCleanup:
         connector.put("s0", "s1", "cleanup_req_42", data)
         assert "cleanup_req_42" in connector._pending_keys
 
-        connector.cleanup("cleanup_req_42")
+        assert connector.cleanup("cleanup_req_42") is True
         assert "cleanup_req_42" not in connector._pending_keys
 
         result = connector.get("s0", "s1", "cleanup_req_42", metadata=None)
@@ -224,7 +224,7 @@ class TestCleanup:
         connector.put("s0", "s1", "consumed_req_99", data)
         connector.get("s0", "s1", "consumed_req_99", metadata=None)
 
-        connector.cleanup("consumed_req_99")
+        assert connector.cleanup("consumed_req_99") is False
         assert "consumed_req_99" not in connector._pending_keys
 
     def test_close_cleans_all_pending(self, connector):
@@ -296,6 +296,24 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch):
     assert key not in connector._pending_keys
 
 
+@pytest.mark.parametrize("suffix", ["_1_2", "_1", "_0", "_suffix"])
+def test_cleanup_prefix_only_matches_own_chunk_keys(connector, suffix):
+    request_id = f"prefix_{uuid.uuid4().hex}"
+    own_keys = [f"{request_id}_0_{i}" for i in range(3)]
+    sibling_key = f"{request_id}{suffix}_0_0"
+    for key in own_keys:
+        assert connector.put("0", "1", key, key)[0]
+    assert connector.put("0", "1", sibling_key, "sibling")[0]
+    # A consumed chunk is no longer counted as reclaimed.
+    assert connector.get("0", "1", own_keys[0])[0] == own_keys[0]
+
+    assert connector.cleanup_prefix(f"{request_id}_0_") == 2
+    for key in own_keys:
+        assert key not in connector._pending_keys
+        assert connector.get("0", "1", key) is None
+    assert connector.get("0", "1", sibling_key)[0] == "sibling"
+
+
 # ── Arrival wakeups ───────────────────────────────────────────────────
 
 

tests/e2e/features/yue2/worker_extension.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Observe YuE2 lifecycle contracts inside a real engine worker."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+
+class Yue2LifecycleWorkerExtension:
+    model_runner: "GPUARModelRunner"
+
+    def start_yue2_probe(self):
+        from unittest.mock import patch
+
+        from vllm.compilation.cuda_graph import CUDAGraphWrapper
+        from vllm.config import CUDAGraphMode
+        from vllm.forward_context import get_forward_context, is_forward_context_available
+
+        from vllm_omni.model_executor.models.yue2.yue2 import ABC_END, HOLD_TOKEN, MUSIC_END
+
+        model = self.model_runner.get_model()
+        self.yue2_probe = {
+            "preemptions": 0,
+            "rollbacks": 0,
+            "checked_rows": 0,
+            "mixed_steps": 0,
+            "synthesis_started": False,
+            "ar_full_replays": 0,
+            "ar_piecewise_replays": 0,
+        }
+        graph_call = CUDAGraphWrapper.__call__
+
+        def count_ar_replays(wrapper, *args, **kwargs):
+            if is_forward_context_available():
+                context = get_forward_context()
+                mode = context.cudagraph_runtime_mode
+                entry = wrapper.concrete_cudagraph_entries.get(context.batch_descriptor)
+                if mode == wrapper.runtime_mode and entry is not None and entry.cudagraph is not None:
+                    if mode == CUDAGraphMode.FULL:
+                        self.yue2_probe["ar_full_replays"] += 1
+                    elif mode == CUDAGraphMode.PIECEWISE:
+                        self.yue2_probe["ar_piecewise_replays"] += 1
+            return graph_call(wrapper, *args, **kwargs)
+
+        # This worker is created for this test and exits with the engine.
+        # Count vLLM AR replays, separately from model-owned sampler/NAR graphs.
+        self.yue2_graph_patch = patch.object(CUDAGraphWrapper, "__call__", count_ar_replays)
+        self.yue2_graph_patch.start()
+        self.yue2_cancel_job = None
+        previous: dict[str, int] = {}
+        prepare = model.prepare_runner_inputs
+        reconcile = model._reconcile_history
+        queue = model._queue_synthesis
+        sample = model.sample
+        resolve = model._resolve_pending
+        emitted: dict[str, int] = {}
+
+        def resolve_pending():
+            if model._pending is not None:
+                event, host, states = model._pending
+                event.synchronize()
+                for state, token in zip(states, host[0].tolist()):
+                    end = ABC_END if state.constants.phase == "abc" else MUSIC_END
+                    expected = HOLD_TOKEN if token == end and not state.constants.skip_synthesis else token
+                    assert emitted[state.request_id] == expected, "engine token and model history disagree"
+            resolve()
+
+        def sample_tokens(logits, sampling_metadata):
+            result = sample(logits, sampling_metadata)
+            phases = {model._states[rid].constants.phase for rid, _, _ in model._step_rows if rid in model._states}
+            if len(phases) == 2:
+                self.yue2_probe["mixed_steps"] += 1
+            for (rid, _, _), row in zip(model._step_rows, result.sampled_token_ids.tolist()):
+                emitted[rid] = row[0]
+            return result
+
+        def prepare_inputs(**kwargs):
+            for rid, computed in zip(kwargs["req_ids"], kwargs["num_computed_tokens"]):
+                if int(computed) < previous.get(rid, 0):
+                    self.yue2_probe["preemptions"] += 1
+                previous[rid] = int(computed)
+            return prepare(**kwargs)
+
+        def reconcile_history(state, accepted):
+            if len(state.history) + int(state.end_drawn) > accepted:
+                self.yue2_probe["rollbacks"] += 1
+            reconcile(state, accepted)
+            assert len(state.history) + int(state.end_drawn) <= accepted
+            if not state.finish_ready and state.job is None and not state.finished:
+                assert len(state.history) == accepted
+                assert not state.truncated and not state.end_drawn
+            self.yue2_probe["checked_rows"] += 1
+
+        def queue_synthesis(state):
+            queue(state)
+            if "cancel-synthesis" in state.request_id:
+                job = state.job
+                assert job.started and job.work is not None and job.keep
+                assert job.event is None, "test must intercept an active synthesis job"
+                # Keep an actual NAR job active until the driver cancels it.
+                # Its first work units ran normally; stop pumping later units
+                # so cancellation cannot race against a short song finishing.
+                self.yue2_cancel_advance = job.advance
+                job.advance = lambda: None
+                self.yue2_cancel_job = job
+                self.yue2_probe["synthesis_started"] = True
+
+        model.prepare_runner_inputs = prepare_inputs
+        model._reconcile_history = reconcile_history
+        model._queue_synthesis = queue_synthesis
+        model.sample = sample_tokens
+        model._resolve_pending = resolve_pending
+        return True
+
+    def get_yue2_probe(self):
+        result = dict(self.yue2_probe)
+        job = self.yue2_cancel_job
+        if job is not None:
+            if job.cancelled:
+                # Exercise the original pump after cleanup: it must not submit
+                # further units even though this test retains the job object.
+                self.yue2_cancel_advance()
+            model = self.model_runner.get_model()
+            result["cancelled"] = job.cancelled
+            result["released"] = job._released and job.work is None and not job.keep and job.stream is None
+            result["removed"] = (
+                job.request_id not in model._states
+                and model._synthesis.active is not job
+                and job not in model._synthesis.waiting
+            )
+        return result

tests/e2e/offline_inference/test_auk_encoder_cudagraph.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK encoder stage: FULL CUDA graph replay matches eager, request by request.
+
+``auk.yaml`` captures the encoder prefill as FULL CUDA graphs padded to a few
+token buckets. A replay reuses static input buffers and the graph pool, so a
+request could in principle read conditioning left behind by an earlier, longer
+prompt. This test runs the same request sequence through an eager encoder and
+through the graphed one and compares the text condition stage 0 hands to
+stage 1 (the fused thinker hidden states) for every request. The sequence
+alternates prompts of different lengths inside one capture bucket, crosses
+into a larger bucket, and then returns to the first prompt.
+
+Like ``test_auk.py`` it needs an assembled checkpoint directory:
+
+    VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk_encoder_cudagraph.py
+"""
+
+from __future__ import annotations
+
+import bisect
+import os
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+import yaml
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.stage_input_processors import auk as auk_input_processor
+from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params
+
+MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR"
+
+# Text-only instruct prompts: the token count, and so the capture bucket,
+# follows the text length. SHORT_A and SHORT_B differ in length but share a
+# bucket; LONG lands in a larger one. The test checks these properties on the
+# recorded condition lengths, so a tokenizer change fails loudly instead of
+# silently shrinking the coverage.
+_TEMPLATE = 'Generate speech based on the following description: "{desc}". The content to speak is: "{text}".'
+SHORT_A = _TEMPLATE.format(
+    desc="A calm young woman speaking warmly.",
+    text="Welcome back, how was your day? I kept some soup warm for you, and the kettle is on if you want tea.",
+)
+SHORT_B = _TEMPLATE.format(
+    desc="An older man with a deep, slightly raspy voice, speaking slowly.",
+    text=(
+        "The train to the coast leaves at seven, so please be on time. Bring a warm coat, "
+        "some cash for the ferry, and the blue umbrella from the hall, because the forecast says rain."
+    ),
+)
+LONG = _TEMPLATE.format(
+    desc=(
+        "A cheerful radio host in her thirties, bright and energetic, with clear articulation, "
+        "a slight smile in her voice, and a brisk but relaxed pace, recorded close to the microphone"
+    ),
+    text=(
+        "Good morning and welcome to the weekend edition. Today we are talking about small gardens, "
+        "the plants that thrive on a sunny balcony, and the neighbours who turned an empty lot into "
+        "a place where the whole street now meets on Saturday afternoons to trade seedlings and stories. "
+        "Stay with us, because after the news we will open the phone lines and hear from you."
+    ),
+)
+SEQUENCE = [SHORT_A, SHORT_B, SHORT_A, LONG, SHORT_B, SHORT_A]
+
+# The DiT output is not inspected; keep stage 1 as cheap as it goes (it also
+# runs eager in both configs, see _deploy).
+GEN_SECONDS = 1.0
+NFE = 1
+
+# bf16 hidden states. Graph replay pads the prefill to the bucket size but runs
+# the same kernels on the real tokens, so the two paths agree to bf16 rounding.
+# Measured on H20: cosine >= 0.99998 and max-abs error <= 2.9e-3 of the
+# condition's max magnitude for every request; the bounds leave about 3x
+# headroom for other GPUs.
+MIN_COSINE = 0.9999
+MAX_REL_ABS = 1e-2
+
+_model_dir = os.environ.get(MODEL_DIR_ENV)
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"),
+]
+
+
+def _deploy(tmp_path: Path, *, graph_encoder: bool) -> str:
+    """``auk.yaml`` with stage 1 eager, and the encoder stage eager or graphed as shipped.
+
+    Only the encoder differs between the two runs. The conditions compared here
+    are recorded before stage 1 runs, so its DiT compile, graph capture and
+    codec bucket warm-up would only add startup time; eager stage 1 skips them
+    and still takes the real stage handoff.
+    """
+    with open(get_deploy_config_path("auk.yaml")) as f:
+        deploy = yaml.safe_load(f)
+    for stage in deploy["stages"]:
+        if stage["stage_id"] == 1:
+            stage["enforce_eager"] = True
+        elif not graph_encoder:
+            stage["enforce_eager"] = True
+            stage.pop("compilation_config", None)
+    path = tmp_path / f"auk_{'graph' if graph_encoder else 'eager'}_encoder.yaml"
+    path.write_text(yaml.safe_dump(deploy))
+    return str(path)
+
+
+def _graph_capture_sizes() -> list[int]:
+    with open(get_deploy_config_path("auk.yaml")) as f:
+        deploy = yaml.safe_load(f)
+    stage = next(s for s in deploy["stages"] if s["stage_id"] == 0)
+    assert stage.get("enforce_eager") is False, "auk.yaml no longer graphs the encoder; update this test"
+    return sorted(stage["compilation_config"]["cudagraph_capture_sizes"])
+
+
+def _run_sequence(deploy_config: str, monkeypatch: pytest.MonkeyPatch) -> list[torch.Tensor]:
+    """Run SEQUENCE one request at a time and return each request's text condition."""
+    conditions: list[torch.Tensor] = []
+    original = auk_input_processor.encoder2dit
+
+    def recording_encoder2dit(*args: Any, **kwargs: Any) -> dict[str, Any] | None:
+        result = original(*args, **kwargs)
+        if result is not None:
+            conditions.append(result["prompt_embeds"].detach().to("cpu", torch.float32).clone())
+        return result
+
+    # The orchestrator runs in this process and resolves the stage input
+    # function by module attribute at startup, so the patch must precede Omni.
+    with monkeypatch.context() as patch:
+        patch.setattr(auk_input_processor, "encoder2dit", recording_encoder2dit)
+        with OmniRunner(str(Path(_model_dir).resolve()), deploy_config=deploy_config) as runner:
+            for text in SEQUENCE:
+                prompt = auk_prompt(text, None, gen_seconds=GEN_SECONDS)
+                outputs = runner.omni.generate(prompt, auk_sampling_params(nfe=NFE, cfg=2.0, seed=0))
+                assert outputs, "no outputs returned"
+    assert len(conditions) == len(SEQUENCE), f"recorded {len(conditions)} conditions for {len(SEQUENCE)} requests"
+    return conditions
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_encoder_graph_replay_matches_eager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    capture_sizes = _graph_capture_sizes()
+    eager = _run_sequence(_deploy(tmp_path, graph_encoder=False), monkeypatch)
+    graphed = _run_sequence(_deploy(tmp_path, graph_encoder=True), monkeypatch)
+
+    lengths = {text: cond.shape[0] for text, cond in zip(SEQUENCE, eager)}
+    for i, (ref, got) in enumerate(zip(eager, graphed)):
+        assert got.shape == ref.shape, f"request {i}: shape {tuple(got.shape)} != eager {tuple(ref.shape)}"
+        assert torch.isfinite(got).all(), f"request {i}: non-finite condition"
+        cosine = torch.nn.functional.cosine_similarity(got.flatten(), ref.flatten(), dim=0).item()
+        rel_abs = ((got - ref).abs().max() / ref.abs().max()).item()
+        assert cosine >= MIN_COSINE and rel_abs <= MAX_REL_ABS, (
+            f"request {i} ({lengths[SEQUENCE[i]]} tokens): cosine {cosine:.6f}, max-abs error {rel_abs:.2e} "
+            "of the condition's magnitude"
+        )
+
+    # The sequence must exercise what it claims to: two different lengths in
+    # one bucket, a second bucket, and every prompt inside the graphed range.
+    assert max(lengths.values()) <= capture_sizes[-1], f"{lengths} exceeds the largest capture size"
+    bucket = {text: capture_sizes[bisect.bisect_left(capture_sizes, n)] for text, n in lengths.items()}
+    assert lengths[SHORT_A] != lengths[SHORT_B] and bucket[SHORT_A] == bucket[SHORT_B], (lengths, bucket)
+    assert bucket[LONG] > bucket[SHORT_A], (lengths, bucket)
+
+    # Repeats of a prompt reproduce its first condition exactly, eager or graphed:
+    # nothing from the requests in between leaks into the replay.
+    for run in (eager, graphed):
+        first: dict[str, torch.Tensor] = {}
+        for text, cond in zip(SEQUENCE, run):
+            if text in first:
+                torch.testing.assert_close(cond, first[text], atol=0, rtol=0)
+            else:
+                first[text] = cond

tests/e2e/online_serving/test_qwen3_omni_shm_leak.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Live /dev/shm leak check for Qwen3-Omni ``async_chunk`` (stage-0-final).
+
+The default L2 Qwen3-Omni job uses ``ci/qwen3_omni_moe.yaml`` with
+``async_chunk: false`` plus ``--no-async-chunk``, so it never exercises the
+thinker→talker SharedMemoryConnector path.
+
+A text-only request is tagged ``omni_final_stage_id=0``, so nothing downstream
+ever ``get()``s what stage-0 ``put()``s for it. Before #7245 the scheduler still
+``put()`` a finished marker on finish and nothing unlinked it, so every such
+request left a ``SharedMemoryConnector`` segment plus lockfile in ``/dev/shm``.
+Stage-0-final finishes now skip the ``put()`` and the orchestrator reclaims any
+undrained segment on request cleanup.
+
+This module starts the production async-chunk deploy and asserts the three
+text-only request ids leave no connector lockfile or POSIX segment in
+``/dev/shm``.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import Path
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1"
+_MODEL = os.environ.get("VLLM_OMNI_TEST_MODEL", "Qwen/Qwen3-Omni-30B-A3B-Instruct")
+_SHM_DIR = Path("/dev/shm")
+_NUM_TEXT_REQUESTS = 3
+_SETTLE_S = 15.0
+
+# CI overlay keeps seq/token limits small; flip async_chunk back on so stage-0
+# still uses SharedMemoryConnector. Production yaml is already async_chunk:true.
+_ASYNC_CHUNK_DEPLOY = modify_stage_config(
+    get_deploy_config_path("ci/qwen3_omni_moe.yaml"),
+    updates={"async_chunk": True},
+)
+
+test_params = [
+    pytest.param(
+        OmniServerParams(
+            model=_MODEL,
+            stage_config_path=_ASYNC_CHUNK_DEPLOY,
+            use_stage_cli=True,
+        ),
+        id="async_chunk",
+    )
+]
+
+
+def _system_prompt() -> dict:
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def _shm_used_mib() -> str:
+    try:
+        st = os.statvfs(_SHM_DIR)
+        used = (st.f_blocks - st.f_bavail) * st.f_frsize
+        return f"{used / (1024 * 1024):.1f}MiB"
+    except OSError:
+        return "unknown"
+
+
+def _is_entry_for_request(name: str, request_id: str) -> bool:
+    """Match a segment or ``shm_`` lockfile keyed by the ``{request_id}-{suffix}`` orchestrator id."""
+    return name.removeprefix("shm_").startswith((f"{request_id}-", f"{request_id}_"))
+
+
+def _entries_for_request_ids(request_ids: list[str]) -> list[str]:
+    if not _SHM_DIR.is_dir():
+        pytest.skip("/dev/shm is not available")
+    found: list[str] = []
+    with os.scandir(_SHM_DIR) as entries:
+        for entry in entries:
+            if any(_is_entry_for_request(entry.name, request_id) for request_id in request_ids):
+                found.append(entry.name)
+    return found
+
+
+def _wait_no_entries_for_request_ids(request_ids: list[str], timeout_s: float = _SETTLE_S) -> list[str]:
+    """Reclaim is fire-and-forget; poll until these ids are gone or *timeout_s*."""
+    deadline = time.monotonic() + timeout_s
+    time.sleep(1.0)
+    leftover: list[str] = _entries_for_request_ids(request_ids)
+    while leftover:
+        if time.monotonic() >= deadline:
+            return leftover
+        time.sleep(2.0)
+        leftover = _entries_for_request_ids(request_ids)
+    return leftover
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_only_async_chunk_does_not_leak_shm(omni_server, online_client) -> None:
+    """Stage-0-final (text-only) requests must not leave connector SHM behind."""
+    print(f"[shm-leak] before requests: /dev/shm_used={_shm_used_mib()}", flush=True)
+
+    messages = dummy_messages_from_mix_data(
+        system_prompt=_system_prompt(),
+        content_text="What is the capital of China? Answer in 20 words.",
+    )
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+        "key_words": {"text": ["beijing"]},
+    }
+    responses = online_client.send_omni_request(request_config, request_num=_NUM_TEXT_REQUESTS)
+    produced = [resp.text_content for resp in responses if getattr(resp, "text_content", None)]
+    assert len(produced) == _NUM_TEXT_REQUESTS, (
+        f"Need {_NUM_TEXT_REQUESTS} text completions to exercise stage-0 finish; "
+        f"got {len(produced)} non-empty texts. A preprocess/engine miss would "
+        f"leave /dev/shm unchanged and hide a leak."
+    )
+    request_ids = [resp.request_id for resp in responses]
+    assert all(request_ids) and len(request_ids) == _NUM_TEXT_REQUESTS, (
+        f"Need chat completion ids to assert per-request SHM cleanup; got {request_ids}"
+    )
+
+    leftover = _wait_no_entries_for_request_ids(request_ids)
+    print(
+        f"[shm-leak] after {_NUM_TEXT_REQUESTS} text-only requests: "
+        f"/dev/shm_used={_shm_used_mib()} leftover={leftover} request_ids={request_ids}",
+        flush=True,
+    )
+    assert not leftover, (
+        f"SharedMemoryConnector leaked /dev/shm entries for {leftover} after stage-0-final request(s) {request_ids}"
+    )

tests/e2e/online_serving/test_yue2.py

@@ -32,9 +32,9 @@ above the byte floor) and passes on that alone, following the MiniMax Music 3
 precedent.
 
 Weekly rather than per-PR: each request is a 16 s song and the verification
-checkpoints add ~5.5 GB of downloads. Validated on a single RTX 4090
-(24 GB); the hardware mark targets an H100 because the CI fleet has no
-consumer cards.
+checkpoints add ~5.5 GB of downloads. The async/graph implementation was
+validated on H200; the hardware mark targets the CI fleet's H100. The older
+RTX 4090 result does not establish this implementation's memory requirements.
 """
 
 import concurrent.futures
@@ -60,8 +60,7 @@ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 900.0
 # 400 frames = 16 s of song. The byte floor sits far below the observed
 # payload (~3.0 MB of 48 kHz stereo 16-bit) to leave room for an early EOS.
 _FRAMES = 400
-# Concurrent pair stays short: it exists to exercise row re-indexing, not
-# to double the wall time.
+# Short concurrent requests exercise row re-indexing as budgets differ.
 _CONCURRENT_FRAMES = 250
 _MIN_BYTES = 1_000_000
 
@@ -173,7 +172,7 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
     # CI runs without the verification assets, so the structural checks below
     # are the only guard against the two real bugs this PR fixed during review:
     # mono/half-speed output (WAV header + duration bounds) and the int32
-    # crash under concurrent row re-indexing (two simultaneous requests).
+    # crash under concurrent row re-indexing (simultaneous requests).
     # Duration bounds: the end token is masked for the first 200 steps
     # (min_tokens), and the 400-frame budget caps the song, so a healthy
     # request lands in [8, 16.5] s; the half-speed bug produced ~32 s.
@@ -190,13 +189,15 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
     # (see _http_speech); the OpenAI SDK response does not expose headers.
     _http_speech(omni_server, seed=SEED, frames=_FRAMES)
 
-    # Two concurrent requests with different seeds: both must succeed and
-    # produce different songs (the model-owned sampler re-indexes rows across
-    # concurrent requests; a dtype mismatch here crashed the engine before).
-    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
-        futures = [pool.submit(_http_speech, omni_server, seed=SEED + i, frames=_CONCURRENT_FRAMES) for i in range(2)]
-        pair = [f.result() for f in futures]
-    assert pair[0] != pair[1], "different seeds produced identical audio"
+    # Different budgets exercise row re-indexing as shorter requests finish.
+    # Compare the same-budget pair so length cannot hide an ignored seed.
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool:
+        futures = [
+            pool.submit(_http_speech, omni_server, seed=SEED + i, frames=frames)
+            for i, frames in enumerate((_CONCURRENT_FRAMES, _FRAMES, _CONCURRENT_FRAMES))
+        ]
+        audios = [f.result() for f in futures]
+    assert audios[0] != audios[2], "different seeds produced identical audio at the same frame budget"
 
     # Without the verification assets the structural assertions above
     # (non-empty payload, byte floor) are the whole test — pass on them.
@@ -235,3 +236,199 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
     assert melody_score > control_score, (
         f"calibration: ABC-conditioned {melody_score:.3f} not above control {control_score:.3f}"
     )
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_yue2_abort_then_next_request_succeeds(omni_server) -> None:
+    """Serving smoke check after a client timeout, not proof of job cleanup."""
+    # The chosen seed and budget time out on the tested cards. Require the
+    # timeout so this cannot silently become another successful request;
+    # whether it lands in AR or synthesis depends on the card.
+    with pytest.raises(httpx.TimeoutException):
+        httpx.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech",
+            json={
+                "model": omni_server.model,
+                "input": LYRICS,
+                "instructions": CAPTION,
+                "seed": SEED + 2,
+                "max_new_tokens": 3000,
+                "stream": False,
+                "response_format": "wav",
+                "extra_params": {"cot": "off"},
+            },
+            timeout=1.0,
+        )
+    _http_speech(omni_server, seed=SEED + 3, frames=_CONCURRENT_FRAMES)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("chunked", [True, False], ids=["chunked", "unchunked"])
+@pytest.mark.parametrize("enforce_eager", [True, False], ids=["eager", "ar-graphs"])
+async def test_yue2_preemption_and_synthesis_abort(chunked, enforce_eager, monkeypatch) -> None:
+    """Force real preemption and abort only after NAR work has been submitted."""
+    import asyncio
+    from pathlib import Path
+
+    import torch
+    from vllm import SamplingParams
+
+    from tests.helpers.runtime import get_model_prefix
+    from tests.helpers.stage_config import modify_stage_config
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.model_executor.models.yue2.yue2 import (
+        ABC_END,
+        EOD,
+        KEY_MAX_AUDIO_FRAMES,
+        KEY_MAX_HOLD_STEPS,
+        KEY_MIN_TOKENS,
+        KEY_PHASE,
+        KEY_PREFIX_IDS,
+        KEY_SEED,
+        KEY_SKIP_SYNTHESIS,
+        STOP_TOKEN_IDS,
+    )
+    from vllm_omni.tokenizers.yue2_prompt import abc_prefix_ids, semantic_frames, semantic_prefix_ids
+    from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
+    from vllm_omni.transformers_utils.repo_utils import hf_api
+
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    model_path = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + MODEL))
+    if not model_path.is_dir():
+        model_path = Path(hf_api().snapshot_download(str(model_path)))
+    tokenizer = YuE2TextTokenizer(model_path / "qwen.tiktoken")
+    prefix = semantic_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "off")
+    frames = 256
+    # Each song fits alone, but four full songs cannot fit in 64 KV blocks.
+    # Chunked recompute also replays output tokens past the original prompt.
+    deploy = modify_stage_config(
+        get_deploy_config_path("yue2.yaml"),
+        updates={
+            "stages": {
+                0: {
+                    "enforce_eager": enforce_eager,
+                    "compilation_config": {
+                        "cudagraph_mode": "NONE" if enforce_eager else "FULL_AND_PIECEWISE",
+                    },
+                    "enable_chunked_prefill": chunked,
+                    "async_scheduling": True,
+                    "max_num_batched_tokens": 256 if chunked else 1024,
+                    "max_model_len": 1000,
+                    "num_gpu_blocks_override": 64,
+                    "block_size": 16,
+                    "max_num_seqs": 4,
+                    "tokenizer": str(model_path),
+                }
+            }
+        },
+    )
+    engine = AsyncOmni(
+        model=str(model_path),
+        deploy_config=deploy,
+        worker_extension_cls="tests.e2e.features.yue2.worker_extension.Yue2LifecycleWorkerExtension",
+        stage_init_timeout=600,
+        init_timeout=900,
+    )
+
+    async def probe():
+        return (await engine.collective_rpc("get_yue2_probe", stage_ids=[0]))[0][0]
+
+    async def song(request_id, seed):
+        params = SamplingParams(
+            max_tokens=1000 - len(prefix),
+            stop_token_ids=STOP_TOKEN_IDS,
+            detokenize=False,
+            extra_args={
+                KEY_PREFIX_IDS: prefix,
+                KEY_SEED: seed,
+                KEY_MAX_AUDIO_FRAMES: frames,
+                KEY_MIN_TOKENS: frames,
+                KEY_MAX_HOLD_STEPS: 1000 - len(prefix) - frames - 1,
+            },
+        )
+        audio = []
+        tokens = []
+        async for output in engine.generate(
+            {"prompt_token_ids": prefix}, request_id=request_id, sampling_params_list=[params]
+        ):
+            tokens = list(output.outputs[0].token_ids)
+            multimodal = getattr(output.outputs[0], "multimodal_output", None)
+            if multimodal is not None and "audio" in multimodal:
+                data = multimodal["audio"]
+                audio.extend(data if isinstance(data, list) else [data])
+        assert audio, "request completed without audio"
+        waveform = torch.cat([part.detach().cpu().reshape(-1) for part in audio])
+        assert torch.isfinite(waveform).all() and waveform.abs().max() > 0
+        assert len(semantic_frames(tokens)) == frames
+        # The VAE convolution trims a small number of boundary samples.
+        assert waveform.numel() == pytest.approx(frames * 1920 * 2, abs=256)
+
+    async def plan():
+        abc_prefix = abc_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "melody")
+        params = SamplingParams(
+            max_tokens=64,
+            stop_token_ids=STOP_TOKEN_IDS,
+            detokenize=False,
+            extra_args={
+                KEY_PHASE: "abc",
+                KEY_PREFIX_IDS: abc_prefix,
+                KEY_SEED: SEED,
+                KEY_SKIP_SYNTHESIS: True,
+            },
+        )
+        tokens = []
+        async for output in engine.generate(
+            {"prompt_token_ids": abc_prefix}, request_id="abc-plan", sampling_params_list=[params]
+        ):
+            tokens = list(output.outputs[0].token_ids)
+        assert tokens and all(t < EOD or t == ABC_END for t in tokens)
+
+    tasks = []
+    try:
+        assert await engine.collective_rpc("start_yue2_probe", stage_ids=[0]) == [[True]]
+        tasks = [asyncio.create_task(plan())]
+        tasks.extend(asyncio.create_task(song(f"preempt-{i}", SEED + i)) for i in range(4))
+        await asyncio.wait_for(asyncio.gather(*tasks), timeout=600)
+        observations = await probe()
+        assert observations["preemptions"] > 0, observations
+        assert observations["rollbacks"] > 0, observations
+        assert observations["checked_rows"] >= 4 * frames, observations
+        assert observations["mixed_steps"] > 0, observations
+        if enforce_eager:
+            assert observations["ar_full_replays"] == 0, observations
+        else:
+            assert observations["ar_full_replays"] > 0, observations
+
+        cancelled = asyncio.create_task(song("cancel-synthesis", SEED + 7))
+        tasks.append(cancelled)
+        deadline = asyncio.get_running_loop().time() + 120
+        while not (await probe())["synthesis_started"]:
+            assert not cancelled.done(), "song finished before the active-job cancellation probe"
+            assert asyncio.get_running_loop().time() < deadline, "NAR job never started"
+            await asyncio.sleep(0.05)
+        cancelled.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await cancelled
+        deadline = asyncio.get_running_loop().time() + 30
+        while True:
+            observations = await probe()
+            if observations.get("removed"):
+                assert observations["cancelled"] and observations["released"], observations
+                break
+            assert asyncio.get_running_loop().time() < deadline, observations
+            await asyncio.sleep(0.05)
+        await asyncio.wait_for(song("after-cancel", SEED + 8), timeout=120)
+        observations = await probe()
+        assert observations["released"]
+        print("YuE2 lifecycle:", observations)
+    finally:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        engine.shutdown()

tests/engine/test_cfg_companion_lifecycle.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Deterministic tests for the CFG-companion output lifecycle.
 
 Regression suite for the companion-output race behind the nightly Bagel
@@ -94,6 +94,9 @@ class _FakePool:
     def release_bindings(self, request_ids):
         pass
 
+    async def release_request_resources(self, request_ids):
+        pass
+
 
 def _make_orchestrator(num_stages: int = 2) -> Orchestrator:
     orch = Orchestrator.__new__(Orchestrator)

tests/engine/test_orchestrator_first_audio.py

@@ -114,10 +114,10 @@ def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)):
     )
 
 
-async def _codec_output(obj, raw):
-    replica = obj.stage_pools[1].replica_id
-    processed = await obj._process_llm_stage_outputs(1, replica, raw, set())
-    await obj._handle_processed_outputs(1, replica, processed)
+async def _codec_output(obj, raw, stage=1):
+    replica = obj.stage_pools[stage].replica_id
+    processed = await obj._process_llm_stage_outputs(stage, replica, raw, set())
+    await obj._handle_processed_outputs(stage, replica, processed)
 
 
 def _messages(obj):
@@ -132,6 +132,57 @@ def _audio(output):
     return torch.cat(audio) if isinstance(audio, list) else audio
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", list(RequestOutputKind))
+@pytest.mark.parametrize("suffix_overtook", [False, True])
+async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kind, suffix_overtook):
+    obj = _orchestrator(kind, final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    suffix = _raw(required=True, samples=(20, 21))
+    terminal = _raw(required=True, terminal=True, samples=(30, 31))
+    if suffix_overtook:
+        await _codec_output(obj, suffix, stage=0)
+        await _codec_output(obj, terminal, stage=0)
+        assert obj.output_async_queue.empty()
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True, samples=(10, 11)))
+    if not suffix_overtook:
+        await _codec_output(obj, suffix, stage=0)
+        await _codec_output(obj, terminal, stage=0)
+    messages = _messages(obj)
+    expected = {
+        RequestOutputKind.DELTA: [[10, 11], [20, 21], [30, 31]],
+        RequestOutputKind.CUMULATIVE: [[10, 11], [10, 11, 20, 21], [10, 11, 20, 21, 30, 31]],
+        RequestOutputKind.FINAL_ONLY: [[10, 11, 20, 21, 30, 31]],
+    }
+    assert [_audio(message.engine_outputs).tolist() for message in messages] == expected[kind]
+    assert all((message.stage_id, message.replica_id) == (0, 2) for message in messages)
+    assert messages[-1].finished
+    assert not obj.request_states["r"].pending_first_audio_outputs
+
+
+@pytest.mark.asyncio
+async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_first_audio():
+    obj = _orchestrator(final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
+    assert obj.request_states["r"].pending_first_audio_outputs
+    await obj._cleanup_request_ids(["r"], abort=True)
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True))
+    assert not obj.request_states and obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_single_stage_first_audio_processing_failure_releases_waiting_terminal():
+    obj = _orchestrator(final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
+    obj.stage_pools[0].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True))
+    error = obj.output_async_queue.get_nowait()
+    assert isinstance(error, ErrorMessage) and error.stage_id == 0
+    assert not obj.request_states
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("kind", list(RequestOutputKind))
 @pytest.mark.parametrize("codec_overtook", [False, True])

tests/engine/test_stage_engine_core_proc.py

@@ -1,9 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 from unittest.mock import patch
 
+import pytest
 from vllm.v1.engine.core import EngineCoreProc
+from vllm.v1.executor.uniproc_executor import UniProcExecutor
 
-from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc
+from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc, _bind_first_audio_sink
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    "first_decoder,stream_decoder,stream_first_audio,expected",
+    [(True, False, False, True), (False, True, False, False), (False, True, True, True), (False, False, True, False)],
+)
+def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder(
+    mocker,
+    first_decoder,
+    stream_decoder,
+    stream_first_audio,
+    expected,
+):
+    executor = UniProcExecutor.__new__(UniProcExecutor)
+    executor.driver_worker = mocker.Mock()
+    runner = executor.driver_worker.worker.model_runner
+    runner.model.first_frame_decoder = object() if first_decoder else None
+    runner.model.stream_decoder = object() if stream_decoder else None
+    runner.model.stream_first_audio = stream_first_audio
+    assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected
+    assert runner.model_state.set_first_audio_sink.call_count == int(expected)
 
 
 def test_preprocess_add_request_preserves_omni_fields():

tests/engine/test_stage_pool_collective_rpc.py

@@ -154,6 +154,46 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_release_request_resources_skips_without_async_chunk():
+    async def run() -> None:
+        call = AsyncMock()
+        pool = StagePool(
+            0,
+            [SimpleNamespace(call_utility_async=call)],  # type: ignore[list-item]
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=False)),
+        )
+
+        await pool.release_request_resources(["req-1"])
+        call.assert_not_awaited()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch):
+    async def run() -> None:
+        async def hang(*_args):
+            await asyncio.Event().wait()
+
+        live = AsyncMock()
+        pool = StagePool(
+            0,
+            [SimpleNamespace(call_utility_async=hang), SimpleNamespace(call_utility_async=live)],  # type: ignore[list-item]
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)),
+        )
+        monkeypatch.setattr(pool, "RELEASE_RPC_TIMEOUT_S", 0.2)
+
+        release = asyncio.create_task(pool.release_request_resources(["req-1"]))
+        await asyncio.sleep(0.05)
+        live.assert_awaited_once_with("omni_release_request_resources", ["req-1"])
+        assert not release.done()
+
+        await asyncio.wait_for(release, timeout=1.0)
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
 @pytest.mark.parametrize("failure", ["error", "timeout", "missing"])

tests/entrypoints/openai_api/test_audio_format.py

@@ -346,3 +346,19 @@ def test_chat_audio_metadata_survives_response_serialization(stream, fmt):
             assert audio_file.samplerate == 22050
     else:
         assert len(raw) == frame_count * 2
+
+
+def test_float32_pcm_fast_path_matches_soundfile():
+    from vllm_omni.entrypoints.openai.audio_utils_mixin import _float32_to_pcm16_bytes
+
+    rng = np.random.default_rng(0)
+    audio = np.concatenate(
+        [
+            rng.uniform(-1.2, 1.2, 4096),
+            np.array([-1.0, 1.0, 0.0, -0.0, 32767 / 32768, -1 / 32768, 0.5 / 32768, -0.5 / 32768]),
+        ]
+    ).astype(np.float32)
+    with BytesIO() as buffer:
+        soundfile.write(buffer, audio, 24000, format="RAW", subtype="PCM_16")
+        expected = buffer.getvalue()
+    assert _float32_to_pcm16_bytes(audio) == expected

tests/entrypoints/openai_api/test_serving_video_stream.py

@@ -7,10 +7,12 @@ from __future__ import annotations
 import asyncio
 import base64
 import gc
+import hashlib
 import io
 import json
 import threading
 import weakref
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
@@ -26,7 +28,9 @@ from vllm_omni.entrypoints.openai.serving_video_stream import (
     QwenOmniStreamingVideoHandler,
     StreamingVideoSessionConfig,
 )
-from vllm_omni.entrypoints.openai.video_stream_base import OmniStreamingVideoHandler
+from vllm_omni.entrypoints.openai.video_stream_base import (
+    OmniStreamingVideoHandler,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -506,17 +510,9 @@ async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot(
         pytest.param(2, ["f0", "f1", "f2"], True, [0, 2], id="repeated-image"),
     ],
 )
-async def test_video_frame_selection_preserves_uncached_frames(
-    monkeypatch, num_frames, frame_ids, repeat_image, expected_indices
-):
-    """Prewarm misses still use image URLs, with metadata paired by buffer position."""
+async def test_video_frame_selection_pairs_metadata_by_position(num_frames, frame_ids, repeat_image, expected_indices):
+    """Queries wait for prewarmed PIL frames, with metadata paired by buffer position."""
     engine_prompts = []
-    pending_decode = asyncio.Event()
-
-    async def blocked_to_thread(*args, **kwargs):
-        await pending_decode.wait()
-
-    monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread)
 
     class OneOutputEngine:
         async def generate(self, **kwargs):
@@ -550,13 +546,13 @@ async def test_video_frame_selection_preserves_uncached_frames(
     assert not [message for message in ws.sent if message["type"] == "error"]
     assert len(engine_prompts) == 1
     content = engine_prompts[0]["messages"][-1]["content"]
-    assert content == [
-        *[
-            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{frames[index]}"}}
-            for index in expected_indices
-        ],
-        {"type": "text", "text": "describe"},
-    ]
+    assert content[-1] == {"type": "text", "text": "describe"}
+    image_parts = content[:-1]
+    assert [part["type"] for part in image_parts] == ["image_pil"] * len(expected_indices)
+    for part, index in zip(image_parts, expected_indices):
+        raw = base64.b64decode(frames[index])
+        assert part["image_pil"].tobytes() == Image.open(io.BytesIO(raw)).convert("RGB").tobytes()
+        assert part["uuid"] == hashlib.md5(raw, usedforsecurity=False).hexdigest()
     consumed_events = [message for message in ws.sent if message["type"] == "video.frames.consumed"]
     if not any(frame_ids):
         assert consumed_events == []
@@ -909,6 +905,9 @@ async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatc
 
     class BlockingEngine:
         async def generate(self, *, request_id, **kwargs):
+            submitted = kwargs.get("submitted")
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
             if request_ids:
                 assert aborted[-1] == request_ids[-1]
             request_ids.append(request_id)
@@ -1004,7 +1003,7 @@ async def test_video_done_waits_for_in_flight_query():
 
 
 @pytest.mark.asyncio
-async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
+async def test_query_waits_for_frame_prewarm(monkeypatch):
     decode_started = threading.Event()
     release_decode = threading.Event()
     query_started = asyncio.Event()
@@ -1035,14 +1034,80 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
     assert decode_started.is_set()
 
     ws.put({"type": "video.query", "text": "describe"})
-    await asyncio.wait_for(query_started.wait(), timeout=2.0)
+    await asyncio.sleep(0.05)
+    assert not query_started.is_set()
 
     release_decode.set()
+    await asyncio.wait_for(query_started.wait(), timeout=2.0)
     ws.put({"type": "video.done"})
     await asyncio.wait_for(task, timeout=2.0)
     assert "session.done" in ws.sent_types()
 
 
+@pytest.mark.asyncio
+async def test_legacy_query_does_not_wait_for_unselected_frame(monkeypatch):
+    """APC off and num_frames=1: a blocked older frame must not delay the query."""
+    frame_a = _make_jpeg(1, 0, 0)
+    frame_b = _make_jpeg(0, 1, 0)
+    decode_a_started = threading.Event()
+    release_a = threading.Event()
+    decode_b_done = threading.Event()
+    query_started = asyncio.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes == frame_a:
+            decode_a_started.set()
+            release_a.wait(timeout=5.0)
+        image = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+        if raw_bytes == frame_b:
+            decode_b_done.set()
+        return image
+
+    class RecordingHandler(QwenOmniStreamingVideoHandler):
+        async def _process_query(self, *args, **kwargs):
+            query_started.set()
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    ws = TimedWebSocket()
+    handler = RecordingHandler(chat_service=object(), engine_client=object(), idle_timeout=5.0)
+    task = asyncio.create_task(handler.handle_session(ws))
+    try:
+        ws.put(
+            {
+                "type": "session.config",
+                "model": "test",
+                "modalities": ["text"],
+                "num_frames": 1,
+                "enable_frame_filter": False,
+            }
+        )
+        await asyncio.sleep(0)
+        ws.put({"type": "video.frame", "data": _b64(frame_a), "frame_id": "A"})
+        for _ in range(100):
+            if decode_a_started.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert decode_a_started.is_set()
+        ws.put({"type": "video.frame", "data": _b64(frame_b), "frame_id": "B"})
+        for _ in range(200):
+            if decode_b_done.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert decode_b_done.is_set()
+
+        ws.put({"type": "video.query", "text": "describe"})
+        await asyncio.wait_for(query_started.wait(), timeout=2.0)
+        assert not release_a.is_set()
+        release_a.set()
+        ws.put({"type": "video.done"})
+        await asyncio.wait_for(task, timeout=2.0)
+    finally:
+        release_a.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("max_frames", [1, 2], ids=["evicted", "retained"])
 async def test_frame_prewarm_only_keeps_retained_images(monkeypatch, max_frames):
@@ -1272,3 +1337,452 @@ def test_build_messages_keeps_recent_history_text_only():
     assert messages[1] == {"role": "assistant", "content": "recent answer"}
     assert messages[2] == user_message
     assert user_message["content"][-1] == {"type": "text", "text": "current question"}
+
+
+# ---------------------------------------------------------------------------
+# Incremental prefill (warmup) lifecycle
+# ---------------------------------------------------------------------------
+
+
+class RecordingReuseEngine:
+    """Engine stub with stage-0 prefix caching on; records generate() and abort()."""
+
+    def __init__(self, warmup_delay: float = 0.0, query_delay: float = 0.0):
+        self.requests: list[dict[str, Any]] = []
+        self.aborted: list[str] = []
+        self._warmup_delay = warmup_delay
+        self._query_delay = query_delay
+        self.engine = SimpleNamespace(
+            stage_vllm_configs=[SimpleNamespace(cache_config=SimpleNamespace(enable_prefix_caching=True))]
+        )
+
+    def warmup_ids(self) -> list[str]:
+        return [r["request_id"] for r in self.requests if r["request_id"].startswith("video-warmup-")]
+
+    def query_ids(self) -> list[str]:
+        return [r["request_id"] for r in self.requests if not r["request_id"].startswith("video-warmup-")]
+
+    def generate(self, **kwargs):
+        self.requests.append(kwargs)
+        submitted = kwargs.get("submitted")
+        if submitted is not None and not submitted.done():
+            submitted.set_result(None)
+        is_warmup = kwargs["request_id"].startswith("video-warmup-")
+        delay = self._warmup_delay if is_warmup else self._query_delay
+
+        async def _gen():
+            if delay:
+                await asyncio.sleep(delay)
+            yield _text_result("ok")
+
+        return _gen()
+
+    async def abort(self, request_id: str) -> None:
+        self.aborted.append(request_id)
+
+
+class PromptRecordingHandler(QwenOmniStreamingVideoHandler):
+    """Records each query turn's user-message content and processor kwargs."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.query_contents: list[list[dict[str, Any]]] = []
+        self.warmup_contents: list[list[dict[str, Any]]] = []
+        self.mm_processor_kwargs: list[Any] = []
+
+    async def _preprocess_to_engine_prompt(self, request):
+        self.mm_processor_kwargs.append(getattr(request, "mm_processor_kwargs", None))
+        return {"prompt": "p"}
+
+    def build_engine_prompt(
+        self, config, frame_buffer, audio_buffer, message_history, query_text, prewarmed_frames, *, frame_indices=None
+    ):
+        messages, user_message = super().build_engine_prompt(
+            config,
+            frame_buffer,
+            audio_buffer,
+            message_history,
+            query_text,
+            prewarmed_frames,
+            frame_indices=frame_indices,
+        )
+        self.query_contents.append(list(user_message["content"]))
+        return messages, user_message
+
+    def build_engine_prompt_prefix(self, config, frame_buffer, message_history, prewarmed_frames):
+        messages = super().build_engine_prompt_prefix(config, frame_buffer, message_history, prewarmed_frames)
+        if messages:
+            self.warmup_contents.append(list(messages[-1]["content"]))
+        return messages
+
+
+async def _poll(predicate, timeout: float = 3.0) -> bool:
+    for _ in range(int(timeout / 0.01)):
+        if predicate():
+            return True
+        await asyncio.sleep(0.01)
+    return bool(predicate())
+
+
+@pytest.mark.asyncio
+async def test_warmup_fires_on_frame_and_query_cancels_it():
+    engine = RecordingReuseEngine(warmup_delay=10.0, query_delay=0.3)
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+
+    assert await _poll(lambda: engine.warmup_ids())
+    warmup_id = engine.warmup_ids()[0]
+    assert engine.requests[0]["sampling_params"].max_tokens == 1
+
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: engine.query_ids())
+    # in-flight warmup was cancelled + aborted, and not restarted mid-query
+    assert warmup_id in engine.aborted
+    assert len(engine.warmup_ids()) == 1
+    assert handler.mm_processor_kwargs[:2] == [
+        {"use_audio_in_video": True},
+        {"use_audio_in_video": True},
+    ]
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    # the committed turn changed the context -> a fresh warmup fires
+    assert await _poll(lambda: len(engine.warmup_ids()) == 2)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_warmup_reaches_stage0_before_abort():
+    """Cancel during submit must not abort a warmup the engine never received."""
+
+    class _HoldSubmitEngine(RecordingReuseEngine):
+        def __init__(self):
+            super().__init__()
+            self.started = asyncio.Event()
+            self.release_submit = asyncio.Event()
+
+        def generate(self, **kwargs):
+            submitted = kwargs.get("submitted")
+
+            async def _gen():
+                self.started.set()
+                await self.release_submit.wait()
+                self.requests.append(kwargs)
+                if submitted is not None and not submitted.done():
+                    submitted.set_result(None)
+                yield _text_result("ok")
+
+            return _gen()
+
+    engine = _HoldSubmitEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(engine.started.is_set)
+    ws.put({"type": "video.query", "text": "describe"})
+    await asyncio.sleep(0.05)
+    assert engine.requests == []
+    assert engine.aborted == []
+    engine.release_submit.set()
+    assert await _poll(lambda: engine.aborted)
+    assert engine.warmup_ids()
+    assert engine.aborted[0] == engine.warmup_ids()[0]
+    assert await _poll(lambda: engine.query_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_query_reaches_stage0_before_abort():
+    """Cancel during query submit must not abort a request the engine never received."""
+
+    class _HoldQueryEngine(RecordingReuseEngine):
+        def __init__(self):
+            super().__init__()
+            self.query_started = asyncio.Event()
+            self.release_submit = asyncio.Event()
+
+        def generate(self, **kwargs):
+            if str(kwargs.get("request_id", "")).startswith("video-warmup-"):
+                return super().generate(**kwargs)
+            submitted = kwargs.get("submitted")
+
+            async def _gen():
+                self.query_started.set()
+                await self.release_submit.wait()
+                self.requests.append(kwargs)
+                if submitted is not None and not submitted.done():
+                    submitted.set_result(None)
+                yield _text_result("ok")
+
+            return _gen()
+
+    engine = _HoldQueryEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(lambda: engine.warmup_ids())
+    ws.put({"type": "video.query", "text": "first"})
+    assert await _poll(engine.query_started.is_set)
+    ws.put({"type": "video.query", "text": "second"})
+    await asyncio.sleep(0.05)
+    assert engine.query_ids() == []
+    assert engine.aborted == []
+    engine.release_submit.set()
+    assert await _poll(lambda: engine.query_ids())
+    first_query = engine.query_ids()[0]
+    assert await _poll(lambda: first_query in engine.aborted)
+    assert await _poll(lambda: len(engine.query_ids()) == 2)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_audio_session_never_starts_warmup():
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text", "audio"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg(9, 9, 9))})
+    await asyncio.sleep(0.3)
+
+    assert engine.warmup_ids() == []
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_warmup_starts_only_after_frames_ready(monkeypatch):
+    decode_started = threading.Event()
+    release_decode = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        decode_started.set()
+        release_decode.wait(timeout=5.0)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(decode_started.is_set)
+    await asyncio.sleep(0.1)
+    assert engine.warmup_ids() == []
+
+    release_decode.set()
+    assert await _poll(lambda: engine.warmup_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_query_cancelling_warmup_keeps_prewarm_alive(monkeypatch):
+    """Cancel warmup during prefill must not drop a later frame still decoding."""
+    frame1 = _b64(_make_jpeg(1, 1, 1))
+    frame2 = _b64(_make_jpeg(2, 2, 2))
+    raw2 = base64.b64decode(frame2)
+    decode_started = threading.Event()
+    release_decode = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes == raw2:
+            decode_started.set()
+            release_decode.wait(timeout=5.0)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine(warmup_delay=10.0)
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame1})
+    assert await _poll(lambda: engine.warmup_ids())
+
+    ws.put({"type": "video.frame", "data": frame2})
+    assert await _poll(decode_started.is_set)
+
+    ws.put({"type": "video.query", "text": "describe"})
+    await asyncio.sleep(0.1)
+    release_decode.set()
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    assert len(handler.query_contents) == 1
+    image_parts = [p for p in handler.query_contents[0] if p.get("type") == "image_pil"]
+    assert len(image_parts) == 2
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_query_keeps_frames_evicted_while_pinned(monkeypatch):
+    """Regression: frames FIFO-evicted while an in-flight query awaits prewarm
+    must keep their PIL entries and stay in the query prompt."""
+    frame_a = _b64(_make_jpeg(1, 1, 1))
+    frame_a2 = _b64(_make_jpeg(2, 2, 2))
+    frame_b = _b64(_make_jpeg(3, 3, 3))
+    raw_a = base64.b64decode(frame_a)
+    decoded: list[bytes] = []
+    release_others = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes != raw_a:
+            release_others.wait(timeout=5.0)
+        decoded.append(raw_bytes)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": ["text"],
+            "max_frames": 2,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame_a})
+    assert await _poll(lambda: raw_a in decoded)
+    ws.put({"type": "video.frame", "data": frame_a2})  # decode gated -> query will wait
+
+    ws.put({"type": "video.query", "text": "describe"})  # pins [a, a2]
+    await asyncio.sleep(0.1)
+    ws.put({"type": "video.frame", "data": frame_b})  # buffer full -> evicts pinned a
+    await asyncio.sleep(0.1)
+    release_others.set()
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    assert len(handler.query_contents) == 1
+    uuids = [p["uuid"] for p in handler.query_contents[0] if p.get("type") == "image_pil"]
+    expected = [hashlib.md5(base64.b64decode(f), usedforsecurity=False).hexdigest() for f in (frame_a, frame_a2)]
+    assert uuids == expected
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+async def _query_image_uuids(modalities: list[str], frames: list[str], *, max_frames: int) -> list[str]:
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": modalities,
+            "max_frames": max_frames,
+            "num_frames": max_frames,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    for frame in frames:
+        ws.put({"type": "video.frame", "data": frame})
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: handler.query_contents)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+    return [part["uuid"] for part in handler.query_contents[0] if part.get("type") == "image_pil"]
+
+
+def _frame_uuid(frame_b64: str) -> str:
+    return hashlib.md5(base64.b64decode(frame_b64), usedforsecurity=False).hexdigest()
+
+
+@pytest.mark.asyncio
+async def test_below_max_frames_incremental_matches_legacy_window():
+    frames = [_b64(_make_jpeg(i, 0, 0)) for i in range(3)]
+    expected = [_frame_uuid(frame) for frame in frames]
+    assert await _query_image_uuids(["text"], frames, max_frames=4) == expected
+    assert await _query_image_uuids(["text", "audio"], frames, max_frames=4) == expected
+
+
+@pytest.mark.asyncio
+async def test_incremental_full_window_drops_oldest_half():
+    frames = [_b64(_make_jpeg(i, 1, 0)) for i in range(5)]
+    kept = await _query_image_uuids(["text"], frames, max_frames=4)
+    legacy = await _query_image_uuids(["text", "audio"], frames, max_frames=4)
+    assert kept == [_frame_uuid(frame) for frame in frames[2:]]
+    assert legacy == [_frame_uuid(frame) for frame in frames[1:]]
+
+
+async def _run_frame_audio_query(modalities: list[str], frame: str, pcm: bytes) -> PromptRecordingHandler:
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": modalities,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame})
+    if "audio" not in modalities:
+        assert await _poll(lambda: engine.warmup_ids())
+    else:
+        await asyncio.sleep(0.05)
+        assert engine.warmup_ids() == []
+    ws.put({"type": "audio.chunk", "data": _b64(pcm)})
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: engine.query_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+    return handler
+
+
+@pytest.mark.asyncio
+async def test_reuse_query_with_audio_matches_legacy_prompt():
+    """Incremental and legacy query prompts match; both set the processor kwarg."""
+    frame = _b64(_make_jpeg())
+    pcm = b"\x01\x00\x02\x00"
+    reuse = await _run_frame_audio_query(["text"], frame, pcm)
+    legacy = await _run_frame_audio_query(["text", "audio"], frame, pcm)
+
+    assert reuse.warmup_contents
+    assert all(part.get("type") != "input_audio" for part in reuse.warmup_contents[0])
+    assert [part.get("type") for part in reuse.warmup_contents[0]] == ["image_pil"]
+
+    assert reuse.query_contents and legacy.query_contents
+    reuse_q = reuse.query_contents[0]
+    legacy_q = legacy.query_contents[0]
+    assert [part.get("type") for part in reuse_q] == ["image_pil", "input_audio", "text"]
+    assert [part.get("type") for part in legacy_q] == ["image_pil", "input_audio", "text"]
+    assert reuse_q[0]["uuid"] == legacy_q[0]["uuid"]
+    assert reuse_q[1] == legacy_q[1]
+    assert reuse_q[2] == legacy_q[2] == {"type": "text", "text": "describe"}
+
+    assert reuse.mm_processor_kwargs[0] == {"use_audio_in_video": True}
+    assert reuse.mm_processor_kwargs[1] == {"use_audio_in_video": True}
+    assert legacy.mm_processor_kwargs[0] == {"use_audio_in_video": True}

tests/entrypoints/openai_api/test_video_audio_delta.py

@@ -174,7 +174,9 @@ async def test_real_producer_deltas_reach_video_client(
         produced.append(OmniRequestOutput.from_stage_output(result, final_output_type="audio"))
 
     class Engine:
-        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None):
+        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None, submitted=None):
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
             if explicit_sampling:
                 assert sampling_params_list is not None
                 assert sampling_params_list[0].output_kind == RequestOutputKind.DELTA

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -11,6 +11,7 @@ from pydantic import ValidationError
 from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs
 from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
+from vllm_omni.config.stage_config import build_stage_runtime_overrides
 from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -945,3 +946,33 @@ def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker):
     assert launched["model"] == "generic-diffusion"
     assert launched["stage_id"] == 0
     assert runtime.stage_pools[0].clients == [client]
+
+
+def test_default_stage_config_preserves_pre_sharded_hsdp_weight_load() -> None:
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"hsdp_weight_load_strategy": "pre_sharded"})[0]
+
+    assert stage_cfg["engine_args"]["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("prefix", ["", "stage_0_"])
+def test_pre_sharded_hsdp_weight_load_survives_stage_overrides(prefix: str) -> None:
+    overrides = build_stage_runtime_overrides(0, {f"{prefix}hsdp_weight_load_strategy": "pre_sharded"})
+
+    assert overrides["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("strategy", ["full", "pre_sharded"])
+def test_typed_default_stage_preserves_hsdp_weight_load_strategy(strategy: str) -> None:
+    config = StageConfigFactory.create_typed_default_diffusion(
+        "generic-diffusion", {"hsdp_weight_load_strategy": strategy}
+    )
+
+    assert config.stage_configs[0].diffusion_config.hsdp_weight_load_strategy == strategy
+
+
+@pytest.mark.parametrize("strategy", ["memory_limited", "invalid"])
+def test_removed_hsdp_loading_strategies_are_rejected(strategy):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    with pytest.raises(ValueError, match="hsdp_weight_load_strategy"):
+        OmniDiffusionConfig(hsdp_weight_load_strategy=strategy)

tests/helpers/client.py

@@ -97,6 +97,7 @@ def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[
 class OmniResponse:
     """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
 
+    request_id: str | None = None
     text_content: str | None = None
     audio_data: list[str] | None = None
     audio_content: str | None = None
@@ -421,6 +422,8 @@ class OnlineOmniClient:
             text_content = ""
             audio_data = []
             for chunk in chat_completion:
+                if result.request_id is None:
+                    result.request_id = getattr(chunk, "id", None)
                 for choice in chunk.choices:
                     content = getattr(getattr(choice, "delta", None), "content", None)
                     modality = getattr(chunk, "modality", None)
@@ -453,6 +456,7 @@ class OnlineOmniClient:
         """Wall clock from *before* ``chat.completions.create`` through response parse + local decode."""
         result = OmniResponse()
         try:
+            result.request_id = getattr(chat_completion, "id", None)
             audio_data = None
             text_content = None
             for choice in chat_completion.choices:

tests/helpers/omni_scheduler.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Scheduler test helpers shared by AR / chunk-transfer unit tests."""
+
+from __future__ import annotations
+
+from types import MethodType
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+
+def bind_omits_transfer_helpers(sched) -> None:
+    """Bind the real omit helpers so a MagicMock scheduler does not skip save_async.
+
+    MagicMock is truthy, so an unbound ``_request_omits_chunk_transfer_to_next_stage``
+    silently skips the chunk put.
+    """
+    sched._omits_kv_transfer_cache = {}
+    sched._omni_final_stage_flags = MethodType(OmniARScheduler._omni_final_stage_flags, sched)
+    sched._request_omits_kv_transfer_to_next_stage = MethodType(
+        OmniARScheduler._request_omits_kv_transfer_to_next_stage,
+        sched,
+    )
+    sched._request_omits_chunk_transfer_to_next_stage = MethodType(
+        OmniARScheduler._request_omits_chunk_transfer_to_next_stage,
+        sched,
+    )

tests/model_executor/models/common/test_fused_sampling.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+def _reference(logits, uniforms, top_k, top_p):
+    if top_k > 0:
+        threshold = logits.topk(top_k, dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < threshold, -torch.inf)
+    if top_p < 1:
+        values, indices = logits.sort(dim=-1, descending=True)
+        probs = values.softmax(dim=-1, dtype=torch.float32)
+        values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf
+        logits = values.scatter(1, indices, values)
+    return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+
+
+@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)])
+def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(173)
+    logits = torch.randn(batch, vocab, device="cuda", dtype=dtype)
+    # The MTP caller supplies a noncontiguous row slice of [B, 15, vocab].
+    uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7)
+    expected = _reference(logits, uniforms, top_k, top_p)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    logits.add_(torch.randn_like(logits) * 0.1)
+    uniforms.uniform_(1e-7, 1 - 1e-7)
+    graph.replay()
+    torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)])
+def test_threshold_ties_and_degenerate_rows(top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    logits = torch.tensor(
+        [[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]],
+        device="cuda",
+        dtype=torch.float32,
+    )
+    uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_native_nucleus_boundaries(dtype):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(1923)
+    logits = torch.randn(1, 2048, device="cuda", dtype=dtype)
+    values, indices = logits.sort(descending=True)
+    probabilities = values.softmax(-1, dtype=torch.float32)
+    before = probabilities.cumsum(-1) - probabilities
+    for rank in [1, 19, 49, 1000, 2000]:
+        boundary = before[0, rank]
+        uniforms = torch.full_like(logits, 0.001, dtype=torch.float32)
+        uniforms[0, indices[0, rank]] = 1 - 1e-7
+        for direction in [-torch.inf, torch.inf]:
+            top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item()
+            actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p)
+            torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_p", [0.8, 1.0])
+def test_strided_logits_and_non_power_of_two_vocabulary(top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    for seed in range(16):
+        torch.manual_seed(seed)
+        logits = torch.randn(128, 2050, device="cuda")[:, ::2]
+        uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7)
+        actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p)
+        torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0)

tests/model_executor/models/common/test_short_kv_attention.py

@@ -16,7 +16,7 @@ pytestmark = [
 @pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)])
 @pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]])
 def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions):
-    from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention
+    from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention
 
     with torch.inference_mode():
         torch.manual_seed(42)

tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""HiFT graphs consume live inputs and return independently owned waveforms."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.model_executor.models.cosyvoice3.test_cosyvoice3_incremental_hift import _make_hift
+
+pytestmark = [pytest.mark.core_model]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("finalize", [False, True])
+@torch.inference_mode()
+def test_hift_graph_preserves_current_and_previous_request(finalize):
+    hift = _make_hift().cuda().eval()
+    hift.remove_weight_norm()
+    hift.f0_predictor.remove_weight_norm()
+    hift.enable_decode_graphs()
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    saved: list[tuple[torch.Tensor, torch.Tensor]] = []
+    for _ in range(5):
+        mel = torch.randn(1, 80, 30, device="cuda", generator=generator)
+        source = torch.randn(1, 1, 30 * 480, device="cuda", generator=generator)
+        expected = hift._decode_eager(mel, source, finalize=finalize)
+        actual = hift.decode(mel, source, finalize=finalize)
+        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+        for previous, snapshot in saved:
+            assert torch.equal(previous, snapshot)
+        if hift._decode_graphs.graphs:
+            saved.append((actual, actual.clone()))
+    assert len(hift._decode_graphs.graphs) == 1
+    hift._decode_graphs.max_graphs = 1
+    mel = torch.randn(1, 80, 31, device="cuda", generator=generator)
+    source = torch.randn(1, 1, 31 * 480, device="cuda", generator=generator)
+    actual = hift.decode(mel, source, finalize=finalize)
+    expected = hift._decode_eager(mel, source, finalize=finalize)
+    assert torch.equal(actual, expected)
+    assert len(hift._decode_graphs.graphs) == 1

tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py

@@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import
     HiggsAudioV3TalkerForConditionalGeneration,
 )
 from vllm_omni.model_executor.models.higgs_audio_v3.model_state import HiggsModelState
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
 from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -34,6 +35,8 @@ def make_state():
         _audio_continuation_id=99,
         update_decode_step_metadata=Mock(),
     )
+    state._stream_pos = {}
+    state._eager_state = EagerMTPState(state)
     return state
 
 

tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("rows", [16, 19])
+@torch.inference_mode()
+def test_first_frame_graph_outputs_own_storage(dtype, rows):
+    class Decoder(torch.nn.Linear):
+        config = SimpleNamespace(num_quantizers=1)
+
+        def forward(self, codes):
+            return codes.to(dtype) * self.weight
+
+    model = Decoder(1, 1, bias=False, device="cuda", dtype=dtype)
+    model.weight.fill_(1)
+    decoder = Qwen3OmniFirstFrameDecoder(model, sample_rate=24000)
+    decoder.capture()
+    codes = torch.arange(rows, device="cuda").reshape(-1, 1)
+    first = decoder.decode(codes)
+    second = decoder.decode(codes + 32)
+    assert first.dtype == second.dtype == torch.float32
+    torch.testing.assert_close(first, codes.float(), rtol=0, atol=0)
+    torch.testing.assert_close(second, (codes + 32).float(), rtol=0, atol=0)

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+def _codec(enabled):
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=8,
+        decoder_dim=32,
+        upsample_rates=[2, 2],
+        upsampling_ratios=[2],
+        codebook_size=16,
+        num_quantizers=2,
+    )
+    config._attn_implementation = "sdpa"
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=config, stage_connector_config={"extra": {"codec_cudnn_benchmark": enabled}}
+        )
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=vllm_config).cuda().eval()
+    model.precompute_snake_caches()
+    return model
+
+
+def _flags():
+    c = torch.backends.cudnn
+    return c.enabled, c.benchmark, c.benchmark_limit, c.deterministic, c.allow_tf32
+
+
+@torch.inference_mode()
+def test_autotuned_codec_graph_replay_restores_flags():
+    torch.manual_seed(21)
+    model = _codec(True)
+    codes = torch.randint(0, 16, (2, 2, 8), device="cuda")
+    original = _flags()
+    for _ in range(3):
+        model(codes)
+    assert _flags() == original
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = model(codes)
+    assert _flags() == original
+    codes.random_(0, 16)
+    graph.replay()
+    expected = model(codes)
+    torch.testing.assert_close(captured, expected)
+    assert _flags() == original
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+@torch.inference_mode()
+def test_codec_failure_restores_cudnn_settings(mocker, enabled):
+    model = _codec(enabled)
+    codes = torch.randint(0, 16, (1, 2, 8), device="cuda")
+    original = _flags()
+
+    def fail(hidden):
+        if enabled:
+            assert torch.backends.cudnn.benchmark
+            assert torch.backends.cudnn.benchmark_limit == 10
+        else:
+            assert _flags() == original
+        raise RuntimeError("decode failed")
+
+    mocker.patch.object(model, "_decode_waveform", side_effect=fail)
+    with pytest.raises(RuntimeError, match="decode failed"):
+        model(codes)
+    assert _flags() == original

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import copy
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeCode2WavDecoderBlock,
+    Qwen3OmniMoeSnakeBeta,
+)
+
+from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav, use_fused_snake
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _blocks():
+    config = Qwen3OmniMoeCode2WavConfig(decoder_dim=32, upsample_rates=[2, 2])
+    torch.manual_seed(0)
+    blocks = torch.nn.ModuleList(Qwen3OmniMoeCode2WavDecoderBlock(config, i) for i in range(2)).eval()
+    with torch.no_grad():
+        for module in blocks.modules():
+            if isinstance(module, Qwen3OmniMoeSnakeBeta):
+                module.alpha.uniform_(-0.5, 0.5)
+                module.beta.uniform_(-0.5, 0.5)
+    return blocks
+
+
+def test_decoder_blocks_swap_every_hf_snake_and_keep_weight_names():
+    blocks = _blocks()
+    hf_count = sum(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    keys = set(blocks.state_dict())
+
+    assert use_fused_snake(blocks) == hf_count > 0
+    assert not any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    assert sum(isinstance(m, SnakeBeta) for m in blocks.modules()) == hf_count
+    assert set(blocks.state_dict()) == keys  # checkpoint weights load unchanged
+
+
+def test_fused_snake_blocks_match_hf_blocks():
+    reference = _blocks()
+    fused = copy.deepcopy(reference)
+    use_fused_snake(fused)
+    fused.load_state_dict(reference.state_dict())
+    x = torch.randn(2, 32, 20)
+    with torch.no_grad():
+        expected = x
+        actual = x
+        for ref_block, fused_block in zip(reference, fused):
+            expected = ref_block(expected)
+            actual = fused_block(actual)
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    ("option", "device", "cuda", "fused"),
+    [
+        (False, "cuda", True, False),
+        (True, "cuda", True, True),
+        (True, "cpu", True, False),
+        (True, "cuda", False, False),
+    ],
+)
+def test_decoder_block_fusion_requires_cuda_opt_in(monkeypatch, option, device, cuda, fused):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: cuda)
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+    )
+    model = Qwen3OmniMoeCode2Wav(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_config=config, stage_connector_config={"extra": {"codec_fused_snake": option}}
+            ),
+            device_config=SimpleNamespace(device=device),
+        )
+    )
+    assert any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in model.decoder.modules()) is (not fused)

tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import (
+    Qwen3OmniMoeCode2Wav,
+    plan_decode_groups,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_Q = 4
+_CODEBOOK = 16
+
+
+def _code2wav(quantizers=_Q) -> Qwen3OmniMoeCode2Wav:
+    torch.manual_seed(0)
+    config = Qwen3OmniMoeCode2WavConfig(
+        codebook_size=_CODEBOOK,
+        num_quantizers=quantizers,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+        sliding_window=8,
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=config)))
+    return model.eval()
+
+
+def test_first_frame_pcm_equals_code2wav_streaming_chunk0():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (3, _Q))
+    with torch.inference_mode():
+        for row in range(codes.shape[0]):
+            pcm = decoder.decode(codes[row : row + 1])
+            chunk0 = code2wav.chunked_decode_streaming(
+                codes[row].reshape(1, _Q, 1), left_context_size=[0], seq_token_counts=[_Q]
+            )[0]
+            assert pcm.dtype == torch.float32
+            assert torch.equal(pcm[0], chunk0.reshape(-1).float())
+            assert 0 < pcm.shape[-1] < int(code2wav.total_upsample)
+        batched = decoder.decode(codes)
+    assert batched.shape == (3, pcm.shape[-1])
+
+
+def test_first_frame_is_a_prefix_of_a_longer_first_chunk():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (_Q, 4))
+    with torch.inference_mode():
+        first = decoder.decode(codes[:, :1].T)[0]
+        chunk0 = code2wav.chunked_decode_streaming(
+            codes.reshape(1, _Q, 4), left_context_size=[0], seq_token_counts=[4 * _Q]
+        )[0]
+    torch.testing.assert_close(chunk0.reshape(-1)[: first.numel()].float(), first)
+
+
+def test_chunk_ramp_adds_exact_code2wav_graph_sizes(monkeypatch):
+    import vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper as wrapper_module
+
+    real = wrapper_module.CUDAGraphDecoderWrapper
+    wrapper = MagicMock(compute_capture_sizes=real.compute_capture_sizes)
+    monkeypatch.setattr(wrapper_module, "CUDAGraphDecoderWrapper", wrapper)
+    extra = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "codec_chunk_ramp": [1, 2, 4, 8, 16, 25]}
+    code2wav = SimpleNamespace(config=SimpleNamespace(num_quantizers=_Q))
+    code2wav.enable_cudagraph = lambda **kw: Qwen3OmniMoeCode2Wav.enable_cudagraph(
+        code2wav, device=torch.device("cuda"), **kw
+    )
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.code2wav = code2wav
+    model.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(async_chunk=True, enforce_eager=False, stage_connector_config={"extra": extra})
+    )
+    for batches in ([], [2, 4, 8]):
+        extra["codec_graph_batch_sizes"] = batches
+        Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+        args = wrapper.call_args.kwargs
+        sizes = args["capture_sizes"]
+        assert {1, 3, 7, 15, 31, 50} <= set(sizes)
+        assert set(real.compute_capture_sizes(codec_chunk_frames=25, codec_left_context_frames=25)) <= set(sizes)
+        assert set(args["extra_capture_shapes"]) == {
+            (b, size) for b in batches for size in sizes if size <= 50 and b * size <= 325
+        }
+    assert code2wav._streaming_batch_sizes[50] == [1, 2, 4]
+    extra.pop("codec_chunk_ramp")
+    extra.pop("codec_graph_batch_sizes")
+    Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+    assert wrapper.call_args.kwargs["capture_sizes"] is None
+
+
+_SIZES = [1, 2, 3, 4, 7, 8, 15, 16, 25, 31, 32, 50, 64]
+
+
+def _bucket(length):
+    return next((size for size in _SIZES if size >= length), None)
+
+
+def test_decode_groups_keep_a_uniform_batch_whole_and_split_by_the_largest_graph():
+    assert plan_decode_groups([50] * 6, _bucket, lambda size: [1, 2, 4, 8]) == [(list(range(6)), 50)]
+    calls = plan_decode_groups([25] * 20, _bucket, lambda size: [1, 2, 4, 8, 16])
+    assert [len(rows) for rows, _length in calls] == [16, 4]
+    calls = plan_decode_groups([1] * 8 + [50] * 8, _bucket, lambda size: [1, 2, 4, 8] if size < 50 else [1, 2, 4])
+    assert [(len(rows), length) for rows, length in calls] == [(8, 1), (4, 50), (4, 50)]
+    assert plan_decode_groups([], _bucket, lambda size: [1, 2]) == []
+
+
+@pytest.mark.parametrize("frames", [65, 129])
+def test_decode_groups_beyond_largest_graph_use_eager(monkeypatch, frames):
+    from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+
+    wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper)
+    wrapper.enabled = wrapper._warmed_up = True
+    wrapper._bucket_sizes = [8, 32, 64]
+    wrapper._compiled_shapes = {(1, 64), (8, 64)}
+    wrapper._compiled_graphs = wrapper.graphs = dict.fromkeys(wrapper._compiled_shapes, MagicMock())
+    wrapper.decoder = MagicMock(side_effect=lambda codes: codes.float() * 2)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    groups = plan_decode_groups([frames, frames], wrapper._get_padded_size, lambda _: (1,))
+    assert groups == [([0], frames), ([1], frames)]
+    codes = torch.arange(2 * frames).reshape(2, 1, frames)
+    for rows, width in groups:
+        actual = wrapper.decode(codes[rows, :, :width])
+        torch.testing.assert_close(actual, codes[rows].float() * 2, rtol=0, atol=0)
+    assert wrapper.decoder.call_count == 2
+    for graph in (*wrapper.graphs.values(), *wrapper._compiled_graphs.values()):
+        graph.replay.assert_not_called()
+
+
+class _PaddingGraphs:
+    def __init__(self, model, batch_sizes):
+        self.model = model
+        self.batch_sizes = sorted(batch_sizes)
+        self.calls = []
+
+    def _get_padded_size(self, length):
+        return _bucket(length)
+
+    def decode(self, codes):
+        rows, _q, frames = codes.shape
+        size = _bucket(frames)
+        batch = next(b for b in self.batch_sizes if b >= rows)
+        self.calls.append((rows, frames))
+        padded = torch.zeros(batch, codes.shape[1], size, dtype=codes.dtype)
+        padded[:rows, :, :frames] = codes
+        out = self.model(padded)
+        return out[:rows, :, : out.shape[-1] - (size - frames) * int(self.model.total_upsample)].clone()
+
+
+def test_grouped_streaming_decode_matches_one_padded_batch():
+    code2wav = _code2wav()
+    lengths = [1, 3, 50, 7, 31]
+    lefts = [0, 1, 25, 3, 15]
+    codes = torch.zeros(len(lengths), _Q, max(lengths), dtype=torch.long)
+    for row, length in enumerate(lengths):
+        codes[row, :, :length] = torch.randint(0, _CODEBOOK, (_Q, length))
+    counts = [length * _Q for length in lengths]
+    with torch.inference_mode():
+        reference = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+        graphs = _PaddingGraphs(code2wav, [1, 2, 4])
+        code2wav._cudagraph_wrapper, code2wav._cudagraph_enabled = graphs, True
+        code2wav._streaming_batch_sizes = {size: [1, 2, 4] for size in _SIZES}
+        grouped = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+
+    for want, got in zip(reference, grouped, strict=True):
+        assert want.shape == got.shape
+        torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize(("frames", "quantizers"), [(1, 16), (4, 16), (4, 4)])
+def test_delivered_first_frame_is_removed_from_code2wav_output(frames, quantizers):
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+
+    codec = _code2wav(quantizers)
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage, model.code2wav, model.code2wav_config = "code2wav", codec, codec.config
+    model.generate_audio = lambda codes, left, counts: codec.chunked_decode_streaming(
+        codes, left_context_size=left, seq_token_counts=counts
+    )
+    codes = torch.randint(0, _CODEBOOK, (2, quantizers, frames))
+    with torch.inference_mode():
+        reference = model.generate_audio(codes, [0, 0], [quantizers * frames] * 2)
+        first = Qwen3OmniFirstFrameDecoder(codec, sample_rate=24000).decode(codes[0, :, :1].T)[0]
+        out = model.forward(
+            input_ids=codes.flatten(),
+            positions=None,
+            seq_token_counts=[quantizers * frames] * 2,
+            runtime_additional_information=[
+                {"meta": {"left_context_size": 0, "first_audio": True}},
+                {"meta": {"left_context_size": 0, "first_audio": False}},
+            ],
+        ).multimodal_outputs
+    torch.testing.assert_close(torch.cat([first, out["model_outputs"][0].flatten()]), reference[0].flatten())
+    torch.testing.assert_close(out["model_outputs"][1], reference[1].reshape(1, -1))
+    assert [bool(flag) for flag in out[FIRST_AUDIO_REQUIRED_KEY]] == [True, False]

tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_H = 4
+_Q = 3
+
+
+class _Talker(nn.Module):
+    num_code_groups = _Q
+
+    def __init__(self):
+        super().__init__()
+        self.text_projection = nn.Identity()
+        self.code_predictor = SimpleNamespace(
+            _top_k=50,
+            _top_p=0.8,
+            set_sampling_params=lambda **kwargs: setattr(self, "sampling", kwargs),
+        )
+
+
+def _talker(*, async_chunk: bool = True) -> Qwen3OmniMoeForConditionalGeneration:
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.model_stage = "talker"
+    model.talker = _Talker()
+    model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=async_chunk, use_v2_model_runner=True))
+    model.tts_eos_embed = torch.full((_H,), -1.0)
+    model.tts_pad_embed = torch.full((_H,), -2.0)
+    model._codec_codebook_size = 2048
+    return model
+
+
+def _rows(*values: float) -> torch.Tensor:
+    return torch.tensor([[value] * _H for value in values])
+
+
+def _step(model, payload: dict) -> torch.Tensor:
+    update: dict = {}
+    text_step = model._thinker_decode_to_talker_decode(payload, torch.device("cpu"), update)
+    for key, value in update.items():
+        payload.setdefault(key, {}).update(value)
+    return text_step.reshape(-1)
+
+
+def test_mtp_frame_valid_rejects_codec_eos_and_special_ids():
+    layer0 = torch.tensor([0, 2047, 2048, 2150, -1])
+    assert _talker().mtp_frame_valid(layer0).tolist() == [True, True, False, False, False]
+
+
+def test_decode_rows_are_consumed_one_per_step_in_arrival_order():
+    model = _talker()
+    # Two rows were already pending (cached at prefill); two more arrive at once.
+    payload = {"embed": {"cached_decode": _rows(1, 2), "decode": _rows(3, 4)}, "meta": {}}
+    steps = [_step(model, payload)[0].item()]
+    for arrival in (None, _rows(5), None, None, None, None):
+        payload["embed"]["decode"] = arrival
+        steps.append(_step(model, payload)[0].item())
+    assert steps == [1, 2, 3, 4, 5, -1, -2]
+    assert payload["embed"]["cached_decode"] is None
+
+
+@pytest.mark.parametrize(("v2", "resumable"), [(True, True), (False, False), (False, True)])
+def test_v1_and_resumable_requests_keep_the_indexed_cache_path(v2, resumable):
+    model = _talker()
+    model.vllm_config.model_config.use_v2_model_runner = v2
+    payload = {"embed": {"cached_decode": _rows(1, 2, 3), "decode": None}, "meta": {"resumable": resumable}}
+    payload["meta"]["num_processed_tokens"] = 2
+    assert _step(model, payload)[0].item() == 3
+
+
+def test_make_omni_output_aligns_codes_to_request_spans():
+    model = _talker()
+    hidden = torch.zeros((8, _H))  # 5 real tokens + graph padding
+    decode_codes = torch.tensor([[5, 6, 7]])
+    buffers = [
+        {"req_id": "_warmup_0_"},  # no preprocess, no codes
+        {"req_id": "a", "codes": {"audio": torch.zeros((3, _Q), dtype=torch.long)}},
+        {"req_id": "b", "codes": {"audio": decode_codes}},
+    ]
+    out = model.make_omni_output(
+        hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 1), (1, 4), (4, 5)]
+    )
+    codes = out.multimodal_outputs["codes"]["audio"]
+    assert codes.shape == (5, _Q)
+    assert codes[4].tolist() == [5, 6, 7]
+    assert out.text_hidden_states.shape[0] == 5
+
+
+def test_make_omni_output_without_spans_keeps_v1_concat():
+    model = _talker()
+    buffers = [{"codes": {"audio": torch.tensor([[1, 2, 3]])}}, {"codes": {"audio": torch.tensor([[4, 5, 6]])}}]
+    out = model.make_omni_output(torch.zeros((4, _H)), model_intermediate_buffer=buffers)
+    assert out.multimodal_outputs["codes"]["audio"].tolist() == [[1, 2, 3], [4, 5, 6]]
+    assert out.text_hidden_states.shape[0] == 2
+
+
+def test_eager_output_is_token_major_with_invalid_rows():
+    model = _talker()
+    model.eager_frames_active = True
+    hidden = torch.zeros((6, _H))
+    buffers = [{"req_id": "a"}, {"req_id": "b"}]
+    out = model.make_omni_output(hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 3), (3, 4)])
+    mm = out.multimodal_outputs
+    assert mm["codes"]["audio"].shape == (4, _Q)
+    assert mm["meta"]["codec_frame_valid"].tolist() == [0, 0, 0, 0]
+    assert out.text_hidden_states is hidden

tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py

@@ -104,6 +104,7 @@ def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, e
     if staged:
         model_config.model_stage = "thinker"
     vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None)
+    vllm_config.speculative_config = None
     thinker = nn.Module()
     thinker.weight = nn.Parameter(torch.ones(1))
     thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start))

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -569,6 +569,27 @@ class TestCodePredictorPerRowGenerators:
             )
 
 
+@pytest.mark.parametrize("fused_requested", [False, True])
+@pytest.mark.parametrize("invalid_input", ["uniforms", "generators"])
+def test_sampling_contract_rejected_before_dispatch(mocker, loaded_target_classes, fused_requested, invalid_input):
+    predictor, talker_config = TestCodePredictorPerRowGenerators()._make_predictor(mocker, loaded_target_classes)
+    predictor._fused_requested = fused_requested
+    predictor._fused = mocker.Mock()
+    setup = mocker.patch.object(predictor, "_setup_compile")
+    uniforms = torch.full((2, 3, 1 if invalid_input == "uniforms" else 64), 0.5)
+    generators = [torch.Generator()] if invalid_input == "generators" else None
+    with pytest.raises(ValueError, match="sample_uniforms must have shape|one entry per row"):
+        predictor(
+            torch.zeros(2, dtype=torch.long),
+            torch.randn(2, talker_config.hidden_size),
+            torch.randn(2, talker_config.hidden_size),
+            sample_uniforms=uniforms,
+            generators=generators,
+        )
+    setup.assert_not_called()
+    predictor._fused.assert_not_called()
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 
@@ -663,6 +684,8 @@ class TestCodePredictorGraphReplay:
 
         predictor = object.__new__(code_predictor_wrapper)
         torch.nn.Module.__init__(predictor)
+        predictor._fused_requested = False
+        predictor._fused = None
         predictor._num_groups = 3
         predictor._model_dtype = torch.float32
         predictor._setup_compile = mocker.Mock()
@@ -1432,3 +1455,24 @@ def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_cl
     expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
     actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms)
     torch.testing.assert_close(actual, expected)
+
+
+def test_tts_cached_predictor_keeps_its_own_step_dispatch(mocker: MockerFixture, loaded_target_classes) -> None:
+    """The TTS cache must not enter Omni's uninitialized incremental loop."""
+    _, _, wrapper, _, _ = loaded_target_classes
+    common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    mocker.patch.object(common_mod.current_omni_platform, "is_cuda", return_value=True)
+    mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+    mocker.patch.object(wrapper, "_stage_connector_extra_config", return_value={"code_predictor_kv_cache": True})
+    cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+    predictor = wrapper(vllm_config=_make_vllm_config(mocker), config=cp_config, talker_config=talker_config)
+    assert predictor._kv_requested
+    predictor._model_dtype = torch.float32
+    predictor._codec_embeds_list = list(predictor.model.codec_embedding)
+    predictor._bucket_sizes = [1]
+    mocker.patch.object(predictor, "_setup_compile")
+    step = mocker.patch.object(predictor, "_predict_step_logits", return_value=torch.ones(1, cp_config.vocab_size))
+    codes = predictor(torch.zeros(1, dtype=torch.long), torch.zeros(1, 32), torch.zeros(1, 32), do_sample=False)
+    assert step.call_count == cp_config.num_code_groups - 1
+    assert codes.shape == (1, cp_config.num_code_groups)
+    assert predictor._kv_buf is None

tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Supplied-uniform samples must match the existing ATen path exactly."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -12,6 +14,79 @@ pytestmark = [
 ]
 
 
+@pytest.mark.parametrize("top_k", [0, 1, 50, 2048])
+@torch.inference_mode()
+def test_fused_sample_kernel_preserves_signed_zero_cutoff_ties(top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.fused_code_predictor import _cp_sample_kernel
+
+    vocab, hidden, groups = 2048, 64, 4
+    logits = torch.zeros(2, vocab, device="cuda", dtype=torch.bfloat16)
+    logits[0, 1] = -0.0
+    logits[1, ::2] = -0.0
+    uniforms = torch.full((2, groups - 1, vocab), 0.01, device="cuda")[:, 1, :]
+    uniforms[0, 1] = uniforms[1, 2] = 0.99
+    codes = torch.full((2, groups), -1, device="cuda", dtype=torch.int64)
+    table = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16)
+    next_input = torch.empty(2, hidden, device="cuda", dtype=table.dtype)
+
+    def sample():
+        _cp_sample_kernel[(2,)](
+            logits, uniforms, uniforms.stride(0), codes, 1, groups, table, next_input, 1.0,
+            V=vocab, TOPK=top_k, HID=hidden, HAS_NEXT=True, num_warps=4,
+        )  # fmt: skip
+
+    scaled = logits.float()
+    if top_k > 0:
+        scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf)
+    expected = (scaled - torch.log(-torch.log(uniforms))).argmax(-1)
+    assert expected.tolist() == [1, 2]
+    sample()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+    assert (codes[:, [0, 2, 3]] == -1).all()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        sample()
+    codes.fill_(-1)
+    graph.replay()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+
+
+@torch.inference_mode()
+def test_fused_predictor_warmup_with_capacity_one():
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    talker = Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=SimpleNamespace(max_num_seqs=1),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=talker
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    hidden = torch.randn(1, 1, 64, device="cuda", dtype=torch.bfloat16)
+    result = predictor(
+        torch.zeros(1, 1, device="cuda", dtype=torch.long), hidden, hidden,
+        sample_uniforms=torch.full((1, 3, 64), 0.5, device="cuda"),
+    )  # fmt: skip
+    assert result.shape == (1, 4)
+    assert ((result >= 0) & (result < 64)).all()
+
+
 @pytest.mark.parametrize("batch", [1, 64])
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
 @pytest.mark.parametrize("top_k", [-1, 1, 50, 2048])
@@ -66,3 +141,47 @@ def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k):
     expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
     actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
     torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("batch", [1, 4])
+@pytest.mark.parametrize("top_k", [1, 50])
+@torch.inference_mode()
+def test_fused_predictor_matches_full_reference_with_fixed_uniforms(mocker, batch, top_k):
+    from vllm.config import VllmConfig
+
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    torch.manual_seed(42)
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=mocker.Mock(max_num_seqs=4),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    ids = torch.randint(0, 64, (batch, 1), device="cuda")
+    embeds = torch.randn(batch, 1, 64, device="cuda", dtype=torch.bfloat16)
+    hidden = torch.randn_like(embeds)
+    uniforms = torch.rand(batch, 3, 64, device="cuda").clamp_(1e-6, 1 - 1e-6)
+    actual = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms).clone()
+    # Use identical parameters through the full re-prefill and ATen sampling
+    # path, bypassing both the fused predictor and its frame-local KV cache.
+    predictor._fused_requested = False
+    predictor._fused = None
+    predictor._frame_cache = None
+    expected = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Prefix graphs must follow per-request first-audio delivery, not a capture hint."""
 
+import math
 from types import SimpleNamespace
 
 import pytest
@@ -18,6 +19,38 @@ pytestmark = [pytest.mark.core_model]
 DEVICE = torch.device("cuda:0")
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("capacity,model_type", [(1, "base"), (128, "base"), (128, "custom_voice")])
+def test_stream_priming_graph_workspace_respects_token_budget(monkeypatch, capacity, model_type):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import streaming_decoder
+
+    captures = []
+
+    def capture(stream, sizes, frames=1):
+        captures.append((sizes, frames))
+        return SimpleNamespace(sizes=sizes, frames=frames)
+
+    monkeypatch.setattr(streaming_decoder, "StreamingDecodeGraphs", capture)
+    model = SimpleNamespace(
+        stream_decoder=SimpleNamespace(max_frames=33),
+        stream_graphs=None,
+        stream_prime_graphs=None,
+        stream_ref_context_frames=25,
+        stream_chunk_frames=25,
+        config=SimpleNamespace(tts_model_type=model_type),
+        vllm_config=SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=512)),
+    )
+    sizes = [size for size in [1, 2, 4, 8, 16, 32, 64, 128] if size <= capacity]
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    expected = [(sizes, 1)]
+    if model_type == "base":
+        expected.append(([size for size in sizes if size * 25 <= 512], 25))
+    assert captures == expected
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    assert captures == expected
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)])
 @pytest.mark.parametrize("initial_frames", [1, 3])
@@ -124,3 +157,134 @@ def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major
             torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4)
             for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):
                 torch.testing.assert_close(cache[key], expected_cache[key], atol=1e-5, rtol=1e-5)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.inference_mode()
+def test_stream_decoder_preemption_restores_all_state_into_a_different_slot():
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import StreamingCodecDecoder
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=2, dtype=torch.bfloat16)
+    codes = torch.randint(0, 32, (1, 8, 2), device=DEVICE)
+    slot = torch.tensor([0], device=DEVICE, dtype=torch.int32)
+    for pos in range(7):
+        stream(codes[:, pos : pos + 1], slot, torch.tensor([pos], device=DEVICE, dtype=torch.int32))
+    saved = stream.save_slot(0)
+    expected = stream(codes[:, 7:], slot, torch.tensor([7], device=DEVICE, dtype=torch.int32)).clone()
+    # Both the freed slot and the destination can hold an unrelated request.
+    other_codes = (codes[:, :3] + 1) % 32
+    for pos in range(3):
+        stream(
+            other_codes[:, pos : pos + 1].expand(2, -1, -1).contiguous(),
+            torch.tensor([0, 1], device=DEVICE, dtype=torch.int32),
+            torch.full((2,), pos, device=DEVICE, dtype=torch.int32),
+        )
+    without_restore = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    ).clone()
+    assert not torch.equal(without_restore, expected)
+    stream.restore_slot(1, saved)
+    actual = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    )
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@torch.inference_mode()
+def test_batched_reference_priming_matches_serial_decoder_continuation(use_graph):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=5, dtype=torch.bfloat16)
+    graphs = StreamingDecodeGraphs(stream, [2], frames=25) if use_graph else None
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.stream_decoder = stream
+    model.stream_prime_graphs = graphs
+    model.stream_chunk_frames = 25
+    refs = [torch.randint(0, 32, (length, 2), device=DEVICE) for length in [25, 3, 51, 25]]
+    slots = [2, 0, 3, 4]
+    frame = torch.randint(0, 32, (4, 1, 2), device=DEVICE)
+    slot_tensor = torch.tensor(slots, device=DEVICE, dtype=torch.int32)
+    positions = torch.tensor([len(ref) for ref in refs], device=DEVICE, dtype=torch.int32)
+    for slot, ref in zip(slots, refs, strict=True):
+        for t0 in range(0, len(ref), 25):
+            stream(
+                ref[None, t0 : t0 + 25].to(torch.int32).contiguous(),
+                torch.tensor([slot], device=DEVICE, dtype=torch.int32),
+                torch.tensor([t0], device=DEVICE, dtype=torch.int32),
+            )
+    expected = stream(frame, slot_tensor, positions).clone()
+    untouched = stream.save_slot(1)
+    # Reuse dirty slots; pos=0 must reset every causal layer for each group.
+    model.prime_stream_decoder(list(zip(slots, refs, strict=True)))
+    actual = stream(frame, slot_tensor, positions)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=3e-5)
+    for got, want in zip(stream.save_slot(1), untouched, strict=True):
+        torch.testing.assert_close(got, want, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@torch.inference_mode()
+def test_stream_frames_match_fp32_exact_decoder(use_graph, dtype):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=dtype)
+    # The tiny HF default initialization shrinks the deep conv stack to
+    # near-zero PCM. Keep the oracle signal large enough to catch lost state.
+    for module in decoder.modules():
+        if isinstance(module, (torch.nn.Conv1d, torch.nn.ConvTranspose1d, torch.nn.Linear)):
+            torch.nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+            if module.bias is not None:
+                module.bias.uniform_(-0.01, 0.01)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    reference = _make_decoder().to(device=DEVICE, dtype=torch.float32)
+    reference.load_state_dict(decoder.state_dict())
+    reference.precompute_snake_caches()
+    # Cross both the attention window and KV ring wrap, checking every frame.
+    codes = torch.randint(0, 32, (2, 136, 2), device=DEVICE)
+    expected = reference._forward_exact(codes.transpose(1, 2))[:, 0]
+    assert expected.square().sum() > 1e-3
+    stream = StreamingCodecDecoder(decoder, num_slots=3, dtype=dtype)
+    decode = StreamingDecodeGraphs(stream, [2]) if use_graph else stream
+    slots = torch.tensor([2, 0], device=DEVICE, dtype=torch.int32)
+    spf = int(stream.spf)
+    for frame in range(codes.shape[1]):
+        positions = torch.full((2,), frame, device=DEVICE, dtype=torch.int32)
+        actual = decode(codes[:, frame : frame + 1].contiguous().to(torch.int32), slots, positions).reshape(2, -1)
+        exact_frame = expected[:, frame * spf : (frame + 1) * spf]
+        relative_rms = (actual - exact_frame).square().mean(1).sqrt() / exact_frame.square().mean(1).sqrt()
+        # BF16 kernels are approximate; compare every frame to the independent
+        # FP32 full decoder, not only one streaming implementation to itself.
+        tolerance = 0.05 if dtype == torch.bfloat16 else 5e-4
+        assert torch.all(relative_rms < tolerance), (frame, relative_rms)

tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Streaming col2im preserves PCM inputs and ping-pong slot state."""
+
+import pytest
+import torch
+from vllm.triton_utils import triton
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import _col2im_kernel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.tts]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("batch,frames,channels,rate", [(1, 1, 5, 2), (3, 3, 129, 3), (9, 1, 64, 4)])
+def test_col2im_partial_tiles_and_reordered_slot_state(batch: int, frames: int, channels: int, rate: int):
+    torch.manual_seed(23)
+    capacity = batch + 2
+    z = torch.randn(batch, frames, 2 * rate, channels).to(torch.bfloat16)
+    bias = torch.randn(channels).to(torch.bfloat16)
+    state = torch.randn(2, capacity, rate, channels).to(torch.bfloat16)
+    expected_state = state.clone()
+    slots = torch.arange(batch - 1, -1, -1, dtype=torch.int32)
+    positions = torch.arange(batch, dtype=torch.int32)
+    parity = torch.arange(capacity, dtype=torch.int32) % 2
+    expected = torch.empty(batch, frames, rate, channels, dtype=torch.bfloat16)
+    for b in range(batch):
+        slot = int(slots[b])
+        side = int(parity[slot])
+        for t in range(frames):
+            previous = z[b, t - 1, rate:] if t else state[side, slot] if positions[b] else torch.zeros(rate, channels)
+            expected[b, t] = (z[b, t, :rate].float() + previous.float() + bias.float()).to(torch.bfloat16)
+        expected_state[1 - side, slot] = z[b, -1, rate:]
+
+    z_gpu, bias_gpu, state_gpu = z.cuda(), bias.cuda(), state.cuda()
+    slots_gpu, positions_gpu, parity_gpu = slots.cuda(), positions.cuda(), parity.cuda()
+    actual = torch.empty_like(expected, device="cuda")
+    width = min(128, triton.next_power_of_2(channels))
+    rows = batch * frames
+    _col2im_kernel[(rows, rate, triton.cdiv(channels, width))](
+        z_gpu,
+        bias_gpu,
+        state_gpu,
+        slots_gpu,
+        positions_gpu,
+        parity_gpu,
+        actual,
+        frames,
+        capacity,
+        R=rate,
+        C=channels,
+        BC=width,
+        num_warps=4,
+    )
+    torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0)
+    torch.testing.assert_close(state_gpu.cpu(), expected_state, rtol=0, atol=0)

tests/model_executor/models/test_talker_first_audio.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("env", "extra", "async_chunk", "v2", "expected"),
+    [
+        (None, {}, True, True, False),
+        ("1", {}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, True),
+        ("0", {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [4, 8]}, True, True, False),
+        (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 1}, True, True, True),
+        (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 4}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1]}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, False, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, False, False),
+    ],
+)
+def test_talker_first_audio_default_and_conditions(monkeypatch, env, extra, async_chunk, v2, expected):
+    if env is None:
+        monkeypatch.delenv("VLLM_OMNI_TALKER_FIRST_AUDIO", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_OMNI_TALKER_FIRST_AUDIO", env)
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            async_chunk=async_chunk, use_v2_model_runner=v2, stage_connector_config={"extra": extra}
+        ),
+        device_config=SimpleNamespace(device="cuda"),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend=None
+        ),
+        cache_config=SimpleNamespace(enable_prefix_caching=False),
+    )
+    assert talker_first_audio_enabled(config) is expected
+    for section, field, value in [
+        ("device_config", "device", "cpu"),
+        ("parallel_config", "tensor_parallel_size", 2),
+        ("parallel_config", "pipeline_parallel_size", 2),
+        ("parallel_config", "distributed_executor_backend", "mp"),
+        ("cache_config", "enable_prefix_caching", True),
+    ]:
+        with monkeypatch.context() as patch:
+            patch.setattr(getattr(config, section), field, value)
+            assert not talker_first_audio_enabled(config)
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+    assert not talker_first_audio_enabled(config)
+
+
+@pytest.mark.parametrize("option", [False, True])
+def test_tts_requires_connector_opt_in(monkeypatch, option):
+    from vllm_omni.model_executor.models.qwen3_tts import first_audio
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor,
+    )
+
+    monkeypatch.setattr(Predictor, "_stage_connector_extra_config", lambda _: {"talker_first_audio": option})
+    monkeypatch.setattr(first_audio, "supports_talker_first_audio", lambda _: True)
+    assert first_audio.talker_first_audio_enabled(None) is option

tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from collections import defaultdict
 from types import SimpleNamespace
 
@@ -18,7 +21,7 @@ def _make_transfer_manager(
     left_context: int = 1,
     initial_chunk_frames: int = 0,
     single_initial_chunk: bool = False,
-    backlog_chunk_frames: int = 0,
+    backlog_chunk_frames: int | str = 0,
     backlog_load_threshold: float = 0.75,
     max_num_seqs: int = 1,
 ):
@@ -227,3 +230,77 @@ def test_invalid_backlog_config_fails_loudly():
             {"audio_codes": torch.tensor([[1, 2, 3]])},
             request,
         )
+
+
+@pytest.mark.parametrize("single_initial_chunk", [False, True])
+@pytest.mark.parametrize("initial_chunk_frames", [0, 4, 10])
+@pytest.mark.parametrize("tensor_payload", [False, True])
+@pytest.mark.parametrize("pattern", ["steady", "drop_at_finish", "rise_midstream", "oscillate"])
+def test_backlog_changes_preserve_every_frame_once(single_initial_chunk, initial_chunk_frames, tensor_payload, pattern):
+    tm = _make_transfer_manager(
+        tensor_payload=tensor_payload,
+        chunk_frames=25,
+        left_context=25,
+        initial_chunk_frames=initial_chunk_frames,
+        single_initial_chunk=single_initial_chunk,
+        backlog_chunk_frames=50,
+        max_num_seqs=64,
+    )
+    request = _make_request()
+    emitted = []
+    for frame in range(1, 141):
+        high = (
+            pattern == "steady"
+            or (pattern == "drop_at_finish" and frame < 140)
+            or (pattern == "rise_midstream" and frame >= 37)
+            or (pattern == "oscillate" and (frame // 13) % 2 == 0)
+        )
+        others = 63 if high else 1
+        tm.code_prompt_token_ids = defaultdict(list, {"req": tm.code_prompt_token_ids["req"]})
+        for idx in range(others):
+            tm.code_prompt_token_ids[f"other-{idx}"] = [None] * 50
+        out = slow_ar_to_dac_decoder_async_chunk(
+            tm,
+            {"audio_codes": torch.tensor([[frame, frame + 1000]]), "audio_code_valid": torch.tensor([True])},
+            request,
+        )
+        if out is not None:
+            codes = out.codes.audio.reshape(2, -1)
+            emitted.extend(codes[0, out.meta.left_context_size :].tolist())
+            tm.put_req_chunk["req"] += 1
+    # EOS may arrive in a separate callback after the last valid frame.
+    out = slow_ar_to_dac_decoder_async_chunk(tm, {}, request, is_finished=True)
+    if out is not None and out.codes.audio.numel():
+        emitted.extend(out.codes.audio.reshape(2, -1)[0, out.meta.left_context_size :].tolist())
+    assert emitted == list(range(1, 141))
+    assert slow_ar_to_dac_decoder_async_chunk(tm, {}, request, is_finished=True) is None
+
+
+def test_final_backlog_drop_flushes_entire_pending_tail():
+    tm = _make_transfer_manager(
+        tensor_payload=True,
+        chunk_frames=25,
+        left_context=25,
+        initial_chunk_frames=4,
+        single_initial_chunk=True,
+        backlog_chunk_frames=50,
+        max_num_seqs=64,
+    )
+    request = _make_request()
+    emitted = []
+    for frame in range(1, 41):
+        if frame < 40:
+            for idx in range(63):
+                tm.code_prompt_token_ids[f"other-{idx}"] = [None] * 30
+        else:
+            for idx in range(32, 63):
+                tm.code_prompt_token_ids.pop(f"other-{idx}")
+        out = slow_ar_to_dac_decoder_async_chunk(
+            tm,
+            {"audio_codes": torch.tensor([[frame]])},
+            request,
+            is_finished=frame == 40,
+        )
+        if out is not None:
+            emitted.extend(out.codes.audio[0, out.meta.left_context_size :].tolist())
+    assert emitted == list(range(1, 41))

tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ramp_decode_windows
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_Q = 2
+_RAMP = [1, 2, 4, 8, 16, 25]
+
+
+def _manager(**extra):
+    config = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "initial_codec_chunk_frames": 4, **extra}
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int),
+        ramp_chunk_count=defaultdict(int),
+        connector=SimpleNamespace(config={"extra": config}),
+    )
+
+
+def _request(request_id="req"):
+    return SimpleNamespace(external_req_id=request_id, sampling_params=SimpleNamespace(stop_token_ids=[]))
+
+
+def _frame(index: int) -> dict:
+    return {
+        "codes": {"audio": torch.tensor([[index, 1000 + index]], dtype=torch.long)},
+        "meta": {"codec_frame_valid": torch.ones(1, dtype=torch.int8)},
+    }
+
+
+def _stream(manager, num_frames, finish_after_last=True):
+    request = _request()
+    chunks = []
+
+    def emit(payload):
+        if payload is None:
+            return
+        codes = payload.codes.audio.reshape(_Q, -1)
+        chunks.append((codes[0].tolist(), int(payload.meta.left_context_size), bool(payload.meta.finished)))
+        # The transport counts every enqueued chunk.
+        manager.put_req_chunk[request.external_req_id] += 1
+        manager.ramp_chunk_count[request.external_req_id] += 1
+
+    for index in range(num_frames):
+        last = finish_after_last and index == num_frames - 1
+        emit(q3.talker2code2wav_async_chunk(manager, _frame(index), request, is_finished=last))
+    return chunks
+
+
+def test_ramp_decode_windows_include_left_context():
+    assert ramp_decode_windows(_RAMP, 25) == [1, 3, 7, 15, 31, 50]
+    assert ramp_decode_windows([1, 2], 0) == [1, 2]
+
+
+def test_ramp_chunks_grow_and_resend_left_context():
+    chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 56 + 25, finish_after_last=False)
+
+    new_frames = [len(window) - left for window, left, _finished in chunks]
+    assert new_frames == [1, 2, 4, 8, 16, 25, 25]
+    assert [left for _window, left, _finished in chunks] == [0, 1, 3, 7, 15, 25, 25]
+    assert [len(window) for window, _left, _finished in chunks] == [1, 3, 7, 15, 31, 50, 50]
+    emitted = 0
+    for window, left, _finished in chunks:
+        # Each window ends at the newest frame and re-sends `left` earlier ones.
+        assert window == list(range(emitted - left, emitted + len(window) - left))
+        emitted += len(window) - left
+    # The first chunk is the stream's first frame alone (no context).
+    assert chunks[0][0] == [0]
+
+
+def test_ramp_flushes_partial_tail_with_left_context():
+    chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 10)
+
+    assert [len(window) - left for window, left, _finished in chunks] == [1, 2, 4, 3]
+    window, left, finished = chunks[-1]
+    assert finished and left == 7 and window == list(range(0, 10))
+
+
+def test_ramp_finish_on_chunk_boundary_leaves_the_marker_to_the_transport():
+    manager = _manager(codec_chunk_ramp=_RAMP)
+    request = _request()
+    for index in range(3):
+        payload = q3.talker2code2wav_async_chunk(manager, _frame(index), request)
+        if payload is not None:
+            manager.ramp_chunk_count[request.external_req_id] += 1
+    # Frames 0 and 1-2 went out as chunks 0 and 1; the finish step has no frame.
+    finish = {"codes": {"audio": torch.zeros((1, _Q), dtype=torch.long)}, "meta": {"codec_frame_valid": torch.zeros(1)}}
+    assert q3.talker2code2wav_async_chunk(manager, finish, request, is_finished=True) is None
+
+
+def test_without_ramp_the_initial_chunk_setting_is_unchanged():
+    chunks = _stream(_manager(), 29 + 3)
+
+    assert [(len(window), left) for window, left, _finished in chunks] == [(4, 0), (29, 4), (28, 25)]
+
+
+def test_first_audio_marker_only_reaches_initial_codec_chunk():
+    manager = _manager(codec_chunk_ramp=[1, 1])
+    for index in range(2):
+        frame = _frame(index)
+        frame["meta"]["first_audio"] = torch.tensor([True])
+        payload = q3.talker2code2wav_async_chunk(manager, frame, _request())
+        assert bool(payload.meta.first_audio) is (index == 0)
+        manager.put_req_chunk["req"] += 1
+        manager.ramp_chunk_count["req"] += 1

tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+ACCEPT = 24
+HIDDEN = 8
+PROMPT = [151644, 872, 5, 6, 151645, 151644, 77091, 198]
+EOS = 151645
+
+
+class _Manager:
+    def __init__(self) -> None:
+        self.put_req_chunk: defaultdict[str, int] = defaultdict(int)
+        self.request_payload: dict[str, dict] = {}
+        self.config = SimpleNamespace(
+            hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=ACCEPT))
+        )
+
+    def process(self, output, request, is_finished=False):
+        payload = q3.thinker2talker_async_chunk(self, output, request, is_finished=is_finished)
+        if payload is not None:
+            self.put_req_chunk[request.external_req_id] += 1
+        return payload
+
+
+def _request(output_token_ids, *, history=True, max_tokens=64):
+    # The native data plane keeps token history only while chunk 0 is pending.
+    return SimpleNamespace(
+        external_req_id="r",
+        prompt_token_ids=list(PROMPT) if history else [],
+        output_token_ids=list(output_token_ids) if history else [],
+        all_token_ids=list(PROMPT) + list(output_token_ids) if history else [],
+        output_token_count=len(output_token_ids),
+        last_output_token_id=output_token_ids[-1] if output_token_ids else None,
+        resumable=False,
+        additional_information=None,
+        sampling_params=SimpleNamespace(all_stop_token_ids={EOS}, max_tokens=max_tokens),
+    )
+
+
+def _embed(token_id: int) -> torch.Tensor:
+    return torch.full((1, HIDDEN), float(token_id), dtype=torch.bfloat16)
+
+
+def _output(token_ids, sampled=None):
+    layers = {
+        0: torch.cat([_embed(t) for t in token_ids]),
+        ACCEPT: torch.randn(len(token_ids), HIDDEN).to(torch.bfloat16),
+    }
+    embed = {"tts_bos": _embed(1), "tts_eos": _embed(2), "tts_pad": _embed(3)}
+    if sampled is not None:
+        embed["sampled"] = sampled
+    return {"hidden_states": {"layers": layers}, "embed": embed}
+
+
+def _run_turn(tokens, *, publish, max_tokens=64, prefill_chunks=(len(PROMPT),)):
+    manager = _Manager()
+    rows, chunks, start = [], 0, 0
+    for i, size in enumerate(prefill_chunks):
+        last = i == len(prefill_chunks) - 1
+        sampled = _embed(tokens[0]) if last else torch.empty(0, dtype=torch.bfloat16)
+        request = _request(tokens[:1] if last else [], max_tokens=max_tokens)
+        payload = manager.process(_output(PROMPT[start : start + size], sampled if publish else None), request)
+        start += size
+        if payload is not None:
+            chunks += 1
+            rows.extend(payload.embed.prefill[len(PROMPT) :])
+    for step in range(1, len(tokens)):
+        history = manager.put_req_chunk["r"] == 0
+        request = _request(tokens[: step + 1], history=history, max_tokens=max_tokens)
+        output = _output([tokens[step - 1]], _embed(tokens[step]) if publish else None)
+        payload = manager.process(output, request)
+        if payload is None:
+            continue
+        chunks += 1
+        rows.extend(payload.embed.prefill[len(PROMPT) :] if payload.embed.prefill is not None else payload.embed.decode)
+    return [int(row[0]) for row in rows], chunks, manager
+
+
+@pytest.mark.parametrize(
+    ("tokens", "max_tokens"),
+    [
+        ([11, 12, 13, EOS], 64),  # ends on a stop token
+        ([11, 12, 13], 3),  # ends on the max_tokens-th token
+        ([11, EOS], 64),
+    ],
+)
+def test_sampled_stream_matches_captured_stream(tokens, max_tokens):
+    late, _chunks, _manager = _run_turn(tokens, publish=False, max_tokens=max_tokens)
+    early, _chunks, manager = _run_turn(tokens, publish=True, max_tokens=max_tokens)
+    # The final token is never processed, so neither path feeds it to the Talker.
+    assert early == late == tokens[:-1]
+    assert q3._SAMPLED_TEXT_STREAM in manager.request_payload["r"]
+
+
+def test_sampled_stream_sends_each_token_one_step_earlier():
+    manager = _Manager()
+    chunk0 = manager.process(_output(PROMPT, sampled=_embed(11)), _request([11]))
+    assert chunk0 is not None  # at the prefill step, not after the next decode step
+    assert torch.equal(chunk0.embed.prefill, torch.cat([_embed(t) for t in [*PROMPT, 11]]))
+    assert chunk0.hidden_states.output.shape[0] == len(PROMPT)
+    assert chunk0.ids.all == [*PROMPT, 11]
+    # Decode step 1 processes 11 and samples 12: it carries 12.
+    chunk1 = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False))
+    assert torch.equal(chunk1.embed.decode, _embed(12))
+
+
+def test_chunked_prefill_uses_sample_of_final_chunk_only():
+    early, chunks, _manager = _run_turn([11, 12, EOS], publish=True, prefill_chunks=(5, 3))
+    assert early == [11, 12] and chunks == 2
+
+
+def test_terminal_first_token_keeps_captured_stream():
+    manager = _Manager()
+    assert manager.process(_output(PROMPT, sampled=_embed(EOS)), _request([EOS])) is None
+    assert q3._SAMPLED_TEXT_STREAM not in manager.request_payload["r"]
+
+
+@pytest.mark.parametrize("prefill_rows", [0, 1, 3])
+def test_missing_sample_on_sampled_stream(prefill_rows):
+    manager = _Manager()
+    manager.process(_output(PROMPT, sampled=_embed(11)), _request([11]))
+    chunk = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False))
+    assert torch.equal(chunk.embed.decode, _embed(12))
+    if not prefill_rows:
+        with pytest.raises(RuntimeError, match="sample embedding missing"):
+            manager.process(_output([12]), _request([11, 12], history=False))
+    else:
+        before = manager.put_req_chunk["r"]
+        for tokens in (PROMPT[:prefill_rows], [11]):
+            partial = _output(tokens, sampled=torch.empty(0))
+            assert manager.process(partial, _request([11, 12], history=False)) is None
+        assert manager.put_req_chunk["r"] == before
+        chunk = manager.process(_output([12], sampled=_embed(13)), _request([11, 12, 13], history=False))
+        assert torch.equal(chunk.embed.decode, _embed(13))

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -1047,3 +1047,80 @@ def test_qwen2_5_omni_thinker2talker_full_payload_noop() -> None:
 
     payload = thinker2talker_full_payload(None, {"any": "thing"}, None)
     assert payload is None
+
+
+def _codec_transfer_manager(request_id: str, *, initial_frames: int = 2, chunk_frames: int = 4):
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int, {request_id: 0}),
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "initial_codec_chunk_frames": initial_frames,
+                    "codec_chunk_frames": chunk_frames,
+                    "codec_left_context_frames": 4,
+                }
+            }
+        ),
+    )
+
+
+def _codec_request(request_id: str) -> SimpleNamespace:
+    return SimpleNamespace(
+        external_req_id=request_id,
+        sampling_params=SimpleNamespace(stop_token_ids=[2150], stop_token_id=None),
+    )
+
+
+def test_talker2code2wav_async_chunk_reads_last_valid_row_of_multi_row_span() -> None:
+    # MRv2 eager frames: a prefill span carries its first frame in its last row.
+    request_id = "eager_prefill"
+    manager = _codec_transfer_manager(request_id)
+    codes = torch.zeros((5, 3), dtype=torch.long)
+    codes[-1] = torch.tensor([11, 12, 13])
+    valid = torch.tensor([0, 0, 0, 0, 1], dtype=torch.int8)
+
+    payload = q3.talker2code2wav_async_chunk(
+        manager, {"codes": {"audio": codes}, "meta": {"codec_frame_valid": valid}}, _codec_request(request_id)
+    )
+
+    assert payload is None  # below the initial chunk size
+    assert [frame.tolist() for frame in manager.code_prompt_token_ids[request_id]] == [[[11, 12, 13]]]
+
+
+def test_talker2code2wav_async_chunk_skips_invalid_frame_and_flushes_on_finish() -> None:
+    request_id = "eager_eos"
+    manager = _codec_transfer_manager(request_id)
+    manager.code_prompt_token_ids[request_id] = [torch.tensor([[7, 8, 9]])]
+    # The step that sampled codec EOS carries an invalid (placeholder) frame.
+    output = {
+        "codes": {"audio": torch.tensor([[2150, 0, 0]])},
+        "meta": {"codec_frame_valid": torch.tensor([0], dtype=torch.int8)},
+    }
+
+    payload = q3.talker2code2wav_async_chunk(manager, output, _codec_request(request_id), is_finished=True)
+
+    assert len(manager.code_prompt_token_ids[request_id]) == 1
+    assert payload is not None
+    assert payload.codes.audio.tolist() == [7, 8, 9]
+    assert payload.meta.finished.item() is True
+
+
+def test_talker2code2wav_async_chunk_without_validity_keeps_one_row_contract() -> None:
+    # V1 / MRv2 deferred frames: one row per step, zero rows are prefill
+    # placeholders and a stop-token CB0 marks the post-EOS step.
+    request_id = "legacy_rows"
+    manager = _codec_transfer_manager(request_id)
+    request = _codec_request(request_id)
+
+    assert (
+        q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.zeros((4, 3), dtype=torch.long)}}, request)
+        is None
+    )
+    assert q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[2150, 1, 1]])}}, request) is None
+    assert manager.code_prompt_token_ids[request_id] == []
+    q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[5, 6, 7]])}}, request)
+    payload = q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[8, 9, 10]])}}, request)
+
+    assert payload is not None  # two frames fill the initial chunk
+    assert payload.codes.audio.tolist() == [5, 8, 6, 9, 7, 10]

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -594,6 +594,40 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():
     torch.testing.assert_close(frames[:3], ref_code[-3:])
 
 
+@pytest.mark.parametrize("tensor_frames", [False, True])
+def test_async_chunks_preserve_codebook_order_with_reference_and_short_tail(tensor_frames):
+    tm = _tm(left_context=3, initial_chunk_frames=1)
+    rid = "r-codebook-order"
+    quantizers = 16
+    reference = [[1500 + f * quantizers + q for q in range(quantizers)] for f in range(5)]
+    tm.request_payload[rid] = torch.tensor(reference, dtype=torch.long)
+    generated = [[1 + f * quantizers + q for q in range(quantizers)] for f in range(30)]
+
+    for start, end, finished in [(0, 1, False), (1, 26, False), (26, 30, True)]:
+        rows = generated[:end]
+        if tensor_frames:
+            # Non-contiguous int32 rows exercise conversion independently of
+            # the native producer's contiguous CPU-long representation.
+            storage = torch.tensor(rows, dtype=torch.int32).repeat_interleave(2, dim=1)
+            tm.code_prompt_token_ids[rid] = list(storage[:, ::2].unbind())
+        else:
+            tm.code_prompt_token_ids[rid] = rows
+        payload = talker2code2wav_async_chunk(
+            transfer_manager=tm,
+            multimodal_output=None if finished else {},
+            request=_req(rid, finished=finished),
+            is_finished=finished,
+        )
+        assert payload is not None
+        expected = (reference[-3:] if start == 0 else []) + generated[start:end]
+        codes = payload.codes.audio
+        assert codes.dtype == torch.long and codes.device.type == "cpu" and codes.is_contiguous()
+        torch.testing.assert_close(codes.reshape(quantizers, -1).T, torch.tensor(expected))
+        assert payload.meta.left_context_size == (3 if start == 0 else 0)
+        assert payload.meta.finished.item() == finished
+        tm.put_req_chunk[rid] += 1
+
+
 def test_ref_code_context_can_be_buffered_before_first_emit():
     tm = _tm()
     rid = "r-ref-buffered"

tests/worker/test_gpu_ar_model_runner.py

@@ -16,6 +16,7 @@ import pytest
 import torch
 from vllm.sampling_params import SamplingParams
 from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.utils import CpuGpuBuffer
 from vllm.v1.worker import gpu_input_batch
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
@@ -1885,3 +1886,63 @@ class TestPreferModelSamplerNoneFallback:
             "least tolerates -- that fallback, then add its directory name to "
             "`expected` above. If you REMOVED one, drop its name."
         )
+
+
+@pytest.mark.parametrize("accepts_extra_args", [False, True])
+def test_prepare_hook_keeps_legacy_signature_and_orders_opt_in_metadata(accepts_extra_args: bool) -> None:
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = torch.nn.Module()
+    runner.model.accepts_runner_sampling_extra_args = accepts_extra_args
+    runner.requests = {
+        rid: CachedRequestState(
+            req_id=rid,
+            prompt_token_ids=[1],
+            mm_features=[],
+            sampling_params=params,
+            generator=None,
+            block_ids=([],),
+            num_computed_tokens=0,
+            output_token_ids=[],
+        )
+        for rid, params in (
+            ("a", SamplingParams(extra_args={"seed": 7})),
+            ("b", None),
+            ("unused", SamplingParams(extra_args={"seed": 9})),
+        )
+    }
+    runner.discard_request_mask = CpuGpuBuffer(3, dtype=torch.bool, device=torch.device("cpu"), pin_memory=False)
+    runner.discard_request_mask.np[:] = [True, False, True]
+    ids, positions = torch.tensor([3, 4]), torch.tensor([0, 1])
+    step_inputs = dict(
+        input_ids=ids,
+        positions=positions,
+        inputs_embeds=None,
+        num_computed_tokens=np.array([0, 1]),
+        num_scheduled_tokens=np.array([1, 1]),
+        input_ids_buffer=ids,
+    )
+    seen = {}
+
+    def legacy_hook(
+        *, input_ids, positions, inputs_embeds, req_ids, num_computed_tokens, num_scheduled_tokens, input_ids_buffer
+    ):
+        seen["req_ids"] = req_ids
+        return input_ids, positions
+
+    def opt_in_hook(*, sampling_extra_args, discard_mask, **kwargs):
+        seen["extra_args"] = sampling_extra_args
+        seen["discard"] = discard_mask.tolist()
+        return legacy_hook(**kwargs)
+
+    out = runner._call_prepare_runner_inputs(
+        opt_in_hook if accepts_extra_args else legacy_hook,
+        req_ids=["b", "a"],
+        **step_inputs,
+    )
+    assert out[0] is ids and out[1] is positions
+    assert seen["req_ids"] == ["b", "a"]
+    if accepts_extra_args:
+        assert seen["extra_args"] == [{}, {"seed": 7}]
+        assert seen["discard"] == [True, False]
+    else:
+        assert set(seen) == {"req_ids"}

tests/worker/test_yue2_request_state.py

@@ -0,0 +1,215 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2's worker-side sampler follows retained engine outputs on recompute."""
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.utils import CpuGpuBuffer
+from vllm.v1.worker.gpu_input_batch import CachedRequestState
+
+from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+from vllm_omni.entrypoints.openai.tts_adapters.yue2 import Yue2Adapter
+from vllm_omni.model_executor.models.yue2 import yue2 as y
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def model(mocker):
+    """Run real state/sampling control flow with CPU staging and fixed draws."""
+    model = object.__new__(y.Yue2ForCausalLM)
+    torch.nn.Module.__init__(model)
+    model._states = {}
+    model._pending = None
+    model._step_rows = []
+    model._step_discard = None
+    model._decode_t0 = {}
+    model._last_mm = None
+    model._audio_queue = []
+    model._max_model_len = 4096
+    model._stage = torch.empty(1024, dtype=torch.long)
+    model._pinned = torch.empty((2, 64), dtype=torch.long)
+    model._synthesis = mocker.Mock(spec=y._SynthesisQueue)
+    mocker.patch.object(model, "_finish_request_safely")
+    mocker.patch.object(model, "_sampler_graph", return_value=None)
+    # CPU copies complete synchronously; this replaces only CUDA completion.
+    event = mocker.Mock(spec=torch.Event)
+    mocker.patch.object(y.torch, "Event", return_value=event)
+    mocker.patch.object(y, "sample_rows", return_value=(torch.tensor([y.CODEC_OFFSET + 7]), torch.tensor([False])))
+    constants = y._RowConstants("r", "semantic", 1.0, 0.9, 100, 1.1, 2, 1, 3, 17, False, 8)
+    model._states["r"] = y._RequestState("r", constants, prompt_len=3, prefix_ids=[1, 2, 3])
+    return model
+
+
+def test_pending_draw_is_discarded_after_chunked_recompute(model) -> None:
+    """A preempted in-flight draw never advances history or the frame limit."""
+    logits = torch.zeros((1, y.VOCAB_SIZE))
+    model._step_rows = [("r", 0, 3)]
+    model.sample(logits, None)
+    state = model._states["r"]
+    # Other requests run while r waits for KV space. Resolving the host copy
+    # may resolve a tentative token, but the resumed row must roll it back.
+    model._step_rows = []
+    model.sample(logits[:0], None)
+    assert state.job is None
+    # A partial recompute must neither accept the pending draw nor draw again.
+    model._step_rows = [("r", 0, 2)]
+    model._step_discard = [True]
+    model.sample(logits, None)
+    assert state.job is None
+    # Completing recompute has retained zero outputs: discard the old draw.
+    model._step_rows = [("r", 2, 1)]
+    model._step_discard = [False]
+    model.sample(logits, None)
+    assert state.history == []
+    # Only this new draw is retained in the next decode input.
+    model._step_rows = [("r", 3, 1)]
+    model.sample(logits, None)
+    assert state.history == [y.CODEC_OFFSET + 7]
+    assert not state.finish_ready and not state.truncated
+
+
+@pytest.mark.parametrize("token", [y.MUSIC_END, y.CODEC_OFFSET + 7])
+def test_discarded_terminal_draw_does_not_start_synthesis(model, mocker, token: int) -> None:
+    """Dropped EOS/frame-limit draws cannot prematurely finish a song."""
+    state = model._states["r"]
+    state.history = [y.CODEC_OFFSET + 2] * 2
+    mocker.patch.object(y, "sample_rows", return_value=(torch.tensor([token]), torch.tensor([False])))
+    logits = torch.zeros((1, y.VOCAB_SIZE))
+    model._step_rows = [("r", 4, 1)]
+    model.sample(logits, None)
+    model._step_rows = []
+    model.sample(logits[:0], None)
+    assert state.job is None
+    model._step_rows = [("r", 0, 5)]
+    model.sample(logits, None)
+    assert state.history == [y.CODEC_OFFSET + 2] * 2
+    assert not state.finish_ready and not state.truncated and not state.end_drawn
+    assert state.job is None and not state.finished
+    model._finish_request_safely.assert_not_called()
+
+
+def test_recompute_clears_terminal_state_and_cancels_stale_job(model) -> None:
+    """Rolling back a drawn end also releases synthesis and revives the row."""
+    state = model._states["r"]
+    state.history = [y.CODEC_OFFSET + 2] * 2
+    state.end_drawn = True
+    state.finished = state.finish_ready = True
+    state.hold_steps = 3
+    state.end_len = 6
+    state.delivered = (torch.zeros((2, 0)), False, False)
+
+    def song(keep):
+        yield
+        return torch.zeros((2, 0))
+
+    job = y._SynthesisJob("r", song)
+    state.job = job
+    model._reconcile_history(state, 2)
+    model._synthesis.cancel.assert_called_once_with(job)
+    assert state.job is None and state.hold_steps == 0
+    assert not state.finished and not state.finish_ready and not state.truncated and not state.end_drawn
+    assert state.end_len is None and state.delivered is None
+
+
+def test_dropped_end_step_redelivers_song_and_end(model, mocker) -> None:
+    """A preemption that drops the step emitting the end token also drops
+    that step's audio; the resumed request must deliver both again."""
+    state = model._states["r"]
+    state.history = [y.CODEC_OFFSET + 2] * 3  # at the 3-frame budget
+    state.truncated = True
+    state.hold_steps = 2
+    audio = torch.ones((2, 5))
+    job = mocker.Mock(t_start=1.0, t_done=2.0)
+    job.done.return_value = True
+    state.job = job
+    model._synthesis.complete.return_value = (audio, False, False)
+    logits = torch.zeros((1, y.VOCAB_SIZE))
+
+    # Engine holds 3 frames + 2 HOLDs; this step ships the song and its end.
+    model._step_rows = [("r", 7, 1)]
+    out = model.sample(logits, None)
+    assert out.sampled_token_ids.tolist() == [[y.MUSIC_END]]
+    assert state.finished and state.end_len == 6
+    assert [entry[:2] for entry in model._audio_queue] == [("r", audio)]
+
+    # An async lookahead row (end token in flight) delivers nothing again.
+    model._step_rows = [("r", 8, 1)]
+    model.sample(logits, None)
+    assert len(model._audio_queue) == 1
+
+    # Preempted with the end step in flight: the resumed row lacks it.
+    model._step_rows = [("r", 0, 8)]
+    model._step_discard = [False]
+    out = model.sample(logits, None)
+    assert out.sampled_token_ids.tolist() == [[y.MUSIC_END]]
+    assert [entry[:2] for entry in model._audio_queue] == [("r", audio), ("r", audio)]
+    assert model._audio_queue[1][2:] == (True, False)  # truncated, not failed
+    assert state.end_len == 6
+    model._synthesis.complete.assert_called_once()
+
+
+def test_accepted_draw_still_reaches_frame_limit(model) -> None:
+    """Confirmed codec frames retain the existing delayed finish behavior."""
+    logits = torch.zeros((1, y.VOCAB_SIZE))
+    for comp, span in ((0, 3), (3, 1), (4, 1)):
+        model._step_rows = [("r", comp, span)]
+        model.sample(logits, None)
+    state = model._states["r"]
+    assert len(state.history) == 2
+    model._resolve_pending()
+    model._reconcile_history(state, 3)
+    assert len(state.history) == 3
+    assert state.finish_ready and state.truncated
+
+
+def test_real_prepare_hook_admits_reordered_requests(model) -> None:
+    """Runner opt-in metadata reaches the actual YuE2 admission hook."""
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = model
+    runner.requests = {}
+    for rid, seed in (("a", 7), ("b", 9)):
+        runner.requests[rid] = CachedRequestState(
+            req_id=rid,
+            prompt_token_ids=[1, 2, 3],
+            mm_features=[],
+            sampling_params=SamplingParams(extra_args={y.KEY_PREFIX_IDS: [1, 2, 3], y.KEY_SEED: seed}),
+            generator=None,
+            block_ids=([],),
+            num_computed_tokens=0,
+            output_token_ids=[],
+        )
+    runner.discard_request_mask = CpuGpuBuffer(2, dtype=torch.bool, device=torch.device("cpu"), pin_memory=False)
+    runner.discard_request_mask.np[:] = [True, False]
+    ids = torch.tensor([1, 2, 1, 2, 3], dtype=torch.int32)
+    positions = torch.tensor([0, 1, 0, 1, 2])
+    runner._call_prepare_runner_inputs(
+        model.prepare_runner_inputs,
+        req_ids=["b", "a"],
+        input_ids=ids,
+        positions=positions,
+        inputs_embeds=None,
+        num_computed_tokens=np.array([0, 0]),
+        num_scheduled_tokens=np.array([2, 3]),
+        input_ids_buffer=ids,
+    )
+    assert model._states["b"].constants.seed == 9
+    assert model._states["a"].constants.seed == 7
+    assert model._states["b"].prompt_len == 3
+    assert model._states["b"].prefix_ids == [1, 2, 3]
+    assert model._step_discard == [True, False]
+
+
+@pytest.mark.parametrize("spare", [0, 1, 4095, 4096, 5000])
+def test_hold_budget_matches_engine_headroom(spare: int) -> None:
+    """Adapter metadata advertises only the HOLD steps left after clipping."""
+    frames = 250
+    prompt = {"prompt_token_ids": [1] * (y.CONTEXT - frames - 1 - spare)}
+    adapter = object.__new__(Yue2Adapter)
+    request = OpenAICreateSpeechRequest(model="m-a-p/YuE2-3B", input="lyrics", max_new_tokens=frames)
+    params = adapter.apply_sampling_overrides([SamplingParams()], request, prompt)[0]
+    assert params.extra_args[y.KEY_MAX_HOLD_STEPS] == min(spare, 4096)
+    assert params.max_tokens == frames + 1 + min(spare, 4096)

tests/worker/test_yue2_sampler.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2's batched sampler draws exactly what the per-row reference draws.
+
+The model samples all rows of a preset as one batch, over the phase's own
+columns only, and writes the multinomial out as ``argmax(p / q)`` with
+``q ~ Exp(1)`` so the draw stays on the device. The reference below is the
+upstream per-row arithmetic over the full vocabulary with
+``torch.multinomial``. Same seeds must give the same tokens, and each
+request's generator must advance by exactly the same amount per step, or a
+request's song changes with the batch it happens to share.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.yue2.yue2 import (
+    ABC_END,
+    ABC_SAMPLING,
+    CODEC_OFFSET,
+    CODEC_SIZE,
+    EOD,
+    MUSIC_END,
+    PHASE_VOCAB,
+    SEMANTIC_SAMPLING,
+    VOCAB_SIZE,
+    sample_rows,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _reference_scores(logits: torch.Tensor, preset: dict, history: list[int], phase: str) -> torch.Tensor:
+    """Upstream per-row scores over the full vocabulary (float32)."""
+    scores = logits.float()  # a new tensor: the logits are bf16
+    end = ABC_END if phase == "abc" else MUSIC_END
+    allowed = torch.full_like(scores, -torch.inf)
+    if phase == "abc":
+        allowed[:EOD] = 0
+    else:
+        allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0
+    allowed[end] = 0
+    scores = scores + allowed
+    if len(history) < preset["min_tokens"]:
+        scores[end] = -torch.inf
+    recent = history[-preset["penalty_window"] :]
+    if preset["repetition_penalty"] != 1.0 and recent:
+        freq = torch.zeros_like(scores)
+        freq.scatter_add_(-1, torch.tensor(recent), torch.ones(len(recent)))
+        alpha = preset["repetition_penalty"] ** freq
+        scores = torch.where(scores < 0, scores * alpha, scores / alpha)
+    if preset["temperature"] == 0:
+        return scores
+    scores = scores / preset["temperature"]
+    threshold = scores.topk(preset["top_k"]).values[-1]
+    scores = scores.masked_fill(scores < threshold, -torch.inf)
+    values, indices = scores.sort(descending=True, stable=True)
+    probabilities = values.softmax(-1)
+    removed = probabilities.cumsum(-1) - probabilities > preset["top_p"]
+    removed[:1] = False
+    values = values.masked_fill(removed, -torch.inf)
+    return values.scatter(-1, indices, values)
+
+
+def _reference_draw(scores: torch.Tensor, generator: torch.Generator, greedy: bool) -> int:
+    if greedy:
+        return int(scores.argmax())
+    return int(torch.multinomial(scores.softmax(-1), 1, generator=generator))
+
+
+def _batched_draw(logits, preset, histories, phase, generators) -> list[int]:
+    vocab = PHASE_VOCAB[phase]
+    window = preset["penalty_window"]
+    recent = [[vocab.col(t) for t in history[-window:]] for history in histories]
+    width = max(map(len, recent))
+    window_cols = torch.tensor([r + [vocab.num_cols] * (width - len(r)) for r in recent]) if width else None
+    block = [len(history) < preset["min_tokens"] for history in histories]
+    ids, bad = sample_rows(
+        logits,
+        vocab,
+        temperature=preset["temperature"],
+        top_p=preset["top_p"],
+        top_k=preset["top_k"],
+        repetition_penalty=preset["repetition_penalty"],
+        window_cols=window_cols,
+        block_end=torch.tensor(block) if any(block) else None,
+        generators=generators,
+    )
+    assert not bad.any()
+    return ids.tolist()
+
+
+@pytest.mark.parametrize(
+    ("phase", "rows", "min_tokens", "temperature"),
+    [
+        ("semantic", 1, None, None),
+        ("semantic", 3, None, None),
+        ("semantic", 3, 0, None),  # the end token may be drawn
+        ("semantic", 2, None, 0.0),  # greedy
+        ("abc", 2, None, None),
+        ("abc", 2, 0, None),
+    ],
+)
+def test_batched_draws_match_per_row_multinomial(phase, rows, min_tokens, temperature) -> None:
+    preset = dict(ABC_SAMPLING if phase == "abc" else SEMANTIC_SAMPLING)
+    if min_tokens is not None:
+        preset["min_tokens"] = min_tokens
+    if temperature is not None:
+        preset["temperature"] = temperature
+    end = ABC_END if phase == "abc" else MUSIC_END
+    reference = [torch.Generator().manual_seed(1000 + row) for row in range(rows)]
+    batched = [torch.Generator().manual_seed(1000 + row) for row in range(rows)]
+    histories: list[list[int]] = [[] for _ in range(rows)]
+    logits_generator = torch.Generator().manual_seed(7)
+    for _ in range(40):
+        logits = (torch.randn((rows, VOCAB_SIZE), generator=logits_generator) * 3).to(torch.bfloat16)
+        # Peaky over a few ids, like a real LM, so the penalty and top-p matter.
+        if phase == "semantic":
+            logits[:, CODEC_OFFSET : CODEC_OFFSET + 300] += 6
+        else:
+            logits[:, :300] += 6
+        logits[:, end] += 4
+        expected = [
+            _reference_draw(
+                _reference_scores(logits[row], preset, histories[row], phase), reference[row], temperature == 0
+            )
+            for row in range(rows)
+        ]
+        assert _batched_draw(logits, preset, histories, phase, batched) == expected
+        for row in range(rows):
+            # A greedy request never reads its stream (its temperature is
+            # fixed for life), so only sampling rows must advance in step.
+            if temperature != 0:
+                assert torch.equal(reference[row].get_state(), batched[row].get_state())
+            if expected[row] != end:
+                histories[row].append(expected[row])
+
+
+def test_a_row_draws_the_same_alone_or_in_a_batch() -> None:
+    """A request's song must not depend on which requests share its step."""
+    preset = SEMANTIC_SAMPLING
+    logits = torch.randn((3, VOCAB_SIZE), generator=torch.Generator().manual_seed(3)).to(torch.bfloat16)
+    histories = [[CODEC_OFFSET + i for i in range(60)], [], [CODEC_OFFSET + 5] * 10]
+    together = _batched_draw(
+        logits, preset, histories, "semantic", [torch.Generator().manual_seed(s) for s in (1, 2, 3)]
+    )
+    alone = [
+        _batched_draw(
+            logits[row : row + 1], preset, histories[row : row + 1], "semantic", [torch.Generator().manual_seed(s)]
+        )[0]
+        for row, s in enumerate((1, 2, 3))
+    ]
+    assert together == alone
+
+
+@pytest.mark.parametrize("phase", ["abc", "semantic"])
+def test_phase_vocab_columns_round_trip(phase) -> None:
+    vocab = PHASE_VOCAB[phase]
+    end = ABC_END if phase == "abc" else MUSIC_END
+    ids = [end, 0, EOD - 1] if phase == "abc" else [end, CODEC_OFFSET, CODEC_OFFSET + CODEC_SIZE - 1]
+    cols = torch.tensor([vocab.col(i) for i in ids])
+    assert vocab.col(end) == vocab.end_col
+    assert cols.max() < vocab.num_cols
+    assert vocab.ids(cols).tolist() == ids
+    logits = torch.randn((2, VOCAB_SIZE))
+    assert torch.equal(vocab.take(logits)[:, cols], logits[:, ids])

tests/worker_v2/test_capture_query_bound_gpu.py

@@ -44,6 +44,7 @@ def _check_prefill():
     spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16)
     kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)])
     state = OmniModelState.__new__(OmniModelState)
+    state.vllm_config = c
     state.max_model_len = 256
     state.supports_mm_inputs = False
     with set_current_vllm_config(c):

tests/worker_v2/test_decode_query_bound_gpu.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Real FA3 graph replay must preserve the decode GQA layout."""
+
+from dataclasses import dataclass, replace
+
+import numpy as np
+import pytest
+import torch
+from transformers import Qwen2Config
+from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.utils import AttentionGroup
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+pytestmark = [pytest.mark.core_model]
+
+
+@dataclass
+class _AttentionScales:
+    _q_scale: torch.Tensor
+    _k_scale: torch.Tensor
+    _v_scale: torch.Tensor
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("graph_mode", [CUDAGraphMode.FULL_AND_PIECEWISE, CUDAGraphMode.FULL_DECODE_ONLY])
+def test_full_decode_graph_matches_fa3_aot_schedule(graph_mode, tmp_path):
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("requires FA3 on Hopper or newer")
+    torch.manual_seed(17)
+    device = torch.device("cuda:0")
+    c = VllmConfig()
+    Qwen2Config(
+        architectures=["Qwen2ForCausalLM"],
+        hidden_size=1024,
+        intermediate_size=2048,
+        num_hidden_layers=1,
+        num_attention_heads=16,
+        num_key_value_heads=2,
+        max_position_embeddings=2048,
+    ).save_pretrained(tmp_path)
+    c.model_config = ModelConfig(model=str(tmp_path), skip_tokenizer_init=True, dtype="bfloat16")
+    c.scheduler_config.max_num_seqs = 64
+    c.compilation_config.cudagraph_mode = graph_mode
+    c.compilation_config.max_cudagraph_capture_size = 128
+    c.attention_config.flash_attn_version = 3
+    spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=64, dtype=torch.bfloat16)
+    kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)])
+    state = OmniModelState.__new__(OmniModelState)
+    state.vllm_config = c
+    state.max_model_len = 2048
+    state.supports_mm_inputs = False
+    with set_current_vllm_config(c):
+        group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0)
+        group.create_metadata_builders(c, device)
+        impl = FlashAttentionImpl(16, 64, 64**-0.5, 2, None, None, "auto")
+        assert impl.vllm_flash_attn_version == 3
+        buffers = InputBuffers(64, 128, device)
+        batch = InputBatch.make_dummy(8, 8, buffers, max_query_len=None)
+        tables = torch.arange(1024, dtype=torch.int32, device=device).reshape(8, 128)
+        slots = torch.zeros((1, 8), dtype=torch.int64, device=device)
+        cache = torch.randn(1024, 2, 16, 128, device=device, dtype=torch.bfloat16)
+        q = torch.randn(8, 16, 64, device=device, dtype=torch.bfloat16)
+        k = v = torch.empty(8, 2, 64, device=device, dtype=torch.bfloat16)
+        layer = _AttentionScales(
+            _q_scale=torch.ones((), device=device),
+            _k_scale=torch.ones((), device=device),
+            _v_scale=torch.ones((), device=device),
+        )
+        md = state.prepare_attn(batch, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig, for_capture=True)[
+            "attn"
+        ]
+        out = torch.full_like(q, float("nan"))
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                impl.forward(layer, q, k, v, cache, md, out)
+        stream.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            impl.forward(layer, q, k, v, cache, md, out)
+        lengths = np.arange(37, 45, dtype=np.int32)
+        batch.seq_lens.copy_(torch.tensor(lengths, device=device))
+        runtime = replace(batch, max_query_len=1, seq_lens_cpu_upper_bound=torch.tensor(lengths))
+        eager_md = state.prepare_attn(runtime, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig)["attn"]
+        expected = torch.empty_like(q)
+        impl.forward(layer, q, k, v, cache, eager_md, expected)
+        state.prepare_attn(runtime, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig)
+        graph.replay()
+        torch.accelerator.synchronize()
+        assert md.max_query_len == 1
+        torch.testing.assert_close(out, expected, atol=0.002, rtol=0.002)

tests/worker_v2/test_index_staging.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from vllm_omni.utils.device_copy import index_to_device
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+def test_index_staging_survives_later_copies_and_graph_replay():
+    stream = torch.cuda.Stream()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.stream(stream):
+        with torch.cuda.graph(graph, stream=stream):
+            captured = index_to_device([17] * 6, "cuda")
+    for _ in range(32):
+        index_to_device([99] * 6, "cuda")
+    graph.replay()
+    torch.accelerator.synchronize()
+    assert captured.tolist() == [17] * 6

tests/worker_v2/test_omni_ar_model_runner.py

@@ -25,6 +25,9 @@ class _FakeStream:
     def wait_stream(self, _stream) -> None:
         pass
 
+    def wait_event(self, _event) -> None:
+        pass
+
 
 class _FakeEvent:
     def record(self, _stream) -> None:
@@ -138,6 +141,7 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history)
 def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     runner = OmniARModelRunner.__new__(OmniARModelRunner)
     runner.model_config = SimpleNamespace(async_chunk=True)
+    runner.model = SimpleNamespace()
     runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None]
     runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0
     waited: list[object] = []
@@ -307,6 +311,231 @@ def test_request_reference_codes_preserve_local_axis(prefill_first, padded):
         assert torch.equal(outputs[i]["codes.audio"], codes[offsets[i] : offsets[i + 1]])
 
 
+def test_stream_audio_cancel_rejects_inflight_output_and_id_reuse() -> None:
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    cancelled = buffer.add("request")
+    buffer.finish({"request"})
+    replacement = buffer.add("request")
+    assert cancelled.push(torch.ones(4), False) is None
+    assert not cancelled.pending
+    assert replacement.push(torch.full((4,), 2.0), False).tolist() == [2.0] * 4
+    buffer.finish({"request"})
+    assert not buffer.requests
+
+
+def test_stream_audio_length_end_flushes_first_and_partial_chunks() -> None:
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("request")
+    output = StreamingAudioOutput(
+        torch.ones(1, 4),
+        torch.tensor([True]),
+        torch.tensor(24000),
+        None,
+        np.array([0, 1]),
+        [state],
+        np.array([True]),
+    )
+    assert output.get_output()[0]["model_outputs"].numel() == 4
+    assert not state.active
+    state = buffer.add("another")
+    assert state.push(torch.ones(4), False).numel() == 4
+    assert state.push(torch.ones(4), False) is None
+    assert state.push(torch.ones(4), True).numel() == 8
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64, torch.bfloat16])
+def test_stream_audio_batch_views_retain_pcm_after_output_release(dtype):
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    states = [buffer.add("a"), buffer.add("b")]
+    expected: list[list[torch.Tensor]] = [[], []]
+    emitted: list[list[torch.Tensor]] = [[], []]
+    for step in range(5):
+        # Mixed prefill/decode spans, invalid rows, and a final partial chunk.
+        wav = (torch.arange(16).reshape(4, 4) + 100 * step).to(dtype)
+        valid = torch.tensor([False, True, True, step != 4])
+        output = StreamingAudioOutput(
+            wav,
+            valid,
+            torch.tensor(24000),
+            None,
+            np.array([0, 3, 4]),
+            states,
+            np.array([step == 4, step == 4]),
+        )
+        for i, (start, end) in enumerate([(0, 3), (3, 4)]):
+            expected[i].append(wav[start:end][valid[start:end]].reshape(-1).clone())
+        for i, payload in enumerate(output.get_output()):
+            if payload is not None:
+                emitted[i].append(payload["model_outputs"])
+        del output, wav  # Pending views must keep the old host allocation alive.
+    for i in range(2):
+        assert torch.equal(torch.cat(emitted[i]), torch.cat(expected[i]))
+        assert torch.cat(emitted[i]).dtype == dtype
+        assert not states[i].active and not states[i].pending
+        assert states[i].pending_samples == 0
+
+
+def test_stream_audio_numpy_direct_first_cancel_and_mixed_fallback():
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("r")
+    state.accept_first_audio()
+    assert state.push(np.arange(8, dtype=np.float32), False) is None
+    assert state.pending_samples == 4
+    # A dtype fallback can arrive while a NumPy frame is pending.
+    chunk = state.push(torch.arange(8, 12, dtype=torch.bfloat16), True)
+    assert chunk.tolist() == list(range(4, 12))
+    assert state.pending_samples == 0
+    buffer.finish({"r"})
+    replacement = buffer.add("r")
+    assert state.push(np.ones(4, dtype=np.float32), False) is None
+    assert replacement.push(np.full(4, 99, dtype=np.float32), False).tolist() == [99] * 4
+    replacement.push(np.ones(4, dtype=np.float32), False)
+    assert replacement.pending_samples == 4
+    buffer.finish({"r"})
+    assert replacement.pending_samples == 0 and not replacement.pending
+
+
+def test_stream_audio_prefill_length_end_reaches_output(mocker) -> None:
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+    from vllm.v1.worker.gpu.states import RequestState
+
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    owner = mocker.Mock()
+    owner.vllm_config.model_config.max_model_len = 4096
+    owner._stream_decode_event = None
+    eager = EagerMTPState(owner)
+    eager._audio_buffer = StreamingAudioBuffer(4, 25)
+    eager._audio_buffer.add("request")
+    eager._frame_requests = {"request"}
+    batch = InputBatch.__new__(InputBatch)
+    batch.num_reqs, batch.req_ids = 1, ["request"]
+    batch.num_computed_tokens_np = np.array([0])
+    batch.num_scheduled_tokens = np.array([3])
+    batch.idx_mapping_np = np.array([0])
+    batch.query_start_loc_np = np.array([0, 3])
+    batch.is_prefilling_np = np.array([True])
+    states = RequestState.__new__(RequestState)
+    states.max_seq_len = np.array([4])
+    mm = {
+        "model_outputs": torch.ones(3, 4),
+        "meta": {"codec_frame_valid": torch.tensor([0, 0, 1])},
+        "sr": [torch.tensor(24000)],
+    }
+    output = eager.prepare_audio_output(batch, states, mm)
+    assert output.get_output()[0]["model_outputs"].numel() == 4
+    assert "model_outputs" not in mm
+
+
+def test_stream_audio_preemption_preserves_position_rng_and_pending_pcm(mocker) -> None:
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    owner = mocker.Mock()
+    owner._stream_pos = {"request": 37}
+    owner._eager_embeds = torch.arange(18).reshape(6, 3).float()
+    owner._eager_ready = {}
+    owner.intermediate_buffer.buffers = [{} for _ in range(6)]
+    generator = torch.Generator().manual_seed(42)
+    owner._mtp_generators = {"request": generator}
+    saved = [torch.arange(6)]
+    owner.model.stream_decoder.save_slot.return_value = saved
+    eager = EagerMTPState(owner)
+    eager._talker_inputs["request"] = TalkerInputs(torch.ones(20, 3), 10)
+    eager._side_stream = mocker.Mock()
+    eager._audio_buffer = StreamingAudioBuffer(4, 25)
+    request = eager._audio_buffer.add("request")
+    request.push(torch.ones(4), False)
+    request.push(torch.ones(4), False)
+    eager.suspend_audio("request", 2)
+    owner._stream_pos.clear()
+    owner._mtp_generators.clear()
+    eager.resume_audio("request", 5)
+    assert owner._stream_pos["request"] == 37
+    assert owner._mtp_generators["request"] is generator
+    assert eager._restore_audio["request"] is saved
+    assert sum(part.numel() for part in request.pending) == 4
+    # A second preemption can happen before the restored slot is used.
+    eager.suspend_audio("request", 5)
+    owner.model.stream_decoder.save_slot.assert_called_once_with(2)
+    eager.finish_audio({"request"})
+    assert not eager._suspended_audio and not eager._restore_audio
+    assert not eager._audio_buffer.requests and not request.active
+
+
+def test_stream_audio_replay_uses_exact_inputs_and_consumes_last_frame_once(mocker):
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs
+
+    owner = mocker.Mock()
+    owner._eager_embeds = torch.tensor([[20.0, 30.0]])
+    owner.model.preprocess.return_value = (None, None, {"mtp_inputs": (None, torch.tensor([[2.0, 3.0]]))})
+    owner.intermediate_buffer.buffers = [{"req_id": "request"}]
+    eager = EagerMTPState(owner)
+    eager._restore_audio["request"] = []
+    history = torch.arange(8).reshape(4, 2).float()
+    eager._talker_inputs["request"] = TalkerInputs(history, 4)
+    first = torch.empty(2, 2)
+    assert eager.replay_inputs("request", 0, 0, torch.ones(2), first)
+    torch.testing.assert_close(first, history[:2])
+    owner.model.preprocess.assert_not_called()
+    last = torch.empty(3, 2)
+    assert eager.replay_inputs("request", 0, 2, torch.ones(3), last)
+    torch.testing.assert_close(last, torch.cat((history[2:], torch.tensor([[22.0, 33.0]]))))
+    owner.model.preprocess.assert_called_once()
+    with pytest.raises(RuntimeError, match="exceeds saved inputs"):
+        eager.replay_inputs("request", 0, 4, torch.ones(2), torch.empty(2, 2))
+
+
+def test_stream_reference_priming_groups_lengths_without_padding_or_slot_aliasing(mocker):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+
+    primes = [(4, torch.full((26, 2), 4)), (1, torch.full((3, 2), 1)), (7, torch.full((26, 2), 7))]
+    calls = []
+
+    def stream(codes, slots, pos):
+        calls.append((codes.clone(), slots.clone(), pos.clone()))
+
+    decoder = mocker.Mock(side_effect=stream, device=torch.device("cpu"))
+    model = mocker.Mock(stream_decoder=decoder, stream_prime_graphs=None, stream_chunk_frames=25)
+    Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, primes)
+    assert len(calls) == 3
+    for (codes, slots, pos), frames, indices, position in zip(
+        calls, [25, 1, 3], [[4, 7], [4, 7], [1]], [0, 25, 0], strict=True
+    ):
+        assert codes.shape == (len(indices), frames, 2)
+        assert codes.dtype == slots.dtype == pos.dtype == torch.int32
+        assert slots.tolist() == indices and pos.tolist() == [position] * len(indices)
+        for row, idx in enumerate(indices):
+            assert torch.all(codes[row] == idx)
+    Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, [])
+    assert len(calls) == 3
+
+
+@pytest.mark.parametrize("context_frames", [25, 72])
+def test_model_owned_reference_context_respects_configured_boundary(mocker, context_frames):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+
+    model = mocker.Mock(stream_ref_context_frames=context_frames, talker_config=mocker.Mock(num_code_groups=16))
+    reference = torch.arange(100 * 16).reshape(100, 16)
+    info = {"codes": {"ref": reference}}
+    got = Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, info)
+    torch.testing.assert_close(got, reference[-context_frames:])
+    assert info["codes"]["ref"] is reference
+    assert Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {}) is None
+    assert (
+        Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {"codes": {"ref": torch.empty(0)}}) is None
+    )
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, streaming):
     monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
@@ -340,6 +569,136 @@ def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch,
     assert (result.pooler_output is None) == streaming
 
 
+def test_stream_audio_history_replays_exact_inputs_after_batch_reordering(mocker):
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+
+    owner = mocker.Mock()
+    owner.vllm_config.model_config.max_model_len = 12
+    eager = EagerMTPState(owner)
+    first = torch.arange(9, dtype=torch.float32).reshape(3, 3)
+    second = torch.arange(9, 18, dtype=torch.float32).reshape(3, 3)
+    expected = {"a": torch.cat((first[:2], second[2:])), "b": torch.cat((first[2:], second[:2]))}
+    batch = InputBatch.__new__(InputBatch)
+    batch.req_ids = ["a", "b"]
+    batch.query_start_loc_np = np.array([0, 2, 3])
+    batch.num_computed_tokens_np = np.array([0, 0])
+    eager.record_inputs(batch, first)
+    batch.req_ids = ["b", "a"]
+    batch.num_computed_tokens_np = np.array([1, 2])
+    eager.record_inputs(batch, second)
+    # Later input-buffer reuse must not change the conditioned replay history.
+    first.fill_(-1)
+    second.fill_(-2)
+    for request_id, wanted in expected.items():
+        eager._restore_audio[request_id] = []
+        replayed = torch.empty_like(wanted)
+        assert eager.replay_inputs(request_id, 0, 0, torch.zeros(3, dtype=torch.long), replayed)
+        torch.testing.assert_close(replayed, wanted)
+
+
+@pytest.mark.parametrize("valid", [False, True])
+def test_stream_audio_direct_first_requires_only_real_pcm_on_terminal(valid):
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("r")
+    state.accept_first_audio()
+    first = StreamingAudioOutput(
+        torch.arange(4).reshape(1, 4),
+        torch.tensor([valid]),
+        torch.tensor(24000),
+        None,
+        np.array([0, 1]),
+        [state],
+        np.array([True]),
+    )
+    payload = first.get_output()[0]
+    if valid:
+        assert payload == {FIRST_AUDIO_REQUIRED_KEY: torch.tensor(True)}
+    else:
+        assert payload is None  # EOS has no promised first PCM to wait for.
+    assert not state.active and not state.pending
+    buffer.finish({"r"})
+    replacement = buffer.add("r")
+    assert replacement.push(torch.full((4,), 99), False).tolist() == [99] * 4
+
+
+def test_stream_audio_direct_first_drops_one_frame_and_retains_partial_suffix():
+    from vllm_omni.worker_v2.streaming_audio import AudioRequest
+
+    state = AudioRequest(samples_per_frame=4, chunk_frames=3)
+    state.accept_first_audio()
+    assert state.push(torch.arange(4), False) is None
+    assert state.first_audio_required and state.emitted
+    assert state.push(torch.arange(4, 8), False) is None
+    chunk = state.push(torch.arange(8, 12), True)
+    assert chunk.tolist() == list(range(4, 12))
+    assert not state.active and not state.pending
+
+
+@pytest.mark.parametrize("consumer", ["pcm_only", "finalizer", "extra"])
+def test_stream_audio_snapshot_skips_discarded_code_partition(monkeypatch, mocker, consumer):
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+
+    from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
+    monkeypatch.setattr(
+        OmniARModelRunner,
+        "_build_async_chunk_outputs_from_mm",
+        lambda *args: pytest.fail("in-stage PCM unnecessarily partitioned code outputs"),
+    )
+    batch = InputBatch.__new__(InputBatch)
+    batch.num_reqs = 3
+    batch.query_start_loc_np = np.array([0, 1, 2, 4])
+    batch.num_scheduled_tokens = np.array([1, 1, 2])
+    batch.num_tokens_after_padding = 4
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    states = [buffer.add(request_id) for request_id in ["frame", "eos", "prefill"]]
+    wav = torch.arange(16, dtype=torch.float32).reshape(4, 4)
+    streaming_audio = StreamingAudioOutput(
+        wav,
+        torch.tensor([True, False, True, False]),
+        torch.tensor(24000),
+        None,
+        batch.query_start_loc_np.copy(),
+        states,
+        np.array([False, False, True]),
+    )
+    copy_mm = mocker.spy(omni_ar_model_runner, "_async_copy_mm")
+
+    def finalize(payload, counts):
+        assert payload["codes"]["audio"].shape == (4, 2)
+        return RequestOutputSnapshot([None] * batch.num_reqs)
+
+    finalizer = mocker.Mock(side_effect=finalize) if consumer == "finalizer" else None
+    extra = ({"extra": torch.tensor([7])}, _FakeEvent()) if consumer == "extra" else None
+    result = _async_output(
+        req_ids=["frame", "eos", "prefill"],
+        sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None),
+        num_sampled_tokens=torch.tensor([1, 1, 1]),
+        multimodal_outputs={"codes": {"audio": torch.ones(4, 2)}},
+        input_batch=batch,
+        async_chunk=True,
+        streaming_audio=streaming_audio,
+        finalize_multimodal=finalizer,
+        extra_multimodal_outputs=extra,
+    ).get_output()
+    assert copy_mm.call_count == {"pcm_only": 0, "finalizer": 1, "extra": 2}[consumer]
+    if finalizer is not None:
+        finalizer.assert_called_once()
+    wav.fill_(-1)
+    assert result.inter_stage_outputs is None
+    assert result.multimodal_outputs[0]["model_outputs"].tolist() == [0, 1, 2, 3]
+    assert result.multimodal_outputs[1] == {}
+    assert result.multimodal_outputs[2]["model_outputs"].tolist() == [8, 9, 10, 11]
+    assert not states[1].active and not states[2].active
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming):
     from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot

tests/worker_v2/test_omni_model_state.py

@@ -4,13 +4,17 @@
 graph/eager paths with per-request seed independence, async snapshot ownership."""
 
 from contextlib import nullcontext
+from dataclasses import replace
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
 import torch
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+from vllm.v1.worker.gpu.states import RequestState
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
 from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _make_safe_get_rope
@@ -36,6 +40,7 @@ class _DummyInputBatch:
 def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False):
     state = object.__new__(OmniModelState)
     model = MagicMock()
+    model.stream_decoder = None
     model.has_preprocess = has_preprocess
     model.has_postprocess = has_postprocess
     model.have_multimodal_outputs = have_multimodal_outputs
@@ -62,6 +67,10 @@ def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, hav
 
     state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs)
     state._static_inputs_embeds = None
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+
+    state._eager_state = EagerMTPState(state)
+    state._stream_pos = {}
     state._mtp_generators = {}
     state._mtp_runner = None
     for name in ("_mtp_input_ids", "_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step", "_mtp_offsets"):
@@ -264,6 +273,7 @@ def test_seed_independence_resolve_once_and_sampling_kwargs():
     # vLLM sampling seed must not produce a talker generator.
     cpu = torch.device("cpu")
     assert state._get_mtp_generator("r1", SimpleNamespace(extra_args={}, seed=42), cpu) is None
+    state._mtp_generators.clear()  # The following rows are new requests.
     # Same model-local seed reproduces identical uniforms regardless of batch makeup.
     state._mtp_sample_uniforms = torch.empty((2, 2, 4))
     assert torch.equal(
@@ -620,6 +630,12 @@ def test_publish_sampled_embeddings_for_rows_whose_sample_is_kept(monkeypatch) -
     assert [tuple(row.shape) for row in sampled] == [(1, 3), (1, 3), (0,), (1, 3)]
     assert [row.tolist() for row in sampled if row.numel()] == [[[v] * 3] for v in (11.0, 12.0, 14.0)]
 
+    batch.is_prefilling_np[:] = True
+    batch.num_computed_prefill_tokens_np[:] = 0
+    batch.num_scheduled_tokens = [1] * 4
+    extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]]))
+    assert all(row.numel() == 0 for row in extra["embed"]["sampled"])
+
     # Speculative steps sample several tokens per row: not published.
     assert state.publish_sampled_embeddings(batch, torch.tensor([[11, 1], [12, 1], [13, 1], [14, 1]])) is None
 
@@ -649,3 +665,32 @@ def test_identity_preprocess_skips_only_decode_rows(prefilling):
     state.run_preprocess(batch, {"input_ids": torch.tensor([1]), "inputs_embeds": embeds})
     assert seen == (["r1"] if prefilling else [])
     assert torch.equal(embeds, torch.ones(1, 4))
+
+
+def test_mm_embeddings_exclude_zero_length_graph_padding_rows():
+    state = object.__new__(OmniModelState)
+    state.supports_mm_inputs = True
+    state.mm_pruner = None
+    state.prompt_embeds_state = None
+    state.encoder_runner = MagicMock(spec=EncoderRunner, inputs_embeds=torch.zeros(8, 2))
+    state.execute_mm_encoder = lambda _: None
+    state.gather_mm_embeddings = lambda _: ([torch.ones(1, 2)], torch.ones(5, dtype=torch.bool))
+
+    def embed(input_ids, *, query_start_loc, multimodal_embeddings, is_multimodal):
+        assert query_start_loc == [0, 2, 5]
+        return input_ids[:, None].expand(-1, 2).float()
+
+    state.model = MagicMock(supports_embed_input_ids_query_start_loc=True, embed_input_ids=embed)
+    buffers = InputBuffers(4, 8, torch.device("cpu"))
+    batch = replace(
+        InputBatch.make_dummy(2, 5, buffers),
+        num_reqs=2,
+        num_reqs_after_padding=4,
+        num_tokens=5,
+        num_tokens_after_padding=8,
+        input_ids=torch.arange(8),
+        query_start_loc_np=np.array([0, 2, 5, 5, 5], dtype=np.int32),
+    )
+    result = state.prepare_inputs_embeds({}, batch, MagicMock(spec=RequestState))
+    torch.testing.assert_close(result[:5], torch.arange(5).float()[:, None].expand(-1, 2))
+    assert result.shape == (8, 2)

vllm_omni/config/environment_variable_inventory.py

@@ -148,11 +148,13 @@ _MODEL_PROMOTE = (
     "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH",
     "COSMOS3_SOUND_TOKENIZER_PATH",
     "COSYVOICE3_BATCH_FLOW",
+    "COSYVOICE3_CACHED_ISTFT",
     "COSYVOICE3_ESTIMATOR_ONNX",
     "COSYVOICE3_FLOW_CUDAGRAPH",
     "COSYVOICE3_FLOW_GRAPH",
     "COSYVOICE3_FLOW_GRAPH_MAX_TOTAL",
     "COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS",
+    "COSYVOICE3_HIFT_GRAPH",
     "COSYVOICE3_PACKED_STREAMING",
     "COSYVOICE3_REFERENCE_PREFETCH",
     "COSYVOICE3_REFERENCE_WORKERS",
@@ -260,6 +262,7 @@ _BENCHMARK_TRANSITIONAL = (
 )
 
 _INTERNAL = (
+    "VLLM_OMNI_TALKER_FIRST_AUDIO",
     "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS",
     "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S",
     "VLLM_OMNI_CONNECTOR_RECV_POLL_MS",

vllm_omni/config/model.py

@@ -5,6 +5,7 @@ from dataclasses import MISSING, field
 from typing import Any
 
 from pydantic import ConfigDict, TypeAdapter
+from transformers import PretrainedConfig
 from vllm.config import ModelConfig
 from vllm.config.utils import config
 from vllm.logger import init_logger
@@ -124,9 +125,12 @@ class OmniModelConfig(ModelConfig):
     """
 
     stage_id: int = 0
+    # Match the upstream derived field while its value is replaced per stage.
+    hf_text_config: PretrainedConfig = field(init=False)
     async_chunk: bool = False
     session_mode: str = "turn"
     retains_state_across_chunks: bool = False
+    supports_running_prefix_cache_reset: bool = True
     use_v2_model_runner: bool = False
     supports_native_mrv2_data_plane: bool = False
     # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin.
@@ -218,7 +222,7 @@ class OmniModelConfig(ModelConfig):
             return convertor.convert()
         return super().get_model_arch_config()
 
-    def draw_hf_text_config(self):
+    def draw_hf_text_config(self) -> PretrainedConfig:
         # transformers' get_text_config method is used to get the text config from thinker_config.
         # to handle the case that each model stage has their own text config,
         # we need to draw the text config from the corresponding model stage.

vllm_omni/config/omni_config.py

@@ -482,6 +482,7 @@ class OmniStageModelConfig(_TrackExplicitConfigFields):
     interleave_mm_strings: bool | None = None
     media_io_kwargs: dict[str, Any] | None = None
     final_output: bool = False
+    supports_running_prefix_cache_reset: bool = True
     active_stream_window: int = Field(default=0, ge=0)
     session_mode: str = "turn"
     duplex_max_sessions: int = Field(default=1, ge=1)
@@ -817,6 +818,7 @@ class _DiffusionConfigProjection:
     prompt_embed_cache_size: int = Field(default=32, ge=1)
     enable_session_state_manager: bool = False
     diffusion_load_format: str = "default"
+    hsdp_weight_load_strategy: str = "full"
     diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_pipeline_cls: Any = None
@@ -1324,6 +1326,7 @@ _DIFFUSION_STAGE_METADATA_FIELDS = frozenset(
         "model_arch",
         "model_stage",
         "retains_state_across_chunks",
+        "supports_running_prefix_cache_reset",
         "scheduler_cls",
         "stage_connector_spec",
         "worker_type",
@@ -1971,6 +1974,8 @@ def _build_model_config(
     if "active_stream_window" not in kwargs:
         kwargs["active_stream_window"] = _copy_value(deploy.active_stream_window)
     kwargs["final_output"] = topology.final_output
+    if not topology.supports_running_prefix_cache_reset:
+        kwargs["supports_running_prefix_cache_reset"] = False
     if "custom_voice_dir" not in kwargs and deploy.custom_voice_dir is not None:
         kwargs["custom_voice_dir"] = _copy_value(deploy.custom_voice_dir)
     stage_runner = resolve_stage_model_runner(deploy, stage_deploy)

vllm_omni/config/pipeline_registry.py

@@ -54,7 +54,7 @@ from vllm_omni.model_executor.models.cosmos3.pipeline import (
     COSMOS3_OMNI_DEPLOY_PIPELINE,
     COSMOS3_POLICY_PIPELINE,
 )
-from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
+from vllm_omni.model_executor.models.cosyvoice3.pipeline import resolve_cosyvoice3_pipeline
 from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
 from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
@@ -117,7 +117,7 @@ from vllm_omni.model_executor.models.qwen3_omni.pipeline import (
     QWEN3_OMNI_THINKER_ONLY_PIPELINE,
     resolve_qwen3_omni_pipeline,
 )
-from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
+from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE, QWEN3_TTS_PIPELINE
 from vllm_omni.model_executor.models.step_audio2.pipeline import (
     STEP_AUDIO2_ASR_PIPELINE,
     STEP_AUDIO2_PIPELINE,
@@ -150,7 +150,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     # key (see deploy/cosmos3_policy_droid.yaml and deploy/cosmos3_omni.yaml).
     "cosmos3_omni_deploy": COSMOS3_OMNI_DEPLOY_PIPELINE,
     "cosmos3_policy": COSMOS3_POLICY_PIPELINE,
-    "cosyvoice3": COSYVOICE3_PIPELINE,
+    "cosyvoice3": resolve_cosyvoice3_pipeline,
     "covo_audio": COVO_AUDIO_PIPELINE,
     "dreamzero": DREAMZERO_PIPELINE,
     "fish_qwen3_omni": FISH_SPEECH_PIPELINE,
@@ -204,6 +204,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "qwen3_omni_moe": resolve_qwen3_omni_pipeline,
     "qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE,
     "qwen3_tts": QWEN3_TTS_PIPELINE,
+    "qwen3_tts_fused": QWEN3_TTS_FUSED_PIPELINE,
     "step_audio_2": STEP_AUDIO2_PIPELINE,
     "step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE,
     "voxcpm2": VOXCPM2_PIPELINE,

vllm_omni/config/stage_config.py

@@ -245,6 +245,8 @@ class StagePipelineConfig:
     # The model keeps per-request execution state while awaiting the next
     # async chunk, so the parked request continues to consume model capacity.
     retains_state_across_chunks: bool = False
+    # Some stateful audio stages cannot roll back already consumed frames.
+    supports_running_prefix_cache_reset: bool = True
     sampling_constraints: dict[str, Any] = field(default_factory=dict)
     custom_process_input_func: str | None = None
     custom_process_next_stage_input_func: str | None = None
@@ -295,6 +297,9 @@ class PipelineConfig:
     model_type: str
     model_arch: str = ""
     stages: tuple[StagePipelineConfig, ...] = ()
+    # A single stage that streams its own final output in async-chunk mode
+    # (e.g. a Talker decoding audio in-stage) keeps deploy.async_chunk.
+    single_stage_async_chunk: bool = False
     # HF architecture aliases: used by StageConfigFactory when the model's
     # HF config reports a generic model_type that collides with a different
     # model (e.g. MiMo Audio reports model_type="qwen2"). The factory
@@ -492,6 +497,7 @@ class StageDeployConfig:
     auxiliary_text_encoder: str | None = None
 
     # Runtime optimizations used by diffusion loading/execution.
+    hsdp_weight_load_strategy: str | None = None
     enable_multithread_weight_load: bool | None = None
     enable_broadcast_weight_load: bool | None = None
     num_weight_load_threads: int | None = None
@@ -1106,6 +1112,8 @@ def _build_engine_args(
     if ps.omni_kv_config:
         engine_args["omni_kv_config"] = dict(ps.omni_kv_config)
     engine_args["requires_full_payload_input"] = ps.requires_full_payload_input
+    if not ps.supports_running_prefix_cache_reset:
+        engine_args["supports_running_prefix_cache_reset"] = False
     return engine_args
 
 
@@ -1171,8 +1179,11 @@ def _resolve_pipeline_async_chunk_enabled(
     and False otherwise. If the user tried to enable async chunk through the deploy
     config, but it's inapplicable or unsupported, it will be disabled with a warning.
     """
-    # Single stage should never use async chunk
     if len(pipeline.stages) <= 1:
+        if pipeline.single_stage_async_chunk:
+            if deploy.async_chunk is False:
+                raise ValueError(f"Pipeline {pipeline.model_type!r} requires async_chunk=True")
+            return True
         if deploy.async_chunk:
             logger.warning(
                 "Deploy config set async_chunk=True, but async chunk is inapplicable "
@@ -1349,6 +1360,8 @@ class StageConfig:
 
         # Terminal-stage ownership comes from topology, not engine overrides.
         engine_args["final_output"] = self.final_output
+        if self.yaml_engine_args.get("supports_running_prefix_cache_reset") is False:
+            engine_args["supports_running_prefix_cache_reset"] = False
 
         # Build runtime config from YAML defaults + CLI overrides
         runtime: dict[str, Any] = dict(self.yaml_runtime)

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -101,6 +101,23 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
     max_num_running_reqs: int
 
+    def reset_prefix_cache(self, reset_running_requests: bool = False, reset_connector: bool = False) -> bool:
+        model_config = self.vllm_config.model_config
+        if (
+            reset_running_requests
+            and self.running
+            and not getattr(model_config, "supports_running_prefix_cache_reset", True)
+        ):
+            # Reset preempts and resumes in the same step while discarding
+            # in-flight tokens. The stateful codec and queued PCM have already
+            # consumed those frames and cannot roll back to that boundary.
+            logger.warning(
+                "This stage cannot reset running requests; wait for "
+                "completion or abort them before resetting the prefix cache."
+            )
+            return False
+        return super().reset_prefix_cache(reset_running_requests, reset_connector)
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         # Track requests that need KV cache transfer when finished
@@ -127,8 +144,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Track requests that have already triggered prefill transfer to avoid duplicates
         self.transfer_triggered_requests: set[str] = set()
 
-        # Cache per-request flag to avoid repeated deserialization of additional_information
-        self._omits_kv_transfer_cache: dict[str, bool] = {}
+        # Cache per-request flags to avoid repeated deserialization of
+        # additional_information. Value is (payload_id, is_stage_zero_final,
+        # force_kv_transfer); payload_id invalidates the entry when the
+        # object is replaced.
+        self._omits_kv_transfer_cache: dict[str, tuple[int, bool, bool]] = {}
 
         # KV-wait start ts for the vllm_omni:kv_wait_s metric; see
         # _emit_kv_wait_output for the engine-core → orchestrator carry.
@@ -247,26 +267,40 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             return config.get(key, default)
         return getattr(config, key, default) if config is not None else default
 
-    def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
-        """True when this stage-zero-final request does not need downstream KV.
-
-        The result is cached per request to avoid repeated deserialization of
-        additional_information on every scheduler tick.
-        """
-        rid = request.request_id
-        cached = self._omits_kv_transfer_cache.get(rid)
-        if cached is not None:
-            return cached
-
+    def _omni_final_stage_flags(self, request: Request) -> tuple[bool, bool]:
+        """Return ``(is_stage_zero_final, force_kv_transfer)`` from request metadata."""
         payload = getattr(request, "additional_information", None)
         if payload is None:
-            result = False
-        else:
-            info = deserialize_additional_information(payload)
-            result = info.get("omni_final_stage_id") == 0 and not bool(info.get("omni_force_kv_transfer", False))
+            return False, False
 
-        self._omits_kv_transfer_cache[rid] = result
-        return result
+        cache = getattr(self, "_omits_kv_transfer_cache", None)
+        rid = request.request_id
+        payload_id = id(payload)
+        if cache is not None:
+            cached = cache.get(rid)
+            if cached is not None and cached[0] == payload_id:
+                return cached[1], cached[2]
+
+        info = deserialize_additional_information(payload)
+        is_final = info.get("omni_final_stage_id") == 0
+        force_kv = bool(info.get("omni_force_kv_transfer", False))
+        if cache is not None:
+            cache[rid] = (payload_id, is_final, force_kv)
+        return is_final, force_kv
+
+    def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
+        """True when this stage-zero-final request does not need downstream KV."""
+        is_final, force_kv = self._omni_final_stage_flags(request)
+        return is_final and not force_kv
+
+    def _request_omits_chunk_transfer_to_next_stage(self, request: Request) -> bool:
+        """True when this request has no downstream chunk consumer.
+
+        CFG companions still force KV transfer but are stage-0-final for
+        ordinary inter-stage chunks, so they omit ``save_async`` here.
+        """
+        is_final, _ = self._omni_final_stage_flags(request)
+        return is_final
 
     def _should_defer_waiting_admission(self) -> bool:
         return False
@@ -517,7 +551,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # domains must consume the old frame before new output passes.
                 request.async_tokens_to_discard = max(0, stale_async_tokens - len(generated_token_ids))
 
-            if output_is_stale or async_output_is_stale:
+            if async_output_is_stale or (output_is_stale and request.drop_stale_output):
                 # Output of a step scheduled before the request's in-flight
                 # tokens were discarded (segment stop / session replacement).
                 # num_computed_tokens was rolled back at the discard site, so
@@ -556,12 +590,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # tokens and rejections. If some tokens are rejected,
                 # num_computed_tokens is decreased by the number of rejected
                 # tokens.
-                if request.num_computed_tokens > 0:
-                    request.num_computed_tokens -= num_rejected
-                # If async scheduling, num_output_placeholders also includes
-                # the scheduled spec tokens count and so is similarly adjusted.
-                if request.num_output_placeholders > 0:
-                    request.num_output_placeholders -= num_rejected
+                if not output_is_stale:
+                    if request.num_computed_tokens > 0:
+                        request.num_computed_tokens -= num_rejected
+                    # If async scheduling, num_output_placeholders also includes
+                    # the scheduled spec tokens count and so is similarly adjusted.
+                    if request.num_output_placeholders > 0:
+                        request.num_output_placeholders -= num_rejected
                 spec_decoding_stats = self.make_spec_decoding_stats(
                     spec_decoding_stats,
                     num_draft_tokens=num_draft_tokens,
@@ -600,7 +635,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Check for stop and update request status.
             if new_token_ids:
                 num_sampled_tokens = len(new_token_ids)
-                new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
+                if output_is_stale:
+                    new_token_ids, stopped = self._update_request_with_output(request, new_token_ids, is_stale=True)
+                else:
+                    new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
                 if new_logprobs is not None and len(new_token_ids) < num_sampled_tokens:
                     # A mid-step stop (e.g. spec-decode tokens sampled past
                     # EOS) trims new_token_ids after the validation slice
@@ -648,6 +686,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             confirmed_num_computed_tokens = None
             boundary_generation = None
+            # Only when this step might save. Read additional_information
+            # before _free_request rewrites it.
+            omits_chunk_transfer = False
+            if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or stopped):
+                omits_chunk_transfer = self._request_omits_chunk_transfer_to_next_stage(request)
             # Capture before resumable stop handling can clear token history.
             output_token_ids: Any = getattr(request, "output_token_ids", None)
             if output_token_ids is None:
@@ -704,6 +747,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     # a queued streaming update, and adding twice swallows the
                     # next duplex unit's listen/speak under async scheduling.
                     if request.num_in_flight_tokens > 0:
+                        request.drop_stale_output = True
                         request.num_stale_output_tokens = request.num_in_flight_tokens
                     if outstanding_async_tokens > 0:
                         # Discard only outputs that are already in flight and
@@ -765,8 +809,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
 
-            if self.chunk_transfer_adapter is not None and (
-                inter_stage_output is not None or is_segment_finished or finished
+            if omits_chunk_transfer and inter_stage_output is not None:
+                logger.warning(
+                    "Skipping inter-stage chunk for request %s: "
+                    "omni_final_stage_id=0 but inter_stage_output is present",
+                    req_id,
+                )
+            if (
+                self.chunk_transfer_adapter is not None
+                and not omits_chunk_transfer
+                and (inter_stage_output is not None or is_segment_finished or finished)
             ):
                 save_kwargs = {
                     "new_token_ids": new_token_ids,
@@ -884,6 +936,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # a resumable stop applies a queued update through this helper.
         in_flight_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0)
         if in_flight_tokens > 0:
+            session.drop_stale_output = True
             session.num_stale_output_tokens = in_flight_tokens
         if outstanding_async_tokens > 0:
             # Async scheduling may already have sampled the previous
@@ -1405,6 +1458,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if status is not None:
                 request.status = status
         if native_transfer and connector_delay_free_blocks:
+            assert transfer_params is not None
             kv_xfer_params = {
                 **(kv_xfer_params or {}),
                 "transfer_id": transfer_params["transfer_id"],

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -222,12 +222,16 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         # num_in_flight_tokens already includes any undrained stale share.
         # Assign instead of accumulating so callers that fenced the same
         # rollover before entering this helper do not count it twice.
+        session.drop_stale_output = True
         session.num_stale_output_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0)
         session.num_output_placeholders = 0
         session.spec_token_ids = []
 
     def _finish_streaming_session_update(self, session: Request, update: StreamingUpdate) -> None:
         """Install payload metadata and send an updated session to admission."""
+        cache = getattr(self, "_omits_kv_transfer_cache", None)
+        if cache is not None:
+            cache.pop(session.request_id, None)
         session.additional_information = update.additional_information or None
         session.model_intermediate_buffer = getattr(
             update,
@@ -316,6 +320,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             # The streaming update may already have fenced this same in-flight
             # frame. Seed idempotently so the replacement does not count it twice.
             request.num_stale_output_tokens = int(getattr(request, "num_in_flight_tokens", 0) or 0)
+            request.drop_stale_output = True
             request.num_output_placeholders = 0
             request.spec_token_ids = []
             self._release_replaced_streaming_prompt_cache(request)

vllm_omni/deploy/auk.yaml

@@ -20,14 +20,20 @@ enable_chunked_prefill: false
 
 stages:
   # Stage 0: Qwen2.5-Omni thinker as a frozen text encoder.
-  # Prefill only — one sampled token, discarded downstream. eager because the
-  # fused-hidden forward walks the decoder layers itself and so bypasses the
-  # compiled graph.
+  # Prefill only: one sampled token, discarded downstream. The fused-hidden
+  # forward walks the decoder layers itself, so torch.compile has nothing to
+  # hook (mode 0); the whole prefill is instead captured as FULL CUDA graphs
+  # padded to the sizes below, which removes most of its launch overhead
+  # (about 20 ms per request on H200). Longer prompts run eager.
   - stage_id: 0
     devices: "0"
     max_num_seqs: 1
     gpu_memory_utilization: 0.45
-    enforce_eager: true
+    enforce_eager: false
+    compilation_config:
+      mode: 0
+      cudagraph_mode: FULL
+      cudagraph_capture_sizes: [64, 128, 192, 256, 384, 512]
     max_num_batched_tokens: 32768
     # AuK takes text plus at most one clip; no image or video tower is built.
     limit_mm_per_prompt: {audio: 1}

vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml

@@ -0,0 +1,28 @@
+# Single-Hopper streaming throughput profile: AR32 / codec32, Flow10.
+# Automatically selected on Hopper devices with at least 140 GiB (H200)
+# when nvidia-cuda-mps-control is available on PATH.
+# Standard sampling is configured below; request stream=true for first audio.
+# F0 and HiFT stay FP32; only deterministic HiFT decode is graphed.
+# HiFT graphs capture on the third use of each exact shape, adding latency then.
+# The codec graph pool is not capped by gpu_memory_utilization; leave memory headroom.
+base_config: cosyvoice3_packed_streaming.yaml
+stages:
+  - stage_id: 0
+    max_num_seqs: 32
+    hf_overrides:
+      cosyvoice3_sampling_mode: standard
+    default_sampling_params:
+      max_tokens: 2048
+      repetition_penalty: 1.21
+      temperature: 0.7
+      top_k: 20
+      top_p: 0.8
+  - stage_id: 1
+    max_num_seqs: 32
+    env:
+      COSYVOICE3_TRT: "0"
+      COSYVOICE3_BATCH_FLOW: "1"
+      COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS: "0"
+      COSYVOICE3_PACKED_STREAMING: "1"
+      COSYVOICE3_CACHED_ISTFT: "1"
+      COSYVOICE3_HIFT_GRAPH: "1"

vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml

@@ -0,0 +1,50 @@
+# Opt-in CUDA Model Runner V2 profile for Qwen3-Omni-MoE: the Thinker (stage 0),
+# Talker (stage 1) and Code2Wav (stage 2) run on MRv2 with its native data
+# plane. The base profile keeps every stage on V1.
+#
+# On MRv2 the Thinker hands each generated token's embedding to the Talker with
+# the step that sampled it (one Thinker step earlier than its next forward would
+# capture the same embedding) and replays decode in FULL CUDA graphs together
+# with the layers it captures for the Talker.
+#
+# When the stages share one GPU, run them under NVIDIA MPS (see
+# docs/configuration/stage_configs.md); startup warns when no MPS server is found.
+#
+# MRv2 downstream stages support turn-based sessions only: realtime (resumable)
+# requests and duplex sessions need the V1 chunk adapter.
+base_config: qwen3_omni_moe.yaml
+
+connectors:
+  connector_of_shared_memory:
+    extra:
+      # First Code2Wav chunk is the stream's first frame alone, then chunks
+      # double up to the steady codec_chunk_frames (replaces the 4-frame
+      # initial chunk). Code2Wav re-decodes up to codec_left_context_frames
+      # before each chunk, so windows are 1, 3, 7, 15, 31, then 50 frames.
+      # A one-frame first chunk is also what lets the Talker deliver it
+      # itself. Opt in explicitly; VLLM_OMNI_TALKER_FIRST_AUDIO=0 disables it.
+      talker_first_audio: true
+      codec_chunk_ramp: [1, 2, 4, 8, 16, 25]
+      # Code2Wav decodes a streaming step as a few length groups on multi-row
+      # CUDA graphs instead of one eager call padded to its longest window
+      # (with the ramp, over half of those padded frames were waste).
+      codec_graph_batch_sizes: [2, 3, 4, 6, 8, 12, 16]
+      codec_fused_snake: true
+
+stages:
+  - stage_id: 0
+    model_runner: v2
+  - stage_id: 1
+    model_runner: v2
+  - stage_id: 2
+    model_runner: v2
+
+platforms:
+  npu:
+    model_runner: v1
+  xpu:
+    model_runner: v1
+  rocm:
+    model_runner: v1
+  musa:
+    model_runner: v1

vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml

@@ -0,0 +1,27 @@
+# Opt-in single-H200 MRv2 profile, measured at client concurrency 32/64.
+# Run all stages under NVIDIA MPS and warm representative shapes first.
+# Retains Triton MoE and native nucleus reductions, fusing only the final
+# sampling mask/Gumbel selection. cuDNN searches convolution algorithms
+# during warmup; short-KV attention and convolution selection change rounding.
+# This is not a claim of bitwise audio or universal P99/throughput improvement.
+# Evidence: docs/design/qwen3_omni_mrv2_performance.md
+base_config: qwen3_omni_moe_mrv2.yaml
+
+connectors:
+  connector_of_shared_memory:
+    extra:
+      code_predictor_kv_cache: true
+      code_predictor_short_kv: true
+      code_predictor_fused_sampling: true
+      codec_cudnn_benchmark: true
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    gpu_memory_utilization: 0.60
+  - stage_id: 1
+    devices: "0"
+    gpu_memory_utilization: 0.22
+  - stage_id: 2
+    devices: "0"
+    gpu_memory_utilization: 0.10

vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml

@@ -0,0 +1,54 @@
+# Single-stage Qwen3-TTS: the Talker decodes every frame with the stateful
+# streaming codec decoder and returns PCM itself (no Code2Wav stage).
+# reset_prefix_cache(reset_running_requests=True) returns False while requests
+# are running; wait for completion or abort them before resetting the cache.
+pipeline: qwen3_tts_fused
+model_runner: v2
+async_chunk: true
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 128
+    # Exact conditioned Talker inputs for preemption replay are allocated
+    # after KV profiling: 16 MiB per active request at 1.7B/max_model_len=4096
+    # (~2 GiB at 128 requests), in addition to codec state and output buffers.
+    # Keep headroom when increasing gpu_memory_utilization or capacity.
+    gpu_memory_utilization: 0.5
+    trust_remote_code: true
+    enable_prefix_caching: false
+    # Model options that multi-stage deployments put on the connector: a
+    # single stage has no connector edge.
+    additional_config:
+      talker_stream_decode: true
+      # Opt in to first-frame delivery on its own queue for lower first-audio
+      # latency; leave disabled when prioritizing peak throughput.
+      talker_stream_first_audio: false
+      # Match the standard two-stage profile's voice-clone context boundary.
+      # A shorter context (e.g. 25) can change timbre continuity.
+      ref_code_context_frames: 72
+      code_predictor_kv_cache: true
+      code_predictor_fused_sampling: true
+      code_predictor_fused: true
+    # Talker decode contexts are short (~100-700 tokens): FA3 split-KV plus
+    # its combine kernel costs more than one pass per head.
+    attention_config:
+      flash_attn_max_num_splits_for_cuda_graph: 1
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 4096
+    compilation_config:
+      cudagraph_mode: FULL
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 144, 160, 176, 192, 208, 224, 240, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512]
+      max_cudagraph_capture_size: 512
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      min_tokens: 2
+      repetition_penalty: 1.05
+    subtalker_sampling_params:
+      do_sample: true
+      temperature: 0.9
+      top_k: 50
+      top_p: 1.0