ci: run focused vLLM 0.30.0 weekly baseline

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/cuda/rebase-pipeline.yml

@@ -1,35 +1,46 @@
+# Focused weekly comparison against unchanged Omni main and official vLLM 0.30.0.
 steps:
-  - label: ":docker: Build image"
-    key: image-build
-    commands:
-      - "bash .buildkite/cuda/scripts/build-ci-image.sh"
-    agents:
-      queue: "cpu_queue_premerge"
-
-  # L2 Test
-  - label: "Upload Ready Pipeline"
-    depends_on: image-build
-    key: upload-ready-pipeline
-    commands:
-      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-ready.yml
-    agents:
-      queue: "cpu_queue_premerge"
-
-  # L3 Test
-  - label: "Upload Merge Pipeline"
-    depends_on: image-build
-    key: upload-merge-pipeline
-    commands:
-      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-merge.yml
-    agents:
-      queue: "cpu_queue_premerge"
-
-  # L4 Test
-  - label: "Upload Nightly Pipeline"
-    depends_on: image-build
-    key: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
-    commands:
-      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-nightly.yml
-    agents:
-      queue: "cpu_queue_premerge"
+- label: ':docker: Build image'
+  key: image-build
+  commands:
+  - bash .buildkite/cuda/scripts/build-ci-image.sh
+  agents:
+    queue: cpu_queue_premerge
+- label: Upload selected weekly baseline
+  depends_on: image-build
+  key: upload-weekly-pipeline
+  commands:
+  - |
+    python3 - <<'PYBASELINE'
+    import tempfile
+    from pathlib import Path
+    import yaml
+    labels = set(['TTS · Audex · CosyVoice · Qwen3-TTS · L4', 'TTS · H100', 'Diffusion · Bagel · H100 · Single-GPU', 'Diffusion · SenseNova-U1 · H100 · Single-GPU', 'Diffusion · DreamZero · H100 · 2-GPU', 'Diffusion · H100 · 4-GPU', 'Diffusion · MAGI-2 Preview Serving Smoke · H100 · 4-GPU'])
+    source = yaml.safe_load(Path('.buildkite/cuda/test-weekly.yml').read_text())
+    selected = set()
+    guard = 'python3 -c \'import vllm, torch, flashinfer, vllm._C_stable_libtorch; assert vllm.__version__ == \'"\'"\'0.30.0\'"\'"\'; assert torch.cuda.is_available(); print(\'"\'"\'Verified baseline runtime:\'"\'"\', vllm.__version__, torch.__version__)\''
+    def select(steps):
+        kept = []
+        for step in steps:
+            if not isinstance(step, dict):
+                continue
+            if 'steps' in step:
+                children = select(step['steps'])
+                if children:
+                    kept.append(dict(step, steps=children))
+            elif step.get('label') in labels:
+                item = dict(step)
+                commands = item['commands']
+                item['commands'] = [guard] + ([commands] if isinstance(commands, str) else commands)
+                kept.append(item)
+                selected.add(item['label'])
+        return kept
+    source['steps'] = select(source['steps'])
+    assert selected == labels, (selected, labels)
+    path = Path(tempfile.gettempdir()) / 'vllm-omni-selected-weekly-baseline.yml'
+    path.write_text(yaml.safe_dump(source, sort_keys=False, allow_unicode=True))
+    print('Selected baseline suites:', sorted(selected), flush=True)
+    PYBASELINE
+    python3 .buildkite/common/scripts/upload_pipeline.py --upload --all /tmp/vllm-omni-selected-weekly-baseline.yml
+  agents:
+    queue: cpu_queue_premerge

docker/Dockerfile.ci

@@ -36,7 +36,8 @@ RUN apt-get update && \
 # for CI
 COPY pyproject.toml setup.py ./
 COPY requirements/ requirements/
-ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
+# Keep source-built dependency versions independent of the Omni CI version.
+ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
 # Editable install: deps + project metadata are cached in this layer, while
 # the package itself is linked to APP_DIR (populated later by COPY --link).
 # Do NOT use a non-editable install here — that copies an empty stub into

docs/configuration/environment_variables.md

@@ -214,7 +214,7 @@ collection, examples, bug reports, or logs.
 
 ## Model-specific variables
 
-The inventory includes 66 variables read by a single model or pipeline family. They are
+The inventory includes 93 variables read by a single model or pipeline family. They are
 not listed as public usage options here because doing so would turn implementation
 escape hatches into an accidental compatibility contract.
 
@@ -223,10 +223,10 @@ Every audited model-specific name has a migration disposition in the
 
 | Disposition | Count | Required outcome |
 | --- | ---: | --- |
-| Promote | 38 | Move a stable setting into typed stage or model configuration. |
+| Promote | 64 | Move a stable setting into typed stage or model configuration. |
 | Request scope | 5 | Move request-varying behavior into a declared request-option schema. |
 | External | 0 | Retain only when a supported third-party library owns the contract. |
-| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. |
+| Internalize | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. |
 | Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. |
 
 The disposition is a migration target, not a statement that the existing

docs/configuration/stage_configs.md

@@ -514,3 +514,50 @@ for that comparison; this switch only controls automatic MPS management.
 
 MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on
 the intended GPU, and warm the complete pipeline before measuring performance.
+
+## MOSS-TTS Local 1.5 with Model Runner V2
+
+`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using
+the shared runtime introduced for Qwen3-TTS:
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
+```
+
+This profile inherits the batching, codec graph buckets and 1-frame/15-frame
+chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local
+Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp`
+capabilities while retaining its V1 `talker_mtp` implementation, sampling
+defaults and explicit request-seed handling. Codec chunks use the native data
+plane; the internal Talker retains its final-only orchestrator output policy.
+
+On CUDA, the profile bounds stage-0 prefill work to 512 tokens per iteration.
+The codec retains its own CUDA Graph capture/replay, including batch buckets
+through 64, but disables Inductor compilation (`compilation_config.mode: 0`)
+of the stateful decoder to avoid lengthy compilation at startup. This is
+independent of `enforce_eager`; the profile keeps `enforce_eager: false`.
+It explicitly selects `cudagraph_mode: FULL`, which works without Inductor;
+the default `FULL_AND_PIECEWISE` would otherwise normalize to `NONE` and clear
+the codec's capture buckets when compilation is disabled.
+
+For sustained C128 serving on a large-memory CUDA GPU, select the
+separate `moss_tts_local_mrv2_high_concurrency.yaml` profile. It uses 128
+stream slots per stage, a 32 GiB Talker KV budget, and the codec's
+`triton_slot` backend with Inductor compilation and codec-owned CUDA graphs.
+The bounded KV budget leaves room for codec state and graphs; it does not
+guarantee that 128 maximum-length prompts fit simultaneously. See the
+[MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention)
+for activation, backend comparisons, memory requirements and benchmark commands.
+
+Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1.
+NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable
+MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio;
+set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and
+`VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM.
+
+Event-driven orchestration remains independently selectable with
+`VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical
+when comparing these modes. Model-runner selection does not change the
+orchestration default or enable experimental reference encoding, chunk ramps,
+generation-output draining or MPS.

recipes/OpenMOSS/MOSS-TTS.md

@@ -35,6 +35,10 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses
 audio-reference generation. Reference transcripts must match the reference
 audio; they are not style instructions.
 
+For the optional CUDA MRV2 runner, see the
+[Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2).
+The default Local deployment continues to use V1.
+
 ## References
 
 - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
@@ -152,3 +156,296 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 - No `ref_audio` required or accepted for MOSS-SoundEffect.
 - Input field maps to the `ambient_sound` parameter in the upstream processor.
 - Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
+
+## Local 1.5 MRV2 and slot attention
+
+`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with
+an explicit deploy profile. The default Local profile continues to use V1.
+Both profiles below preserve 1-frame initial and 15-frame steady codec chunks
+and the model's sampling defaults.
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
+```
+
+This C64 profile bounds Talker prefill to 512 tokens and retains codec-owned
+CUDA graphs without compiling the codec with Inductor.
+
+For sustained high concurrency on a large-memory CUDA GPU, use:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+The high-concurrency profile places both stages on one GPU, sets each stage's
+capacity to 128 and fixes the Talker KV budget at 32 GiB. The Talker uses
+mixed FULL CUDA graphs with both its token budget and capture limit set to
+512, so new prompt chunks remain within graph coverage. It also enables
+`hf_overrides.mrv2_gpu_slot_state`: Local hidden states, audio codes and
+continuation control stay in GPU request slots, gathered through MRV2's
+batch mapping. Prefill conditioning and explicit per-request sampling seeds
+retain their original behavior. Published audio codes own their storage, so
+slot reuse cannot overwrite an in-flight output. Set the override to `false`
+to compare with the generic Omni model state. The profile selects
+`codec_attention_backend: triton_slot`, Inductor mode 3 with combo kernels
+disabled, and codec CUDA graph buckets through 128. The configuration was
+validated on one H200. The earlier capacity-256 variant used about 110 GiB of
+sampled peak GPU memory, including loading and warmup. Reduce the stage capacities,
+graph buckets and coalescing target together when adapting it to smaller GPUs. Cold codec
+compilation can take several minutes; subsequent starts can reuse the AOT
+cache. Graph capture is still performed at startup.
+
+The high-concurrency profile also coalesces ready codec streams with
+`connectors.shm.extra.generation_min_batch_size: 32` and
+`generation_max_wait_ms: 12`. The target controls when to stop waiting; it
+neither caps execution at 32 nor guarantees every codec group has 32 rows.
+First and final chunks share the bounded window at high concurrency. When
+fewer than 32 requests remain, dispatch is immediate. Pending outputs are
+retired before waiting, with the original deadline preserved, so waiting does
+not prevent the previous batch from releasing its in-flight state. Cancellation
+and input notifications keep their ordinary scheduler bookkeeping. This option
+requires a stateful native MRV2 generation stage with TP1/PP1. Set the wait to
+`0` to disable coalescing; other profiles keep the immediate-dispatch default.
+
+The high-concurrency profile enables two additional Local model-state
+optimizations under the CUDA stage-0 `hf_overrides`, alongside
+`mrv2_gpu_slot_state: true`:
+
+- `mrv2_batch_prefill: true` reuses the batch's text embeddings and combines
+  CPU reference-code slices and their destination positions into one pinned
+  upload. Host staging is reused only after the upload event completes;
+  GPU-resident references keep the original path.
+- `mrv2_direct_tokens: true` returns the already determined text token after
+  the Local audio/stop decision. It avoids constructing and sampling the
+  full text vocabulary. Audio-code and binary-stop sampling are unchanged.
+  Requests requiring distribution metadata or token constraints, and
+  unsupported execution modes, retain the normal sampler.
+
+Other profiles leave both flags disabled. Set both flags explicitly when
+comparing their combined effect. Deployment
+inheritance replaces a stage's `hf_overrides` mapping, so an overlay must also
+retain `mrv2_gpu_slot_state: true`. Measure first-packet latency as well as
+throughput: faster prompt preparation can change competition between the
+Talker and codec sharing the GPU.
+
+For the single-H200 capacity-128 throughput baseline, keep both switches on.
+In complete Seed-TTS EN1088 runs at client concurrency 128, four combined-switch
+rounds pooled 344.11 audio-s/s versus 335.78 for batch prefill alone; direct
+tokens alone had no reliable gain. The combined path had 685 ms mean first
+audio versus 679 ms across six control rounds, with lower mean completion time.
+A fixed-seed 128-row audio sample showed mean Whisper WER of 3.68% versus
+3.78% for the control. This sample does not establish speaker similarity or
+exclude small quality changes. At concurrency 256, the combined path had
+raised mean first-audio latency by about 0.33 s, so other profiles leave the
+switches off until their latency and quality are checked.
+
+The optional connector setting `generation_coalescing_policy: idle_wait`
+waits for an inbox notification when no stream is runnable, all receivers
+are parked and registered, and no output needs retirement. The first ready
+arrival starts a fresh coalescing window, preserving the batch budget.
+Control messages, including cancellations, can wait up to two windows
+instead of one (24 ms at the 12 ms setting). The default policy remains
+`fixed`; neither policy changes the global orchestration default.
+
+The codec backends differ in state access:
+
+- `sdpa` uses PyTorch attention after gathering and updating the ring cache.
+- `triton` replaces the attention calculation, retaining the ring-cache
+  gather/copy and explicit mask construction.
+- `triton_slot` writes surviving K/V directly into request slots, attends
+  directly to the ring, and advances the active slot offsets in three
+  ordered kernels. Graph-padding rows do not advance persistent state.
+
+The slot path preserves the existing chunk-complete ring semantics, including
+retaining the final cache-capacity tokens when a chunk exceeds the ring.
+It does not introduce another cache owner or change request-slot lifetime.
+Zero-length padding rows skip the attention computation and emit zeros.
+
+The high-concurrency profile includes intermediate codec batch buckets 6, 12
+and 24. To compare with power-of-two buckets, change the stage-1
+`cudagraph_capture_sizes` to `[1, 2, 4, 8, 16, 32, 64, 128]`.
+Keep the maximum bucket equal to the state capacity to retain terminal-tail
+coalescing. Smaller buckets reduce padding work but require more graphs;
+measure complete serving runs before selecting them for a deployment.
+
+For a codec output-path comparison, set
+`connectors.shm.extra.codec_gpu_stream_output: 0` to select the synchronous
+codec output path; `1` selects GPU output snapshots on a private codec
+stream. Both configurations use MRV2 and keep streaming audio responses.
+The option changes input ownership, metadata staging and snapshot handling
+as well as transfer scheduling, so its timing difference is not just D2H
+copy time.
+
+Those chunk-complete semantics truncate the causal window: every ring holds
+exactly `context` entries and a whole chunk is written before attention runs,
+so token `i` of a `T`-token chunk sees `capacity - T + i + 1` keys. With the
+tokenizer-v2 decoder's per-layer contexts (10/10/8/4/2/1 s, i.e. 125/250/400/
+400/400/400 tokens) a 15-frame chunk is 480 tokens at the last transformer, so
+its first 80 tokens attend to nothing and no chunk sees the previous one there.
+Padding a terminal tail to 15 frames therefore decodes it differently from an
+exact-length decode; the difference is deterministic (identical in fp32 and
+fp64), not bf16 noise. The connector option `codec_ring_headroom: 1` sizes each
+ring as `context + max_chunk_frames * tokens_per_frame` (including ramp shapes); chunked streaming
+then reproduces whole-sequence decoding bit-for-bit in fp64 and padded tails
+equal exact tails. It increases state/graph memory and capture time; historical C256 experiments
+observed about a 2% throughput cost. It is off by default and its perceptual
+quality benefit has not been established. Terminal tails are padded
+into the regular 15-frame graph bucket in either mode.
+The slot kernel is specific to the CUDA tokenizer-v2 decoder; other paths
+retain their existing attention implementation. Event-driven orchestration
+remains independently selectable and its default is unchanged.
+
+### Optional progressive chunks
+
+Set `codec_chunk_ramp` under `connectors.shm.extra` to insert smaller chunks
+before steady decoding, for example:
+
+```yaml
+connectors:
+  shm:
+    extra:
+      codec_chunk_ramp: [1, 4, 15]
+```
+
+The ramp's first entry overrides `initial_codec_chunk_frames`. After the last
+entry, chunks use `codec_chunk_frames`. Each request advances independently;
+the final partial chunk is flushed and an empty terminal carries no codec
+tokens. The codec captures every ramp length, sizes the maximum execution
+step for the largest entry, and uses the ramp's first entry for its dedicated
+first-chunk graph. Additional shapes increase startup and graph memory.
+
+Ramps are opt-in. They can reduce early playback underrun while increasing
+codec calls, first-audio latency or total generation time. The measured C128
+throughput presets retain 1→15; their published figures do not establish an
+additional benefit from enabling ramps. The processor and graph-shape tests
+run without GPU or model weights:
+
+```bash
+CUDA_VISIBLE_DEVICES= PYTHONPATH=. python -m pytest -q \
+  tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py \
+  -m 'core_model and cpu' --run-level core_model
+```
+
+### Low-latency and reference-encoding options
+
+`moss_tts_local_mrv2_low_latency.yaml` keeps both stage capacities at 128 and
+uses prefill 2048, GPU slot state, batch prefill, direct tokens and eager MTP.
+Its codec first-chunk fast path uses dedicated graphs and a stream handoff
+before the regular decoder reuses the request slot. A gate limits contention
+with regular codec work. Slot waits are bounded at 30 seconds. A failed or
+stalled decode raises an error rather than replaying a partially advanced
+slot; a failed output enqueue retains owned PCM for regular delivery.
+Closing the decoder rejects new jobs. Its regular-batch limit of 16 is internal dispatch
+policy; it does not change the client concurrency or stage capacities.
+
+Reference encoding runs in the API layer, independently of MRV2. Enable
+reference graphs explicitly with `VLLM_OMNI_MOSS_REF_GRAPHS=1`; the default
+eager encoder is retained when the option is absent. The graph path uses the
+loaded tokenizer's encoder/quantizer, length buckets and optional compilation
+(`VLLM_OMNI_MOSS_REF_COMPILE=0` disables compilation). It also uses windowed
+attention; `VLLM_OMNI_MOSS_REF_ATTN=sdpa` retains the original attention.
+Compilation and attention changes need not produce bit-identical codes.
+
+For multiple API processes, `VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR` enables shared
+reference-code storage and, by default, a single encoder host with four
+workers. Use a dedicated directory per service, checkpoint and encoding
+configuration. Workers have their own graph resources; all graph captures
+complete before serving begins. `VLLM_OMNI_MOSS_REF_SHARED_ENCODER=0` retains
+separate encoders while sharing codes. More workers/graphs consume memory;
+they do not imply more GPUs. Example after selecting an available GPU and
+configuring any operator-managed MPS daemon:
+
+```bash
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
+VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 \
+VLLM_OMNI_MOSS_REF_GRAPHS=1 \
+VLLM_OMNI_MOSS_REF_ENCODER_WORKERS=4 \
+VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS=0 \
+VLLM_OMNI_MOSS_REF_INFLIGHT=4 \
+VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS=2 \
+VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/moss-local-service \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --api-server-count 4 \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+This example is a configurable serving profile, not the complete launch
+configuration of a historical benchmark. In particular, the low-latency YAML
+uses synchronous stage-1 scheduling; the retained source46 experiment used
+asynchronous stage-1 scheduling, FP8 backbone, MPS and additional encoder
+environment settings. Do not infer a throughput result from the YAML alone.
+
+### Cold versus hot reference measurements
+
+A cold reference misses the encoded-reference cache: the request pays for
+reference parsing/preparation and GPU encoding before speech generation.
+A hot reference reuses reference codes but still synthesizes new speech;
+generated audio is not cached. Pre-registering a reference moves encoding
+into registration and can make subsequent synthesis hot, but registration
+latency is still part of the first-use cost.
+
+Model startup/compilation is separate from reference coldness. EN1088 contains
+repeated references, so its first pass is not uniformly cold. For an all-unique
+cold pass, use its 666 first-occurrence references, keep order, warm execution
+with disjoint references, then repeat the same 666 requests. A second pass
+through independent API-local caches is not guaranteed hot; verify cache
+coverage before comparing it with a shared-cache result. Report requests/s
+alongside audio-s/s and retain failed, empty, long and near-silent outputs.
+
+### Reproduce the serving benchmark
+
+Use the complete Seed-TTS English test set, including reference audio and text.
+Set `SEED_TTS_DATA` to the directory containing `en/meta.lst` and its 1088
+entries. Run the native benchmark after the server is ready:
+
+```bash
+export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000
+export VLLM_OMNI_BENCH_AUDIO_CHANNELS=2
+export SEED_TTS_WER_EVAL=0
+
+for concurrency in 128; do
+  for phase in warm r1 r2; do
+    vllm bench serve --omni \
+      --model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
+      --backend openai-audio-speech --endpoint /v1/audio/speech \
+      --dataset-name seed-tts --dataset-path "$SEED_TTS_DATA" \
+      --seed-tts-locale en --disable-shuffle --num-prompts 1088 \
+      --num-warmups 0 --ready-check-timeout-sec 0 \
+      --output-len 256 --max-concurrency "$concurrency" \
+      --request-rate inf --seed 42 \
+      --extra-body '{"task_type":"Base","max_new_tokens":256}' \
+      --save-result --save-detailed --result-dir results/moss-local \
+      --result-filename "c${concurrency}-${phase}.json"
+  done
+done
+```
+
+Discard the complete `warm` pass and combine measured runs as total generated
+audio seconds divided by total benchmark duration. Require 1088 successes,
+zero failures and 1088 nonempty-audio metric samples in each run. The output
+cap matches the benchmark protocol; success alone does not establish speech
+quality or that every sentence ended before the cap. Dataset seed 42 does
+not fix an independent sampling seed for every request.
+
+For an attention-only comparison, copy the high-concurrency YAML beside the
+original to preserve relative `base_config` resolution. Change its
+codec `compilation_config.mode` to `0` while retaining `cudagraph_mode: FULL`,
+and compare `codec_attention_backend: triton` against `triton_slot`. Keep all
+other settings, warmup and client concurrency identical. Restart the server
+between configurations and use separate result directories. Comparing the
+compiled slot profile to uncompiled Triton includes both changes.
+
+Kernel and MHA regression tests require CUDA, compatible vLLM/Triton packages,
+and no model weights:
+
+```bash
+python -m pytest -q tests/model_executor/models/moss_tts/test_slot_attention.py \
+  tests/model_executor/models/moss_tts/test_streaming_attention.py \
+  -m 'core_model and cuda' --run-level=core_model
+```

tests/config/test_config_factory.py

@@ -2928,9 +2928,17 @@ class TestPlatformOverrides:
 
         assert deploy.stages[0].engine_extras["block_size"] == 128
 
-    @pytest.mark.parametrize("platform", ["cuda", "npu"])
-    def test_recommended_native_runner_platform_defaults(self, platform):
-        filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
+    @pytest.mark.parametrize(
+        "filename,pipeline_key",
+        [
+            ("qwen3_tts_mrv2.yaml", "qwen3_tts"),
+            ("moss_tts_local_mrv2.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"),
+        ],
+    )
+    @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+    def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key):
         deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
         deploy = _apply_platform_overrides(deploy, platform=platform)
         expect_v2 = platform == "cuda"
@@ -2946,6 +2954,60 @@ class TestPlatformOverrides:
             # v2 only engages the native plane on stages declaring support.
             assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)
 
+    def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name))
+        candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml"))
+        base = load_deploy_config(base_path)
+        candidate = load_deploy_config(candidate_path)
+        assert base.model_runner == "v1"
+        assert candidate.stages == base.stages
+        assert candidate.connectors == base.connectors
+        assert candidate.async_chunk == base.async_chunk
+        cuda = _apply_platform_overrides(candidate, platform="cuda")
+        stages = merge_pipeline_deploy(pipeline, cuda)
+        assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512
+        codec_args = stages[1].yaml_engine_args
+        assert codec_args["max_num_seqs"] == 64
+        assert codec_args["enforce_eager"] is False
+        assert codec_args["compilation_config"] == {
+            "mode": 0,
+            "cudagraph_mode": "FULL",
+            "cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64],
+            "cudagraph_num_of_warmups": 1,
+        }
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            # Platform resolution mutates the deployment; each server loads
+            # a fresh profile before applying its own platform overrides.
+            fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform)
+            original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert fallback.stages == original.stages
+        for variant in ("moss_tts_delay", "moss_tts_realtime"):
+            assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages)
+
+    def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml"))
+        deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda")
+        talker, codec = merge_pipeline_deploy(pipeline, deploy)
+        assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True
+        assert talker.yaml_engine_args["max_num_batched_tokens"] == 512
+        assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3
+        assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot"
+        compilation = codec.yaml_engine_args["compilation_config"]
+        assert compilation["mode"] == 3 and compilation["backend"] == "inductor"
+        assert compilation["cudagraph_mode"] == "FULL"
+        assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+        assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False}
+        base_path = Path(get_deploy_config_path("moss_tts_local.yaml"))
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            actual = _apply_platform_overrides(load_deploy_config(path), platform=platform)
+            expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert actual.stages == expected.stages
+
     @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)])
     def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native):
         from unittest.mock import Mock

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 32,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 77,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 95,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
         EnvironmentVariableCategory.INTERNAL: 7,
     }
@@ -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: 48,
+        ModelEnvironmentVariableDisposition.PROMOTE: 66,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

tests/core/sched/test_generation_batch_coalescing.py

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded readiness coalescing must not allocate or consume request state."""
+
+import queue
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+from vllm_omni.outputs import OmniConnectorOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch):
+    events = [notification("0"), notification("1", "2", "3")]
+    s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])])
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.006)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012])
+
+
+@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"])
+def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    if pending == "in_flight":
+        s.requests["0"].num_in_flight_tokens = 1
+    elif pending == "admission":
+        s.requests["0"].status = RequestStatus.WAITING
+    elif pending == "registration":
+        s.input_coordinator.pending_chunk_registrations = [object()]
+    elif pending == "terminal":
+        s._pending_data_plane_terminal_req_ids = {"0"}
+    else:
+        s.requests.clear()
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == 10
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_idle_wait_without_arrivals_is_bounded(monkeypatch):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == pytest.approx(10.012)
+
+
+def notification(*ids, terminal=()):
+    return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal))
+
+
+class Inbox:
+    def __init__(self, clock, initial=(), delayed=()):
+        self.clock = clock
+        self.initial = list(initial)
+        self.delayed = list(delayed)
+        self.waits = []
+
+    def get_nowait(self):
+        if self.initial:
+            return self.initial.pop(0)
+        raise queue.Empty
+
+    def get(self, timeout):
+        self.waits.append(timeout)
+        if not self.delayed or self.delayed[0][0] > timeout:
+            self.clock[0] += timeout
+            raise queue.Empty
+        delay, value = self.delayed.pop(0)
+        self.clock[0] += delay
+        return value
+
+
+def scheduler(monkeypatch, initial=(), delayed=()):
+    clock = [10.0]
+    monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0])
+    s = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    s._generation_max_wait_s = 0.012
+    s._generation_min_batch_size = 4
+    s._pause_state = PauseState.UNPAUSED
+    s.requests = {
+        str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8)
+    }
+    s.running = list(s.requests.values())
+    s.input_coordinator = SimpleNamespace(
+        requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set()
+    )
+    s._latest_omni_connector_output = None
+    s._omni_connector_output_inbox = Inbox(clock, initial, delayed)
+    return s, clock
+
+
+def test_arrivals_form_target_batch_in_order(monkeypatch):
+    events = [notification("0"), notification("1", "2"), notification("3")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.005)
+    assert not s.input_coordinator.requests_with_ready_chunks
+    assert all(r.num_in_flight_tokens == 0 for r in s.requests.values())
+
+
+def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch):
+    events = [notification("0", "gone", "1", "2"), notification("0")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])])
+    s.requests["1"].is_finished = lambda: True
+    s.requests["2"].is_finished = lambda: True
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009])
+
+
+@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"])
+def test_immediate_dispatch(monkeypatch, reason):
+    event = notification("0")
+    s, clock = scheduler(monkeypatch, [event])
+    if reason == "enough":
+        event.chunk_ready_req_ids.update(["1", "2", "3"])
+    elif reason == "empty":
+        event.chunk_ready_req_ids.clear()
+    elif reason == "low":
+        s.requests = {"0": s.requests["0"]}
+    elif reason == "paused":
+        s._pause_state = PauseState.PAUSED_ALL
+    elif reason == "off":
+        s._generation_max_wait_s = 0
+    assert s._drain_omni_connector_outputs() == [event]
+    assert not s._omni_connector_output_inbox.waits
+    assert clock[0] == 10
+
+
+def test_terminal_arrival_keeps_original_deadline(monkeypatch):
+    events = [notification("0"), notification("1", terminal=["1"])]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_pending_ready_and_legacy_notification_preserved(monkeypatch):
+    s, _ = scheduler(monkeypatch)
+    s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"])
+    last = notification("3")
+    s._latest_omni_connector_output = last
+    assert s._drain_omni_connector_outputs() == [last]
+    assert s._latest_omni_connector_output is None
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))])
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.012)
+    s._omni_connector_output_inbox.initial.append(notification("3"))
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.024)
+
+
+def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker):
+    event = notification("0", "cancelled")
+    event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}}
+    s, _ = scheduler(monkeypatch, [event])
+    s._generation_max_wait_s = 0
+    s.waiting = []
+    s.input_coordinator._async_chunk = True
+    s.input_coordinator.update_request_metadata = mocker.Mock()
+    s.input_coordinator.process_pending_chunks = mocker.Mock()
+    s._consume_pending_connector_output("generation")
+    s.input_coordinator.update_request_metadata.assert_called_once_with(
+        s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation"
+    )
+    s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set())
+
+
+@pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))])
+def test_invalid_config_rejected(monkeypatch, target, wait):
+    from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler
+
+    def initialize(s, *args, **kwargs):
+        s.max_num_running_reqs = 256
+        s.vllm_config = SimpleNamespace(
+            model_config=SimpleNamespace(
+                retains_state_across_chunks=True,
+                stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}},
+            )
+        )
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", initialize)
+    monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None)
+    with pytest.raises(ValueError):
+        OmniGenerationScheduler()
+
+
+@pytest.mark.parametrize("terminal", [False, True])
+def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal):
+    event = notification("0", terminal=["0"] if terminal else [])
+    s, clock = scheduler(monkeypatch, [event])
+    s.running = s.running[1:]
+    assert s._drain_omni_connector_outputs() == [event]
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch):
+    events = [notification("0"), notification("1")]
+    s, clock = scheduler(monkeypatch, events[:1])
+    s.requests["7"].num_in_flight_tokens = 1
+    assert s._drain_omni_connector_outputs() == events[:1]
+    assert s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+    s.input_coordinator.requests_with_ready_chunks.add("0")
+    # Simulate EngineCore retiring the previous batch after five milliseconds.
+    clock[0] += 0.005
+    s.requests["7"].num_in_flight_tokens = 0
+    s._omni_connector_output_inbox.initial.append(events[1])
+    assert s._drain_omni_connector_outputs() == events[1:]
+    assert not s._generation_defer_batch
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.007])
+    assert clock[0] == pytest.approx(10.012)
+    assert s._generation_batch_deadline is None
+
+
+def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")])
+    s.requests["7"].num_in_flight_tokens = 1
+    s._generation_batch_deadline = clock[0] - 0.001
+    assert len(s._drain_omni_connector_outputs()) == 1
+    assert not s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker):
+    from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler
+
+    request = _chunk_request("pending")
+    s = _make_generation_scheduler(request)
+    s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True)
+    allocate = mocker.spy(s.kv_cache_manager, "allocate_slots")
+    output = s.schedule()
+    assert output.total_num_scheduled_tokens == 0
+    assert not output.num_scheduled_tokens
+    allocate.assert_not_called()
+    assert list(s.waiting) == [request]
+    assert request.num_computed_tokens == request.num_in_flight_tokens == 0
+    assert s.chunk_transfer_adapter.restore_called

tests/core/sched/test_generation_scheduler_init.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise generation constructor policy after actual deploy/platform resolution."""
+
+from types import SimpleNamespace
+
+import pytest
+
+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
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def construct_scheduler(monkeypatch, mocker):
+    # Only upstream cache allocation and connector construction are omitted.
+    # Keep the actual Omni I/O initialization, native decision and constructor.
+    mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter")
+
+    def upstream_init(self, config):
+        self.vllm_config = config
+        self.max_num_running_reqs = config.scheduler_config.max_num_seqs
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init)
+    return OmniGenerationScheduler
+
+
+def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            use_v2_model_runner=native,
+            supports_native_mrv2_data_plane=True,
+            retains_state_across_chunks=stateful,
+            async_chunk=True,
+            stage_id=1,
+            stage_connector_config={"extra": extras or {}},
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp),
+        scheduler_config=SimpleNamespace(max_num_seqs=capacity),
+    )
+
+
+@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"])
+@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+def test_moss_profile_generation_constructor_after_platform_resolution(
+    construct_scheduler, monkeypatch, mocker, profile, platform
+):
+    # merge_pipeline_deploy owns platform resolution. Mock the detected device
+    # instead of applying an override first and then resolving the worker's
+    # actual platform a second time.
+    monkeypatch.setattr(current_omni_platform, "device_name", platform)
+    deploy = load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml"))
+    pipeline = resolve_pipeline_config("moss_tts_local")
+    codec = merge_pipeline_deploy(pipeline, deploy)[1]
+    extras = deploy.connectors["shm"]["extra"]
+    config = _config(
+        native=codec.yaml_engine_args["use_v2_model_runner"],
+        stateful=pipeline.stages[1].retains_state_across_chunks,
+        capacity=codec.yaml_engine_args["max_num_seqs"],
+        extras=extras,
+    )
+    warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning")
+    scheduler = construct_scheduler(config)
+    # The raw extras persist through fallback; the live constructor must cope.
+    assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12
+    if platform == "cuda":
+        assert scheduler._native_data_plane and scheduler.input_coordinator is not None
+        assert scheduler._generation_min_batch_size == 32
+        assert scheduler._generation_max_wait_s == 0.012
+        assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0)
+        warning.assert_not_called()
+        if profile == "low_latency":
+            assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1
+    else:
+        assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None
+        assert scheduler._generation_min_batch_size == 1
+        assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0
+        warning.assert_called_once()
+
+
+@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)])
+def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp):
+    config = _config(
+        stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12}
+    )
+    with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"):
+        construct_scheduler(config)
+
+
+@pytest.mark.parametrize(
+    "extras,message",
+    [
+        ({"generation_min_batch_size": 0}, "generation_min_batch_size"),
+        ({"generation_min_batch_size": 129}, "generation_min_batch_size"),
+        ({"generation_max_wait_ms": -1}, "generation_max_wait_ms"),
+        ({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"),
+        ({"generation_max_regular_batch": -1}, "generation_max_regular_batch"),
+        ({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"),
+    ],
+)
+def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message):
+    with pytest.raises(ValueError, match=message):
+        construct_scheduler(_config(native=False, extras=extras))

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -126,6 +126,11 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
     sched.max_num_running_reqs = 8
     sched._native_data_plane = False
     sched._retains_state_across_chunks = False
+    # Unset MagicMock attributes are truthy. Keep optional scheduling modes
+    # disabled so this fixture actually schedules its one-shot request.
+    sched._generation_defer_batch = False
+    sched._generation_first_chunk_only = False
+    sched._generation_max_regular_batch = 0
     sched._pause_state = PauseState.UNPAUSED
     sched.requests = {request.request_id: request}
     sched.running = [request]

tests/engine/test_stage_engine_core_proc.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import patch
 
 import pytest
@@ -25,11 +26,15 @@ def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder(
     expected,
 ):
     executor = UniProcExecutor.__new__(UniProcExecutor)
+    executor.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)
+    )
     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
+    runner._omni_data_plane = None
     assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected
     assert runner.model_state.set_first_audio_sink.call_count == int(expected)
 
@@ -54,3 +59,64 @@ def test_preprocess_add_request_preserves_omni_fields():
     assert current_wave == 3
     assert result.external_req_id == "external"
     assert result.additional_information == {"conditioning": "payload"}
+
+
+def test_codec_uses_generic_first_audio_binding(monkeypatch):
+    import queue
+
+    import torch
+
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+    from vllm_omni.engine import stage_engine_core_proc as module
+
+    calls: dict[str, Any] = {}
+
+    def hook(*args):
+        return False
+
+    def bind(sink):
+        calls["sink"] = sink
+        return hook
+
+    plane = SimpleNamespace(set_first_chunk_hook=lambda value: calls.update(hook=value))
+    model = SimpleNamespace(bind_first_chunk_fast_path=bind)
+    runner = SimpleNamespace(model=model, get_model=lambda: model, _omni_data_plane=plane)
+    executor = SimpleNamespace(
+        vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
+        driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
+    )
+    monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
+    outputs: queue.Queue = queue.Queue()
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=2)})
+    assert module._bind_first_audio_sink(executor, outputs, scheduler)
+    assert calls["hook"] is hook
+    calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
+    client, batch = outputs.get_nowait()
+    assert client == 2 and FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
+
+
+def test_first_audio_binding_preserves_talker_marker(monkeypatch):
+    import queue
+
+    import torch
+
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+    from vllm_omni.engine import stage_engine_core_proc as module
+
+    calls: dict[str, Any] = {}
+    runner = SimpleNamespace(
+        model=SimpleNamespace(first_frame_decoder=object()),
+        model_state=SimpleNamespace(set_first_audio_sink=lambda sink: calls.update(sink=sink)),
+    )
+    executor = SimpleNamespace(
+        vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
+        driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
+    )
+    monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
+    outputs: queue.Queue = queue.Queue()
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=0)})
+    assert module._bind_first_audio_sink(executor, outputs, scheduler)
+    calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
+    assert outputs.get_nowait()[1].outputs[0].multimodal_output[FIRST_AUDIO_KEY]
+    executor.vllm_config.parallel_config.tensor_parallel_size = 2
+    assert not module._bind_first_audio_sink(executor, outputs, scheduler)

tests/entrypoints/openai_api/test_moss_local_prompt.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The Local-v1.5 continuation prompt fast path equals the processor layout."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Processor:
+    """Segment API of the model's processor with a deterministic tokenizer."""
+
+    def __init__(self):
+        self.model_config = SimpleNamespace(
+            n_vq=3,
+            im_start_token_id=1,
+            im_end_token_id=2,
+            audio_start_token_id=3,
+            audio_pad_token_id=1024,
+            audio_assistant_slot_token_id=5,
+        )
+        self.encoded = []
+
+    def _encode_text(self, text):
+        self.encoded.append(text)
+        return [100 + ord(c) % 50 for c in text]
+
+    def _user_prompt_prefix_ids(self):
+        return [1] + self._encode_text("user\n") + self._encode_text("<ref>")
+
+    def _user_prompt_after_reference_ids(self, language, fields):
+        return self._encode_text(f"lang={fields.get('language', language)}")
+
+    def _assistant_prompt_prefix_ids(self):
+        return self._encode_text("</user>") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant")
+
+    def reference(self, text, language, codes):
+        # Mirrors the processor's _build_continuation_codes + _pad (one item).
+        fields = {} if language is None else {"language": language}
+        ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language, fields)
+            + self._encode_text(text)
+            + self._assistant_prompt_prefix_ids()
+            + [3]
+        )
+        text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long)
+        text_rows[:, 0] = torch.tensor(ids)
+        audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long)
+        audio_rows[:, 0] = 5
+        audio_rows[:, 1:] = codes
+        unified = torch.cat([text_rows, audio_rows])
+        return unified[:, 0].tolist(), unified[:, 1:].contiguous()
+
+
+@pytest.mark.parametrize("language", [None, "English"])
+def test_matches_processor_and_caches_constant_segments(language):
+    proc = _Processor()
+    codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32)
+    kwargs = {"text": "ref words target words"}
+    if language is not None:
+        kwargs["language"] = language
+    ids, audio = _local_continuation_prompt(proc, kwargs, [codes])
+    expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long())
+    assert ids == expected_ids
+    assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio)
+
+    proc.encoded.clear()
+    _local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes])
+    assert proc.encoded == ["other"]
+
+
+def test_unsupported_layouts_fall_back():
+    proc = _Processor()
+    codes = torch.zeros((2, 3), dtype=torch.long)
+    assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -691,6 +691,45 @@ def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker)
     audio_tokenizer.to.assert_called_once()
 
 
+@pytest.mark.parametrize(
+    "variant,graphs_env,prepares",
+    [("local", None, False), ("local", "1", True), ("local", "0", False), ("realtime", "1", False)],
+)
+def test_moss_warmup_prepares_the_reference_encoder(variant, graphs_env, prepares, mocker, monkeypatch):
+    if graphs_env is None:
+        monkeypatch.delenv("VLLM_OMNI_MOSS_REF_GRAPHS", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", graphs_env)
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    adapter._moss_variant = variant
+    encoder = mocker.Mock()
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    asyncio.run(adapter.warmup())
+    assert encoder.prepare.called is prepares
+
+
+def test_moss_warmup_failure_does_not_stop_startup(mocker, monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    encoder = mocker.Mock()
+    encoder.prepare.side_effect = RuntimeError("capture failed")
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    asyncio.run(adapter.warmup())
+    encoder.prepare.assert_called_once_with()
+
+
+def test_moss_shared_encoder_startup_failure_stops_api_startup(mocker, monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
+    from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError
+
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    encoder = mocker.Mock()
+    encoder.prepare.side_effect = SharedReferenceEncoderStartupError("still compiling")
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    with pytest.raises(SharedReferenceEncoderStartupError):
+        asyncio.run(adapter.warmup())
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

tests/entrypoints/test_pd_disaggregation.py

@@ -288,6 +288,9 @@ def _setup_log_mocks(monkeypatch):
         def build_and_log_summary(self):
             return "Fake summary"
 
+        def log_timing_summary(self):
+            pass
+
     monkeypatch.setattr(
         "vllm_omni.entrypoints.omni.OrchestratorAggregator",
         _FakeOrchestratorAggregator,

tests/metrics/test_emit_calls.py

@@ -111,7 +111,7 @@ class TestFailureCounterWiring:
         obj.request_states["req-1"] = SimpleNamespace(
             metrics=SimpleNamespace(
                 e2e_done=set(),
-                build_and_log_summary=lambda: None,
+                log_timing_summary=lambda: None,
             ),
             failure_recorded=False,
         )
@@ -127,7 +127,7 @@ class TestFailureCounterWiring:
         obj.request_states["req-1"] = SimpleNamespace(
             metrics=SimpleNamespace(
                 e2e_done=set(),
-                build_and_log_summary=lambda: None,
+                log_timing_summary=lambda: None,
             ),
             failure_recorded=False,
         )

tests/metrics/test_stats.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import pytest
@@ -155,3 +158,39 @@ def test_build_and_log_summary_multiple_requests() -> None:
     r2_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r2")
     assert len(r1_stage_entry["stages"]) == 2
     assert len(r2_stage_entry["stages"]) == 1
+
+
+def test_log_timing_summary_skips_tables_unless_debug(monkeypatch) -> None:
+    from vllm_omni.metrics import stats
+
+    agg = OrchestratorAggregator(num_stages=1, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=0)
+    agg.stage_first_ts[0] = 0.0
+    agg.stage_last_ts[0] = 0.02
+    agg.on_stage_metrics(
+        0,
+        "r1",
+        StageRequestStats(
+            batch_id=1,
+            batch_size=1,
+            num_tokens_in=3,
+            num_tokens_out=3,
+            stage_gen_time_ms=20.0,
+            rx_transfer_bytes=0,
+            rx_decode_time_ms=0.0,
+            rx_in_flight_time_ms=0.0,
+            stage_stats=StageStats(),
+        ),
+    )
+    agg.on_finalize_request(0, "r1", req_start_ts=0.0)
+    tables, lines = [], []
+
+    def record_table(*args, **kwargs):
+        tables.append(args)
+        return ""
+
+    monkeypatch.setattr(stats, "_format_table", record_table)
+    monkeypatch.setattr(stats.logger, "info", lambda msg, *args: lines.append(msg % args))
+    monkeypatch.setattr(stats.logger, "isEnabledFor", lambda level: False)
+    agg.log_timing_summary()
+    assert not tables
+    assert len(lines) == 1 and lines[0].startswith("[OmniTiming] req=r1 total=")

tests/model_executor/models/moss_tts/test_codec_async_output.py

@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Waveform ownership across graph replays and asynchronous output copies."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+
+pytestmark = [pytest.mark.core_model]
+
+
+class CausalCodec(nn.Module):
+    downsample_rate = 1
+
+    def __init__(self, channels):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(()))
+        self.channels = channels
+        self.config = SimpleNamespace(codebook_size=1024)
+
+    def initialize_decoder_state_pool(self, capacity, scratch):
+        self.state = torch.zeros(capacity + scratch, device=self.weight.device)
+
+    def reset_decoder_state_slots(self, slots):
+        self.state[slots] = 0
+
+    def decode_streaming_batch(self, codes, lengths, slots, valid_rows):
+        audio = codes.sum(0).float().cumsum(-1) + self.state[slots, None]
+        self.state[slots] = audio[:, -1]
+        channels = torch.arange(1, self.channels + 1, device=audio.device)
+        return SimpleNamespace(audio=audio[:, None] * channels[None, :, None], audio_lengths=lengths)
+
+
+class ReplayGraph:
+    """A real CUDA graph with deliberately reused output storage."""
+
+    batch_sizes = [1]
+
+    def __init__(self, codec):
+        self.codes = torch.ones((2, 1, 3), device="cuda", dtype=torch.long)
+        self.slots = torch.zeros(1, device="cuda", dtype=torch.long)
+        self.lengths = torch.full((1,), 3, device="cuda", dtype=torch.long)
+        self.valid = torch.ones(1, device="cuda", dtype=torch.bool)
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid)
+        torch.cuda.current_stream().wait_stream(stream)
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self.output = codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid)
+        codec.state.zero_()
+
+    def decode(self, codes, slots, *, allow_frame_padding):
+        if codes.shape[1:] != (1, 3):
+            return None
+        self.codes.copy_(codes)
+        self.slots.copy_(slots)
+        self.graph.replay()
+        return self.output.audio, self.output.audio_lengths, 1
+
+
+def make_decoder(*, v2=True, channels=2, device="cpu", model_type="moss_tts_local", streaming=True):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(n_vq=2, model_type=model_type),
+            async_chunk=streaming,
+            use_v2_model_runner=v2,
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8),
+    )
+    decoder = MossTTSCodecDecoder(vllm_config=config)
+    decoder._codec = CausalCodec(channels).to(device)
+    decoder._n_channels = channels
+    decoder._stream_max_step_frames = 3
+    return decoder
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("ramp,frames", [(None, [1, 15]), ([1, 4, 8, 15], [1, 4, 8, 15]), ([1, 4, 4, 15], [1, 4, 15])])
+def test_ramp_chunks_are_captured_at_exact_frame_counts(ramp, frames):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(n_vq=2, model_type="moss_tts_local"),
+            async_chunk=True,
+            use_v2_model_runner=True,
+            stage_connector_config={
+                "extra": {
+                    "initial_codec_chunk_frames": 1,
+                    "codec_chunk_frames": 15,
+                    "codec_chunk_ramp": ramp,
+                }
+            },
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8),
+    )
+    decoder = MossTTSCodecDecoder(vllm_config=config)
+    assert decoder._streaming_graph_frame_sizes == frames
+    assert decoder._stream_max_step_frames == 15
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "v2,streaming,model_type",
+    [
+        (False, True, "moss_tts_local"),
+        (True, False, "moss_tts_local"),
+        (True, True, "moss_tts_realtime"),
+        (True, True, "moss_tts_local"),
+    ],
+)
+def test_gpu_output_requires_local_streaming_mrv2_and_cuda(v2, streaming, model_type):
+    decoder = make_decoder(v2=v2, streaming=streaming, model_type=model_type)
+    assert decoder._gpu_stream_output == (v2 and streaming and model_type == "moss_tts_local")
+    session = decoder._ensure_stream_session()
+    slot = session.acquire()
+    assert session.step({slot: torch.ones(2, 3)})[slot].device.type == "cpu"
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("v2", [False, True])
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("graph", [False, True])
+@torch.no_grad()
+def test_replays_terminal_tails_and_slot_reuse_preserve_pending_waveforms(v2, channels, graph):
+    from vllm_omni.outputs import OmniModelRunnerOutput
+    from vllm_omni.worker_v2.omni_generation_model_runner import (
+        OmniGenerationAsyncOutput,
+        OmniGenerationModelRunner,
+        _contains_cuda_tensor,
+    )
+
+    decoder = make_decoder(v2=v2, channels=channels, device="cuda")
+    # Real Local-v1.5 weights replace the class's 24 kHz default at load time.
+    decoder._sr_tensor = torch.tensor(48_000, dtype=torch.int32)
+    session = decoder._ensure_stream_session()
+    if graph:
+        session._cudagraph_wrapper = ReplayGraph(decoder._codec)
+    copy_stream = torch.cuda.Stream()
+    pending = []
+    # Do not materialize any output until later forwards have replayed the
+    # graph, reset terminal/cancelled state, and reused those state slots.
+    for round_index in range(4):
+        value = round_index + 1
+        codes = torch.cat(
+            [
+                torch.full((14,), value, device="cuda", dtype=torch.long),
+                torch.full((6,), value + 10, device="cuda", dtype=torch.long),
+            ]
+        )
+        info = [{"request_id": rid, "meta": {"finished": rid == "tail"}} for rid in ["tail", "empty", "live"]]
+        output = decoder.forward(codes, runtime_additional_information=info, seq_token_counts=[14, 0, 6])
+        # The runner may overwrite its input buffer before codec completion.
+        codes.fill_(-100)
+        payload = output.multimodal_outputs
+        assert _contains_cuda_tensor(payload) == v2
+        assert isinstance(payload, PackedOutputSnapshot) == v2
+        if v2:
+            # Exactly one waveform slab, regardless of request/chunk count.
+            assert sum(slab.is_cuda for _, slab in payload._groups) == 1
+            result = OmniModelRunnerOutput(
+                req_ids=["tail", "empty", "live"],
+                req_id_to_index={"tail": 0, "empty": 1, "live": 2},
+                sampled_token_ids=[[], [], []],
+            )
+            pending.append(
+                OmniGenerationAsyncOutput(
+                    model_runner_output=result,
+                    multimodal_outputs=payload,
+                    num_reqs=3,
+                    main_stream=torch.cuda.current_stream(),
+                    copy_stream=copy_stream,
+                )
+            )
+        else:
+            pending.append(OmniGenerationModelRunner._build_pooler_output(output, 3))
+        decoder.on_requests_finished(["live"])
+        assert decoder._stream_req_slots == {}
+        assert not session._leased_slots
+
+    for round_index, item in enumerate(pending):
+        rows = item.get_output().multimodal_outputs if v2 else item
+        for index, (frames, value) in enumerate([(7, round_index + 1), (0, 0), (3, round_index + 11)]):
+            expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)
+            if channels > 1:
+                expected = expected[None] * torch.arange(1, channels + 1)[:, None]
+            torch.testing.assert_close(rows[index]["model_outputs"], expected, rtol=0, atol=0)
+            assert rows[index]["sr"].item() == 48_000
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_owned_input_copy_clamps_entire_ragged_batch():
+    from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm
+
+    decoder = make_decoder(device="cuda", channels=2)
+    codes = torch.tensor([-1, 1024, 1, -10, 2, 1025, 5000, -2], device="cuda")
+    output = decoder.forward(
+        codes,
+        runtime_additional_information=[{"request_id": name} for name in ("a", "empty", "b")],
+        seq_token_counts=[6, 0, 2],
+    )
+    codes.fill_(-100)
+    cpu = _async_copy_mm(output.multimodal_outputs, 0, pin_memory=True)
+    torch.accelerator.synchronize()
+    expected = torch.tensor([[0.0, 1025.0, 2049.0], [0.0, 2050.0, 4098.0]])
+    torch.testing.assert_close(cpu["model_outputs"][0], expected, rtol=0, atol=0)
+    assert cpu["model_outputs"][1].numel() == 0
+    torch.testing.assert_close(cpu["model_outputs"][2], torch.tensor([[1023.0], [2046.0]]), rtol=0, atol=0)
+    decoder.on_requests_finished(["a", "b"])
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_pinned_metadata_reuse_waits_for_dma_and_preserves_mixed_terminals():
+    decoder = make_decoder(device="cuda")
+    session = decoder._ensure_stream_session()
+    # Force reuse every call, with an outstanding transfer on the first call.
+    # This delay tests host-buffer ownership, never performance or overlap.
+    session._metadata_ring = session._metadata_ring[:1]
+    slots = [session.acquire() for _ in range(3)]
+    pending = []
+    for step in range(8):
+        order = slots[step % 3 :] + slots[: step % 3]
+        terminal = {order[1]}
+        torch.cuda._sleep(2_000_000)
+        ids, tails = session._stage_slot_ids(order, terminal)
+        pending.append((ids.clone(), tails.clone(), order, list(terminal)))
+    session.close()
+    for ids, tails, expected_ids, expected_tails in pending:
+        assert ids.cpu().tolist() == expected_ids
+        assert tails.cpu().tolist() == expected_tails
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_mixed_terminal_reset_does_not_reset_live_slot():
+    decoder = make_decoder(device="cuda", channels=1)
+    session = decoder._ensure_stream_session()
+    first, second = session.acquire(), session.acquire()
+    codes = torch.ones(2, 3, device="cuda")
+    initial = session.step({first: codes, second: 2 * codes}, terminal_slots={first})
+    following = session.step({second: codes, first: codes})
+    torch.testing.assert_close(initial[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0)
+    torch.testing.assert_close(following[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0)
+    torch.testing.assert_close(following[second].cpu(), torch.tensor([[14.0, 16.0, 18.0]]), rtol=0, atol=0)
+    session.close()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("explicit_copy_stream", [False, True])
+@torch.no_grad()
+def test_output_copy_waits_for_private_codec_stream(explicit_copy_stream):
+    from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm
+
+    decoder = make_decoder(device="cuda", channels=1)
+    decoder._codec_stream = torch.cuda.Stream()
+    with torch.cuda.stream(decoder._codec_stream):
+        torch.cuda._sleep(20_000_000)
+    codes = torch.ones(6, device="cuda", dtype=torch.long)
+    output = decoder.forward(
+        codes,
+        runtime_additional_information=[{"request_id": "live"}],
+        seq_token_counts=[6],
+    )
+    codes.fill_(-100)
+    payload = output.multimodal_outputs
+    assert isinstance(payload, PackedOutputSnapshot)
+    assert payload.producer_event is not None
+    # No explicit consumer wait here: the shared copier must honor readiness.
+    copy_stream = torch.cuda.Stream()
+    with torch.cuda.stream(copy_stream):
+        cpu = _async_copy_mm(payload, 0, copy_stream=copy_stream if explicit_copy_stream else None, pin_memory=True)
+    decoder.on_requests_finished(["live"])
+    copy_stream.synchronize()
+    torch.testing.assert_close(cpu["model_outputs"][0], torch.tensor([2.0, 4.0, 6.0]), rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("channels", [1, 2])
+@torch.no_grad()
+def test_short_terminal_row_crops_each_channel_into_output_slab(channels):
+    decoder = make_decoder(device="cuda", channels=channels)
+    session = decoder._ensure_stream_session()
+    short, live = session.acquire(), session.acquire()
+    codes = {short: torch.ones(2, 2, device="cuda"), live: 3 * torch.ones(2, 5, device="cuda")}
+    slab = torch.empty(channels * 7, device="cuda")
+    buffers = {
+        short: slab[: channels * 2].view(channels, 2),
+        live: slab[channels * 2 :].view(channels, 5),
+    }
+    result = session.step(codes, terminal_slots={short}, output_buffers=buffers)
+    for slot, frames, value in [(short, 2, 1), (live, 5, 3)]:
+        expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)[None]
+        expected = expected * torch.arange(1, channels + 1)[:, None]
+        assert result[slot].data_ptr() == buffers[slot].data_ptr()
+        torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0)
+    session.close()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("layout", ["adjacent", "strided", "gapped"])
+@torch.no_grad()
+def test_batched_output_copy_preserves_rows_and_slab_ownership(channels, layout):
+    decoder = make_decoder(device="cuda", channels=channels)
+    session = decoder._ensure_stream_session()
+    slots = [session.acquire() for _ in range(4)]
+    slab = torch.full((6 if layout == "gapped" else 4, channels, 6), -1.0, device="cuda")
+    if layout == "gapped":
+        buffers = {slot: slab[row] for row, slot in zip([0, 1, 3, 4], slots, strict=True)}
+        frames = 6
+    elif layout == "adjacent":
+        buffers = {slot: slab.reshape(4, channels * 6)[row].view(channels, 6) for row, slot in enumerate(slots)}
+        frames = 6
+    else:
+        # Interleaved with other requests' output: the short views must not
+        # overwrite the untouched part of another row/channel.
+        buffers = {slot: slab[row, :, :3] for row, slot in enumerate(slots)}
+        frames = 3
+    codes = {slot: torch.full((2, frames), row + 1, device="cuda") for row, slot in enumerate(slots)}
+    result = session.step(codes, output_buffers=buffers)
+    for row, slot in enumerate(slots):
+        expected = 2 * (row + 1) * torch.arange(1, frames + 1, dtype=torch.float32)[None]
+        expected = expected * torch.arange(1, channels + 1)[:, None]
+        assert result[slot].data_ptr() == buffers[slot].data_ptr()
+        torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0)
+    if layout == "strided":
+        assert torch.all(slab[:, :, 3:] == -1)
+    elif layout == "gapped":
+        assert torch.all(slab[[2, 5]] == -1)
+    session.close()

tests/model_executor/models/moss_tts/test_codec_ring_headroom.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Chunked streaming decode must match whole-sequence decode.
+
+``RingKVCache.complete`` writes a whole chunk before attending. With a ring of
+exactly ``context`` entries, a chunk of T tokens evicts up to T keys that its
+own queries should still see; for T > context it evicts part of the chunk
+itself. ``initialize_decoder_state_pool(chunk_frames=...)`` adds one chunk of
+headroom per layer so streaming reproduces the non-streaming decoder.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerModel
+from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import (
+    MossAudioTokenizerConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TRANSFORMER = {
+    "module_type": "Transformer",
+    "d_model": 16,
+    "num_heads": 2,
+    "num_layers": 1,
+    "dim_feedforward": 32,
+    "causal": True,
+    "norm": "layer_norm",
+    "positional_embedding": "rope",
+    "max_period": 10000,
+    "gating": "none",
+    "layer_scale": 1.0,
+    "conv_layout": True,
+}
+
+
+def tiny_codec() -> MossAudioTokenizerModel:
+    # 64 Hz mono input, 4-sample code frames (16 code frames per second).
+    # Decoder: transformer @16 Hz (context 0.5 s = 8 tokens) -> x2 -> transformer
+    # @32 Hz (context 0.25 s = 8 tokens) -> x2 -> 64 Hz. A 6-frame chunk is 12
+    # tokens at the second transformer, more than its 8-token ring.
+    config = MossAudioTokenizerConfig(
+        sampling_rate=64,
+        downsample_rate=4,
+        number_channels=1,
+        enable_channel_interleave=False,
+        encoder_kwargs=[
+            {"module_type": "PatchedPretransform", "patch_size": 4},
+            {**TRANSFORMER, "input_dimension": 4, "output_dimension": 8, "context_duration": 0.5},
+        ],
+        decoder_kwargs=[
+            {**TRANSFORMER, "input_dimension": 8, "output_dimension": 16, "context_duration": 0.5},
+            {"module_type": "PatchedPretransform", "patch_size": 2},
+            {**TRANSFORMER, "input_dimension": 8, "output_dimension": 2, "context_duration": 0.25},
+            {"module_type": "PatchedPretransform", "patch_size": 2},
+        ],
+        quantizer_kwargs={
+            "input_dim": 8,
+            "rvq_dim": 8,
+            "output_dim": 8,
+            "num_quantizers": 2,
+            "codebook_size": 16,
+            "codebook_dim": 4,
+            "quantizer_type": "rlfq",
+        },
+    )
+    torch.manual_seed(0)
+    model = MossAudioTokenizerModel(config).eval()
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.copy_(torch.randn_like(parameter) * 0.5)
+    model.quantizer.build_decode_lut(2, dtype=torch.float32)
+    return model
+
+
+def stream_decode(model, codes, chunk_frames, *, headroom):
+    n_q, _, total = codes.shape
+    model.initialize_decoder_state_pool(1, 0, chunk_frames=chunk_frames if headroom else 0)
+    try:
+        slot = torch.zeros(1, dtype=torch.long)
+        valid = torch.ones(1, dtype=torch.bool)
+        parts = []
+        for start in range(0, total, chunk_frames):
+            chunk = codes[:, :, start : start + chunk_frames]
+            lengths = torch.full((1,), chunk.shape[-1], dtype=torch.long)
+            parts.append(model.decode_streaming_batch(chunk, lengths, slot, valid).audio)
+        return torch.cat(parts, dim=-1)
+    finally:
+        model.close_decoder_state_pool()
+
+
+@torch.no_grad()
+def test_ring_headroom_sizes_each_layer_for_one_chunk():
+    model = tiny_codec()
+    transformers = [m for m in model.decoder if hasattr(m, "tokens_per_input_frame")]
+    assert [m.tokens_per_input_frame for m in transformers] == [1, 2]
+    assert [m.transformer.layers[0].self_attn.context for m in transformers] == [8, 8]
+    model.initialize_decoder_state_pool(1, 0, chunk_frames=6)
+    try:
+        rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers]
+        assert rings == [8 + 6, 8 + 12]
+    finally:
+        model.close_decoder_state_pool()
+    model.initialize_decoder_state_pool(1, 0)
+    try:
+        rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers]
+        assert rings == [8, 8]
+    finally:
+        model.close_decoder_state_pool()
+
+
+@torch.no_grad()
+def test_chunked_streaming_matches_whole_sequence_only_with_headroom():
+    model = tiny_codec()
+    torch.manual_seed(1)
+    codes = torch.randint(0, 16, (2, 1, 18))
+    reference = model._decode_frame(codes).audio
+    with_headroom = stream_decode(model, codes, 6, headroom=True)
+    legacy = stream_decode(model, codes, 6, headroom=False)
+    assert with_headroom.shape == reference.shape == legacy.shape
+    torch.testing.assert_close(with_headroom, reference, rtol=1e-5, atol=1e-5)
+    # The legacy context-sized ring lets the 12-token chunk evict keys inside
+    # its own window, so streaming diverges from whole-sequence decoding.
+    assert not torch.allclose(legacy, reference, rtol=1e-3, atol=1e-3)
+
+
+@torch.no_grad()
+def test_padded_terminal_tail_matches_exact_tail_with_headroom():
+    model = tiny_codec()
+    torch.manual_seed(2)
+    codes = torch.randint(0, 16, (2, 1, 14))
+    history, tail = codes[:, :, :12], codes[:, :, 12:]
+    padded_tail = torch.nn.functional.pad(tail, (0, 4))  # 2 real frames padded to a 6-frame step
+    outputs = {}
+    for name, last in [("exact", tail), ("padded", padded_tail)]:
+        model.initialize_decoder_state_pool(1, 0, chunk_frames=6)
+        try:
+            slot, valid = torch.zeros(1, dtype=torch.long), torch.ones(1, dtype=torch.bool)
+            for chunk in (history[:, :, :6], history[:, :, 6:]):
+                model.decode_streaming_batch(chunk, torch.tensor([6]), slot, valid)
+            audio = model.decode_streaming_batch(last, torch.tensor([last.shape[-1]]), slot, valid).audio
+            outputs[name] = audio[..., : 2 * model.downsample_rate]
+        finally:
+            model.close_decoder_state_pool()
+    torch.testing.assert_close(outputs["padded"], outputs["exact"], rtol=1e-5, atol=1e-5)

tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py

@@ -0,0 +1,547 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The first-chunk fast path emits audio out of band; the main path resumes the stream."""
+
+import queue
+import threading
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from tests.model_executor.models.moss_tts.test_streaming_terminal_batch import session as cpu_session
+from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import MossFirstChunkFastPath, _SlotHandoff
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
+
+pytestmark = pytest.mark.core_model
+
+N_VQ, SAMPLES, CHANNELS = 2, 4, 2
+
+
+class _Session:
+    def __init__(self, capacity):
+        self.free = list(range(capacity))
+        self.reset_waits = []
+
+    def acquire(self, *, oldest=False):
+        return self.free.pop(0 if oldest else -1) if self.free else None
+
+    def order_after_reset(self, slot, stream):
+        self.reset_waits.append(slot)
+
+    def release(self, slot, **kwargs):
+        self.free.append(slot)
+
+
+class _Wrapper:
+    batch_sizes = [1, 2, 4, 8]
+
+    def __init__(self, device):
+        self.state = torch.zeros(16, device=device)
+        self.calls = []
+        self.gate = threading.Event()
+        self.gate.set()
+
+    def decode(self, codes, slots):
+        self.gate.wait()
+        self.calls.append(int(codes.shape[1]))
+        base = codes.sum(0)[:, 0].float() + 100 * slots.float()
+        channel = torch.arange(CHANNELS, device=codes.device).float()
+        audio = base[:, None, None] + channel[None, :, None] + torch.zeros(1, 1, SAMPLES, device=codes.device)
+        self.state[slots] = 1
+        return audio.to(torch.bfloat16), None, int(codes.shape[1])
+
+
+@pytest.fixture
+def cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda")
+
+
+def _fast(device, capacity=4):
+    session, wrapper = _Session(capacity), _Wrapper(device)
+    fast = MossFirstChunkFastPath(
+        session,
+        wrapper,
+        n_vq=N_VQ,
+        frames=1,
+        codebook_size=1024,
+        samples_per_frame=SAMPLES,
+        n_channels=CHANNELS,
+        sample_rate=torch.tensor(24000, dtype=torch.int32),
+        device=device,
+    )
+    return fast, session, wrapper
+
+
+class _CallbackSink:
+    def __init__(self, callback):
+        self.callback = callback
+
+    def prepare(self, request_ids):
+        callback = self.callback
+
+        class Delivery:
+            routes = {rid: (3 if rid == "r0" else 0) for rid in request_ids}
+
+            def __call__(self, ids, rows, sr):
+                for rid, row in zip(ids, rows, strict=True):
+                    callback(self.routes[rid], rid, {"model_outputs": row, "sr": sr})
+
+            def fail(self, ids):
+                pass
+
+        return Delivery()
+
+
+@pytest.mark.cuda
+def test_emits_each_first_chunk_and_hands_slot_to_main_path(cuda):
+    fast, session, wrapper = _fast(cuda)
+    emitted = {}
+    done = threading.Event()
+
+    def sink(client, req_id, payload):
+        emitted[req_id] = (client, payload["model_outputs"], int(payload["sr"]))
+        if len(emitted) == 2:
+            done.set()
+
+    slots: dict[str, int] = {}
+    assert not fast.submit("r0", "k0", [1, 2], slots)  # unbound: main path owns it
+    fast.bind(_CallbackSink(sink))
+    wrapper.gate.clear()
+    assert fast.submit("r0", "k0", torch.tensor([1, 2]), slots)
+    assert fast.submit("r1", "k1", [5, 7], slots)
+    assert not fast.submit("r0", "k0", [1, 2], slots)  # already has a slot
+    assert not fast.submit("r2", "k2", [1, 2, 3, 4], slots)  # not a 1-frame chunk
+    assert slots == {"k0": 0, "k1": 1} and session.free == [2, 3]
+    wrapper.gate.set()
+    assert done.wait(10)
+    fast.close()
+
+    assert wrapper.calls in ([1, 1], [2])
+    for req_id, client, code_sum, slot in (("r0", 3, 3, 0), ("r1", 0, 12, 1)):
+        got_client, wav, sr = emitted[req_id]
+        assert got_client == client and sr == 24000 and wav.device.type == "cpu"
+        expected = (code_sum + 100 * slot + torch.arange(CHANNELS).float())[:, None].expand(CHANNELS, SAMPLES)
+        torch.testing.assert_close(wav, expected)
+    assert sorted(session.reset_waits) == [0, 1]
+
+    # The main path consumes the marker once and is ordered after the decode.
+    assert fast.take_decoded("k0") and not fast.take_decoded("k0")
+    fast.order_after(0)
+    fast.order_after(0)  # no pending handoff: no-op
+    assert wrapper.state[:2].tolist() == [1, 1]
+
+
+@pytest.mark.cuda
+def test_order_after_blocks_until_decode_finishes(cuda):
+    fast, _, wrapper = _fast(cuda)
+    fast.bind(_CallbackSink(lambda *args: None))
+    wrapper.gate.clear()
+    assert fast.submit("r0", "k0", [1, 2], {})
+    finished = threading.Event()
+
+    def wait_for_slot():
+        fast.order_after(0)
+        finished.set()
+
+    waiter = threading.Thread(target=wait_for_slot)
+    waiter.start()
+    assert not finished.wait(0.2)
+    wrapper.gate.set()
+    assert finished.wait(10)
+    waiter.join()
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_no_free_slot_leaves_chunk_to_main_path(cuda):
+    fast, _, _ = _fast(cuda, capacity=0)
+    fast.bind(_CallbackSink(lambda *args: None))
+    slots: dict[str, int] = {}
+    assert not fast.submit("r0", "k0", [1, 2], slots)
+    assert slots == {} and not fast.take_decoded("k0")
+    fast.close()
+
+
+class _FakeFast:
+    frames = 1
+
+    def __init__(self, decoded, unsent=None):
+        self.decoded = set(decoded)
+        self.unsent = dict(unsent or {})
+        self.ordered = []
+        self.forgotten = []
+
+    def take_decoded(self, key):
+        if key in self.decoded:
+            self.decoded.remove(key)
+            return True
+        return False
+
+    def get_request_slot(self, key, slots):
+        return slots.get(key)
+
+    def order_after(self, slot):
+        self.ordered.append(slot)
+
+    def take_unsent_audio(self, key):
+        return self.unsent.pop(key, None)
+
+    def forget(self, key):
+        self.forgotten.append(key)
+
+    def gate(self):
+        self.gated = getattr(self, "gated", 0) + 1
+
+
+def _decoder(s, fast):
+    decoder = object.__new__(MossTTSCodecDecoder)
+    nn.Module.__init__(decoder)
+    decoder._stream_max_step_frames = 15
+    decoder._stream_state_capacity = 8
+    decoder._stream_req_slots = {}
+    decoder._ensure_stream_session = lambda: s
+    decoder._first_chunk_fast_path = fast
+    decoder._stream_session = s
+    decoder._codec_stream = None
+    return decoder
+
+
+@pytest.mark.cpu
+def test_main_path_skips_fast_decoded_chunk_and_resumes_state():
+    s = cpu_session()
+    reference = cpu_session()
+    fast = _FakeFast({"a"})
+    decoder = _decoder(s, fast)
+    # The fast path leased slot 0 and decoded frame 0 of "a" into it.
+    slot = s.acquire(oldest=True)
+    decoder._stream_req_slots["a"] = slot
+    s.step({slot: torch.full((2, 1), 3.0)})
+    ref_slot = reference.acquire()
+    reference.step({ref_slot: torch.full((2, 1), 3.0)})
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False), (1, "b", torch.ones(2, 1), False)])
+    assert 0 not in out and out[1].tolist() == [[2.0]]
+    assert fast.ordered == [slot]
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.ones(2, 3), True)])
+    expected = reference.step({ref_slot: torch.ones(2, 3)})[ref_slot]
+    torch.testing.assert_close(out[0], expected)
+    assert "a" not in decoder._stream_req_slots and fast.ordered == [slot, slot]
+
+
+@pytest.mark.cpu
+def test_abort_orders_release_after_fast_decode():
+    s = cpu_session()
+    fast = _FakeFast({"a"})
+    decoder = _decoder(s, fast)
+    slot = decoder._stream_req_slots["a"] = s.acquire(oldest=True)
+    decoder.on_requests_finished({"a"})
+    assert fast.ordered == [slot] and fast.forgotten == ["a"] and decoder._stream_req_slots == {}
+    assert slot in s._free_stream_slots
+
+
+@pytest.mark.cpu
+def test_main_path_returns_unsent_first_chunk_without_decoding_twice():
+    s = cpu_session()
+    first = torch.full((1, 1), 7.0)
+    fast = _FakeFast({"a"}, {"a": first})
+    decoder = _decoder(s, fast)
+    slot = s.acquire(oldest=True)
+    decoder._stream_req_slots["a"] = slot
+    s.step({slot: torch.full((2, 1), 3.0)})
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False)])
+    torch.testing.assert_close(out[0], first)
+    assert fast.unsent == {} and fast.ordered == [slot]
+
+    # A payload containing a later frame must concatenate it after the
+    # preserved first samples, while advancing state only for the later frame.
+    first = torch.full((1, 1), 7.0)
+    fast.decoded.add("b")
+    fast.unsent["b"] = first
+    slot_b = s.acquire(oldest=True)
+    decoder._stream_req_slots["b"] = slot_b
+    s.step({slot_b: torch.full((2, 1), 3.0)})
+    out = decoder._decode_streaming_batch([(0, "b", torch.full((2, 2), 3.0), False)])
+    torch.testing.assert_close(out[0][..., :1], first)
+    assert out[0].shape[-1] == 2
+
+
+@pytest.mark.cpu
+def test_oldest_acquire_reuses_least_recently_released_slot():
+    s = cpu_session()
+    a, b = s.acquire(), s.acquire()
+    s.release(a)
+    s.release(b)
+    assert s.acquire() == b
+    assert s.acquire(oldest=True) == 7  # never-used slots were freed before any release
+    assert s._free_stream_slots[-1] == a
+
+
+@pytest.mark.cuda
+def test_failed_decode_reports_error_without_replaying_state(cuda):
+    fast, _, wrapper = _fast(cuda)
+
+    def fail(codes, slots):
+        raise RuntimeError("boom")
+
+    wrapper.decode = fail
+    fast.bind(_CallbackSink(lambda *args: pytest.fail("no output expected")))
+    slots: dict[str, int] = {}
+    assert fast.submit("r0", "k0", [1, 2], slots)
+    with pytest.raises(RuntimeError, match="first-chunk decode failed"):
+        fast.order_after(slots["k0"])
+    assert fast.take_decoded("k0")
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_failed_output_handoff_retains_decoded_samples(cuda):
+    fast, _, _ = _fast(cuda)
+
+    def reject(*args):
+        raise RuntimeError("output queue closed")
+
+    fast.bind(_CallbackSink(reject))
+    slots: dict[str, int] = {}
+    assert fast.submit("r0", "k0", [1, 2], slots)
+    fast.order_after(slots["k0"])
+    assert fast.take_decoded("k0")
+    wav = fast.take_unsent_audio("k0")
+    assert wav is not None and wav.shape == (CHANNELS, SAMPLES)
+    assert fast.take_unsent_audio("k0") is None
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_gate_orders_main_stream_after_inflight_fast_decode(cuda):
+    fast, _, _ = _fast(cuda)
+    side = torch.cuda.Stream()
+    with torch.cuda.stream(side):
+        torch.cuda._sleep(200_000_000)
+        pending = torch.cuda.Event()
+        pending.record(side)
+    fast._inflight = pending
+    fast.gate()  # disabled: no wait
+    unordered = torch.cuda.Event()
+    unordered.record()
+    fast._gate_main = True
+    fast.gate()
+    ordered = torch.cuda.Event()
+    ordered.record()
+    unordered.synchronize()
+    assert not pending.query() and not ordered.query()
+    ordered.synchronize()
+    assert pending.query()
+
+
+def _cpu_fast(monkeypatch, *, scheduler=None):
+    # Replace only CUDA allocations, retaining the real lifecycle and worker.
+    import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module
+
+    monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: tensor)
+    monkeypatch.setattr(torch.cuda, "Stream", lambda **kwargs: None)
+    monkeypatch.setattr(torch.cuda, "Event", lambda: None)
+    monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=lambda device: None))
+    session = _Session(2)
+    fast = MossFirstChunkFastPath(
+        session,
+        SimpleNamespace(batch_sizes=[1, 2]),
+        n_vq=2,
+        frames=1,
+        codebook_size=1024,
+        samples_per_frame=4,
+        n_channels=1,
+        sample_rate=torch.tensor(24000),
+        device=torch.device("cpu:0"),
+        handoff_timeout_s=0.01,
+    )
+    outputs: queue.Queue = queue.Queue()
+    scheduler = scheduler or SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    sink = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False)
+    return fast, session, sink, outputs
+
+
+def _cancel_during_admission(fast, session, scheduler, sink):
+    entered, release = threading.Event(), threading.Event()
+
+    class PausedSink:
+        def prepare(self, ids):
+            delivery = sink.prepare(ids)
+            if delivery.routes:
+                entered.set()
+                assert release.wait(5)
+            return delivery
+
+    fast.bind(PausedSink())
+    decoder = _decoder(session, fast)
+    errors = []
+    finished = threading.Event()
+    finishing = threading.Event()
+
+    def submit():
+        try:
+            assert fast.submit("r", "r", [1, 2], decoder._stream_req_slots)
+        except BaseException as error:
+            errors.append(error)
+
+    def cancel():
+        try:
+            finishing.set()
+            decoder.on_requests_finished({"r"})
+            finished.set()
+        except BaseException as error:
+            errors.append(error)
+
+    producer = threading.Thread(target=submit)
+    cleaner = threading.Thread(target=cancel)
+    producer.start()
+    try:
+        assert entered.wait(5)
+        assert decoder._stream_req_slots == {}
+        # Scheduler cancellation retires the route before the runner hook.
+        scheduler.requests.pop("r")
+        cleaner.start()
+        assert finishing.wait(5)
+        assert not finished.wait(0.1)
+        release.set()
+        producer.join(5)
+        cleaner.join(5)
+        assert not producer.is_alive() and not cleaner.is_alive()
+        assert not errors, errors
+        assert finished.is_set()
+        assert decoder._stream_req_slots == {}
+        assert fast._decoded == {} and fast._handoffs == {}
+        assert sorted(session.free) == [0, 1]
+        assert not fast.submit("r", "r", [1, 2], decoder._stream_req_slots)
+
+        # A fresh request can use and return the same capacity immediately.
+        scheduler.requests["reuse"] = SimpleNamespace(client_index=3)
+        release.set()
+        assert fast.submit("reuse", "reuse", [1, 2], decoder._stream_req_slots)
+        scheduler.requests.pop("reuse")
+        decoder.on_requests_finished({"reuse"})
+        assert decoder._stream_req_slots == {} and sorted(session.free) == [0, 1]
+    finally:
+        release.set()
+        producer.join(5)
+        if cleaner.ident is not None:
+            cleaner.join(5)
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_cancel_waits_for_first_chunk_admission_and_reclaims_slot(monkeypatch):
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    fast, session, sink, _ = _cpu_fast(monkeypatch, scheduler=scheduler)
+
+    # Keep the actual queue/worker/handoff lifecycle, omitting CUDA decode.
+    def decode(jobs):
+        for job in jobs:
+            job.handoff.done.set()
+
+    monkeypatch.setattr(fast, "_decode", decode)
+    fast._handoff_timeout_s = 5.0
+    _cancel_during_admission(fast, session, scheduler, sink)
+
+
+@pytest.mark.cuda
+def test_cancel_during_first_chunk_admission_orders_cuda_handoff_and_reuse(cuda):
+    fast, session, _ = _fast(cuda, capacity=2)
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    sink = engine_output_queue_sink(queue.Queue(), scheduler, upstream_first_audio=False)
+    _cancel_during_admission(fast, session, scheduler, sink)
+
+
+@pytest.mark.cpu
+def test_slot_timeout_keeps_handoff_and_prevents_slot_reuse(monkeypatch):
+    fast, _, _, _ = _cpu_fast(monkeypatch)
+    handoff = _SlotHandoff()
+    fast._handoffs[0] = handoff
+    for _ in range(2):
+        with pytest.raises(TimeoutError, match="slot 0"):
+            fast.order_after(0)
+        assert fast._handoffs[0] is handoff
+    handoff.done.set()
+    fast.order_after(0)
+    assert fast._handoffs == {}
+
+
+@pytest.mark.cpu
+def test_closed_decoder_rejects_jobs_and_cannot_restart(monkeypatch):
+    fast, session, sink, _ = _cpu_fast(monkeypatch)
+    fast.bind(sink)
+    fast.close()
+    slots: dict[str, int] = {}
+    assert not fast.submit("r", "external", [1, 2], slots)
+    assert slots == {} and session.free == [0, 1]
+    with pytest.raises(RuntimeError, match="closed"):
+        fast.bind(sink)
+
+
+@pytest.mark.cpu
+def test_missing_route_does_not_claim_slot(monkeypatch):
+    fast, session, sink, outputs = _cpu_fast(monkeypatch, scheduler=SimpleNamespace(requests={}))
+    fast.bind(sink)
+    try:
+        assert not fast.submit("gone", "external", [1, 2], {})
+        assert session.free == [0, 1] and outputs.empty()
+    finally:
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_decode_failure_wakes_current_and_queued_handoffs(monkeypatch):
+    from vllm.v1.engine import FinishReason
+
+    fast, _, sink, outputs = _cpu_fast(monkeypatch)
+    entered, release = threading.Event(), threading.Event()
+
+    def fail(jobs):
+        entered.set()
+        assert release.wait(2)
+        raise RuntimeError("partial replay failed")
+
+    fast._decode = fail
+    fast.bind(sink)
+    slots: dict[str, int] = {}
+    try:
+        assert fast.submit("r", "a", [1, 2], slots)
+        assert entered.wait(2)
+        assert fast.submit("r", "b", [3, 4], slots)
+        release.set()
+        fast._thread.join(2)
+        assert not fast._thread.is_alive()
+        for slot in slots.values():
+            with pytest.raises(RuntimeError, match="decode failed"):
+                fast.order_after(slot)
+        assert all(fast.take_decoded(key) for key in slots)
+        assert not fast.submit("r", "c", [1, 2], slots)
+        assert outputs.qsize() == 2
+        while not outputs.empty():
+            assert outputs.get_nowait()[1].outputs[0].finish_reason == FinishReason.ERROR
+    finally:
+        release.set()
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_platform_failure_wakes_pending_handoff(monkeypatch):
+    import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module
+
+    fast, _, _, _ = _cpu_fast(monkeypatch)
+    fast._handoffs[0] = _SlotHandoff()
+
+    def fail(device):
+        raise RuntimeError("device unavailable")
+
+    monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=fail))
+    fast._run()
+    assert fast._closed
+    with pytest.raises(RuntimeError, match="decode failed"):
+        fast.order_after(0)

tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Per-request code rows keep the runner payload layout with one host copy."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.local_model_state import _CodeRowsSnapshot
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_rows_map_to_batch_positions_and_copy_once():
+    codes = torch.arange(6, dtype=torch.long).reshape(2, 3)
+    snapshot = _CodeRowsSnapshot(codes, [2, 0], 4)
+    assert isinstance(snapshot, PackedOutputSnapshot) and snapshot.producer_event is None
+    rows = snapshot["codes"]["audio"]
+    assert [tuple(row.shape) for row in rows] == [(1, 3), (0, 3), (1, 3), (0, 3)]
+    assert torch.equal(rows[2], codes[:1]) and torch.equal(rows[0], codes[1:])
+
+    copies = []
+
+    def copy(tensor):
+        copies.append(tensor)
+        return tensor.clone()
+
+    host = snapshot.copy_to_cpu(copy)["codes"]["audio"]
+    assert len(copies) == 1 and copies[0] is codes
+    codes.zero_()
+    assert host[2].tolist() == [[0, 1, 2]] and host[0].tolist() == [[3, 4, 5]]
+    assert host[1].numel() == 0 and host[1].dtype == torch.long

tests/model_executor/models/moss_tts/test_local_compiled_sampler.py

@@ -0,0 +1,37 @@
+# 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.tts,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+def test_explicit_generator_bypasses_compiled_audio_sampler():
+    from transformers import GPT2Config
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+        MossTTSLocalDepthTransformer,
+    )
+
+    model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=80, n_head=1, n_inner=160)).cuda().bfloat16().eval()
+    heads = torch.nn.ModuleList([torch.nn.Linear(80, 32) for _ in range(2)]).cuda().bfloat16()
+    embeddings = torch.nn.ModuleList([torch.nn.Embedding(32, 80) for _ in range(2)]).cuda().bfloat16()
+    binary = torch.nn.Linear(80, 2).cuda().bfloat16()
+    hidden = torch.randn(1, 80, device="cuda", dtype=torch.bfloat16)
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    expected = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator)
+
+    def forbidden(*args, **kwargs):
+        raise AssertionError("compiled path selected")
+
+    model._compiled_audio_sampler = forbidden
+    generator.manual_seed(17)
+    actual = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator)
+    assert all(torch.equal(a, b) for a, b in zip(expected, actual))
+    with pytest.raises(AssertionError, match="compiled path selected"):
+        model.generate_frame(hidden, heads, embeddings, binary, n_vq=2)

tests/model_executor/models/moss_tts/test_local_determined_tokens.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Determined text tokens retain the upstream sampling/counting contract."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.states import RequestState
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("seed", [None, 17])
+@pytest.mark.parametrize("temperature", [0.0, 1.7])
+def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed, temperature):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda")
+    reqs = RequestState(5, 32, 16, 0, 8, device)
+    sampler = Sampler(5, 8, device, reqs)
+    for i, length in enumerate([4, 1, 3]):
+        reqs.add_request(str(i), length, [2] * length, 0, 8)
+        slot = reqs.req_id_to_index[str(i)]
+        sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed))
+    reqs.apply_staged_writes()
+    sampler.apply_staged_writes()
+    state = _state(MossLocalModelState, device)
+    state._direct_tokens = True
+    state.model._batch_state = None
+    batch = _batch(device, [2, 4, 3], [1, 2, 1])
+    batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32)
+    batch.cu_num_logits_np = np.arange(4, dtype=np.int32)
+    batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device)
+    batch.expanded_idx_mapping = batch.idx_mapping
+    batch.expanded_local_pos = torch.zeros(3, dtype=torch.int32, device=device)
+    batch.logits_indices = batch.query_start_loc[1:].long() - 1
+    batch.positions = torch.arange(batch.num_tokens, device=device)
+    state.model._batch_should_continue = torch.tensor([False, True, True], device=device)
+    with torch.inference_mode():
+        logits = state.model.compute_logits(torch.ones(3, 4, device=device))
+        expected = sampler(logits, batch)
+        actual = state.sample_determined_tokens(batch, sampler)
+    assert actual is not None
+    for field in ["sampled_token_ids", "num_sampled", "num_rejected"]:
+        torch.testing.assert_close(getattr(actual, field), getattr(expected, field), rtol=0, atol=0)
+    assert actual.num_sampled.tolist() == [1, 0, 1]
+    snapshot = actual.sampled_token_ids.clone()
+    state.model._batch_should_continue.logical_not_()
+    torch.testing.assert_close(actual.sampled_token_ids, snapshot, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize(
+    "constraint",
+    [
+        {"min_tokens": 2, "stop_token_ids": [1]},
+        {"allowed_token_ids": [1, 2]},
+        {"logit_bias": {1: -10.0}},
+    ],
+)
+def test_real_sampler_constraints_disable_determined_tokens(constraint):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda")
+    reqs = RequestState(5, 32, 16, 0, 8, device)
+    sampler = Sampler(5, 8, device, reqs)
+    reqs.add_request("constrained", 1, [2], 0, 8)
+    slot = reqs.req_id_to_index["constrained"]
+    sampler.add_request(slot, 1, SamplingParams(**constraint))
+    state = _state(MossLocalModelState, device)
+    state._direct_tokens = True
+    # Use the actual upstream request registration, including its combined
+    # min-token/allowed-ID/bias flag, rather than mocking that flag ourselves.
+    assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([slot])), sampler) is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"])
+def test_distribution_features_fall_back_before_using_gpu(feature):
+    sampler = object.__new__(Sampler)
+    sampler.compute_nans = feature == "nans"
+    sampler.return_sampling_mask = feature == "mask"
+    sampler.trace_replay_state = object() if feature == "trace" else None
+    sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None
+    sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"]))
+    sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"]))
+    sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"])))
+    sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"]))
+    state = object.__new__(MossLocalModelState)
+    state._direct_tokens = True
+    assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([0])), sampler) is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("feature", ["grammar", "sharding", "pipeline", "speculative", "empty", "declined"])
+def test_runner_retains_upstream_path_for_unsupported_execution(mocker, feature):
+    hook = mocker.Mock(return_value=None)
+
+    class State:
+        def sample_determined_tokens(self, batch, sampler):
+            return hook(batch, sampler)
+
+    runner = object.__new__(OmniARModelRunner)
+    runner.model_state = State()
+    runner.sampler = object()
+    runner.batch_sharder = object() if feature == "sharding" else None
+    runner.pp_handler = object() if feature == "pipeline" else None
+    batch = SimpleNamespace(num_draft_tokens=int(feature == "speculative"), num_reqs=0 if feature == "empty" else 1)
+    grammar = object() if feature == "grammar" else None
+    fallback = mocker.patch.object(OmniGPUModelRunner, "sample", return_value="upstream")
+    hidden = object()
+    assert runner.sample(hidden, batch, grammar) == "upstream"
+    fallback.assert_called_once_with(hidden, batch, grammar)
+    assert hook.call_count == int(feature == "declined")

tests/model_executor/models/moss_tts/test_local_eager_mtp.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Eager MTP emits the same frames one step earlier with identical backbone inputs."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _admit, _batch, _state, _step
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def device(request):
+    if request.param == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device(request.param)
+
+
+def _run(state, device, schedule, stop_step):
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[3], prompt[0] = 4, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    embeds = []
+    frames: dict[int, list[torch.Tensor]] = {3: [], 0: []}
+    for index, (slots, counts) in enumerate(schedule):
+        state.model.force_stop = index == stop_step
+        embed, payload, _ = _step(state, _batch(device, slots, counts), req_states)
+        embeds.append(embed)
+        rows = payload.get("codes", {}).get("audio", [])
+        for slot, row in zip(slots, rows):
+            if row.numel() and bool(row.ne(state.model.audio_pad_token_id).any()):
+                frames[slot].append((index, row.clone()))
+        for slot, count in zip(slots, counts):
+            computed[slot] += count
+    return embeds, frames
+
+
+def test_eager_frames_match_canonical_one_step_earlier(device):
+    schedule = [([3, 0], [2, 2]), ([3, 0], [2, 1]), ([3, 0], [1, 1]), ([0, 3], [1, 1]), ([3, 0], [1, 1])]
+    canonical, eager = (_state(MossLocalModelState, device) for _ in range(2))
+    eager._local_eager_mtp = True
+    for state in (canonical, eager):
+        _admit(state, 3, "long", 17)
+        _admit(state, 0, "short", 17)
+    # Stop both streams at the canonical step that would draw the 4th step's frames.
+    c_embeds, c_frames = _run(canonical, device, schedule, stop_step=4)
+    e_embeds, e_frames = _run(eager, device, schedule, stop_step=3)
+
+    for c, e in zip(c_embeds, e_embeds):
+        torch.testing.assert_close(c, e, rtol=0, atol=0)
+    for slot in (3, 0):
+        assert [f for _, f in c_frames[slot]] and len(c_frames[slot]) == len(e_frames[slot])
+        for (ci, cf), (ei, ef) in zip(c_frames[slot], e_frames[slot]):
+            assert ei == ci - 1
+            torch.testing.assert_close(cf, ef, rtol=0, atol=0)
+
+
+def test_stopped_stream_does_not_emit_again(device):
+    state = _state(MossLocalModelState, device)
+    state._local_eager_mtp = True
+    _admit(state, 0, "short", 17)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0] = 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    emitted = []
+    for index, count in enumerate([2, 1, 1, 1]):
+        state.model.force_stop = index == 1
+        _, payload, _ = _step(state, _batch(device, [0], [count]), req_states)
+        row = payload["codes"]["audio"][0]
+        emitted.append(bool(row.ne(state.model.audio_pad_token_id).any()))
+        computed[0] += count
+        state.model.force_stop = False
+    # Frame at the prefill step, stop at step 1, then asynchronous overrun steps.
+    assert emitted == [True, False, False, False]
+
+
+def test_unmarked_decode_row_ignores_stale_slot_state(device):
+    canonical, eager = (_state(MossLocalModelState, device) for _ in range(2))
+    eager._local_eager_mtp = True
+    for state in (canonical, eager):
+        _admit(state, 0, "short", 17)
+    # A decode row whose completing prefill was not observed in this state.
+    eager._keep_pool[0] = False
+    eager._eager_emb_pool[0].fill_(3.0)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0], computed[0] = 2, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    for state in (canonical, eager):
+        state.intermediate_buffer.buffers[0]["audio_state"] = {"is_stopping": False}
+    c_embed, c_payload, _ = _step(canonical, _batch(device, [0], [1]), req_states)
+    e_embed, e_payload, _ = _step(eager, _batch(device, [0], [1]), req_states)
+    torch.testing.assert_close(c_embed, e_embed, rtol=0, atol=0)
+    torch.testing.assert_close(c_payload["codes"]["audio"][0], e_payload["codes"]["audio"][0], rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatch, mocker):
+    """The Local config must not enable the parent's incompatible eager hook."""
+    from vllm_omni.worker_v2 import omni_ar_model_runner as runner_module
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+    from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+    from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+
+    device = torch.device("cpu")
+    template = _state(MossLocalModelState, device)
+    model = template.model
+    model.config = SimpleNamespace(mrv2_eager_mtp=True)
+
+    def initialize_parent(state, *args):
+        state.__dict__.update(template.__dict__)
+        state._eager_mtp, state._eager_rows = False, None
+        state._eager_state = EagerMTPState(state)
+
+    monkeypatch.setattr(OmniModelState, "__init__", initialize_parent)
+    state = MossLocalModelState(SimpleNamespace(), model, None, device)
+    _admit(state, 0, "req", 17)
+    batch = _batch(device, [0], [2])
+    batch.req_ids = ["req"]
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0] = 2
+    _step(state, batch, SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed))
+    hook = mocker.spy(state, "run_eager_mtp")
+
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.execute_model_state = SimpleNamespace(
+        input_batch=batch,
+        hidden_states=torch.zeros(2, 4),
+        finished_req_ids=set(),
+        ec_connector_output=None,
+        routed_experts=None,
+    )
+    runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None
+    runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False
+    runner.model_config = SimpleNamespace(async_chunk=False)
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
+    runner.model_state, runner.model = state, model
+    runner.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.zeros(5, 2)),
+        num_computed_tokens=SimpleNamespace(gpu=torch.from_numpy(computed)),
+        prompt_len=SimpleNamespace(np=prompt, gpu=torch.from_numpy(prompt)),
+    )
+    runner.main_stream = runner.output_copy_stream = mocker.Mock()
+    runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = mocker.Mock()
+    runner.sampler = None
+    runner.sample = mocker.Mock(
+        return_value=(
+            SimpleNamespace(sampled_token_ids=torch.tensor([[2]])),
+            torch.ones(1),
+            torch.zeros(1),
+        )
+    )
+    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={}))
+    runner.postprocess_sampled = mocker.Mock()
+    runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None))
+    mock_out = SimpleNamespace(copy_event=None)
+    monkeypatch.setattr(runner_module, "OmniAsyncOutput", mocker.Mock(return_value=mock_out))
+
+    assert runner.sample_tokens(None) is mock_out
+    hook.assert_called_once()
+    assert state._eager_mtp is False and state._eager_rows is None
+    assert state._local_eager_mtp is True

tests/model_executor/models/moss_tts/test_local_model_state.py

@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare slot lifecycle with canonical Local hooks, including output ownership."""
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+from vllm_omni.worker_v2.model_states import init_omni_model_state
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def device(request):
+    if request.param == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device(request.param)
+
+
+def _model(device):
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+    model.model = torch.nn.Module()
+    weight = torch.arange(32, dtype=dtype, device=device).reshape(8, 4) / 16
+    model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=weight)
+    model.hidden_size, model.n_vq = 4, 2
+    model.audio_pad_token_id, model.text_vocab_size = 8, 8
+    model.audio_assistant_slot_token_id, model.im_end_token_id = 2, 3
+    model.talker_mtp_output_key = ("audio_codes", "current")
+    model.talker_mtp_graph_safe = False
+    model.gpu_resident_buffer_keys = {("hidden_states", "last"), ("audio_codes", "current")}
+    model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
+    model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).to(dtype) / 8
+    model.force_stop = False
+
+    def frame(hidden, *_args, generator=None, temperature=None, top_k=None, top_p=None, **_kwargs):
+        assert (temperature, top_k, top_p) == (1.7, 25, 0.8)
+        noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator)
+        codes = (hidden[:, :2].mul(16).long() + noise) % 8
+        return torch.full((hidden.shape[0],), not model.force_stop, device=device), codes
+
+    model.local_transformer = SimpleNamespace(generate_frame=frame)
+    return model
+
+
+def _state(cls, device):
+    state = object.__new__(cls)
+    state.model = _model(device)
+    state.device, state.dtype = device, state.model.model.embed_tokens.weight.dtype
+    state.has_preprocess = state.has_postprocess = state.have_multimodal_outputs = True
+    state.scheduler_config = SimpleNamespace(max_num_seqs=5)
+    state.vllm_config = SimpleNamespace()
+    state.intermediate_buffer = OmniIntermediateBuffer(5)
+    state._static_inputs_embeds = torch.zeros(16, 4, device=device, dtype=state.dtype)
+    state._mtp_input_ids = torch.zeros(5, device=device, dtype=torch.long)
+    state._mtp_offsets = torch.zeros(5, device=device, dtype=torch.long)
+    for name in ("_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step"):
+        setattr(state, name, torch.zeros(5, 4, device=device, dtype=state.dtype))
+    state._mtp_runner = state._mtp_sample_uniforms = None
+    state._mtp_generators = {}
+    state._stream_pos = {}
+    state._stream_decode_event = None
+    state._eager_state = EagerMTPState(state)
+    if cls is MossLocalModelState:
+        state._init_slot_buffers(5, 4, device, state.dtype)
+    return state
+
+
+def _admit(state, slot, name, seed):
+    params = SimpleNamespace(extra_args={"tts_local_seed": seed} if seed is not None else {}, seed=None)
+    with patch.object(DefaultModelState, "add_request", return_value=None):
+        state.add_request(slot, SimpleNamespace(req_id=name, mm_features=[], sampling_params=params))
+    state.intermediate_buffer.buffers[slot]["codes"] = {"ref": torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])}
+
+
+def _batch(device, slots, counts):
+    starts = np.array([0, *np.cumsum(counts)], dtype=np.int32)
+    return SimpleNamespace(
+        idx_mapping_np=np.array(slots),
+        idx_mapping=torch.tensor(slots, device=device),
+        num_reqs=len(slots),
+        num_tokens=int(starts[-1]),
+        num_scheduled_tokens=np.array(counts),
+        query_start_loc_np=starts,
+        query_start_loc=torch.tensor(starts, device=device),
+        # Upstream MRV2 uses int32 token buffers, unlike the MTP graph inputs.
+        input_ids=torch.full((int(starts[-1]),), 2, dtype=torch.int32, device=device),
+    )
+
+
+def _step(state, batch, req_states, dispatcher=None):
+    inputs = {"input_ids": batch.input_ids, "inputs_embeds": state._static_inputs_embeds[: batch.num_tokens]}
+    with (
+        patch("vllm.forward_context.set_forward_context", return_value=nullcontext()),
+        patch(
+            "vllm_omni.model_executor.models.moss_tts.local_model_state.set_forward_context", return_value=nullcontext()
+        ),
+        torch.inference_mode(),
+    ):
+        state.run_preprocess(batch, inputs, req_states, dispatcher)
+        hidden = inputs["inputs_embeds"] + 0.25
+        state.run_postprocess(hidden, batch)
+        _, payload = state.postprocess_model_output(hidden, batch, req_states)
+        last = batch.query_start_loc[1:] - 1
+        tokens = state.model.compute_logits(hidden.index_select(0, last)).argmax(-1)
+        result = inputs["inputs_embeds"].clone(), payload, tokens
+        hidden.fill_(-100)  # next graph replay can overwrite its output immediately
+    return result
+
+
+@pytest.mark.parametrize("seed", [None, 17])
+@pytest.mark.parametrize("batch_prefill", [False, True])
+def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill):
+    reference, candidate = (_state(cls, device) for cls in (OmniModelState, MossLocalModelState))
+    candidate._batch_prefill = batch_prefill
+    for state in (reference, candidate):
+        _admit(state, 3, "long", seed)
+        _admit(state, 0, "short", seed)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[3], prompt[0] = 4, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    saved = []
+    for index, (slots, counts) in enumerate(
+        [
+            ([3, 0], [2, 2]),
+            ([0, 3], [1, 1]),
+            ([3, 0], [1, 1]),
+            ([3, 0], [1, 1]),
+            ([0, 3], [1, 1]),
+            ([3, 0], [1, 1]),
+        ]
+    ):
+        if index == 4:
+            # Cancellation/removal followed by reuse must not retain hidden,
+            # reference offset, control, or the old request's RNG stream.
+            for state in (reference, candidate):
+                state.remove_request("short")
+                assert "short" not in state._mtp_generators
+                _admit(state, 0, "replacement", seed)
+            computed[0], prompt[0] = 0, 1
+        outputs = []
+        for state in (reference, candidate):
+            state.model.force_stop = index == 3
+            torch.manual_seed(91 + index)
+            outputs.append(_step(state, _batch(device, slots, counts), req_states))
+        a, b = outputs
+        torch.testing.assert_close(a[0], b[0], rtol=0, atol=0)
+        torch.testing.assert_close(a[2], b[2], rtol=0, atol=0)
+        if index == 3:
+            assert b[2].tolist() == [3, 3]
+        assert a[1].keys() == b[1].keys()
+        for ar, br in zip(a[1].get("codes", {}).get("audio", []), b[1].get("codes", {}).get("audio", [])):
+            torch.testing.assert_close(ar, br, rtol=0, atol=0)
+            saved.append((br, br.clone()))
+        for slot, count in zip(slots, counts):
+            computed[slot] += count
+            old = reference.intermediate_buffer.buffers[slot]
+            new = candidate.intermediate_buffer.buffers[slot]
+            assert old["ref_offset"] == new["ref_offset"]
+            torch.testing.assert_close(candidate._hidden_pool[slot], old["hidden_states"]["last"], rtol=0, atol=0)
+            assert "hidden_states" not in new and "audio_codes" not in new
+    candidate._codes_pool.fill_(8)
+    candidate._hidden_pool.zero_()
+    for current, snapshot in saved:
+        torch.testing.assert_close(current, snapshot, rtol=0, atol=0)
+
+
+def test_slot_honors_external_stopping_control(device):
+    states = [_state(cls, device) for cls in (OmniModelState, MossLocalModelState)]
+    for state in states:
+        _admit(state, 0, "stop", None)
+        req = SimpleNamespace(prompt_len=np.ones(5), num_computed_tokens=np.zeros(5))
+        _step(state, _batch(device, [0], [1]), req)
+        state.intermediate_buffer.buffers[0]["audio_state"]["is_stopping"] = True
+        req.num_computed_tokens[:] = 1
+        output = _step(state, _batch(device, [0], [1]), req)
+        assert output[2].tolist() == [3]
+        assert output[1]["codes"]["audio"][0].tolist() == [[8, 8]]
+
+
+@pytest.mark.cuda
+def test_slot_mtp_graph_padding_reorder_and_owned_output():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    states = [_state(cls, torch.device("cuda")) for cls in (OmniModelState, MossLocalModelState)]
+    for state in states:
+
+        def frame(hidden, *_args, **_kwargs):
+            return hidden[:, 0] < 3, hidden[:, :2].mul(16).long() % 8
+
+        state.model.local_transformer.generate_frame = frame
+        for slot in [0, 3, 4]:
+            _admit(state, slot, f"req{slot}", None)
+    candidate = states[1]
+    buffers = [
+        candidate._mtp_input_ids[:4],
+        candidate._mtp_input_embeds[:4],
+        candidate._mtp_hidden[:4],
+        candidate._mtp_text_step[:4],
+    ]
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream), torch.inference_mode():
+        for _ in range(3):
+            candidate.model.mtp(*buffers)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph), torch.inference_mode():
+        graph_output = candidate.model.mtp(*buffers)
+
+    def replay(*args, **kwargs):
+        assert all(a.data_ptr() == b.data_ptr() for a, b in zip(args, buffers))
+        assert args[0].shape[0] == 4
+        graph.replay()
+        return graph_output
+
+    candidate._mtp_runner = replay
+    candidate._is_mtp_graph_runner = lambda: True
+    from vllm.config import CUDAGraphMode
+
+    def dispatcher(bsz):
+        return SimpleNamespace(num_tokens=4, cg_mode=CUDAGraphMode.FULL)
+
+    req = SimpleNamespace(prompt_len=np.ones(5, dtype=np.int32), num_computed_tokens=np.zeros(5, dtype=np.int32))
+    retained = []
+    for slots in [[0, 3, 4], [3, 0], [4, 0, 3], [3, 4]]:
+        out = [
+            _step(
+                state,
+                _batch(torch.device("cuda"), slots, [1] * len(slots)),
+                req,
+                dispatcher if state is candidate else None,
+            )
+            for state in states
+        ]
+        torch.testing.assert_close(out[0][0], out[1][0], rtol=0, atol=0)
+        torch.testing.assert_close(out[0][2], out[1][2], rtol=0, atol=0)
+        for a, b in zip(out[0][1].get("codes", {}).get("audio", []), out[1][1].get("codes", {}).get("audio", [])):
+            torch.testing.assert_close(a, b, rtol=0, atol=0)
+            retained.append((b, b.clone()))
+        req.num_computed_tokens[slots] += 1
+    with torch.inference_mode():
+        for tensor in graph_output:
+            tensor.fill_(7)
+    for tensor, copy in retained:
+        torch.testing.assert_close(tensor, copy, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_factory_requires_explicit_class_capability_and_valid_state(monkeypatch):
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    dynamic = MagicMock(has_preprocess=True, create_omni_model_state=None)
+    assert isinstance(init_omni_model_state(None, dynamic, None, torch.device("cpu")), OmniModelState)
+    dynamic.create_mrv2_model_state.assert_not_called()
+
+    class Invalid:
+        has_preprocess = True
+
+        def create_mrv2_model_state(self, *args):
+            return object()
+
+    with pytest.raises(TypeError, match="must return a ModelState"):
+        init_omni_model_state(None, Invalid(), None, torch.device("cpu"))
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("enabled", [False, True])
+def test_local_factory_default_and_option(monkeypatch, enabled):
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    monkeypatch.setattr(MossLocalModelState, "__init__", lambda *args: None)
+    model = _model(torch.device("cpu"))
+    model.config = SimpleNamespace(mrv2_gpu_slot_state=enabled)
+    state = init_omni_model_state(None, model, None, torch.device("cpu"))
+    assert type(state) is (MossLocalModelState if enabled else OmniModelState)

tests/model_executor/models/moss_tts/test_local_prefill_batch.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference batching preserves coordinates and asynchronous source ownership."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("flat", [False, True])
+def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mocker, flat):
+    state = _state(MossLocalModelState, torch.device("cpu"))
+    state._batch_prefill = True
+    codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])
+    for slot, offset in [(3, 1), (0, 2)]:
+        state.intermediate_buffer.buffers[slot] = {
+            "req_id": str(slot),
+            "codes": {"ref": codes.flatten() if flat else codes},
+            "ref_offset": offset,
+        }
+    batch = _batch(torch.device("cpu"), [3, 0], [2, 1])
+    req = SimpleNamespace(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5))
+    embeds = state._static_inputs_embeds[:3]
+    expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[[1, 2, 2]])
+    embed = mocker.spy(state.model.model.embed_tokens, "forward")
+    with torch.inference_mode():
+        state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req)
+    torch.testing.assert_close(embeds, expected, rtol=0, atol=0)
+    assert embed.call_count == 1
+    assert state.intermediate_buffer.buffers[3]["ref_offset"] == 3
+    assert state.intermediate_buffer.buffers[0]["ref_offset"] == 3
+
+
+@pytest.mark.cuda
+def test_pending_reference_uploads_survive_staging_reuse_and_overflow():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    state = _state(MossLocalModelState, torch.device("cuda"))
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    outputs = []
+    with torch.cuda.stream(stream), torch.inference_mode():
+        # Keep the first H2D pending while CPU prepares all subsequent batches.
+        torch.cuda._sleep(100_000_000)
+        for value in range(1, 8):
+            embeds = torch.zeros(4, 4, device="cuda", dtype=state.dtype)
+            codes = torch.full((2, 2), value)
+            state._apply_reference_batch(embeds, [(1, codes)])
+            outputs.append((embeds, value))
+            codes.fill_(-1)
+    stream.synchronize()
+    assert len(state._prefill_staging) <= 3
+    for output, value in outputs:
+        expected = torch.zeros_like(output)
+        expected[1:3] = value / 8
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+def test_batched_real_audio_embedding_matches_scalar_bfloat16_reduction():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    state = _state(MossLocalModelState, torch.device("cuda"))
+    model = state.model
+    model.n_vq = 12
+    model.audio_vocab_size = model.audio_pad_token_id = 1024
+    model._stacked_audio_emb_w = torch.randn(12, 1024, 3072, device="cuda", dtype=torch.bfloat16)
+    model._audio_embed = MossTTSLocalTalkerForGeneration._audio_embed.__get__(model)
+    state._static_inputs_embeds = torch.zeros(512, 3072, device="cuda", dtype=torch.bfloat16)
+    references = [(start, torch.randint(0, 1025, (count, 12))) for start, count in [(0, 169), (201, 11), (220, 257)]]
+    with torch.inference_mode():
+        actual = torch.randn_like(state._static_inputs_embeds)
+        expected = actual.clone()
+        for start, codes in references:
+            expected[start : start + len(codes)] += model._audio_embed(codes.to("cuda"))
+        state._apply_reference_batch(actual, references)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/model_executor/models/moss_tts/test_local_short_attention.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.tts,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+@pytest.mark.parametrize("batch,heads,dim", [(1, 4, 64), (16, 32, 80), (128, 32, 80), (3, 8, 128)])
+def test_all_depth_positions_with_strided_cache_and_dirty_future(batch, heads, dim):
+    from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+    torch.manual_seed(43)
+    query = torch.randn(batch, heads, 1, dim, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(batch, heads, 12, dim, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    for position in range(12):
+        # Non-contiguous prefixes retain the stride of the full 12-slot cache.
+        k, v = key[:, :, : position + 1], value[:, :, : position + 1]
+        before_k, before_v = key.clone(), value.clone()
+        expected = F.scaled_dot_product_attention(query, k, v)
+        actual = local_short_attention(query, k, v)
+        assert torch.equal(key, before_k) and torch.equal(value, before_v)
+        if position == 0:
+            assert torch.equal(actual, v)
+        torch.testing.assert_close(actual, expected, atol=0.015625, rtol=0.01)
+        # Future cache contents must never affect an earlier position.
+        key[:, :, position + 1 :] += 100
+        value[:, :, position + 1 :] -= 100
+        assert torch.equal(local_short_attention(query, k, v), actual)
+        key.copy_(before_k)
+        value.copy_(before_v)
+
+
+def test_compile_and_graph_replay_reads_fresh_inputs():
+    from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+    query = torch.randn(16, 32, 1, 96, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(16, 32, 12, 96, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    compiled = torch.compile(local_short_attention, fullgraph=True)
+    for _ in range(3):
+        compiled(query, key, value)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        result = compiled(query, key, value)
+    for _ in range(3):
+        query.normal_()
+        key.normal_()
+        value.normal_()
+        graph.replay()
+        torch.testing.assert_close(result, F.scaled_dot_product_attention(query, key, value), atol=0.015625, rtol=0.01)
+
+
+def test_depth_block_cache_overwrite_and_compiled_positions(monkeypatch):
+    from transformers import GPT2Config
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+        MossTTSLocalDepthTransformer,
+    )
+
+    monkeypatch.setenv("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "1")
+    model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=320, n_head=4, n_inner=640)).cuda().bfloat16().eval()
+    attn = model.h[0].attn
+    attn.prepare_rope_cache(12, torch.device("cuda"), torch.bfloat16)
+    cache = tuple(torch.empty(4, 4, 12, 80, device="cuda", dtype=torch.bfloat16) for _ in range(2))
+    forward = torch.compile(model._forward_prefix, fullgraph=True, dynamic=True)
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(4, 12, 320, device="cuda", dtype=torch.bfloat16)
+            # The second frame overwrites the previous frame's K/V before use.
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = forward(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=0.03125, rtol=0.02)

tests/model_executor/models/moss_tts/test_reference_encoder.py

@@ -6,6 +6,7 @@ single-flight, and micro-batched encoding."""
 import asyncio
 import threading
 import time
+from tempfile import TemporaryDirectory
 
 import pytest
 import torch
@@ -598,3 +599,239 @@ def test_reference_resampler_is_cached_and_matches_functional_resample():
         assert info.misses == 1 and info.hits == 1
     finally:
         _reference_resampler.cache_clear()
+
+
+# --------------------------------------------------------------------------- #
+# Inline data: URI fast path                                                  #
+# --------------------------------------------------------------------------- #
+
+
+class _DigestAudio(_FakeAudio):
+    """Resolver whose key for data: URIs is the SHA-1 digest, like serving."""
+
+    def __init__(self, *, digest_keys=True):
+        super().__init__()
+        self._digest_keys = digest_keys
+
+    async def resolve(self, ref_str):
+        wav_list, sr, _ = await super().resolve(ref_str)
+        if not self._digest_keys:
+            return wav_list, sr, "rk:" + ref_str
+        import hashlib
+
+        return wav_list, sr, hashlib.sha1(ref_str.encode("utf-8")).hexdigest()
+
+    def artifact_key(self, cache_key):
+        return None
+
+
+async def test_inline_reference_skips_repeat_resolution(make_encoder):
+    proc, audio = _FakeProcessor(), _DigestAudio()
+    ref = "data:audio/wav;base64,QUJD"
+    audio.register(ref, 5)
+    enc = make_encoder(proc)
+
+    first, key1 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+    first.fill_(0)  # callers own their copy
+    second, key2 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+
+    assert audio.resolve_calls == [ref]
+    assert key1 == key2
+    assert torch.equal(second, torch.full((3, _N_VQ), 5, dtype=torch.long))
+
+
+async def test_inline_index_requires_matching_resolve_key(make_encoder):
+    proc, audio = _FakeProcessor(), _DigestAudio(digest_keys=False)
+    ref = "data:audio/wav;base64,REVG"
+    audio.register(ref, 6)
+    enc = make_encoder(proc)
+
+    for _ in range(2):
+        await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+    assert audio.resolve_calls == [ref, ref]
+
+
+async def test_inline_reference_shared_across_encoders(make_encoder, monkeypatch):
+    # Keep the AF_UNIX socket path independent of pytest's nested base path.
+    with TemporaryDirectory(prefix="moss-ref-") as shared_dir:
+        monkeypatch.setenv("VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", shared_dir)
+        ref = "data:audio/wav;base64,R0hJ"
+        first_audio, second_audio = _DigestAudio(), _DigestAudio()
+        first_audio.register(ref, 9)
+        second_audio.register(ref, 9)
+        first, second = make_encoder(_FakeProcessor()), make_encoder(_FakeProcessor())
+
+        a, key_a = await first.encode(
+            ref, resolve_ref_audio=first_audio.resolve, get_artifact_key=first_audio.artifact_key
+        )
+        b, key_b = await second.encode(
+            ref, resolve_ref_audio=second_audio.resolve, get_artifact_key=second_audio.artifact_key
+        )
+
+        assert first_audio.resolve_calls == [ref] and second_audio.resolve_calls == []
+        assert key_a == key_b and torch.equal(a, b)
+
+
+# --------------------------------------------------------------------------- #
+# Reference-encoder CUDA graphs: selection and processor-equivalent prep      #
+# --------------------------------------------------------------------------- #
+
+
+class _GraphsStub:
+    def __init__(self, result):
+        self.result = result
+        self.calls: list[list[torch.Tensor]] = []
+
+    def encode(self, wavs):
+        self.calls.append(wavs)
+        return self.result
+
+
+class _LoudnessProcessor(_FakeProcessor):
+    @staticmethod
+    def loudness_normalize(wav):
+        return wav * 2.0
+
+
+async def test_encode_prepared_prefers_graphs_and_falls_back_to_processor(make_encoder):
+    proc = _LoudnessProcessor()
+    enc = make_encoder(proc)
+    graphs_codes = [torch.full((3, _N_VQ), 7, dtype=torch.long)]
+    enc._graphs_enabled, enc._graphs = True, _GraphsStub(graphs_codes)
+    prepared = [torch.ones(1, 100)]
+    assert enc._encode_prepared(prepared) is graphs_codes and proc.attempt_sizes == []
+    enc._graphs = _GraphsStub([None])  # no graph fits this clip
+    out = enc._encode_prepared(prepared)
+    assert proc.attempt_sizes == [1] and out[0].tolist() == [[1] * _N_VQ] * 3
+
+
+async def test_processor_prepare_matches_processor_channel_and_loudness_handling(make_encoder):
+    enc = make_encoder(_LoudnessProcessor())
+    mono = enc._processor_prepare(torch.ones(100))
+    assert mono.shape == (2, 100) and torch.equal(mono, torch.full((2, 100), 2.0))
+    many = enc._processor_prepare(torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4))
+    assert torch.equal(many, torch.arange(8, dtype=torch.float32).reshape(2, 4) * 2.0)
+
+
+async def test_processor_prepare_resamples_to_the_tokenizer_rate(make_encoder):
+    torchaudio = pytest.importorskip("torchaudio")
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    enc = make_encoder(proc)  # works at 24 kHz, the Local-v1.5 reference rate
+    wav = torch.randn(1, 2400)
+    expected = torchaudio.functional.resample(wav.repeat(2, 1), 24000, 48000) * 2.0
+    # The encoder resamples the mono clip before duplicating it. Some CPU conv
+    # kernels round a one-row batch differently from a two-row one, so the two
+    # orders agree to float32 precision, not bit for bit.
+    torch.testing.assert_close(enc._processor_prepare(wav), expected, rtol=0, atol=1e-5)
+
+
+@pytest.mark.parametrize("role", ["", "host", "client"])
+async def test_prepare_captures_graphs_or_starts_the_shared_host(make_encoder, monkeypatch, tmp_path, role):
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    enc = make_encoder(_FakeProcessor())
+    enc._shares_encoder, enc._shared_codes_dir = True, str(tmp_path)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: role)
+    captured, hosts, ready = [], [], []
+    monkeypatch.setattr(enc, "_reference_graphs", lambda: captured.append(True))
+    monkeypatch.setattr(enc, "_host_workers", lambda: [])
+    helper = SimpleNamespace(wait_until_ready=lambda: ready.append(True))
+
+    def make_host(*args, **kwargs):
+        hosts.append(args)
+        return helper
+
+    monkeypatch.setattr(re_mod, "SharedReferenceEncoderHost", make_host)
+    monkeypatch.setattr(re_mod, "SharedReferenceEncoderClient", lambda *args: helper)
+    enc.prepare()
+    enc.prepare()
+    # Every API process waits for the shared host before startup completes.
+    assert (len(captured), len(hosts)) == {"": (2, 0), "host": (0, 1), "client": (0, 0)}[role]
+    assert len(ready) == (2 if role else 0)
+
+
+async def test_clips_reach_the_encoder_at_the_tokenizer_rate(make_encoder, monkeypatch):
+    torchaudio = pytest.importorskip("torchaudio")
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    enc = make_encoder(proc)  # works at 24 kHz
+    sent = []
+
+    def record_prepared(prepared):
+        sent.extend(prepared)
+        return [None] * len(prepared)
+
+    monkeypatch.setattr(enc, "_encode_shared_or_local", record_prepared)
+    wav = [float(v) for v in torch.randn(2400)]
+    enc._encode_batch_sync([(wav, _SR)])
+    # Resampled where the request arrived (a shared encoder receives it this way); still one channel.
+    expected = torchaudio.functional.resample(torch.tensor(wav)[None], 24000, 48000)
+    torch.testing.assert_close(sent[0], expected, rtol=0, atol=0)
+    # The eager path is told the clip is already at the tokenizer rate.
+    calls = []
+
+    def record_encoding(wavs, sampling_rate, n_vq=None):
+        calls.append(sampling_rate)
+        return [None]
+
+    monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding)
+    enc._graphs_enabled = False
+    enc._encode_prepared(sent)
+    assert calls == [48000]
+
+
+async def test_batcher_runs_up_to_max_inflight_batches_at_once():
+    inside = threading.Barrier(2, timeout=5)
+    seen: list[int] = []
+
+    def encode_batch(payload):
+        inside.wait()  # two batches are encoding at the same time
+        seen.append(len(payload))
+        return [torch.zeros(1) for _ in payload]
+
+    batcher = _RefEncodeBatcher(encode_batch, window_ms=0.0, max_batch=8, max_inflight=2)
+    try:
+        first = asyncio.ensure_future(batcher.submit([1.0], _SR))
+        await asyncio.sleep(0.05)  # the first batch is in flight before the second clip arrives
+        await asyncio.gather(first, batcher.submit([2.0], _SR))
+    finally:
+        await batcher.aclose()
+    assert seen == [1, 1]
+
+
+def test_reference_stream_priority(monkeypatch):
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    monkeypatch.delenv(re_mod._REF_ENCODE_PRIORITY_ENV, raising=False)
+    assert re_mod.reference_stream_priority() == 0
+    monkeypatch.setenv(re_mod._REF_ENCODE_PRIORITY_ENV, "high")
+    expected = torch.cuda.Stream.priority_range()[1] if torch.cuda.is_available() else 0
+    assert re_mod.reference_stream_priority() == expected
+
+
+async def test_device_prep_resamples_on_the_tokenizer_device(make_encoder, monkeypatch):
+    torchaudio = pytest.importorskip("torchaudio")
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    monkeypatch.setattr(proc, "audio_tokenizer", torch.nn.Linear(1, 1), raising=False)  # CPU placement
+    enc = make_encoder(proc)
+    enc._gpu_prep = True
+    wav = torch.randn(1, 2400)
+    clip = enc._encoder_input(wav)
+    assert clip.shape == (1, 2400)  # left at the working rate for the encoder's device
+    expected = torchaudio.functional.resample(wav, 24000, 48000).repeat(2, 1) * 2.0
+    torch.testing.assert_close(enc._finish_prepare(clip), expected, rtol=1e-5, atol=1e-6)
+    # The eager fallback is told the clip's own rate.
+    calls = []
+
+    def record_encoding(wavs, sampling_rate, n_vq=None):
+        calls.append(sampling_rate)
+        return [None]
+
+    monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding)
+    enc._graphs_enabled = False
+    enc._encode_prepared([clip])
+    assert calls == [24000]
+    assert re_mod._REF_ENCODE_GPU_PREP_ENV == "VLLM_OMNI_MOSS_REF_GPU_PREP"

tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py

@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference-encoder CUDA graphs reproduce the tokenizer's eager encode."""
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import (
+    MossReferenceEncoderGraphs,
+    _cached_reference_rope,
+    _reference_rope_factors,
+    _validate_unmasked_reference_encoder,
+    _windowed_attention,
+    install_windowed_attention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+FRAME, N_VQ = 8, 3
+
+
+class _Stage(nn.Module):
+    """Causal per-frame stage: frames never mix, like a causal encoder layer."""
+
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out)
+
+    def forward(self, x, lengths):
+        return torch.tanh(self.proj(x)), lengths
+
+
+class _Quantizer(nn.Module):
+    def forward(self, hidden, lengths, n_q):
+        # hidden (B, frames, D): codes (n_q, B, frames) from per-frame argmax slices.
+        codes = torch.stack([hidden[..., q::n_q].argmax(-1) for q in range(n_q)])
+        return hidden, codes, lengths
+
+
+class _Tokenizer(nn.Module):
+    number_channels = 2
+    downsample_rate = FRAME
+    sampling_rate = 64
+
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.ModuleList([_Stage(FRAME * 2, 16), _Stage(16, 12)])
+        self.quantizer = _Quantizer()
+
+    def _flatten_channels_for_codec(self, audio, lengths):
+        batch, channels, samples = audio.shape
+        frames = samples // FRAME
+        values = audio.reshape(batch, channels, frames, FRAME).permute(0, 2, 1, 3).reshape(batch, frames, -1)
+        return values, lengths // FRAME
+
+    def _codec_inference_autocast(self):
+        return nullcontext()
+
+    def eager(self, wavs):
+        out = []
+        for wav in wavs:
+            values, lengths = self._flatten_channels_for_codec(wav[None], torch.tensor([wav.shape[-1]]))
+            for module in self.encoder:
+                values, lengths = module(values, lengths)
+            _, codes, _ = self.quantizer(values, lengths, N_VQ)
+            out.append(codes[:, 0, : wav.shape[-1] // FRAME].transpose(0, 1).long())
+        return out
+
+
+@pytest.fixture
+def cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda", torch.accelerator.current_device_index())
+
+
+def test_cached_rope_reuses_longest_table_without_modifying_inputs(cuda):
+    from types import SimpleNamespace
+
+    module = SimpleNamespace(rope=SimpleNamespace(max_period=10000.0))
+    cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
+    torch.manual_seed(17)
+    # Prime once; all subsequent sequences and batches use an immutable prefix.
+    q = torch.randn(1, 3, 128, 32, device=cuda, dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    _cached_reference_rope(module, q, k, cache)
+    stored = next(iter(cache.values()))
+    for batch, length in [(1, 17), (4, 64), (2, 128)]:
+        q = torch.randn(batch, 3, length, 32, device=cuda, dtype=torch.bfloat16)
+        k = torch.randn_like(q)
+        before = (q.clone(), k.clone())
+        got = _cached_reference_rope(module, q, k, cache)
+        cos, sin = _reference_rope_factors(length, 32, 10000.0, cuda)
+        torch.testing.assert_close(stored[0][:length], cos, rtol=0, atol=0)
+        torch.testing.assert_close(stored[1][:length], sin, rtol=0, atol=0)
+        for value, source, saved in zip(got, (q, k), before):
+            real, imag = source.float().reshape(batch, 3, length, 16, 2).unbind(-1)
+            expected = torch.stack((real * cos - imag * sin, real * sin + imag * cos), -1).to(source.dtype)
+            torch.testing.assert_close(value, expected.reshape_as(source), rtol=0, atol=0)
+            torch.testing.assert_close(source, saved, rtol=0, atol=0)
+        assert len(cache) == 1
+        assert next(iter(cache.values())) is stored
+
+
+def test_cached_rope_distinguishes_period_and_head_width(cuda):
+    from types import SimpleNamespace
+
+    cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
+    for width, period in [(32, 10000.0), (64, 10000.0), (32, 1000.0)]:
+        module = SimpleNamespace(rope=SimpleNamespace(max_period=period))
+        q = torch.randn(1, 2, 32, width, device=cuda, dtype=torch.bfloat16)
+        _cached_reference_rope(module, q, q, cache)
+    assert len(cache) == 3
+
+
+@pytest.mark.parametrize("batched_transfer", [False, True])
+def test_graph_codes_match_eager_per_clip_and_fall_back_outside_buckets(cuda, batched_transfer):
+    torch.manual_seed(0)
+    tokenizer = _Tokenizer().to(cuda)
+    # Buckets of 2 s and 4 s at 64 Hz: 128 and 256 samples.
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), batched_transfer=batched_transfer
+    )
+    graphs.capture()
+    assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128), (4, 256)]
+    assert graphs.max_batch == 4
+
+    for count, frames in [(1, [5]), (3, [16, 3, 9]), (4, [32, 1, 20, 17]), (2, [16, 16])]:
+        wavs = [torch.randn(2, n * FRAME) for n in frames]
+        codes = graphs.encode(wavs)
+        expected = tokenizer.eager([w.to(cuda) for w in wavs])
+        assert len(codes) == count
+        for got, want in zip(codes, expected):
+            assert got.device.type == "cpu" and got.dtype == torch.long
+            torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+    # Longer than every bucket: left to the caller; the rest still use graphs.
+    long_clip, short_clip = torch.randn(2, 33 * FRAME), torch.randn(2, 4 * FRAME)
+    codes = graphs.encode([long_clip, short_clip])
+    assert codes[0] is None
+    torch.testing.assert_close(codes[1], tokenizer.eager([short_clip.to(cuda)])[0].cpu(), rtol=0, atol=0)
+
+    # Wider than every batch: split, longest first, and returned in input order.
+    wavs = [torch.randn(2, n * FRAME) for n in (3, 30, 7, 12, 1)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+    # Result ownership includes the degenerate single-frame shape. Reusing
+    # the pinned output slab must not overwrite a previous caller's codes.
+    previous = graphs.encode([torch.randn(2, FRAME)])[0]
+    snapshot = previous.clone()
+    graphs.encode([torch.randn(2, 16 * FRAME) for _ in range(4)])
+    torch.testing.assert_close(previous, snapshot, rtol=0, atol=0)
+
+
+def test_batch_audio_budget_limits_captured_graphs(cuda):
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), max_batch_seconds=8.0
+    )
+    graphs.capture()
+    # 4 x 4 s exceeds the 8 s budget; a single clip is always captured.
+    assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128)]
+    wavs = [torch.randn(2, 30 * FRAME) for _ in range(3)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+def test_finer_singleton_buckets_preserve_multiclip_batching(cuda):
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer,
+        n_vq=N_VQ,
+        batch_sizes=(1, 2, 4),
+        bucket_seconds=(2.0, 4.0),
+        singleton_bucket_seconds=(1.0, 3.0),
+    )
+    graphs.capture()
+    assert (1, 64) in graphs.captured and (1, 192) in graphs.captured
+    assert (2, 64) not in graphs.captured and (4, 192) not in graphs.captured
+    replay = graphs._replay
+    shapes = []
+
+    def record(entry, wavs):
+        shapes.append(tuple(entry.audio.shape))
+        return replay(entry, wavs)
+
+    graphs._replay = record
+    one = torch.randn(2, 20 * FRAME)
+    out = graphs.encode([one])
+    assert shapes == [(1, 2, 192)]
+    torch.testing.assert_close(out[0], tokenizer.eager([one.to(cuda)])[0].cpu(), rtol=0, atol=0)
+    shapes.clear()
+    wavs = [one, one[:, : 8 * FRAME], one[:, : 12 * FRAME]]
+    out = graphs.encode(wavs)
+    assert shapes == [(4, 2, 256)]
+    for got, want in zip(out, tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+def test_compiled_capture_matches_eager(cuda):
+    torch.manual_seed(2)
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 4), bucket_seconds=(2.0, 4.0), compile_core=True
+    )
+    graphs.capture()
+    assert graphs.captured == [(1, 128), (1, 256), (4, 128), (4, 256)]
+    wavs = [torch.randn(2, n * FRAME) for n in (5, 30, 12)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+class _Attention(nn.Module):
+    """The attention surface the tokenizer's MHA exposes to the windowed kernel."""
+
+    def __init__(self, embed_dim=128, heads=2, context=5):
+        super().__init__()
+        self.embed_dim, self.num_heads, self.context, self.causal = embed_dim, heads, context, True
+        self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
+
+    def _project_qkv(self, x):
+        b, t, _ = x.shape
+        return self.in_proj(x).reshape(b, t, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0)
+
+    def _apply_dense_rope(self, q, k):
+        return q, k
+
+    def _forward_non_streaming_sdpa(self, x, input_lengths):
+        # Reference semantics: causal, distance < context, padded rows zeroed.
+        b, t, _ = x.shape
+        q, k, v = self._project_qkv(x)
+        pos = torch.arange(t, device=x.device)
+        delta = pos.view(-1, 1) - pos.view(1, -1)
+        mask = (delta >= 0) & (delta < self.context)
+        out = F.scaled_dot_product_attention(q, k, v, mask[None, None])
+        valid = (pos.view(1, t) < input_lengths.view(-1, 1)).view(b, 1, t, 1)
+        out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype))
+        return out.transpose(1, 2).reshape(b, t, self.embed_dim)
+
+
+@pytest.mark.parametrize("fa_version", [2, 3])
+def test_windowed_attention_matches_masked_sdpa(cuda, fa_version):
+    if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9:
+        pytest.skip("FA3 requires Hopper")
+    torch.manual_seed(1)
+    model = nn.Module()
+    model.encoder = nn.ModuleList([_Attention().to(cuda, torch.bfloat16)])
+    attention = model.encoder[0]
+    x = torch.randn(3, 40, 128, device=cuda, dtype=torch.bfloat16)
+    lengths = torch.tensor([40, 17, 3], device=cuda)
+    expected = attention._forward_non_streaming_sdpa(x, lengths)
+    assert install_windowed_attention(model, fa_version=fa_version) == 1
+    actual = attention._forward_non_streaming_sdpa(x, lengths)
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    assert actual[1, 17:].count_nonzero() == 0 and actual[2, 3:].count_nonzero() == 0
+    # Non-bf16 inputs keep the tokenizer's own SDPA path.
+    x32 = x.float()
+    attention.float()
+    torch.testing.assert_close(
+        attention._forward_non_streaming_sdpa(x32, lengths), expected.float(), rtol=3e-2, atol=3e-2
+    )
+
+
+@pytest.mark.parametrize("fa_version", [2, 3])
+def test_unmasked_causal_queries_preserve_valid_prefix_and_ignore_future_padding(cuda, fa_version):
+    if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9:
+        pytest.skip("FA3 requires Hopper")
+    torch.manual_seed(44)
+    attention = _Attention().to(cuda, torch.bfloat16)
+    x = torch.randn(4, 40, 128, device=cuda, dtype=torch.bfloat16)
+    lengths = torch.tensor([40, 17, 1, 0], device=cuda)
+    kwargs = dict(sdpa=attention._forward_non_streaming_sdpa, fa_version=fa_version)
+    with torch.no_grad():
+        masked = _windowed_attention(attention, x, lengths, **kwargs)
+        unmasked = _windowed_attention(attention, x, lengths, skip_padded_query_mask=True, **kwargs)
+        poisoned = x.clone()
+        for row, length in enumerate([40, 17, 1, 0]):
+            poisoned[row, length:] = torch.randn_like(poisoned[row, length:]) * 30
+        changed = _windowed_attention(attention, poisoned, lengths, skip_padded_query_mask=True, **kwargs)
+        for row, length in enumerate([40, 17, 1, 0]):
+            torch.testing.assert_close(unmasked[row, :length], masked[row, :length], rtol=0, atol=0)
+            torch.testing.assert_close(changed[row, :length], masked[row, :length], rtol=0, atol=0)
+            assert masked[row, length:].count_nonzero() == 0
+        assert unmasked[1, 17:].count_nonzero() > 0
+        # The FP32 fallback retains its original masking contract.
+        attention.float()
+        got = _windowed_attention(attention, x.float(), lengths, skip_padded_query_mask=True, **kwargs)
+        expected = attention._forward_non_streaming_sdpa(x.float(), lengths)
+        torch.testing.assert_close(got, expected, rtol=0, atol=0)
+
+
+def test_unmasked_queries_reject_unknown_encoder_before_patching():
+    tokenizer = _Tokenizer()
+    with pytest.raises(ValueError, match="patch-downsample"):
+        install_windowed_attention(tokenizer, skip_padded_query_mask=True)
+
+
+@pytest.mark.parametrize("bad_part", ["upsample", "noncausal", "quantizer", "convolution"])
+def test_unmasked_queries_reject_structures_that_can_read_invalid_rows(bad_part):
+    from types import SimpleNamespace
+
+    patch = type("MossAudioTokenizerPatchedPretransform", (), {})()
+    patch.is_downsample = bad_part != "upsample"
+    projected = type("MossAudioTokenizerProjectedTransformer", (), {})()
+    projected.transformer = SimpleNamespace(
+        layers=[SimpleNamespace(self_attn=SimpleNamespace(causal=bad_part != "noncausal"))]
+    )
+    quantizer_cls = type(
+        "MossAudioTokenizerResidualVQ" if bad_part != "quantizer" else "UnknownQuantizer", (nn.Module,), {}
+    )
+    quantizer = quantizer_cls()
+    quantizer.proj = nn.Conv1d(2, 2, 3 if bad_part == "convolution" else 1)
+    tokenizer = SimpleNamespace(encoder=[patch, projected], quantizer=quantizer)
+    with pytest.raises(ValueError):
+        _validate_unmasked_reference_encoder(tokenizer)

tests/model_executor/models/moss_tts/test_shared_reference_encoder.py

@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""API processes share one reference encoder through a Unix socket."""
+
+import threading
+from tempfile import TemporaryDirectory
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts import shared_reference_encoder as shared
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _codes(wav):
+    # Deterministic per clip: frames from the clip length, values from its first sample.
+    return torch.full((wav.shape[-1] // 10, 4), int(wav.reshape(-1)[0]), dtype=torch.long)
+
+
+class _Recorder:
+    def __init__(self, gate=None):
+        self.batches = []
+        self.gate = gate
+
+    def __call__(self, wavs):
+        if self.gate is not None:
+            self.gate.wait(5)
+        self.batches.append(len(wavs))
+        return [RuntimeError("bad clip") if int(w.reshape(-1)[0]) < 0 else _codes(w) for w in wavs]
+
+
+@pytest.fixture
+def host_dir(monkeypatch):
+    monkeypatch.setattr(shared, "_host_lock_handle", None)
+    # AF_UNIX limits the full socket path, including pytest directory names.
+    with TemporaryDirectory(prefix="moss-ref-") as path:
+        yield path
+
+
+def test_only_one_process_per_directory_hosts(host_dir):
+    assert shared.elect_host(host_dir) is True
+    assert shared.elect_host(host_dir) is False  # a second open file description cannot take the lock
+
+
+def test_close_releases_queued_and_active_callers(host_dir):
+    entered, release = threading.Event(), threading.Event()
+
+    def encode(wavs):
+        entered.set()
+        assert release.wait(5)
+        return [_codes(wav) for wav in wavs]
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [(encode, None)])
+    results: dict[str, list[torch.Tensor | BaseException]] = {}
+    threads = [
+        threading.Thread(target=lambda key=key: results.setdefault(key, host.encode([torch.ones(1, 20)])), daemon=True)
+        for key in ("active", "queued")
+    ]
+    try:
+        threads[0].start()
+        assert entered.wait(2)
+        threads[1].start()
+        with host._ready:
+            assert host._ready.wait_for(lambda: len(host._pending) == 1, timeout=2)
+        host.close()
+        for thread in threads:
+            thread.join(1)
+            assert not thread.is_alive(), "closed encoder left a request waiting"
+        assert all(isinstance(results[key][0], RuntimeError) for key in results)
+    finally:
+        release.set()
+        host.close()
+        for thread in threads:
+            if thread.ident is not None:
+                thread.join(1)
+
+
+def test_worker_result_count_mismatch_fails_entire_batch(host_dir):
+    host = shared.SharedReferenceEncoderHost(host_dir, [(lambda wavs: [_codes(wavs[0])], None)])
+    results = []
+    caller = threading.Thread(target=lambda: results.extend(host.encode([torch.ones(1, 20)] * 2)), daemon=True)
+    try:
+        caller.start()
+        caller.join(2)
+        assert not caller.is_alive(), "missing encoder result left a request waiting"
+        assert len(results) == 2
+        assert all(isinstance(result, RuntimeError) for result in results)
+    finally:
+        host.close()
+        caller.join(1)
+
+
+def test_client_round_trip_and_per_clip_errors(host_dir):
+    host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), None)])
+    try:
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        wavs = [torch.full((1, 40), 3.0), torch.full((1, 70), -1.0), torch.full((1, 20), 5.0)]
+        out = client.encode(wavs)
+        assert out[0].dtype == torch.long and torch.equal(out[0], _codes(wavs[0]))
+        assert isinstance(out[1], RuntimeError) and "bad clip" in str(out[1])
+        assert torch.equal(out[2], _codes(wavs[2]))
+        # The host process's own clips skip the socket.
+        assert torch.equal(host.encode([wavs[2]])[0], _codes(wavs[2]))
+    finally:
+        host.close()
+
+
+def test_host_merges_clips_from_concurrent_clients(host_dir):
+    gate = threading.Event()
+    recorder = _Recorder(gate)
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)])
+    try:
+        results = {}
+
+        def send(name, value):
+            results[name] = shared.SharedReferenceEncoderClient(host_dir).encode([torch.full((1, 30), value)] * 2)
+
+        # Encoding is held until both clients' clips are queued.
+        warm = threading.Thread(target=lambda: host.encode([torch.full((1, 10), 1.0)]))
+        warm.start()
+        threads = [threading.Thread(target=send, args=(n, v)) for n, v in (("a", 2.0), ("b", 4.0))]
+        for thread in threads:
+            thread.start()
+        threading.Timer(0.5, gate.set).start()
+        for thread in (warm, *threads):
+            thread.join(10)
+        # Both clients' clips were encoded together, not one request at a time.
+        assert sum(recorder.batches) == 5 and max(recorder.batches) >= 4
+        assert all(int(c[0, 0]) == 2 for c in results["a"]) and all(int(c[0, 0]) == 4 for c in results["b"])
+    finally:
+        host.close()
+
+
+def test_client_without_host_raises(host_dir, monkeypatch):
+    monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3)
+    with pytest.raises((FileNotFoundError, ConnectionRefusedError)):
+        shared.SharedReferenceEncoderClient(host_dir).encode([torch.zeros(1, 10)])
+
+
+def test_host_window_collects_across_workers_without_reserving_singletons(host_dir):
+    recorder = _Recorder()
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)] * 3, batch_window_ms=500)
+    try:
+        barrier = threading.Barrier(3)
+        results = []
+
+        def send(value):
+            barrier.wait()
+            results.extend(host.encode([torch.full((1, 20), value)]))
+
+        threads = [threading.Thread(target=send, args=(v,)) for v in (1.0, 2.0)]
+        for thread in threads:
+            thread.start()
+        barrier.wait()
+        for thread in threads:
+            thread.join(5)
+            assert not thread.is_alive()
+        assert recorder.batches == [2]
+        assert sorted(int(x[0, 0]) for x in results) == [1, 2]
+    finally:
+        host.close()
+
+
+def test_reference_encoder_client_falls_back_to_local_encode(host_dir, monkeypatch):
+    from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+    from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+    monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client")
+    monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3)
+    proc = _FakeProcessor()
+    enc = re_mod.MossReferenceEncoder(
+        proc, variant="local", n_vq=_N_VQ, sr_target=_SR, speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20)
+    )
+    placed = []
+    monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", lambda: placed.append(True))
+    out = enc._encode_batch_sync([([7.0] * 50, _SR)])
+    assert placed == [True] and proc.attempt_sizes == [1] and int(out[0][0, 0]) == 7
+
+
+def test_workers_encode_concurrently_after_their_warmups(host_dir):
+    inside, release = threading.Barrier(3, timeout=5), threading.Event()
+    warmed = []
+
+    def worker(name):
+        def encode(wavs):
+            inside.wait()  # both workers are inside an encode at the same time
+            release.wait(5)
+            return [_codes(w) for w in wavs]
+
+        return encode, lambda: warmed.append(name)
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [worker("a"), worker("b")])
+    try:
+        out = []
+
+        def submit(value):
+            out.append(host.encode([torch.full((1, 20), value)]))
+
+        first = threading.Thread(target=submit, args=(1.0,))
+        first.start()
+        # The first worker is busy by now, so the second clip goes to the other one.
+        threading.Event().wait(0.1)
+        second = threading.Thread(target=submit, args=(2.0,))
+        second.start()
+        inside.wait()
+        release.set()
+        for thread in (first, second):
+            thread.join(5)
+        assert sorted(warmed) == ["a", "b"] and sorted(int(o[0][0, 0]) for o in out) == [1, 2]
+    finally:
+        host.close()
+
+
+def test_client_requests_in_flight_use_their_own_connections(host_dir):
+    entered, release = threading.Semaphore(0), threading.Event()
+
+    def worker():
+        def encode(wavs):
+            entered.release()
+            release.wait(5)
+            return [_codes(w) for w in wavs]
+
+        return encode, None
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [worker(), worker()])
+    try:
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        out = []
+        threads = [
+            threading.Thread(target=lambda v=v: out.append(client.encode([torch.full((1, 20), v)]))) for v in (1.0, 2.0)
+        ]
+        for thread in threads:
+            thread.start()
+            # A worker holds this request, still unanswered, before the next is
+            # sent. Sent together, one free worker may take both as one batch.
+            assert entered.acquire(timeout=5)
+        release.set()  # both requests reached the host at once
+        for thread in threads:
+            thread.join(5)
+        assert sorted(int(o[0][0, 0]) for o in out) == [1, 2]
+        assert len(client._idle) == 2  # both connections are kept for reuse
+    finally:
+        host.close()
+
+
+def test_readiness_waits_for_every_worker_warmup(host_dir, monkeypatch):
+    entered, release = threading.Event(), threading.Event()
+    recorder = _Recorder()
+    warmed = []
+
+    def last_warmup():
+        entered.set()
+        assert release.wait(5)
+        warmed.append("last")
+
+    host = shared.SharedReferenceEncoderHost(
+        host_dir, [(recorder, lambda: warmed.append("first")), (recorder, last_warmup)]
+    )
+    client = shared.SharedReferenceEncoderClient(host_dir)
+    try:
+        assert entered.wait(5)
+        monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 0.03)
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            client.wait_until_ready()
+        assert warmed == ["first"] and recorder.batches == [] and not client._idle
+        release.set()
+        monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5)
+        host.wait_until_ready()
+        client.wait_until_ready()
+        assert warmed == ["first", "last"] and recorder.batches == []
+        assert len(client._idle) == 1
+        assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20)))
+    finally:
+        release.set()
+        host.close()
+
+
+def test_slow_startup_does_not_use_encode_timeout(host_dir, monkeypatch):
+    release = threading.Event()
+    monkeypatch.setattr(shared, "_REQUEST_TIMEOUT_S", 0.05)
+    monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5, raising=False)
+    host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), lambda: release.wait(5))])
+    timer = threading.Timer(0.2, release.set)
+    timer.start()
+    try:
+        # The first socket waits for readiness before starting its much shorter
+        # encode deadline. No API warmup call is required to protect this path.
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20)))
+    finally:
+        release.set()
+        timer.join(5)
+        host.close()
+
+
+def test_failed_host_warmup_is_not_ready(host_dir):
+    def fail():
+        raise RuntimeError("warmup failed")
+
+    recorder = _Recorder()
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, fail)])
+    try:
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            shared.SharedReferenceEncoderClient(host_dir).wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.encode([torch.ones(1, 20)])
+        assert recorder.batches == []
+    finally:
+        host.close()
+
+
+def test_startup_failure_does_not_trigger_local_capture(host_dir, monkeypatch):
+    from unittest.mock import Mock
+
+    from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+    from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+    monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client")
+    enc = re_mod.MossReferenceEncoder(
+        _FakeProcessor(),
+        variant="local",
+        n_vq=_N_VQ,
+        sr_target=_SR,
+        speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20),
+    )
+    enc._shared_client = Mock()
+    enc._shared_client.encode.side_effect = shared.SharedReferenceEncoderStartupError("still compiling")
+    local = Mock()
+    monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", local)
+    with pytest.raises(shared.SharedReferenceEncoderStartupError):
+        enc._encode_shared_or_local([torch.ones(1, 20)])
+    local.assert_not_called()

tests/model_executor/models/moss_tts/test_slot_attention.py

@@ -0,0 +1,301 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import copy
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
+    MossAudioTokenizerMultiheadAttention,
+    MossAudioTokenizerRotaryEmbedding,
+    RingKVCache,
+    StreamingExecutionContext,
+)
+from vllm_omni.model_executor.models.moss_tts.slot_attention import slot_ring_attention, slot_ring_attention_rows
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize(
+    "frames,capacity", [(1, 125), (2, 250), (15, 125), (30, 250), (32, 400), (60, 400), (480, 400)]
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("execution", ["graph", "compiled"])
+def test_empty_tile_skip_matches_static_ring_across_reuse(frames, capacity, dtype, execution):
+    torch.manual_seed(45)
+    batch, heads, states, dim = 3, 2, 5, 64
+    projected = torch.randn(batch, frames, 3, heads, dim, device="cuda", dtype=dtype)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, states, heads, capacity, dim, device="cuda", dtype=dtype)
+    expected_cache = cache.clone()
+    end = torch.zeros(states, 2, device="cuda", dtype=torch.long)[:, 0]
+    expected_end = end.clone()
+    slots = torch.tensor([0, 4, 2], device="cuda")
+    lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0]
+
+    def candidate():
+        return slot_ring_attention(q, k, v, cache, end, slots, lengths, 17, skip_empty_tiles=True)
+
+    graph = None
+    if execution == "compiled":
+        run = torch.compile(candidate, fullgraph=True)
+    else:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            candidate()
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            actual = candidate()
+
+    for step, start in enumerate([0, 1, capacity - 1, capacity + 9, 0, 3 * capacity]):
+        end.fill_(start)
+        expected_end.fill_(start)
+        cache.copy_(expected_cache)
+        projected.normal_()
+        slots.copy_(torch.roll(torch.arange(states, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        expected = slot_ring_attention(q, k, v, expected_cache, expected_end, slots, lengths, 17)
+        if graph is not None:
+            graph.replay()
+        else:
+            actual = run()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, expected_end, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("frames,capacity", [(1, 125), (15, 140), (480, 400), (480, 880)])
+def test_all_padding_preserves_state_and_replay_can_become_active(dtype, frames, capacity):
+    q, k, v = [torch.randn(2, 2, frames, 64, device="cuda", dtype=dtype) for _ in range(3)]
+    cache = torch.randn(2, 4, 2, capacity, 64, device="cuda", dtype=dtype)
+    end = torch.tensor([0, capacity + 3, 9, 0], device="cuda")
+    slots = torch.tensor([1, 3], device="cuda")
+    lengths = torch.zeros(2, device="cuda", dtype=torch.int32)
+    before_cache, before_end = cache.clone(), end.clone()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity)
+    for active in [False, True, False, True]:
+        lengths.fill_(frames if active else 0)
+        expected = _ring_reference(q, k, v, before_cache, before_end, slots, lengths, capacity)
+        graph.replay()
+        torch.testing.assert_close(cache, before_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, before_end, rtol=0, atol=0)
+        torch.testing.assert_close(out, expected, rtol=0.025, atol=0.025)
+        if not active:
+            assert out.count_nonzero() == 0
+
+
+def test_dense_ring_retains_last_tokens_when_chunk_exceeds_capacity():
+    torch.manual_seed(19)
+    capacity, frames = 400, 480
+    state = RingKVCache(3, 2, 64, capacity, device=torch.device("cuda"), dtype=torch.float32)
+    slots = torch.tensor([1, 0], device="cuda")
+    context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=torch.ones(2, device="cuda", dtype=torch.bool))
+    for start in [0, frames]:
+        k = torch.randn(2, 2, frames, 64, device="cuda")
+        v = torch.randn_like(k)
+        state.complete(k, v, execution_context=context)
+        retained = torch.arange(frames - capacity, frames, device="cuda")
+        destinations = (start + retained) % capacity
+        torch.testing.assert_close(state.cache[0, slots][:, :, destinations], k[:, :, retained], rtol=0, atol=0)
+        torch.testing.assert_close(state.cache[1, slots][:, :, destinations], v[:, :, retained], rtol=0, atol=0)
+        torch.testing.assert_close(state.end_offset[slots], torch.full_like(slots, start + frames), rtol=0, atol=0)
+
+
+def _ring_reference(q, k, v, cache, end, slots, lengths, context):
+    """Independent chronological ring oracle with FP32 attention arithmetic."""
+    capacity = cache.shape[-2]
+    output = torch.zeros_like(q)
+    for row, (slot, length) in enumerate(zip(slots.tolist(), lengths.tolist())):
+        if length == 0:
+            continue
+        start = int(end[slot])
+        for token in range(max(0, length - capacity), length):
+            cache[0, slot, :, (start + token) % capacity] = k[row, :, token]
+            cache[1, slot, :, (start + token) % capacity] = v[row, :, token]
+        stop = start + length
+        physical_positions = torch.full((capacity,), -1, device=q.device, dtype=torch.long)
+        available = torch.arange(max(0, stop - capacity), stop, device=q.device)
+        physical_positions[available % capacity] = available
+        queries = torch.arange(start, stop, device=q.device)
+        distance = queries[:, None] - physical_positions[None]
+        mask = (physical_positions[None] >= 0) & (distance >= 0)
+        if context > 0:
+            mask &= distance < context
+        output[row, :, :length] = F.scaled_dot_product_attention(
+            q[row, :, :length].float(), cache[0, slot].float(), cache[1, slot].float(), mask[None]
+        ).to(q.dtype)
+        end[slot] = stop
+    return output
+
+
+@pytest.mark.parametrize("frames,capacity,head_dim", [(1, 125, 64), (15, 125, 64), (480, 400, 64), (33, 65, 32)])
+@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"])
+def test_slot_attention_preserves_ring_state_and_padding(frames, capacity, head_dim, execution):
+    torch.manual_seed(29)
+    batch, heads, state_capacity = 3, 2, 5
+    projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype)
+    # Strided metadata is allowed by the kernel contract.
+    end = torch.zeros(state_capacity, 2, device="cuda", dtype=torch.long)[:, 0]
+    slots = torch.arange(batch * 2, device="cuda")[::2]
+    lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0]
+    context = capacity
+    attention = slot_ring_attention
+    if execution == "compiled":
+        # Functionalization must preserve mutations of the persistent ring
+        # and its offsets between calls, including slot reuse.
+        attention = torch.compile(slot_ring_attention, backend="inductor", fullgraph=True)
+    if execution == "graph":
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                slot_ring_attention(q, k, v, cache, end, slots, lengths, context)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = slot_ring_attention(q, k, v, cache, end, slots, lengths, context)
+    end.zero_()
+    expected_cache, expected_end = cache.clone(), end.clone()
+    for step in range(8):
+        slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        projected.normal_()
+        if step == 4:
+            end[1:3].zero_()
+            expected_end[1:3].zero_()
+        expected = _ring_reference(q, k, v, expected_cache, expected_end, slots, lengths, context)
+        if execution == "graph":
+            graph.replay()
+        else:
+            actual = attention(q, k, v, cache, end, slots, lengths, context)
+        torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, expected_end, rtol=0, atol=0)
+        torch.testing.assert_close(actual, expected, rtol=0.025, atol=0.025)
+        padding = torch.arange(frames, device="cuda")[None] >= lengths[:, None]
+        assert actual.permute(0, 2, 1, 3)[padding].count_nonzero() == 0
+
+
+@pytest.mark.parametrize("frames,context", [(1, 9), (3, 9), (15, 125), (480, 400)])
+def test_slot_attention_integrates_with_current_mha(frames, context):
+    torch.manual_seed(9)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128,
+        2,
+        causal=True,
+        context=context,
+        rope=MossAudioTokenizerRotaryEmbedding(),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    candidate = copy.deepcopy(reference)
+    candidate._slot_attention = slot_ring_attention
+    for model in (reference, candidate):
+        model._streaming_state = model._init_streaming_state(4)
+    for step in range(8):
+        slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda")
+        valid = torch.tensor([True, True, False], device="cuda")
+        execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=execution_context)
+            actual = candidate(x, x, x, execution_context=execution_context)
+        torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02)
+        cs, rs = candidate._streaming_state, reference._streaming_state
+        torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.cache[:, :3], rs.kv_cache.cache[:, :3], rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))
+
+
+@pytest.mark.parametrize(
+    "frames,capacity,head_dim", [(1, 125, 64), (15, 140, 64), (40, 125, 64), (480, 400, 64), (33, 65, 32)]
+)
+@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"])
+def test_fused_rows_attention_matches_three_launch_op_exactly(frames, capacity, head_dim, execution):
+    torch.manual_seed(31)
+    batch, heads, state_capacity = 3, 2, 6
+    projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype)
+    reference_cache = cache.clone()
+    end = torch.zeros(state_capacity, device="cuda", dtype=torch.long)
+    slots = torch.zeros(batch, device="cuda", dtype=torch.long)
+    lengths = torch.zeros(batch, device="cuda", dtype=torch.int32)
+    starts = torch.zeros(batch, device="cuda", dtype=torch.long)
+    attention = slot_ring_attention_rows
+    if execution == "compiled":
+        attention = torch.compile(slot_ring_attention_rows, backend="inductor", fullgraph=True)
+    if execution == "graph":
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity)
+        cache.copy_(reference_cache)
+    for step in range(8):
+        slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        starts.copy_(end[slots])
+        projected.normal_()
+        expected = slot_ring_attention(q, k, v, reference_cache, end, slots, lengths, capacity)
+        if execution == "graph":
+            graph.replay()
+        else:
+            actual = attention(q, k, v, cache, starts, slots, lengths, capacity)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        torch.testing.assert_close(cache, reference_cache, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("frames,context", [(1, 9), (15, 125), (480, 400)])
+def test_fused_rows_attention_integrates_with_mha_offsets(frames, context):
+    torch.manual_seed(11)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128,
+        2,
+        causal=True,
+        context=context,
+        rope=MossAudioTokenizerRotaryEmbedding(),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    candidate = copy.deepcopy(reference)
+    reference._slot_attention = slot_ring_attention
+    candidate._slot_attention_rows = slot_ring_attention_rows
+    for model in (reference, candidate):
+        model._streaming_state = model._init_streaming_state(4)
+    for step in range(8):
+        slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda")
+        valid = torch.tensor([True, True, False], device="cuda")
+        execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=execution_context)
+            actual = candidate(x, x, x, execution_context=execution_context)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        cs, rs = candidate._streaming_state, reference._streaming_state
+        torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.cache, rs.kv_cache.cache, rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))

tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""A second streaming-decoder graph set can replay concurrently with the first."""
+
+import threading
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
+    CUDAGraphStreamingDecoderWrapper,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+CAP, DIM = 8, 64
+
+
+class _Codec(nn.Module):
+    def __init__(self, device):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(DIM, DIM, device=device) / DIM)
+        self.state = torch.zeros(CAP + 16, DIM, device=device)
+
+    def decode_streaming_tensors(self, codes, lengths, slots, valid_rows):
+        x = codes.float().sum(0)[..., None].expand(-1, -1, DIM) + self.state[slots][:, None]
+        y = x @ self.weight
+        self.state.index_copy_(0, slots, y[:, -1])
+        return y.transpose(1, 2).contiguous(), lengths
+
+    def reset_decoder_state_slots(self, slots):
+        self.state.index_fill_(0, slots.to(self.state.device), 0)
+
+
+def _wrapper(codec, **kwargs):
+    return CUDAGraphStreamingDecoderWrapper(
+        codec,
+        state_capacity=CAP,
+        num_quantizers=2,
+        vllm_config=None,
+        compiled_decode=codec.decode_streaming_tensors,
+        **kwargs,
+    )
+
+
+def test_private_wrapper_captures_on_its_own_stream_and_replays_concurrently():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    codec = _Codec(device)
+    main = _wrapper(codec, batch_sizes=[4], frame_sizes=[3])
+    fast = _wrapper(codec, batch_sizes=[1, 2], frame_sizes=[1], scratch_base=CAP + 4, private_pool=True)
+    main.warmup(device)
+    fast.warmup(device)
+    assert main.is_ready and fast.is_ready
+    assert main._capture_stream is None and fast._capture_stream is not None
+
+    codes_main = torch.randint(0, 5, (2, 4, 3), device=device)
+    codes_fast = torch.randint(0, 5, (2, 2, 1), device=device)
+    errors = []
+
+    def run(wrapper, codes, slots, iters):
+        stream = torch.cuda.Stream(device=device)
+        slot_ids = torch.tensor(slots, device=device)
+        try:
+            with torch.cuda.stream(stream):
+                for _ in range(iters):
+                    codec.reset_decoder_state_slots(slot_ids)
+                    audio = wrapper.decode(codes, slot_ids)[0][: len(slots)].clone()
+                    expected = (codes.float().sum(0)[..., None].expand(-1, -1, DIM) @ codec.weight).transpose(1, 2)
+                    stream.synchronize()
+                    torch.testing.assert_close(audio, expected, rtol=1e-3, atol=1e-3)
+        except Exception as error:  # surfaced below
+            errors.append(error)
+
+    threads = [
+        threading.Thread(target=run, args=(main, codes_main, [0, 1, 2, 3], 200)),
+        threading.Thread(target=run, args=(fast, codes_fast, [4, 5], 200)),
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(timeout=60)
+    assert not any(thread.is_alive() for thread in threads), "concurrent graph replay did not finish"
+    assert not errors, errors

tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py

@@ -54,8 +54,19 @@ def test_terminal_padding_preserves_audio_live_state_and_reused_slots():
     assert [resumed[i].item() for i in (0, 1)] == [2, 36]
 
 
-@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)])
-def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capacity, expected_calls):
+@pytest.mark.parametrize(
+    "graph_capacity,expected_groups",
+    [
+        # No graph, or graphs that do not cover the state capacity: exact tails.
+        (None, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}),
+        (4, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}),
+        # Full graph coverage: finished tails (3, 7) join the padded 8-frame
+        # bucket with the live 8-frame row; the finished 1-frame row keeps the
+        # exact 1-frame graph and the live 15-frame row stays alone.
+        (8, {8: [0, 1, 2], 1: [3], 15: [4]}),
+    ],
+)
+def test_terminal_tails_coalesce_into_padded_frame_buckets(mocker, graph_capacity, expected_groups):
     s = session()
     s._cudagraph_wrapper = (
         None
@@ -81,9 +92,16 @@ def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capa
             for i, (length, done) in enumerate(zip(lengths, finished, strict=True))
         ]
     )
-    assert step.call_count == expected_calls
+    # Slots are leased in request order, so slot == row here.
+    groups = {}
+    for call in step.call_args_list:
+        plan = call.args[0]
+        groups[max(int(codes.shape[1]) for codes in plan.values())] = sorted(plan)
+        terminal = call.kwargs["terminal_slots"] or set()
+        assert terminal == {slot for slot in plan if finished[slot]}
+    assert groups == expected_groups
+    assert step.call_count == len(expected_groups)
+    # Output is always cropped to the exact frame count, padded or not.
     assert [result[i].shape[-1] for i in range(5)] == lengths
-    if graph_capacity == 8:
-        assert list(step.call_args_list[0].args[0]) == [0, 1, 2]
-        assert step.call_args_list[0].kwargs["terminal_slots"] == {0, 1}
-        assert list(step.call_args_list[1].args[0]) == [3]  # Preserve the exact one-frame graph.
+    # Finished rows released their slots; live rows keep theirs.
+    assert set(decoder._stream_req_slots) == {"2", "4"}

tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MOSS raw codec rows are buffered into first/regular chunks in codebook-major order."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.moss_tts import (
+    _MOSS_AUDIO_PAD_CODE,
+    talker2codec_raw_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _manager(first=1, regular=3):
+    config = {"extra": {"codec_chunk_frames": regular, "initial_codec_chunk_frames": first}}
+    return SimpleNamespace(connector=SimpleNamespace(config=config))
+
+
+def _step(manager, frames, finished=False):
+    request = SimpleNamespace(request_id="r", external_req_id="r")
+    payload = talker2codec_raw_async_chunk(manager, {"codes": {"audio": frames}}, request, is_finished=finished)
+    if payload is not None and payload.codes is not None:
+        manager.put_req_chunk["r"] += 1
+        manager.ramp_chunk_count["r"] += 1
+    return payload
+
+
+def test_first_regular_and_terminal_chunks():
+    manager = _manager()
+    rows = torch.arange(6 * 4, dtype=torch.int32).reshape(6, 4)
+    first = _step(manager, rows[:1])
+    assert first.meta.codec_chunk_frames == 1
+    assert first.codes.audio.dtype == torch.int64
+    assert torch.equal(first.codes.audio, rows[0].long())
+
+    assert _step(manager, rows[1:2]) is None
+    padded = torch.full((1, 4), _MOSS_AUDIO_PAD_CODE, dtype=torch.int32)
+    assert _step(manager, padded) is None  # PAD rows are not audio
+    assert _step(manager, rows[2:3]) is None
+    regular = _step(manager, rows[3:4])
+    assert regular.meta.codec_chunk_frames == 3
+    assert torch.equal(regular.codes.audio, rows[1:4].long().T.reshape(-1))
+
+    tail = _step(manager, rows[4:6], finished=True)
+    assert tail.meta.codec_chunk_frames == 2
+    assert torch.equal(tail.codes.audio, rows[4:6].long().T.reshape(-1))
+    assert bool(tail.meta.finished)
+    assert "r" not in manager.code_prompt_token_ids
+
+
+def test_buffered_rows_do_not_alias_producer_storage():
+    manager = _manager(first=2, regular=2)
+    frames = torch.tensor([[7, 8, 9]])
+    assert _step(manager, frames) is None
+    frames.fill_(0)
+    chunk = _step(manager, torch.tensor([[1, 2, 3]]))
+    assert torch.equal(chunk.codes.audio, torch.tensor([7, 1, 8, 2, 9, 3]))
+
+
+@pytest.mark.parametrize("total", [1, 3, 5, 10, 13, 28, 31])
+def test_ramp_preserves_every_frame_and_terminal(total):
+    manager = _manager(first=1, regular=15)
+    manager.connector.config["extra"]["codec_chunk_ramp"] = [1, 4, 8, 15]
+    rows = torch.arange(total * 4, dtype=torch.int32).reshape(total, 4)
+    packets = []
+    for index in range(total):
+        packet = _step(manager, rows[index : index + 1], finished=index == total - 1)
+        if packet is not None:
+            packets.append(packet)
+    rebuilt = torch.cat([packet.codes.audio.reshape(4, -1).T for packet in packets])
+    assert torch.equal(rebuilt, rows.long())
+    assert all(not bool(packet.meta.finished) for packet in packets[:-1])
+    assert bool(packets[-1].meta.finished)
+    assert "r" not in manager.code_prompt_token_ids
+    if total == 31:
+        assert [packet.meta.codec_chunk_frames for packet in packets] == [1, 4, 8, 15, 3]

tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Empty MOSS terminals must not execute prewarmed or previous codec tokens."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator
+from vllm_omni.data_entry_keys import to_dict
+from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import (
+    _OmniConnectorPayloadTransportMixin as Transport,
+)
+from vllm_omni.model_executor.stage_input_processors.moss_tts import (
+    _MOSS_AUDIO_PAD_CODE,
+    talker2codec_raw_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _request(tokens):
+    return SimpleNamespace(
+        request_id="r",
+        external_req_id="r",
+        prompt_token_ids=list(tokens),
+        num_prompt_tokens=len(tokens),
+        _all_token_ids=list(tokens),
+        _output_token_ids=[],
+        num_computed_tokens=0,
+    )
+
+
+@pytest.mark.parametrize("after_audio", [False, True])
+def test_empty_terminal_clears_native_codec_prompt(after_audio):
+    manager = SimpleNamespace(
+        connector=SimpleNamespace(config={"extra": {"initial_codec_chunk_frames": 1, "codec_chunk_frames": 3}})
+    )
+    source = _request([0] * 88)
+    previous = [0] * 88
+    if after_audio:
+        packet = talker2codec_raw_async_chunk(manager, {"codes": {"audio": torch.ones(1, 12)}}, source)
+        manager.put_req_chunk["r"] += 1
+        previous = packet.codes.audio.tolist()
+    terminal = talker2codec_raw_async_chunk(
+        manager,
+        {"codes": {"audio": torch.full((1, 12), _MOSS_AUDIO_PAD_CODE)}},
+        source,
+        is_finished=True,
+    )
+    payload = to_dict(terminal)
+    assert not Transport._payload_is_consumable(payload)
+    metadata = Transport._extract_scheduling_metadata(payload)
+    receiver = _request(previous)
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.update_request_metadata({"r": receiver}, {"r": metadata}, model_mode="generation")
+    assert receiver.prompt_token_ids == []
+    assert receiver._all_token_ids == []
+    assert receiver.num_prompt_tokens == receiver.num_computed_tokens == 0
+    assert "r" in coordinator.input_terminal_req_ids
+    assert "r" not in manager.code_prompt_token_ids
+
+
+@pytest.mark.parametrize("empty", [[], torch.empty(0, dtype=torch.long)])
+def test_explicit_empty_codes_clear_generation_prompt(empty):
+    request = _request([17, 18, 19])
+    request._output_token_ids = [23]
+    request.num_computed_tokens = 3
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.update_request_metadata(
+        {"r": request}, {"r": {"code_predictor_codes": empty, "input_terminal": True}}, model_mode="generation"
+    )
+    assert request.prompt_token_ids == request._all_token_ids == request._output_token_ids == []
+    assert request.num_prompt_tokens == request.num_computed_tokens == 0
+    assert "r" in coordinator.input_terminal_req_ids
+
+
+@pytest.mark.parametrize("metadata", [{"input_terminal": True}, {"code_predictor_codes": None}])
+def test_absent_codes_preserve_generation_prompt(metadata):
+    request = _request([17, 18, 19])
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.update_request_metadata({"r": request}, {"r": metadata}, model_mode="generation")
+    assert request.prompt_token_ids == request._all_token_ids == [17, 18, 19]
+
+
+def test_ar_prompt_is_not_replaced_by_empty_audio_codes():
+    request = _request([17, 18, 19])
+    coordinator = OmniSchedulingCoordinator(stage_id=1)
+    coordinator.update_request_metadata(
+        {"r": request}, {"r": {"code_predictor_codes": [], "input_terminal": True}}, model_mode="ar"
+    )
+    assert request.prompt_token_ids == request._all_token_ids == [17, 18, 19]

tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Progressive Local codec chunks preserve frames and agree with graph shapes."""
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from unittest.mock import Mock
+
+import pytest
+import torch
+from transformers import PretrainedConfig
+from vllm.config import CompilationConfig, ModelConfig, SchedulerConfig, VllmConfig
+
+from vllm_omni.data_entry_keys import OmniPayloadStruct
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.model_executor.stage_input_processors.moss_tts import talker2codec_raw_async_chunk
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class _Connector:
+    config: dict[str, dict[str, object]]
+
+
+@dataclass
+class _TransferManager:
+    connector: _Connector
+    code_prompt_token_ids: dict[str, list[object]] = field(default_factory=lambda: defaultdict(list))
+    request_payload: dict[str, object] = field(default_factory=dict)
+    put_req_chunk: dict[str, int] = field(default_factory=lambda: defaultdict(int))
+    ramp_chunk_count: dict[str, int] = field(default_factory=lambda: defaultdict(int))
+
+
+@dataclass
+class _Request:
+    external_req_id: str
+
+
+def _manager(ramp: object = None, *, first: int = 1) -> _TransferManager:
+    extra = {"initial_codec_chunk_frames": first, "codec_chunk_frames": 15, "codec_chunk_ramp": ramp}
+    return _TransferManager(_Connector({"extra": extra}))
+
+
+def _emit(
+    manager: _TransferManager, req_id: str, value: int | None, finished: bool = False
+) -> OmniPayloadStruct | None:
+    output = None if value is None else {"codes": {"audio": torch.full((1, 4), value)}}
+    payload = talker2codec_raw_async_chunk(manager, output, _Request(req_id), finished)
+    if payload is not None:
+        manager.put_req_chunk[req_id] += 1
+        manager.ramp_chunk_count[req_id] += 1
+    return payload
+
+
+@pytest.mark.parametrize(
+    "ramp,first,expected",
+    [
+        (None, 1, [1, 15, 15, 9]),
+        ([1, 2, 4, 8, 15], 1, [1, 2, 4, 8, 15, 10]),
+        ([1, 4, 15], 7, [1, 4, 15, 15, 5]),
+        ([4, 8, 15], 1, [4, 8, 15, 13]),
+        ([1, 20, 15], 1, [1, 20, 15, 4]),
+        ("1,4,15", 1, [1, 4, 15, 15, 5]),
+        ([1, 0, 15], 1, [1, 15, 15, 9]),
+        ([1], 1, [1, 15, 15, 9]),
+        ("invalid", 1, [1, 15, 15, 9]),
+    ],
+)
+def test_order_and_final_flush(ramp: object, first: int, expected: list[int]) -> None:
+    manager = _manager(ramp, first=first)
+    packets = []
+    for index in range(40):
+        packet = _emit(manager, "a", index, finished=index == 39)
+        if packet is not None:
+            packets.append(packet)
+    assert [p.meta.codec_chunk_frames for p in packets] == expected
+    actual = torch.cat([p.codes.audio.reshape(4, -1).T for p in packets])
+    assert torch.equal(actual, torch.arange(40).unsqueeze(1).expand(-1, 4))
+    assert all(not bool(p.meta.finished) for p in packets[:-1])
+    assert bool(packets[-1].meta.finished)
+    assert "a" not in manager.code_prompt_token_ids
+
+
+def test_requests_progress_independently_and_empty_finish() -> None:
+    manager = _manager([1, 2, 4, 8, 15])
+    for name in ("a", "b"):
+        first = _emit(manager, name, 0)
+        assert first is not None
+        assert first.meta.codec_chunk_frames == 1
+    assert _emit(manager, "a", 1) is None
+    final_a = _emit(manager, "a", None, finished=True)
+    assert final_a is not None
+    assert final_a.meta.codec_chunk_frames == 1
+    assert bool(final_a.meta.finished)
+    final_b = _emit(manager, "b", None, finished=True)
+    assert final_b is not None
+    assert final_b.meta.codec_chunk_frames == 0
+    assert bool(final_b.meta.finished)
+    assert final_b.codes.audio.numel() == 0
+    assert not manager.code_prompt_token_ids
+
+
+@pytest.mark.parametrize(
+    "ramp,first,expected,max_step,effective_first",
+    [
+        (None, 1, [1, 15], 15, 1),
+        ([1, 2, 4, 8, 15], 1, [1, 2, 4, 8, 15], 15, 1),
+        ([1, 20, 15], 1, [1, 15, 20], 20, 1),
+        ([4, 8, 15], 1, [4, 8, 15], 15, 4),
+        ([1, 4, 15], 7, [1, 4, 15], 15, 1),
+        ([1, 0, 15], 1, [1, 15], 15, 1),
+    ],
+)
+def test_graph_shapes_and_fast_path_follow_sender(
+    ramp: object, first: int, expected: list[int], max_step: int, effective_first: int
+) -> None:
+    connector = _manager(ramp, first=first).connector.config
+    connector["extra"]["codec_first_chunk_fast_path"] = 1
+    config = Mock(
+        spec=VllmConfig,
+        model_config=Mock(
+            spec=ModelConfig,
+            hf_config=PretrainedConfig(model_type="moss_tts_local"),
+            async_chunk=True,
+            use_v2_model_runner=True,
+            enforce_eager=False,
+            stage_connector_config=connector,
+        ),
+        scheduler_config=Mock(spec=SchedulerConfig, max_num_seqs=8),
+        compilation_config=Mock(spec=CompilationConfig, cudagraph_capture_sizes=[1, 2, 4, 8]),
+    )
+    codec = MossTTSCodecDecoder(vllm_config=config)
+    assert codec._streaming_graph_frame_sizes == expected
+    assert codec._stream_max_step_frames == max_step
+    assert codec._first_chunk_fast
+    assert codec._initial_stream_chunk_frames == effective_first

tests/worker/test_omni_connector_mixin.py

@@ -39,6 +39,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class MockConnector:
+    sender_host: str | None = None
+    sender_zmq_port: int | None = None
+
     """In-memory connector for testing (mimics OmniConnectorBase)."""
 
     def __init__(self, stage_id: int = 0):
@@ -187,7 +190,7 @@ def test_init_payload_connector_ownership(role, custom_func, expected):
 
 @pytest.mark.parametrize("async_chunk", [False, True])
 @pytest.mark.parametrize("rank", [0, 1])
-def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
+def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank, mocker):
     host = MixinHost()
     host.init_omni_connectors(_make_model_config(async_chunk=async_chunk))
     # No background threads: deterministic interleaved polling.
@@ -196,7 +199,7 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
     connector = host._omni_connector
     connector.sender_host = "configured-host"
     connector.sender_zmq_port = 50051
-    connector.get = MagicMock(return_value=None)
+    get = mocker.patch.object(connector, "get", return_value=None)
     first = _make_request("r1", "external1")
     first.payload_sender_info = {"host": "producer-a", "zmq_port": "50101"}
     second = _make_request("r2", "external2")
@@ -209,14 +212,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
             for req_id in ("r2", "r1", "r2"):
                 assert not host._poll_single_request(req_id)
             if rank == 1:
-                connector.get.assert_not_called()
+                get.assert_not_called()
             else:
-                assert connector.get.call_args_list == [
+                assert get.call_args_list == [
                     unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}),
                     unittest.mock.call("0", "1", "external1_0_0", {"source_host": "producer-a", "source_port": 50101}),
                     unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}),
                 ]
-                connector.get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1)
+                get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1)
                 assert host._poll_single_request("r2")
                 assert "r2" not in host._pending_load_reqs
                 assert "r1" in host._pending_load_reqs
@@ -225,14 +228,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
             assert not host._pending_load_reqs
             # Reused internal ID without sender info cannot inherit either endpoint.
             replacement = _make_request("r1", "external-new")
-            connector.get.reset_mock(return_value=True)
-            connector.get.return_value = None
+            get.reset_mock(return_value=True)
+            get.return_value = None
             host.register_chunk_recv(replacement)
             assert not host._poll_single_request("r1")
             if rank == 0:
-                connector.get.assert_called_once_with("0", "1", "external-new_0_0")
+                get.assert_called_once_with("0", "1", "external-new_0_0")
             else:
-                connector.get.assert_not_called()
+                get.assert_not_called()
             assert connector.sender_host == "configured-host"
             assert connector.sender_zmq_port == 50051
             assert first.payload_sender_info == {"host": "producer-a", "zmq_port": "50101"}
@@ -1437,6 +1440,58 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
 
         host.shutdown_omni_connectors()
 
+    def test_first_chunk_hook_sees_only_chunk_zero_before_staging(self):
+        host = MixinHost()
+        host.init_omni_connectors(
+            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"),
+        )
+        host._omni_connector = MagicMock()
+        host._stage_id = 1
+        host._async_chunk = True
+        host._model_mode = "gen"
+        seen = []
+
+        def hook(req_id, handle, payload):
+            # The chunk is not yet visible to the scheduler when the hook runs.
+            seen.append((req_id, handle.external_req_id, payload["codes"]["audio"], "r1" in host._finished_load_reqs))
+
+        host.set_first_chunk_hook(hook)
+        host.register_chunk_recv(_make_request("r1", "ext-r1"))
+        chunk = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}}
+        host._omni_connector.get.return_value = (chunk, 1)
+        self.assertTrue(host._poll_single_request("r1"))
+        self.assertEqual(seen, [("r1", "ext-r1", [1, 2], False)])
+        self.assertIn("r1", host._finished_load_reqs)
+
+        host.get_omni_connector_output()
+        host._local_stage_payload_cache.clear()
+        host._local_request_metadata.clear()
+        host._omni_connector.get.return_value = ({"codes": {"audio": [3, 4]}, "meta": {}}, 1)
+        self.assertTrue(host._poll_single_request("r1"))
+        self.assertEqual(len(seen), 1)
+
+        host.shutdown_omni_connectors()
+
+    def test_first_chunk_hook_skips_terminal_first_chunk(self):
+        host = MixinHost()
+        host.init_omni_connectors(
+            model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"),
+        )
+        host._omni_connector = MagicMock()
+        host._stage_id = 1
+        host._async_chunk = True
+        host._model_mode = "gen"
+        hook = MagicMock()
+        host.set_first_chunk_hook(hook)
+        host.register_chunk_recv(_make_request("r1"))
+        host._omni_connector.get.return_value = (
+            {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(True)}},
+            1,
+        )
+        self.assertTrue(host._poll_single_request("r1"))
+        hook.assert_not_called()
+        host.shutdown_omni_connectors()
+
     def test_non_ar_recv_waits_for_scheduler_handoff_before_fetching_next_chunk(self):
         host = MixinHost()
         host.init_omni_connectors(
@@ -1480,6 +1535,39 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
         host.shutdown_omni_connectors()
 
 
+@pytest.mark.parametrize("cancel_at", ["connector_get", "first_chunk_hook"])
+def test_cancelled_chunk_receive_does_not_recreate_delivery_state(mocker, cancel_at):
+    host = MixinHost()
+    host.init_omni_connectors(_make_model_config(async_chunk=True, worker_type="gen"))
+    host._stage_id = 1
+    host._omni_connector = mocker.MagicMock()
+    host.register_chunk_recv(_make_request("r1"))
+    seen = []
+    payload = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}}
+
+    def get(*args, **kwargs):
+        if cancel_at == "connector_get":
+            host.cleanup_finished_request("r1")
+        return payload, 1
+
+    def hook(req_id, handle, chunk):
+        seen.append(req_id)
+        host.cleanup_finished_request(req_id)
+
+    host._omni_connector.get.side_effect = get
+    host.set_first_chunk_hook(hook)
+    try:
+        assert not host._poll_single_request("r1")
+        assert seen == (["r1"] if cancel_at == "first_chunk_hook" else [])
+        assert host.get_local_stage_payload("r1") is None
+        assert "r1" not in host._get_req_chunk
+        assert "r1" not in host._pending_load_reqs
+        output = host.get_omni_connector_output()
+        assert output.chunk_ready_req_ids == set() and output.request_metadata == {}
+    finally:
+        host.shutdown_omni_connectors()
+
+
 class TestRankAwareKVRouting(unittest.TestCase):
     def _make_host(self, *, from_tp: int, to_tp: int, local_rank: int) -> MixinHost:
         host = MixinHost()

tests/worker_v2/test_first_audio_sender.py

@@ -235,3 +235,19 @@ def test_sender_copy_failure_reports_all_prepared_routes():
     assert client == 3
     assert [(o.request_id, o.finish_reason) for o in outputs.outputs] == [("r", FinishReason.ERROR)]
     assert output_queue.empty()
+
+
+def test_final_codec_delivery_omits_upstream_marker_and_freezes_route():
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+
+    outputs: queue.Queue = queue.Queue()
+    scheduler = _scheduler(r=3)
+    delivery = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False).prepare(["r"])
+    scheduler.requests.clear()
+    pcm = torch.ones(4)
+    delivery(["r"], [pcm], torch.tensor(24000))
+    client, batch = outputs.get_nowait()
+    assert client == 3
+    assert FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
+    assert batch.outputs[0].multimodal_output["model_outputs"] is pcm
+    assert batch.outputs[0].finish_reason is None

tests/worker_v2/test_omni_ar_model_runner.py

@@ -15,6 +15,7 @@ from vllm.v1.outputs import RoutedExpertsTensors
 from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
 
 import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
 from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput
 from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot
 
@@ -144,6 +145,7 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     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
+    runner._last_multimodal_snapshot_slot = None
     waited: list[object] = []
     runner.main_stream = SimpleNamespace(wait_event=waited.append)
     source = torch.tensor([[7, 8]], dtype=torch.long)
@@ -161,6 +163,35 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     assert waited == [copy_event]
 
 
+def test_producer_snapshot_is_not_repacked() -> 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
+    runner._last_multimodal_snapshot_slot = None
+    source = torch.tensor([[7, 8]], dtype=torch.float32)
+    snapshot = pack_output_snapshot({"audio": source}, {}, max_buckets=1)
+    assert isinstance(snapshot, PackedOutputSnapshot)
+
+    retained = runner._retain_multimodal_outputs(snapshot)
+
+    assert retained is snapshot
+    assert runner._last_multimodal_snapshot_slot is None
+    assert runner._async_mm_snapshot_pending == [False]
+
+
+def test_fresh_per_step_outputs_are_not_repacked() -> None:
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.model_config = SimpleNamespace(async_chunk=True)
+    runner.model = SimpleNamespace(mm_outputs_fresh_per_step=True)
+    runner._last_multimodal_snapshot_slot = None
+    outputs = {"model_outputs": torch.ones(2, 4)}
+
+    assert runner._retain_multimodal_outputs(outputs) is outputs
+    assert runner._last_multimodal_snapshot_slot is None
+
+
 def test_snapshot_slots_bounded_by_shape_and_packed_grouping_isolation(monkeypatch) -> None:
     slot: dict[tuple[Any, ...], torch.Tensor] = {}
     omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(1, 2), slot)

tests/worker_v2/test_omni_model_state.py

@@ -224,6 +224,61 @@ def test_static_decode_embeddings_refresh_from_input_ids():
     assert OmniModelState._preprocess_result_needs_writeback(original, original.view_as(original)) is True
 
 
+def test_moss_local_decode_runs_depth_predictor_and_routes_eos(mocker):
+    """Exercise MRV2 dispatch through the real Local hook, output and logits."""
+    from vllm.sampling_params import SamplingParams
+    from vllm.v1.worker.gpu.states import RequestState
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import MossTTSLocalDepthTransformer
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    model.model = torch.nn.Module()
+    model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=torch.zeros(8, 4))
+    model.hidden_size = 4
+    model.n_vq = 2
+    model.audio_pad_token_id = 8
+    model.audio_assistant_slot_token_id = 2
+    model.im_end_token_id = 3
+    model.text_vocab_size = 8
+    model.talker_mtp_output_key = ("audio_codes", "current")
+    model.talker_mtp_graph_safe = False
+    model.gpu_resident_buffer_keys = set()
+    model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
+    model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).float()
+    frame = mocker.Mock(return_value=(torch.tensor([True, False]), torch.tensor([[1, 2], [3, 4]])))
+    model.local_transformer = mocker.Mock(spec=MossTTSLocalDepthTransformer, generate_frame=frame)
+
+    state = _make_state(max_num_reqs=2, has_preprocess=True, have_multimodal_outputs=True)
+    state.model = model
+    for idx in range(2):
+        state.intermediate_buffer.buffers[idx] = {
+            "req_id": f"r{idx}",
+            "audio_state": {"is_stopping": False},
+            "hidden_states": {"last": torch.full((4,), float(idx + 1))},
+        }
+    batch = _DummyInputBatch([1, 0], num_computed_tokens_cpu=[1, 1])
+    inputs = {"input_ids": torch.tensor([2, 2]), "inputs_embeds": torch.zeros(2, 4)}
+    request_state = mocker.Mock(spec=RequestState, prompt_len=np.array([1, 1]), num_computed_tokens=None)
+    state.run_preprocess(batch, inputs, request_state)
+
+    frame.assert_called_once()
+    torch.testing.assert_close(frame.call_args.args[0], torch.tensor([[2.0] * 4, [1.0] * 4]))
+    assert frame.call_args.kwargs["temperature"] == 1.7
+    assert frame.call_args.kwargs["top_k"] == 25
+    assert frame.call_args.kwargs["top_p"] == 0.8
+    torch.testing.assert_close(inputs["inputs_embeds"], torch.tensor([[1.0] * 4, [0.0] * 4]))
+    _, payload = state.postprocess_model_output(torch.zeros(2, 4), batch, request_state)
+    assert [codes.tolist() for codes in payload["codes"]["audio"]] == [[[1, 2]], [[8, 8]]]
+    assert model.compute_logits(torch.zeros(2, 4)).argmax(-1).tolist() == [2, 3]
+    # An explicit local seed must reach the request-owned MRV2 generator.
+    params = SamplingParams(extra_args={"tts_local_seed": 17}, seed=99)
+    assert state._get_mtp_generator("seeded", params, torch.device("cpu")).initial_seed() == 17
+
+
 @pytest.mark.parametrize("owned", [False, True])
 def test_batched_postprocess_gpu_snapshot_writeback(owned):
     # batch=[1, 0]: last-token indices follow the reordered batch, the scalar

vllm_omni/config/environment_variable_inventory.py

@@ -182,6 +182,24 @@ _MODEL_PROMOTE = (
     "VLLM_GEPARD_FIRST_CHUNK_FRAMES",
     "VLLM_GEPARD_LOOKBACK_FRAMES",
     "VLLM_OMNI_FISH_KVCACHE_ATTN",
+    "VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER",
+    "VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN",
+    "VLLM_OMNI_MOSS_REF_ATTN",
+    "VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER",
+    "VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS",
+    "VLLM_OMNI_MOSS_REF_CACHE_ROPE",
+    "VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR",
+    "VLLM_OMNI_MOSS_REF_COMPILE",
+    "VLLM_OMNI_MOSS_REF_ENCODER_WORKERS",
+    "VLLM_OMNI_MOSS_REF_FA_VERSION",
+    "VLLM_OMNI_MOSS_REF_GPU_PREP",
+    "VLLM_OMNI_MOSS_REF_GRAPHS",
+    "VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS",
+    "VLLM_OMNI_MOSS_REF_INFLIGHT",
+    "VLLM_OMNI_MOSS_REF_SHARED_ENCODER",
+    "VLLM_OMNI_MOSS_REF_SINGLETON_BUCKET_MS",
+    "VLLM_OMNI_MOSS_REF_SKIP_PADDED_QUERY_MASK",
+    "VLLM_OMNI_MOSS_REF_STREAM_PRIORITY",
     "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
     "VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS",
     "VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES",

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -3,7 +3,9 @@
 
 from __future__ import annotations
 
+import math
 import os
+import queue
 import time
 from collections import defaultdict
 from typing import TYPE_CHECKING, Any
@@ -58,6 +60,48 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._init_omni_io_scheduling_state()
         self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False))
         self._pending_finish_reqs: list[Request] = []
+        connector_config = getattr(model_config, "stage_connector_config", None)
+        extra = (
+            connector_config.get("extra", connector_config)
+            if isinstance(connector_config, dict)
+            else getattr(connector_config, "extra", None)
+        ) or {}
+        self._generation_min_batch_size = int(extra.get("generation_min_batch_size", 1))
+        self._generation_max_wait_s = float(extra.get("generation_max_wait_ms", 0)) / 1000
+        self._generation_coalescing_policy = extra.get("generation_coalescing_policy", "fixed")
+        self._generation_max_regular_batch = int(extra.get("generation_max_regular_batch", 0) or 0)
+        if self._generation_max_regular_batch < 0:
+            raise ValueError("generation_max_regular_batch must be nonnegative")
+        if self._generation_coalescing_policy not in ("fixed", "idle_wait"):
+            raise ValueError("generation_coalescing_policy must be fixed or idle_wait")
+        if not 1 <= self._generation_min_batch_size <= self.max_num_running_reqs:
+            raise ValueError("generation_min_batch_size must be in [1, max_num_seqs]")
+        if not math.isfinite(self._generation_max_wait_s) or self._generation_max_wait_s < 0:
+            raise ValueError("generation_max_wait_ms must be finite and nonnegative")
+        if self._generation_max_wait_s and self._generation_min_batch_size > 1:
+            parallel = self.vllm_config.parallel_config
+            if not self._native_data_plane:
+                # CUDA-oriented profiles can inherit a platform's V1 runner.
+                # Their connector extras survive that override, but batching
+                # must revert to the fallback runner's default scheduling.
+                logger.warning(
+                    "Generation batch waiting and regular-batch limit ignored: "
+                    "requires native MRV2; using default generation scheduling."
+                )
+                self._generation_min_batch_size = 1
+                self._generation_max_wait_s = 0.0
+                self._generation_max_regular_batch = 0
+            elif not (
+                self._retains_state_across_chunks
+                and parallel.tensor_parallel_size == parallel.pipeline_parallel_size == 1
+            ):
+                raise ValueError("Generation batch waiting requires a stateful native MRV2 TP1/PP1 stage")
+            else:
+                logger.info(
+                    "Generation input coalescing: target batch=%d, max wait=%.3f ms",
+                    self._generation_min_batch_size,
+                    self._generation_max_wait_s * 1000,
+                )
         self._first_chunk_express = os.environ.get("VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS") == "1"
         self._last_step_express = False
         # Streams that already had a chunk scheduled; the rest await their first.
@@ -76,6 +120,101 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 self._express_min_slack_s,
             )
 
+    def _drain_omni_connector_outputs(self):
+        self._generation_defer_batch = False
+        outputs = super()._drain_omni_connector_outputs()
+        max_wait = getattr(self, "_generation_max_wait_s", 0)
+        target = getattr(self, "_generation_min_batch_size", 1)
+        coordinator = self.input_coordinator
+        if not max_wait or target <= 1 or coordinator is None or self._pause_state != PauseState.UNPAUSED:
+            return outputs
+
+        # Count distinct runnable streams, not notifications, padding rows, or
+        # chunks whose previous execution still owns their state. Waiting here
+        # precedes allocation and metadata replacement. Do not schedule a batch
+        # and then discard it to try to accumulate more input.
+        ready = set(coordinator.requests_with_ready_chunks)
+        terminal = set(coordinator.finished_requests) | set(coordinator.input_terminal_req_ids)
+        deadline = getattr(self, "_generation_batch_deadline", None)
+        if deadline is None:
+            deadline = time.monotonic() + max_wait
+        in_flight = any(
+            not request.is_finished() and request.num_in_flight_tokens > 0 for request in self.requests.values()
+        )
+        new_outputs = outputs
+        waited_without_input = False
+        while True:
+            for output in new_outputs:
+                ready.update(output.chunk_ready_req_ids)
+                terminal.update(output.chunk_finished_req_ids)
+                terminal.update(
+                    req_id for req_id, metadata in output.request_metadata.items() if metadata.get("input_terminal")
+                )
+            runnable = {
+                req_id
+                for req_id in ready | terminal
+                if (request := self.requests.get(req_id)) is not None
+                and not request.is_finished()
+                and request.num_in_flight_tokens == 0
+            }
+            if runnable and waited_without_input:
+                # Idle waiting must not eat the ready batch's coalescing budget.
+                # Start that budget at its first arrival, once only. The two
+                # bounded windows can delay control processing by at most 2W.
+                deadline = time.monotonic() + max_wait
+                waited_without_input = False
+            # A first/final chunk participates in the same bounded window.
+            # At high concurrency almost every batch contains one; bypassing
+            # the entire batch in that case defeats coalescing. Low concurrency
+            # still dispatches immediately, including single-request tails.
+            if not runnable:
+                # With no output to retire and all receivers already parked,
+                # a wakeable inbox wait avoids repeatedly executing empty steps.
+                # Admission/registration/terminal work must reach the runner
+                # first; waiting for a not-yet-registered receiver would stall it.
+                can_wait_idle = (
+                    getattr(self, "_generation_coalescing_policy", "fixed") == "idle_wait"
+                    and not in_flight
+                    and bool(self.requests)
+                    and not getattr(coordinator, "pending_chunk_registrations", ())
+                    and not getattr(coordinator, "pending_input_registrations", ())
+                    and not getattr(self, "_pending_data_plane_terminal_req_ids", ())
+                    and all(request.status == RequestStatus.WAITING_FOR_CHUNK for request in self.requests.values())
+                )
+                if not can_wait_idle:
+                    break
+                waited_without_input = True
+            elif len(runnable) >= target or len(self.requests) < target:
+                break
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                break
+            if in_flight:
+                # Let EngineCore retire the prior output before waiting for
+                # more input. Blocking this thread also blocks the in-flight
+                # fence that makes those streams eligible for the next batch.
+                # Emit a normal empty scheduler step (including connector and
+                # cancellation bookkeeping), never discard an allocated step.
+                self._generation_defer_batch = True
+                self._generation_batch_deadline = deadline
+                return outputs
+            try:
+                output = self._omni_connector_output_inbox.get(timeout=remaining)
+            except queue.Empty:
+                break
+            outputs.append(output)
+            new_outputs = [output]
+            # Drain a burst before re-evaluating, retaining notification order.
+            while True:
+                try:
+                    output = self._omni_connector_output_inbox.get_nowait()
+                except queue.Empty:
+                    break
+                outputs.append(output)
+                new_outputs.append(output)
+        self._generation_batch_deadline = None
+        return outputs
+
     def _build_generation_scheduler_output(
         self,
         new_reqs_data: list[OmniNewRequestData],
@@ -256,6 +395,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._drop_aborted_queued_requests()
         self._requeue_completed_native_chunks()
         self._process_pending_omni_inputs(model_mode="generation")
+        if getattr(self, "_generation_defer_batch", False):
+            token_budget = 0
+        max_regular = getattr(self, "_generation_max_regular_batch", 0)
+        if max_regular:
+            execution_batch_size = min(execution_batch_size, max_regular)
         self._drop_aborted_queued_requests()
         self._resync_streaming_input_counter()
         async_chunk_transport = self._async_chunk_transport_enabled()

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -331,13 +331,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             # ready-chunk marker remains until scheduler admission succeeds.
             replaced_ids.discard(request_id)
 
-    def _consume_pending_connector_output(self, model_mode: str) -> None:
-        """Drain ``self._latest_omni_connector_output`` into the coordinator.
-
-        Called at the top of every ``schedule()`` cycle.  Identical between
-        AR and generation schedulers except for the ``model_mode`` argument
-        forwarded to ``update_request_metadata``.
-        """
+    def _drain_omni_connector_outputs(self) -> list[OmniConnectorOutput]:
+        """Collect control messages in order, without changing request state."""
         connector_outputs: list[OmniConnectorOutput] = []
         inbox = getattr(self, "_omni_connector_output_inbox", None)
         if inbox is not None:
@@ -350,6 +345,11 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         self._latest_omni_connector_output = None
         if connector_output is not None:
             connector_outputs.append(connector_output)
+        return connector_outputs
+
+    def _consume_pending_connector_output(self, model_mode: str) -> None:
+        """Drain input notifications into the coordinator on the scheduler thread."""
+        connector_outputs = self._drain_omni_connector_outputs()
         input_coordinator = getattr(self, "input_coordinator", None)
         if input_coordinator is None:
             return

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -379,14 +379,17 @@ class OmniSchedulingCoordinator:
                             )
 
             if model_mode != "ar":
-                new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes"))
+                new_codes = metadata.get("code_predictor_codes")
+                new_ids = self._flatten_prompt_token_ids(new_codes)
                 runtime_seed = None
                 if "left_context_size" in metadata:
                     runtime_seed = {
                         "meta": {"left_context_size": metadata["left_context_size"]},
                     }
                 request._omni_initial_model_buffer = runtime_seed
-                if new_ids:
+                # An explicit empty snapshot clears a previous codec chunk
+                # or prewarmed placeholder. Missing codes leave it intact.
+                if new_codes is not None:
                     request.prompt_token_ids = new_ids
                     request.num_prompt_tokens = len(new_ids)
                     request._all_token_ids.clear()

vllm_omni/deploy/moss_tts_local.yaml

@@ -18,6 +18,9 @@ connectors:
       connector_get_max_wait: 300
       initial_codec_chunk_frames: 1   # 1 frame * 80 ms/frame = 80 ms first audio chunk
       codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s steady chunks
+      # Optional startup ramp; its first entry overrides initial chunk size.
+      # Extra graph shapes/calls trade codec cost and latency for less underrun.
+      # codec_chunk_ramp: [1, 4, 15]
       codec_left_context_frames: 0
       # Chunk ramp (opt-in): gradual chunk-size ladder replacing the
       # initial_codec_chunk_frames -> codec_chunk_frames cliff. The last

vllm_omni/deploy/moss_tts_local_mrv2.yaml

@@ -0,0 +1,30 @@
+# Opt-in CUDA MRV2 for MOSS-TTS-Local-Transformer-v1.5.
+# Keep the base profile's batching, chunk geometry and V1 default.
+base_config: moss_tts_local.yaml
+model_runner: v2
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 0
+        # Bound prefill work so new prompts do not monopolize streaming decode.
+        max_num_batched_tokens: 512
+      - stage_id: 1
+        compilation_config:
+          # The codec owns CUDA Graph capture/replay. Skip Inductor compilation
+          # of its stateful decoder; plain CUDA graphs remain enabled.
+          mode: 0
+          # FULL works without Inductor; the default FULL_AND_PIECEWISE would
+          # normalize to NONE and erase the codec's capture sizes.
+          cudagraph_mode: FULL
+          # Platform overrides replace compilation_config as a whole.
+          cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64]
+          cudagraph_num_of_warmups: 1
+  npu:
+    model_runner: v1
+  xpu:
+    model_runner: v1
+  rocm:
+    model_runner: v1
+  musa:
+    model_runner: v1

vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml

@@ -0,0 +1,53 @@
+# Opt-in CUDA profile for MOSS-TTS-Local-Transformer-v1.5 on a large-memory GPU.
+# Both stages share GPU 0. Validated with one H200 and Seed-TTS EN1088.
+base_config: moss_tts_local_mrv2.yaml
+
+connectors:
+  shm:
+    extra:
+      # Coalesce ready codec streams before submitting a small execution batch.
+      # This is a dispatch target, not a cap: execution capacity remains 128.
+      generation_min_batch_size: 32
+      generation_max_wait_ms: 12
+      # Optional; preserves the first/steady sizes while adding a middle chunk.
+      # codec_chunk_ramp: [1, 4, 15]
+
+# Optional codec attention-ring headroom (see recipes/OpenMOSS/MOSS-TTS.md):
+# sizes each ring as context + one chunk so chunked streaming matches
+# whole-sequence decoding. Requires additional state/graph memory; evaluate
+# throughput and audio quality before enabling for your context and capacity.
+# Add codec_ring_headroom: 1 under connectors.shm.extra above.
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 0
+        max_num_seqs: 128
+        max_num_batched_tokens: 512
+        hf_overrides:
+          mrv2_gpu_slot_state: true
+          # Throughput-oriented pair; measure first-packet latency for your load.
+          mrv2_batch_prefill: true
+          mrv2_direct_tokens: true
+        # Mixed prefill/decode fits the same bounded token budget as capture.
+        compilation_config:
+          mode: 3
+          cudagraph_mode: FULL
+          max_cudagraph_capture_size: 512
+          cudagraph_num_of_warmups: 1
+        # Leave room for the codec state pool and its batch-128 CUDA graphs.
+        kv_cache_memory_bytes: 34359738368  # 32 GiB
+      - stage_id: 1
+        max_num_seqs: 128
+        hf_overrides:
+          codec_attention_backend: triton_slot
+        compilation_config:
+          mode: 3
+          backend: inductor
+          # The codec owns capture/replay after Inductor compiles its decoder.
+          cudagraph_mode: FULL
+          cudagraph_capture_sizes: [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+          cudagraph_num_of_warmups: 1
+          inductor_compile_config:
+            combo_kernels: false
+            benchmark_combo_kernel: false

vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml

@@ -0,0 +1,66 @@
+# Opt-in CUDA profile for MOSS-TTS-Local-Transformer-v1.5 tuned for first-packet
+# latency at client concurrency 128 on one H200 (Seed-TTS EN1088, closed loop).
+# Both stages share GPU 0. Use moss_tts_local_mrv2_high_concurrency.yaml for
+# throughput-first serving.
+#
+# Optional launch-level settings (see the full command in recipes/OpenMOSS/MOSS-TTS.md):
+#   - CUDA MPS for the GPU (stage 0 below lowers its own MPS client priority);
+#   - VLLM_OMNI_EVENT_DRIVEN_ORCH=1;
+#   - --api-server-count 4 with VLLM_WORKER_MULTIPROC_METHOD=spawn and
+#     VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/<dir> so frontends share reference codes;
+#   - VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 so Stage 1 picks up first chunks sooner.
+base_config: moss_tts_local_mrv2.yaml
+
+connectors:
+  shm:
+    extra:
+      generation_min_batch_size: 32
+      generation_max_wait_ms: 12
+      # Stage 1 decodes each first chunk from its receive thread on a small
+      # dedicated graph set and emits it directly, so first chunks no longer
+      # wait for the engine loop or a regular codec step.
+      codec_first_chunk_fast_path: 1
+      # Regular codec replays wait for an in-flight first-chunk decode.
+      codec_first_chunk_gate: 1
+      # Short regular codec steps bound contention with first-chunk decodes.
+      generation_max_regular_batch: 16
+      # Optional ramp. Both the first-chunk graph and regular graphs follow it.
+      # codec_chunk_ramp: [1, 4, 15]
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 0
+        max_num_seqs: 128
+        # Larger prefill budget admits bursts of new requests in fewer steps.
+        max_num_batched_tokens: 2048
+        hf_overrides:
+          mrv2_gpu_slot_state: true
+          mrv2_batch_prefill: true
+          mrv2_direct_tokens: true
+          # Sample each frame right after the forward that produced it.
+          mrv2_eager_mtp: true
+        compilation_config:
+          mode: 3
+          cudagraph_mode: FULL
+          max_cudagraph_capture_size: 2048
+          cudagraph_num_of_warmups: 1
+        kv_cache_memory_bytes: 34359738368  # 32 GiB
+        env:
+          # Under MPS, let codec kernels schedule ahead of the talker's.
+          CUDA_MPS_CLIENT_PRIORITY: "1"
+      - stage_id: 1
+        max_num_seqs: 128
+        # Keep synchronous scheduling; benchmark async scheduling separately.
+        async_scheduling: false
+        hf_overrides:
+          codec_attention_backend: triton_slot
+        compilation_config:
+          mode: 3
+          backend: inductor
+          cudagraph_mode: FULL
+          cudagraph_capture_sizes: [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+          cudagraph_num_of_warmups: 1
+          inductor_compile_config:
+            combo_kernels: false
+            benchmark_combo_kernel: false

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py

@@ -1169,7 +1169,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 self._payload_finished(payload_data),
             )
 
-        self._get_req_chunk[req_id] += 1
+        with self._lock:
+            if self._async_chunk and request is not None and req_id not in self._pending_load_reqs:
+                # A connector get can return after cancellation removed its
+                # receiver. Do not recreate delivery state for that request.
+                return False
+            self._get_req_chunk[req_id] += 1
 
         if self._async_chunk:
             is_finished = self._payload_finished(payload_data)
@@ -1180,8 +1185,16 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 if not self._payload_value_has_content(new_ids) and not is_finished:
                     return False
                 payload_consumable = self._payload_is_consumable(payload_data)
+                first_chunk_hook = getattr(self, "_first_chunk_hook", None)
+                if first_chunk_hook is not None and chunk_id == 0 and not is_finished and payload_consumable:
+                    first_chunk_hook(req_id, request, payload_data)
 
             with self._lock:
+                if request is not None and req_id not in self._pending_load_reqs:
+                    # The first-chunk hook may have overlapped cancellation.
+                    # Its slot cleanup is ordered by the model; this payload
+                    # must not republish readiness after receive teardown.
+                    return False
                 if self._model_mode == "ar":
                     # Accumulation, staging, and model-side consume/ack share
                     # this lock. Keeping the transition atomic prevents the
@@ -1723,3 +1736,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         sink: Callable[[OmniConnectorOutput], None] | None,
     ) -> None:
         self._omni_connector_output_sink = sink
+
+    def set_first_chunk_hook(self, hook: Callable[[str, Any, OmniPayload], bool]) -> None:
+        """Claim eligible first chunks before publishing them to the model loop."""
+        self._first_chunk_hook = hook

vllm_omni/engine/stage_engine_core_proc.py

@@ -80,9 +80,12 @@ def _signal_exit_code(signum: int) -> int:
 
 
 def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: Any) -> bool:
-    """Let a TP1 in-process runner that decodes first audio deliver it as its own output."""
+    """Bind in-process first audio to the generic engine output sink."""
     if not isinstance(model_executor, UniProcExecutor):
         return False
+    parallel_config = model_executor.vllm_config.parallel_config
+    if parallel_config.tensor_parallel_size != 1 or parallel_config.pipeline_parallel_size != 1:
+        return False
     worker = getattr(getattr(model_executor, "driver_worker", None), "worker", None)
     model_runner = getattr(worker, "model_runner", None)
     model_state = getattr(model_runner, "model_state", None)
@@ -90,12 +93,23 @@ def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: An
     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
 
-    assert model_state is not None
-    model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler))
+    if decodes_audio and hasattr(model_state, "set_first_audio_sink"):
+        assert model_state is not None
+        model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler))
+        return True
+    data_plane = getattr(model_runner, "_omni_data_plane", None)
+    get_model = getattr(model_runner, "get_model", None)
+    if data_plane is None or not callable(get_model):
+        return False
+    bind = getattr(get_model(), "bind_first_chunk_fast_path", None)
+    if not callable(bind):
+        return False
+    hook = bind(engine_output_queue_sink(output_queue, scheduler, upstream_first_audio=False))
+    if hook is None:
+        return False
+    data_plane.set_first_chunk_hook(hook)
     return True
 
 

vllm_omni/entrypoints/omni_base.py

@@ -441,7 +441,7 @@ class OmniBase(PDDisaggregationMixin):
             if self.log_stats:
                 # Emit per-request orchestrator timing (including e2e_total_ms)
                 # before dropping request state.
-                req_state.metrics.build_and_log_summary()
+                req_state.metrics.log_timing_summary()
         except Exception:
             logger.exception(
                 "[%s] Failed to build/log summary for req=%s",

vllm_omni/entrypoints/openai/tts_adapters/base.py

@@ -81,13 +81,24 @@ def conditioning_cache_salt(
     This replaces the uploaded data URI in the hash without changing the request.
     """
     h = hashlib.sha256()
+    ref_audio: str | list[str] | tuple[str, int] | None = request.ref_audio
+    if (
+        isinstance(ref_audio, str)
+        and len(ref_audio) > 256
+        and tts_params is not None
+        and tts_params.get("ref_audio_cache_key") is not None
+    ):
+        # The resolve key below is already a content-aware digest of this
+        # locator. Hashing a repr of a large inline data URI again dominates
+        # the per-request serving cost without adding information.
+        ref_audio = ("ref_audio_cache_key", len(ref_audio))
     for part in (
         request.input,
         request.task_type,
         request.language,
         request.voice,
         request.ref_text,
-        registered_voice if registered_voice is not None else request.ref_audio,
+        registered_voice if registered_voice is not None else ref_audio,
         request.instructions,
         request.x_vector_only_mode,
         request.speaker_embedding,

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -32,6 +32,59 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _local_continuation_prompt(proc: Any, user_kwargs: dict[str, Any], reference: Any) -> tuple[list[int], Any] | None:
+    """Build the Local-v1.5 continuation prompt without the generic processor call.
+
+    Mirrors ``_build_continuation_codes`` of the model's processor: every
+    template segment is encoded separately there, so the constant segments
+    (keyed by ``language``) are cached and only the request text is
+    tokenized. Returns ``(text_ids, audio_codes)`` equal to
+    ``unified[:, 0].tolist()`` / ``unified[:, 1:]``, or ``None`` when the
+    request or processor does not match that exact layout.
+    """
+    import torch
+
+    if set(user_kwargs) - {"text", "language"}:
+        return None
+    if not isinstance(reference, list) or len(reference) != 1 or not isinstance(reference[0], torch.Tensor):
+        return None
+    cfg = getattr(proc, "model_config", None)
+    n_vq = getattr(cfg, "n_vq", None)
+    if cfg is None or not isinstance(n_vq, int) or n_vq <= 0:
+        return None
+    codes = reference[0]
+    if codes.ndim != 2 or int(codes.shape[1]) != n_vq or codes.shape[0] == 0:
+        return None
+    required = (
+        "_user_prompt_prefix_ids",
+        "_encode_text",
+        "_user_prompt_after_reference_ids",
+        "_assistant_prompt_prefix_ids",
+    )
+    if not all(callable(getattr(proc, name, None)) for name in required):
+        return None
+    language = user_kwargs.get("language")
+    cache = proc.__dict__.setdefault("_vllm_omni_continuation_segments", {})
+    segments = cache.get(language)
+    if segments is None:
+        fields = {} if language is None else {"language": language}
+        head = (
+            proc._user_prompt_prefix_ids()
+            + proc._encode_text("None")
+            + proc._user_prompt_after_reference_ids(language, fields)
+        )
+        tail = proc._assistant_prompt_prefix_ids() + [int(cfg.audio_start_token_id)]
+        segments = cache[language] = ([int(t) for t in head], [int(t) for t in tail])
+    head, tail = segments
+    text = user_kwargs.get("text")
+    prompt_ids = head + [int(t) for t in proc._encode_text("" if text is None else str(text))] + tail
+    num_audio = int(codes.shape[0])
+    text_ids = prompt_ids + [int(cfg.audio_assistant_slot_token_id)] * num_audio
+    audio_codes = torch.full((len(text_ids), n_vq), int(cfg.audio_pad_token_id), dtype=torch.int64)
+    audio_codes[len(prompt_ids) :] = codes.to(dtype=torch.int64, device="cpu")
+    return text_ids, audio_codes
+
+
 class _MossTTSAdapterBase(ARTTSAdapter):
     accumulate_nonstreaming: bool = False
 
@@ -82,6 +135,28 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         self._moss_ref_encoder = encoder
         return encoder
 
+    async def warmup(self) -> None:
+        """Set up the reference encoder at startup, not on the first cloning request.
+
+        Its CUDA graph capture takes seconds per API process (about a minute for
+        a shared encoder's workers).
+        """
+        if self._moss_variant not in ("tts", "local", "ttsd"):
+            return
+        from vllm_omni.model_executor.models.moss_tts.reference_encoder import reference_graphs_enabled
+        from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError
+
+        if not reference_graphs_enabled():
+            return
+        try:
+            await asyncio.to_thread(self._get_moss_ref_encoder().prepare)
+        except SharedReferenceEncoderStartupError:
+            # Do not advertise readiness and trigger concurrent local captures
+            # while the shared host is still compiling or has failed startup.
+            raise
+        except Exception:  # noqa: BLE001 - requests still encode (and capture) on demand
+            logger.warning("MOSS reference encoder warmup failed", exc_info=True)
+
     async def _encode_moss_references(
         self,
         request: "OpenAICreateSpeechRequest",
@@ -269,8 +344,16 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
         if hasattr(proc, "audio_tokenizer"):
             device = self._resolve_ref_encoder_device()
-            proc.audio_tokenizer = proc.audio_tokenizer.to(device).eval()
-            logger.info("MOSS reference-audio encoder (audio_tokenizer) placed on %s", device)
+            proc._vllm_omni_ref_encoder_device = device
+            from vllm_omni.model_executor.models.moss_tts.reference_encoder import shared_encoder_role
+
+            if shared_encoder_role() == "client":
+                # Another API process hosts the shared encoder on the GPU.
+                proc.audio_tokenizer = proc.audio_tokenizer.eval()
+                logger.info("MOSS reference-audio encoder: using the shared encoder of another API process")
+            else:
+                proc.audio_tokenizer = proc.audio_tokenizer.to(device).eval()
+                logger.info("MOSS reference-audio encoder (audio_tokenizer) placed on %s", device)
         self._moss_processor_cache = proc
         return proc
 
@@ -405,18 +488,24 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special
         # tokens and cols 1..n_vq are the delay-pattern audio code grid (mostly
         # audio_pad_code outside the reference block).
+        fast = None
         if v == "local" and request.ref_text and request.ref_text.strip():
             reference = user_kwargs.pop("reference")
             user_kwargs["text"] = request.ref_text.strip() + " " + (request.input or "")
-            user_msg = proc.build_user_message(**user_kwargs)
-            assistant_msg = proc.build_assistant_message(audio_codes_list=reference)
-            batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation")
+            fast = _local_continuation_prompt(proc, user_kwargs, reference)
+            if fast is None:
+                user_msg = proc.build_user_message(**user_kwargs)
+                assistant_msg = proc.build_assistant_message(audio_codes_list=reference)
+                batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation")
         else:
             user_msg = proc.build_user_message(**user_kwargs)
             batch = proc(conversations=[[user_msg]], mode="generation")
-        unified = batch["input_ids"][0]  # torch.LongTensor (L, 1+n_vq)
-        text_ids: list[int] = unified[:, 0].tolist()
-        audio_codes: torch.Tensor = unified[:, 1:].contiguous().to(torch.int64)
+        if fast is not None:
+            text_ids, audio_codes = fast
+        else:
+            unified = batch["input_ids"][0]  # torch.LongTensor (L, 1+n_vq)
+            text_ids = unified[:, 0].tolist()
+            audio_codes = unified[:, 1:].contiguous().to(torch.int64)
 
         params = {
             "prompt_token_ids": text_ids,

vllm_omni/metrics/stats.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import copy
+import logging
 import time
 from collections import defaultdict
 from collections.abc import Callable, Sequence
@@ -851,9 +852,62 @@ class OrchestratorAggregator:
             for cached_key in [k for k in self._replica_cache if k[1] == rid_key]:
                 self._replica_cache.pop(cached_key, None)
 
+    def log_timing_summary(self) -> None:
+        """Emit only the concise per-request ``[OmniTiming]`` lines.
+
+        The serving path uses this unless DEBUG tables are requested; it skips
+        the table rows that ``build_and_log_summary`` builds for its result.
+        """
+        if not self.log_stats:
+            return
+        if logger.isEnabledFor(logging.DEBUG):
+            self.build_and_log_summary()
+            return
+        for rid in sorted(set(self.stage_events.keys()) | {e.request_id for e in self.e2e_events}):
+            e2e_evt = next((e for e in self.e2e_events if e.request_id == rid), None)
+            stage_evts = sorted(
+                self.stage_events.get(rid, []),
+                key=lambda e: e.stage_id if e.stage_id is not None else -1,
+            )
+            self._log_omni_timing(rid, e2e_evt, stage_evts)
+
+    def _log_omni_timing(self, rid: str, e2e_evt: RequestE2EStats | None, stage_evts: list[StageRequestStats]) -> None:
+        pt: dict[str, float] = {}
+        if stage_evts:
+            pt = stage_evts[-1].pipeline_timings or {}
+        if pt or e2e_evt:
+            parts = [f"req={rid}"]
+            if e2e_evt:
+                parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s")
+            if "preprocess_ms" in pt:
+                parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s")
+            if e2e_evt:
+                engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0)
+                parts.append(f"engine={engine_ms / 1000.0:.2f}s")
+            stage_parts = []
+            for evt in stage_evts:
+                sid = evt.stage_id if evt.stage_id is not None else "?"
+                t = evt.stage_gen_time_ms / 1000.0
+                stage_parts.append(f"{sid}:{t:.2f}s")
+            if stage_parts:
+                parts.append(f"stages=[{','.join(stage_parts)}]")
+            transfer_parts = []
+            for te in self.transfer_events.values():
+                if te.request_id == rid:
+                    transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms")
+            if transfer_parts:
+                parts.append(f"transfers=[{','.join(transfer_parts)}]")
+            if "ar2diffusion_ms" in pt:
+                parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms")
+            logger.info("[OmniTiming] %s", " ".join(parts))
+
     def build_and_log_summary(self) -> dict[str, Any]:
         if not self.log_stats:
             return {}
+        # Per-request tables are diagnostic. Formatting them on the serving
+        # event loop costs more than the request's own bookkeeping, so they
+        # are emitted at DEBUG; the concise [OmniTiming] line stays at INFO.
+        log_tables = logger.isEnabledFor(logging.DEBUG)
         wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
         e2e_avg_req = (wall_time_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
         e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / wall_time_ms) if wall_time_ms > 0 else 0.0
@@ -895,8 +949,8 @@ class OrchestratorAggregator:
             v = overall_summary.get(k, None)
             if v not in (0, 0.0, 0.000, None, ""):
                 overall_fields.append(k)
-        if overall_fields:
-            logger.info(
+        if log_tables and overall_fields:
+            logger.debug(
                 "\n%s",
                 _format_table("Overall Summary", overall_summary, overall_fields),
             )
@@ -921,8 +975,8 @@ class OrchestratorAggregator:
                         nonzero_e2e_fields.add(k)
                 value_fields_e2e = sorted(nonzero_e2e_fields)
 
-                if value_fields_e2e:
-                    logger.info(
+                if log_tables and value_fields_e2e:
+                    logger.debug(
                         "\n%s",
                         _format_table(
                             f"RequestE2EStats [request_id={rid}]",
@@ -936,34 +990,7 @@ class OrchestratorAggregator:
                 self.stage_events.get(rid, []),
                 key=lambda e: e.stage_id if e.stage_id is not None else -1,
             )
-            pt = {}
-            if stage_evts:
-                pt = stage_evts[-1].pipeline_timings or {}
-            if pt or e2e_evt:
-                parts = [f"req={rid}"]
-                if e2e_evt:
-                    parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s")
-                if "preprocess_ms" in pt:
-                    parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s")
-                if e2e_evt:
-                    engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0)
-                    parts.append(f"engine={engine_ms / 1000.0:.2f}s")
-                stage_parts = []
-                for evt in stage_evts:
-                    sid = evt.stage_id if evt.stage_id is not None else "?"
-                    t = evt.stage_gen_time_ms / 1000.0
-                    stage_parts.append(f"{sid}:{t:.2f}s")
-                if stage_parts:
-                    parts.append(f"stages=[{','.join(stage_parts)}]")
-                transfer_parts = []
-                for te in self.transfer_events.values():
-                    if te.request_id == rid:
-                        transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms")
-                if transfer_parts:
-                    parts.append(f"transfers=[{','.join(transfer_parts)}]")
-                if "ar2diffusion_ms" in pt:
-                    parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms")
-                logger.info("[OmniTiming] %s", " ".join(parts))
+            self._log_omni_timing(rid, e2e_evt, stage_evts)
 
             # === Stage table (columns = stage_id) ===
             # if any stage has diffusion_metrics, remove postprocess_time_ms field
@@ -987,7 +1014,7 @@ class OrchestratorAggregator:
 
             result_stage_table.append({"request_id": rid, "stages": stage_rows})
 
-            if stage_rows:
+            if log_tables and stage_rows:
                 # filter out all-zero fields for logging
                 all_value_fields = set()
                 for row in stage_rows:
@@ -1006,7 +1033,7 @@ class OrchestratorAggregator:
                         value_fields_list.append(field)
 
                 if value_fields_list:
-                    logger.info(
+                    logger.debug(
                         "\n%s",
                         _format_table(
                             f"StageRequestStats [request_id={rid}]",
@@ -1027,7 +1054,7 @@ class OrchestratorAggregator:
             ]
             result_trans_table.append({"request_id": rid, "transfers": transfer_rows})
 
-            if transfer_rows:
+            if log_tables and transfer_rows:
                 # filter out all-zero fields for logging
                 all_value_fields = set()
                 for row in transfer_rows:
@@ -1046,7 +1073,7 @@ class OrchestratorAggregator:
                         value_fields_list.append(field)
 
                 if value_fields_list:
-                    logger.info(
+                    logger.debug(
                         "\n%s",
                         _format_table(
                             f"TransferEdgeStats [request_id={rid}]",

vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py

@@ -539,11 +539,15 @@ class RingKVCache:
             end_offset = self.end_offset.index_select(0, slots)
             row_cache = self.cache.index_select(1, slots)
 
-            indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype)
+            # Upsampling can make a chunk longer than the ring. Only the last
+            # capacity tokens survive; writing all T tokens gives scatter_
+            # duplicate destinations and nondeterministic cache contents.
+            write_length = min(T, self.capacity)
+            indexes = torch.arange(T - write_length, T, device=end_offset.device, dtype=end_offset.dtype)
             indexes = (indexes + end_offset.view(-1, 1)) % self.capacity
-            scatter_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D)
-            row_cache[0].scatter_(2, scatter_indexes, k)
-            row_cache[1].scatter_(2, scatter_indexes, v)
+            scatter_indexes = indexes.view(B, 1, write_length, 1).expand(-1, H, write_length, D)
+            row_cache[0].scatter_(2, scatter_indexes, k[:, :, -write_length:])
+            row_cache[1].scatter_(2, scatter_indexes, v[:, :, -write_length:])
             # Live and graph-padding rows always map to distinct slots. The
             # latter map only to scratch state, so this write cannot corrupt a
             # request even though dense graph operators still execute it.
@@ -729,7 +733,13 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
         if self.context is None:
             capacity = self.weights_per_step if self.weights_per_step else 1024
         else:
-            capacity = self.context
+            # RingKVCache.complete writes a whole chunk before attending, so a
+            # ring of exactly ``context`` entries lets a T-token chunk evict up
+            # to T tokens of history that the chunk's own queries should still
+            # see (and, for T > context, evict part of the chunk itself).
+            # ``_ring_headroom`` (max chunk tokens at this layer) keeps every
+            # in-window key resident; 0 preserves the legacy truncation.
+            capacity = self.context + int(getattr(self, "_ring_headroom", 0))
 
         kv_cache = RingKVCache(
             batch_size,
@@ -791,29 +801,63 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
         if self.rope:
             q, k = self.rope(q, k, offset, time_before_heads=False)
 
-        k, v, pos_k = self._complete_kv(k, v, execution_context)
-        pos_k = pos_k[:, None]
-
-        if self.causal:
-            pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1)
-            delta = pos_q - pos_k
-            attn_bias = (pos_k >= 0) & (delta >= 0)
-            if self.context is not None:
-                attn_bias = attn_bias & (delta < self.context)
-            attn_bias = attn_bias[:, None]
+        slot_attention = getattr(self, "_slot_attention", None)
+        slot_attention_rows = getattr(self, "_slot_attention_rows", None)
+        advance_kv_offset = False
+        slot_path = state is not None and execution_context is not None and self.causal
+        if slot_path and slot_attention_rows is not None:
+            # One launch: the gathered MHA offset equals the ring's end_offset
+            # for every slot, and both advance together below.
+            valid_lengths = execution_context.valid_rows.to(dtype=torch.int32) * T
+            x = slot_attention_rows(
+                q,
+                k,
+                v,
+                state.kv_cache.cache,
+                offset,
+                execution_context.state_slot_ids,
+                valid_lengths,
+                self.context if self.context is not None else -1,
+            )
+            advance_kv_offset = True
+        elif slot_attention is not None and slot_path:
+            # Current streaming batches use one exact T. Padding rows own
+            # scratch slots and must not advance persistent attention state.
+            valid_lengths = execution_context.valid_rows.to(dtype=torch.int32) * T
+            x = slot_attention(
+                q,
+                k,
+                v,
+                state.kv_cache.cache,
+                state.kv_cache.end_offset,
+                execution_context.state_slot_ids,
+                valid_lengths,
+                self.context if self.context is not None else -1,
+            )
         else:
-            attn_bias = None
+            k, v, pos_k = self._complete_kv(k, v, execution_context)
+            pos_k = pos_k[:, None]
 
-        streaming_attention = getattr(self, "_streaming_attention", None)
-        if (
-            streaming_attention is not None
-            and attn_bias is not None
-            and q.dtype == torch.bfloat16
-            and q.shape[-1] == 64
-        ):
-            x = streaming_attention(q, k, v, attn_bias)
-        else:
-            x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+            if self.causal:
+                pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1)
+                delta = pos_q - pos_k
+                attn_bias = (pos_k >= 0) & (delta >= 0)
+                if self.context is not None:
+                    attn_bias = attn_bias & (delta < self.context)
+                attn_bias = attn_bias[:, None]
+            else:
+                attn_bias = None
+
+            streaming_attention = getattr(self, "_streaming_attention", None)
+            if (
+                streaming_attention is not None
+                and attn_bias is not None
+                and q.dtype == torch.bfloat16
+                and q.shape[-1] == 64
+            ):
+                x = streaming_attention(q, k, v, attn_bias)
+            else:
+                x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
         x = x.transpose(1, 2).reshape(B, T, self.embed_dim)
         x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu)
 
@@ -824,6 +868,8 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
             else:
                 next_offset = torch.where(execution_context.valid_rows, offset + T, offset)
                 state.offset.index_copy_(0, execution_context.state_slot_ids, next_offset)
+                if advance_kv_offset:
+                    state.kv_cache.end_offset.index_copy_(0, execution_context.state_slot_ids, next_offset)
         return x
 
 
@@ -1624,6 +1670,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         decoder_kwargs_list = copy.deepcopy(config.decoder_kwargs)
         self.decoder = nn.ModuleList()
 
+        decoder_input_frame_rate = current_frame_rate
         for decoder_kwargs_i in decoder_kwargs_list:
             decoder_kwargs_i = dict(decoder_kwargs_i)
             if decoder_kwargs_i["module_type"] == "PatchedPretransform":
@@ -1638,6 +1685,11 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
                         context=int(round(current_frame_rate * context_duration)),
                     )
                 )
+                # Tokens this transformer sees per input code frame. Streaming
+                # ring headroom is sized from it (see initialize_decoder_state_pool).
+                self.decoder[-1].tokens_per_input_frame = max(
+                    1, int(round(current_frame_rate / decoder_input_frame_rate))
+                )
             current_frame_rate *= self.decoder[-1].downsample_ratio
 
         expected_output_frame_rate = float(self.sampling_rate * channel_interleave_factor)
@@ -1731,10 +1783,29 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         self._decoder_state_capacity = 0
         self._decoder_slot_offsets = None
 
-    def initialize_decoder_state_pool(self, state_capacity: int, scratch_capacity: int = 0) -> None:
-        """Allocate persistent decoder state independently of execution B."""
-        if state_capacity <= 0 or scratch_capacity < 0:
-            raise ValueError(f"Invalid decoder state capacities: state={state_capacity}, scratch={scratch_capacity}.")
+    def initialize_decoder_state_pool(
+        self, state_capacity: int, scratch_capacity: int = 0, chunk_frames: int = 0
+    ) -> None:
+        """Allocate persistent decoder state independently of execution B.
+
+        ``chunk_frames`` > 0 sizes every decoder attention ring as
+        ``context + chunk_frames * tokens_per_input_frame`` so that a chunk of
+        up to ``chunk_frames`` code frames never evicts keys inside its own
+        causal window; chunked streaming then matches whole-sequence decoding.
+        ``chunk_frames`` = 0 keeps the legacy ``context``-sized rings.
+        """
+        if state_capacity <= 0 or scratch_capacity < 0 or chunk_frames < 0:
+            raise ValueError(
+                "Invalid decoder state capacities: "
+                f"state={state_capacity}, scratch={scratch_capacity}, chunk_frames={chunk_frames}."
+            )
+        for module in self.decoder:
+            if not isinstance(module, MossAudioTokenizerProjectedTransformer):
+                continue
+            tokens_per_frame = int(getattr(module, "tokens_per_input_frame", 1))
+            for attention in module.modules():
+                if isinstance(attention, MossAudioTokenizerMultiheadAttention):
+                    attention._ring_headroom = chunk_frames * tokens_per_frame
         self._start_streaming(state_capacity + scratch_capacity, decoder_only=True)
         self._decoder_state_capacity = state_capacity + scratch_capacity
 

vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py

@@ -81,15 +81,31 @@ class CUDAGraphStreamingDecoderWrapper:
         frame_sizes: list[int],
         num_quantizers: int,
         vllm_config: VllmConfig,
+        scratch_base: int | None = None,
+        private_pool: bool = False,
+        compiled_decode: nn.Module | None = None,
     ) -> None:
         self.codec = codec
         self.state_capacity = int(state_capacity)
+        # Graph padding rows use scratch slots [scratch_base, scratch_base+B).
+        # A wrapper replaying concurrently with another needs its own scratch
+        # range and graph memory pool.
+        self.scratch_base = self.state_capacity if scratch_base is None else int(scratch_base)
+        self._private_pool = bool(private_pool)
+        # Graphs bake in the cuBLAS workspace of the stream they were captured
+        # on, and torch.cuda.graph shares one default capture stream. Split-K
+        # GEMMs synchronize through that workspace, so two graphs replaying
+        # concurrently from one capture stream can spin on each other forever.
+        self._capture_stream: torch.cuda.Stream | None = None
         self.batch_sizes = sorted({int(size) for size in batch_sizes if 0 < int(size) <= state_capacity})
         self.frame_sizes = sorted({int(size) for size in frame_sizes if int(size) > 0})
         self.num_quantizers = int(num_quantizers)
         self.graphs: dict[tuple[int, int], _CapturedStreamingDecodeGraph] = {}
         self._pool = None
         self._warmed_up = False
+        if compiled_decode is not None:
+            self._compiled_decode: nn.Module | None = compiled_decode
+            return
         # vLLM owns Inductor compilation; this wrapper remains the sole owner
         # of CUDA Graph capture/replay because it understands persistent codec
         # state slots. Disable vLLM's CUDA Graph layer to avoid nested graphs.
@@ -98,7 +114,7 @@ class CUDAGraphStreamingDecoderWrapper:
         compile_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
         compile_config.compilation_config.static_forward_context = {}
         with set_current_vllm_config(compile_config):
-            self._compiled_decode: nn.Module | None = _MossStreamingDecodeCompileAdapter(
+            self._compiled_decode = _MossStreamingDecodeCompileAdapter(
                 codec,
                 vllm_config=compile_config,
             )
@@ -175,7 +191,7 @@ class CUDAGraphStreamingDecoderWrapper:
                     frame_size,
                     exc_info=True,
                 )
-                scratch_slots = self.state_capacity + torch.arange(
+                scratch_slots = self.scratch_base + torch.arange(
                     batch_size,
                     dtype=torch.long,
                     device=device,
@@ -206,10 +222,13 @@ class CUDAGraphStreamingDecoderWrapper:
             device=device,
         )
         lengths = torch.zeros(batch_size, dtype=torch.long, device=device)
-        scratch_slots = self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=device)
+        scratch_slots = self.scratch_base + torch.arange(batch_size, dtype=torch.long, device=device)
         valid_rows = torch.zeros(batch_size, dtype=torch.bool, device=device)
 
-        stream = torch.cuda.Stream()
+        if self._private_pool and self._capture_stream is None:
+            self._capture_stream = torch.cuda.Stream(device=device)
+        # Warm up on the capture stream so its workspace exists before capture.
+        stream = self._capture_stream or torch.cuda.Stream()
         stream.wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(stream):
             for _ in range(3):
@@ -219,9 +238,11 @@ class CUDAGraphStreamingDecoderWrapper:
         self.codec.reset_decoder_state_slots(scratch_slots)
 
         if self._pool is None:
-            self._pool = current_platform.get_global_graph_pool()
+            self._pool = (
+                torch.cuda.graph_pool_handle() if self._private_pool else current_platform.get_global_graph_pool()
+            )
         graph = CUDAGraph()
-        with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"):
+        with torch.cuda.graph(graph, pool=self._pool, stream=self._capture_stream, capture_error_mode="thread_local"):
             audio, audio_lengths = decode(codes, lengths, scratch_slots, valid_rows)
 
         self.graphs[(batch_size, frame_size)] = _CapturedStreamingDecodeGraph(
@@ -272,7 +293,7 @@ class CUDAGraphStreamingDecoderWrapper:
         entry.static_lengths.zero_()
         entry.static_lengths[:actual_batch_size].fill_(int(frame_size))
         entry.static_state_slot_ids.copy_(
-            self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device)
+            self.scratch_base + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device)
         )
         entry.static_state_slot_ids[:actual_batch_size].copy_(state_slot_ids, non_blocking=True)
         entry.static_valid_rows.zero_()

vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py

@@ -0,0 +1,333 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Decode a stream's first codec chunk as soon as Stage 1 receives it.
+
+The Stage-1 engine loop only sees a new chunk between steps, and a step waits
+for its codec graph to finish. A first chunk that arrives mid-step therefore
+waits for the rest of that step, then for its own replay. This path takes the
+first chunk from the receive thread, replays a small dedicated graph on its
+own stream while the main step runs, and emits the audio straight to the
+engine output queue. The main path later receives the same chunk, finds it
+already decoded, and continues the stream from the state written here.
+"""
+
+from __future__ import annotations
+
+import math
+import queue
+import threading
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker_v2.first_audio_sender import FirstAudioSink, _PreparedDelivery
+
+logger = init_logger(__name__)
+
+
+class _SlotHandoff:
+    """Orders the main path's first use of a slot after the fast decode."""
+
+    __slots__ = ("done", "event", "error")
+
+    def __init__(self) -> None:
+        self.done = threading.Event()
+        self.event: torch.cuda.Event | None = None
+        self.error: Exception | None = None
+
+
+@dataclass
+class _Job:
+    request_id: str
+    request_key: str
+    delivery: _PreparedDelivery
+    slot: int
+    codes: torch.Tensor  # (n_vq, T) int64 on host
+    handoff: _SlotHandoff
+
+
+class MossFirstChunkFastPath:
+    def __init__(
+        self,
+        session: Any,
+        wrapper: Any,
+        *,
+        n_vq: int,
+        frames: int,
+        codebook_size: int,
+        samples_per_frame: int,
+        n_channels: int,
+        sample_rate: torch.Tensor,
+        device: torch.device,
+        gate_main: bool = False,
+        handoff_timeout_s: float = 30.0,
+    ) -> None:
+        if not math.isfinite(handoff_timeout_s) or handoff_timeout_s <= 0:
+            raise ValueError("handoff_timeout_s must be positive and finite")
+        self._handoff_timeout_s = float(handoff_timeout_s)
+        self._session = session
+        # When set, the main codec stream waits for an in-flight fast decode
+        # before its next replay, so the small first-chunk graph runs mostly
+        # uncontended instead of interleaving with a large regular replay.
+        self._gate_main = bool(gate_main)
+        self._inflight: torch.cuda.Event | None = None
+        self._wrapper = wrapper
+        self._n_vq = int(n_vq)
+        self.frames = int(frames)
+        self._codebook_size = int(codebook_size)
+        self._samples = int(frames) * int(samples_per_frame)
+        self._n_channels = int(n_channels)
+        self._sr = sample_rate
+        if device.index is None:
+            device = torch.device(device.type, torch.accelerator.current_device_index())
+        self._device = device
+        self._max_batch = max(wrapper.batch_sizes)
+        self._stream = torch.cuda.Stream(device=device, priority=-1)
+        self._host_codes = torch.zeros((self._n_vq, self._max_batch, self.frames), dtype=torch.long).pin_memory()
+        self._host_slots = torch.zeros(self._max_batch, dtype=torch.long).pin_memory()
+        audio_shape = (self._max_batch, self._n_channels, self._samples)
+        self._host_audio = torch.empty(audio_shape, dtype=torch.float32).pin_memory()
+        self._copied = torch.cuda.Event()
+        self._sink: FirstAudioSink | None = None
+        # Request key -> slot for chunks decoded here and not yet seen by the main path.
+        self._decoded: dict[str, int] = {}
+        # A failed output-queue handoff must be replayed from the already
+        # decoded samples. Re-decoding would advance the streaming state twice.
+        self._unsent_audio: dict[str, torch.Tensor] = {}
+        self._handoffs: dict[int, _SlotHandoff] = {}
+        self._lock = threading.Lock()
+        self._jobs: queue.SimpleQueue[_Job | None] = queue.SimpleQueue()
+        self._thread: threading.Thread | None = None
+        self._closed = False
+        self._worker_error: Exception | None = None
+        self._emitted = 0
+        self._batches = 0
+        self._decode_s = 0.0
+
+    def bind(self, sink: FirstAudioSink) -> None:
+        if not callable(getattr(sink, "prepare", None)):
+            raise TypeError("The codec first-chunk sink requires prepare(request_ids)")
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("The codec first-chunk decoder is closed")
+            self._sink = sink
+            if self._thread is None:
+                self._thread = threading.Thread(target=self._run, name="moss-codec-first-chunk", daemon=True)
+                self._thread.start()
+
+    def close(self) -> None:
+        with self._lock:
+            if not self._closed:
+                self._closed = True
+                self._jobs.put(None)
+            thread = self._thread
+        if thread is not None:
+            thread.join(timeout=5.0)
+            # A stalled decode may still own GPU state: keep its thread and
+            # handoff, rather than permitting a restart or slot reuse.
+            if thread.is_alive():
+                logger.warning("Codec first-chunk decoder is still stopping; slot waits remain bounded")
+
+    # ------------------------------------------------------------------
+    # Receive thread
+    # ------------------------------------------------------------------
+    def submit(
+        self,
+        request_id: str,
+        request_key: str,
+        codes_flat: Any,
+        req_slots: dict[str, int],
+    ) -> bool:
+        """Claim a first chunk for fast decode; False leaves it to the main path."""
+        # Own the codes: the job outlives this call and the payload it came from.
+        codes = torch.as_tensor(codes_flat, dtype=torch.long).reshape(-1).clone()
+        if codes.numel() != self._n_vq * self.frames:
+            return False
+        with self._lock:
+            if (
+                self._closed
+                or self._sink is None
+                or self._thread is None
+                or not self._thread.is_alive()
+                or request_key in req_slots
+            ):
+                return False
+            delivery = self._sink.prepare([request_id])
+            if request_id not in delivery.routes:
+                return False
+            slot = self._session.acquire(oldest=True)
+            if slot is None:
+                return False
+            handoff = _SlotHandoff()
+            self._handoffs[slot] = handoff
+            self._decoded[request_key] = slot
+            req_slots[request_key] = slot
+            codes = codes.reshape(self._n_vq, self.frames)
+            # Enqueue under the admission lock: close's sentinel cannot
+            # overtake an accepted job.
+            self._jobs.put(_Job(request_id, request_key, delivery, slot, codes, handoff))
+        return True
+
+    # ------------------------------------------------------------------
+    # Main (engine) thread
+    # ------------------------------------------------------------------
+    def get_request_slot(self, request_key: str, req_slots: dict[str, int]) -> int | None:
+        """Observe a completed admission before deciding whether cleanup is needed.
+
+        The scheduler removes cancelled routes before the runner's finish hook.
+        A submit that already froze its route must finish leasing its slot
+        before that hook checks it. Waiting for decode happens outside this
+        lock, so the worker can still complete the slot handoff.
+        """
+        with self._lock:
+            return req_slots.get(request_key)
+
+    def take_decoded(self, request_key: str) -> bool:
+        """True once per request whose first chunk this path already decoded."""
+        with self._lock:
+            return self._decoded.pop(request_key, None) is not None
+
+    def take_unsent_audio(self, request_key: str) -> torch.Tensor | None:
+        """Return first-chunk samples when the out-of-band sink rejected them."""
+        with self._lock:
+            return self._unsent_audio.pop(request_key, None)
+
+    def forget(self, request_key: str) -> None:
+        with self._lock:
+            self._decoded.pop(request_key, None)
+            self._unsent_audio.pop(request_key, None)
+
+    def gate(self) -> None:
+        """Queue the current stream behind the latest fast decode, if still running."""
+        event = self._inflight
+        if self._gate_main and event is not None and not event.query():
+            torch.cuda.current_stream(self._device).wait_event(event)
+
+    def order_after(self, slot: int) -> None:
+        """Make the current stream's next use of ``slot`` follow the fast decode."""
+        with self._lock:
+            handoff = self._handoffs.get(slot)
+        if handoff is None:
+            return
+        if not handoff.done.wait(timeout=self._handoff_timeout_s):
+            # Keep the handoff. The worker may still be writing the slot, so
+            # neither cancellation nor a retry may release/re-decode it.
+            raise TimeoutError(f"Codec first-chunk slot {slot} did not complete within {self._handoff_timeout_s:g}s")
+        if handoff.error is not None:
+            raise RuntimeError(f"Codec first-chunk decode failed for slot {slot}") from handoff.error
+        if handoff.event is not None:
+            torch.cuda.current_stream(self._device).wait_event(handoff.event)
+        with self._lock:
+            self._handoffs.pop(slot, None)
+
+    # ------------------------------------------------------------------
+    # Decode thread
+    # ------------------------------------------------------------------
+    def _run(self) -> None:
+        try:
+            current_omni_platform.set_device(self._device)
+            self._run_jobs()
+        except Exception as error:
+            logger.exception("MOSS codec first-chunk worker failed")
+            self._worker_error = error
+        finally:
+            with self._lock:
+                self._closed = True
+                for handoff in self._handoffs.values():
+                    if not handoff.done.is_set():
+                        handoff.error = self._worker_error or RuntimeError("Codec first-chunk worker stopped")
+                        handoff.done.set()
+            # Accepted, unprocessed routes must not be left waiting for PCM.
+            while True:
+                try:
+                    job = self._jobs.get_nowait()
+                except queue.Empty:
+                    break
+                if job is not None:
+                    self._fail_delivery(job)
+
+    @staticmethod
+    def _fail_delivery(job: _Job) -> None:
+        try:
+            job.delivery.fail([job.request_id])
+        except Exception:
+            logger.exception("Codec first-chunk error delivery failed for %s", job.request_id)
+
+    def _run_jobs(self) -> None:
+        while True:
+            job = self._jobs.get()
+            if job is None:
+                return
+            jobs = [job]
+            while len(jobs) < self._max_batch:
+                try:
+                    extra = self._jobs.get_nowait()
+                except queue.Empty:
+                    break
+                if extra is None:
+                    self._jobs.put(None)
+                    break
+                jobs.append(extra)
+            try:
+                start = time.perf_counter()
+                self._decode(jobs)
+                self._decode_s += time.perf_counter() - start
+                self._batches += 1
+                before, self._emitted = self._emitted, self._emitted + len(jobs)
+                if self._emitted // 1000 != before // 1000:
+                    logger.info(
+                        "MOSS codec first-chunk fast path: %d chunks in %d batches, %.2f ms/batch",
+                        self._emitted,
+                        self._batches,
+                        1e3 * self._decode_s / self._batches,
+                    )
+            except Exception as error:
+                logger.exception("MOSS codec first-chunk fast decode failed for %s", [j.request_id for j in jobs])
+                # A partial graph replay may have advanced the slot. Surface
+                # the failure instead of re-decoding and advancing it twice.
+                for item in jobs:
+                    item.handoff.error = error
+                    item.handoff.done.set()
+                    self._fail_delivery(item)
+                raise
+
+    @torch.no_grad()
+    def _decode(self, jobs: list[_Job]) -> None:
+        n = len(jobs)
+        for row, item in enumerate(jobs):
+            self._host_codes[:, row].copy_(item.codes)
+            self._host_slots[row] = item.slot
+        stream = self._stream
+        with torch.cuda.stream(stream):
+            for item in jobs:
+                self._session.order_after_reset(item.slot, stream)
+            codes = self._host_codes[:, :n].to(self._device, non_blocking=True).clamp_(0, self._codebook_size - 1)
+            slots = self._host_slots[:n].to(self._device, non_blocking=True)
+            result = self._wrapper.decode(codes, slots)
+            if result is None:
+                raise RuntimeError(f"no first-chunk codec graph for B={n}, T={self.frames}")
+            audio = result[0][:n, ..., : self._samples].float()
+            state_written = torch.cuda.Event()
+            state_written.record(stream)
+            self._inflight = state_written
+            self._host_audio[:n].view_as(audio).copy_(audio, non_blocking=True)
+            self._copied.record(stream)
+        for item in jobs:
+            item.handoff.event = state_written
+        self._copied.synchronize()
+        for row, item in enumerate(jobs):
+            wav = self._host_audio[row].clone()
+            if self._n_channels == 1:
+                wav = wav.reshape(-1)
+            try:
+                item.delivery([item.request_id], [wav], self._sr)
+            except Exception:
+                logger.exception("MOSS codec first-chunk output failed for %s", item.request_id)
+                with self._lock:
+                    self._unsent_audio[item.request_key] = wav
+            finally:
+                item.handoff.done.set()

vllm_omni/model_executor/models/moss_tts/local_model_state.py

@@ -0,0 +1,447 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MOSS Local request-slot state, independent of each step's batch order.
+
+State writes, gathers and MTP replay run on the runner stream. Output gathers
+own their storage; the runner retains its normal asynchronous output snapshot.
+No request dictionary holds a view into a mutable slot or CUDA graph buffer.
+"""
+
+import numpy as np
+import torch
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import set_forward_context
+from vllm.logger import init_logger
+from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.v1.worker.gpu.input_batch import get_num_sampled_and_rejected
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+logger = init_logger(__name__)
+
+
+def _metadata_to_device(data: np.ndarray, device: torch.device) -> torch.Tensor:
+    # CPU state tests need no CUDA allocator. The serving CUDA path retains
+    # vLLM's asynchronous pinned upload and its ownership/lifetime handling.
+    if device.type == "cpu":
+        return torch.from_numpy(data)
+    return async_tensor_h2d(data, device=device)
+
+
+class _CodeRowsSnapshot(PackedOutputSnapshot):
+    """Owned per-request code rows backed by one gathered device tensor.
+
+    The generic snapshot flattens every per-request leaf, concatenates them
+    into a slab and slices each back, then slices each again after the D2H.
+    Here the gathered ``[decode_rows, n_vq]`` tensor is already owned, so the
+    copy is a single D2H and the per-request views come from one ``split``.
+    """
+
+    def __init__(self, codes: torch.Tensor, decode_rows: list[int], num_reqs: int) -> None:
+        dict.__init__(self, {"codes": {"audio": self._rows(codes, decode_rows, num_reqs)}})
+        self._codes = codes
+        self._decode_rows = decode_rows
+        self._num_reqs = num_reqs
+        self.producer_event = None
+
+    @staticmethod
+    def _rows(codes: torch.Tensor, decode_rows: list[int], num_reqs: int) -> list[torch.Tensor]:
+        outputs = [codes.new_empty((0, codes.shape[-1]))] * num_reqs
+        for batch_row, row in zip(decode_rows, codes.split(1), strict=True):
+            outputs[batch_row] = row
+        return outputs
+
+    def copy_to_cpu(self, copy_tensor):
+        cpu = copy_tensor(self._codes)
+        return {"codes": {"audio": self._rows(cpu, self._decode_rows, self._num_reqs)}}
+
+
+class MossLocalModelState(OmniModelState):
+    def __init__(self, vllm_config, model, encoder_cache, device):
+        super().__init__(vllm_config, model, encoder_cache, device)
+        self._init_slot_buffers(self.scheduler_config.max_num_seqs, model.hidden_size, device, self.dtype)
+        self._batch_prefill = bool(getattr(model.config, "mrv2_batch_prefill", False))
+        self._direct_tokens = bool(getattr(model.config, "mrv2_direct_tokens", False))
+        self._local_eager_mtp = bool(getattr(model.config, "mrv2_eager_mtp", False))
+        logger.info(
+            "MOSS Local MRV2 batch prefill=%s, determined tokens=%s, eager MTP=%s",
+            self._batch_prefill,
+            self._direct_tokens,
+            self._local_eager_mtp,
+        )
+        logger.info("MOSS Local MRV2 GPU slot state enabled: capacity=%d", self.scheduler_config.max_num_seqs)
+
+    def _init_slot_buffers(self, capacity, hidden_size, device, dtype):
+        self._hidden_pool = torch.zeros((capacity, hidden_size), dtype=dtype, device=device)
+        self._codes_pool = torch.full(
+            (capacity, self.model.n_vq), self.model.audio_pad_token_id, dtype=torch.long, device=device
+        )
+        # Execution eligibility mirrors audio_state.is_stopping. Binary MTP
+        # continuation is a per-step output, not a sticky request-state update.
+        self._active_pool = torch.zeros(capacity, dtype=torch.bool, device=device)
+        self._active_host = [False] * capacity
+        self._decode_rows = []
+        self._prefill_staging = []
+        self._batch_prefill = False
+        self._direct_tokens = False
+        # Eager MTP: a frame is sampled right after the forward that produced
+        # its hidden state, one step earlier than the canonical preprocess
+        # path. Its input-embedding contribution (audio embedding times emit
+        # mask) is added to the next step's text embedding; ``_keep_pool``
+        # chains emission so a stopped stream cannot emit again.
+        self._local_eager_mtp = False
+        self._eager_emb_pool = torch.zeros((capacity, hidden_size), dtype=dtype, device=device)
+        self._keep_pool = torch.zeros(capacity, dtype=torch.bool, device=device)
+        self._completing_rows = []
+        self._local_eager_rows = []
+        # A slot takes the eager path only after its final prefill chunk ran
+        # an eager MTP in this state; others keep the canonical path, so a
+        # reused slot never reads a previous request's frame state.
+        self._eager_host = [False] * capacity
+        self._mtp_dispatcher = None
+
+    def sample_determined_tokens(self, batch, sampler):
+        """Skip full-vocabulary sampling only for unconstrained Local tokens.
+
+        The audio/binary MTP sampler has already made the random decision. The
+        outer text distribution has exactly one finite logit. Keep the original
+        path for processors that can mask it or request distribution metadata.
+        """
+        if not self._direct_tokens or type(sampler) is not Sampler:
+            return None
+        rows = batch.idx_mapping_np
+        if (
+            sampler.compute_nans
+            or sampler.return_sampling_mask
+            or sampler.trace_replay_state is not None
+            or sampler.get_logprobs_dims(rows) is not None
+            or np.any(sampler.logit_bias_state.use_logit_bias[rows])
+            or np.any(sampler.penalties_state.use_penalty[rows])
+            or np.any(sampler.bad_words_state.num_bad_words.np[rows])
+            or (
+                sampler.thinking_budget_state.enabled
+                and np.any(sampler.thinking_budget_state.use_thinking_budget[rows])
+            )
+        ):
+            return None
+        keep = self.model._batch_should_continue
+        if keep is None or keep.numel() != batch.num_reqs:
+            return None
+        # Owned storage: subsequent forwards can replace the continue mask
+        # while these token IDs are being copied or delivered asynchronously.
+        tokens = torch.where(keep, self.model.audio_assistant_slot_token_id, self.model.im_end_token_id)
+        num_sampled, num_rejected = get_num_sampled_and_rejected(
+            batch.seq_lens.new_ones(batch.num_reqs),
+            batch.seq_lens,
+            batch.cu_num_logits,
+            batch.idx_mapping,
+            sampler.req_states.prefill_len.gpu,
+        )
+        return SamplerOutput(tokens.reshape(-1, 1), None, None, num_sampled, num_rejected)
+
+    def _apply_reference_batch(self, embeds, references):
+        """One owned pinned upload for reference codes and their token positions.
+
+        Host storage may be reused only after its H2D event completes. Device
+        storage is consumed and overwritten in runner-stream order. If all
+        cached slots are busy, allocate a transient snapshot instead of waiting.
+        """
+        if not references:
+            return
+        count = sum(codes.shape[0] for _, codes in references)
+        capacity = max(count, self._static_inputs_embeds.shape[0])
+        staging = None
+        for item in self._prefill_staging:
+            if item[0].shape[0] >= count and (item[2] is None or item[2].query()):
+                staging = item
+                break
+        if staging is None:
+            host = torch.empty((capacity, self.model.n_vq + 1), dtype=torch.long, pin_memory=embeds.is_cuda)
+            device = torch.empty_like(host, device=embeds.device) if embeds.is_cuda else host
+            staging = [host, device, torch.cuda.Event() if embeds.is_cuda else None]
+            if len(self._prefill_staging) < 3:
+                self._prefill_staging.append(staging)
+        host, device, event = staging
+        offset = 0
+        for start, codes in references:
+            end = offset + codes.shape[0]
+            host[offset:end, 0].numpy()[:] = np.arange(start, start + codes.shape[0])
+            host[offset:end, 1:].copy_(codes)
+            offset = end
+        if embeds.is_cuda:
+            device[:count].copy_(host[:count], non_blocking=True)
+            event.record(torch.cuda.current_stream(embeds.device))
+        packed = device[:count]
+        positions = packed[:, 0]
+        audio = self.model._audio_embed(packed[:, 1:])
+        embeds.index_copy_(0, positions, embeds.index_select(0, positions) + audio)
+
+    def add_request(self, req_index, new_req_data):
+        super().add_request(req_index, new_req_data)
+        self._hidden_pool[req_index].zero_()
+        self._codes_pool[req_index].fill_(self.model.audio_pad_token_id)
+        self._active_pool[req_index].fill_(True)
+        self._active_host[req_index] = True
+        self._eager_host[req_index] = False
+        self._keep_pool[req_index].fill_(False)
+        self._eager_emb_pool[req_index].zero_()
+
+    def remove_request(self, req_index_or_id):
+        # Freeing is CPU bookkeeping. Reinitialization happens on admission,
+        # in stream order, so pending output snapshots remain independent.
+        super().remove_request(req_index_or_id)
+
+    def _select_rows(self, tensor, rows):
+        if rows == list(range(len(rows))):
+            return tensor[: len(rows)]
+        indices = _metadata_to_device(np.asarray(rows, dtype=np.int64), tensor.device)
+        return tensor.index_select(0, indices)
+
+    def run_preprocess(self, input_batch, model_inputs, req_states=None, mtp_batch_descriptor_dispatcher=None):
+        input_ids = model_inputs.get("input_ids")
+        if input_ids is None:
+            input_ids = input_batch.input_ids
+        embeds = model_inputs.get("inputs_embeds")
+        if embeds is None:
+            embeds = self.model.embed_input_ids(input_ids[: input_batch.num_tokens])
+            model_inputs["inputs_embeds"] = embeds
+        elif self._static_inputs_embeds is not None and embeds.data_ptr() == self._static_inputs_embeds.data_ptr():
+            embeds[: input_batch.num_tokens].copy_(self.model.embed_input_ids(input_ids[: input_batch.num_tokens]))
+
+        decode_rows = []
+        completing_rows = []
+        references = []
+        for row in range(input_batch.num_reqs):
+            slot = int(input_batch.idx_mapping_np[row])
+            buf = self.intermediate_buffer.buffers[slot]
+            if not buf or str(buf.get("req_id", "")).startswith("_warmup_"):
+                continue
+            start = int(input_batch.query_start_loc_np[row])
+            count = int(input_batch.num_scheduled_tokens[row])
+            prompt_len = (
+                None if req_states is None else self._get_req_state_value(getattr(req_states, "prompt_len", None), slot)
+            )
+            computed = (
+                None
+                if req_states is None
+                else self._get_req_state_value(getattr(req_states, "num_computed_tokens", None), slot)
+            )
+            if computed is None:
+                computed = self._get_input_batch_num_computed(input_batch, slot, row)
+            prefill = computed < prompt_len if computed is not None and prompt_len is not None else count > 1
+            if prefill and computed is not None and prompt_len is not None and computed + count >= prompt_len:
+                completing_rows.append(row)
+            if count == 1 and not prefill and isinstance(buf.get("audio_state"), dict):
+                active = not bool(buf["audio_state"].get("is_stopping"))
+                if active != self._active_host[slot]:
+                    self._active_pool[slot].fill_(active)
+                    self._active_host[slot] = active
+                decode_rows.append(row)
+                continue
+            if self._batch_prefill:
+                ref = (buf.get("codes", {}) or {}).get("ref")
+                ref_offset = int(buf.get("ref_offset", 0))
+                # Device-side references keep the canonical path; do not add
+                # an implicit D2H just to place them in a host staging buffer.
+                if not isinstance(ref, torch.Tensor) or ref.device.type == "cpu":
+                    if isinstance(ref, torch.Tensor) and ref.numel():
+                        if ref.dim() == 1 and ref.numel() % self.model.n_vq == 0:
+                            ref = ref.view(-1, self.model.n_vq)
+                        if ref.dim() == 2:
+                            chunk = ref[ref_offset : ref_offset + count]
+                            if chunk.numel() and chunk.shape[0] == count:
+                                references.append((start, chunk))
+                    # Text embeddings have already been computed for all input
+                    # tokens. Match the canonical prefill state transition.
+                    self.intermediate_buffer.update(
+                        slot,
+                        {
+                            "audio_state": {"is_stopping": False},
+                            "ref_offset": ref_offset + count,
+                        },
+                        self.model.gpu_resident_buffer_keys,
+                    )
+                    if not self._active_host[slot]:
+                        self._active_pool[slot].fill_(True)
+                        self._active_host[slot] = True
+                    continue
+            # Reference conditioning and chunked-prefill coordinates retain
+            # the model's canonical scalar implementation.
+            info = dict(buf)
+            info["_omni_is_prefill"] = prefill
+            if prompt_len is not None:
+                info["_omni_prompt_len"] = prompt_len
+            if computed is not None:
+                info["_omni_num_computed_tokens"] = computed
+            ids = input_ids[start : start + count]
+            previous = embeds[start : start + count]
+            new_ids, new_embeds, updates = self.model.preprocess(ids, previous, **info)
+            if "mtp_inputs" in updates:
+                raise RuntimeError("MOSS slot state requires an admitted prefill before decode")
+            if self._preprocess_result_needs_writeback(ids, new_ids):
+                ids.copy_(new_ids)
+            if self._preprocess_result_needs_writeback(previous, new_embeds):
+                previous.copy_(new_embeds)
+            self.intermediate_buffer.update(slot, updates, self.model.gpu_resident_buffer_keys)
+            active = not bool(updates.get("audio_state", {}).get("is_stopping", False))
+            if active != self._active_host[slot]:
+                self._active_pool[slot].fill_(active)
+                self._active_host[slot] = active
+
+        self._apply_reference_batch(embeds, references)
+        self._decode_rows = decode_rows
+        if self._local_eager_mtp:
+            self._completing_rows = completing_rows
+            self._mtp_dispatcher = mtp_batch_descriptor_dispatcher
+            eager_rows = [row for row in decode_rows if self._eager_host[int(input_batch.idx_mapping_np[row])]]
+            canonical_rows = [row for row in decode_rows if not self._eager_host[int(input_batch.idx_mapping_np[row])]]
+            self._local_eager_rows = eager_rows
+            if eager_rows:
+                # Their frame was sampled after the previous forward.
+                slots = self._select_rows(input_batch.idx_mapping, eager_rows)
+                offsets = self._mtp_offsets[: len(eager_rows)]
+                offsets.copy_(self._select_rows(input_batch.query_start_loc, eager_rows))
+                embeds.index_copy_(
+                    0, offsets, embeds.index_select(0, offsets) + self._eager_emb_pool.index_select(0, slots)
+                )
+            if canonical_rows:
+                self._run_slot_mtp(input_batch, input_ids, embeds, canonical_rows, mtp_batch_descriptor_dispatcher)
+        elif decode_rows:
+            self._run_slot_mtp(input_batch, input_ids, embeds, decode_rows, mtp_batch_descriptor_dispatcher)
+
+    def _run_slot_mtp(self, batch, input_ids, embeds, rows, dispatcher):
+        bsz = len(rows)
+        slots = self._select_rows(batch.idx_mapping, rows)
+        # Upstream query_start_loc is int32; index_copy_ requires int64.
+        offsets = self._mtp_offsets[:bsz]
+        offsets.copy_(self._select_rows(batch.query_start_loc, rows))
+        ids, emb = self._mtp_input_ids[:bsz], self._mtp_input_embeds[:bsz]
+        hidden, step = self._mtp_hidden[:bsz], self._mtp_text_step[:bsz]
+        ids.copy_(input_ids.index_select(0, offsets))  # MRV2 token inputs are int32.
+        torch.index_select(embeds, 0, offsets, out=emb)
+        torch.index_select(self._hidden_pool, 0, slots, out=hidden)
+        step[:, 0].copy_(self._active_pool.index_select(0, slots))
+
+        buffers = [self.intermediate_buffer.buffers[int(batch.idx_mapping_np[row])] for row in rows]
+        req_ids = [str(buf["req_id"]) for buf in buffers]
+        generators = [
+            self._get_mtp_generator(req_id, buf.get("sampling_params"), ids.device)
+            for req_id, buf in zip(req_ids, buffers, strict=True)
+        ]
+        can_graph = self._is_mtp_graph_runner() and not any(g is not None for g in generators)
+        desc = dispatcher(bsz) if can_graph and dispatcher is not None else None
+        mode = getattr(desc, "cg_mode", CUDAGraphMode.FULL) if can_graph else CUDAGraphMode.NONE
+        num_tokens = int(desc.num_tokens) if desc is not None else bsz
+        with set_forward_context(
+            None, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=mode, batch_descriptor=desc
+        ):
+            new_emb, codes = self._call_mtp_with_sampling(
+                self._mtp_input_ids[:num_tokens],
+                self._mtp_input_embeds[:num_tokens],
+                self._mtp_hidden[:num_tokens],
+                self._mtp_text_step[:num_tokens],
+                buffers=buffers,
+                req_ids=req_ids,
+                generators=generators,
+            )
+        embeds.index_copy_(0, offsets, new_emb[:bsz])
+        self._codes_pool.index_copy_(0, slots, codes[:bsz])
+
+    def _run_local_eager_mtp(self, batch, rows, completing):
+        """Sample the next frame from this forward's hidden state.
+
+        Same MTP inputs as the canonical next-step preprocess (hidden state
+        and active flag); the input embedding is zero, so the returned
+        embedding is exactly the audio contribution the next step adds to its
+        text embedding. A decode row emits only if its previous frame did.
+        """
+        bsz = len(rows)
+        slots = self._select_rows(batch.idx_mapping, rows)
+        ids, emb = self._mtp_input_ids[:bsz], self._mtp_input_embeds[:bsz]
+        hidden, step = self._mtp_hidden[:bsz], self._mtp_text_step[:bsz]
+        ids.zero_()  # the MTP does not read token ids
+        emb.zero_()
+        torch.index_select(self._hidden_pool, 0, slots, out=hidden)
+        gate = self._keep_pool.index_select(0, slots)
+        if completing:
+            first = _metadata_to_device(
+                np.asarray([row in completing for row in rows], dtype=np.bool_),
+                gate.device,
+            )
+            gate = gate | first
+        step[:, 0].copy_(self._active_pool.index_select(0, slots) & gate)
+
+        buffers = [self.intermediate_buffer.buffers[int(batch.idx_mapping_np[row])] for row in rows]
+        req_ids = [str(buf["req_id"]) for buf in buffers]
+        generators = [
+            self._get_mtp_generator(req_id, buf.get("sampling_params"), ids.device)
+            for req_id, buf in zip(req_ids, buffers, strict=True)
+        ]
+        dispatcher = self._mtp_dispatcher
+        can_graph = self._is_mtp_graph_runner() and not any(g is not None for g in generators)
+        desc = dispatcher(bsz) if can_graph and dispatcher is not None else None
+        mode = getattr(desc, "cg_mode", CUDAGraphMode.FULL) if can_graph else CUDAGraphMode.NONE
+        num_tokens = int(desc.num_tokens) if desc is not None else bsz
+        with set_forward_context(
+            None, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=mode, batch_descriptor=desc
+        ):
+            contribution, codes = self._call_mtp_with_sampling(
+                self._mtp_input_ids[:num_tokens],
+                self._mtp_input_embeds[:num_tokens],
+                self._mtp_hidden[:num_tokens],
+                self._mtp_text_step[:num_tokens],
+                buffers=buffers,
+                req_ids=req_ids,
+                generators=generators,
+            )
+        codes = codes[:bsz]
+        self._eager_emb_pool.index_copy_(0, slots, contribution[:bsz])
+        self._codes_pool.index_copy_(0, slots, codes)
+        self._keep_pool.index_copy_(0, slots, codes.ne(self.model.audio_pad_token_id).any(dim=-1))
+
+    def run_postprocess(self, hidden_states, input_batch):
+        # query_start_loc is already device-resident and authoritative for
+        # mixed/chunked-prefill batches. index_copy snapshots graph outputs.
+        last = input_batch.query_start_loc[1 : input_batch.num_reqs + 1] - 1
+        self._hidden_pool.index_copy_(
+            0, input_batch.idx_mapping[: input_batch.num_reqs], hidden_states.index_select(0, last)
+        )
+        if self._local_eager_mtp:
+            completing = self._completing_rows
+            rows = sorted(self._local_eager_rows + completing)
+            # Canonical decode rows emit their preprocess frame; eager rows
+            # and completed prefills emit the frame sampled below.
+            self._decode_rows = sorted(set(self._decode_rows) | set(completing))
+            if rows:
+                self._run_local_eager_mtp(input_batch, rows, set(completing))
+                for row in completing:
+                    self._eager_host[int(input_batch.idx_mapping_np[row])] = True
+
+    def postprocess_model_output(self, model_output, input_batch, req_states):
+        # Local backbone returns a tensor; keep the standard runner's output
+        # snapshot/materializer contract and omit hidden payload transport.
+        if not isinstance(model_output, torch.Tensor):
+            raise TypeError("MOSS Local slot state requires a tensor backbone output")
+        model = self.model
+        model._batch_state = None
+        model._batch_state_spans = None
+        model._batch_should_continue = torch.ones(input_batch.num_reqs, device=model_output.device, dtype=torch.bool)
+        if not self._decode_rows:
+            return model_output, {}
+        # This gather is an owned snapshot. Later slot writes cannot change
+        # data already handed to the runner, even before its D2H completes.
+        slots = self._select_rows(input_batch.idx_mapping, self._decode_rows)
+        codes = self._codes_pool.index_select(0, slots)
+        keep = codes.ne(model.audio_pad_token_id).any(dim=-1)
+        if self._decode_rows == list(range(input_batch.num_reqs)):
+            model._batch_should_continue = keep
+        else:
+            rows = _metadata_to_device(
+                np.asarray(self._decode_rows, dtype=np.int64),
+                codes.device,
+            )
+            model._batch_should_continue.index_copy_(0, rows, keep)
+        return model_output, _CodeRowsSnapshot(codes, list(self._decode_rows), input_batch.num_reqs)

vllm_omni/model_executor/models/moss_tts/local_short_attention.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Single-query attention for the frame-local MOSS depth transformer.
+
+The caller has already written this token's K/V and provides only its valid
+prefix. There is no persistent state here. BF16 probability rounding matches
+the FlashAttention P/V multiply, but reduction order is not bit-identical.
+"""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _short_attention(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    out_ptr,
+    qs0: tl.constexpr,
+    qs1: tl.constexpr,
+    ks0: tl.constexpr,
+    ks1: tl.constexpr,
+    ks2: tl.constexpr,
+    vs0: tl.constexpr,
+    vs1: tl.constexpr,
+    vs2: tl.constexpr,
+    heads: tl.constexpr,
+    dim: tl.constexpr,
+    tokens: tl.constexpr,
+    block_d: tl.constexpr,
+    block_t: tl.constexpr,
+):
+    row = tl.program_id(0)
+    batch, head = row // heads, row % heads
+    d, t = tl.arange(0, block_d), tl.arange(0, block_t)
+    query = tl.load(q_ptr + batch * qs0 + head * qs1 + d, d < dim, 0).to(tl.float32)
+    keys = tl.load(
+        k_ptr + batch * ks0 + head * ks1 + t[:, None] * ks2 + d[None, :],
+        (t[:, None] < tokens) & (d[None, :] < dim),
+        0,
+    ).to(tl.float32)
+    logits = tl.sum(keys * query[None, :], axis=1) * (dim**-0.5 * 1.4426950408889634)
+    logits = tl.where(t < tokens, logits, float("-inf"))
+    probs = tl.exp2(logits - tl.max(logits, axis=0))
+    denominator = tl.sum(probs, axis=0)
+    values = tl.load(
+        v_ptr + batch * vs0 + head * vs1 + t[:, None] * vs2 + d[None, :],
+        (t[:, None] < tokens) & (d[None, :] < dim),
+        0,
+    )
+    probs = probs.to(values.dtype).to(tl.float32)
+    output = tl.sum(values.to(tl.float32) * probs[:, None], axis=0) / denominator
+    tl.store(out_ptr + row * dim + d, output, d < dim)
+
+
+@torch.library.custom_op("vllm_omni::moss_local_short_attention", mutates_args=())
+def local_short_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    batch, heads, seq_len, dim = query.shape
+    if not (
+        query.is_cuda
+        and key.device == query.device
+        and value.device == query.device
+        and query.dtype == key.dtype == value.dtype == torch.bfloat16
+        and seq_len == 1
+        and 0 < dim <= 128
+        and key.shape == value.shape
+        and key.shape[:2] == query.shape[:2]
+        and key.shape[-1] == dim
+        and 0 < key.shape[2] <= 16
+        and query.stride(-1) == key.stride(-1) == value.stride(-1) == 1
+    ):
+        raise ValueError("MOSS short attention requires CUDA BF16 single-query inputs and 1..16 valid keys")
+    output = query.new_empty((batch, heads, 1, dim))
+    _short_attention[(batch * heads,)](
+        query,
+        key,
+        value,
+        output,
+        query.stride(0),
+        query.stride(1),
+        key.stride(0),
+        key.stride(1),
+        key.stride(2),
+        value.stride(0),
+        value.stride(1),
+        value.stride(2),
+        heads,
+        dim,
+        key.shape[2],
+        triton.next_power_of_2(dim),
+        triton.next_power_of_2(key.shape[2]),
+        num_warps=1,
+    )
+    return output
+
+
+@local_short_attention.register_fake
+def _fake_local_short_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    return query.new_empty(query.shape)

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -8,7 +8,10 @@
 
 from __future__ import annotations
 
+import threading
 from collections.abc import Iterable
+from contextlib import nullcontext
+from functools import partial
 from typing import Any
 
 import torch
@@ -32,11 +35,16 @@ from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer
 from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
     CUDAGraphStreamingDecoderWrapper,
 )
+from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import (
+    MossFirstChunkFastPath,
+)
 from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
     MossTTSCUDAGraphCodecWrapper,
 )
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot, pack_output_snapshot
 from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp
+from vllm_omni.worker_v2.first_audio_sender import FirstAudioSink
 
 logger = init_logger(__name__)
 
@@ -53,18 +61,46 @@ class _MossCodecStreamSession:
         vllm_config: VllmConfig,
         graph_batch_sizes: list[int] | None = None,
         graph_frame_sizes: list[int] | None = None,
+        gpu_output: bool = False,
+        chunk_frames: int = 0,
+        ring_headroom: bool = False,
+        fast_batch_sizes: list[int] | None = None,
+        fast_frames: int = 0,
     ) -> None:
         self._codec = codec
         self._state_capacity = int(state_capacity)
+        self._chunk_frames = int(chunk_frames)
+        self._ring_headroom = bool(ring_headroom)
         self._n_vq = int(n_vq)
         self._device = next(codec.parameters()).device
+        self._gpu_output = gpu_output and self._device.type == "cuda"
         self._free_stream_slots = list(reversed(range(self._state_capacity)))
         self._leased_slots: set[int] = set()
+        # The first-chunk fast path leases slots from its decode thread.
+        self._slot_lock = threading.Lock()
+        self._slot_reset_events: dict[int, torch.cuda.Event] = {}
         self._closed = False
+        self._metadata_index = 0
+        self._metadata_ring = []
+        if self._gpu_output:
+            # The metadata DMA is tiny, but a short ring forces the producer
+            # CPU thread to synchronize with the codec stream while it waits
+            # for a host slab to become reusable.  Keep enough entries for a
+            # scheduler burst so staging remains enqueue-only in practice.
+            for _ in range(64):
+                host = torch.empty(2 * self._state_capacity, dtype=torch.long, pin_memory=True)
+                gpu = torch.empty_like(host, device=self._device)
+                self._metadata_ring.append((host, host.numpy(), gpu, torch.cuda.Event(), False))
         self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
+        self._fast_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
         batch_sizes = sorted({int(size) for size in (graph_batch_sizes or []) if 0 < int(size) <= self._state_capacity})
         frame_sizes = sorted({int(size) for size in (graph_frame_sizes or []) if int(size) > 0})
         scratch_capacity = max(batch_sizes, default=0) if self._device.type in ("cuda", "npu") else 0
+        fast_batch_sizes = sorted({int(size) for size in (fast_batch_sizes or []) if 0 < int(size)})
+        if not (fast_frames > 0 and fast_batch_sizes and batch_sizes and frame_sizes and self._device.type == "cuda"):
+            fast_batch_sizes = []
+        fast_scratch_base = self._state_capacity + scratch_capacity
+        scratch_capacity += max(fast_batch_sizes, default=0)
         self._total_state_capacity = self._state_capacity + scratch_capacity
         self._state_slot_ids = torch.arange(
             self._total_state_capacity,
@@ -76,7 +112,14 @@ class _MossCodecStreamSession:
         if not callable(initialize_state_pool):
             raise RuntimeError("The streaming codec does not implement a decoder state pool.")
         with torch.no_grad():
-            initialize_state_pool(self._state_capacity, scratch_capacity)
+            if self._ring_headroom:
+                if self._chunk_frames <= 0:
+                    raise ValueError("codec_ring_headroom requires a positive codec chunk frame count.")
+                # Size every attention ring for context + one full chunk so a
+                # chunk write never evicts keys inside its own causal window.
+                initialize_state_pool(self._state_capacity, scratch_capacity, chunk_frames=self._chunk_frames)
+            else:
+                initialize_state_pool(self._state_capacity, scratch_capacity)
         if batch_sizes and frame_sizes and self._device.type in ("cuda", "npu"):
             if self._device.type == "npu":
                 from vllm_omni.platforms.npu.models.moss_tts_streaming_decode_wrapper import (
@@ -101,16 +144,34 @@ class _MossCodecStreamSession:
                     vllm_config=vllm_config,
                 )
             self._cudagraph_wrapper.warmup(self._device)
+            if fast_batch_sizes and self._cudagraph_wrapper.is_ready:
+                main_compiled = self._cudagraph_wrapper._compiled_decode
+                self._fast_wrapper = CUDAGraphStreamingDecoderWrapper(
+                    codec,
+                    state_capacity=self._state_capacity,
+                    batch_sizes=fast_batch_sizes,
+                    frame_sizes=[int(fast_frames)],
+                    num_quantizers=self._n_vq,
+                    vllm_config=vllm_config,
+                    scratch_base=fast_scratch_base,
+                    private_pool=True,
+                    compiled_decode=main_compiled if main_compiled is not None else codec.decode_streaming_tensors,
+                )
+                self._fast_wrapper.warmup(self._device)
+                if not self._fast_wrapper.is_ready:
+                    self._fast_wrapper = None
             self.reset_slots(list(range(self._state_capacity + scratch_capacity)))
             if not self._cudagraph_wrapper.is_ready:
                 self._cudagraph_wrapper = None
 
-    def acquire(self) -> int | None:
-        if not self._free_stream_slots:
-            return None
-        slot = self._free_stream_slots.pop()
-        self._leased_slots.add(slot)
-        return slot
+    def acquire(self, *, oldest: bool = False) -> int | None:
+        """Lease a slot; ``oldest`` takes the one whose reset was queued first."""
+        with self._slot_lock:
+            if not self._free_stream_slots:
+                return None
+            slot = self._free_stream_slots.pop(0 if oldest else -1)
+            self._leased_slots.add(slot)
+            return slot
 
     def release(self, slot: int, *, state_already_reset: bool = False) -> None:
         if self._closed:
@@ -119,8 +180,23 @@ class _MossCodecStreamSession:
             return
         if not state_already_reset:
             self.reset_slots([slot])
-        self._leased_slots.remove(slot)
-        self._free_stream_slots.append(slot)
+        reset_event = None
+        if self._fast_wrapper is not None:
+            # The fast path decodes on another stream; it must not read this
+            # slot before the reset queued on the current stream has run.
+            reset_event = torch.cuda.Event()
+            reset_event.record(torch.cuda.current_stream(self._device))
+        with self._slot_lock:
+            if reset_event is not None:
+                self._slot_reset_events[slot] = reset_event
+            self._leased_slots.remove(slot)
+            self._free_stream_slots.append(slot)
+
+    def order_after_reset(self, slot: int, stream: torch.cuda.Stream) -> None:
+        with self._slot_lock:
+            event = self._slot_reset_events.pop(slot, None)
+        if event is not None and not event.query():
+            stream.wait_event(event)
 
     def _reset_slot_ids(self, slot_ids: torch.Tensor) -> None:
         reset_streaming_slots = getattr(self._codec, "reset_decoder_state_slots", None)
@@ -139,9 +215,29 @@ class _MossCodecStreamSession:
             raise ValueError(f"Decoder state slots must be contiguous, got {slots}")
         self._reset_slot_ids(self._state_slot_ids[start : start + len(slots)])
 
+    def _stage_slot_ids(self, slots: list[int], terminal_slots: set[int]) -> tuple[torch.Tensor, torch.Tensor]:
+        # Host slabs cannot be overwritten until their DMA completes. Device
+        # slabs are reused in producer-stream order, after all previous reads.
+        index = self._metadata_index
+        host, view, gpu, event, recorded = self._metadata_ring[index]
+        if recorded and not event.query():
+            event.synchronize()
+        terminal_ids = [slot for slot in slots if slot in terminal_slots]
+        n, end = len(slots), len(slots) + len(terminal_ids)
+        view[:n] = slots
+        view[n:end] = terminal_ids
+        gpu[:end].copy_(host[:end], non_blocking=True)
+        event.record(torch.cuda.current_stream(self._device))
+        self._metadata_ring[index] = (host, view, gpu, event, True)
+        self._metadata_index = (index + 1) % len(self._metadata_ring)
+        return gpu[:n], gpu[n:end]
+
     def close(self) -> None:
         if self._closed:
             return
+        for _, _, _, event, recorded in self._metadata_ring:
+            if recorded:
+                event.synchronize()
         close_state_pool = getattr(self._codec, "close_decoder_state_pool", None)
         with torch.no_grad():
             if callable(close_state_pool):
@@ -154,6 +250,7 @@ class _MossCodecStreamSession:
         slot_codes: dict[int, torch.Tensor],
         *,
         terminal_slots: set[int] | None = None,
+        output_buffers: dict[int, torch.Tensor] | None = None,
     ) -> dict[int, torch.Tensor]:
         if not slot_codes:
             return {}
@@ -178,17 +275,22 @@ class _MossCodecStreamSession:
             ],
             dim=1,
         )
-        state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long)
+        staged_terminal_ids = None
+        if self._metadata_ring:
+            state_slot_ids, staged_terminal_ids = self._stage_slot_ids(slots, terminal_slots)
+        else:
+            state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long)
 
         graph_output: tuple[torch.Tensor, torch.Tensor, int] | None = None
         if self._cudagraph_wrapper is not None:
-            # A terminal tail can safely pad T to a larger graph: decoder
-            # attention is causal, audio_lengths trims the padding, and every
-            # affected state slot is reset immediately after this step.
+            # An all-terminal group may borrow the next larger frame graph.
+            # The decoder is causal, so cropping the padded tail on the time
+            # axis recovers the exact-T waveform up to attention-ring
+            # eviction (see codec_ring_headroom); the state is reset below.
             graph_output = self._cudagraph_wrapper.decode(
                 codes_step,
                 state_slot_ids,
-                allow_frame_padding=len(terminal_slots) == len(slots),
+                allow_frame_padding=bool(terminal_slots) and len(terminal_slots) == len(slots),
             )
 
         used_cudagraph = graph_output is not None
@@ -214,11 +316,69 @@ class _MossCodecStreamSession:
             audio_tensor = result.audio
 
         if terminal_slots:
-            terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots]
-            terminal_slot_ids = state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows]
+            if staged_terminal_ids is not None:
+                terminal_slot_ids = staged_terminal_ids
+            else:
+                terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots]
+                terminal_slot_ids = (
+                    state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows]
+                )
             self._reset_slot_ids(terminal_slot_ids)
 
-        audio = audio_tensor.detach().to("cpu", torch.float32)
+        if self._gpu_output and output_buffers is not None:
+            # Write directly into the caller-owned per-forward slab.  This
+            # preserves graph output before the next replay without creating
+            # a temporary waveform tensor that would be packed again later.
+            out: dict[int, torch.Tensor] = {}
+            row = 0
+            while row < len(slots):
+                destination = output_buffers.get(slots[row])
+                if destination is None:
+                    row += 1
+                    continue
+                end = row + 1
+                if end < len(slots) and destination.is_contiguous():
+                    storage = destination.untyped_storage().data_ptr()
+                    offset = destination.storage_offset()
+                    row_elements = destination.numel()
+                    while end < len(slots):
+                        following = output_buffers.get(slots[end])
+                        if not (
+                            following is not None
+                            and following.shape == destination.shape
+                            and following.dtype == destination.dtype
+                            and following.device == destination.device
+                            and following.is_contiguous()
+                            and following.untyped_storage().data_ptr() == storage
+                            and following.storage_offset() == offset + (end - row) * row_elements
+                        ):
+                            break
+                        end += 1
+                if end > row + 1:
+                    # A mixed-frame forward can leave gaps between equal-sized
+                    # rows in its slab. Copy each adjacent run as one batch,
+                    # without overwriting the intervening requests' output.
+                    destination = destination.as_strided(
+                        (end - row, *destination.shape), (destination.numel(), *destination.stride())
+                    )
+                    source = audio_tensor[row:end]
+                else:
+                    source = audio_tensor[row]
+                # Crop time before reshaping: flattening first mixes stereo
+                # channels whenever a row is shorter than the decode batch.
+                destination.copy_(source[..., : destination.shape[-1]].view_as(destination))
+                out.update((slot, output_buffers[slot]) for slot in slots[row:end])
+                row = end
+            return out
+
+        if self._gpu_output:
+            # A graph's static output is overwritten on every replay, including
+            # multiple replays within one forward. Snapshot on the producer
+            # stream before returning; record_stream alone cannot prevent this
+            # overwrite. MRV2 owns the subsequent pinned, asynchronous D2H.
+            audio = audio_tensor.detach().to(dtype=torch.float32, copy=True)
+        else:
+            audio = audio_tensor.detach().to("cpu", torch.float32)
         out: dict[int, torch.Tensor] = {}
         for row, slot in enumerate(slots):
             audio_length = int(slot_codes[slot].shape[1]) * self._samples_per_frame
@@ -270,6 +430,7 @@ class MossTTSCodecDecoder(nn.Module):
     requires_exact_input_shape: bool = True
 
     _OUTPUT_SAMPLE_RATE: int = 24_000
+    _first_chunk_fast_path: MossFirstChunkFastPath | None = None
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         super().__init__()
@@ -292,20 +453,38 @@ class MossTTSCodecDecoder(nn.Module):
         self._n_channels: int = 1
         self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
         self._stream_session: _MossCodecStreamSession | None = None
+        self._codec_stream: torch.cuda.Stream | None = None
         scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None)
         self._stream_state_capacity = max(1, int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1))
         self._initial_stream_chunk_frames: int = self._connector_int("initial_codec_chunk_frames", default=0)
         self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0)
         self._stream_max_step_frames: int = self._stream_chunk_frames or 100
+        # Opt-in: size codec attention rings as context + one chunk so chunked
+        # streaming reproduces whole-sequence decoding (costs ring memory).
+        self._stream_ring_headroom: bool = bool(self._connector_int("codec_ring_headroom", default=0))
         self._stream_req_slots: dict[str, int] = {}
+        # Opt-in: decode each stream's first chunk from the receive thread.
+        self._first_chunk_fast = bool(self._connector_int("codec_first_chunk_fast_path", default=0))
+        self._first_chunk_fast_path: MossFirstChunkFastPath | None = None
         self._async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False))
+        self._gpu_stream_output = (
+            self._async_chunk
+            and bool(getattr(vllm_config.model_config, "use_v2_model_runner", False))
+            and getattr(cfg, "model_type", None) == "moss_tts_local"
+            # Diagnostic switch: 0 copies each decode group's batch to the host synchronously.
+            and bool(self._connector_int("codec_gpu_stream_output", default=1))
+        )
         self._streaming_graph_batch_sizes = self._streaming_graph_batch_sizes_from_compilation_config()
-        # codec_chunk_ramp ladder sizes join the streaming graph frame-size
-        # set so every ladder (B, T) combination is captured during warmup.
+        ramp = parse_chunk_ramp(self._connector_extra(), steady=self._stream_chunk_frames)
+        if ramp:
+            # The sender gives the ramp precedence over initial chunk size.
+            # The dedicated first-chunk graph must use that same first shape.
+            self._initial_stream_chunk_frames = ramp[0]
+        self._stream_max_step_frames = max(
+            self._stream_max_step_frames, self._initial_stream_chunk_frames, *(ramp or [])
+        )
         self._streaming_graph_frame_sizes = _resolve_streaming_graph_frame_sizes(
-            self._initial_stream_chunk_frames,
-            self._stream_chunk_frames,
-            self._connector_extra_cfg(),
+            self._initial_stream_chunk_frames, self._stream_chunk_frames, self._connector_extra()
         )
 
     # ------------------------------------------------------------------
@@ -337,6 +516,41 @@ class MossTTSCodecDecoder(nn.Module):
         inputs_embeds: torch.Tensor | None = None,
         runtime_additional_information: list[dict[str, Any]] | None = None,
         **kwargs: Any,
+    ) -> OmniOutput:
+        if self._gpu_stream_output and input_ids is not None and input_ids.is_cuda:
+            if self._codec_stream is None:
+                self._codec_stream = torch.cuda.Stream(device=input_ids.device)
+            main_stream = torch.cuda.current_stream(input_ids.device)
+            # Runner input storage is reused by the next batch on main_stream.
+            # Own it before allowing that preparation to overlap codec decode.
+            owned_input_ids = input_ids.clone()
+            with torch.cuda.stream(self._codec_stream):
+                self._codec_stream.wait_stream(main_stream)
+                owned_input_ids.record_stream(self._codec_stream)
+                output = self._forward(
+                    owned_input_ids,
+                    positions,
+                    intermediate_tensors,
+                    inputs_embeds,
+                    runtime_additional_information,
+                    **kwargs,
+                )
+                if isinstance(output.multimodal_outputs, PackedOutputSnapshot):
+                    output.multimodal_outputs.record_producer_event(self._codec_stream)
+            return output
+        return self._forward(
+            input_ids, positions, intermediate_tensors, inputs_embeds, runtime_additional_information, **kwargs
+        )
+
+    @torch.no_grad()
+    def _forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
     ) -> OmniOutput:
         """Decode audio VQ codes to waveform.
 
@@ -493,13 +707,37 @@ class MossTTSCodecDecoder(nn.Module):
             audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
 
         if streaming_work:
-            for i, wav in self._decode_streaming_batch(streaming_work).items():
+            stream_output_slab = None
+            stream_output_buffers: dict[int, torch.Tensor] = {}
+            if self._gpu_stream_output and device.type == "cuda":
+                total_elements = 0
+                for output_index, _, codes_nq_t, _ in streaming_work:
+                    samples = int(codes_nq_t.shape[1]) * self._codec.downsample_rate
+                    total_elements += samples * (self._n_channels if self._n_channels > 1 else 1)
+                if total_elements:
+                    stream_output_slab = torch.empty(total_elements, dtype=torch.float32, device=device)
+                    offset = 0
+                    for output_index, _, codes_nq_t, _ in streaming_work:
+                        samples = int(codes_nq_t.shape[1]) * self._codec.downsample_rate
+                        elements = samples * (self._n_channels if self._n_channels > 1 else 1)
+                        shape = (self._n_channels, samples) if self._n_channels > 1 else (samples,)
+                        stream_output_buffers[output_index] = stream_output_slab.narrow(0, offset, elements).view(shape)
+                        offset += elements
+            for i, wav in self._decode_streaming_batch(
+                streaming_work,
+                output_buffers=stream_output_buffers or None,
+            ).items():
                 audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
 
-        return OmniOutput(
-            text_hidden_states=None,
-            multimodal_outputs={"model_outputs": audios, "sr": srs},
-        )
+        payload = {"model_outputs": audios, "sr": srs}
+        if self._gpu_stream_output and device.type == "cuda":
+            # One transfer per dtype/device preserves ragged request lengths
+            # without issuing one D2H per waveform. A fresh slot owns each
+            # output until the runner completes its copy; no reusable ring.
+            packed = pack_output_snapshot(payload, {}, max_buckets=0, reuse_existing_storage=True)
+            if packed is not None:
+                payload = packed
+        return OmniOutput(text_hidden_states=None, multimodal_outputs=payload)
 
     @staticmethod
     def _normalize_seq_token_counts(value: Any) -> list[int] | None:
@@ -546,25 +784,104 @@ class MossTTSCodecDecoder(nn.Module):
             vllm_config=self.vllm_config,
             graph_batch_sizes=self._streaming_graph_batch_sizes,
             graph_frame_sizes=self._streaming_graph_frame_sizes,
+            gpu_output=self._gpu_stream_output,
+            chunk_frames=self._stream_max_step_frames,
+            ring_headroom=self._stream_ring_headroom,
+            fast_batch_sizes=[1, 2, 4, 8] if self._first_chunk_fast else None,
+            fast_frames=self._initial_stream_chunk_frames if self._first_chunk_fast else 0,
         )
+        if self._stream_session._fast_wrapper is not None:
+            self._first_chunk_fast_path = MossFirstChunkFastPath(
+                self._stream_session,
+                self._stream_session._fast_wrapper,
+                n_vq=self._n_vq,
+                frames=self._initial_stream_chunk_frames,
+                codebook_size=int(self._codec.config.codebook_size),
+                samples_per_frame=int(self._codec.downsample_rate),
+                n_channels=self._n_channels,
+                sample_rate=self._sr_tensor,
+                device=next(self._codec.parameters()).device,
+                gate_main=bool(self._connector_int("codec_first_chunk_gate", default=0)),
+            )
+            logger.info(
+                "MOSS-TTS codec first-chunk fast path ready: B=%s T=%d",
+                self._stream_session._fast_wrapper.capture_sizes,
+                self._initial_stream_chunk_frames,
+            )
+        elif self._first_chunk_fast:
+            logger.warning("MOSS-TTS codec first-chunk fast path requested but unavailable; using the main path.")
+        if self._stream_ring_headroom:
+            logger.info(
+                "MOSS-TTS codec attention rings sized with headroom for %d-frame chunks",
+                self._stream_max_step_frames,
+            )
         return self._stream_session
 
+    def bind_first_chunk_fast_path(self, sink: FirstAudioSink) -> Any:
+        """Start the fast path with an output sink; returns the receive hook."""
+        fast = self._first_chunk_fast_path
+        if fast is None:
+            return None
+        fast.bind(sink)
+        return self._submit_first_chunk
+
+    def _submit_first_chunk(self, request_id: str, handle: Any, payload: Any) -> bool:
+        fast = self._first_chunk_fast_path
+        if fast is None or not isinstance(payload, dict):
+            return False
+        codes = payload.get("codes")
+        codes = codes.get("audio") if isinstance(codes, dict) else None
+        if codes is None:
+            return False
+        meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {}
+        key = meta.get("req_id")
+        if isinstance(key, (list, tuple)):
+            key = key[0] if key else None
+        if key is None:
+            key = getattr(handle, "external_req_id", None) or request_id
+        return fast.submit(request_id, str(key), codes, self._stream_req_slots)
+
     def _decode_streaming_batch(
         self,
         items: list[tuple[int, str, torch.Tensor, bool]],
+        *,
+        output_buffers: dict[int, torch.Tensor] | None = None,
     ) -> dict[int, torch.Tensor]:
         session = self._ensure_stream_session()
         if session is None:
             return {}
 
         outputs: dict[int, torch.Tensor] = {}
+        unsent_first_audio: dict[int, torch.Tensor] = {}
         grouped: dict[int, list[tuple[int, str, int, torch.Tensor, bool]]] = {}
         max_step_frames = max(1, int(self._stream_max_step_frames))
         graph = session._cudagraph_wrapper
+        # Terminal tails join the regular padded frame bucket when every
+        # execution batch width has a graph, so a step stays one replay
+        # instead of one eager or plain-graph decode per distinct tail length.
         coalesce_tails = graph is not None and max(graph.batch_sizes, default=0) >= self._stream_state_capacity
 
+        fast = self._first_chunk_fast_path
         for output_index, request_id, codes_nq_t, finished in items:
             slot = self._stream_req_slots.get(request_id)
+            if fast is not None and slot is not None:
+                fast.order_after(slot)
+                if fast.take_decoded(request_id):
+                    unsent = fast.take_unsent_audio(request_id)
+                    if unsent is not None:
+                        unsent_first_audio[output_index] = unsent
+                    # The fast path already emitted this chunk's audio and
+                    # advanced the slot state past it.
+                    codes_nq_t = codes_nq_t[:, fast.frames :]
+                    if codes_nq_t.shape[1] == 0:
+                        if unsent is not None:
+                            outputs[output_index] = unsent
+                        if finished:
+                            self._finish_stream_request(request_id, session, slot)
+                        continue
+                    if output_buffers is not None:
+                        # Sized for the untrimmed chunk; decode into a fresh tensor.
+                        output_buffers.pop(output_index, None)
             if slot is None:
                 slot = session.acquire()
                 if slot is None:
@@ -580,9 +897,11 @@ class MossTTSCodecDecoder(nn.Module):
                     slot,
                     codes_nq_t,
                     terminal=finished,
+                    output_buffer=output_buffers.get(output_index) if output_buffers is not None else None,
                 )
                 if wav is not None:
-                    outputs[output_index] = wav
+                    prefix = unsent_first_audio.pop(output_index, None)
+                    outputs[output_index] = torch.cat((prefix, wav), dim=-1) if prefix is not None else wav
                 if finished:
                     self._finish_stream_request(
                         request_id,
@@ -593,21 +912,33 @@ class MossTTSCodecDecoder(nn.Module):
                 continue
 
             frame_count = int(codes_nq_t.shape[1])
-            # Only combine tails that already used the same padded graph.
-            # Preserve exact-size graphs (including T=1) and the eager path.
             group_frames = frame_count
             if coalesce_tails and finished:
+                # Only terminal rows may be padded: their state is reset after
+                # the step, and the output is cropped to the exact frame count.
                 group_frames = graph._select_frame_size(frame_count, allow_padding=True) or frame_count
             grouped.setdefault(group_frames, []).append((output_index, request_id, slot, codes_nq_t, finished))
 
+        if fast is not None and grouped:
+            fast.gate()
         for group in grouped.values():
             plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group}
             terminal_slots = {slot for _, _, slot, _, finished in group if finished}
-            decoded = session.step(plan, terminal_slots=terminal_slots)
+            step_buffers = (
+                {
+                    slot: output_buffers[output_index]
+                    for output_index, _, slot, _, _ in group
+                    if output_buffers is not None and output_index in output_buffers
+                }
+                if output_buffers is not None
+                else None
+            )
+            decoded = session.step(plan, terminal_slots=terminal_slots, output_buffers=step_buffers)
             for output_index, request_id, slot, _, finished in group:
                 wav = decoded.get(slot)
                 if wav is not None:
-                    outputs[output_index] = wav
+                    prefix = unsent_first_audio.pop(output_index, None)
+                    outputs[output_index] = torch.cat((prefix, wav), dim=-1) if prefix is not None else wav
                 if finished:
                     self._finish_stream_request(
                         request_id,
@@ -625,6 +956,7 @@ class MossTTSCodecDecoder(nn.Module):
         codes_nq_t: torch.Tensor,
         *,
         terminal: bool = False,
+        output_buffer: torch.Tensor | None = None,
     ) -> torch.Tensor | None:
         if codes_nq_t.numel() == 0:
             return None
@@ -633,13 +965,24 @@ class MossTTSCodecDecoder(nn.Module):
         for start in range(0, int(codes_nq_t.shape[1]), max_step_frames):
             chunk = codes_nq_t[:, start : start + max_step_frames]
             is_terminal_chunk = terminal and start + max_step_frames >= int(codes_nq_t.shape[1])
+            chunk_samples = int(chunk.shape[1]) * self._codec.downsample_rate
+            chunk_buffer = (
+                output_buffer[
+                    ..., start * self._codec.downsample_rate : start * self._codec.downsample_rate + chunk_samples
+                ]
+                if output_buffer is not None
+                else None
+            )
             decoded = session.step(
                 {slot: chunk},
                 terminal_slots={slot} if is_terminal_chunk else None,
+                output_buffers={slot: chunk_buffer} if chunk_buffer is not None else None,
             )
             wav = decoded.get(slot)
             if wav is not None:
                 parts.append(wav)
+        if output_buffer is not None:
+            return output_buffer
         if not parts:
             return None
         return torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0]
@@ -665,29 +1008,30 @@ class MossTTSCodecDecoder(nn.Module):
         terminal payload, so the runner calls this hook from its finished-request
         path to avoid leaking stream slots and buffered codes.
         """
-        session = self._stream_session
-        for req_id in finished_req_ids:
-            request_id = str(req_id)
-            slot = self._stream_req_slots.get(request_id)
-            if slot is None:
-                continue
-            if session is not None:
-                self._finish_stream_request(request_id, session, slot)
-            else:
-                self._stream_req_slots.pop(request_id, None)
+        # Cancellation and terminal cleanup share the decoder's stream order;
+        # a reused slot must not be reset ahead of an outstanding decode.
+        context = torch.cuda.stream(self._codec_stream) if self._codec_stream is not None else nullcontext()
+        with context:
+            session = self._stream_session
+            for req_id in finished_req_ids:
+                request_id = str(req_id)
+                fast = self._first_chunk_fast_path
+                slot = (
+                    fast.get_request_slot(request_id, self._stream_req_slots)
+                    if fast is not None
+                    else self._stream_req_slots.get(request_id)
+                )
+                if slot is None:
+                    continue
+                if fast is not None:
+                    fast.order_after(slot)
+                    fast.forget(request_id)
+                if session is not None:
+                    self._finish_stream_request(request_id, session, slot)
+                else:
+                    self._stream_req_slots.pop(request_id, None)
 
-    def _connector_int(self, name: str, default: int = 0) -> int:
-        model_cfg = getattr(self.vllm_config, "model_config", None)
-        connector_cfg = getattr(model_cfg, "stage_connector_config", None)
-        if isinstance(connector_cfg, dict):
-            extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg)
-        else:
-            extra_cfg = getattr(connector_cfg, "extra", None)
-        if isinstance(extra_cfg, dict) and name in extra_cfg:
-            return int(extra_cfg[name])
-        return default
-
-    def _connector_extra_cfg(self) -> dict:
+    def _connector_extra(self) -> dict[str, Any]:
         model_cfg = getattr(self.vllm_config, "model_config", None)
         connector_cfg = getattr(model_cfg, "stage_connector_config", None)
         if isinstance(connector_cfg, dict):
@@ -696,6 +1040,9 @@ class MossTTSCodecDecoder(nn.Module):
             extra_cfg = getattr(connector_cfg, "extra", None)
         return extra_cfg if isinstance(extra_cfg, dict) else {}
 
+    def _connector_int(self, name: str, default: int = 0) -> int:
+        return int(self._connector_extra().get(name, default))
+
     def _streaming_graph_batch_sizes_from_compilation_config(self) -> list[int]:
         if getattr(self.vllm_config.model_config, "enforce_eager", True):
             return []
@@ -816,15 +1163,33 @@ class MossTTSCodecDecoder(nn.Module):
             codec.decoder.to(dtype=torch.bfloat16)
         attention_backend = getattr(self.vllm_config.model_config.hf_config, "codec_attention_backend", "sdpa")
         if attention_backend != "sdpa":
-            if attention_backend != "triton" or device.type != "cuda":
+            if attention_backend not in {"triton", "triton_slot"} or device.type != "cuda":
                 raise ValueError(f"Unsupported codec attention backend/device: {attention_backend}/{device.type}")
             from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerMultiheadAttention
             from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention
 
+            if attention_backend == "triton_slot":
+                from vllm_omni.model_executor.models.moss_tts.slot_attention import (
+                    slot_ring_attention,
+                    slot_ring_attention_rows,
+                )
+
+                skip_empty_tiles = bool(self._connector_int("codec_skip_empty_attention_tiles", default=0))
+                fused_slots = bool(self._connector_int("codec_fused_slot_attention", default=0))
+                if skip_empty_tiles and fused_slots:
+                    raise ValueError("codec_skip_empty_attention_tiles requires unfused slot attention")
+                if skip_empty_tiles:
+                    slot_ring_attention = partial(slot_ring_attention, skip_empty_tiles=True)
+                    logger.info("MOSS codec slot attention: skipping empty tiles for T <= 32")
+
             for module in codec.decoder.modules():
                 if isinstance(module, MossAudioTokenizerMultiheadAttention):
                     module._streaming_attention = masked_attention
-            logger.info("Enabled Triton masked attention for the streaming codec decoder")
+                    if attention_backend == "triton_slot":
+                        module._slot_attention = slot_ring_attention
+                        if self._connector_int("codec_fused_slot_attention", default=0):
+                            module._slot_attention_rows = slot_ring_attention_rows
+            logger.info("Enabled codec attention backend=%s", attention_backend)
         build_decode_lut = getattr(codec.quantizer, "build_decode_lut", None)
         if callable(build_decode_lut):
             lut_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -21,7 +21,8 @@ checkpoint 1:1 so ``load_weights()`` needs no remapping.
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+from collections.abc import Callable, Sequence
 
 import torch
 import torch.nn as nn
@@ -52,6 +53,7 @@ class _MossTTSLocalAttention(nn.Module):
         self.n_head = n_head
         self.head_dim = hidden_size // n_head
         self.embed_dim = hidden_size
+        self._short_attention = os.environ.get("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "0") == "1"
         self.c_attn = nn.Linear(hidden_size, 3 * hidden_size, bias=True)
         self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True)
         inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim))
@@ -134,7 +136,19 @@ class _MossTTSLocalAttention(nn.Module):
             v_cache[:, :, position : position + 1].copy_(value)
             key = k_cache[:, :, : position + 1]
             value = v_cache[:, :, : position + 1]
-        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None)
+        if (
+            self._short_attention
+            and kv_cache is not None
+            and query.is_cuda
+            and query.dtype == torch.bfloat16
+            and self.head_dim <= 128
+            and key.shape[2] <= 16
+        ):
+            from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+            attn_output = local_short_attention(query, key, value)
+        else:
+            attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None)
         attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
         return self.c_proj(attn_output)
 
@@ -193,6 +207,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
         self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)])
         self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps)
         self._compiled_forward_prefix = None
+        self._compiled_audio_sampler = None
 
     def _forward_prefix(
         self,
@@ -216,6 +231,14 @@ class MossTTSLocalDepthTransformer(nn.Module):
             dynamic=True,
             options={"epilogue_fusion": False},
         )
+        if os.environ.get("VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", "0") == "1":
+            # Keep the existing top-k/top-p algorithm and torch RNG. Explicit
+            # per-request generators use the original helper below. The
+            # binary continue/stop head is deliberately unchanged.
+            self._compiled_audio_sampler = torch.compile(
+                _sample_token, fullgraph=True, dynamic=True, options={"fallback_random": True}
+            )
+            logger.info("MOSS-TTS Local compiled audio-channel sampler enabled")
         logger.info("MOSS-TTS local depth frame-local KV execution enabled with torch.compile")
 
     def _run_prefix(
@@ -239,6 +262,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
         do_sample: bool,
         generator: torch.Generator | None,
         generators: Sequence[torch.Generator | None] | None = None,
+        sample_token: Callable[..., torch.Tensor] | None = None,
     ) -> torch.Tensor:
         """Compute channel logits, apply repetition penalty, sample, store."""
         channel_logits = audio_lm_heads[channel_index](local_hidden).float()
@@ -250,7 +274,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
                 pos = sel > 0
                 sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty)
                 channel_logits.index_copy_(-1, hist_t, sel)
-        channel_token = _sample_token(
+        channel_token = (sample_token or _sample_token)(
             channel_logits,
             temperature,
             top_k,
@@ -351,6 +375,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
                 do_sample=do_sample,
                 generator=generator,
                 generators=generators,
+                sample_token=self._compiled_audio_sampler if generator is None and generators is None else None,
             )
 
             if channel_index + 1 < n_vq:

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -37,6 +37,7 @@ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth impo
 )
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker.sampling_utils import get_tts_local_seed
 
 logger = init_logger(__name__)
 
@@ -1675,6 +1676,31 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         return input_embeds_out, output_codes
 
     # ------------------------------------------------------------------
+    # MRV2 reads model-owned capabilities; retain the V1 hook and platform
+    # graph-safety setting as the canonical implementation.
+    mtp = talker_mtp
+    get_mtp_seed = staticmethod(get_tts_local_seed)
+
+    def create_mrv2_model_state(self, vllm_config, encoder_cache, device):
+        from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+        state_cls = OmniModelState
+        if getattr(self.config, "mrv2_gpu_slot_state", False):
+            from .local_model_state import MossLocalModelState
+
+            state_cls = MossLocalModelState
+        return state_cls(vllm_config, self, encoder_cache, device)
+
+    @property
+    def mtp_output_key(self) -> tuple[str, str]:
+        """Return the request-buffer destination for one local audio frame."""
+        return self.talker_mtp_output_key
+
+    @property
+    def mtp_graph_safe(self) -> bool:
+        """Honor the same platform graph-safety override as the V1 runner."""
+        return self.talker_mtp_graph_safe
+
     # Package runner-generated audio frames
     # ------------------------------------------------------------------
 

vllm_omni/model_executor/models/moss_tts/pipeline.py

@@ -97,6 +97,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
         StagePipelineConfig(
             stage_id=0,
             model_stage="moss_tts_local",
+            supports_native_mrv2_data_plane=True,
             execution_type=StageExecutionType.LLM_AR,
             input_sources=(),
             owns_tokenizer=True,
@@ -114,6 +115,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
         StagePipelineConfig(
             stage_id=1,
             model_stage="moss_tts_local_codec",
+            supports_native_mrv2_data_plane=True,
             execution_type=StageExecutionType.LLM_GENERATION,
             input_sources=(0,),
             final_output=True,

vllm_omni/model_executor/models/moss_tts/reference_encoder.py

@@ -25,7 +25,10 @@ so importing it from the API-server process does not pull the talker/codec.
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import hashlib
+import os
+from collections import OrderedDict
 from collections.abc import Awaitable, Callable
 from functools import lru_cache
 from typing import Any
@@ -35,22 +38,93 @@ import torch
 from numpy.typing import NDArray
 from vllm.logger import init_logger
 
+from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import (
+    FINE_BATCH_SIZES,
+    FINE_BUCKET_SECONDS,
+    MossReferenceEncoderGraphs,
+    install_windowed_attention,
+)
+from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import (
+    SharedReferenceEncoderClient,
+    SharedReferenceEncoderHost,
+    SharedReferenceEncoderStartupError,
+    elect_host,
+)
+
 logger = init_logger(__name__)
 
 # Coalescing defaults. Kept as module constants (not env/CLI) to match the
 # ``_REF_AUDIO_RESOLVE_CACHE_MAX_*`` convention in serving_speech.py.
 _REF_ENCODE_BATCH_WINDOW_MS = 10.0
-_REF_ENCODE_MAX_BATCH = 8
+_REF_ENCODE_MAX_BATCH = 32
+# Experiment switches for cold-reference latency (see the A/B notes):
+# batch window in ms; batches in flight per API process with a shared encoder;
+# "high" puts reference encodes on the highest-priority CUDA stream.
+_REF_ENCODE_WINDOW_ENV = "VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS"
+_REF_ENCODE_INFLIGHT_ENV = "VLLM_OMNI_MOSS_REF_INFLIGHT"
+_REF_ENCODE_PRIORITY_ENV = "VLLM_OMNI_MOSS_REF_STREAM_PRIORITY"
+# "1" resamples and loudness-normalizes clips on the encoder's GPU instead of
+# the API process's CPU (not bit-identical to the CPU path).
+_REF_ENCODE_GPU_PREP_ENV = "VLLM_OMNI_MOSS_REF_GPU_PREP"
 
 _INT32_MAX = 2**31
+# Inline (data: URI) references already resolved and encoded, keyed by the
+# serving resolve key. Codes are small; the URI itself is not retained.
+_INLINE_INDEX_MAX_ENTRIES = 4096
+# Optional directory (e.g. under /dev/shm) that API processes of one server
+# share, so an inline reference encoded by any frontend is reused by all.
+_SHARED_CODES_DIR_ENV = "VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR"
+# "0" keeps one reference encoder per API process even with a shared directory.
+_SHARED_ENCODER_ENV = "VLLM_OMNI_MOSS_REF_SHARED_ENCODER"
+_shared_role: str | None = None
+# Concurrent encode workers in the shared host, each with its own CUDA graphs.
+_SHARED_WORKERS_ENV = "VLLM_OMNI_MOSS_REF_ENCODER_WORKERS"
+_SHARED_WORKER_BATCH_SECONDS = 96.0
+
+
+def shared_encoder_role() -> str:
+    """``"host"``, ``"client"`` or ``""`` (no shared encoder) for this process.
+
+    Decided once, before the processor places its tokenizer on the GPU, so
+    client processes can leave it on the CPU.
+    """
+    global _shared_role
+    if _shared_role is None:
+        shared_dir = os.environ.get(_SHARED_CODES_DIR_ENV)
+        if not shared_dir or os.environ.get(_SHARED_ENCODER_ENV, "1") == "0":
+            _shared_role = ""
+        else:
+            os.makedirs(shared_dir, exist_ok=True)
+            _shared_role = "host" if elect_host(shared_dir) else "client"
+    return _shared_role
+
+
+# Opt in to reference graphs; existing deployments keep the eager encoder.
+_REF_ENCODE_GRAPHS_ENV = "VLLM_OMNI_MOSS_REF_GRAPHS"
+
+
+def reference_graphs_enabled() -> bool:
+    return os.environ.get(_REF_ENCODE_GRAPHS_ENV, "0") == "1"
+
+
+def reference_stream_priority() -> int:
+    """CUDA stream priority for reference encodes (lower is more urgent)."""
+    if os.environ.get(_REF_ENCODE_PRIORITY_ENV, "") == "high" and torch.cuda.is_available():
+        return torch.cuda.Stream.priority_range()[1]
+    return 0
+
+
+# "sdpa" keeps the tokenizer's dense masked SDPA attention. By default the
+# reference encoder attends through a local-window flash kernel, the backend
+# the tokenizer config selects (flash_attention_2) when flash-attn is present.
+_REF_ENCODE_ATTN_ENV = "VLLM_OMNI_MOSS_REF_ATTN"
+# "0" captures the encoder's eager kernels. By default the captured encoder is
+# compiled first, which fuses its many small elementwise kernels.
+_REF_ENCODE_COMPILE_ENV = "VLLM_OMNI_MOSS_REF_COMPILE"
 
 _Waveform = NDArray[np.float32] | list[float]
 
 
-def _sha1(s: str) -> str:
-    return hashlib.sha1((s or "").encode("utf-8")).hexdigest()
-
-
 @lru_cache(maxsize=16)
 def _reference_resampler(sr: int, sr_target: int) -> torch.nn.Module:
     import torchaudio
@@ -72,6 +146,13 @@ def _prep_wav_sync(waveform: _Waveform, sr: int, sr_target: int) -> torch.Tensor
     return wav
 
 
+def _inline_reference_key(ref_str: str) -> str | None:
+    """SHA-1 of a ``data:`` URI, matching the serving resolve key for it."""
+    if not isinstance(ref_str, str) or ref_str[:5].lower() != "data:":
+        return None
+    return hashlib.sha1(ref_str.encode("utf-8")).hexdigest()
+
+
 def _to_compact_codes(codes: torch.Tensor) -> torch.Tensor:
     """Downcast RVQ codes to int32 for compact caching."""
     codes = codes.detach().cpu().contiguous()
@@ -119,18 +200,26 @@ class _RefEncodeBatcher:
         *,
         window_ms: float,
         max_batch: int,
+        max_inflight: int = 1,
     ):
         # encode_batch_fn: sync, takes [(waveform, sr), ...], returns a list of
         # (codes_tensor | Exception) aligned to the input order.
         self._encode_batch_fn = encode_batch_fn
         self._window_s = max(0.0, float(window_ms) / 1000.0)
         self._max_batch = max(1, int(max_batch))
+        # Batches encoded at once. Clips that arrive while every slot is busy
+        # are queued and form the next batch.
+        self._max_inflight = max(1, int(max_inflight))
         self._queue: asyncio.Queue | None = None
+        self._slots: asyncio.Semaphore | None = None
+        self._inflight: set[asyncio.Task] = set()
         self._drainer: asyncio.Task | None = None
 
     def _ensure_started(self) -> None:
         if self._queue is None:
             self._queue = asyncio.Queue()
+        if self._slots is None:
+            self._slots = asyncio.Semaphore(self._max_inflight)
         if self._drainer is None or self._drainer.done():
             self._drainer = asyncio.create_task(self._drain_loop())
 
@@ -141,16 +230,27 @@ class _RefEncodeBatcher:
         return await fut
 
     async def _drain_loop(self) -> None:
-        assert self._queue is not None
+        assert self._queue is not None and self._slots is not None
         while True:
-            first = await self._queue.get()
-            jobs = None
+            await self._slots.acquire()
             try:
-                jobs = await self._coalesce(first)
-                await self._run_batch(jobs)
-            finally:
-                # Do not retain the completed batch while waiting for new work.
-                del first, jobs
+                jobs = await self._coalesce(await self._queue.get())
+            except BaseException:
+                self._slots.release()
+                raise
+            task = asyncio.create_task(self._run_and_release(jobs))
+            self._inflight.add(task)
+            task.add_done_callback(self._inflight.discard)
+            # Do not retain the batch here while waiting for new work.
+            del jobs
+
+    async def _run_and_release(self, jobs: list[tuple[_Waveform, int, asyncio.Future]]) -> None:
+        try:
+            await self._run_batch(jobs)
+        finally:
+            del jobs
+            assert self._slots is not None
+            self._slots.release()
 
     async def _coalesce(self, first: tuple) -> list[tuple]:
         """Group ``first`` with jobs arriving within the batch window."""
@@ -200,6 +300,8 @@ class _RefEncodeBatcher:
             except (asyncio.CancelledError, Exception):  # noqa: BLE001
                 pass
             self._drainer = None
+        for task in list(self._inflight):
+            task.cancel()
 
 
 class MossReferenceEncoder:
@@ -213,7 +315,7 @@ class MossReferenceEncoder:
         n_vq: int,
         sr_target: int,
         speaker_cache: Any,
-        window_ms: float = _REF_ENCODE_BATCH_WINDOW_MS,
+        window_ms: float | None = None,
         max_batch: int = _REF_ENCODE_MAX_BATCH,
     ):
         self._processor = processor
@@ -224,8 +326,33 @@ class MossReferenceEncoder:
         # model_type namespaces the whole family so a moss_tts server never
         # collides with another model's speaker-cache entries.
         self._model_type = f"moss_tts_{variant}_nq{int(n_vq)}"
-        self._inflight: dict[str, asyncio.Task] = {}
-        self._batcher = _RefEncodeBatcher(self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch)
+        self._inflight: dict[str | tuple[str, str], asyncio.Task] = {}
+        self._inline_index: OrderedDict[str, torch.Tensor] = OrderedDict()
+        self._shared_codes_dir = os.environ.get(_SHARED_CODES_DIR_ENV) or None
+        if self._shared_codes_dir:
+            os.makedirs(self._shared_codes_dir, exist_ok=True)
+        # Realtime encodes with its own codec and never shares the encoder.
+        self._shares_encoder = bool(self._shared_codes_dir) and type(self)._encode_prepared is (
+            MossReferenceEncoder._encode_prepared
+        )
+        self._shared_host: SharedReferenceEncoderHost | None = None
+        self._shared_client: SharedReferenceEncoderClient | None = None
+        self._graphs: MossReferenceEncoderGraphs | None = None
+        self._attention_installed = False
+        self._graphs_enabled = reference_graphs_enabled()
+        if window_ms is None:
+            window_ms = float(os.environ.get(_REF_ENCODE_WINDOW_ENV, _REF_ENCODE_BATCH_WINDOW_MS))
+        # A shared encoder serves several batches at once (one per host
+        # worker); a process's own graphs replay one batch at a time.
+        inflight = 1
+        if self._shares_encoder and os.environ.get(_SHARED_ENCODER_ENV, "1") != "0":
+            inflight = max(1, int(os.environ.get(_REF_ENCODE_INFLIGHT_ENV, "1")))
+        self._local_stream: torch.cuda.Stream | None = None
+        self._gpu_prep = os.environ.get(_REF_ENCODE_GPU_PREP_ENV, "0") == "1"
+        self._gpu_resamplers: dict[torch.device, torch.nn.Module] = {}
+        self._batcher = _RefEncodeBatcher(
+            self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch, max_inflight=inflight
+        )
 
     def _make_cache_key(self, name: str, created_at: int) -> tuple:
         return self._speaker_cache.make_cache_key(name, model_type=self._model_type, created_at=int(created_at))
@@ -251,6 +378,7 @@ class MossReferenceEncoder:
         caller salts those requests with ``voice_created_at`` instead).
         """
         voice_name, created_at = _registered_voice(voice_name, voice_created_at)
+        inline_key: str | None = None
 
         if voice_name:
             # A named voice has a stable key that does not depend on the
@@ -260,6 +388,16 @@ class MossReferenceEncoder:
             if cached is not None:
                 return _clone_out(cached["codes"]), None
         else:
+            inline_key = _inline_reference_key(ref_str)
+            if inline_key is not None:
+                codes = self._inline_index.get(inline_key)
+                if codes is None:
+                    codes = self._load_shared_codes(inline_key)
+                if codes is not None:
+                    # A data: URI is immutable content, and its serving
+                    # resolve key is this same digest: skip the resolve hop.
+                    self._inline_index.move_to_end(inline_key)
+                    return _clone_out(codes), inline_key
             # Anonymous refs have no pre-resolve hot path: the content key must
             # come from the resolve itself (mtime/size from a single stat) so an
             # on-disk edit invalidates the cached codes; the flight body
@@ -267,17 +405,55 @@ class MossReferenceEncoder:
             # when the resolve cache is warm. The flight key is the request-side
             # reference (not the content hash), so concurrent requests for the
             # same ref_str also share the resolve/download, not just the encode.
-            flight_key = "ref:" + _sha1(ref_str)
+            # Key by the locator itself: equality is exact, and hashing a
+            # large inline data URI once per request object is much cheaper
+            # than a separate SHA-1 of its UTF-8 copy on the event loop.
+            flight_key = ("ref", ref_str)
 
         codes, resolve_key = await self._single_flight(
             flight_key,
             lambda: self._resolve_and_encode(ref_str, resolve_ref_audio, get_artifact_key, voice_name, created_at),
         )
+        if not voice_name and inline_key is not None and resolve_key == inline_key:
+            # Only index when the resolver's key is exactly this digest.
+            self._remember_inline(inline_key, codes)
+            self._store_shared_codes(inline_key, codes)
         return _clone_out(codes), resolve_key
 
+    def _remember_inline(self, key: str, codes: torch.Tensor) -> None:
+        self._inline_index[key] = codes
+        self._inline_index.move_to_end(key)
+        while len(self._inline_index) > _INLINE_INDEX_MAX_ENTRIES:
+            self._inline_index.popitem(last=False)
+
+    def _load_shared_codes(self, key: str) -> torch.Tensor | None:
+        if not self._shared_codes_dir:
+            return None
+        try:
+            array = np.load(os.path.join(self._shared_codes_dir, key + ".npy"), allow_pickle=False)
+        except (OSError, ValueError):
+            return None
+        if array.ndim != 2 or array.shape[1] != self._n_vq or array.dtype not in (np.int32, np.int64):
+            return None
+        codes = torch.from_numpy(array)
+        self._remember_inline(key, codes)
+        return codes
+
+    def _store_shared_codes(self, key: str, codes: torch.Tensor) -> None:
+        if not self._shared_codes_dir:
+            return
+        path = os.path.join(self._shared_codes_dir, key + ".npy")
+        tmp = f"{path}.{os.getpid()}.tmp"
+        try:
+            with open(tmp, "wb") as handle:
+                np.save(handle, codes.detach().cpu().numpy(), allow_pickle=False)
+            os.replace(tmp, path)
+        except OSError:
+            logger.warning("MOSS ref codes: could not write shared entry %s", path, exc_info=True)
+
     async def _single_flight(
         self,
-        flight_key: str,
+        flight_key: str | tuple[str, str],
         start_flight: Callable[[], Awaitable[tuple[torch.Tensor, str]]],
     ) -> tuple[torch.Tensor, str]:
         """Join the in-flight encode for ``flight_key``, starting one if absent.
@@ -300,7 +476,7 @@ class MossReferenceEncoder:
         # a duplicate resolve/encode instead of joining this one. Identity
         # guard: only drop the slot if it still holds *our* task (a later
         # request may have replaced it after ours completed).
-        def _retire(t: asyncio.Task, key: str = flight_key) -> None:
+        def _retire(t: asyncio.Task, key: str | tuple[str, str] = flight_key) -> None:
             if self._inflight.get(key) is t:
                 self._inflight.pop(key, None)
 
@@ -350,29 +526,228 @@ class MossReferenceEncoder:
         prepared_idx: list[int] = []
         for i, (waveform, sr) in enumerate(payload):
             try:
-                prepared.append(_prep_wav_sync(waveform, sr, self._sr_target))
+                prepared.append(self._encoder_input(_prep_wav_sync(waveform, sr, self._sr_target)))
                 prepared_idx.append(i)
             except Exception as exc:  # noqa: BLE001 — isolate this clip's failure
                 results[i] = exc
         if not prepared:
             return results
 
-        try:
-            codes_list = self._encode_prepared(prepared)
-            for local_i, orig_i in enumerate(prepared_idx):
-                results[orig_i] = codes_list[local_i]
-        except Exception:  # noqa: BLE001 — batch forward failed; retry item-by-item
-            logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared))
-            for local_i, orig_i in enumerate(prepared_idx):
-                try:
-                    results[orig_i] = self._encode_prepared([prepared[local_i]])[0]
-                except Exception as exc:  # noqa: BLE001 — isolate this clip's failure
-                    results[orig_i] = exc
+        for orig_i, codes in zip(prepared_idx, self._encode_shared_or_local(prepared)):
+            results[orig_i] = codes
         return results
 
-    def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]:
-        with torch.no_grad():
-            return self._processor.encode_audios_from_wav(prepared, sampling_rate=self._sr_target, n_vq=self._n_vq)
+    def prepare(self) -> None:
+        """Capture the CUDA graphs (or start the shared encoder) before the first request."""
+        role = shared_encoder_role() if self._shares_encoder else ""
+        if role == "host":
+            self._host().wait_until_ready()
+        elif role == "client":
+            if self._shared_client is None:
+                self._shared_client = SharedReferenceEncoderClient(self._shared_codes_dir)
+            self._shared_client.wait_until_ready()
+        elif role == "":
+            self._reference_graphs()
+
+    def _host(self) -> SharedReferenceEncoderHost:
+        if self._shared_host is None:
+            # Graph capture runs in the background; clips queue until it ends.
+            self._shared_host = SharedReferenceEncoderHost(
+                self._shared_codes_dir,
+                self._host_workers(),
+                stream_priority=reference_stream_priority(),
+                batch_window_ms=float(os.environ.get("VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS", "0")),
+            )
+        return self._shared_host
+
+    def _encode_shared_or_local(self, prepared: list[torch.Tensor]) -> list:
+        role = shared_encoder_role() if self._shares_encoder else ""
+        if role == "host":
+            return self._host().encode(prepared)
+        if role == "client":
+            if self._shared_client is None:
+                self._shared_client = SharedReferenceEncoderClient(self._shared_codes_dir)
+            try:
+                return self._shared_client.encode(prepared)
+            except SharedReferenceEncoderStartupError:
+                raise
+            except Exception:  # noqa: BLE001 — the host is unreachable: encode here
+                logger.warning("MOSS shared reference encoder unavailable; encoding locally", exc_info=True)
+                self._place_tokenizer_for_local_encode()
+        with self._local_stream_context():
+            return self._encode_local(prepared)
+
+    def _local_stream_context(self):
+        tokenizer = getattr(self._processor, "audio_tokenizer", None)
+        params = next(tokenizer.parameters(), None) if isinstance(tokenizer, torch.nn.Module) else None
+        if params is None or params.device.type != "cuda":
+            return contextlib.nullcontext()
+        if self._local_stream is None:
+            self._local_stream = torch.cuda.Stream(device=params.device, priority=reference_stream_priority())
+        return torch.cuda.stream(self._local_stream)
+
+    def _host_workers(self) -> list:
+        """Encode workers of the shared host, each with its own graphs and stream."""
+        count = max(1, int(os.environ.get(_SHARED_WORKERS_ENV, "4")))
+        if not self._graph_capable():
+            return [(self._encode_local, None)] * count
+        workers = []
+        for _ in range(count):
+            graphs = MossReferenceEncoderGraphs(
+                self._processor.audio_tokenizer,
+                n_vq=self._n_vq,
+                batch_sizes=FINE_BATCH_SIZES,
+                bucket_seconds=FINE_BUCKET_SECONDS,
+                max_batch_seconds=_SHARED_WORKER_BATCH_SECONDS,
+                compile_core=self._compile_graphs(),
+                batched_transfer=os.environ.get("VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", "0") == "1",
+                singleton_bucket_seconds=self._singleton_bucket_seconds(),
+            )
+            workers.append(((lambda wavs, g=graphs: self._encode_local(wavs, g)), graphs.capture))
+        return workers
+
+    def _encode_local(self, prepared: list[torch.Tensor], graphs: MossReferenceEncoderGraphs | None = None) -> list:
+        """One batched forward, retried per clip so one bad clip fails alone."""
+        try:
+            return list(self._encode_prepared(prepared, graphs))
+        except Exception:  # noqa: BLE001 — batch forward failed; retry item-by-item
+            logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared))
+        results: list = []
+        for wav in prepared:
+            try:
+                results.append(self._encode_prepared([wav], graphs)[0])
+            except Exception as exc:  # noqa: BLE001 — isolate this clip's failure
+                results.append(exc)
+        return results
+
+    def _place_tokenizer_for_local_encode(self) -> None:
+        """A client process keeps the tokenizer off the GPU until it must encode."""
+        tokenizer = getattr(self._processor, "audio_tokenizer", None)
+        device = getattr(self._processor, "_vllm_omni_ref_encoder_device", None)
+        if tokenizer is not None and device is not None and next(tokenizer.parameters()).device != device:
+            self._processor.audio_tokenizer = tokenizer.to(device)
+
+    def _encode_prepared(
+        self, prepared: list[torch.Tensor], graphs: MossReferenceEncoderGraphs | None = None
+    ) -> list[torch.Tensor]:
+        # ``prepared`` clips come from _encoder_input, at ``_clip_rate()``.
+        graphs = graphs or self._reference_graphs()
+        if graphs is None:
+            with torch.no_grad():
+                return self._processor.encode_audios_from_wav(
+                    prepared, sampling_rate=self._clip_rate(), n_vq=self._n_vq
+                )
+        codes = graphs.encode([self._finish_prepare(wav) for wav in prepared])
+        missing = [i for i, c in enumerate(codes) if c is None]
+        if missing:
+            with torch.no_grad():
+                eager = self._processor.encode_audios_from_wav(
+                    [prepared[i] for i in missing], sampling_rate=self._clip_rate(), n_vq=self._n_vq
+                )
+            for i, c in zip(missing, eager):
+                codes[i] = c
+        return codes
+
+    def _graph_capable(self) -> bool:
+        """True when the tokenizer can run graphs; installs windowed attention once."""
+        if not self._graphs_enabled:
+            return False
+        tokenizer = getattr(self._processor, "audio_tokenizer", None)
+        device = next(tokenizer.parameters()).device if tokenizer is not None else None
+        if device is None or device.type != "cuda" or not hasattr(tokenizer, "_codec_inference_autocast"):
+            self._graphs_enabled = False
+            return False
+        if not self._attention_installed and os.environ.get(_REF_ENCODE_ATTN_ENV, "flash") == "flash":
+            logger.info(
+                "MOSS reference encoder: %d attention layers use local-window flash attention",
+                install_windowed_attention(
+                    tokenizer,
+                    fa_version=int(os.environ.get("VLLM_OMNI_MOSS_REF_FA_VERSION", "2")),
+                    cache_rope=os.environ.get("VLLM_OMNI_MOSS_REF_CACHE_ROPE", "0") == "1",
+                    skip_padded_query_mask=os.environ.get("VLLM_OMNI_MOSS_REF_SKIP_PADDED_QUERY_MASK", "0") == "1",
+                ),
+            )
+        self._attention_installed = True
+        return True
+
+    @staticmethod
+    def _compile_graphs() -> bool:
+        return os.environ.get(_REF_ENCODE_COMPILE_ENV, "1") != "0"
+
+    @staticmethod
+    def _singleton_bucket_seconds() -> tuple[float, ...] | None:
+        step_ms = int(os.environ.get("VLLM_OMNI_MOSS_REF_SINGLETON_BUCKET_MS", "0"))
+        if step_ms == 0:
+            return None
+        if not 80 <= step_ms <= 1000:
+            raise ValueError("MOSS reference singleton bucket step must be 80..1000 ms, or 0 to disable")
+        return tuple(ms / 1000 for ms in range(3000, 10001, step_ms))
+
+    def _reference_graphs(self) -> MossReferenceEncoderGraphs | None:
+        """Capture on first use, on the encode worker thread that replays them."""
+        if self._graphs is None:
+            if not self._graph_capable():
+                return None
+            graphs = MossReferenceEncoderGraphs(
+                self._processor.audio_tokenizer,
+                n_vq=self._n_vq,
+                compile_core=self._compile_graphs(),
+                batched_transfer=os.environ.get("VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", "0") == "1",
+                singleton_bucket_seconds=self._singleton_bucket_seconds(),
+            )
+            graphs.capture()
+            self._graphs = graphs
+        return self._graphs
+
+    def _tokenizer_rate(self) -> int:
+        return int(getattr(getattr(self._processor, "model_config", None), "sampling_rate", self._sr_target))
+
+    def _encoder_input(self, wav: torch.Tensor) -> torch.Tensor:
+        """A prepped clip at the tokenizer rate, one or two channels.
+
+        This is the costly half of the processor's per-clip preparation; it runs
+        in the process that received the request, even with a shared encoder.
+        """
+        import torchaudio
+
+        if wav.ndim == 1:
+            wav = wav.unsqueeze(0)
+        if wav.shape[0] > 2:
+            wav = wav[:2]
+        if self._clip_rate() != self._sr_target:
+            # Channels resample independently, so a mono clip is resampled
+            # once before it is duplicated; the result is the same.
+            wav = torchaudio.functional.resample(wav, self._sr_target, self._tokenizer_rate())
+        return wav
+
+    def _clip_rate(self) -> int:
+        """Sample rate of the clips _encoder_input hands to the encoder."""
+        return self._sr_target if self._gpu_prep else self._tokenizer_rate()
+
+    def _finish_prepare(self, wav: torch.Tensor) -> torch.Tensor:
+        """The rest of the processor's preparation: two channels, normalized loudness."""
+        if self._gpu_prep:
+            device = next(self._processor.audio_tokenizer.parameters()).device
+            wav = wav.to(device)
+            if self._clip_rate() != self._tokenizer_rate():
+                wav = self._gpu_resampler(device)(wav)
+        if wav.shape[0] == 1:
+            wav = wav.repeat(2, 1)
+        return self._processor.loudness_normalize(wav)
+
+    def _gpu_resampler(self, device: torch.device) -> torch.nn.Module:
+        resampler = self._gpu_resamplers.get(device)
+        if resampler is None:
+            import torchaudio
+
+            # The same windowed-sinc filter as torchaudio.functional.resample's defaults.
+            resampler = torchaudio.transforms.Resample(self._sr_target, self._tokenizer_rate()).to(device)
+            self._gpu_resamplers[device] = resampler
+        return resampler
+
+    def _processor_prepare(self, wav: torch.Tensor) -> torch.Tensor:
+        """Per-clip preparation of the processor's ``encode_audios_from_wav``."""
+        return self._finish_prepare(self._encoder_input(wav))
 
     async def aclose(self) -> None:
         """Release the batcher drainer + any in-flight encodes (tests/shutdown)."""
@@ -385,7 +760,10 @@ class MossReferenceEncoder:
 class MossRealtimeReferenceEncoder(MossReferenceEncoder):
     """Use the realtime codec with the shared reference cache and batcher."""
 
-    def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]:
+    def _encoder_input(self, wav: torch.Tensor) -> torch.Tensor:
+        return wav  # the realtime codec takes the clip at the working rate
+
+    def _encode_prepared(self, prepared: list[torch.Tensor], graphs: Any = None) -> list[torch.Tensor]:
         with torch.no_grad():
             encoded = self._processor.batch_encode([wav.squeeze(0) for wav in prepared], num_quantizers=self._n_vq)
         return [

vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py

@@ -0,0 +1,399 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CUDA graphs for batched MOSS reference-audio encoding.
+
+Reference clips are encoded by the upstream MOSS audio tokenizer that the
+processor loads. Its eager forward is dominated by kernel launches for the
+small batches that serving produces. This runner replays graphs of the same
+encoder and quantizer modules, captured per ``(batch, length)`` bucket. For a
+given padded input the codes are bit-identical to the tokenizer's own
+``_encode_frame``. Only the per-clip crop, which needs host values, runs
+outside the graph.
+
+The encoder is causal, so padding a clip to its length bucket leaves its valid
+frames unchanged up to kernel selection for the padded shape, the same effect
+that batching already has when the tokenizer pads a batch to its longest clip.
+"""
+
+from __future__ import annotations
+
+import math
+import types
+from dataclasses import dataclass
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+DEFAULT_BATCH_SIZES = (1, 2, 4, 8, 16, 32)
+DEFAULT_BUCKET_SECONDS = (4.0, 6.0, 8.0, 10.0, 12.0, 16.0, 20.0)
+# Fine steps keep padding small: a padded row or padded time is encoded in
+# full. Worth the extra graphs where several encoders share one GPU.
+FINE_BATCH_SIZES = (1, 2, 3, 4, 5, 6, 8)
+FINE_BUCKET_SECONDS = (3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0, 16.0, 20.0)
+# Graphs are captured only while batch * bucket stays within this much audio,
+# which bounds the shared graph memory pool by its largest batch.
+DEFAULT_MAX_BATCH_SECONDS = 256.0
+
+
+@torch.compile(dynamic=False)
+def _reference_rope_factors(seqlen: int, head_dim: int, max_period: float, device: torch.device):
+    # Use the same compiled FP32 arithmetic as the encoder. Eager libdevice
+    # sin/cos and compiled sin/cos can round differently near a VQ boundary.
+    ds = torch.arange(head_dim // 2, device=device, dtype=torch.float32)
+    freqs = torch.exp(ds * (-math.log(max_period) * 2 / head_dim))
+    angles = torch.arange(seqlen, device=device, dtype=torch.float32)[:, None] * freqs[None, :]
+    return torch.cos(angles), torch.sin(angles)
+
+
+def _cached_reference_rope(self, q: torch.Tensor, k: torch.Tensor, cache: dict):
+    if self.rope is None:
+        return q, k
+    batch, heads, seqlen, dim = q.shape
+    key = (dim, self.rope.max_period, q.device)
+    if key not in cache or cache[key][0].shape[0] < seqlen:
+        cache[key] = _reference_rope_factors(seqlen, dim, self.rope.max_period, q.device)
+    cos, sin = cache[key]
+    cos, sin = cos[:seqlen], sin[:seqlen]
+    qr, qi = q.view(batch, heads, seqlen, dim // 2, 2).unbind(-1)
+    kr, ki = k.view(batch, heads, seqlen, dim // 2, 2).unbind(-1)
+    qr, qi, kr, ki = qr.float(), qi.float(), kr.float(), ki.float()
+    # Preserve the original arithmetic, casts, and B/H/T/D output layout.
+    qo = torch.stack(((qr * cos - qi * sin).to(q.dtype), (qr * sin + qi * cos).to(q.dtype)), dim=-1)
+    ko = torch.stack(((kr * cos - ki * sin).to(k.dtype), (kr * sin + ki * cos).to(k.dtype)), dim=-1)
+    return qo.view(batch, heads, seqlen, dim), ko.view(batch, heads, seqlen, dim)
+
+
+def _windowed_attention(
+    self,
+    x: torch.Tensor,
+    input_lengths: torch.Tensor,
+    *,
+    sdpa,
+    fa_version: int = 2,
+    rope_cache: dict | None = None,
+    skip_padded_query_mask: bool = False,
+) -> torch.Tensor:
+    """The tokenizer's non-streaming attention with a local-window flash kernel.
+
+    Same projection, RoPE and padded-row zeroing as the tokenizer's SDPA path.
+    That path masks each query to ``context`` keys (distance 0 .. context-1)
+    inside a dense block; the flash kernel visits only that window.
+    """
+    batch, seqlen, _ = x.shape
+    q, k, v = self._project_qkv(x)
+    q, k = self._apply_dense_rope(q, k) if rope_cache is None else _cached_reference_rope(self, q, k, rope_cache)
+    if q.dtype not in (torch.bfloat16, torch.float16):
+        return sdpa(x, input_lengths)
+    from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+    heads, dim = q.shape[1], q.shape[-1]
+
+    def packed(t: torch.Tensor) -> torch.Tensor:
+        return t.transpose(1, 2).reshape(batch * seqlen, heads, dim)
+
+    starts = torch.arange(0, (batch + 1) * seqlen, seqlen, device=x.device, dtype=torch.int32)
+    window = [self.context - 1, 0] if self.causal and self.context is not None else [-1, -1]
+    out = flash_attn_varlen_func(
+        packed(q),
+        packed(k),
+        packed(v),
+        max_seqlen_q=seqlen,
+        cu_seqlens_q=starts,
+        max_seqlen_k=seqlen,
+        cu_seqlens_k=starts,
+        causal=self.causal,
+        window_size=window,
+        fa_version=fa_version,
+    ).view(batch, seqlen, heads, dim)
+    if not skip_padded_query_mask:
+        valid = (torch.arange(seqlen, device=x.device).view(1, seqlen) < input_lengths.view(-1, 1)).view(
+            batch, seqlen, 1, 1
+        )
+        out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype))
+    return out.reshape(batch, seqlen, self.embed_dim)
+
+
+def _validate_unmasked_reference_encoder(tokenizer: torch.nn.Module) -> None:
+    """Restrict the optimization to the audited HF encoder's valid-prefix contract.
+
+    Causal attention and pointwise transformer operations cannot read future
+    padded rows. Its patch-downsample floors lengths, so every valid output
+    group consists entirely of valid input rows. The quantizer is pointwise
+    and the caller crops codes by those lengths. Padded *outputs* may differ.
+    This argument does not apply to the streaming decoder or to upsampling.
+    """
+    transformers = 0
+    for module in tokenizer.encoder:
+        name = type(module).__name__
+        if name == "MossAudioTokenizerPatchedPretransform" and module.is_downsample:
+            continue
+        if name == "MossAudioTokenizerProjectedTransformer":
+            layers = module.transformer.layers
+            if layers and all(layer.self_attn.causal for layer in layers):
+                transformers += 1
+                continue
+        raise ValueError("Unmasked reference queries require the causal MOSS patch-downsample encoder")
+    if not transformers:
+        raise ValueError("Unmasked reference queries require at least one causal MOSS transformer")
+    if type(tokenizer.quantizer).__name__ not in {
+        "MossAudioTokenizerResidualVQ",
+        "MossAudioTokenizerResidualLFQ",
+    }:
+        raise ValueError("Unmasked reference queries require the pointwise MOSS quantizer")
+    for module in tokenizer.quantizer.modules():
+        if isinstance(module, torch.nn.Conv1d) and module.kernel_size != (1,):
+            raise ValueError("Unmasked reference queries do not support temporal quantizer convolutions")
+
+
+def install_windowed_attention(
+    tokenizer: torch.nn.Module,
+    *,
+    fa_version: int = 2,
+    cache_rope: bool = False,
+    skip_padded_query_mask: bool = False,
+) -> int:
+    """Route the encoder's non-streaming attention through a local-window flash kernel."""
+    if skip_padded_query_mask:
+        _validate_unmasked_reference_encoder(tokenizer)
+        logger.info("MOSS reference encoder: skipping padded query masks (valid-prefix outputs only)")
+    count = 0
+    rope_cache = {} if cache_rope else None
+    if cache_rope:
+        tokenizer._moss_reference_rope_cache = rope_cache
+        tokenizer._moss_reference_rope_prepared_samples = 0
+    for module in tokenizer.encoder.modules():
+        sdpa = getattr(module, "_forward_non_streaming_sdpa", None)
+        if sdpa is None or not hasattr(module, "_project_qkv") or not getattr(module, "causal", False):
+            continue
+
+        def attention(self, x, input_lengths, _sdpa=sdpa):
+            return _windowed_attention(
+                self,
+                x,
+                input_lengths,
+                sdpa=_sdpa,
+                fa_version=fa_version,
+                rope_cache=rope_cache,
+                skip_padded_query_mask=skip_padded_query_mask,
+            )
+
+        module._forward_non_streaming_sdpa = types.MethodType(attention, module)
+        count += 1
+    return count
+
+
+@dataclass
+class _Graph:
+    graph: torch.cuda.CUDAGraph
+    audio: torch.Tensor  # (B, C, T) float32 input
+    lengths: torch.Tensor  # (B,) int64 input, in samples per channel
+    codes: torch.Tensor  # (n_vq, B, frames) output
+    code_lengths: torch.Tensor  # (B,) output
+
+
+class MossReferenceEncoderGraphs:
+    """Encode prepared reference waveforms with captured tokenizer graphs."""
+
+    def __init__(
+        self,
+        tokenizer: torch.nn.Module,
+        *,
+        n_vq: int,
+        batch_sizes: tuple[int, ...] = DEFAULT_BATCH_SIZES,
+        bucket_seconds: tuple[float, ...] = DEFAULT_BUCKET_SECONDS,
+        max_batch_seconds: float = DEFAULT_MAX_BATCH_SECONDS,
+        compile_core: bool = False,
+        batched_transfer: bool = False,
+        singleton_bucket_seconds: tuple[float, ...] | None = None,
+    ) -> None:
+        self._tokenizer = tokenizer
+        self._compile_core = compile_core
+        self._batched_transfer = batched_transfer
+        self._n_vq = int(n_vq)
+        self._device = next(tokenizer.parameters()).device
+        self._channels = int(tokenizer.number_channels)
+        frame = int(tokenizer.downsample_rate)
+        sampling_rate = int(tokenizer.sampling_rate)
+        # Whole frames, so the tokenizer's own frame padding is a no-op.
+        self._buckets = sorted({int(math.ceil(s * sampling_rate / frame)) * frame for s in bucket_seconds})
+        self._singleton_buckets = sorted(
+            set(self._buckets)
+            | {int(math.ceil(s * sampling_rate / frame)) * frame for s in (singleton_bucket_seconds or ())}
+        )
+        self._batch_sizes = sorted({int(b) for b in batch_sizes if int(b) > 0})
+        self._max_batch_samples = float(max_batch_seconds) * sampling_rate
+        self._graphs: dict[tuple[int, int], _Graph] = {}
+        self._pool = None
+        self._transfer_buffers = None
+
+    @property
+    def max_batch(self) -> int:
+        return self._batch_sizes[-1]
+
+    @property
+    def captured(self) -> list[tuple[int, int]]:
+        return sorted(self._graphs)
+
+    def _core(self, audio: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # Mirrors the tokenizer's _encode_frame without its host-side crop.
+        tokenizer = self._tokenizer
+        values, value_lengths = tokenizer._flatten_channels_for_codec(audio, lengths)
+        with tokenizer._codec_inference_autocast():
+            hidden, hidden_lengths = values, value_lengths
+            for module in tokenizer.encoder:
+                hidden, hidden_lengths = module(hidden, hidden_lengths)
+        _, codes, code_lengths = tokenizer.quantizer(hidden.float(), hidden_lengths, self._n_vq)
+        return codes, code_lengths
+
+    @torch.no_grad()
+    def capture(self) -> None:
+        if self._device.type != "cuda" or self._graphs:
+            return
+        self._pool = torch.cuda.graph_pool_handle()
+        stream = torch.cuda.Stream(device=self._device)
+        prepared = getattr(self._tokenizer, "_moss_reference_rope_prepared_samples", None)
+        max_samples = max(self._singleton_buckets)
+        if prepared is not None and prepared < max_samples:
+            # Prime immutable FP32 factors at the largest length before any
+            # capture or compilation. Smaller lengths use a prefix of the
+            # same table; the dictionary stays fixed throughout capture and
+            # workers share tables without introducing dynamic-shape guards
+            # for a separate dictionary key at each sequence length.
+            audio = torch.zeros(1, self._channels, max_samples, device=self._device)
+            lengths = torch.full((1,), max_samples, dtype=torch.long, device=self._device)
+            self._core(audio, lengths)
+            self._tokenizer._moss_reference_rope_prepared_samples = max_samples
+        core = self._core
+        if self._compile_core:
+            # Fuses the encoder's many small elementwise kernels (casts, RoPE,
+            # residuals). Only capture runs the compiled code; replays do not.
+            core = torch.compile(self._core, dynamic=True)
+        # Each encoder stage is its own module shape: allow a compile per stage.
+        with torch._dynamo.config.patch(recompile_limit=64):
+            self._capture_all(core, stream)
+        if self._batched_transfer and self._graphs:
+            # One reusable slab per worker, rather than one per graph bucket.
+            entries = list(self._graphs.values())
+            self._transfer_buffers = (
+                torch.empty(max(e.audio.numel() for e in entries), dtype=entries[0].audio.dtype, pin_memory=True),
+                torch.empty(max(e.lengths.numel() for e in entries), dtype=entries[0].lengths.dtype, pin_memory=True),
+                torch.empty(max(e.codes.numel() for e in entries), dtype=entries[0].codes.dtype, pin_memory=True),
+                torch.empty(
+                    max(e.code_lengths.numel() for e in entries), dtype=entries[0].code_lengths.dtype, pin_memory=True
+                ),
+                torch.cuda.Event(),
+            )
+        logger.info(
+            "MOSS reference encoder CUDA graphs: batch=%s buckets(s)=%s compiled=%s captured=%d",
+            self._batch_sizes,
+            [round(s / int(self._tokenizer.sampling_rate), 2) for s in self._buckets],
+            self._compile_core,
+            len(self._graphs),
+        )
+        if self._singleton_buckets != self._buckets:
+            logger.info(
+                "MOSS reference encoder singleton buckets(s): %s",
+                [round(s / int(self._tokenizer.sampling_rate), 2) for s in self._singleton_buckets],
+            )
+
+    def _capture_all(self, core, stream: torch.cuda.Stream) -> None:
+        for batch in reversed(self._batch_sizes):
+            buckets = self._singleton_buckets if batch == 1 else self._buckets
+            for samples in reversed(buckets):
+                if batch > 1 and batch * samples > self._max_batch_samples:
+                    continue
+                audio = torch.zeros(batch, self._channels, samples, device=self._device)
+                lengths = torch.full((batch,), samples, dtype=torch.long, device=self._device)
+                try:
+                    stream.wait_stream(torch.cuda.current_stream(self._device))
+                    with torch.cuda.stream(stream):
+                        # Lazy library setup must happen before capture; one run does it.
+                        core(audio, lengths)
+                    torch.cuda.current_stream(self._device).wait_stream(stream)
+                    graph = torch.cuda.CUDAGraph()
+                    with torch.cuda.graph(graph, pool=self._pool, stream=stream):
+                        codes, code_lengths = core(audio, lengths)
+                except Exception:
+                    logger.warning(
+                        "MOSS reference encoder graph (B=%d, samples=%d) not captured; using eager",
+                        batch,
+                        samples,
+                        exc_info=True,
+                    )
+                    continue
+                self._graphs[(batch, samples)] = _Graph(graph, audio, lengths, codes, code_lengths)
+
+    @torch.no_grad()
+    def encode(self, wavs: list[torch.Tensor]) -> list[torch.Tensor | None]:
+        """Codes ``(frames, n_vq)`` int64 on CPU per clip; None where no graph fits.
+
+        ``wavs`` are prepared the way the processor prepares them: ``(C, T)``
+        float32 after channel handling, resampling and loudness normalization.
+        Clips are grouped longest first, so each replay pads to the bucket of
+        similar-length clips.
+        """
+        out: list[torch.Tensor | None] = [None] * len(wavs)
+        order = sorted(range(len(wavs)), key=lambda i: -int(wavs[i].shape[-1]))
+        pos = 0
+        while pos < len(order):
+            longest = int(wavs[order[pos]].shape[-1])
+            # Fine singleton buckets reduce padding without fragmenting an
+            # already assembled batch into separate graph replays.
+            buckets = self._singleton_buckets if len(order) - pos == 1 else self._buckets
+            samples = next((s for s in buckets if s >= longest and (1, s) in self._graphs), None)
+            if samples is None:
+                pos += 1  # longer than every bucket: left to the caller
+                continue
+            widest = max(b for b in self._batch_sizes if (b, samples) in self._graphs)
+            rows = order[pos : pos + widest]
+            batch = next(b for b in self._batch_sizes if b >= len(rows) and (b, samples) in self._graphs)
+            for row, codes in zip(rows, self._replay(self._graphs[(batch, samples)], [wavs[i] for i in rows])):
+                out[row] = codes
+            pos += len(rows)
+        return out
+
+    def _replay(self, entry: _Graph, wavs: list[torch.Tensor]) -> list[torch.Tensor]:
+        if self._transfer_buffers is not None and all(w.device.type == "cpu" for w in wavs):
+            return self._replay_batched_transfer(entry, wavs)
+        count, batch, samples = len(wavs), entry.audio.shape[0], entry.audio.shape[-1]
+        # Unused rows encode full-length silence: a fully masked row would be
+        # all-NaN inside attention, and rows never mix.
+        entry.audio.zero_()
+        lengths = [int(w.shape[-1]) for w in wavs] + [samples] * (batch - count)
+        entry.lengths.copy_(torch.tensor(lengths, dtype=torch.long))
+        for row, wav in enumerate(wavs):
+            entry.audio[row, :, : wav.shape[-1]].copy_(wav)
+        entry.graph.replay()
+        codes = entry.codes[:, :count].cpu()
+        code_lengths = entry.code_lengths[:count].tolist()
+        return [codes[:, row, : int(code_lengths[row])].transpose(0, 1).contiguous().long() for row in range(count)]
+
+    def _replay_batched_transfer(self, entry: _Graph, wavs: list[torch.Tensor]) -> list[torch.Tensor]:
+        """Stage one batch on pinned host memory, then wait once for its outputs."""
+        audio_slab, lengths_slab, codes_slab, code_lengths_slab, copied = self._transfer_buffers
+        count, batch, samples = len(wavs), entry.audio.shape[0], entry.audio.shape[-1]
+        audio = audio_slab[: entry.audio.numel()].view_as(entry.audio)
+        lengths = lengths_slab[:batch]
+        audio.zero_()
+        lengths.fill_(samples)
+        for row, wav in enumerate(wavs):
+            audio[row, :, : wav.shape[-1]].copy_(wav)
+            lengths[row] = wav.shape[-1]
+        entry.audio.copy_(audio, non_blocking=True)
+        entry.lengths.copy_(lengths, non_blocking=True)
+        entry.graph.replay()
+        result = entry.codes[:, :count]
+        codes = codes_slab[: result.numel()].view_as(result)
+        code_lengths = code_lengths_slab[:count]
+        codes.copy_(result, non_blocking=True)
+        code_lengths.copy_(entry.code_lengths[:count], non_blocking=True)
+        copied.record(torch.cuda.current_stream(self._device))
+        copied.synchronize()
+        # Returned codes must outlive this worker's next replay, including
+        # one-frame clips whose transpose is already contiguous.
+        return [
+            codes[:, row, : int(code_lengths[row])].transpose(0, 1).to(dtype=torch.long, copy=True).contiguous()
+            for row in range(count)
+        ]

vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py

@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""One reference-audio encoder shared by the API processes of a server.
+
+Each API process otherwise loads its own copy of the audio tokenizer on the
+GPU and encodes only the clips it receives. When the processes share a
+directory (``VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR``), the first one to take the lock in it
+hosts the encoder and serves the others over a Unix socket in that directory.
+Requests from all processes are merged into batches, and only the host keeps
+the tokenizer and its CUDA graphs on the GPU.
+
+Messages carry raw arrays and a JSON header; nothing is unpickled. A client
+that cannot reach the host encodes locally instead.
+"""
+
+from __future__ import annotations
+
+import collections
+import fcntl
+import json
+import os
+import threading
+import time
+from collections.abc import Callable
+from contextlib import nullcontext
+from multiprocessing.connection import Client, Connection, Listener
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_LOCK_NAME = "ref-encoder.lock"
+_SOCKET_NAME = "ref-encoder.sock"
+# Clips one worker merges into an encode, across all API processes. Several
+# workers encode concurrently on their own streams: one serial queue of large
+# batches makes a burst of new references wait behind each other. A free
+# worker takes what is queued without waiting for more: API processes already
+# gather their clips before sending them.
+_MAX_BATCH = 8
+# A client gives up on the host after this long and encodes locally.
+_CONNECT_TIMEOUT_S = 120.0
+_REQUEST_TIMEOUT_S = 120.0
+# First-time compilation can outlast a normal encode request. Readiness has a
+# separate deadline and must finish before the API starts accepting requests.
+_STARTUP_TIMEOUT_S = 1200.0
+
+EncodeFn = Callable[[list[torch.Tensor]], list]
+
+# Open for the life of the process: closing it would release the host lock.
+_host_lock_handle = None
+
+
+class SharedReferenceEncoderStartupError(RuntimeError):
+    """The shared encoder cannot safely accept requests yet."""
+
+
+def elect_host(shared_dir: str) -> bool:
+    """True in the one process that holds the encoder lock for ``shared_dir``.
+
+    The lock is held for the life of the process, so a restarted server
+    elects a new host even if the old socket file is still present.
+    """
+    global _host_lock_handle
+    path = os.path.join(shared_dir, _LOCK_NAME)
+    handle = open(path, "a+")  # noqa: SIM115 - kept open to hold the lock
+    try:
+        fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
+    except OSError:
+        handle.close()
+        return False
+    _host_lock_handle = handle
+    return True
+
+
+def _send(conn: Connection, header: dict, arrays: list[np.ndarray]) -> None:
+    conn.send_bytes(json.dumps(header).encode())
+    for array in arrays:
+        conn.send_bytes(np.ascontiguousarray(array).tobytes())
+
+
+def _recv(conn: Connection) -> tuple[dict, list[bytes]]:
+    header = json.loads(conn.recv_bytes().decode())
+    return header, [conn.recv_bytes() for _ in range(int(header["count"]))]
+
+
+class _Job:
+    __slots__ = ("wav", "done", "result", "queued_at")
+
+    def __init__(self, wav: torch.Tensor) -> None:
+        self.wav = wav
+        self.done = threading.Event()
+        self.result: torch.Tensor | BaseException | None = None
+        self.queued_at = time.monotonic()
+
+
+class SharedReferenceEncoderHost:
+    """Serves reference encodes for every API process sharing ``shared_dir``.
+
+    ``workers`` pairs each worker's encode function with an optional warmup
+    (for example, capturing that worker's CUDA graphs).
+    """
+
+    def __init__(
+        self,
+        shared_dir: str,
+        workers: list[tuple[EncodeFn, Callable[[], object] | None]],
+        *,
+        stream_priority: int = 0,
+        batch_window_ms: float = 0.0,
+    ) -> None:
+        self._stream_priority = stream_priority
+        self._batch_window_s = max(0.0, float(batch_window_ms)) / 1000
+        self._pending: collections.deque[_Job] = collections.deque()
+        # Includes jobs already taken by a worker, so shutdown can release
+        # every caller without waiting for an in-flight GPU operation.
+        self._outstanding: set[_Job] = set()
+        self._ready = threading.Condition()
+        self._path = os.path.join(shared_dir, _SOCKET_NAME)
+        if os.path.exists(self._path):
+            os.unlink(self._path)  # left by a previous host; we hold the lock
+        self._listener = Listener(self._path, family="AF_UNIX")
+        os.chmod(self._path, 0o600)
+        self._closed = False
+        self._started = threading.Event()
+        self._startup_error: Exception | None = None
+        threading.Thread(
+            target=self._start_workers, args=(workers,), name="moss-ref-encoder-start", daemon=True
+        ).start()
+        threading.Thread(target=self._accept_loop, name="moss-ref-encoder-accept", daemon=True).start()
+        logger.info("MOSS shared reference encoder serving at %s with %d workers", self._path, len(workers))
+
+    def encode(self, wavs: list[torch.Tensor]) -> list:
+        """Encode on behalf of the host process itself."""
+        self.wait_until_ready()
+        jobs = [_Job(wav) for wav in wavs]
+        with self._ready:
+            if self._closed:
+                raise RuntimeError("shared reference encoder is closed")
+            self._outstanding.update(jobs)
+            self._pending.extend(jobs)
+            self._ready.notify(len(jobs))
+        for job in jobs:
+            job.done.wait()
+        return [job.result for job in jobs]
+
+    def close(self) -> None:
+        with self._ready:
+            self._closed = True
+            self._started.set()
+            for job in self._outstanding:
+                job.result = RuntimeError("shared reference encoder closed before encoding completed")
+                job.done.set()
+            self._outstanding.clear()
+            self._pending.clear()
+            self._ready.notify_all()
+        try:
+            self._listener.close()
+        except OSError:
+            pass
+
+    def wait_until_ready(self) -> None:
+        if not self._started.wait(_STARTUP_TIMEOUT_S):
+            raise SharedReferenceEncoderStartupError("shared reference encoder warmup timed out")
+        if self._closed or self._startup_error is not None:
+            raise SharedReferenceEncoderStartupError("shared reference encoder startup failed") from self._startup_error
+
+    def _accept_loop(self) -> None:
+        while not self._closed:
+            try:
+                conn = self._listener.accept()
+            except OSError:
+                return
+            threading.Thread(target=self._serve, args=(conn,), name="moss-ref-encoder-conn", daemon=True).start()
+
+    def _serve(self, conn: Connection) -> None:
+        with conn:
+            while True:
+                try:
+                    header, payloads = _recv(conn)
+                except (EOFError, OSError):
+                    return
+                if header.get("op") == "ready":
+                    try:
+                        self.wait_until_ready()
+                        error = None
+                    except SharedReferenceEncoderStartupError as exc:
+                        error = str(exc)
+                    try:
+                        _send(conn, {"count": 0, "ready": error is None, "error": error}, [])
+                    except OSError:
+                        return
+                    continue
+                wavs = [
+                    torch.from_numpy(np.frombuffer(data, dtype=np.float32).reshape(shape).copy())
+                    for data, shape in zip(payloads, header["shapes"])
+                ]
+                results = self.encode(wavs)
+                errors = [None if isinstance(r, torch.Tensor) else repr(r) for r in results]
+                codes = [r.to(torch.int32).numpy() for r in results if isinstance(r, torch.Tensor)]
+                reply = {"count": len(codes), "shapes": [list(c.shape) for c in codes], "errors": errors}
+                try:
+                    _send(conn, reply, codes)
+                except OSError:
+                    return
+
+    def _start_workers(self, workers: list[tuple[EncodeFn, Callable[[], object] | None]]) -> None:
+        # Every warmup (graph capture) finishes before any worker encodes: CUDA
+        # work in another thread would invalidate a capture in progress.
+        try:
+            for _, warmup in workers:
+                if warmup is not None:
+                    warmup()
+            if self._closed:
+                return
+            for index, (encode_fn, _) in enumerate(workers):
+                threading.Thread(
+                    target=self._encode_loop, args=(encode_fn,), name=f"moss-ref-encoder-{index}", daemon=True
+                ).start()
+            logger.info("MOSS shared reference encoder ready: %d workers", len(workers))
+        except Exception as exc:  # noqa: BLE001 - wake readiness waiters with the failure
+            self._startup_error = exc
+            logger.exception("MOSS shared reference encoder startup failed")
+        finally:
+            self._started.set()
+
+    def _encode_loop(self, encode_fn: EncodeFn) -> None:
+        stream = torch.cuda.Stream(priority=self._stream_priority) if torch.cuda.is_available() else None
+        with torch.cuda.stream(stream) if stream is not None else nullcontext():
+            self._serve_jobs(encode_fn)
+
+    def _serve_jobs(self, encode_fn: EncodeFn) -> None:
+        while True:
+            with self._ready:
+                while True:
+                    if self._closed:
+                        return
+                    if not self._pending:
+                        self._ready.wait()
+                        continue
+                    remaining = self._pending[0].queued_at + self._batch_window_s - time.monotonic()
+                    if remaining > 0 and len(self._pending) < _MAX_BATCH:
+                        # Keep clips in the common queue while collecting;
+                        # workers must not each reserve a singleton first.
+                        self._ready.wait(remaining)
+                        continue
+                    break
+                jobs = [self._pending.popleft() for _ in range(min(_MAX_BATCH, len(self._pending)))]
+            try:
+                results = encode_fn([job.wav for job in jobs])
+                if len(results) != len(jobs):
+                    raise RuntimeError(
+                        f"shared reference encoder returned {len(results)} results for {len(jobs)} clips"
+                    )
+            except Exception as error:  # noqa: BLE001 - every waiter gets the failure
+                results = [error] * len(jobs)
+            with self._ready:
+                for job, result in zip(jobs, results, strict=True):
+                    # close() may already have completed this caller with an
+                    # error. A late GPU result must not replace that outcome.
+                    if not job.done.is_set():
+                        job.result = result
+                        job.done.set()
+                    self._outstanding.discard(job)
+
+
+class SharedReferenceEncoderClient:
+    """Sends reference clips to the host process of ``shared_dir``."""
+
+    def __init__(self, shared_dir: str) -> None:
+        self._path = os.path.join(shared_dir, _SOCKET_NAME)
+        # One connection per request in flight; idle ones are reused.
+        self._idle: list[Connection] = []
+        self._lock = threading.Lock()
+        self._connected = False
+
+    def _connect(self) -> Connection:
+        # The host may still be starting until the first connection succeeds;
+        # later reconnects are quick.
+        deadline = time.monotonic() + (1.0 if self._connected else _CONNECT_TIMEOUT_S)
+        while True:
+            try:
+                conn = Client(self._path, family="AF_UNIX")
+                self._connected = True
+                break
+            except (FileNotFoundError, ConnectionRefusedError):
+                if time.monotonic() >= deadline:
+                    raise
+                time.sleep(0.2)
+        try:
+            _send(conn, {"op": "ready", "count": 0}, [])
+            if not conn.poll(_STARTUP_TIMEOUT_S):
+                raise SharedReferenceEncoderStartupError("shared reference encoder readiness timed out")
+            header, _ = _recv(conn)
+            if not header.get("ready"):
+                raise SharedReferenceEncoderStartupError(header.get("error", "shared reference encoder not ready"))
+            return conn
+        except BaseException:
+            conn.close()
+            raise
+
+    def wait_until_ready(self) -> None:
+        """Startup handshake, without encoding or populating the reference cache."""
+        try:
+            conn = self._connect()
+        except Exception as exc:
+            raise SharedReferenceEncoderStartupError("shared reference encoder is not ready") from exc
+        with self._lock:
+            self._idle.append(conn)
+
+    def encode(self, wavs: list[torch.Tensor]) -> list:
+        """Codes (or an exception) per clip; raises when the host is unreachable."""
+        with self._lock:
+            conn = self._idle.pop() if self._idle else None
+        try:
+            if conn is None:
+                conn = self._connect()
+            arrays = [w.detach().to(torch.float32).cpu().numpy() for w in wavs]
+            _send(conn, {"count": len(arrays), "shapes": [list(a.shape) for a in arrays]}, arrays)
+            if not conn.poll(_REQUEST_TIMEOUT_S):
+                raise TimeoutError(f"no reply from the shared reference encoder in {_REQUEST_TIMEOUT_S}s")
+            header, payloads = _recv(conn)
+        except BaseException:
+            if conn is not None:
+                conn.close()
+            raise
+        with self._lock:
+            self._idle.append(conn)
+        codes = iter(
+            torch.from_numpy(np.frombuffer(data, dtype=np.int32).reshape(shape).copy()).long()
+            for data, shape in zip(payloads, header["shapes"])
+        )
+        return [next(codes) if error is None else RuntimeError(error) for error in header["errors"]]

vllm_omni/model_executor/models/moss_tts/slot_attention.py

@@ -0,0 +1,475 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Direct-slot ring attention for MOSS streaming codec.
+
+Inputs q/k/v are (execution_batch, heads, padded_T, head_dim). cache is
+contiguous (2, state_capacity, heads, ring_capacity, head_dim). end_offset
+contains absolute *pre-call* positions per state slot. slot_ids and valid_lengths
+contain one entry per execution row. valid_lengths=0 is a scratch/padding row.
+
+This op writes only valid tokens, attends to the resulting ring with the same
+position semantics as RingKVCache.complete, then advances only active offsets.
+Slots must be unique and in bounds; 0 <= valid_lengths <= padded_T. Chunks
+longer than the ring retain only their final ring_capacity tokens.
+These device-value preconditions are caller-owned to avoid host synchronization.
+The op deliberately preserves chunk-complete ring overwrite semantics; it does
+not recover entries overwritten by later tokens of the same input chunk.
+"""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _write(
+    k_ptr,
+    v_ptr,
+    cache_ptr,
+    end_ptr,
+    slots_ptr,
+    lengths_ptr,
+    ks0: tl.constexpr,
+    ks1: tl.constexpr,
+    ks2: tl.constexpr,
+    ks3: tl.constexpr,
+    vs0: tl.constexpr,
+    vs1: tl.constexpr,
+    vs2: tl.constexpr,
+    vs3: tl.constexpr,
+    end_stride: tl.constexpr,
+    slot_stride: tl.constexpr,
+    length_stride: tl.constexpr,
+    state_capacity: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_frames: tl.constexpr,
+    ring_capacity: tl.constexpr,
+    head_dim: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    row = tl.program_id(0)
+    slot = tl.load(slots_ptr + row * slot_stride)
+    length = tl.load(lengths_ptr + row * length_stride)
+    start = tl.load(end_ptr + slot * end_stride)
+    x = tl.program_id(1) * block_size + tl.arange(0, block_size)
+    dim = x % head_dim
+    token = (x // head_dim) % num_frames
+    head = x // (num_frames * head_dim)
+    # For chunks longer than the ring, write only the final C tokens. This
+    # avoids duplicate physical addresses (and write races) in one launch.
+    skip = tl.maximum(length - ring_capacity, 0)
+    source_token = token + skip
+    mask = (head < num_heads) & (source_token < length)
+    pos = (start + source_token) % ring_capacity
+    dst = ((slot * num_heads + head) * ring_capacity + pos) * head_dim + dim
+    kval = tl.load(k_ptr + row * ks0 + head * ks1 + source_token * ks2 + dim * ks3, mask, 0)
+    vval = tl.load(v_ptr + row * vs0 + head * vs1 + source_token * vs2 + dim * vs3, mask, 0)
+    tl.store(cache_ptr + dst, kval, mask)
+    tl.store(cache_ptr + state_capacity * num_heads * ring_capacity * head_dim + dst, vval, mask)
+
+
+@triton.jit
+def _attend(
+    q_ptr,
+    cache_ptr,
+    end_ptr,
+    slots_ptr,
+    lengths_ptr,
+    out_ptr,
+    qs0: tl.constexpr,
+    qs1: tl.constexpr,
+    qs2: tl.constexpr,
+    qs3: tl.constexpr,
+    end_stride: tl.constexpr,
+    slot_stride: tl.constexpr,
+    length_stride: tl.constexpr,
+    state_capacity: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_frames: tl.constexpr,
+    ring_capacity: tl.constexpr,
+    head_dim: tl.constexpr,
+    context_size: tl.constexpr,
+    block_m: tl.constexpr,
+    block_n: tl.constexpr,
+    skip_empty_tiles: tl.constexpr = False,
+):
+    rows = tl.program_id(0) * block_m + tl.arange(0, block_m)
+    bh = tl.program_id(1)
+    batch, head = bh // num_heads, bh % num_heads
+    slot = tl.load(slots_ptr + batch * slot_stride)
+    length = tl.load(lengths_ptr + batch * length_stride)
+    start = tl.load(end_ptr + slot * end_stride)
+    end = start + length
+    dims = tl.arange(0, head_dim)
+    q = tl.load(
+        q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, rows[:, None] < num_frames, 0
+    )
+    acc = tl.full((block_m, head_dim), 0.0, tl.float32)
+    maximum = tl.full((block_m,), -float("inf"), tl.float32)
+    denominator = tl.full((block_m,), 0.0, tl.float32)
+    # Padding rows own scratch slots and must emit zeros without reading KV.
+    # Keep valid-row accumulation and the ring overwrite semantics unchanged.
+    if skip_empty_tiles:
+        block_start = tl.program_id(0) * block_m
+        live = (length > 0) & (block_start < length)
+        live = live & (block_start + block_m > tl.maximum(length - ring_capacity, 0))
+    else:
+        live = length > 0
+    if live:
+        # Physical-ring order matches the existing masked_attention accumulation.
+        last = end - 1
+        last_index = last % ring_capacity
+        if skip_empty_tiles:
+            # Before the first wrap, only physical keys [0, end) exist.
+            # Omitted blocks have zero probability for every query; keep the
+            # physical order and arithmetic of all remaining blocks unchanged.
+            key_blocks = tl.cdiv(tl.minimum(end, ring_capacity), block_n)
+        else:
+            key_blocks = tl.cdiv(ring_capacity, block_n)
+        for first in range(key_blocks):
+            cols = first * block_n + tl.arange(0, block_n)
+            delta = cols - last_index
+            positions = tl.where(delta <= 0, last + delta, last + delta - ring_capacity)
+            distance = start + rows[:, None] - positions[None, :]
+            allowed = (
+                (rows[:, None] < length)
+                & (cols[None, :] < ring_capacity)
+                & (positions[None, :] >= 0)
+                & (positions[None, :] < end)
+                & (distance >= 0)
+            )
+            if context_size > 0:
+                allowed = allowed & (distance < context_size)
+            base = (slot * num_heads + head) * ring_capacity * head_dim
+            k = tl.load(cache_ptr + base + cols[None, :] * head_dim + dims[:, None], cols[None, :] < ring_capacity, 0)
+            score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5)
+            score = tl.where(allowed, score, -float("inf"))
+            next_max = tl.maximum(maximum, tl.max(score, 1))
+            safe_max = tl.where(next_max == -float("inf"), 0.0, next_max)
+            rescale = tl.exp(maximum - safe_max)
+            p = tl.exp(score - safe_max[:, None])
+            denominator = denominator * rescale + tl.sum(p, 1)
+            v = tl.load(
+                cache_ptr
+                + state_capacity * num_heads * ring_capacity * head_dim
+                + base
+                + cols[:, None] * head_dim
+                + dims[None, :],
+                cols[:, None] < ring_capacity,
+                0,
+            )
+            acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v)
+            maximum = next_max
+    out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None]
+    tl.store(
+        out_ptr + ((batch * num_heads + head) * num_frames + rows[:, None]) * head_dim + dims[None, :],
+        out,
+        rows[:, None] < num_frames,
+    )
+
+
+@triton.jit
+def _advance(
+    end_ptr,
+    slots_ptr,
+    lengths_ptr,
+    end_stride: tl.constexpr,
+    slot_stride: tl.constexpr,
+    length_stride: tl.constexpr,
+    batch_size: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    row = tl.arange(0, block_size)
+    slot = tl.load(slots_ptr + row * slot_stride, row < batch_size, 0)
+    length = tl.load(lengths_ptr + row * length_stride, row < batch_size, 0)
+    active = (row < batch_size) & (length > 0)
+    start = tl.load(end_ptr + slot * end_stride, active, 0)
+    tl.store(end_ptr + slot * end_stride, start + length, active)
+
+
+@triton.jit
+def _write_attend(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    cache_ptr,
+    starts_ptr,
+    slots_ptr,
+    lengths_ptr,
+    out_ptr,
+    qs0: tl.constexpr,
+    qs1: tl.constexpr,
+    qs2: tl.constexpr,
+    qs3: tl.constexpr,
+    ks0: tl.constexpr,
+    ks1: tl.constexpr,
+    ks2: tl.constexpr,
+    ks3: tl.constexpr,
+    vs0: tl.constexpr,
+    vs1: tl.constexpr,
+    vs2: tl.constexpr,
+    vs3: tl.constexpr,
+    start_stride: tl.constexpr,
+    slot_stride: tl.constexpr,
+    length_stride: tl.constexpr,
+    state_capacity: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_frames: tl.constexpr,
+    ring_capacity: tl.constexpr,
+    head_dim: tl.constexpr,
+    context_size: tl.constexpr,
+    write_frames: tl.constexpr,
+    block_m: tl.constexpr,
+    block_n: tl.constexpr,
+    block_w: tl.constexpr,
+):
+    rows = tl.program_id(0) * block_m + tl.arange(0, block_m)
+    bh = tl.program_id(1)
+    batch, head = bh // num_heads, bh % num_heads
+    slot = tl.load(slots_ptr + batch * slot_stride)
+    length = tl.load(lengths_ptr + batch * length_stride)
+    start = tl.load(starts_ptr + batch * start_stride)
+    end = start + length
+    dims = tl.arange(0, head_dim)
+    base = (slot * num_heads + head) * ring_capacity * head_dim
+    v_base = state_capacity * num_heads * ring_capacity * head_dim + base
+    # Every program of this (row, head) writes the whole chunk for its head,
+    # so it only reads ring entries it wrote itself (duplicate writes across
+    # query blocks store identical values). Chunks longer than the ring keep
+    # only their final ring_capacity tokens, as in _write.
+    skip = tl.maximum(length - ring_capacity, 0)
+    for first in range(tl.cdiv(write_frames, block_w)):
+        token = first * block_w + tl.arange(0, block_w)
+        source = token + skip
+        mask = (source < length)[:, None]
+        pos = (start + source) % ring_capacity
+        kval = tl.load(k_ptr + batch * ks0 + head * ks1 + source[:, None] * ks2 + dims[None, :] * ks3, mask, 0)
+        vval = tl.load(v_ptr + batch * vs0 + head * vs1 + source[:, None] * vs2 + dims[None, :] * vs3, mask, 0)
+        dst = pos[:, None] * head_dim + dims[None, :]
+        tl.store(cache_ptr + base + dst, kval, mask)
+        tl.store(cache_ptr + v_base + dst, vval, mask)
+    tl.debug_barrier()
+    q = tl.load(
+        q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, rows[:, None] < num_frames, 0
+    )
+    acc = tl.full((block_m, head_dim), 0.0, tl.float32)
+    maximum = tl.full((block_m,), -float("inf"), tl.float32)
+    denominator = tl.full((block_m,), 0.0, tl.float32)
+    if length > 0:
+        last = end - 1
+        last_index = last % ring_capacity
+        for first in range(tl.cdiv(ring_capacity, block_n)):
+            cols = first * block_n + tl.arange(0, block_n)
+            delta = cols - last_index
+            positions = tl.where(delta <= 0, last + delta, last + delta - ring_capacity)
+            distance = start + rows[:, None] - positions[None, :]
+            allowed = (
+                (rows[:, None] < length)
+                & (cols[None, :] < ring_capacity)
+                & (positions[None, :] >= 0)
+                & (positions[None, :] < end)
+                & (distance >= 0)
+            )
+            if context_size > 0:
+                allowed = allowed & (distance < context_size)
+            k = tl.load(cache_ptr + base + cols[None, :] * head_dim + dims[:, None], cols[None, :] < ring_capacity, 0)
+            score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5)
+            score = tl.where(allowed, score, -float("inf"))
+            next_max = tl.maximum(maximum, tl.max(score, 1))
+            safe_max = tl.where(next_max == -float("inf"), 0.0, next_max)
+            rescale = tl.exp(maximum - safe_max)
+            p = tl.exp(score - safe_max[:, None])
+            denominator = denominator * rescale + tl.sum(p, 1)
+            v = tl.load(
+                cache_ptr + v_base + cols[:, None] * head_dim + dims[None, :],
+                cols[:, None] < ring_capacity,
+                0,
+            )
+            acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v)
+            maximum = next_max
+    out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None]
+    tl.store(
+        out_ptr + ((batch * num_heads + head) * num_frames + rows[:, None]) * head_dim + dims[None, :],
+        out,
+        rows[:, None] < num_frames,
+    )
+
+
+@torch.library.custom_op("vllm_omni::moss_slot_ring_attention", mutates_args=("cache", "end_offset"))
+def slot_ring_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: torch.Tensor,
+    end_offset: torch.Tensor,
+    slot_ids: torch.Tensor,
+    valid_lengths: torch.Tensor,
+    context: int = -1,
+    skip_empty_tiles: bool = False,
+) -> torch.Tensor:
+    """Fused address/mask path, with three ordered launches and no KV gather.
+
+    q/k are already rotary-position encoded. The caller still advances its
+    separate MHA/RoPE offset state; this op advances only end_offset in place.
+    context=-1 means unbounded causal attention over available ring entries.
+    Outputs beyond each row's valid_lengths are exactly zero.
+    skip_empty_tiles skips wholly masked tiles for small (T <= 32) chunks.
+    Larger chunks keep the static loop, which benchmarks faster after wrap.
+    """
+    b, h, t, d = q.shape
+    if k.shape != q.shape or v.shape != q.shape:
+        raise ValueError("q/k/v must have identical shapes")
+    if q.dtype not in (torch.bfloat16, torch.float16) or k.dtype != q.dtype or v.dtype != q.dtype:
+        raise ValueError("q/k/v must share bf16 or fp16 dtype")
+    if d not in (32, 64, 128) or t < 1:
+        raise ValueError("requires head_dim 32/64/128 and positive padded_T")
+    if cache.ndim != 5 or cache.shape[0] != 2 or cache.shape[2] != h or cache.shape[4] != d:
+        raise ValueError("cache must be (2, state_capacity, heads, capacity, head_dim)")
+    s, c = cache.shape[1], cache.shape[3]
+    if not cache.is_contiguous() or cache.dtype != q.dtype:
+        raise ValueError("cache must be contiguous and share q dtype")
+    if end_offset.shape != (s,) or slot_ids.shape != (b,) or valid_lengths.shape != (b,):
+        raise ValueError("invalid metadata shapes")
+    if any(x.dtype not in (torch.int32, torch.int64) for x in (end_offset, slot_ids, valid_lengths)):
+        raise ValueError("metadata must be int32/int64")
+    if q.device.type != "cuda" or any(x.device != q.device for x in (k, v, cache, end_offset, slot_ids, valid_lengths)):
+        raise ValueError("all tensors must share one CUDA device")
+    if context == 0 or context < -1:
+        raise ValueError("context must be -1 or positive")
+    out = torch.empty(q.shape, device=q.device, dtype=q.dtype)
+    if b == 0:
+        return out
+    strides = (end_offset.stride(0), slot_ids.stride(0), valid_lengths.stride(0))
+    _write[(b, triton.cdiv(h * t * d, 256))](
+        k, v, cache, end_offset, slot_ids, valid_lengths, *k.stride(), *v.stride(), *strides, s, h, t, c, d, 256
+    )
+    bm = 16 if t <= 32 else 64
+    _attend[(triton.cdiv(t, bm), b * h)](
+        q,
+        cache,
+        end_offset,
+        slot_ids,
+        valid_lengths,
+        out,
+        *q.stride(),
+        *strides,
+        s,
+        h,
+        t,
+        c,
+        d,
+        context,
+        bm,
+        64,
+        skip_empty_tiles=skip_empty_tiles and t <= 32,
+        num_warps=8 if c <= 256 else 4,
+        num_stages=2,
+    )
+    _advance[(1,)](end_offset, slot_ids, valid_lengths, *strides, b, triton.next_power_of_2(b))
+    return out
+
+
+@slot_ring_attention.register_fake
+def _(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: torch.Tensor,
+    end_offset: torch.Tensor,
+    slot_ids: torch.Tensor,
+    valid_lengths: torch.Tensor,
+    context: int = -1,
+    skip_empty_tiles: bool = False,
+) -> torch.Tensor:
+    return torch.empty(q.shape, device=q.device, dtype=q.dtype)
+
+
+def _check_inputs(q, k, v, cache, metadata, context):
+    b, h, t, d = q.shape
+    if k.shape != q.shape or v.shape != q.shape:
+        raise ValueError("q/k/v must have identical shapes")
+    if q.dtype not in (torch.bfloat16, torch.float16) or k.dtype != q.dtype or v.dtype != q.dtype:
+        raise ValueError("q/k/v must share bf16 or fp16 dtype")
+    if d not in (32, 64, 128) or t < 1:
+        raise ValueError("requires head_dim 32/64/128 and positive padded_T")
+    if cache.ndim != 5 or cache.shape[0] != 2 or cache.shape[2] != h or cache.shape[4] != d:
+        raise ValueError("cache must be (2, state_capacity, heads, capacity, head_dim)")
+    if not cache.is_contiguous() or cache.dtype != q.dtype:
+        raise ValueError("cache must be contiguous and share q dtype")
+    if any(x.dtype not in (torch.int32, torch.int64) for x in metadata):
+        raise ValueError("metadata must be int32/int64")
+    if q.device.type != "cuda" or any(x.device != q.device for x in (k, v, cache, *metadata)):
+        raise ValueError("all tensors must share one CUDA device")
+    if context == 0 or context < -1:
+        raise ValueError("context must be -1 or positive")
+
+
+@torch.library.custom_op("vllm_omni::moss_slot_ring_attention_rows", mutates_args=("cache",))
+def slot_ring_attention_rows(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: torch.Tensor,
+    starts: torch.Tensor,
+    slot_ids: torch.Tensor,
+    valid_lengths: torch.Tensor,
+    context: int = -1,
+) -> torch.Tensor:
+    """``slot_ring_attention`` as one launch, with per-row start positions.
+
+    ``starts`` holds each execution row's pre-call absolute position (the
+    caller's gathered offset). The ring is written and attended exactly as by
+    ``slot_ring_attention``; no offset is advanced here, so the caller owns
+    every offset update and can fuse them with its own.
+    """
+    b, h, t, d = q.shape
+    _check_inputs(q, k, v, cache, (starts, slot_ids, valid_lengths), context)
+    if starts.shape != (b,) or slot_ids.shape != (b,) or valid_lengths.shape != (b,):
+        raise ValueError("invalid metadata shapes")
+    out = torch.empty(q.shape, device=q.device, dtype=q.dtype)
+    if b == 0:
+        return out
+    s, c = cache.shape[1], cache.shape[3]
+    bm = 16 if t <= 32 else 64
+    _write_attend[(triton.cdiv(t, bm), b * h)](
+        q,
+        k,
+        v,
+        cache,
+        starts,
+        slot_ids,
+        valid_lengths,
+        out,
+        *q.stride(),
+        *k.stride(),
+        *v.stride(),
+        starts.stride(0),
+        slot_ids.stride(0),
+        valid_lengths.stride(0),
+        s,
+        h,
+        t,
+        c,
+        d,
+        context,
+        min(t, c),
+        bm,
+        64,
+        16,
+        num_warps=8 if c <= 256 else 4,
+        num_stages=2,
+    )
+    return out
+
+
+@slot_ring_attention_rows.register_fake
+def _(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: torch.Tensor,
+    starts: torch.Tensor,
+    slot_ids: torch.Tensor,
+    valid_lengths: torch.Tensor,
+    context: int = -1,
+) -> torch.Tensor:
+    return torch.empty(q.shape, device=q.device, dtype=q.dtype)

vllm_omni/model_executor/output_snapshot.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Owned, dtype-grouped snapshots for asynchronous runner outputs."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch.utils._pytree import tree_flatten, tree_unflatten
+
+
+@dataclass(frozen=True)
+class RequestOutputSnapshot:
+    """Already partitioned, CPU-owned payloads in the current request order.
+
+    Model finalizers may return this after D2H completes to avoid recursively
+    cloning and partitioning data whose request ownership is already known.
+    """
+
+    inter_stage: list[dict[str, Any] | None]
+    client: list[dict[str, Any] | None] | None = None
+
+
+class PackedOutputSnapshot(dict):
+    """A normal payload mapping with an internal batched-copy plan.
+
+    The runner's snapshot-slot event protects the device slabs until D2H is
+    complete. Host slabs are allocated per output, so downstream consumers can
+    retain their views without depending on reuse of the device ring slot.
+    """
+
+    def __init__(self, leaves: list[Any], spec: Any, groups: list[tuple[list[int], torch.Tensor]]) -> None:
+        super().__init__(tree_unflatten(leaves, spec))
+        self._leaves = leaves
+        self._spec = spec
+        self._groups = groups
+        self.producer_event: torch.cuda.Event | None = None
+
+    def record_producer_event(self, stream: torch.cuda.Stream) -> None:
+        """Publish readiness when a model produces on a separate CUDA stream."""
+        self.producer_event = torch.cuda.Event()
+        self.producer_event.record(stream)
+
+    def copy_to_cpu(self, copy_tensor: Callable[[torch.Tensor], torch.Tensor]) -> dict[str, Any]:
+        leaves = list(self._leaves)
+        for indices, slab in self._groups:
+            cpu = copy_tensor(slab)
+            offset = 0
+            for index in indices:
+                value = leaves[index]
+                size = value.numel()
+                leaves[index] = cpu.narrow(0, offset, size).view(value.shape)
+                offset += size
+        return tree_unflatten(leaves, self._spec)
+
+
+def pack_output_snapshot(
+    payload: dict[str, Any],
+    slot: dict[tuple[Any, ...], torch.Tensor],
+    *,
+    max_buckets: int,
+    reuse_existing_storage: bool = False,
+) -> PackedOutputSnapshot | None:
+    """Copy tensor leaves once per dtype/device, retaining their exact values.
+
+    No CUDA synchronization is introduced here. The caller must select the
+    producer stream and wait for the previous consumer before reusing a slot.
+    """
+    leaves, spec = tree_flatten(payload)
+    grouped: dict[tuple[Any, ...], list[int]] = defaultdict(list)
+    for index, value in enumerate(leaves):
+        if isinstance(value, torch.Tensor):
+            # Preserve the existing per-tensor path for unusual layouts.
+            if value.layout != torch.strided or value.is_quantized:
+                return None
+            grouped[(value.dtype, value.device)].append(index)
+    groups = []
+    for (dtype, device), indices in grouped.items():
+        values = [leaves[index].detach().reshape(-1) for index in indices]
+        size = sum(value.numel() for value in values)
+        # A producer may have already written the leaves into one contiguous
+        # per-forward slab.  Reuse that storage directly instead of copying
+        # it into a second packed slab.  The strict offset/stride check keeps
+        # this optimization limited to views whose order matches the copy
+        # plan; arbitrary payloads retain the existing copy path.
+        if reuse_existing_storage and values and all(value.is_contiguous() for value in values):
+            storage = values[0].untyped_storage()
+            offset = values[0].storage_offset()
+            cursor = offset
+            shares_storage = True
+            for value in values:
+                if value.untyped_storage().data_ptr() != storage.data_ptr() or value.storage_offset() != cursor:
+                    shares_storage = False
+                    break
+                cursor += value.numel()
+            if shares_storage:
+                slab = values[0].as_strided((size,), (1,), storage_offset=offset)
+                offset = 0
+                for index, value in zip(indices, values):
+                    shape = leaves[index].shape
+                    leaves[index] = slab.narrow(0, offset, value.numel()).view(shape)
+                    offset += value.numel()
+                groups.append((indices, slab))
+                continue
+        key = ("packed-output", dtype, device, size)
+        slab = slot.get(key)
+        if slab is None:
+            slab = torch.empty(size, dtype=dtype, device=device)
+            if len(slot) < max_buckets:
+                slot[key] = slab
+        if len(values) == 1:
+            slab.copy_(values[0])
+        else:
+            torch.cat(values, out=slab)
+        offset = 0
+        for index, value in zip(indices, values):
+            shape = leaves[index].shape
+            leaves[index] = slab.narrow(0, offset, value.numel()).view(shape)
+            offset += value.numel()
+        groups.append((indices, slab))
+    return PackedOutputSnapshot(leaves, spec, groups)

vllm_omni/model_executor/stage_input_processors/moss_tts.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
@@ -9,15 +12,13 @@ from collections import defaultdict
 from collections.abc import Mapping
 from typing import Any
 
+import numpy as np
 import torch
 from vllm.inputs import TokensPrompt as OmniTokensPrompt
 from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
-from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
-    parse_chunk_ramp,
-    ramp_chunk_size,
-)
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_chunk_size
 
 logger = init_logger(__name__)
 
@@ -237,7 +238,8 @@ def talker2codec_raw_async_chunk(
     Stage 0 emits newly generated raw codec rows shaped ``[T, n_vq]`` (normally
     ``[1, n_vq]`` per decode step). This processor buffers those new rows until
     a codec chunk is ready, then forwards the chunk to Stage 1. No delay-pattern
-    de-delay is applied on this path.
+    de-delay is applied on this path. An optional ``codec_chunk_ramp`` selects
+    successive chunk sizes and takes precedence over the initial chunk size.
     """
     external_req_id = getattr(request, "external_req_id", None)
     req_id = str(external_req_id if external_req_id is not None else getattr(request, "request_id", id(request)))
@@ -248,6 +250,8 @@ def talker2codec_raw_async_chunk(
         transfer_manager.request_payload = {}
     if not hasattr(transfer_manager, "put_req_chunk"):
         transfer_manager.put_req_chunk = defaultdict(int)
+    if not hasattr(transfer_manager, "ramp_chunk_count"):
+        transfer_manager.ramp_chunk_count = defaultdict(int)
 
     pending_frames = transfer_manager.code_prompt_token_ids[req_id]
 
@@ -255,14 +259,16 @@ def talker2codec_raw_async_chunk(
         codes_dict = multimodal_output.get("codes", {}) or {}
         new_frames = codes_dict.get("audio")
         if isinstance(new_frames, torch.Tensor) and new_frames.numel() > 0:
-            frames_cpu = new_frames.detach().to("cpu", torch.long).contiguous()
-            if frames_cpu.ndim == 1:
-                frames_cpu = frames_cpu.reshape(1, -1)
-            if frames_cpu.ndim != 2:
-                raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_cpu.shape)}")
-            valid_rows = frames_cpu.ne(_MOSS_AUDIO_PAD_CODE).any(dim=1)
-            for frame in frames_cpu[valid_rows]:
-                pending_frames.append(frame.clone())
+            # This runs per request on every decode step. Tiny torch ops cost
+            # tens of microseconds each here; NumPy on the host copy does not.
+            frames_np = new_frames.detach().cpu().numpy()
+            if frames_np.ndim == 1:
+                frames_np = frames_np.reshape(1, -1)
+            if frames_np.ndim != 2:
+                raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_np.shape)}")
+            valid_rows = (frames_np != _MOSS_AUDIO_PAD_CODE).any(axis=1)
+            owned = np.array(frames_np if valid_rows.all() else frames_np[valid_rows], dtype=np.int64)
+            pending_frames.extend(owned)
         # Raw/local streaming should mirror the non-streaming path: the codec
         # decodes only generated audio rows. Reference audio conditions the
         # talker, but feeding its codes into the codec streaming state adds a
@@ -288,26 +294,27 @@ def talker2codec_raw_async_chunk(
         initial_chunk_frames = chunk_frames
 
     pending = len(pending_frames)
-    # codec_chunk_ramp (opt-in): static ladder of chunk sizes replacing the
-    # initial_codec_chunk_frames -> codec_chunk_frames cliff. The chunk index
-    # is maintained by the connector (ramp_chunk_count, incremented per
-    # emitted chunk); past the ladder, ramp_chunk_size falls back to steady.
-    if not hasattr(transfer_manager, "_moss_ramp_parsed"):
-        transfer_manager._moss_ramp_parsed = parse_chunk_ramp(cfg, steady=chunk_frames)
-    _ramp = transfer_manager._moss_ramp_parsed
-    if _ramp is not None:
-        if not hasattr(transfer_manager, "ramp_chunk_count"):
-            transfer_manager.ramp_chunk_count = defaultdict(int)
-        _chunk_idx = int(transfer_manager.ramp_chunk_count.get(req_id, 0))
-        threshold = ramp_chunk_size(_chunk_idx, _ramp, chunk_frames)
-    else:
-        emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0
-        threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames
+    emitted_chunks = int(transfer_manager.put_req_chunk.get(req_id, 0))
+    emitted_any = emitted_chunks > 0
+    threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames
+    # The connector configuration is fixed for this transfer manager's lifetime.
+    if not hasattr(transfer_manager, "_moss_chunk_ramp"):
+        transfer_manager._moss_chunk_ramp = parse_chunk_ramp(cfg, steady=chunk_frames)
+    ramp = transfer_manager._moss_chunk_ramp
+    if ramp is not None:
+        # The ladder is indexed by the connector's segment-local counter, which
+        # restarts at each segment boundary; put_req_chunk is request-global.
+        ramp_index = int(transfer_manager.ramp_chunk_count.get(req_id, 0))
+        threshold = ramp_chunk_size(ramp_index, ramp, chunk_frames)
     if pending <= 0:
         if is_finished:
             transfer_manager.code_prompt_token_ids.pop(req_id, None)
             transfer_manager.request_payload.pop(req_id, None)
             return OmniPayloadStruct(
+                # Explicitly replace the previous chunk (or the prewarmed
+                # placeholder) with no codes. Metadata alone is not an
+                # empty audio snapshot in the native connector contract.
+                codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
                 meta=MetaStruct(
                     req_id=[req_id],
                     left_context_size=0,
@@ -325,13 +332,10 @@ def talker2codec_raw_async_chunk(
     emit_frames = pending if is_finished else threshold
     chunk_rows = pending_frames[:emit_frames]
     del pending_frames[:emit_frames]
-    chunk_codes = torch.stack(
-        [row.to(torch.long).cpu() for row in chunk_rows],
-        dim=0,
-    ).contiguous()
+    chunk_np = np.stack([np.asarray(row, dtype=np.int64) for row in chunk_rows])
     finished = bool(is_finished and len(pending_frames) == 0)
 
-    codec_flat = chunk_codes.transpose(0, 1).contiguous().reshape(-1).to(torch.long)
+    codec_flat = torch.from_numpy(np.ascontiguousarray(chunk_np.T).reshape(-1))
 
     if finished:
         transfer_manager.code_prompt_token_ids.pop(req_id, None)
@@ -342,7 +346,7 @@ def talker2codec_raw_async_chunk(
         meta=MetaStruct(
             req_id=[req_id],
             left_context_size=0,
-            codec_chunk_frames=int(chunk_codes.shape[0]),
+            codec_chunk_frames=int(chunk_np.shape[0]),
             codec_left_context_frames=0,
             code_flat_numel=int(codec_flat.numel()),
             stream_finished=torch.tensor(finished, dtype=torch.bool),

vllm_omni/worker_v2/first_audio_sender.py

@@ -111,9 +111,10 @@ class _PreparedDelivery:
 
 
 class _EngineOutputSink:
-    def __init__(self, output_queue: Any, scheduler: Any) -> None:
+    def __init__(self, output_queue: Any, scheduler: Any, *, upstream_first_audio: bool = True) -> None:
         self.output_queue = output_queue
         self.scheduler = scheduler
+        self.upstream_first_audio = upstream_first_audio
 
     def prepare(self, request_ids: list[str]) -> _PreparedDelivery:
         routes = {}
@@ -130,15 +131,14 @@ class _EngineOutputSink:
         for request_id, pcm in zip(request_ids, pcm_rows, strict=True):
             if request_id not in routes:
                 continue
+            payload = {"model_outputs": pcm, "sr": sample_rate}
+            if self.upstream_first_audio:
+                payload[FIRST_AUDIO_KEY] = torch.tensor(True)
             by_client.setdefault(routes[request_id], []).append(
                 OmniEngineCoreOutput(
                     request_id=request_id,
                     new_token_ids=[],
-                    multimodal_output={
-                        "model_outputs": pcm,
-                        "sr": sample_rate,
-                        FIRST_AUDIO_KEY: torch.tensor(True),
-                    },
+                    multimodal_output=payload,
                 )
             )
         for client_index, outputs in by_client.items():
@@ -163,5 +163,8 @@ class _EngineOutputSink:
                 routes.pop(request_id)
 
 
-def engine_output_queue_sink(output_queue: Any, scheduler: Any) -> _EngineOutputSink:
-    return _EngineOutputSink(output_queue, scheduler)
+def engine_output_queue_sink(
+    output_queue: Any, scheduler: Any, *, upstream_first_audio: bool = True
+) -> _EngineOutputSink:
+    """Route prepared audio; final-stage PCM must not carry the upstream marker."""
+    return _EngineOutputSink(output_queue, scheduler, upstream_first_audio=upstream_first_audio)

vllm_omni/worker_v2/model_states/__init__.py

@@ -38,6 +38,12 @@ def init_omni_model_state(
         getattr(model, flag, False) is True for flag in ("has_preprocess", "has_postprocess", "have_multimodal_outputs")
     )
     if uses_omni_lifecycle:
+        factory = getattr(type(model), "create_mrv2_model_state", None)
+        if callable(factory):
+            state = factory(model, vllm_config, encoder_cache, device)
+            if not isinstance(state, ModelState):
+                raise TypeError("create_mrv2_model_state must return a ModelState")
+            return state
         from vllm_omni.worker_v2.model_states.omni_model_state import (
             OmniModelState,
         )

vllm_omni/worker_v2/omni_ar_model_runner.py

@@ -138,6 +138,23 @@ class OmniARModelRunner(OmniGPUModelRunner):
     # sample_tokens: OmniOutput handling + pooler_output + async D2H
     # ------------------------------------------------------------------
 
+    def sample(self, hidden_states, input_batch, grammar_output):
+        # An explicitly declared model-state hook may return already determined
+        # tokens. Unsupported sampling features retain the upstream path.
+        sample_determined = getattr(type(self.model_state), "sample_determined_tokens", None)
+        if (
+            sample_determined is not None
+            and grammar_output is None
+            and self.batch_sharder is None
+            and self.pp_handler is None
+            and input_batch.num_draft_tokens == 0
+            and input_batch.num_reqs > 0
+        ):
+            output = sample_determined(self.model_state, input_batch, self.sampler)
+            if output is not None:
+                return output, output.num_sampled, output.num_rejected
+        return super().sample(hidden_states, input_batch, grammar_output)
+
     @torch.inference_mode()
     @step_eplb_after()
     def sample_tokens(
@@ -311,6 +328,11 @@ class OmniARModelRunner(OmniGPUModelRunner):
             # Freshly allocated each step and filled after sampling; nothing
             # overwrites them before the host copy.
             return outputs
+        # A producer-side PackedOutputSnapshot already owns its device slabs
+        # for this forward and carries the event that makes them visible to
+        # the output copy stream. Repacking it would add a second GPU copy.
+        if isinstance(outputs, PackedOutputSnapshot):
+            return outputs
         slot_index = self._async_mm_snapshot_cursor
         if self._async_mm_snapshot_pending[slot_index]:
             raise RuntimeError("Async multimodal snapshot ring exhausted before sample_tokens().")
@@ -567,6 +589,8 @@ def _async_copy_mm(
     if not mm_outputs:
         return {}
     if isinstance(mm_outputs, PackedOutputSnapshot):
+        if mm_outputs.producer_event is not None:
+            mm_outputs.producer_event.wait(copy_stream)
         return mm_outputs.copy_to_cpu(
             lambda tensor: _async_copy_tensor(tensor, copy_stream=copy_stream, pin_memory=pin_memory)
         )

vllm_omni/worker_v2/output_snapshot.py

@@ -1,91 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Owned, dtype-grouped snapshots for asynchronous runner outputs."""
+"""Compatibility imports for the shared model/runner output snapshot contract."""
 
-from __future__ import annotations
+from vllm_omni.model_executor.output_snapshot import (
+    PackedOutputSnapshot,
+    RequestOutputSnapshot,
+    pack_output_snapshot,
+)
 
-from collections import defaultdict
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Any
-
-import torch
-from torch.utils._pytree import tree_flatten, tree_unflatten
-
-
-@dataclass(frozen=True)
-class RequestOutputSnapshot:
-    """Already partitioned, CPU-owned payloads in the current request order.
-
-    Model finalizers may return this after D2H completes to avoid recursively
-    cloning and partitioning data whose request ownership is already known.
-    """
-
-    inter_stage: list[dict[str, Any] | None]
-    client: list[dict[str, Any] | None] | None = None
-
-
-class PackedOutputSnapshot(dict):
-    """A normal payload mapping with an internal batched-copy plan.
-
-    The runner's snapshot-slot event protects the device slabs until D2H is
-    complete. Host slabs are allocated per output, so downstream consumers can
-    retain their views without depending on reuse of the device ring slot.
-    """
-
-    def __init__(self, leaves: list[Any], spec: Any, groups: list[tuple[list[int], torch.Tensor]]) -> None:
-        super().__init__(tree_unflatten(leaves, spec))
-        self._leaves = leaves
-        self._spec = spec
-        self._groups = groups
-
-    def copy_to_cpu(self, copy_tensor: Callable[[torch.Tensor], torch.Tensor]) -> dict[str, Any]:
-        leaves = list(self._leaves)
-        for indices, slab in self._groups:
-            cpu = copy_tensor(slab)
-            offset = 0
-            for index in indices:
-                value = leaves[index]
-                size = value.numel()
-                leaves[index] = cpu.narrow(0, offset, size).view(value.shape)
-                offset += size
-        return tree_unflatten(leaves, self._spec)
-
-
-def pack_output_snapshot(
-    payload: dict[str, Any], slot: dict[tuple[Any, ...], torch.Tensor], *, max_buckets: int
-) -> PackedOutputSnapshot | None:
-    """Copy tensor leaves once per dtype/device, retaining their exact values.
-
-    No CUDA synchronization is introduced here. The caller must select the
-    producer stream and wait for the previous consumer before reusing a slot.
-    """
-    leaves, spec = tree_flatten(payload)
-    grouped: dict[tuple[Any, ...], list[int]] = defaultdict(list)
-    for index, value in enumerate(leaves):
-        if isinstance(value, torch.Tensor):
-            # Preserve the existing per-tensor path for unusual layouts.
-            if value.layout != torch.strided or value.is_quantized:
-                return None
-            grouped[(value.dtype, value.device)].append(index)
-    groups = []
-    for (dtype, device), indices in grouped.items():
-        values = [leaves[index].detach().reshape(-1) for index in indices]
-        size = sum(value.numel() for value in values)
-        key = ("packed-output", dtype, device, size)
-        slab = slot.get(key)
-        if slab is None:
-            slab = torch.empty(size, dtype=dtype, device=device)
-            if len(slot) < max_buckets:
-                slot[key] = slab
-        if len(values) == 1:
-            slab.copy_(values[0])
-        else:
-            torch.cat(values, out=slab)
-        offset = 0
-        for index, value in zip(indices, values):
-            shape = leaves[index].shape
-            leaves[index] = slab.narrow(0, offset, value.numel()).view(shape)
-            offset += value.numel()
-        groups.append((indices, slab))
-    return PackedOutputSnapshot(leaves, spec, groups)
+__all__ = ["PackedOutputSnapshot", "RequestOutputSnapshot", "pack_output_snapshot"]