[CI] Retrigger CosyVoice3 ready validation

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-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.

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,

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_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_stale_drain.py

@@ -106,7 +106,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
     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 +114,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 +134,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 +246,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/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/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/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/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_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/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/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

@@ -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,
@@ -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
@@ -517,7 +534,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 +573,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 +618,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
@@ -704,6 +725,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
@@ -884,6 +906,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 +1428,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,6 +222,7 @@ 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 = []
@@ -316,6 +317,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/cosyvoice3_packed_streaming_optimized_standard.yaml

@@ -0,0 +1,24 @@
+# Opt-in single-Hopper streaming throughput profile: AR32 / codec32, Flow10.
+# Use --hf-overrides '{"cosyvoice3_sampling_mode": "standard"}' and stream=true.
+# 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
+    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

vllm_omni/deploy/yue2.yaml

@@ -1,5 +1,5 @@
 # YuE2-3B deploy: single-stage native-AR text-to-music (~3.6B bf16 + VAE).
-# Target 1x 24 GB card; verified on a single RTX 4090.
+# Four-slot CUDA configuration; validate memory on the target card/workload.
 #
 # The architectures pin is required: the checkpoint's config.json names
 # "YuE2ForCausalLM" with an auto_map that only makes sense under
@@ -13,16 +13,17 @@ dtype: bfloat16
 
 stages:
   - stage_id: 0
+    # A finishing song keeps its slot during side-stream synthesis,
+    # emitting HOLD tokens until its audio is ready.
     max_num_seqs: 4
-    # Measured on a 24 GB RTX 4090 with 4 concurrent requests pinned to the
-    # 9000-frame cap (worst case): the engine reserves ~15.8 GiB at 0.70 and
-    # peak usage during the terminal NAR/VAE finishing pass reaches 20.9 GiB,
-    # leaving ~3 GiB of headroom. 0.85 reserves ~19.4 GiB and the same
-    # workload OOMs inside the finishing pass.
+    # Reserve synthesis headroom: NAR K/V, acoustic graph buffers and VAE
+    # activations are allocated beyond the engine budget. The earlier 4090
+    # memory measurement predates async/compiled synthesis; 0.70 is a starting
+    # point, not a guarantee for 9000 frames or long ABC prefixes. Lower it if
+    # the target workload needs more synthesis memory (see the model recipe).
     gpu_memory_utilization: 0.70
-    # Phase 1 keeps the backbone eager: the model-owned sampler runs after the
-    # forward, and the terminal NAR/VAE pass runs inside the finishing step.
-    # CUDA-graph capture of the standard Qwen3 backbone is a perf follow-up.
+    # Keep AR capture off by default. Sampler/NAR graphs still need
+    # separate synthesis headroom.
     enforce_eager: true
     # Prefix caching stays off: the omni prefix-cache save snapshots the
     # step's mm payload before the model-owned sampler runs, but YuE2
@@ -42,14 +43,17 @@ stages:
     hf_overrides:
       architectures: ["Yue2ForCausalLM"]
     default_sampling_params:
-      # Placeholder values: the model's own sampler overrides everything from
-      # extra_args. Their presence turns on has_sampling_extra_args, which is
-      # what threads the per-request yue2_* keys into the model's forward.
+      # Defaults for callers that omit model-owned sampling options. The
+      # opt-in prepare_runner_inputs hook reads each request's extra_args;
+      # it does not depend on the forward's has_sampling_extra_args flag.
       extra_args:
         yue2_phase: "semantic"
         yue2_seed: 831001
         yue2_max_audio_frames: 200
+        yue2_max_hold_steps: 4096
       detokenize: false
-      # Engine-side backstop; the model stops on its own end token first.
-      max_tokens: 9000
+      stop_token_ids: [151848, 151852]  # ABC_END, MUSIC_END
+      # 200 frames + 4096 synthesis holds + the final end token. Callers
+      # changing the frame budget must also reserve this token headroom.
+      max_tokens: 4297
       temperature: 1.0

vllm_omni/diffusion/attention/backends/abstract.py

@@ -7,6 +7,10 @@ from typing import Any, Generic, Literal, TypeVar
 
 import torch
 
+from vllm_omni.diffusion.attention.capabilities import (
+    ExecutionContext,
+    ExecutionPathResult,
+)
 from vllm_omni.platforms import current_omni_platform
 
 
@@ -132,6 +136,15 @@ class AttentionBackend(ABC):
 
         return False
 
+    @classmethod
+    def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult:
+        """Return conservative pre-construction capabilities.
+
+        Unmigrated backends remain advisory so introducing the contract cannot
+        reject an execution path that works today.
+        """
+        return ExecutionPathResult.unmigrated(cls.get_name(), context)
+
 
 @dataclass(frozen=True, slots=True)
 class QueryRange:
@@ -218,6 +231,9 @@ class AttentionMetadata:
     #   "valid_kv_length": int — contiguous valid K/V prefix length. Local
     #     backends may slice it directly; ring backends trim each circulated
     #     K/V block from the same global prefix.
+    #   "attention_mask_mode": "none" | "padding" | "arbitrary" — producer
+    #     assertion used by capability planning. "none" means a present mask
+    #     is semantically a no-op; absent means its contents are runtime-dependent.
     #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
     #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
     #     Requires the [real, pad] two-document packing contract; see
@@ -277,6 +293,21 @@ class AttentionImpl(ABC, Generic[T]):
             return True
         return kv_cache_dtype in cls._supported_kv_cache_dtypes.get(platform_key, set())
 
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None,
+    ) -> ExecutionPathResult:
+        """Resolve the complete path after backend initialization.
+
+        Subclasses report verified results; this inspection API does not enforce them.
+        """
+        del query, key, value, attn_metadata
+        return ExecutionPathResult.unmigrated(type(self).__name__, context)
+
     def forward(
         self,
         query: torch.Tensor,

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -3,7 +3,9 @@
 
 import os
 from collections.abc import Callable
+from dataclasses import replace
 from functools import cache, partial
+from typing import NamedTuple
 
 import torch
 from vllm.logger import init_logger
@@ -14,6 +16,16 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
     run_paged_piecewise_plan,
 )
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    ExecutionPathResult,
+    MaskMode,
+    PackingMode,
+    ParallelStrategy,
+    SupportStatus,
+)
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.platforms import current_omni_platform
 
@@ -29,6 +41,133 @@ def _get_npu_compressed_causal_mask(device: torch.device) -> torch.Tensor:
     ).contiguous()
 
 
