Skip to content

[Core] Fail recompute preemption on AR stages that cannot replay their decode state - #9249

Merged
Gaohan123 merged 9 commits into
vllm-project:mainfrom
amy-why-3459:fix/recompute-preemption-fail
Oct 24, 2026
Merged

Gaohan123 merged 9 commits into
vllm-project:mainfrom
amy-why-3459:fix/recompute-preemption-fail

Conversation

@amy-why-3459

@amy-why-3459 amy-why-3459 commented Oct 22, 2026 •

Copy link
Copy Markdown
Collaborator

Contributor intent (required human answers)

Fix #6179

1. What outcome do you want to achieve?

Make preemption safe for every stage whose decode step carries more state than
the sampled token. When the scheduler preempts such a request, the caller must
see an explicit truncation instead of audio that silently splices or restarts.

Today it does neither. _preempt_request() zeroes num_computed_tokens, so the
replay is scheduled as a prefill of prompt + already-generated frames, but the
model rebuilds only the prompt part. Each affected stage then fails differently:

  • Qwen3-Omni talker: _preprocess copies
    min(span_len, req_embeds.shape[0]) rows into a torch.empty buffer
    (vllm_omni/worker/gpu_model_runner.py); the rest reach attention
    uninitialised. Nothing raises, because the layer-0 codec id stays inside the
    codec vocab.
  • VoxCPM2: audio token ids are embedded through the text table and fault
    (EngineDeadError, as reported in [Bug]: recompute preemption restarts native-AR TTS generation from frame 0, splicing unrelated audio #6179).
  • Gepard: replay restarts at frame 0; the caller gets two unrelated halves.
  • MTP-style stages (audio8_tts, fish_speech, qwen3_tts, moss_tts,
    personaplex): the prefill path pads the unbuilt rows with a pad
    token embedding.

2. What is the highest-leverage variable for that outcome?

Whether a stage can declare, at the stage level, that recompute preemption is
unsafe for it -- and whether the scheduler honours that declaration.

Per-model warnings (the shape #5666 added to Gepard at gepard_talker.py)
only describe the damage after it happens, and they cover one model at a time.
The variable that decides the outcome is a declaration on StagePipelineConfig
plus one branch in the scheduler's preemption path: every affected model then
opts in with one line, and every model that genuinely can replay keeps today's
behavior untouched.

3. How will this work move that variable?

  • StagePipelineConfig.recompute_preemption: Literal["allow", "fail"] defaults
    to "allow", so no existing pipeline changes behavior.
  • Plumbed through OmniStageModelConfig (AR stages only) into OmniModelConfig.
    Diffusion stages have no KV-recompute preemption and reject unknown fields, so
    the field is not forwarded to them.
  • OmniSchedulerMixin._preempt_request reads it once at init. Under "fail" it
    frees the request's resources and emits exactly one terminal ERROR instead of
    requeueing, so the failure is visible at the API boundary.
  • Eleven stages declare "fail": Qwen3-Omni talker, VoxCPM2 latent_generator,
    Gepard, audio8_tts, fish_speech, Qwen3-TTS (both topologies),
    personaplex, and all three MOSS-TTS topologies.
  • Stages that keep "allow": every thinker, and Qwen2.5-Omni's talker -- it
    keeps all codebook layers in the token stream, so its replay is faithful.
    Fish Omni also keeps allow: its MRv2 model state saves generated codebooks
    on preemption and reconstructs the full frame rows on resume.
  • An explicit reset_prefix_cache scope keeps upstream's intentional resume
    path (subject to supports_running_prefix_cache_reset). Capacity preemption
    still fails when KV delivery sets drop_stale_output=True.
  • Both config projection paths forward the capability only for AR stages.
  • The terminal error describes truncated output without assuming audio.

4. Whose customer or user needs does this address, and what are those needs?

Users serving speech models under memory pressure. When the KV cache fills, the
scheduler preempts, and these users currently get spliced audio, silently
corrupted frames, or a dead stage replica -- with nothing in the response saying
so. They need a preemption they can see and act on (raise
gpu_memory_utilization, lower max_num_seqs), not a stream that quietly
contains garbage.

Changes and review coverage

Capacity preemption of a non-replayable AR stage now emits one terminal ERROR,
releases request/KV state, and never requeues it. The reset exception is scoped
with try/finally around the explicit reset operation; drop_stale_output is
not a reset discriminator. Default allow stages retain upstream behavior.

The scheduler regression uses the production OmniARScheduler constructor,
schedule(), finish_requests(), KV allocator and update_from_output() /
terminal assembly. A real two-block pool (one null block, one usable block)
fits a 16-token prefill and cannot allocate its first decode token. Connector
I/O and runner output are synthetic; scheduler lifecycle methods are not mocked.

Coverage includes both KV-delivery settings, client routing, exactly one ERROR
(no later duplicate ABORT), no requeue, request removal, KV block reclamation,
explicit reset/resume, reset exception cleanup, and subsequent capacity failure
after reset. Config tests cover allow/fail across all execution types.

Updated docs/design/feature/prefix_caching.md to document the fail-mode and
explicit-reset exception. No new deployment knob or user API is introduced.
Related: #6179.

Rebase validation

Rebased onto a147aaf27e10ea4e68b2cc257e3820ac5e7358d7; rebase commit:
4467d555fa5b3d254cc0930db4d83eb5af9e7d63; final head: 6da4d1bc6845f975a17027c4cbb86b39f66886d9. Resolved the Fish Omni test-file
conflict by preserving both main's NPU device-selection / TP noise-broadcast
tests and this PR's capacity-preemption / slot-reuse / resume regression.
git range-diff confirms the four PR commits retain their changes; only the
insertion context of the appended Fish Omni test changed.

source /data/why/scripts/env.sh
python -m pytest \
  tests/core/sched \
  tests/model_executor/models/fish_omni/test_fish_omni_model_state.py \
  tests/model_executor/models/test_recompute_preemption_declarations.py \
  tests/model_executor/models/qwen3_omni/test_qwen3_omni_preempt_replay.py \
  tests/config/test_omni_config.py \
  -q -p no:xdist -o addopts='' \
  -k 'not test_pending_generation_capability_reaches_both_engine_config_paths'
pre-commit run --from-ref a147aaf27e10ea4e68b2cc257e3820ac5e7358d7 --to-ref HEAD

At the rebase commit: 789 passed, 1 skipped, 2 deselected. Pre-commit
then exposed two mypy errors introduced by main's optional-platform lookup in
_apply_platform_overrides. The final head adds an explicit non-None guard
before string operations. After that change, the structured config tests passed
(315 passed, 1 skipped, 2 deselected) and the platform override suite passed
(61 passed). All applicable pre-commit hooks passed.

The additional platform suite command was:

source /data/why/scripts/env.sh
python -m pytest tests/config/test_config_factory.py::TestPlatformOverrides \
  -q -p no:xdist -o addopts=''
``` The two excluded configuration cases retain the previously
recorded offline-model-cache limitation. No GPU/NPU end-to-end inference or
full repository UT was run. Earlier validation below is historical and names
its original tested commits.

## Latest review fixes

- Preserve Fish Omni's replay-capable MRv2 path by removing the blanket `fail`
  declaration. The regression takes the policy from `FISH_OMNI_PIPELINE`, forces
  capacity preemption through the production scheduler, saves codebooks,
  overwrites the old slot with another request, resumes into a different slot,
  and verifies complete restored frame rows/history plus non-terminal output.
- Preserve upstream auxiliary-output cleanup before removing a fail-mode victim.
  The exhausted-pool test now enables the actual `AuxOutputSchedulerConnector`
  and verifies one ERROR, KV reclamation, removal from its active hash tracking,
  one finished notification and a clean subsequent step, with both KV-delivery
  settings.
- Document these contracts in `docs/design/feature/prefix_caching.md` and move
  Fish Omni from the fail declarations to the replayable controls (11 fail stages).

### Latest validation

Fix commit: `813668f9eb3404da238a7a6ac3d82561c11a674a` (parent `0a75edf8b12d0b155a3b77d64c8a397bfa9d1a31`).
Same environment as below. New regressions were first run against the parent
with only the test changes present:

```bash
source /data/why/scripts/env.sh
python -m pytest \
  tests/core/sched/test_omni_recompute_preemption.py \
  tests/model_executor/models/fish_omni/test_fish_omni_model_state.py \
  -q -p no:xdist -o addopts='' -k capacity_preemption

Before: 3 failed, 4 passed. Both auxiliary-enabled fail cases raised
KeyError: 'req-preempt' in build_connector_meta; the Fish Omni case was
FINISHED_ERROR instead of PREEMPTED. After: all corresponding cases pass;
the complete preemption, Fish model-state and declaration files give 42 passed.

Expanded scheduler validation on the final code:

source /data/why/scripts/env.sh
python -m pytest \
  tests/core/sched \
  tests/model_executor/models/fish_omni/test_fish_omni_model_state.py \
  tests/model_executor/models/test_recompute_preemption_declarations.py \
  tests/model_executor/models/qwen3_omni/test_qwen3_omni_preempt_replay.py \
  -q -p no:xdist -o addopts=''

471 passed. All applicable pre-commit hooks passed, including mypy, ruff, SPDX, test marks and markdownlint.

The Fish regression executes its actual model-state save/restore on CPU with
synthetic prompt/model buffers and runner outputs. It does not execute GPU
inference, GPU-to-host event timing, Mamba kernels, or end-to-end audio decoding.
Full repository UT and named-model GPU E2E were not run. Earlier configuration
results and their offline-cache limitations are recorded below against their
original tested revision.

Previous validation (head 0a75edf8b1)

Rebased onto 59fedaf594bc4c13442f37668f2c7f33e8db1892; fix head: 0a75edf8b12d0b155a3b77d64c8a397bfa9d1a31.
Environment: Python 3.12.13, vLLM 0.32.0, PyTorch 2.13.0+cu130,
Transformers 5.14.1. CPU scheduler tests; no model weights or GPU inference.

The same fix-head test file was copied into isolated checkouts for the two
older revisions. Run from each checkout:

source /data/why/scripts/env.sh
PYTHONPATH=. python -m pytest \
  tests/core/sched/test_omni_recompute_preemption.py \
  -q -p no:xdist -o addopts='' -k capacity_preemption

To reproduce against an older checkout, copy the test from the fix first:

git show 0a75edf8b12d0b155a3b77d64c8a397bfa9d1a31:tests/core/sched/test_omni_recompute_preemption.py > \
  tests/core/sched/test_omni_recompute_preemption.py
Revision Result Observed behavior
Requested baseline 717f8a3499fb7dc6b3128b58a8baeacded1a7bbb 2 failed, 2 passed Both fail-mode cases remain PREEMPTED instead of FINISHED_ERROR; allow controls pass.
Reviewed PR head 84464c53381a05473b610c4e8aeba68b2c5bbce4 1 failed, 3 passed KV-delivery fail-mode still requeues; ordinary fail-mode and allow controls pass.
Fix 0a75edf8b12d0b155a3b77d64c8a397bfa9d1a31 4 passed Both fail-mode cases emit one ERROR and release state; allow controls pass.

Expanded validation:

source /data/why/scripts/env.sh
python -m pytest \
  tests/core/sched \
  tests/model_executor/models/test_recompute_preemption_declarations.py \
  tests/model_executor/models/qwen3_omni/test_qwen3_omni_preempt_replay.py \
  tests/config/test_omni_config.py \
  -q -p no:xdist -o addopts='' \
  -k 'not test_pending_generation_capability_reaches_both_engine_config_paths'

765 passed, 1 skipped, 2 deselected on the clean fix head. All seven production-scheduler preemption/reset tests and
six config projection cases pass. The two deselected existing
test_pending_generation_capability_reaches_both_engine_config_paths cases
were also attempted: both fail because offline mode cannot resolve the
test-model snapshot. The same two failures were reproduced on baseline
717f8a3499fb7dc6b3128b58a8baeacded1a7bbb.

All applicable pre-commit hooks passed on the final head, including mypy 3.10,
ruff, SPDX, test marks, forbidden imports and markdownlint:

source /data/why/scripts/env.sh
pre-commit run --from-ref 59fedaf594bc4c13442f37668f2c7f33e8db1892 --to-ref HEAD

Not run: full repository UT, GPU/model end-to-end pressure testing, or real
network KV delivery. These results establish the scheduler lifecycle contract;
they do not claim reproduced audio corruption or corrected audio quality on
every declared model. The installed editable Omni package also reports a
0.31/0.32 version-metadata mismatch warning with vLLM 0.32.0.

@chatgpt-codex-connector

Copy link
Copy Markdown

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-22T09:24:51.548208Z
We couldn’t complete this review. Comment @codex review to try again.
b82be16 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

Copy link
Copy Markdown

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.

@vllm-omni-review-bot

Copy link
Copy Markdown

This PR appears to belong to: docs/design/module/vllm_omni_config.md.

Module owners: @lishunyang12 @alex-jw-brooks @tzhouam

Routing: @lishunyang12 via module of the changed files, CODEOWNERS; @alex-jw-brooks via module of the changed files; @tzhouam via CODEOWNERS

@amy-why-3459, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

@vllm-omni-review-bot

vllm-omni-review-bot commented Oct 22, 2026 •

Copy link
Copy Markdown

Omni ReviewBot triage note

Automated triage of commit a7012ea6779b produced:

  • Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

@amy-why-3459 amy-why-3459 changed the title [Core][Qwen3-Omni] Fail recompute preemption on stages that cannot replay [Core] Fail recompute preemption on AR stages that cannot replay their decode state Oct 22, 2026

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@amy-why-3459, the claimed fail-fast/cleanup behavior still needs a regression through the production scheduler boundary on 84464c53381a05473b610c4e8aeba68b2c5bbce4:

  • Force capacity preemption through schedule() and actual finish/terminal-output assembly, rather than mocking finish_requests and calling _preempt_request directly. Acceptance: a non-replayable request emits exactly one terminal ERROR, is not requeued, and releases its request/KV state; preserve intentional reset behavior.
  • Include the KV-delivery case. The new guard treats drop_stale_output=True as a reset bypass, but pinned vLLM 0.32.0 capacity preemption also passes that flag when requires_kv_delivery is true. Please distinguish the two reasons explicitly or demonstrate and test why this topology is excluded from the fail-mode contract.
  • Resolve the current merge conflict, then supply the same reproducer/test commands and results against baseline 717f8a3499fb7dc6b3128b58a8baeacded1a7bbb and the fix head, showing the original failure and corrected terminal behavior.

This is a specific contract/evidence gap: I have not established a failing named-model deployment. Source assessment covered the diff and scheduler/cleanup boundary, not all model replay implementations; no tests were executed by me.

@amy-why-3459

Copy link
Copy Markdown
Collaborator Author

@amy-why-3459

@amy-why-3459 amy-why-3459 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed 84464c53381a (self-@). The stage declaration + fail-path design is the right shape for #6179, and the pending-output reuse of _pending_input_timeout_outputs / _attach_finished_request_sets matches the existing timeout contract.

Two contract gaps still block landing this as-is (one already called out on this head, verified against pinned vLLM 0.32):

  1. drop_stale_output is not a reliable "reset only" signal. Capacity preemption in upstream Scheduler.schedule passes drop_stale_output=self.requires_kv_delivery (vllm@5df63a72 L~833). With any connector that sets requires_kv_delivery, the new guard falls through to super()._preempt_request and requeues — so fail stages still splice under KV-delivery topologies. Distinguish reset vs capacity (or document/test why these pipelines never enable that connector).
  2. Plumbing asymmetry. _build_model_config gates on LLM_AR, but _build_engine_args forwards whenever != "allow". A mistaken fail on a diffusion stage would hit EngineArgs / unknown-field rejection on one path and be silently dropped on the other. Mirror the AR gate (or reject at StagePipelineConfig validation).

Also: the terminal reason hard-codes "the audio generated so far is truncated" — fine for talker/codec stages, but odd if a non-audio AR stage later opts in. Prefer a modality-neutral wording.

Coverage: unit tests mock finish_requests and call _preempt_request directly; still need one path through schedule() → finish → _attach_finished_request_sets that asserts a single ERROR and no requeue (including the KV-delivery case above). Merge conflict vs current main also still open.

# ``drop_stale_output=True`` also marks reset_prefix_cache's own
# running-request reset, which intends to resume. Only a capacity
# preemption is the unrecoverable case.
if not getattr(self, "recompute_preemption_fails", False) or drop_stale_output:

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

drop_stale_output=True is also set for capacity preemption when self.requires_kv_delivery is true (upstream Scheduler.schedule passes drop_stale_output=self.requires_kv_delivery). Treating every True as a reset-prefix-cache bypass means fail-mode never fires on those topologies — the request is requeued and splices. Please gate on the reset path specifically (or prove these stages never run with a delivery-requiring connector) and cover it with a test.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 0a75edf. The bypass now checks an explicit reset_prefix_cache scope, cleared with try/finally, rather than drop_stale_output. Capacity preemption therefore takes the fail path even when requires_kv_delivery=True; intentional reset/resume is preserved.

The regression now runs the production OmniARScheduler schedule() and update_from_output() paths with a real exhausted KV block pool, without mocking finish_requests or allocation. It checks exactly one terminal ERROR, no requeue, request removal and KV reclamation with and without KV delivery, plus reset/resume and reset-scope cleanup. All seven preemption/reset tests passed. The same capacity tests fail both fail-mode cases on baseline 717f8a3, only the KV-delivery case on prior head 84464c5, and pass on the fix head. Commands and results are in the PR description. Connector I/O and runner output are synthetic; this is scheduler lifecycle coverage, not real-network or named-model GPU E2E validation.

reason = (
"Request was preempted and this stage cannot replay its decode state; "
"the audio generated so far is truncated. Raise the stage's "
"gpu_memory_utilization or lower max_num_seqs to avoid preemption."

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"the audio generated so far is truncated" assumes an audio stage. The same path covers any AR stage that opts into fail. Prefer something modality-neutral, e.g. "output generated so far is truncated / cannot be replayed".

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 0a75edf. The terminal reason now says "the output generated so far is truncated", so it does not assume an audio stage. The production-scheduler regression verifies that the terminal ERROR carries the replay/truncation reason.

Comment thread vllm_omni/config/stage_config.py Outdated
engine_args["requires_full_payload_input"] = ps.requires_full_payload_input
if not ps.supports_running_prefix_cache_reset:
engine_args["supports_running_prefix_cache_reset"] = False
if ps.recompute_preemption != "allow":

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_build_model_config only forwards recompute_preemption for LLM_AR, but this path forwards for any stage with != "allow". A diffusion stage that accidentally sets fail would then hit EngineArgs / unknown-field rejection here while the model-config path silently drops it. Please mirror the AR gate (or validate at StagePipelineConfig).

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 0a75edf. _build_engine_args now uses the same StageExecutionType.LLM_AR gate as _build_model_config. Added test_recompute_preemption_projection_is_ar_only, parametrized over allow/fail and all three execution types, to check both projections and ensure non-AR legacy arguments omit the field. All six cases passed.

@amy-why-3459 amy-why-3459 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed delta 0a75edf8 (vs prior 84464c53). The three contract gaps from the last pass look closed.

What this commit fixes

  1. Reset vs capacity. _resetting_prefix_cache scoped around reset_prefix_cache (try/finally) is the right discriminator. Capacity preemption — including when a KV-delivery connector sets drop_stale_output=True — now takes the fail path under recompute_preemption="fail", while explicit reset still resumes. That matches pinned vLLM 0.32’s dual use of the flag.

  2. AR-only plumbing. _build_engine_args now mirrors _build_model_config (LLM_AR gate). The new test_recompute_preemption_projection_is_ar_only locks both projection paths.

  3. Production-scheduler coverage. test_omni_recompute_preemption.py no longer stubs finish_requests / calls _preempt_request alone. Prefill → schedule() → update_from_output with a two-block pool, parametrized requires_kv_delivery, one ERROR + KV free, allow-mode requeue, reset resume, and flag cleared on reset error cover the lifecycle contract. Reason text is modality-neutral (“output … truncated”).

Residual notes (non-blocking)

  • Doc + body validation table against baseline 717f8a34 / prior head 84464c53 / fix head is clear; still no named-model end-to-end under real memory pressure (explicitly out of scope here).
  • Mergeable vs current main; reviews still needed from the requested set.

Coverage this pass: mixin preemption/reset, both config projection paths, and the rewritten scheduler tests. Looks ready from the contract side.

@amy-why-3459

Copy link
Copy Markdown
Collaborator Author

@natureofnature PTAL

@hsliuustc0106 hsliuustc0106 added the high priority high priority issue, needs to be done asap label Oct 23, 2026 — with ChatGPT Codex Connector

@natureofnature natureofnature left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The reset/capacity distinction, AR-only config projection, and production-scheduler regression address the earlier concerns. Two remaining issues are noted inline: Fish Omni's existing replay path is disabled, and fail-mode preemption skips auxiliary-output cleanup. I inspected the current scheduler, model-state and downstream error paths; I did not execute tests.

【Conclusion: Found blocking code issues in commit 0a75edf.】
【Action recommended: Preserve Fish Omni's replayable CUDA path and auxiliary-output cleanup; verify Fish Omni resumes after preemption and auxiliary-enabled fail-mode emits one ERROR without crashing before merging.】

# The backbone input sums the codebook embeddings of every sampled
# code but codebook 0 (`codebook_sum`, fish_omni_model.py); a replay
# from token ids rebuilds only the text part (#6179).
recompute_preemption="fail",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Fish Omni's default CUDA MRv2 path already preserves the missing codebooks: on_request_preempted saves them, and add_request rebuilds the full generated-frame rows on resume. This blanket fail terminates those requests under KV pressure instead of using that supported continuation. Could we retain allow for the replay-capable path and cover the actual pipeline policy with a preemption/slot-reuse/resume regression, checking saved frames are restored without a terminal ERROR?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in the current rebased head 6da4d1b. Fish Omni retains the default allow policy. The regression takes the policy from FISH_OMNI_PIPELINE, forces real scheduler capacity preemption, saves the generated codebooks, overwrites the original slot with another request and resumes into a different slot. It verifies full frame rows/history are restored and output remains non-terminal. The declaration controls now also require Fish Omni to stay allow.

The rebased scheduler/model/config run passed 789 tests (1 skipped, 2 known offline-cache cases deselected), including this regression and main's new Fish Omni NPU selection / TP tests. This exercises the model-state hooks on CPU with synthetic buffers; GPU Mamba, device-copy event timing and end-to-end audio were not tested.

# The caller already popped this request from ``running``; freeing here
# is what upstream's requeue path would otherwise do later.
self.finish_requests([request.request_id], RequestStatus.FINISHED_ERROR)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] With routed-expert auxiliary outputs enabled, this bypasses upstream preemption's aux_output_connector.request_finished(request). Omni's _free_request does not call that hook, so the victim remains in _sent_hash_counts after removal from self.requests; the same scheduling pass accesses the deleted request in build_connector_meta and raises KeyError before delivering the ERROR. Could we preserve auxiliary cleanup and extend the exhausted-pool regression with the actual auxiliary connector enabled, checking one ERROR, reclaimed KV, its finished notification and a clean subsequent step?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in the current rebased head 6da4d1b. Fail-mode preemption calls aux_output_connector.request_finished(request) before finish_requests removes the request, preserving upstream auxiliary cleanup.

The regression enables the actual AuxOutputSchedulerConnector and checks one terminal ERROR, reclaimed KV, removal from _sent_hash_counts, exactly one finished notification and a clean subsequent step, with and without KV delivery. Before the fix both auxiliary cases reproduced KeyError in build_connector_meta; both pass after the fix, including on the rebased code. The rebased scheduler/model/config run passed 789 tests (1 skipped, 2 known offline-cache cases deselected). All applicable pre-commit checks pass on the final head.

@amy-why-3459 amy-why-3459 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed delta 813668f9 (vs prior tip 0a75edf8). This addresses the two P2s from the previous pass.

What this commit fixes

  1. Fish Omni stays allow. Dropping the blanket recompute_preemption="fail" on FISH_OMNI_PIPELINE is correct given MRv2’s on_request_preempted / add_request codebook save-restore. Declaration tests now put Fish Omni in _ALLOW, and test_pipeline_capacity_preemption_restores_frames_after_slot_reuse forces capacity preemption through the production scheduler, overwrites the old slot, resumes into a new one, and checks restored frame rows / history with a non-terminal output. That matches the intended contract.

  2. Auxiliary-output cleanup on fail-mode. Calling aux_output_connector.request_finished(request) before finish_requests closes the _sent_hash_counts / build_connector_meta KeyError window. The exhausted-pool regression is now parametrized on auxiliary and asserts hash-count removal, one finished notification, and a clean subsequent step with both KV-delivery settings.

Docs in prefix_caching.md cover both points.

Residual notes (non-blocking)

  • Tip is currently merge-conflicted vs main (mergeable_state=dirty); rebase/resolve before landing.
  • The Fish Omni regression drives scheduler + model-state hooks explicitly (CPU, synthetic buffers); it does not run GPU Mamba / end-to-end audio. That scope is clear in the PR body and is fine for the lifecycle contract.

Coverage this pass: mixin fail-path aux cleanup, Fish Omni declaration + slot-reuse resume test, and the auxiliary-enabled exhausted-pool cases. Looks good from the contract side once the conflict is cleared.

A talker decode step feeds PagedAttention the summed embedding of all eight
RVQ codec codes plus the thinker's text step, but only the layer-0 codec id
enters the token stream. The other seven codes are sampled by the code
predictor, and `last_talker_hidden` / `text_step` live in the per-request
intermediate buffer.

`_preempt_request()` zeroes `num_computed_tokens` and replays the prompt, so a
preempted talker cannot rebuild that KV. The replay is scheduled as a prefill
of prompt + already-generated frames, but `talker_preprocess_prefill` rebuilds
only the prompt part, and `_preprocess` copies
`min(span_len, req_embeds.shape[0])` rows into a `torch.empty` buffer -- the
remaining rows reach attention uninitialised. Nothing raises: the layer-0 id
stays inside the codec vocab, so unlike VoxCPM2 the replay does not fault on a
text-embedding lookup.

Add `recompute_preemption: Literal["allow", "fail"]` to `StagePipelineConfig`
(the stage-level capability vllm-project#6179 converged on) and declare `"fail"` on
Qwen3-Omni's talker. Under `"fail"` the scheduler reclaims the request's
resources and emits exactly one terminal error instead of requeueing it, so the
caller learns the audio is truncated rather than receiving output that silently
restarted at frame 0. `reset_prefix_cache`'s own running-request reset passes
`drop_stale_output=True` and keeps upstream's resume behavior.

Gepard and VoxCPM2 can opt in with the same declaration; their replay paths are
unchanged here.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Extends the stage capability to the remaining native-AR stages whose decode
step feeds PagedAttention embeddings that `output_token_ids` cannot rebuild:

- VoxCPM2: per-request `curr_embed_for_next` / `prev_feat_embed` plus CFM-VAE
  progress; replaying also embeds audio ids through the text table and faults.
- Gepard-1.0: 31 of 32 FSQ codes per frame ride a side channel. This replaces
  the warning at gepard_talker.py that only reported the damage afterwards.
- fish_omni: the backbone input sums every sampled codebook but codebook 0.
- audio8_tts, fish_speech, qwen3_tts, moss_tts (all three topologies),
  personaplex: `talker_mtp` folds the sampled codebooks and the per-frame
  temporal hidden into the next step's `inputs_embeds`.

Qwen2.5-Omni's talker stays `"allow"`: it keeps every codebook layer in the
token stream, so its replay is faithful. Thinkers stay `"allow"` for the same
reason -- plain text AR.

Tests: one cross-model declaration test parametrized over all twelve (pipeline,
stage) pairs plus the three that must keep upstream recompute, and the Qwen3-Omni
row-shortfall arithmetic kept separate from the declarations.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
@amy-why-3459
amy-why-3459 force-pushed the fix/recompute-preemption-fail branch from 813668f to 6da4d1b Compare October 23, 2026 08:57
@amy-why-3459

Copy link
Copy Markdown
Collaborator Author

Addressed the remaining merge-conflict note in the latest review: rebased onto main a147aaf and pushed 6da4d1b.

The Fish Omni test conflict is resolved with both main's NPU selection / TP broadcast tests and the PR's slot-reuse resume regression retained. Rebase validation: 789 passed, 1 skipped, 2 deselected. A small follow-up fixes two mypy errors from main by narrowing an optional platform name; after that change, structured config tests returned 315 passed / 1 skipped / 2 deselected, and platform-override tests returned 61 passed. All applicable pre-commit hooks passed. The two excluded tests have the previously documented offline model-cache limitation. No full-repository UT or GPU/NPU E2E claim is made. Commands and exact revisions are recorded in the PR description.

@amy-why-3459 amy-why-3459 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed rebase tip 6da4d1bc (delta since prior tip 813668f9). Empty-body reviews on this SHA were thread replies only.

Rebase / tip delta

  1. Fish Omni test conflict resolution. Tip keeps main’s NPU device-selection and TP noise-broadcast cases and appends test_pipeline_capacity_preemption_restores_frames_after_slot_reuse. That matches the intended merge: both sides preserved, PR coverage still exercises production-scheduler capacity preemption → slot overwrite → resume into a new slot with restored frame rows / history and a non-terminal output.

  2. Mypy optional-platform guard. _apply_platform_overrides now gates with if uses_v2 and platform is not None before the unsupported-platform lookup / .upper(). That is the right fix for main’s optional platform type; no behavioral change for a real platform string.

  3. Prior contracts still hold on this tip. Fish Omni stage stays default allow (no recompute_preemption="fail"). Fail-mode still calls aux_output_connector.request_finished(request) before finish_requests, and the exhausted-pool regression remains parametrized on auxiliary + KV-delivery.

Residual notes (non-blocking)

  • Conflict vs main looks cleared (mergeable); still needs other requested reviews before landing.
  • Validation remains CPU scheduler / config / Fish model-state scope; no named-model GPU E2E under real KV pressure (already disclosed in the body).

Coverage this pass: rebase tip vs 813668f9, Fish Omni test file coexistence, platform guard, and re-check of Fish allow + aux cleanup on 6da4d1bc. Looks good from the contract side.

@amy-why-3459 amy-why-3459 added the ready label to trigger buildkite CI label Oct 23, 2026
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
@amy-why-3459

Copy link
Copy Markdown
Collaborator Author

Fixed the Model Executor Test failure from Buildkite #17918 in 13f322d.

Full model-executor test collection registers a custom MiniCPM KV cache spec before the scheduler fixture runs. A nonempty registry skips vLLM's lazy built-in registration, leaving FullAttentionSpec without a manager. The fixture now calls register_all_kvcache_specs(config), matching EngineCore initialization and retaining the real scheduler/KV allocator coverage.

Reproduced the original assertion locally with full model-executor collection. After the fix:

pytest -q tests/model_executor tests/core/sched/test_omni_recompute_preemption.py -m 'core_model and cpu' -k 'test_fish_omni_model_state or test_omni_recompute_preemption'

29 passed, 8063 deselected. All applicable pre-commit hooks passed. The full CI shard was not rerun locally; remote CI still needs to validate the new commit.

@amy-why-3459 amy-why-3459 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed delta 13f322d8 (vs prior reviewed tip 6da4d1bc; intermediate 5416bdf was merge-from-main only).

What this commit fixes

Model Executor Test flake / registry poisoning. Full model-executor collection can register a custom MiniCPM KV cache spec before this fixture builds OmniARScheduler. A nonempty registry skips vLLM’s lazy built-in registration, so FullAttentionSpec has no manager and the shard fails. Calling register_all_kvcache_specs(config) before constructing the scheduler matches EngineCore init and keeps the real allocator path.

Local note in the PR comment (29 passed under the model-executor + preempt filter) matches the intended fix shape; remote CI still needs to green the tip.

Residual notes (non-blocking)

  • Prior contracts on this tip still hold: Fish Omni stays allow, fail-mode still does aux request_finished before finish_requests, reset vs capacity still discriminated by _resetting_prefix_cache.
  • Still waiting on other requested reviews / CI on 13f322d8 before land.

Coverage this pass: tip delta only (tests/core/sched/test_omni_recompute_preemption.py fixture). Looks good from the fixture/registry side.

config.cache_config.num_gpu_blocks = num_blocks
# Match EngineCore initialization before constructing the scheduler. Custom
# specs registered during test collection can suppress lazy built-in registration.
register_all_kvcache_specs(config)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good placement: registering before OmniARScheduler(...) matches EngineCore and avoids the nonempty-registry skip of lazy built-ins when MiniCPM (or other) specs were registered during collection. No further change needed here.

@Gaohan123 Gaohan123 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks. And please cherry pick it into branch v0.32.0

@Gaohan123 Gaohan123 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Oct 24, 2026
@Gaohan123
Gaohan123 enabled auto-merge (squash) October 24, 2026 09:01
@Gaohan123
Gaohan123 merged commit c42398c into vllm-project:main Oct 24, 2026
9 of 11 checks passed
tzhouam pushed a commit to tzhouam/vllm-omni that referenced this pull request Oct 24, 2026
…r decode state (vllm-project#9249)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
(cherry picked from commit c42398c)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

high priority high priority issue, needs to be done asap ready label to trigger buildkite CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: recompute preemption restarts native-AR TTS generation from frame 0, splicing unrelated audio

5 participants