[gcanlin]

Purpose

Reduce persistent streaming-decoder state memory for MOSS-TTS-Local v1.5 / MOSS Audio Tokenizer v2. The new V2 decoder-pool path is enabled by default, without an environment-variable switch. This is a memory optimization, not an E2E throughput improvement claim.

• Share per-step offsets, ring positions, scatter indices and attention masks within each resolution group; advance offsets once after its final layer.
• Replace one persistent scratch slot per graph-padding row with a single non-leaseable null slot. Mask persistent KV/offset writes for invalid rows.
• Write back only touched ring positions with Triton custom ops. Preserve the reference scatter, physical key order and native SDPA, including existing oversized-chunk scatter behavior.
• Give the shared-state compile adapter a separate AOT identity.
• V2 decoder pools use the new path automatically. Fixed-width/reference streaming retains independent state, and V1 codecs retain their existing behavior. The former VLLM_OMNI_MOSS_CODEC_SHARED_KV environment variable is no longer read.

This does not implement paged attention or remove the complete-ring gather. K/V data are never shared across layers or live requests. Earlier Local lookup/fused-kernel experiments are not included.

Related to #7202, but targets persistent codec state allocation and per-resolution metadata rather than its Local/MTP and masked-attention changes; not dependent on that PR.

Test Plan

vLLM Version: 0.28.0; PyTorch 2.13.0+cu130; Triton 3.7.1.

vLLM-Omni Commit: measured baseline 677265b64; measured implementation a2f90fd9 with the new path enabled. Follow-ups remove local-only artifacts from the PR, align test conventions, and make that same path the V2 default. Serving performance has not been remeasured after the default-switch change.

Hardware: one H200 (143771 MiB reported total), both AR and codec stages on GPU 0, BF16. Shared host. Memory fractions 0.60/0.30; reference encoder on CPU; chunk sizes 1/15 frames. No precision/sampling/scheduling changes.

Regression coverage includes reset/reorder, wraparound, oversized chunks, all-padding batches, null-slot immutability, group offset aliases, close/reinitialize and legacy lifecycle, CUDA Graph replay, strict dynamic compilation, and real wrapper staging with a small codec.

Local regression (requires the repo/vLLM environment and a CUDA GPU for the GPU cases; no checkpoint download):

python -m pytest tests/model_executor/models/moss_tts/test_shared_codec_kv.py -o addopts= -q --run-level core_model

CI-like CPU selection (existing L1 CPU sweep; no Buildkite configuration changes):

python -m pytest tests/model_executor/models/moss_tts/test_shared_codec_kv.py -o addopts= -q -m 'core_model and cpu' --run-level core_model

Microbenchmark numbers below were collected with a local-only harness; new benchmark scripts are intentionally excluded from this PR. The existing repository serving benchmark can reproduce the E2E workload below.

Serving reproduction:

CUDA_VISIBLE_DEVICES=0 vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni --host 0.0.0.0 --port 8124 --trust-remote-code --deploy-config vllm_omni/deploy/moss_tts_local.yaml --allowed-local-media-path /path/to/data
VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000 VLLM_OMNI_BENCH_AUDIO_CHANNELS=2 python benchmarks/tts/bench_tts.py --host 127.0.0.1 --port 8124 --model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --task voice_clone --locale en --dataset-path /path/to/seedtts_testset --num-prompts 1088 --concurrency 64 --output-len 256 --output-dir results/shared-kv -- --save-detailed

For C128 memory reproduction, change both stage max_num_seqs values to 128 and append 128 to codec capture buckets. Main/PR used identical resulting configs. The checked-in deploy YAML is unchanged from the baseline (C64).

Test Result

C128 memory

Post-capture idle values: all 16 codec graphs captured; three samples five seconds apart were identical for each variant. GPU was empty before each serial launch.

Memory (MiB)	Main	PR
Whole GPU used	118859	107189
AR process	83436	83436
Codec process	35406	23736
Whole GPU free	24298	35968

11670 MiB / 11.40 GiB less service memory, entirely in the codec process. The tensor-shape KV saving is 10.49 GiB; allocator/graph state also contributes to process memory. Main compiled fresh while PR reused an AOT artifact. These are not exact peak or matched fully warmed measurements: the initial main warmup was intentionally interrupted, and the PR arm used startup-only measurement.

C64 E2E

Full1088 warmup excluded, then two full1088 measured rounds per variant:

Audio-s/s	Round 2	Round 3	Mean
Main	144.1106	146.5795	145.3450
PR	144.1143	146.3224	145.2184

-0.087%: no measurable throughput benefit. All completed phases: 1088 successful / 0 failed. Within-arm spans were 1.70% and 1.52%. Ordered input lengths match. Runs were sequential, not counterbalanced; generated lengths were stochastic.

Plain CUDA Graph microbenchmarks improved decode time, but compiled B1/T15 showed only 6.797 -> 6.503 ms (~4.3% less time). Do not extrapolate plain-graph gains to serving.

Validation boundaries / draft follow-up

• Default-path regression rerun: 11 passed, 15 warnings in 28.84s on H200, including old environment-variable unset/0/1 cases. Related codec CUDA Graph, reference encoder, per-request codes and local-depth RoPE tests: 67 passed, 15 warnings in 1.47s. All applicable local pre-commit hooks pass. Existing deprecation/version-mismatch warnings remain.
• Real-weight synthetic-code waveform checks are diagnostic only. The reference T=480/C=400 CUDA scatter has duplicate destinations and is not assumed bitwise deterministic. Some runs show candidate/control differences; no blanket bit-exactness or perceptual-equivalence claim.
• WER/SIM/UTMOS and real-code quality acceptance remain pending. The new V2 path now defaults on at the contributor's request; this default change does not constitute quality acceptance. Keep the PR draft pending review and validation.
• DCO sign-off awaits contributor confirmation; this draft is not merge-ready.
• Based on the measured main revision, not rebased onto the latest main. A local git merge-tree check against 0e602509 found no text conflicts; integration/CI validation is pending.
• Repository quick precheck performed; separate AOT adapters intentionally retained because state mutation/aliasing differs. Public helper typing remains a reviewable follow-up. No new model-specific Python examples, no allowlist expansion.

Implementation notes, reports, raw artifacts, and experimental benchmark scripts remain local. No design Markdown, results/, or benchmarks/ changes are included in the final PR diff.

[hsliuustc0106]

This PR touches vllm_omni/model_executor/, tests/model_executor/ (5 files). Based on CODEOWNERS coverage of the changed files, the most-related reviewers appear to be:

@tzhouam @Sy0307 @amy-why-3459

Could one of you take a look when you get a chance? Thanks!

[hsliuustc0106]

@gcanlin this PR is labeled ready + high priority, but CI is failing on the latest commit:

• DCO

Could you please take a look and push an update to get CI green? Once the checks pass we can proceed with review/merge. Thanks!