+if not hasattr(torch.ops.vllm_omni, "fa4_dense_attention"):
+
+    @torch.library.custom_op("vllm_omni::fa4_dense_attention", mutates_args=())
+    def _fa4_dense_attention_op(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        softmax_scale: float,
+        causal: bool,
+        deterministic: bool,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.utils.fa import flash_attn_func
+
+        if flash_attn_func is None:
+            raise RuntimeError("CuTe FlashAttention-4 is unavailable")
+        kwargs = {
+            "causal": causal,
+            "softmax_scale": softmax_scale,
+        }
+        if deterministic:
+            kwargs["deterministic"] = True
+        out = flash_attn_func(query, key, value, **kwargs)
+        return out[0] if isinstance(out, tuple) else out
+
+    @_fa4_dense_attention_op.register_fake
+    def _fa4_dense_attention_fake(
+        query,
+        key,
+        value,
+        softmax_scale,
+        causal,
+        deterministic,
+    ):
+        # FA4 returns contiguous output even when Q is noncontiguous; V sets its head dimension.
+        return query.new_empty((*query.shape[:-1], value.shape[-1]))
+
+
+_fa4_dense_attention_op = torch.ops.vllm_omni.fa4_dense_attention
+
+
+_PACKED_KEYS = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
+
+
+class _FlashAttentionMetadataPlan(NamedTuple):
+    attention_mask: torch.Tensor | None
+    mask_mode: MaskMode
+    full_attn_spans: list[list[tuple[int, int]]] | None
+    extra: dict
+    packing_mode: PackingMode
+
+
+def _normalize_flash_attention_metadata(
+    attn_metadata: AttentionMetadata | None,
+) -> _FlashAttentionMetadataPlan:
+    attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+    full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
+    extra = attn_metadata.extra if attn_metadata is not None else {}
+    published_mask_mode = extra.get("attention_mask_mode")
+    if attention_mask is None:
+        mask_mode = MaskMode.NONE
+    elif published_mask_mode is None:
+        mask_mode = MaskMode.UNKNOWN
+    else:
+        try:
+            mask_mode = MaskMode(published_mask_mode)
+        except ValueError as error:
+            raise ValueError(f"Unknown attention_mask_mode {published_mask_mode!r}") from error
+        if mask_mode is MaskMode.UNKNOWN:
+            raise ValueError("attention_mask_mode='unknown' is reserved for unpublished semantics")
+    # Piecewise dispatch takes precedence and does not consume packed metadata.
+    present_packed_keys = [key for key in _PACKED_KEYS if key in extra] if full_attn_spans is None else []
+    if present_packed_keys and len(present_packed_keys) != len(_PACKED_KEYS):
+        missing = sorted(set(_PACKED_KEYS) - set(present_packed_keys))
+        raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
+
+    packing_mode = PackingMode.NONE
+    if present_packed_keys:
+        cu_seqlens_q = extra["cu_seqlens_q"]
+        packing_mode = PackingMode.MULTI_DOCUMENT if cu_seqlens_q.shape[0] > 3 else PackingMode.PACKED_PADDING
+    return _FlashAttentionMetadataPlan(
+        attention_mask=attention_mask,
+        mask_mode=mask_mode,
+        full_attn_spans=full_attn_spans,
+        extra=extra,
+        packing_mode=packing_mode,
+    )
+
+
+def _flash_attention_execution_path(
+    context: ExecutionContext,
+) -> ExecutionPathResult:
+    if context.mask_mode is MaskMode.UNKNOWN:
+        return ExecutionPathResult.unmigrated(
+            "FLASH_ATTN",
+            context,
+            path="runtime_mask_dependent",
+        )
+    if not (
+        context.platform == "cuda"
+        and context.kernel_variant == "fa4"
+        and context.dtype in {"bfloat16", "torch.bfloat16"}
+        and context.causal is False
+        and context.mask_mode is MaskMode.NONE
+        and context.packing_mode is PackingMode.NONE
+        and not context.piecewise
+        and not context.paged_kv
+        and context.kv_cache_dtype is None
+        and context.parallel_strategy is ParallelStrategy.NONE
+        and not context.outer_boundaries
+    ):
+        return ExecutionPathResult.unmigrated(
+            "FLASH_ATTN",
+            context,
+            path="unverified",
+        )
+
+    return ExecutionPathResult(
+        backend="FLASH_ATTN",
+        path="fa4_dense",
+        support=CapabilityResult.supported(),
+        compilation_mode=CompilationMode.CUSTOM_OP,
+        platform=context.platform,
+        kernel_variant=context.kernel_variant,
+        parallel_strategy=context.parallel_strategy,
+    )
+
+
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supports_piecewise_spans: bool = True
@@ -74,6 +213,10 @@ class FlashAttentionBackend(AttentionBackend):
     def get_impl_cls() -> type["FlashAttentionImpl"]:
         return FlashAttentionImpl
 
+    @classmethod
+    def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult:
+        return _flash_attention_execution_path(replace(context, kernel_variant=None))
+
 
 class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
     # Per-platform attention quantization support.
@@ -108,11 +251,135 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         self.softmax_scale = softmax_scale
         self.qkv_layout = qkv_layout
         self.is_cross_attn = role == "cross"
+        from vllm_omni.diffusion.attention.backends.utils.fa import IS_AITER, IS_FLASH_ATTN_4
+
+        self._kernel_variant = "fa4" if IS_FLASH_ATTN_4 else "aiter" if IS_AITER else None
         cfg = get_current_diffusion_config_or_none()
         self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False
         if backend_kwargs:
             logger.warning("FlashAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        if context.platform == "npu":
+            return self._resolve_npu_execution_path(context, query, key, attn_metadata)
+        metadata = _normalize_flash_attention_metadata(attn_metadata)
+        context = replace(
+            context,
+            kernel_variant=self._kernel_variant,
+            dtype=str(query.dtype).removeprefix("torch."),
+            causal=self.causal,
+            mask_mode=metadata.mask_mode,
+            packing_mode=metadata.packing_mode,
+            piecewise=metadata.full_attn_spans is not None,
+            kv_cache_dtype=metadata.extra.get("kv_cache_dtype"),
+        )
+        if context.platform == "rocm":
+            return self._resolve_rocm_execution_path(context)
+        result = _flash_attention_execution_path(context)
+        if result.support.status is not SupportStatus.SUPPORTED:
+            return result
+
+        from vllm_omni.diffusion.attention.backends.utils.fa import validate_fa4_head_dims
+
+        try:
+            if query.dtype != key.dtype or query.dtype != value.dtype:
+                raise ValueError("Q, K, and V dtypes must match")
+            if query.device != key.device or query.device != value.device:
+                raise ValueError("Q, K, and V devices must match")
+            if query.shape[-1] != key.shape[-1]:
+                raise ValueError("Q and K head dimensions must match")
+            verified = validate_fa4_head_dims(query.shape[-1], value.shape[-1], 16 // value.element_size())
+        except (AssertionError, ValueError) as error:
+            return replace(
+                result,
+                support=CapabilityResult.unsupported(f"FA4: {error} Select compatible inputs or another backend."),
+            )
+        if not verified:
+            return replace(
+                result,
+                support=CapabilityResult.unmigrated("Selected FA4 kernel has no available head-dimension validator"),
+                compilation_mode=CompilationMode.EAGER_ONLY,
+            )
+        return result
+
+    @staticmethod
+    def _resolve_rocm_execution_path(context: ExecutionContext) -> ExecutionPathResult:
+        """Describe AITER dispatch; device and compiler validation is pending."""
+        result = ExecutionPathResult.unmigrated("FLASH_ATTN", context, path="rocm_unverified")
+        if (
+            context.kernel_variant != "aiter"
+            or context.parallel_strategy is not ParallelStrategy.NONE
+            or context.outer_boundaries
+            or context.paged_kv
+            or context.piecewise
+            or context.kv_cache_dtype is not None
+        ):
+            return result
+        if context.packing_mode is not PackingMode.NONE:
+            path = "rocm_packed_varlen"
+        elif context.mask_mode is MaskMode.UNKNOWN:
+            path = "rocm_runtime_mask_dependent"
+        elif context.mask_mode is not MaskMode.NONE:
+            path = "rocm_masked_varlen"
+        else:
+            path = "rocm_dense"
+        return replace(
+            result,
+            path=path,
+            support=CapabilityResult.unmigrated("AITER routing only; kernel and compilation validation require ROCm"),
+        )
+
+    def _resolve_npu_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        """Describe existing MindIE routing without claiming hardware validation.
+
+        NPU packing has its own opt-in and [real, pad] contract. In particular,
+        incomplete packed metadata can fall back to a mask, so CUDA metadata
+        normalization must not run here. No tensor values are inspected.
+        """
+        context = replace(context, kernel_variant="mindiesd")
+        result = ExecutionPathResult.unmigrated("FLASH_ATTN", context, path="npu_unverified")
+        extra = attn_metadata.extra if attn_metadata else {}
+        if (
+            context.parallel_strategy is not ParallelStrategy.NONE
+            or context.outer_boundaries
+            or context.paged_kv
+            or self.causal
+            or extra.get("kv_cache_dtype") not in (None, "float", "auto")
+            or (attn_metadata is not None and attn_metadata.full_attn_spans is not None)
+        ):
+            return result
+
+        path = "npu_masked" if attn_metadata is not None and attn_metadata.attn_mask is not None else "npu_dense"
+        if extra.get("npu_attn_varlen", False):
+            if self._resolve_packed_seq_npu(query, key, extra) is not None:
+                path = (
+                    "npu_prefix_kv_slice"
+                    if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention"
+                    else "npu_packed_varlen"
+                )
+            else:
+                # This route may reject the request if neither an explicit mask
+                # nor usable valid_kv_length is supplied; it is not support.
+                path = "npu_masked_fallback"
+        return replace(
+            result,
+            path=path,
+            support=CapabilityResult.unmigrated("NPU routing only; kernel and compilation validation require Ascend"),
+        )
+
     def _warn_fa_deterministic_non_dense(self, path: str) -> None:
         if not self.fa_deterministic:
             return
@@ -401,6 +668,7 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         """CUDA/ROCm/MUSA flash attention implementation."""
         from vllm_omni.diffusion.attention.backends.utils.fa import (
             HAS_FLASH_ATTN,
+            IS_FLASH_ATTN_4,
             flash_attn_func,
             flash_attn_varlen_func,
         )
@@ -412,9 +680,10 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
             )
 
-        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
-        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
-        extra = attn_metadata.extra if attn_metadata is not None else {}
+        metadata_plan = _normalize_flash_attention_metadata(attn_metadata)
+        attention_mask = metadata_plan.attention_mask
+        full_attn_spans = metadata_plan.full_attn_spans
+        extra = metadata_plan.extra
 
         # Try piecewise attention
         if full_attn_spans is not None:
@@ -443,12 +712,7 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 query_ranges=None if attn_metadata is None else attn_metadata.query_ranges,
             )
 
-        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
-        present_packed_keys = [key for key in packed_keys if key in extra]
-        if present_packed_keys:
-            if len(present_packed_keys) != len(packed_keys):
-                missing = sorted(set(packed_keys) - set(present_packed_keys))
-                raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
+        if metadata_plan.packing_mode is not PackingMode.NONE:
             self._warn_fa_deterministic_non_dense("packed-varlen")
             return self._forward_varlen_packed(
                 query,
@@ -460,7 +724,14 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 max_seqlen_k=extra["max_seqlen_k"],
             )
 
-        if attention_mask is not None and torch.any(~attention_mask):
+        use_masked_path = False
+        if attention_mask is not None:
+            use_masked_path = (
+                torch.any(~attention_mask)
+                if metadata_plan.mask_mode is MaskMode.UNKNOWN
+                else metadata_plan.mask_mode is not MaskMode.NONE
+            )
+        if use_masked_path:
             self._warn_fa_deterministic_non_dense("masked-varlen")
             return self._forward_varlen_masked(
                 query,
@@ -470,6 +741,15 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
             )
 
         if flash_attn_func is not None:
+            if IS_FLASH_ATTN_4:
+                return _fa4_dense_attention_op(
+                    query,
+                    key,
+                    value,
+                    self.softmax_scale,
+                    self.causal,
+                    self.fa_deterministic,
+                )
             fa_kwargs = {
                 "causal": self.causal,
                 "softmax_scale": self.softmax_scale,

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -59,6 +59,8 @@ def is_flash_attn_4_available() -> bool:
 # Bind via aliases so mypy does not treat each candidate import as a redefinition.
 flash_attn_func: FlashAttnFn | None = None
 flash_attn_varlen_func: FlashAttnFn | None = None
+IS_FLASH_ATTN_4 = False
+IS_AITER = False
 
 if current_omni_platform.is_rocm():
     # ROCm: try Aiter first
@@ -71,6 +73,7 @@ if current_omni_platform.is_rocm():
 
             flash_attn_func = _fa_func
             flash_attn_varlen_func = _fa_varlen
+            IS_AITER = True
     except (ImportError, ModuleNotFoundError):
         pass
 elif current_omni_platform.is_xpu():
@@ -102,6 +105,7 @@ else:
 
                 flash_attn_func = _fa_func
                 flash_attn_varlen_func = _fa_varlen
+                IS_FLASH_ATTN_4 = True
                 logger.info("Using CuTe FlashAttention-4 on Blackwell")
             except Exception as exc:
                 # Optional FA4 dependencies may be present but ABI-incompatible
@@ -166,6 +170,27 @@ else:
 HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None
 
 
+def validate_fa4_head_dims(head_dim: int, head_dim_v: int, alignment: int) -> bool:
+    """Use the selected FA4 implementation's constraints during path resolution.
+
+    Return False when its validator is unavailable; propagate kernel validation
+    errors for the backend to report. Keep the private FA4 API dependency here,
+    rather than duplicating architecture/version-specific dimension rules.
+    This is a planning-time call, never part of compiled tensor execution.
+    """
+    try:
+        from flash_attn.cute.interface import _get_device_arch, _validate_head_dims
+    except ImportError:
+        return False
+
+    arch = _get_device_arch() // 10
+    # FA4 routes SM80/SM120 through separate kernels and bypasses this validator.
+    if arch not in (9, 10, 11):
+        return False
+    _validate_head_dims(head_dim, head_dim_v, arch, alignment)
+    return True
+
+
 def _choose_vllm_flash_attn_version(
     device_major: int,
     requested: str | int | None,

vllm_omni/diffusion/attention/capabilities.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+
+class SupportStatus(str, Enum):
+    SUPPORTED = "supported"
+    UNSUPPORTED = "unsupported"
+    UNMIGRATED = "unmigrated"
+
+
+class CompilationMode(str, Enum):
+    TRACEABLE = "traceable"
+    CUSTOM_OP = "custom_op"
+    EAGER_ONLY = "eager_only"
+
+
+class PackingMode(str, Enum):
+    NONE = "none"
+    PACKED_PADDING = "packed_padding"
+    MULTI_DOCUMENT = "multi_document"
+
+
+class MaskMode(str, Enum):
+    NONE = "none"
+    PADDING = "padding"
+    ARBITRARY = "arbitrary"
+    UNKNOWN = "unknown"
+
+
+class ParallelStrategy(str, Enum):
+    NONE = "none"
+    ULYSSES = "ulysses"
+    RING = "ring"
+    HYBRID_ULYSSES_RING = "hybrid_ulysses_ring"
+    ALLGATHER_KV = "allgather_kv"
+
+
+class OuterBoundary(str, Enum):
+    HSDP = "hsdp"
+
+
+@dataclass(frozen=True, slots=True)
+class CapabilityResult:
+    status: SupportStatus
+    reason: str | None = None
+
+    def __post_init__(self) -> None:
+        if self.status is SupportStatus.UNSUPPORTED and not self.reason:
+            raise ValueError("UNSUPPORTED capability results require an actionable reason")
+
+    @classmethod
+    def supported(cls) -> CapabilityResult:
+        return cls(SupportStatus.SUPPORTED)
+
+    @classmethod
+    def unsupported(cls, reason: str) -> CapabilityResult:
+        return cls(SupportStatus.UNSUPPORTED, reason)
+
+    @classmethod
+    def unmigrated(cls, reason: str | None = None) -> CapabilityResult:
+        return cls(SupportStatus.UNMIGRATED, reason)
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionContext:
+    platform: str
+    kernel_variant: str | None = None
+    dtype: str | None = None
+    causal: bool | None = None
+    mask_mode: MaskMode = MaskMode.NONE
+    packing_mode: PackingMode = PackingMode.NONE
+    piecewise: bool = False
+    paged_kv: bool = False
+    kv_cache_dtype: str | None = None
+    parallel_strategy: ParallelStrategy = ParallelStrategy.NONE
+    require_fullgraph: bool = False
+    outer_boundaries: frozenset[OuterBoundary] = frozenset()
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionPathResult:
+    backend: str
+    path: str
+    support: CapabilityResult
+    compilation_mode: CompilationMode
+    platform: str
+    kernel_variant: str | None
+    parallel_strategy: ParallelStrategy
+
+    @classmethod
+    def unmigrated(
+        cls,
+        backend: str,
+        context: ExecutionContext,
+        *,
+        path: str = "unmigrated",
+    ) -> ExecutionPathResult:
+        return cls(
+            backend=backend,
+            path=path,
+            support=CapabilityResult.unmigrated(),
+            compilation_mode=CompilationMode.EAGER_ONLY,
+            platform=context.platform,
+            kernel_variant=context.kernel_variant,
+            parallel_strategy=context.parallel_strategy,
+        )
+
+    def requested_support(self, context: ExecutionContext) -> CapabilityResult:
+        """Validate request-level guarantees without changing path migration state."""
+        if context.require_fullgraph and self.support.status is SupportStatus.UNMIGRATED:
+            return CapabilityResult.unsupported(
+                f"{self.backend} path {self.path!r} has no verified fullgraph declaration; "
+                "use a verified path or disable the guaranteed-fullgraph request"
+            )
+        if self.support.status is not SupportStatus.SUPPORTED:
+            return self.support
+        if context.require_fullgraph and self.compilation_mode is CompilationMode.EAGER_ONLY:
+            return CapabilityResult.unsupported(
+                f"{self.backend} path {self.path!r} does not support guaranteed fullgraph compilation"
+            )
+        return CapabilityResult.supported()

vllm_omni/diffusion/attention/layer.py

@@ -19,7 +19,13 @@ from vllm.model_executor.models.utils import extract_layer_index
 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
-from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend, SDPAImpl
+from vllm_omni.diffusion.attention.capabilities import (
+    ExecutionContext,
+    ExecutionPathResult,
+    OuterBoundary,
+    ParallelStrategy,
+)
 from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
@@ -119,6 +125,8 @@ class Attention(nn.Module):
 
         config = get_current_diffusion_config_or_none()
         attention_config = config.diffusion_attention_config if config is not None else None
+        parallel_config = getattr(config, "parallel_config", None)
+        self._hsdp_compile_boundary_enabled = bool(getattr(parallel_config, "use_hsdp", False))
 
         from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 
@@ -131,6 +139,9 @@ class Attention(nn.Module):
             is DiffusionKVCacheMode.PAGED_SCHEDULER
         )
         self._scheduler_paged_kv = scheduler_paged_kv
+        self.attn_backend: type[AttentionBackend] | None
+        self.attention: AttentionImpl | nn.Module
+        self.sdpa_fallback: SDPAImpl | None
         if custom_attention is None:
             attn_backend_cls, spec = get_attn_backend_for_role(
                 role=role,
@@ -166,7 +177,6 @@ class Attention(nn.Module):
                     attn_backend_cls.get_name(),
                     dense_backend_name,
                 )
-            parallel_config = getattr(config, "parallel_config", None)
             allgather_degree = getattr(parallel_config, "allgather_degree", 1)
             # TODO: Move AllGather-KV compatibility into an AttentionBackend capability
             # so validation does not depend on backend names.
@@ -279,6 +289,7 @@ class Attention(nn.Module):
 
         if self.paged_kv_cache_role is None:
             return None
+        assert self.attn_backend is not None  # Custom attention cannot opt into paged KV.
         dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
         # Keep backend layout discovery under the same config context used by
         # upstream vLLM's attention-spec collector.  vLLM 0.29 moved the
@@ -378,17 +389,68 @@ class Attention(nn.Module):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
-        if torch.compiler.is_compiling() and is_forward_context_available():
-            od_config = get_forward_context().omni_diffusion_config
-            parallel_config = getattr(od_config, "parallel_config", None)
-            if getattr(parallel_config, "use_hsdp", False):
-                # Keep HSDP/FSDP2 parameter all-gather outside Inductor's
-                # attention graph; otherwise scheduler dependency analysis can
-                # fail on the fused attention region.
-                return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata)
+        if torch.compiler.is_compiling() and self._uses_hsdp_compile_boundary():
+            # Keep HSDP/FSDP2 parameter all-gather outside Inductor's
+            # attention graph; otherwise scheduler dependency analysis can
+            # fail on the fused attention region.
+            return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata)
 
         return self._forward_impl(query, key, value, attn_metadata)
 
