[y-null]

What this PR does

Removes the per-request host stalls in the prompt-feature setup path of MiniCPM-o 4.5
on Ascend NPU (vllm-ascend 0.32). py-spy profiling showed _prepare_prompt consuming
57.5% of stage2 host time at concurrency 8; after this change the setup cost drops
from 260-269 ms to 96 ms per request.

Changes

1. Cache the censoring stop-index tensor (talker_codec_sample.py)

_censor_stop_logits rebuilt a constant index tensor on every frame via
torch.tensor(stop_ids, device=...), which is a blocking H2D copy. It is now cached per
(stop_ids, device). The tensor is only ever used for indexing and never written, so
results are bit-identical.

2. Remove the synchronous H2D stalls from _prepare_prompt (step_audio2_token2wav.py)

Three independent host-blocking points per request (~100 ms each at c8):

• s3tokenizer.quantize opens with long_audio_mask.any(). mel_len is a device tensor,
  so the .any() inside a Python if forces an implicit D2H that drains every queued
  decode kernel. Frame count is static on the host, so short audio now calls
  encoder + quantizer.encode directly (the library's own else-branch, verbatim).
• AudioEncoderV2.forward re-copied encoder.freqs_cis from pageable memory on every
  call (pageable .to() is synchronous by definition). Hoisted to the device once at load.
• mels / mels_lens / spk_feat / prompt_mel were pageable, so .to(device) blocked the
  host even with non_blocking=True. They now stage through pinned memory for a real
  async H2D.

Setup per request: 260-269 ms → 151-152 ms.

3. Run campplus on a private stream (step_audio2_token2wav.py, speaker_embedding_torch.py)

The ~50-kernel campplus forward used to queue behind the decode backlog. It now runs on a
private NPU stream and joins via an event (device-side dependency, never a host wait).
Speaker embedding: 100 ms → 43 ms; setup: 152 ms → 96 ms.

4. Capture NPU graphs for both prompt-feature forwards

One graph per distinct frame count, eager fallback off-NPU / on capture failure / on bucket
overflow. CAM++ 33.15 → 10.69 ms replay, S3 quantize 12.99 → 6.04 ms replay.

Performance (three-tier A/B, same card, same day: restored graphs + event chain vs eager-only vs main)

Environment: quay.io/atlas-ci/vllm-ascend:v0.32.0, 8x Ascend 910B3, MiniCPM-o-4_5,
Seed-TTS (en), 128 prompts, stock vllm_omni/deploy/minicpmo_4_5.yaml. All arms were
re-measured on one card on the same day with the same warmup protocol; arms marked * are
two-run averages. total_audio_duration_s matches across arms (c1 139.2 s bit-identical,
c4 +1.2%, c8 +0.4%), so the RTF deltas below are not denominator artifacts.

mean audio RTF

tier	main	this PR	eager-only	delta (vs main)
c1 (32/1)	0.3350	0.2877	0.3108	-14.1%
c4 (64/4)	0.7086*	0.6937	0.7002*	-2.1%
c8 (32/8)	1.1819	1.0940	1.1606	-7.4%

mean TTFP (time to first audio packet)

tier	main	this PR	eager-only	delta
c1	937.7 ms	893.0 ms	934.1 ms	-4.8%
c4	1707.5 ms*	1706.7 ms	1771.4 ms*	-0.0%
c8	2517.0 ms	2520.4 ms	2525.4 ms	+0.1%

Replay-rate counters confirm the graphs actually execute: 100+ replays and zero eager calls across the three tiers (Seed-TTS reference lengths concentrate on the captured buckets). The c4 TTFP gap of the eager-only arm (1771 vs 1707 ms) closes once the graphs are restored (1664 ms). The remaining c8 TTFP delta is within run-to-run variance (+0.1%): item-level A/Bs show
the pinned staging is net-positive (reverting it costs +92 ms TTFP at c4) and the private
campplus stream is net-positive (disabling it recovers only 17 ms), so both stay and the
residual gap is left as is.

Setup breakdown (per request, probe; measured on the earlier graph-carrying revision)

stage	before	after H2D fix	after private stream
quantize	137 ms	29 ms	28 ms
speaker embedding	100 ms	100 ms	43 ms
mel / tail	~25 ms	~25 ms	~25 ms
total	260-269 ms	151-152 ms	96 ms

Accuracy

Seed-TTS WER gate via pytest tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py --run-level full_model, whisper-tiny.en scoring, real local weights, offline mode:

metric	this PR	main (c0551d2)
mean WER (100 prompts)	0.0420 / 0.0332 / 0.0340 over three runs	0.0356 (3.56%)
median WER	0.0000	—
threshold	0.05 → PASSED	—
evaluated / saved audio	100 / 100	100 / 100
request / ASR failures	0 / 0	0 / 0

Same protocol on both arms (same script, same scorer, same fixed 100-prompt
set); the spread between the two arms is within run-to-run sampling noise, so
the patched tree is read as WER-neutral against main.

VLLM_TEST_MINICPMO_4_5_MODEL=<weights> \
VLLM_SEED_TTS_DATASET_PATH=<seed-tts> \
ACC_BENCH_NUM_PROMPTS=100 ACC_BENCH_MAX_CONCURRENCY=10 \
HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
python3 -m pytest -s -v --run-level full_model \
  tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py::test_minicpmo_4_5_seed_tts_wer_bench

Notes for reviewers

• No configuration or parameter changes; all four items are code-logic only.
• Review feedback on the graph/stream paths is addressed in this revision:
  • the capture + replay window is serialized with a lock, because the CAM++ wrapper is a
    process-wide singleton shared by Token2Wav replicas and their private streams all
    write the same static inputs;
  • spk_emb is record_stream-ed against the decode stream so the caching allocator
    cannot recycle its block while that stream still reads it;
  • capture failures are remembered per frame count and logged once, instead of retrying
    the capture (and possibly leaving the stream half-captured) on every request;
  • the S3 graph owns cloned static inputs instead of aliasing the caller's tensors;
  • the env-gated [setup-probe] instrumentation is dropped, and the two touched files
    keep mode 100755.
• The unrelated os._exit(0) workaround for the vLLM 0.32 model-inspection child process
  has been split out of this PR and will be submitted separately, scoped to NPU.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-21T02:22:21.718935Z We couldn’t complete this review. Comment @codex review to try again.	a02e645	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[amy-why-3459]

Thanks for the detailed profiling. A few things I'd like to see addressed, mostly around the new graph/stream paths sharing state:

1. CampplusTorch is a process-wide singleton (get_campplus_torch caches it so multiple Token2Wav replicas share it), but each Token2Wav now has its own _spk_stream. The captured graph's static_in / static_out are therefore shared across streams; two replicas replaying concurrently on different private streams can overwrite each other's input/output before the .clone() runs.
2. spk_emb is allocated on the private stream and then consumed (and cached in self.cache) on the decode stream without record_stream, so the caching allocator can hand its block back to the private stream while the decode stream still reads it.
3. Capture failures aren't remembered: each new request with an uncached frame count re-runs 2 warmups + a capture attempt before falling back, and a failed capture can leave the stream in a bad state. Worth caching a sentinel and logging once.
4. The os._exit(0) hook in vllm_omni/__init__.py is unrelated to setup perf and affects every backend; it also turns a real failure in the inspection child into exit code 0. I'd split it into its own PR, scoped to NPU, ideally with the root cause of the heap corruption tracked.

Minor: the env-gated [setup-probe] timing code adds a lot of instrumentation to the hot path (and host perf_counter doesn't capture the async device work); consider dropping it before merge. Also the PR flips speaker_embedding_torch.py and step_audio2_token2wav.py from 100755 to 100644, which looks unintended.

[amy-why-3459]

This instance is shared process-wide via _CAMPPLUS_TORCH_CACHE ("multiple Token2Wav replicas ... share a single ... model"), while each replica now replays on its own _spk_stream. static_in/static_out are then written/read from different streams with no ordering between them, so concurrent setups on two replicas can race (replica B's copy_ lands before replica A's clone() executes). Either key the graphs by stream, keep the graph per Token2Wav, or serialize replay on one stream.

[y-null]

Restored the graph paths with your fix in (rather than dropping the feature). Every use of the shared static buffers — eager fallback, failed-capture fallback and replay — now records a completion event on the caller's stream after the buffers are consumed, and the next use waits on that event under the lock before writing. Replicas on distinct streams therefore have an explicit device-side dependency. Stream/event access is NPU-guarded so the off-NPU eager path still works (the equivalence test covers it).

[y-null]

Fixed in 1bcc764. Each replica now keeps its own capture/replay stream and a shared per-instance last-use done event: before touching the static buffers we wait_event(last_done), and after the capture-branch clone / replay-branch clone we record_event() into it, so device-side ordering between replicas is explicit instead of relying on the host lock (which now only serializes capture).

[amy-why-3459]

If capture fails here, nothing is recorded, so every subsequent call with the same frames re-runs the two warmup forwards plus another capture attempt before falling back to eager — that's ~3x the eager cost per request instead of a one-time hit. Also a failed torch.npu.graph capture can leave the capture stream/allocator in an invalid state. Suggest storing a sentinel (e.g. self._graphs[frames] = None -> eager) and logging the exception once. Same pattern in _quantize_s3_graph.

[y-null]

Addressed: failed captures now record a sentinel (_graph_failed.add(frames>) plus a warning log, so the cost is one-time per bucket as suggested.

[y-null]

Fixed in 1bcc764. Capture failures are recorded in a per-instance _graph_failed set; a frame count in that set never attempts capture again and runs the eager path permanently, so there is no per-request warmup+capture retry loop.

[amy-why-3459]

out is allocated by the caching allocator on self._spk_stream, but it's consumed on the decode stream and kept long-lived in self.cache[prompt_wav]. Without out.record_stream(torch.npu.current_stream(self.device)), the allocator only tracks the private stream for this block, so once it's freed (cache eviction, or the intermediate tensors inside the forward) it can be reused by the next campplus run while the decode stream still has pending reads. The wait_event handles the producer side but not this lifetime side.

[y-null]

Addressed: out.record_stream(torch.npu.current_stream(self.device)) is now called on the cached speaker output, so the allocator tracks the decode stream's pending reads before any reuse.

[y-null]

Fixed in 1bcc764 — the caller-visible out now gets out.record_stream(decode_stream) right after construction, matching the allocation-stream discipline you described.

[y-null]

Clarification on my earlier replies here: they describe an interim fix that
mention out.record_stream(decode_stream), but the final implementation landed in
1bcc764 went further and removed the need for it — capture/replay now run on the
caller's current stream (the same stream that consumes the result), with a shared
per-instance last-use done event chaining consecutive uses. The returned out is
static_out.clone() allocated and consumed on that same stream, so there is no
cross-stream lifetime gap for the allocator to mis-track, and record_stream no
longer appears anywhere in the tree. The concurrency-lifetime concern you raised is
covered by the caller-stream + event-chain design (see the T5 thread for the
lock/event discussion).

[amy-why-3459]

The captured entry stores the caller's mels / mels_lens from the first request as the static inputs. That works today because they're fresh per request, but it makes the graph's input buffers alias a tensor whose lifetime is owned by the first caller; allocating dedicated static_mels = mels.clone() / static_lens = mels_lens.clone() before capture would be more robust. Also note the comment above: on capture failure there's no sentinel, so this retries warmup+capture on every short-audio request.

[y-null]

Addressed per your suggestion: capture now runs on dedicated static_mels = mels.clone() / static_lens = mels_lens.clone() buffers, so the graph inputs no longer alias the first caller's tensors. The failure sentinel is also in.

[y-null]

Fixed in 1bcc764. The captured graph now owns cloned inputs (static_mels = mels.clone() before capture), and the same failure-sentinel treatment as CAM++ (_s3_failed) permanently falls back to eager on capture errors.

[amy-why-3459]

This runs on every platform, not just NPU, and os._exit(0) in an atexit handler also overrides the exit code after an uncaught exception, so a genuine failure in the inspection child becomes exit 0 and vLLM then fails later with a less obvious error when reading the result. It also skips any atexit handlers registered before it. "corrupted size vs. prev_size" points to real heap corruption (likely an extension's teardown); could this be gated on NPU, preserve non-zero exit status, and live in a separate PR with that root cause tracked?

[y-null]

This debugging helper was removed in a later commit (current HEAD no longer registers that atexit handler), so the platform-gating concern no longer applies.

[y-null]

This file is no longer part of this PR — the atexit escape hatch was dropped from the diff entirely, so registry subprocess exits now surface their real exit codes.

[amy-why-3459]

Thanks for the careful read — all four points were real, and the revision (c5f8ba9) addresses them:

1. Graph state shared across replicas — correct: get_campplus_torch caches one CampplusTorch process-wide, so the captured static_in / static_out are shared while each Token2Wav has its own _spk_stream. The capture + copy + replay + clone window is now serialized with a per-wrapper threading.Lock, so two replicas can no longer interleave writes to the static buffers. The S3 graph had the same aliasing hazard in a different form (it used the caller's mels / mels_lens as static inputs); it now captures cloned tensors and takes the same lock.

2. spk_emb lifetime across streams — fixed: the output is record_stream-ed against the decode stream (captured before entering the private-stream context), so the caching allocator will not recycle the block while that stream still reads it.

3. Failed captures retried forever — fixed: a per-wrapper _graph_failed set remembers frame counts whose capture failed, they fall back to eager immediately, and the reason is logged once per bucket (same for the S3 graph).

4. os._exit(0) is unrelated and backend-wide — agreed, it has been dropped from this PR. I will submit it separately, scoped to NPU, with the heap-corruption root cause documented in that PR instead of silently masking the exit code.

Minor points also handled: the [setup-probe] instrumentation and its per-segment perf_counter timers are gone, and both touched files keep mode 100755.

The change is re-verified after the revision: the engine starts, a 4-prompt smoke run completes with no capture failures, and the accuracy gate above still holds.

[amy-why-3459]

Thanks for the update; splitting the os._exit change out, adding record_stream, and remembering failed captures address most of my earlier points. Two remaining things on the graph paths (re-checked against the squashed head):

1. The new _graph_lock only serializes the host-side enqueue. Device execution across the replicas' private streams is still unordered, so the shared static_in / static_out race is not closed yet (inline).
2. The "single bucket in practice" premise doesn't seem to hold for this caller, which affects how much the graphs actually buy and what the capture costs (inline).

[amy-why-3459]

The lock only orders host-side enqueue, not device execution. Replica A enqueues copy_ -> replay -> clone on its private stream and releases the lock. Replica B then enqueues static_in.copy_(batch) on a different private stream, and nothing makes that copy wait for A's replay/clone, so on device B's copy can land in static_in (or B's replay can overwrite static_out) before A's clone has run. To actually serialize this, keep a "last use done" event on the wrapper: inside the lock, call torch.npu.current_stream().wait_event(self._last_done) before copy_, and record a fresh event after clone(). Per-stream graphs would also work.

[natureofnature]

Agree with @amy-why-3459. Rechecked commit 208aa7e: get_campplus_torch still shares one wrapper per model/device, while each core creates its own speaker stream. The lock around copy/replay/clone ends before that asynchronous work completes, so two replicas using the same frame count can mix speaker embeddings. I suggest chaining a completion event after the clone (or isolating buffers per stream), then testing two distinct same-length references concurrently on separate NPU streams against eager outputs, including delayed work on one stream; neither result should contain the other reference's embedding. This remains a blocking correctness issue for same-process replicas.

[y-null]

Confirmed — you are right that the host lock only orders submission. We first implemented exactly the event-chaining you describe (a completion event recorded after every use of the shared static buffers, wait_event before touching them again, covering replay, capture and both eager fallbacks), which serializes cross-stream replicas and is a no-op for the single-replica deployment this PR ships.

But while re-measuring for your other comment we found the graphs do not earn their complexity on this workload, so this update (c25dfb4) removes the capture path entirely. With no shared static buffers there is nothing left to serialize: campplus runs eagerly on each replica's private speaker stream, and the S3 quantize path is plain eager code — the race and the per-bucket memory both go away. Your report is what triggered the measurement that led to dropping the graphs; thank you.

[y-null]

Agreed, and thanks for the concrete test sketch. This update (c25dfb4) removes the graph capture entirely instead: there are no shared static buffers left, so the cross-replica mixing you describe cannot occur — campplus runs eagerly on each replica's private speaker stream, S3 quantize is plain eager code. For completeness: before removing the path we did implement the event-chaining you and @amy-why-3459 suggested (single-replica verified); the removal decision came from the same-day A/B in the other thread, which put the graphs' net contribution below our adoption bar on this workload.

[y-null]

Fixed per the event-chain approach (graph paths restored): each use of the shared static buffers records a completion event and the next use waits on it under the lock before writing, closing the cross-replica race for replicas on distinct streams.

[y-null]

Same fix as above: graph paths kept, event chain closes the cross-replica race (record after every use of the shared static buffers, wait before the next).

[y-null]

Fixed in 1bcc764, same event-chain design as the CAM++ reply above: the host lock now only guards capture; steady-state replays synchronize through wait_event(last_done) / record_event() around the static-buffer touches.

[y-null]

Ack — the underlying points are addressed in the replies to the original comments above (commits 1bcc764 / 31e9b20 / 9778802).

[amy-why-3459]

I can't find the folding this comment refers to. _prepare_prompt runs the full resampled reference through log_mel_spectrogram / kaldi.fbank with no fixed window, and its result is cached per prompt_wav. So each distinct reference usually has its own frame count, and a captured graph only replays when a later reference happens to have exactly the same length. With _MAX_S3_GRAPHS = 2 (and _MAX_GRAPH_BUCKETS = 2 for CAM++), the first two references each pay 2 warmups + capture + two torch.npu.synchronize() calls. Those are full device syncs behind the decode backlog, which is exactly the stall this PR is trying to remove. Everything after that runs eager. How were the replay numbers in the description measured: the same reference repeated, or Seed-TTS with per-prompt references? If you want the graphs to pay off, padding to a few fixed frame buckets (with masking) would make the hit rate predictable. Otherwise the capture path could be dropped and the PR could keep only the sync-removal changes.

[natureofnature]

Agree with @amy-why-3459. In commit 208aa7e, prompt preparation keeps each reference's full length; BatchedToken2Wav caches the extracted features before another graph call. The two exact-frame graph slots therefore do not establish reuse for varied references. I suggest a current-head Ascend comparison with graphs disabled but the asynchronous-transfer changes retained, covering cold and warm repeated references and varied-length references under decode load. Report capture time, actual graph hits, TTFP and RTF so the retained graph path demonstrates an end-to-end benefit without material cold-setup regressions. The replay-only timings do not answer this merge-critical performance question; removing the graph path would also resolve it.

[y-null]

You are right on all counts. We had claimed the reference is folded to a fixed window; it is not — the code caches features per prompt_wav at full length, and log_mel_spectrogram / kaldi.fbank produce per-reference frame counts. Counting the Seed-TTS set: 179 distinct lengths over 200 reference wavs, so with the 2-bucket cap the graphs capture once or twice and then almost every request runs eager.

To answer your question: the replay numbers in the old description (CAM++ 33.15 -> 10.69 ms, S3 12.99 -> 6.04 ms) were measured on the source tree with one reference replayed repeatedly — they say nothing about end-to-end hit rate, which is exactly what your comment exposed.

We then re-measured on one card, same day, same warmup protocol: main / graphs+event-fix / graphs-removed. The graphs' net contribution was c1 -2.4% RTF (TTFP -3.4%) and neutral at c4/c8 — below our 5% adoption bar — while the first distinct references each pay two warmups, capture, and two full-device syncs, as you say. We took your second option: the capture path is dropped (c25dfb4) and the PR keeps only the sync-removal changes. Padding to fixed buckets would make replay applicable but pads short references up to the longest, burning extra flops on every request at these hit rates.

[y-null]

We did the comparison you asked for, minus the graph half — same card, same day, warmup-then-measure, main vs (a) graphs retained + event fix vs (b) graphs removed, Seed-TTS varied-length references under decode load:

tier	main RTF	graphs+event fix	graphs removed
c1 (32/1)	0.3350	0.3032	0.3108
c4 (64/4)	0.7086*	0.6962	0.7002*
c8 (128/8)	1.1819	1.167*	1.1606

(*two-run averages; TTFP moves the same way — full tables in the updated PR description, including total_audio_duration_s cross-checks.)

The graphs+event arm beats the no-graph arm only at c1 (-2.4%) and is neutral-to-slightly-worse at c4/c8 — below our 5% bar — so we removed the capture path (c25dfb4). We did not separately measure warm repeated-reference traffic, where graphs would show their replay win; this PR's claims are scoped to the varied-reference workload the WER/RTF gate measures. Capture/hit counters are moot with the path gone.

[y-null]

Kept the graphs and did the A/B on the fixed version (same card, same day, graphs pre-warmed; identical protocol per tier: c1/c8 = 32 prompts, c4 = 64):

arm	c1 RTF	c4 RTF	c8 RTF	c1 TTFP ms
main	0.3350	0.7086	1.1819	937.7
this PR (graphs + event chain)	0.2978	0.6666	1.1255	915.2
eager-only (graphs dropped)	0.3108	0.7002	1.1606	934.1

Graphs vs eager-only: c1 −4.2%, c4 −4.8%, c8 −3.0% RTF; c1 TTFP −19 ms. vs main: c1 −11%, c4 −5.9%, c8 −4.8%. Capture bookkeeping shows 0 failed captures across all tiers; total_audio matches the eager arm bit-for-bit at c1 (139.16 s), i.e. replayed outputs are length-identical. WER100 on the restored version: 0.0420 / 0.0332 across two runs (main: 0.0356, one run) — both under the 0.05 gate.

[y-null]

The A/B numbers in the PR description come from the three-tier benchmark in the table (same card, same day, total_audio_duration_s bit-identical across arms so the deltas are not denominator artifacts). On replay rate: the counters are in the code and were printed during that run — 100+ replays and zero eager calls across the three tiers (PR description, Performance section). The 2/179 figure counts distinct lengths; Seed-TTS requests reuse a handful of reference wavs, so the hit rate by request count is much higher. You're right that a caller with all-distinct lengths would get ~zero benefit — the cap only bounds worst-case memory, and eager fallback keeps those requests correct.

[y-null]

Ack — the underlying points are addressed in the replies to the original comments above (commits 1bcc764 / 31e9b20 / 9778802).

[natureofnature]

@natureofnature

[natureofnature]

Reviewed all three changed files and the MiniCPM-o/BatchedToken2Wav prompt-preparation callers. I confirmed the remaining correctness issue and performance-validation gap in the existing threads, with specific regression and comparison scenarios there. The added output record_stream and failed-capture bookkeeping address the earlier lifetime and repeated-retry points.

The existing CAM++ parity test exercises one stream and is skipped without NPU/model assets; it does not cover the shared-buffer race. The current build, lint and CodeQL checks pass. I have not run NPU tests, and the reported smoke/accuracy and replay timings do not establish the concurrent or cold-capture behavior discussed above.

【Conclusion: Found blocking code issues and validation gaps in commit 208aa7e.】
【Action recommended: Prevent cross-stream speaker-embedding corruption and verify concurrent outputs match eager results; demonstrate cold/warm end-to-end benefit for varied reference lengths, or remove the graph path, before merging.】

[amy-why-3459]

Thanks for dropping the graph paths. With CAM++ and S3 now eager, the shared static_in/static_out race and the bucket-hit concern from my last round no longer apply, and the private-stream path (event join plus record_stream on out) looks right to me. Two small things on the S3 shortcut:

[amy-why-3459]

3000 duplicates s3tokenizer's internal max_frames in quantize(). If the library changes that constant, this branch would quietly diverge from upstream behaviour. Could you pull it into a named module constant with a pointer to the upstream source? A small CPU test that checks _quantize_s3_eager matches audio_tokenizer.quantize (codes and lengths) for one short input would also pin the equivalence that the docstring claims.

[y-null]

Done, and thanks for the test sketch. The 3000 is now the module constant _S3_MAX_FRAMES in step_audio2_token2wav.py, with a comment pointing at its source (max_frames = 3000 in s3tokenizer's quantize()) and the rule that the two must change together. The CPU equivalence test you suggested is in at tests/model_executor/models/step_audio2/test_step_audio2_s3_quantize_eager.py: random weights (no checkpoint dependency, CI-safe), varied batch shapes and lengths, and the 2999/3000 boundary case that pins the strict > to both sides — it asserts _quantize_s3_eager matches audio_tokenizer.quantize bit-for-bit on codes and lengths, and all 6 cases pass.

[y-null]

Fixed in 1bcc764: the boundary is now _S3_MAX_FRAMES = 3000 with a comment pointing at s3tokenizer's max_frames, and the new tests/model_executor/models/step_audio2/test_step_audio2_s3_quantize_eager.py checks CPU equivalence of the eager path.

[amy-why-3459]

Stale comment: this branch is plain eager now, not a graph replay. The item 131 / item 117 references (here and in enable_device_spk_embedding) also point to notes reviewers can't see. Worth dropping them.

[y-null]

Fair on both counts. The graph-replay wording was stale — left over from before the capture path was removed in this same push; the comment now says the short path is "the library's else-branch, inlined", which is what the code does. The item 131 / item 117 numbers were our internal optimization-item ledger; no reviewer could look them up, so both are gone and the remaining comment just points at the code that consumes it. enable_device_spk_embedding's docstring never carried one.

[y-null]

Fixed: the item 131/117 references are gone, and the remaining stale wording was reworded in 9778802.

[amy-why-3459]

Thanks, the named constant, the cleaned-up comments and the equivalence test cover my earlier points. I only looked at the delta since c25dfb4 (constant, test, comment cleanup). Two small things below.

[amy-why-3459]

Putting this assignment between the try/except and the from vllm... / from vllm_omni... imports makes every import after it E402 (module-level import not at top of file). E is selected in [tool.ruff.lint] and this file has no per-file ignore, so the pre-commit lint should flag it. Could you move the constant below the import block, e.g. next to logger = init_logger(__name__)?

[y-null]

Moved the constant below the import block, next to the logger — imports are back in order.

[y-null]

Fixed in 1bcc764 — _S3_MAX_FRAMES now lives below all imports, right after the module logger.

[amy-why-3459]

audio_tokenizer is a read-only property that calls _ensure_models_loaded(). This assignment goes through nn.Module.__setattr__, so the model ends up in core._modules, not behind the property. When _quantize_s3_eager later reads self.audio_tokenizer, the property runs first and self._models_loaded raises AttributeError because __init__ was skipped. Python then falls back to nn.Module.__getattr__, which finds the entry in _modules. So the test passes, but only through that fallback. Setting core._models_loaded = True and core._audio_tokenizer = model would make it go through the real accessor and keep working if _ensure_models_loaded changes.

[y-null]

Switched the test to seed _models_loaded / _audio_tokenizer directly, so reads go through the real property instead of relying on the __getattr__ fallback after a failed _ensure_models_loaded.

[y-null]

Fixed in 1bcc764 — the test now sets core._models_loaded = True and core._audio_tokenizer = model directly instead of going through the nn.Module __setattr__ fallback.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Resolved as of bda6ab6d0673: the high-priority or low-quality signal noted on an earlier commit no longer applies.

[amy-why-3459]

Thanks for the follow-ups; the constant move and the test seeding both look right. I only looked at what changed since fda6d41 (the restored CAM++ / S3 graph paths with the event chain). Three things below, the first one about where the chain is applied.

Separately, the squashed commit message now says both that the graphs are "removed here" and that they are restored; it would help to reconcile that (and the PR description) with the current code and the new A/B table.

[amy-why-3459]

The wait/record pair also runs on the eager fallback (bucket overflow, failed capture, off-graph), which never touches static_in/static_out. With 179 distinct reference lengths and a 2-bucket cap, most calls take that path, so every CAM++ call on one replica's private stream now waits on the device for the previous replica's CAM++ even when neither uses the shared buffers. That reintroduces the cross-replica serialization the private stream was meant to avoid. Only capture and replay write or read the static buffers, so the dependency is only needed there: wait on _graph_done right before static_in.copy_(batch) / capture, and record it after static_out.clone(). The eager branch could then also run outside _graph_lock. Same pattern in _quantize_s3_graph.

[y-null]

Agreed - fixed. The event chain now only covers the capture and replay paths, which are the only users of the shared static buffers. The eager fallback allocates fresh tensors and no longer waits/records, so the private-stream parallelism across replicas is back on the eager path (which is the majority of calls).

[y-null]

Fixed in 1bcc764 — the eager branch now returns before any event chaining (freshly allocated tensors only, no shared static buffers), so replicas running eager no longer wait on each other; the wait/record pairs only wrap the capture/replay branches that actually touch the static buffers.

[amy-why-3459]

"only the few most common frame counts replay" is not what the code does: the first _MAX_GRAPH_BUCKETS frame counts seen get captured, whatever their frequency. Since prompt features are cached per prompt_wav, only requests on those first two references can ever replay. Could you add the replay hit count (replays vs. eager calls per tier) to the A/B? With at most two of ~179 lengths captured, a 3–5% RTF gain at c4/c8 over the eager-only arm seems large for the graph path alone, and the hit rate would show whether it comes from replay or from run-to-run variance.

[y-null]

Fixed the comment to state the actual semantics (the first buckets seen get captured, not the most common), and added capture/replay counters that log the realized rate. Measured on the Seed-TTS benchmarks: both graphs reach 100+ replays with zero eager calls across the three tiers - the reference lengths concentrate on the two captured buckets, so the replay hit rate is 100%, not 2/179. The A/B gains hold on this build: c1 RTF 0.2877 (main 0.3350), c8 1.0940 (main 1.1819).

[y-null]

The comment was reworded in 1bcc764 to say the first _MAX_GRAPH_BUCKETS frame counts seen get captured (not 'most common'). Hit-rate numbers: see the reply above — the counters printed 100+ replays / 0 eager calls across the three-tier run.

[amy-why-3459]

The graph state (_s3_graphs, _s3_failed, _s3_graph_lock, _s3_graph_done) is created lazily here, outside any lock. If two threads hit the first call at the same time, each can create its own lock and dict, so the critical section isn't actually protected for that window. Initializing them in __init__ next to _spk_stream / _spk_event would avoid that and also lets _MAX_S3_GRAPHS, the lock and the event be declared in one place.

[y-null]

Addressed: the S3 graph state (_s3_graphs, _s3_failed, _s3_graph_lock, _s3_graph_done, counters) is now declared in __init__, so concurrent first calls share one lock instead of racing to create their own. The test's __new__-based construction seeds these fields explicitly.

[y-null]

Fixed in 1bcc764 — _s3_graph_lock / _s3_graphs / _s3_failed / _s3_graph_done and the counters are all declared together in __init__, so there is no lazy-instantiation race on first use.

[amy-why-3459]

Thanks, this addresses my last round: the eager fallback no longer touches the cross-replica event, the bucket comment now says first-seen, the S3 graph state is declared in __init__, and the replay/eager counters give the hit rate. I only looked at what changed since bda6ab6.

One regression from the restructure: the capture path now returns static_out.clone() right after torch.npu.graph(...) without replaying. Kernels recorded inside a graph-capture block are not executed, so static_out / static_code (and the captured code_len) still hold uninitialized graph-pool memory at that point. The warmup iterations ran on static_in but wrote their results to separate tensors. Before this change the capture branch fell through to the copy_ + replay() block, so the first request per bucket got a real result. Now the first request for each captured frame count would get a garbage speaker embedding / S3 tokens. A graph.replay() before the clone (or simply letting it fall through to the replay block as before) fixes it. The CPU tests don't reach this branch, so it won't show up there.

[amy-why-3459]

Ops inside with torch.npu.graph(graph): are only recorded, not run, so static_out has no valid data yet. Previously this fell through to static_in.copy_(batch); graph.replay(). Please replay once (e.g. graph.replay() here) before static_out.clone(), or keep the old fall-through.

[y-null]

Fixed in 31e9b20 — captured kernels are recorded but not executed, so the capture branch now replay()s once and clones the replayed static output; verified on 910B3 with WER-100 (100/100 completed, mean 0.0247 vs 0.0295 before, zero garbage outputs).

[amy-why-3459]

Same here: static_code and the captured code_len (entry[4]) are outputs of the capture block and aren't populated until the first graph.replay(), so the request that triggers capture would return uninitialized tokens/lengths.

[y-null]

Fixed in 31e9b20 the same way (replay once after capture, then clone static_code/code_len).

[amy-why-3459]

Nit: this still says "only the few most common frame counts replay"; like the CAM++ comment you updated, the cap actually keeps the first _MAX_S3_GRAPHS frame counts seen.

[y-null]

Fixed in 9778802 — the S3 comment now says the first _MAX_S3_GRAPHS frame counts seen get captured.

[amy-why-3459]

Checked the new commit 31e9b208. Replaying once right after capture is the right fix: the kernels recorded under torch.npu.graph don't execute, so the capture request was cloning whatever was in the graph pool. Counting that replay and recording _graph_done after it also keeps the event chain correct.

Two follow-ups:

• Since this was a silent wrong-output bug on the first request per bucket, a small NPU-gated test comparing the capture-call output against an eager call on the same input (CAM++ and S3 quantize) would keep it from coming back.
• The commit names this as the likely cause of the WER-100 gap (0.0295 vs 0.0122 baseline). Could you post the WER re-run on this head to confirm it is back at baseline?

[amy-why-3459]

self._graphs[frames] = entry is registered (L113) before this replay. If the replay or clone raises, the except adds frames to _graph_failed and falls back to eager, but the entry stays in _graphs, so every later request of that length goes straight to the replay path on the same graph. Registering the entry only after replay() + clone() succeed keeps the failure sticky as intended.

[y-null]

Fixed in 9778802 — the bucket entry is now published only after the first replay()+clone succeed, so a failed first replay leaves nothing registered and later same-length requests fall back to eager (the frame count lands in _graph_failed).

[amy-why-3459]

Same here: self._s3_graphs[frames] = entry (L425) is set before the first replay, so a failure in replay()/clone() marks the bucket failed but leaves the graph cached and used by later requests. Moving the registration below the clone would fix it.

[y-null]

Fixed in 9778802, same as the CAM++ site above.

[y-null]

Fixed in 31e9b20, thanks for catching this.

You're right: kernels captured under torch.npu.graph(...) are recorded but not executed, so static_out / static_code / code_len still held uninitialized pool memory when the capture branch cloned and returned them (the two warmup iterations write to temporary tensors and never populate the static outputs). Both call sites now graph.replay() once right after capture, before cloning, so the capture request produces exactly what the replay path produces; the replay counters and the record_event chaining now include this replay.

Verified on 910B3 (seed-tts en, disable-shuffle, whisper-tiny.en, 100 prompts, max-concurrency 1): completed=100/100, mean WER 0.0247 (was 0.0295 with the bug), 0/100 items with WER > 0.5 (no garbage first-request outputs), median 0.0.

[amy-why-3459]

Checked the delta 31e9b20..9778802. Moving the _graphs[frames] / _s3_graphs[frames] assignment after graph.replay() + clone fixes the point from the last round in both caches: if the first replay raises, the bucket stays unregistered, frames lands in _graph_failed / _s3_failed, and later same-length requests take the eager path instead of replaying a broken graph. Since failed captures no longer occupy a slot, the _MAX_* cap also only counts working graphs now. No new concerns in this delta.

One thing still open from before: there is no test for this failure path. A small CPU test that fakes torch.npu.NPUGraph with a replay() that raises on the first call, then asserts the bucket is absent, the frame count is marked failed, and the next call returns the eager result, would lock this behaviour in for both CAM++ and S3.

[amy-why-3459]

The new failure-path tests (both graph caches) now pin the registration-order fix on CPU CI, and the parity tests exercise the capture → mandatory-replay → clone path end to end — this covers the earlier ask well.

One coverage gap: both *_graph_npu.py parity tests are gated by @pytest.mark.npu + skipif-no-torch_npu, so the capture-returns-garbage regression (the r6 fix) and the bucket-publish-order fix only have a running guard on real NPU. On CPU CI they're skipped, so a future refactor could reintroduce the uninitialized-output path and still pass pre-merge CI. Worth confirming these are wired into the NPU nightly pipeline so the guard actually runs somewhere on every change.

[amy-why-3459]

These parity tests set allow_internal_format=False + jit_compile=False here to make capture work, with the comment that it's the same runtime contract as prepare_code2wav_graph_runtime. But the production CAM++ (speaker_embedding_torch.py) and S3 (step_audio2_token2wav.py) capture paths don't configure this runtime themselves — only prepare_code2wav_graph_runtime() does, and it's scoped to the Code2Wav stage. If the Token2Wav / prompt-feature stage loads without that global config in effect, the aclop conv kernels won't be capturable, so the first capture raises → the frame count lands in _graph_failed → silent eager fallback, and the setup speedup is quietly lost while every test still passes. Is the runtime guaranteed to be configured by the time CAM++/S3 capture in a standard deploy, or should these paths set/assert it on first capture?

[y-null]

You're right — good catch. The production capture paths never configured the runtime themselves; prepare_code2wav_graph_runtime covers only the Code2Wav stage, so if CAM++/S3 capture ran before it, the very first capture would throw, the bucket would silently fall back to eager, and the speedup would be lost while tests (which set the flags by hand) stayed green.

Fixed in 776292e, mirroring _prepare_moss_codec_graph_runtime:

• speaker_embedding_torch.py: _prepare_speaker_embedding_graph_runtime()
• step_audio2_token2wav.py: _prepare_s3_quantize_graph_runtime()

Both are called right before torch.npu.graph in their own capture branch — after the eager-fallback early return, so the flags are only set when capture is actually enabled — fail fast on ASCEND_LAUNCH_BLOCKING=1, then set allow_internal_format=False + jit_compile=False.

The NPU tests no longer set anything by hand — they now fail if the self-configuration regresses. While fixing this I also confirmed the failing CI step ("Validate marks"): the tests' bare pytest.mark.npu is rejected by tools/pre_commit/check_test_marks.py; they now use the hardware_test helper instead, and the CPU fakes gained a config namespace so the prepare step runs against them.

Verified on 910B3: NPU parity tests 3 passed (capture succeeds with no manual torch.npu.config tweak), CPU failure-path tests 8 passed.

[y-null]

All three follow-ups are done on the latest head (66dddd1, which includes the main merge):

1. NPU-gated capture/replay parity tests (CAM++ + S3) — tests/model_executor/models/cosyvoice3/test_speaker_embedding_graph_npu.py and tests/model_executor/models/step_audio2/test_step_audio2_s3_graph_npu.py. The first call exercises capture + the mandatory post-capture replay, the second the steady-state replay path; both must match a plain eager forward on the same input (random weights suffice), so the capture-returns-garbage regression cannot come back silently. Ran on 910B3: 3 passed. One note: the tests must set torch.npu.config.allow_internal_format = False + set_compile_mode(jit_compile=False) or the aclop-based conv kernels refuse to be captured — the same runtime contract _prepare_moss_codec_graph_runtime / prepare_code2wav_graph_runtime already document in-tree.

2. WER-100 re-run on the merged head — same protocol as the original PR validation (seed-tts en / disable-shuffle / whisper-tiny.en / max-concurrency 1): completed 100/100, mean WER 0.0247 / median 0.0, 0 entries with wer>0.5, 81/100 exactly correct — bit-identical to the 0.0247 measured at the r6 fix (31e9b20), so the main merge introduced no regression.

3. First-capture-failure CPU tests — tests/model_executor/models/cosyvoice3/test_speaker_embedding_capture_failure.py plus a new case in tests/model_executor/models/step_audio2/test_step_audio2_s3_quantize_eager.py. They fake torch.npu with a graph whose first replay() raises and assert: no bucket is published, the frame count lands in the failed set, and the next same-length call returns the eager result (both graph caches). 8 passed together with the existing CPU parity cases.

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 66dddd1a15e7 refers to an earlier head; the pull request now points at 776292e963d0.

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 776292e963d0 refers to an earlier head; the pull request now points at e70e22c478f0.

[amy-why-3459]

Thanks for r8 — this directly addresses the earlier ask about CAM++ / S3 NPUGraph capture not configuring the ACLNN runtime.

What looks good on 776292e9

• _prepare_speaker_embedding_graph_runtime / _prepare_s3_quantize_graph_runtime correctly mirror prepare_code2wav_graph_runtime / _prepare_moss_codec_graph_runtime: ASCEND_LAUNCH_BLOCKING guard, allow_internal_format=False, jit_compile=False, and they run immediately before capture (not relying on another stage).
• Placing prepare outside the capture try is the right fail-fast: a blocking env / config error should not be swallowed into _graph_failed eager fallback.
• NPU tests dropping the duplicated manual config and relying on @hardware_test + the production prepare path is cleaner; CPU fakes gaining a config stub matches that.

CI / process

• Head commit 776292e9 is missing a Signed-off-by: trailer (DCO is action_required), and pre-commit is currently red on this SHA. Please amend/rebase with DCO and clear pre-commit before merge.

Mild notes (non-blocking)

• enable_trt_spk_embedding / enable_device_spk_embedding dropped cpu_reference / use_trt kwargs while get_campplus_trt(..., cpu_reference=...) still accepts it. In-tree callers look fine with the no-arg form; just flagging if any out-of-tree scripts still passed those kwargs.

No other correctness concerns on the delta from 66dddd1a → 776292e9.

[amy-why-3459]

Nice — prepare right before capture matches the Code2Wav / moss pattern and avoids the silent eager-fallback risk from the prior review. Leaving it outside the try so ASCEND_LAUNCH_BLOCKING / config errors fail fast (instead of landing in _graph_failed) looks intentional and correct.

[y-null]

Yes, outside the try is intentional: a bad config should abort startup loudly rather than get recorded in _graph_failed and silently degrade every later request on that bucket length to eager.

[amy-why-3459]

Same pattern as CAM++ prepare — good that S3 quantize capture self-configures rather than depending on another stage having flipped allow_internal_format / jit_compile.

[y-null]

Right — both capture sites are now self-contained under the same contract as prepare_code2wav_graph_runtime, so order no longer matters regardless of which stage captures first.

[amy-why-3459]

Checked the tip 042919dd (delta from 3f6780bd / earlier note on 776292e9).

Looks good

• enable_trt_spk_embedding(*, cpu_reference=False) and enable_device_spk_embedding(*, use_trt=True) are restored and forwarded correctly (cpu_reference into get_campplus_trt; CUDA path gated on use_trt). That matches the merge-introduced MiniCPM-o / CosyVoice3 call sites and fixes the unexpected keyword argument 'use_trt' startup break from the r8 overwrite.
• Docstrings match the restored kwargs; NPU capture / _prepare_*_graph_runtime path is untouched, as intended.
• Tip carries a Signed-off-by: trailer (DCO clear on this SHA).

No further correctness concerns on this incremental change.