+    def _uses_hsdp_compile_boundary(self) -> bool:
+        if self._hsdp_compile_boundary_enabled:
+            return True
+        if not is_forward_context_available():
+            return False
+        od_config = get_forward_context().omni_diffusion_config
+        parallel_config = getattr(od_config, "parallel_config", None)
+        return bool(getattr(parallel_config, "use_hsdp", False))
+
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        """Compose backend capabilities with outer Attention boundaries."""
+        boundaries = set(context.outer_boundaries)
+        if self._uses_hsdp_compile_boundary():
+            boundaries.add(OuterBoundary.HSDP)
+        active_strategy = self._get_active_parallel_strategy()
+        strategy_name = active_strategy.name
+        if self.use_ring and active_strategy.enabled and strategy_name == "ulysses":
+            parallel_strategy = ParallelStrategy.HYBRID_ULYSSES_RING
+        elif self.use_ring and active_strategy.enabled:
+            parallel_strategy = ParallelStrategy.RING
+        else:
+            parallel_strategy = {
+                "allgather_kv": ParallelStrategy.ALLGATHER_KV,
+                "ulysses": ParallelStrategy.ULYSSES,
+            }.get(strategy_name, ParallelStrategy.NONE)
+        resolved_context = replace(
+            context,
+            outer_boundaries=frozenset(boundaries),
+            paged_kv=self.is_paged_kv_active(),
+            parallel_strategy=parallel_strategy,
+        )
+        resolver = getattr(self.attention, "resolve_execution_path", None)
+        if not callable(resolver):
+            return ExecutionPathResult.unmigrated(
+                type(self.attention).__name__,
+                resolved_context,
+            )
+        if not resolved_context.paged_kv:
+            attn_metadata = self._with_kv_cache_dtype(attn_metadata)
+        return resolver(
+            resolved_context,
+            query,
+            key,
+            value,
+            attn_metadata,
+        )
+
     @torch.compiler.disable
     def _forward_hsdp_compile_boundary(
         self,
@@ -422,10 +484,11 @@ class Attention(nn.Module):
             )
         use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
         if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
-            assert self.attn_backend is not None
+            backend_name = (
+                self.attn_backend.get_name() if self.attn_backend is not None else type(self.attention).__name__
+            )
             raise NotImplementedError(
-                f"Diffusion paged KV requires an Omni backend with paged support; "
-                f"selected {self.attn_backend.get_name()}"
+                f"Diffusion paged KV requires an Omni backend with paged support; selected {backend_name}"
             )
         if use_paged_attention and strategy is not self._no_parallel_strategy:
             strategy_name = strategy.name
@@ -554,9 +617,11 @@ class Attention(nn.Module):
             self._scheduler_paged_kv
             and self.paged_kv_cache_role is not None
             and in_kv_memory_profile
-            and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN"
+            and self.attn_backend is not None
+            and self.attn_backend.get_name() == "FLASH_ATTN"
             and not current_omni_platform.supports_diffusion_dense_flash_attention()
         ):
+            assert self.sdpa_fallback is not None
             logger.warning_once(
                 "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. "
                 "Formal paged requests still use the platform-native paged attention backend."

vllm_omni/diffusion/data.py

@@ -983,6 +983,9 @@ class OmniDiffusionConfig:
     enable_broadcast_weight_load: bool = False
     num_weight_load_threads: int = 4
 
+    # Shard meta parameters before reading rank-local HF safetensors slices.
+    hsdp_weight_load_strategy: str = "full"
+
     # Enable sleep mode
     enable_sleep_mode: bool = False
 
@@ -1196,6 +1199,10 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if self.hsdp_weight_load_strategy not in {"full", "pre_sharded"}:
+            raise ValueError(
+                f"hsdp_weight_load_strategy must be 'full' or 'pre_sharded', got {self.hsdp_weight_load_strategy!r}"
+            )
         from vllm_omni.diffusion.offloader.config import (
             OffloadStrategy,
             materialize_legacy_offload_flags,

vllm_omni/diffusion/distributed/hsdp.py

@@ -46,6 +46,14 @@ class HSDPInferenceConfig:
     reshard_after_forward: bool = True
 
 
+@dataclass
+class HSDPShardContext:
+    """Reusable state for incrementally sharding one HSDP model."""
+
+    hsdp_kwargs: dict[str, Any]
+    ignored_params: set[nn.Parameter] | None = None
+
+
 def _create_hsdp_mesh(
     device_type: str,
     replicate_size: int,
@@ -111,6 +119,38 @@ def apply_hsdp_to_model(
     Returns:
         HSDP-wrapped model ready for inference
     """
+    context = prepare_hsdp_shard_context(
+        model,
+        hsdp_config,
+        target_device=target_device,
+    )
+
+    hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)
+    if not hsdp_shard_conditions:
+        raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined")
+
+    # Apply HSDP sharding, this will automatically handle weight distribution
+    shard_model(
+        model,
+        hsdp_shard_conditions=hsdp_shard_conditions,
+        context=context,
+    )
+
+    logger.info("HSDP applied to model: %s", type(model).__name__)
+    return model
+
+
+def prepare_hsdp_shard_context(
+    model: nn.Module,
+    hsdp_config: HSDPInferenceConfig,
+    target_device: torch.device | None = None,
+) -> HSDPShardContext:
+    """Prepare mesh, precision, and ignored-parameter state for HSDP.
+
+    The returned context can be reused to shard completed child modules while
+    checkpoint weights are still streaming, then to shard the root once all
+    children have been loaded.
+    """
     if not hsdp_config.enabled:
         raise ValueError("HSDP is not enabled in config")
 
@@ -160,7 +200,7 @@ def apply_hsdp_to_model(
     )
 
     hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)
-    if not hsdp_shard_conditions or len(hsdp_shard_conditions) == 0:
+    if not hsdp_shard_conditions:
         raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined")
 
     # Collect parameters of any modules the model wants excluded from FSDP sharding.
@@ -188,7 +228,11 @@ def apply_hsdp_to_model(
         except AttributeError:
             logger.warning("_hsdp_ignored_modules entry %r not found on model", mod_name)
             continue
-        sub_mod.to(target_device)
+        has_meta_tensor = any(tensor.device.type == "meta" for tensor in (*sub_mod.parameters(), *sub_mod.buffers()))
+        if has_meta_tensor:
+            sub_mod.to_empty(device=target_device)
+        else:
+            sub_mod.to(target_device)
         ignored_params.update(sub_mod.parameters())
     if ignored_params:
         logger.info(
@@ -224,21 +268,36 @@ def apply_hsdp_to_model(
             target_device,
         )
 
-    # Apply HSDP sharding, this will automatically handle weight distribution
-    shard_model(
-        model,
-        reshard_after_forward=hsdp_config.reshard_after_forward,
-        mp_policy=mp_policy,
-        mesh=device_mesh,
-        hsdp_shard_conditions=hsdp_shard_conditions,
-        ignored_params=ignored_params if ignored_params else None,
+    return HSDPShardContext(
+        hsdp_kwargs={
+            "reshard_after_forward": hsdp_config.reshard_after_forward,
+            "mesh": device_mesh,
+            "mp_policy": mp_policy,
+        },
+        ignored_params=ignored_params or None,
     )
 
-    for param in model.parameters():
+
+def shard_hsdp_module(module: nn.Module, context: HSDPShardContext) -> None:
+    """Shard one completed child module with a prepared HSDP context."""
+    module_kwargs = dict(context.hsdp_kwargs)
+    if context.ignored_params:
+        module_ignored_params = context.ignored_params.intersection(module.parameters())
+        if module_ignored_params:
+            module_kwargs["ignored_params"] = module_ignored_params
+    fully_shard(module, **module_kwargs)
+    for param in module.parameters():
         param.requires_grad = False
 
-    logger.info("HSDP applied to model: %s", type(model).__name__)
-    return model
+
+def finalize_hsdp_root(model: nn.Module, context: HSDPShardContext) -> None:
+    """Shard the model root after all selected child modules are sharded."""
+    root_kwargs = dict(context.hsdp_kwargs)
+    if context.ignored_params:
+        root_kwargs["ignored_params"] = context.ignored_params
+    fully_shard(model, **root_kwargs)
+    for param in model.parameters():
+        param.requires_grad = False
 
 
 def shard_model(
@@ -249,6 +308,7 @@ def shard_model(
     mesh: DeviceMesh | None = None,
     hsdp_shard_conditions: list[Callable[[str, nn.Module], bool]],
     ignored_params: set[nn.Parameter] | None = None,
+    context: HSDPShardContext | None = None,
 ) -> None:
     """Apply HSDP sharding to model modules based on shard conditions.
 
@@ -259,32 +319,30 @@ def shard_model(
     This is required for packed integer parameters inside sharded transformer
     blocks; the root wrap receives the full set for all remaining parameters.
     """
-    hsdp_kwargs: dict[str, Any] = {
-        "reshard_after_forward": reshard_after_forward,
-        "mesh": mesh,
-        "mp_policy": mp_policy,
-    }
+    if context is None:
+        context = HSDPShardContext(
+            hsdp_kwargs={
+                "reshard_after_forward": reshard_after_forward,
+                "mesh": mesh,
+                "mp_policy": mp_policy,
+            },
+            ignored_params=ignored_params,
+        )
+    elif any(value is not None for value in (mp_policy, mesh, ignored_params)) or not reshard_after_forward:
+        raise ValueError("Pass either context or individual HSDP shard options, not both")
 
     num_sharded = 0
     for name, module in reversed(list(model.named_modules())):
         if any(cond(name, module) for cond in hsdp_shard_conditions):
-            module_kwargs = dict(hsdp_kwargs)
-            if ignored_params:
-                module_ignored_params = ignored_params.intersection(module.parameters())
-                if module_ignored_params:
-                    module_kwargs["ignored_params"] = module_ignored_params
-            fully_shard(module, **module_kwargs)
+            shard_hsdp_module(module, context)
             num_sharded += 1
 
     if num_sharded == 0:
         raise ValueError("No modules were sharded")
 
-    root_kwargs = dict(hsdp_kwargs)
-    if ignored_params:
-        root_kwargs["ignored_params"] = ignored_params
-    fully_shard(model, **root_kwargs)
+    finalize_hsdp_root(model, context)
     logger.info(
         "Sharded %d modules + root (ignored_params=%d)",
         num_sharded,
-        len(ignored_params) if ignored_params else 0,
+        len(context.ignored_params) if context.ignored_params else 0,
     )

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -16,6 +16,7 @@ import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
@@ -30,7 +31,14 @@ from vllm.utils.torch_utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp_to_model
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    HSDPShardContext,
+    apply_hsdp_to_model,
+    finalize_hsdp_root,
+    prepare_hsdp_shard_context,
+    shard_hsdp_module,
+)
 from vllm_omni.diffusion.lora.manager import LoRABackend
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
@@ -38,6 +46,8 @@ from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
 from vllm_omni.diffusion.model_loader.host_weight_loader import HWRLoaderMixin, _HWRCommitError
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
+    TensorBinding,
+    build_checkpoint_binding_plan,
     build_checkpoint_mmap_plan,
     has_online_quantization,
 )
@@ -49,7 +59,7 @@ from vllm_omni.diffusion.offloader.config import (
     OffloadStrategy,
     resolve_offload,
 )
-from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery, PipelineModules
 from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan
 from vllm_omni.diffusion.registry import initialize_model
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
@@ -138,6 +148,36 @@ def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> typ
     )
 
 
+def _is_distributed_rank_zero() -> bool:
+    return (
+        not torch.distributed.is_available()
+        or not torch.distributed.is_initialized()
+        or torch.distributed.get_rank() == 0
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class _HSDPLoadRoot:
+    name: str
+    module: nn.Module
+    context: HSDPShardContext
+
+
+@dataclasses.dataclass(frozen=True)
+class _HSDPModuleLoadGroup:
+    name: str
+    module: nn.Module
+    context: HSDPShardContext
+
+
+@dataclasses.dataclass(frozen=True)
+class _PreShardedHSDPLoadPlan:
+    roots: list[_HSDPLoadRoot]
+    groups: list[_HSDPModuleLoadGroup]
+    bindings: dict[str, TensorBinding]
+    binding_names: frozenset[str]
+
+
 class DiffusersPipelineLoader(HWRLoaderMixin):
     """Model loader that can load diffusers pipeline components from disk."""
 
@@ -1222,6 +1262,451 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                     raise ValueError(f"Unknown load_format: {load_format}")
         return model
 
+    def _load_model_with_pre_sharded_hsdp(
+        self,
+        model: nn.Module,
+        discovered_modules: PipelineModules,
+        hsdp_config: HSDPInferenceConfig,
+        target_device: torch.device,
+    ) -> nn.Module:
+        """Shard meta DiT parameters before reading their rank-local checkpoint slices."""
+        from torch.distributed.checkpoint import HuggingFaceStorageReader
+        from torch.distributed.checkpoint import load as dcp_load
+
+        if self.counter_before_loading_weights == 0.0:
+            self.counter_before_loading_weights = time.perf_counter()
+
+        load_plan = self._prepare_pre_sharded_hsdp_load(
+            model,
+            discovered_modules,
+            hsdp_config,
+            target_device,
+        )
+        self._check_no_unplanned_meta_tensors(model, load_plan.binding_names)
+        self._validate_pre_sharded_hsdp_bindings(load_plan)
+
+        self._process_pre_sharded_hsdp_groups_on_meta(load_plan)
+
+        for group in load_plan.groups:
+            logger.debug("Applying meta-first HSDP shard to %s", group.name)
+            shard_hsdp_module(group.module, group.context)
+        for root in load_plan.roots:
+            logger.debug("Finalizing meta-first HSDP root %s", root.name)
+            finalize_hsdp_root(root.module, root.context)
+
+        self._materialize_pre_sharded_hsdp_state(model, load_plan, target_device)
+
+        targets = self._get_pre_sharded_hsdp_targets(model, load_plan)
+        checkpoint_states, local_payload_bytes = self._prepare_pre_sharded_checkpoint_states(
+            load_plan,
+            targets,
+        )
+
+        thread_count = int(getattr(self.od_config, "num_weight_load_threads", 4))
+        if _is_distributed_rank_zero():
+            directory_label = "directory" if len(checkpoint_states) == 1 else "directories"
+            logger.info(
+                "Pre-sharded HSDP loading %.2f GiB of rank-local tensor slices directly into "
+                "FSDP-owned storage from %d HF safetensors %s with %d reader thread(s) per rank",
+                local_payload_bytes / 1024**3,
+                len(checkpoint_states),
+                directory_label,
+                thread_count,
+            )
+
+        for checkpoint_dir in sorted(checkpoint_states, key=str):
+            reader = HuggingFaceStorageReader(str(checkpoint_dir), thread_count=thread_count)
+            # Each worker has the full checkpoint available and each DTensor
+            # already describes its rank-local destination slice. Avoid DCP's
+            # coordinator collectives, which can leave rank 0 with less GPU
+            # memory headroom than the other HSDP workers.
+            dcp_load(checkpoint_states[checkpoint_dir], storage_reader=reader, no_dist=True)
+
+        remaining_meta = self._pre_sharded_meta_target_names(model, load_plan)
+        if remaining_meta:
+            raise ValueError(
+                "Pre-sharded HSDP checkpoint loading left tensors on meta; "
+                f"first missing tensors: {remaining_meta[:16]}"
+            )
+
+        self._finalize_pre_sharded_hsdp_loaded_weights(load_plan)
+
+        self.counter_after_loading_weights = time.perf_counter()
+        logger.info_once(
+            "Loading weights took %.2f seconds",
+            self.counter_after_loading_weights - self.counter_before_loading_weights,
+        )
+        self._move_non_hsdp_modules(discovered_modules, target_device)
+        return model
+
+    @staticmethod
+    def _validate_pre_sharded_hsdp_bindings(load_plan: _PreShardedHSDPLoadPlan) -> None:
+        unsupported = [name for name, binding in load_plan.bindings.items() if binding.transform is not None]
+        if unsupported:
+            raise ValueError(
+                "Pre-sharded HSDP requires checkpoint tensors in runtime layout; "
+                f"tensor transforms are unsupported: {unsupported[:16]}"
+            )
+        non_safetensors = sorted(
+            {
+                binding.file_path
+                for binding in load_plan.bindings.values()
+                if not binding.file_path.endswith(".safetensors")
+            }
+        )
+        if non_safetensors:
+            raise ValueError(
+                f"Pre-sharded HSDP requires existing HF safetensors files; unsupported files: {non_safetensors[:8]}"
+            )
+
+    @staticmethod
+    def _get_pre_sharded_hsdp_targets(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> dict[str, torch.Tensor]:
+        targets: dict[str, torch.Tensor] = {
+            **dict(model.named_buffers()),
+            **dict(model.named_parameters()),
+        }
+        missing_targets = sorted(load_plan.binding_names - targets.keys())
+        if missing_targets:
+            raise ValueError(
+                f"Pre-sharded HSDP changed checkpoint target names; first missing targets: {missing_targets[:16]}"
+            )
+        return targets
+
+    @staticmethod
+    def _materialize_pre_sharded_hsdp_state(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+        target_device: torch.device,
+    ) -> None:
+        """Materialize the installed FSDP rank-local state before checkpoint I/O.
+
+        ``fully_shard`` supports meta initialization through ``Module.to_empty``:
+        its registered ``_apply`` hooks update the internal sharded storage to
+        point at the newly allocated rank-local tensors. Modules are materialized
+        in the same HSDP-group order and child-first traversal used by ordinary
+        ``fully_shard`` initialization instead of the forward traversal used by
+        a recursive ``to_empty`` call.
+        PyTorch's swap-on-conversion mode preserves Parameter identity so that
+        post-load kernel objects created before sharding cannot retain stale
+        references to the original meta Parameters.
+
+        A complete checkpoint plan covers every parameter and persistent buffer
+        in each DiT.  Non-persistent buffers are not checkpoint entries, so keep
+        their initialized values across ``to_empty``.
+        """
+        unplanned_parameters: list[str] = []
+        saved_buffers: dict[str, torch.Tensor] = {}
+        for root in load_plan.roots:
+            prefix = f"{root.name}."
+            for local_name, _parameter in root.module.named_parameters():
+                full_name = prefix + local_name
+                if full_name not in load_plan.binding_names:
+                    unplanned_parameters.append(full_name)
+            for local_name, buffer in root.module.named_buffers():
+                full_name = prefix + local_name
+                if full_name in load_plan.binding_names:
+                    continue
+                if buffer.device.type == "meta":
+                    raise ValueError(f"Pre-sharded HSDP cannot preserve an unplanned meta buffer: {full_name}")
+                saved_buffers[full_name] = buffer.detach().cpu().clone()
+
+        if unplanned_parameters:
+            raise ValueError(
+                "Pre-sharded HSDP requires checkpoint coverage for every DiT parameter; "
+                f"first unplanned parameters: {unplanned_parameters[:16]}"
+            )
+
+        previous_swap_mode = torch.__future__.get_swap_module_params_on_conversion()
+        torch.__future__.set_swap_module_params_on_conversion(True)
+        try:
+            sharded_group_roots = {group.module for group in load_plan.groups}
+            materialized_modules: set[nn.Module] = set()
+            for group in load_plan.groups:
+                for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+                    group.module,
+                    nested_hsdp_roots=sharded_group_roots - {group.module},
+                ):
+                    if module not in materialized_modules:
+                        module.to_empty(device=target_device, recurse=False)
+                        materialized_modules.add(module)
+            for root in load_plan.roots:
+                for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+                    root.module,
+                    nested_hsdp_roots=sharded_group_roots,
+                ):
+                    if module not in materialized_modules:
+                        module.to_empty(device=target_device, recurse=False)
+                        materialized_modules.add(module)
+        finally:
+            torch.__future__.set_swap_module_params_on_conversion(previous_swap_mode)
+
+        if saved_buffers:
+            materialized_buffers = dict(model.named_buffers())
+            with torch.no_grad():
+                for name, value in saved_buffers.items():
+                    target = materialized_buffers[name]
+                    target.copy_(value.to(device=target.device, dtype=target.dtype))
+
+    @staticmethod
+    def _hsdp_managed_modules_post_order(
+        root: nn.Module,
+        *,
+        nested_hsdp_roots: set[nn.Module],
+    ) -> list[nn.Module]:
+        """Mirror FSDP's child-first managed-module traversal.
+
+        Nested FSDP roots own their own parameters and are materialized in their
+        corresponding load group, so the enclosing traversal must not descend
+        into them.
+        """
+        modules: list[nn.Module] = []
+        visited: set[nn.Module] = set()
+
+        def visit(module: nn.Module) -> None:
+            if module in visited:
+                return
+            visited.add(module)
+            for child in module.children():
+                if child not in nested_hsdp_roots:
+                    visit(child)
+            modules.append(module)
+
+        visit(root)
+        return modules
+
+    @staticmethod
+    def _prepare_pre_sharded_checkpoint_states(
+        load_plan: _PreShardedHSDPLoadPlan,
+        targets: dict[str, torch.Tensor],
+    ) -> tuple[dict[Path, dict[str, torch.Tensor]], int]:
+        from torch.distributed.tensor import DTensor
+
+        checkpoint_states: dict[Path, dict[str, torch.Tensor]] = {}
+        checkpoint_owners: dict[tuple[Path, str], str] = {}
+        local_payload_bytes = 0
+        for name, binding in load_plan.bindings.items():
+            target = targets[name]
+
+            checkpoint_dir = Path(binding.file_path).parent
+            owner_key = (checkpoint_dir, binding.checkpoint_key)
+            if previous_owner := checkpoint_owners.get(owner_key):
+                raise ValueError(
+                    "Pre-sharded HSDP cannot bind one checkpoint tensor to multiple runtime tensors: "
+                    f"{binding.checkpoint_key!r} maps to both {previous_owner!r} and {name!r}"
+                )
+            checkpoint_owners[owner_key] = name
+
+            if target.device.type == "meta":
+                raise ValueError(f"Pre-sharded HSDP target {name!r} was not materialized")
+
+            local_target = target.to_local() if isinstance(target, DTensor) else target
+            local_payload_bytes += local_target.numel() * local_target.element_size()
+            checkpoint_states.setdefault(checkpoint_dir, {})[binding.checkpoint_key] = target
+        return checkpoint_states, local_payload_bytes
+
+    @staticmethod
+    def _pre_sharded_meta_target_names(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> list[str]:
+        from torch.distributed.tensor import DTensor
+
+        targets = {
+            **dict(model.named_buffers()),
+            **dict(model.named_parameters()),
+        }
+        return [
+            name
+            for name in load_plan.binding_names
+            if (target := targets[name]).device.type == "meta"
+            or (isinstance(target, DTensor) and target.to_local().device.type == "meta")
+        ]
+
+    @staticmethod
+    def _process_pre_sharded_hsdp_groups_on_meta(load_plan: _PreShardedHSDPLoadPlan) -> None:
+        modules_to_process: list[tuple[nn.Module, QuantizeMethodBase]] = []
+        for group in load_plan.groups:
+            for _, module in group.module.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
+                    continue
+                modules_to_process.append((module, quant_method))
+
+        if not modules_to_process:
+            return
+        for root in load_plan.roots:
+            if not getattr(root.module, "_hsdp_pre_sharded_meta_post_load", False):
+                raise ValueError(
+                    f"Model {type(root.module).__name__} has not declared its post-load processing safe "
+                    "for pre-sharded meta initialization"
+                )
+        for module, quant_method in modules_to_process:
+            if type(quant_method) is not UnquantizedLinearMethod:
+                raise ValueError(
+                    "Pre-sharded HSDP does not support the selected post-load processing; "
+                    f"{type(module).__name__} uses {type(quant_method).__name__}"
+                )
+        for module, quant_method in modules_to_process:
+            quant_method.process_weights_after_loading(module)
+
+    def _prepare_pre_sharded_hsdp_load(
+        self,
+        model: nn.Module,
+        discovered_modules: PipelineModules,
+        hsdp_config: HSDPInferenceConfig,
+        target_device: torch.device,
+    ) -> _PreShardedHSDPLoadPlan:
+        outer_dit_names, outer_dits = discovered_modules.outermost_dits()
+        if not outer_dits:
+            raise ValueError("No DiT modules discovered for HSDP sharding")
+
+        sources = self._get_weight_sources(model)
+        plan_result = build_checkpoint_binding_plan(
+            model,
+            dit_modules=tuple(zip(outer_dit_names, outer_dits)),
+            sources=sources,
+            model_path=str(getattr(self.od_config, "model", "")) or None,
+            tensor_parallel_size=int(getattr(self.parallel_config, "tensor_parallel_size", 1)),
+            online_quantization=False,
+        )
+        if plan_result.plan is None:
+            raise ValueError(f"Pre-sharded HSDP checkpoint loading is incompatible: {plan_result.fallback_reason}")
+        plan = plan_result.plan
+        unsupported_sources = [source.prefix for source in sources if source.prefix not in plan.planned_source_prefixes]
+        if unsupported_sources:
+            raise ValueError(
+                "Pre-sharded HSDP checkpoint loading requires dedicated DiT sources; "
+                f"unsupported source prefixes: {unsupported_sources}"
+            )
+
+        roots: list[_HSDPLoadRoot] = []
+        groups: list[_HSDPModuleLoadGroup] = []
+        for outer_name, outer_dit in zip(outer_dit_names, outer_dits):
+            conditions = getattr(outer_dit, "_hsdp_shard_conditions", None)
+            if not conditions:
+                raise ValueError(f"Model {type(outer_dit).__name__} has no _hsdp_shard_conditions defined")
+
+            context = prepare_hsdp_shard_context(
+                outer_dit,
+                hsdp_config,
+                target_device=target_device,
+            )
+            roots.append(_HSDPLoadRoot(outer_name, outer_dit, context))
+            for local_name, module in reversed(list(outer_dit.named_modules())):
+                if not any(condition(local_name, module) for condition in conditions):
+                    continue
+                if not local_name:
+                    raise ValueError("_hsdp_shard_conditions must not select the DiT root")
+                full_name = f"{outer_name}.{local_name}" if local_name else outer_name
+                groups.append(_HSDPModuleLoadGroup(full_name, module, context))
+
+        parameters = dict(model.named_parameters())
+        targets: dict[str, torch.Tensor] = {**dict(model.named_buffers()), **parameters}
+        preserved_params = {param for root in roots for param in root.context.ignored_params or ()}
+        binding_names = frozenset(plan.bindings)
+        self._release_hsdp_checkpoint_targets_to_meta(
+            model,
+            targets,
+            parameters,
+            binding_names,
+            preserved_params,
+        )
+
+        return _PreShardedHSDPLoadPlan(
+            roots=roots,
+            groups=groups,
+            bindings=plan.bindings,
+            binding_names=binding_names,
+        )
+
+    @staticmethod
+    def _check_no_unplanned_meta_tensors(model: nn.Module, planned_names: frozenset[str]) -> None:
+        unplanned_meta = [
+            name
+            for name, tensor in (*model.named_parameters(), *model.named_buffers())
+            if tensor.device.type == "meta" and name not in planned_names
+        ]
+        if unplanned_meta:
+            raise ValueError(
+                "Pre-sharded HSDP initialized tensors on meta that are not covered by the "
+                f"checkpoint plan. First unplanned tensors: {unplanned_meta[:16]}"
+            )
+
+    @staticmethod
+    def _release_hsdp_checkpoint_targets_to_meta(
+        model: nn.Module,
+        targets: dict[str, torch.Tensor],
+        parameters: dict[str, nn.Parameter],
+        binding_names: frozenset[str],
+        preserved_params: set[nn.Parameter],
+    ) -> None:
+        released: dict[str, tuple[int, float]] = {}
+        for name in binding_names:
+            target = targets[name]
+            if target in preserved_params or target.device.type == "meta":
+                continue
+            logical_gib = target.numel() * target.element_size() / 1024**3
+            device_type = target.device.type
+            count, gib = released.get(device_type, (0, 0.0))
+            released[device_type] = (count + 1, gib + logical_gib)
+            if name in parameters:
+                replacement = DiffusersPipelineLoader._make_parameter_like(
+                    target,
+                    torch.empty_like(target, device="meta"),
+                )
+                torch.utils.swap_tensors(target, replacement)
+            else:
+                replacement = torch.empty_like(target, device="meta")
+                torch.utils.swap_tensors(target, replacement)
+        if released:
+            logger.info("Released HSDP checkpoint-covered tensors to meta before pre-sharded load: %s", released)
+
+    @staticmethod
+    def _make_parameter_like(template: nn.Parameter, tensor: torch.Tensor) -> nn.Parameter:
+        """Create a Parameter on new storage while preserving loader metadata."""
+        if type(template) is nn.Parameter:
+            replacement = nn.Parameter(tensor, requires_grad=template.requires_grad)
+            replacement.__dict__.update(getattr(template, "__dict__", {}))
+            return replacement
+
+        replacement = template.__class__.__new__(template.__class__, tensor)
+        replacement.__dict__.update(getattr(template, "__dict__", {}))
+        replacement.requires_grad_(template.requires_grad)
+        return replacement
+
+    @staticmethod
+    def _finalize_pre_sharded_hsdp_loaded_weights(
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> None:
+        """Run model post-load hooks for tensors materialized by pre-sharded loading."""
+        for root in load_plan.roots:
+            post_load = getattr(root.module, "post_load_weights", None)
+            if callable(post_load):
+                post_load()
+            root.module.eval()
+            validate = getattr(root.module, "validate_loaded_weights", None)
+            if callable(validate):
+                prefix = f"{root.name}."
+                validate({name.removeprefix(prefix) for name in load_plan.binding_names if name.startswith(prefix)})
+
+    @staticmethod
+    def _move_non_hsdp_modules(
+        discovered_modules: PipelineModules,
+        target_device: torch.device,
+    ) -> None:
+        modules_to_move: list[nn.Module] = []
+        if discovered_modules.vaes is not None:
+            modules_to_move.extend(discovered_modules.vaes)
+        if discovered_modules.encoders is not None:
+            modules_to_move.extend(discovered_modules.encoders)
+        if discovered_modules.resident_modules is not None:
+            modules_to_move.extend(discovered_modules.resident_modules)
+        for module in modules_to_move:
+            module.to(target_device)
+
     def _load_model_with_hsdp(
         self,
         target_device: torch.device,
@@ -1251,6 +1736,29 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         # mapping (QKV fusion, etc.).
         if load_format == "diffusers":
             raise ValueError("HSDP is not supported with the diffusers adapter load format")
+        strategy = getattr(self.od_config, "hsdp_weight_load_strategy", "full")
+        if strategy == "pre_sharded":
+            if load_format != "default":
+                raise ValueError(
+                    f"hsdp_weight_load_strategy={strategy!r} currently supports only diffusion_load_format='default'"
+                )
+            if self.quant_config is not None:
+                raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support quantization yet")
+            if getattr(self.od_config, "lora_path", None):
+                raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support LoRA yet")
+        elif strategy != "full":
+            raise ValueError(f"Unknown hsdp_weight_load_strategy: {strategy!r}")
+
+        if strategy == "pre_sharded":
+            model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
+            discovered_modules = ModuleDiscovery.discover(model)
+            return self._load_model_with_pre_sharded_hsdp(
+                model,
+                discovered_modules,
+                hsdp_config,
+                target_device,
+            )
+
         # Same host-fallback bound as the ordinary path: with a quant config the
         # HSDP path initializes on the accelerator (hsdp_defer_to_cpu=False), so
         # over-wide NPU-unquantizable fallback weights would otherwise round-trip

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Loader-owned host-weight plans shared with diffusion offload backends."""
+"""Loader-owned host-weight plans shared by direct-checkpoint consumers."""
 
 from __future__ import annotations
 
@@ -205,7 +205,7 @@ def _build_source_map(
     if not model_to_ckpt:
         raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources")
     logger.info(
-        "Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight",
+        "Indexed %d runtime tensor names from %d checkpoint keys for checkpoint preflight",
         len(model_to_ckpt),
         indexed_keys,
     )
@@ -313,28 +313,28 @@ def _validate_source_metadata(
                     )
 
 
-def build_checkpoint_mmap_plan(
+def build_checkpoint_binding_plan(
     pipeline: nn.Module,
     *,
     dit_modules: Sequence[tuple[str, nn.Module]],
     sources: Sequence[object],
     model_path: str | None,
     tensor_parallel_size: int,
-    use_hsdp: bool,
     online_quantization: bool,
+    use_pipeline_key_remap: bool = True,
     has_distilled_lora: bool = False,
 ) -> HostWeightPlanResult:
-    """Build a complete direct-checkpoint plan or return a fallback reason."""
+    """Build a complete binding plan for the pre-sharded HSDP consumer."""
     if tensor_parallel_size != 1:
         return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
-    if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
     if online_quantization:
         return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
     if has_distilled_lora:
         return HostWeightPlanResult(None, "distilled LoRA requires the ordinary loader")
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
+    if not callable(remap_fn) and use_pipeline_key_remap:
+        remap_fn = getattr(pipeline, "remap_checkpoint_key", None)
     if not callable(remap_fn):
         remap_fn = None
     adapter = get_direct_mmap_adapter(pipeline)
@@ -387,10 +387,37 @@ def build_checkpoint_mmap_plan(
     )
 
 
+def build_checkpoint_mmap_plan(
+    pipeline: nn.Module,
+    *,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    sources: Sequence[object],
+    model_path: str | None,
+    tensor_parallel_size: int,
+    use_hsdp: bool,
+    online_quantization: bool,
+    has_distilled_lora: bool = False,
+) -> HostWeightPlanResult:
+    """Build the DLO mmap plan without changing its HSDP eligibility."""
+    if use_hsdp:
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
+    return build_checkpoint_binding_plan(
+        pipeline,
+        dit_modules=dit_modules,
+        sources=sources,
+        model_path=model_path,
+        tensor_parallel_size=tensor_parallel_size,
+        online_quantization=online_quantization,
+        use_pipeline_key_remap=False,
+        has_distilled_lora=has_distilled_lora,
+    )
+
+
 __all__ = [
     "HostWeightPlan",
     "HostWeightPlanResult",
     "TensorBinding",
+    "build_checkpoint_binding_plan",
     "build_checkpoint_mmap_plan",
     "has_online_quantization",
 ]

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -38,6 +38,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
+from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
 from vllm_omni.diffusion.offloader.config import (
     OffloadStrategy,
     resolve_offload_strategy,
@@ -1084,6 +1085,7 @@ class Cosmos3LanguageModel(nn.Module):
         rope_theta: float,
         mrope_section: list[int],
         quant_config: QuantizationConfig | None = None,
+        release_completed_blocks_to_meta: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -1093,21 +1095,21 @@ class Cosmos3LanguageModel(nn.Module):
             rope_theta=rope_theta,
             mrope_section=mrope_section,
         )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3UndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
+        self.layers = nn.ModuleList()
+        for i in range(num_hidden_layers):
+            layer = Cosmos3UndDecoderLayer(
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=head_dim,
+                rms_norm_eps=rms_norm_eps,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.layers.append(layer)
         # TODO: Not used right now, will be used in the future for prompt upsampler.
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
@@ -1228,6 +1230,10 @@ class Cosmos3VFMTransformer(nn.Module):
 
     _hsdp_shard_conditions = [_is_transformer_block]
 
+    # Standard unquantized vLLM linears need no value-dependent post-load
+    # processing on accelerators. Edge inherits the same loading contract.
+    _hsdp_pre_sharded_meta_post_load = True
+
     # Modules whose parameters must NOT be FSDP-sharded at the root level.
     # time_embedder is cast to fp32 by post_load_weights for precision; if it
     # were swept into the root flat-parameter under MixedPrecisionPolicy(param_dtype=bf16),
@@ -1346,6 +1352,10 @@ class Cosmos3VFMTransformer(nn.Module):
         language_model_quant_config, gen_layers_quant_config = _resolve_cosmos3_quant_configs(
             getattr(od_config, "quantization_config", None)
         )
+        release_completed_blocks_to_meta = bool(
+            getattr(getattr(od_config, "parallel_config", None), "use_hsdp", False)
+            and getattr(od_config, "hsdp_weight_load_strategy", "full") == "pre_sharded"
+        )
         mixed_precision_config, mixed_precision_source = resolve_mixed_precision_config(od_config)
         if mixed_precision_config is None:
             if mixed_precision_source == "additional_config_disabled":
@@ -1372,6 +1382,7 @@ class Cosmos3VFMTransformer(nn.Module):
             rope_theta=self.rope_theta,
             mrope_section=self.mrope_section,
             quant_config=language_model_quant_config,
+            release_completed_blocks_to_meta=release_completed_blocks_to_meta,
             prefix="language_model",
             **self._language_model_kwargs(),
         )
@@ -1399,24 +1410,24 @@ class Cosmos3VFMTransformer(nn.Module):
             self.audio_proj_out = nn.Linear(self.hidden_size, self.sound_dim)
             self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
 
-        self.gen_layers = nn.ModuleList(
-            [
-                Cosmos3GenDecoderLayer(
-                    layer_idx=i,
-                    hidden_size=self.hidden_size,
-                    intermediate_size=self.intermediate_size,
-                    num_attention_heads=self.num_attention_heads,
-                    num_key_value_heads=self.num_key_value_heads,
-                    head_dim=self.head_dim,
-                    rms_norm_eps=self.rms_norm_eps,
-                    quant_config=gen_layers_quant_config,
-                    mlp_cls=self._gen_mlp_cls,
-                    qk_norm=self.qk_norm_for_diffusion,
-                    prefix=f"gen_layers.{i}",
-                )
-                for i in range(self.num_hidden_layers)
-            ]
-        )
+        self.gen_layers = nn.ModuleList()
+        for i in range(self.num_hidden_layers):
+            layer = Cosmos3GenDecoderLayer(
+                layer_idx=i,
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size,
+                num_attention_heads=self.num_attention_heads,
+                num_key_value_heads=self.num_key_value_heads,
+                head_dim=self.head_dim,
+                rms_norm_eps=self.rms_norm_eps,
+                quant_config=gen_layers_quant_config,
+                mlp_cls=self._gen_mlp_cls,
+                qk_norm=self.qk_norm_for_diffusion,
+                prefix=f"gen_layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.gen_layers.append(layer)
 
         self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None
         if mixed_precision_config is not None:

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelL
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
+from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
 
 from .transformer_cosmos3 import (
     Cosmos3VFMTransformer,
@@ -258,6 +259,7 @@ class Cosmos3EdgeLanguageModel(nn.Module):
         mrope_section: list[int],
         use_und_k_norm_for_gen: bool,
         quant_config: QuantizationConfig | None = None,
+        release_completed_blocks_to_meta: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -267,22 +269,22 @@ class Cosmos3EdgeLanguageModel(nn.Module):
             rope_theta=rope_theta,
             mrope_section=mrope_section,
         )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3EdgeUndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    use_und_k_norm_for_gen=use_und_k_norm_for_gen,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
+        self.layers = nn.ModuleList()
+        for i in range(num_hidden_layers):
+            layer = Cosmos3EdgeUndDecoderLayer(
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=head_dim,
+                rms_norm_eps=rms_norm_eps,
+                use_und_k_norm_for_gen=use_und_k_norm_for_gen,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.layers.append(layer)
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
     def forward(

vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py

@@ -288,13 +288,13 @@ def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig):
             output_key = output_keys[0]
             frames = frames[output_key]
         output_type = getattr(sampling_params, "output_type", None) or "pt"
-        # The image serving path currently accepts PIL images or NumPy arrays,
-        # while LingBot decodes T2I outputs to tensors. Keep that compatibility
-        # conversion without discarding the model's image/video payload key.
+        # Decoded LingBot pixels are already in [0, 1]. The video API treats
+        # floating tensors as [-1, 1], so use its NumPy path to avoid normalizing
+        # twice. Images also need NumPy for serving; keep latent tensors intact.
         if (
             isinstance(frames, torch.Tensor)
             and output_type != "latent"
-            and (output_type == "np" or output_key == "image")
+            and (output_type == "np" or output_key in {"image", "video"})
         ):
             frames = frames.float().cpu().numpy()
         return {output_key: frames} if output_key is not None else frames

vllm_omni/diffusion/models/utils.py

@@ -168,6 +168,28 @@ def init_parameters(
         init_parameters(child, dtype, device)
 
 
+def release_module_parameters_to_meta(module: nn.Module) -> None:
+    """Release a fully constructed module's parameter storage to meta.
+
+    The module must finish construction first so initialization code and vLLM
+    weight-loader metadata follow the ordinary CPU path. ``swap_tensors`` keeps
+    each registered Parameter object's identity while dropping only its storage.
+    Persistent buffers intentionally remain materialized.
+    """
+    for param in module.parameters():
+        if param.device.type == "meta":
+            continue
+
+        meta_tensor = torch.empty_like(param, device="meta")
+        if type(param) is nn.Parameter:
+            meta_param = nn.Parameter(meta_tensor, requires_grad=param.requires_grad)
+        else:
+            meta_param = param.__class__.__new__(param.__class__, meta_tensor)
+            meta_param.requires_grad_(param.requires_grad)
+        meta_param.__dict__.update(getattr(param, "__dict__", {}))
+        torch.utils.swap_tensors(param, meta_param)
+
+
 def init_parameters_preserving_vllm_linear(
     module: nn.Module,
     dtype: torch.dtype | None,

vllm_omni/engine/arg_utils.py

@@ -198,6 +198,7 @@ class OmniEngineArgs(EngineArgs):
     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
     # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything.
@@ -430,6 +431,7 @@ class OmniEngineArgs(EngineArgs):
             async_chunk=self.async_chunk,
             session_mode=self.session_mode,
             retains_state_across_chunks=self.retains_state_across_chunks,
+            supports_running_prefix_cache_reset=self.supports_running_prefix_cache_reset,
             use_v2_model_runner=self.use_v2_model_runner,
             supports_native_mrv2_data_plane=self.supports_native_mrv2_data_plane,
             active_stream_window=self.active_stream_window,
@@ -558,6 +560,7 @@ class OrchestratorArgs:
     # === Diffusion model config ===
     num_gpus: int | None = None
     model_class_name: str | None = None
+    hsdp_weight_load_strategy: str | None = None
     diffusion_load_format: str | None = None
     lora_path: list[str] | None = None
     lora_backend: str | None = None

vllm_omni/engine/orchestrator.py

@@ -898,7 +898,7 @@ class OrchestratorBase:
                     req_state is None
                     or req_state.upstream_first_audio
                     or req_state.pending_upstream_first_audio is not None
-                    or stage_id + 1 > req_state.final_stage_id
+                    or (stage_id if final_output else stage_id + 1) > req_state.final_stage_id
                 ):
                     continue
                 audio = mm.get("model_outputs")
@@ -935,7 +935,7 @@ class OrchestratorBase:
         if pending is None or self.request_states.get(req_state.request_id) is not req_state:
             return
         source_stage, first_output = pending
-        codec_stage = source_stage + 1
+        codec_stage = source_stage if self.stage_pools[source_stage].final_output else source_stage + 1
         pool = self.stage_pools[codec_stage]
         replica_id = pool.get_bound_replica_id(req_state.request_id)
         if replica_id is None or req_state.request_id not in pool.output_processor.request_states:

vllm_omni/engine/stage_engine_core_proc.py

@@ -87,7 +87,10 @@ def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: An
     model_runner = getattr(worker, "model_runner", None)
     model_state = getattr(model_runner, "model_state", None)
     model = getattr(model_runner, "model", None)
-    if getattr(model, "first_frame_decoder", None) is None or not hasattr(model_state, "set_first_audio_sink"):
+    decodes_audio = getattr(model, "first_frame_decoder", None) is not None or (
+        getattr(model, "stream_decoder", None) is not None and bool(getattr(model, "stream_first_audio", False))
+    )
+    if not decodes_audio or not hasattr(model_state, "set_first_audio_sink"):
         return False
     from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
 

vllm_omni/entrypoints/async_omni.py

@@ -128,6 +128,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         reasoning_ended: bool | None = None,
         reasoning_parser_kwargs: dict[str, Any] | None = None,
         arrival_time: float | None = None,
+        submitted: asyncio.Future[None] | None = None,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         """Generate outputs for the given prompt(s) asynchronously.
 
@@ -296,6 +297,9 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
             if admitting:
                 await self._release_generate_admission()
                 admitting = False
+            # Callers wait on this before abort so Stage0 observes the sender cache.
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
             # Refresh gauges on arrival.
             self._publish_request_gauges(len(self.request_states))
 

vllm_omni/entrypoints/cli/serve.py

@@ -560,6 +560,15 @@ class OmniServeCommand(CLISubcommand):
                 "custom_pipeline, dummy, or diffusers for the HF diffusers adapter."
             ),
         )
+        omni_config_group.add_argument(
+            "--hsdp-weight-load-strategy",
+            choices=["full", "pre_sharded"],
+            default=None,
+            help=(
+                "HSDP checkpoint loading strategy: 'full' loads then shards; "
+                "'pre_sharded' shards meta parameters first and reads rank-local HF safetensors slices."
+            ),
+        )
         omni_config_group.add_argument(
             "--lora-path",
             type=str,

vllm_omni/entrypoints/openai/api_server.py

@@ -249,6 +249,33 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
     await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
 
 
+def _own_reuseport_socket(sock: socket.socket, client_count: int) -> socket.socket:
+    """Give this API server its own SO_REUSEPORT listener on the shared address.
+
+    All API servers accepting from the one inherited socket lets whichever
+    wakes first take a whole burst of connections (the event loop accepts
+    in a loop), and keep-alive then pins that imbalance for the connections'
+    lifetime. Separate reuseport listeners let the kernel spread connections
+    by hash. The inherited socket is bound but never listened on, so it
+    receives nothing.
+    """
+    if client_count <= 1:
+        return sock
+    if sock.family not in (socket.AF_INET, socket.AF_INET6):
+        return sock
+    try:
+        if not sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT):
+            return sock
+        own = socket.socket(family=sock.family, type=socket.SOCK_STREAM)
+        own.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        own.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
+        own.bind(sock.getsockname()[:2])
+    except OSError as e:
+        logger.warning("Per-server reuseport listener unavailable (%s); sharing the inherited socket", e)
+        return sock
+    return own
+
+
 def run_omni_api_server_worker_proc(
     listen_address: str,
     sock: socket.socket,
@@ -278,6 +305,7 @@ def run_omni_api_server_worker_proc(
 
     set_process_title("APIServer", str(client_index))
     decorate_logs("APIServer", skip_if_decorated=True)
+    sock = _own_reuseport_socket(sock, int(omni_client_config["client_count"]))
     uvloop.run(
         omni_run_server_worker(
             listen_address,
@@ -512,7 +540,7 @@ async def build_async_omni(
     # Ensures everything is shutdown and cleaned up on error/exit
     async with build_async_omni_from_stage_config(
         args,
-        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
+        disable_frontend_multiprocessing=bool(disable_frontend_multiprocessing),
         client_config=client_config,
     ) as async_omni:
         yield async_omni

vllm_omni/entrypoints/openai/audio_utils_mixin.py

@@ -46,6 +46,15 @@ except ImportError:
 logger = init_logger(__name__)
 
 
+def _float32_to_pcm16_bytes(audio: np.ndarray) -> bytes:
+    """PCM_16 bytes of float32 samples, bit-identical to libsndfile's RAW writer.
+
+    libsndfile (1.2) converts with ``floor(x * 0x8000)`` clipped to int16.
+    """
+    scaled = np.floor(audio.astype(np.float64) * 32768.0)
+    return np.clip(scaled, -32768, 32767).astype("<i2").tobytes()
+
+
 class AudioMixin:
     """Mixin class to add audio-related utilities."""
 
@@ -93,9 +102,15 @@ class AudioMixin:
 
         soundfile_format, media_type, kwargs = supported_formats[response_format]
 
-        with BytesIO() as buffer:
-            soundfile.write(buffer, audio_tensor, sample_rate, format=soundfile_format, **kwargs)
-            audio_data = buffer.getvalue()
+        audio_data: bytes | str
+        if response_format == "pcm" and isinstance(audio_tensor, np.ndarray) and audio_tensor.dtype == np.float32:
+            # Through BytesIO, soundfile makes a Python callback per block;
+            # streamed chunks made that the API workers' largest cost.
+            audio_data = _float32_to_pcm16_bytes(audio_tensor)
+        else:
+            with BytesIO() as buffer:
+                soundfile.write(buffer, audio_tensor, sample_rate, format=soundfile_format, **kwargs)
+                audio_data = buffer.getvalue()
 
         if base64_encode:
             import base64

vllm_omni/entrypoints/openai/serving_video_stream.py

@@ -66,29 +66,10 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase):
     ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
         prewarmed = prewarmed_frames or {}
         if frame_indices is None:
-            frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed)
-        user_content: list[dict] = []
-        for index in frame_indices:
-            frame_b64 = frame_buffer[index]
-            cached = prewarmed.get(frame_b64)
-            if cached is not None:
-                # The shared selector excludes failed decodes before this point.
-                assert isinstance(cached, tuple)
-                pil, pil_uuid = cached
-                user_content.append(
-                    {
-                        "type": "image_pil",
-                        "image_pil": pil,
-                        "uuid": pil_uuid,
-                    }
-                )
-            else:
-                user_content.append(
-                    {
-                        "type": "image_url",
-                        "image_url": {"url": f"data:image/jpeg;base64,{frame_b64}"},
-                    }
-                )
+            frame_indices = self._select_frame_indices(config, frame_buffer, prewarmed)
+        frames = [frame_buffer[index] for index in frame_indices]
+
+        user_content: list[dict] = self._build_frame_image_parts(frames, prewarmed)
 
         if len(audio_buffer) > 0:
             wav_b64 = self._pcm_to_wav_b64(bytes(audio_buffer))
@@ -107,14 +88,7 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase):
 
         user_message: dict[str, Any] = {"role": "user", "content": user_content}
 
-        messages: list[dict[str, Any]] = []
-        if config.system_prompt:
-            messages.append({"role": "system", "content": config.system_prompt})
-
-        recent_history = message_history[-2:] if len(message_history) > 2 else message_history
-        for hist_msg in recent_history:
-            messages.append(self._text_only_message(hist_msg))
-
+        messages = self._history_prefix_messages(config, message_history)
         messages.append(user_message)
 
         return messages, user_message

vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py

@@ -397,34 +397,60 @@ class Qwen3TTSAdapter(ARTTSAdapter):
 
         return params
 
+    def _load_text_tokenizer(self) -> None:
+        server = self.ctx.server
+        if server._tts_tokenizer is not None:
+            return
+        import transformers
+        from transformers import AutoTokenizer
+
+        # Same construction as the engine's text tokenizer
+        # (Qwen3TTSPromptEmbedsBuilder.get_text_tokenizer), so the ids computed
+        # here can be handed to the engine.
+        kwargs: dict[str, Any] = dict(trust_remote_code=True, use_fast=True, padding_side="left")
+        if transformers.__version__ < "5":
+            kwargs["fix_mistral_regex"] = True
+        assert self.ctx.engine_client is not None
+        model_name = self.ctx.engine_client.model_config.model
+        server._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, **kwargs)
+
     def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int:
         """Estimate prompt length so the placeholder matches model-side embeddings."""
         try:
             from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+                PRECOMPUTED_TEXT_IDS_KEY,
                 Qwen3TTSPromptEmbedsBuilder,
+                build_assistant_text,
             )
 
             server = self.ctx.server
-            if server._tts_tokenizer is None:
-                from transformers import AutoTokenizer
-
-                model_name = self.ctx.engine_client.model_config.model
-                server._tts_tokenizer = AutoTokenizer.from_pretrained(
-                    model_name,
-                    trust_remote_code=True,
-                    padding_side="left",
-                )
+            self._load_text_tokenizer()
+            assert self.ctx.engine_client is not None
             hf_config = self.ctx.engine_client.model_config.hf_config
             talker_config = hf_config.talker_config
             task_type = (tts_params.get("task_type") or ["CustomVoice"])[0]
-            return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+            text = (tts_params.get("text") or [""])[0]
+            assistant_text = build_assistant_text(text) if isinstance(text, str) else None
+            assistant_ids: list[list[int]] = []
+
+            def tokenize_prompt(t: str) -> list[int]:
+                ids = server._tts_tokenizer(t, padding=False)["input_ids"]
+                if t == assistant_text:
+                    assistant_ids.append(ids)
+                return ids
+
+            prompt_len = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
                 additional_information=tts_params,
                 task_type=task_type,
-                tokenize_prompt=lambda t: server._tts_tokenizer(t, padding=False)["input_ids"],
+                tokenize_prompt=tokenize_prompt,
                 codec_language_id=getattr(talker_config, "codec_language_id", None),
                 spk_is_dialect=getattr(talker_config, "spk_is_dialect", None),
                 estimate_ref_code_len=self._estimate_ref_code_len,
             )
+            if assistant_ids:
+                # Spares the engine tokenizing the same text on its step thread.
+                tts_params[PRECOMPUTED_TEXT_IDS_KEY] = [assistant_ids[0]]
+            return prompt_len
         except Exception as e:
             logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e)
             return 2048