[maithilijoshi20]

Fixes #8756.

Purpose

What broke

When a streaming AR diffusion pipeline enabled supports_chunk_step_grouping, one worker invocation ran every denoise step through the chunk boundary. An interaction submitted while that invocation was running could not be processed until after the runner had already applied the boundary and prepared the next chunk, so the interaction became active one chunk late.

Reproduction

The multi-step regression is covered by a deterministic runner-level test:

python -m pytest -q tests/diffusion/test_diffusion_step_pipeline.py \
  -k ar_runner_applies_interaction_queued_after_grouped_work_at_the_next_boundary

Before this fix, the grouped call reaches the boundary immediately and the test cannot queue the interaction before boundary application.

Root cause

ARDiffusionModelRunner._execute_stepwise_core grouped the final step of the current chunk. The base runner applies pending interactions and calls prepare_next_chunk as part of that final step, before control returns to the scheduler/RPC loop.

Fix and scope

Stop grouped execution after reaching the penultimate step. The boundary-finishing step runs in a fresh scheduler invocation, allowing pending interaction RPCs to be handled first.

This latency improvement applies only to chunks with at least two denoise steps. A one-step chunk has no pre-boundary yield point and retains the baseline one-call interaction latency. A dedicated regression test documents that limitation.

Tests cover boundary yielding, interaction timing, the one-step scope restriction, cached continuations, grouped failures, grouped call patterns for two- and three-step chunks, and output equivalence.

Test Plan

vLLM Version: 0.31.0

vLLM-Omni Base: 233f08421

• Run the complete stepwise runner test module.
• Run the adjacent streaming interaction test module.
• Run Ruff lint, Ruff format, and whitespace checks.
• Verify CUDA tensor execution on 1x NVIDIA GeForce RTX 3050 OEM (8 GiB) under WSL2.

Test Result

Local WSL2

• 91 passed, 3 deselected across the expanded stepwise runner and streaming interaction suites (-k "not parallel").
• Ruff check, Ruff format, and changed-file whitespace checks passed.
• CUDA smoke test passed on the RTX 3050.
• Changed-file pre-commit hooks passed except mypy-3.10, which reports existing errors on unchanged lines in both touched files (for example runner lines 142/366 and test lines 272/1157 on upstream/main). No new mypy error is introduced by this diff.

A800 reproduction (Windyskr)

• On one NVIDIA A800, the original [Bug]: Chunk step grouping applies a mid-chunk interaction one chunk late #8756 WebSocket reproduction activates the prompt update at generation_chunk=1 with this PR; the grouped base first activates it at generation_chunk=2.
• Streamed latents are bitwise equal to request-mode outputs across grouped and ungrouped base/PR runs. Without an update, WebSocket chunk bytes are identical across all four runs.
• Grouped throughput shows no regression: median end-to-end time is 9.89 s for this PR versus 9.96 s for the base (five runs after one warmup); first-chunk latency is 1.13 s for both.
• The A800 validation passed all 92 relevant tests, including the three two-GPU stepwise tests, plus 554 broader diffusion tests. These runs stacked [Model] Add Self-Forcing (causal Wan2.1-T2V-1.3B, 4-step DMD) with stepwise streaming #8754 on both the PR and its base to exercise Self-Forcing.

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

[vllm-omni-review-bot]

This PR matches CODEOWNERS paths: /tests/.

Code owners: @NickCao @yenuo26

Routing: @NickCao via CODEOWNERS; @yenuo26 via CODEOWNERS

@maithilijoshi20, 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.

[maithilijoshi20]

Hey @Windyskr , could you validate #8810 with your original Self-Forcing/A800 reproduction and confirm the interaction becomes active at chunk 1 without output or throughput regressions?

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 926b2ff24226 produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[Windyskr]

@maithilijoshi20 thanks for picking this up. I validated #8810 with the #8756 reproduction on one A800 (SM80). #8754 (Self-Forcing) was applied on top of this PR (54f4ef7c) and, for comparison, on top of its base 233f0842. As in the issue, the grouped runs add supports_chunk_step_grouping = True to SelfForcingPipeline; this is a test-only change, since #8754 does not declare it.

Interaction latency. Setup: WS /v1/realtime/video, 832x480, 81 frames, seed 0, eager. The prompt update is sent at video.start; two runs per row.

Tree	Grouping	Update first active at	Chunk 1 vs. the run without the update
base 233f0842	on	generation_chunk=2	byte-identical
#8810	on	generation_chunk=1	differs
#8810	off	generation_chunk=1	differs
base 233f0842	off	generation_chunk=1	differs

Output. Streamed latents (AsyncOmni, output_type="latent", 7 chunks) are bitwise equal to request mode on all four trees. The WS chunk bytes of the runs without an update are identical across all four.

Throughput, with grouping on, which is the path this PR changes. End of the 81-frame WS stream without an update, median of 5 runs after 1 warmup, same GPU:

• base: 9.96 s (range 9.91-10.09)
• [Bugfix] Fix grouped chunk interaction latency #8810: 9.89 s (range 9.88-9.94)

The first chunk arrives at 1.13 s in both. I left out timings for the ungrouped runs: the PR does not change that path, and another job started on the shared GPU while they ran.

Tests on this PR + #8754:

• tests/diffusion/test_diffusion_step_pipeline.py tests/diffusion/streaming_interactive_generation: 92 passed. This includes the three 2-GPU stepwise tests (Ulysses, ring, CFG) that were skipped on your machine.
• tests/diffusion/models/self_forcing tests/diffusion/models/lingbot_world tests/diffusion/ar_diffusion -m "not parallel": 554 passed.

Not run: LingBot World itself (no weights here). It goes through the same runner path.

One optional suggestion. test_ar_runner_leaves_the_chunk_boundary_step_for_a_fresh_call uses the two-step stub. There, grouped and ungrouped calls are now identical ([1, 2, 3, 4]), so the test no longer shows that grouping still merges the steps before the boundary. With _ThreeStepChunkPipeline, the grouped pattern is [2, 3, 5, 6] and the ungrouped one [1, 2, 3, 4, 5, 6], with the same emitted chunks (I checked both). Parametrizing the test over both stubs would cover that.

From the reporter's side, this fixes #8756 in the original reproduction.

AI assistance: the validation runs and this comment were done with Claude Code.

[Windyskr]

@maithilijoshi20 a small follow-up to my comment above. With the boundary step left for a fresh call, the grouping tests whose stubs use two-step chunks no longer reach the grouped continuation. For example, dropping a continuation step's output in the grouped loop still passes the whole module. I opened maithilijoshi20#1 against your branch. It:

• switches those tests to _ThreeStepChunkPipeline;
• parametrizes the call-pattern test, as in the suggestion above;
• aligns the _execute_stepwise_core docstring.

It touches tests and one docstring only (91 passed). Feel free to merge it, cherry-pick it or ignore it.

AI assistance: drafted with Claude Code.

[kunkunblueberry]

The fix direction and implementation look correct: leaving the chunk-boundary step for the next scheduler call should remove the one-chunk interaction delay, while preserving the one-step-chunk limitation.

One minor test-quality concern: test_ar_runner_applies_interaction_queued_after_grouped_work_at_the_next_boundary does not exercise the real submit_interaction/interaction-coordinator path. It replaces _interaction_coordinator with a MagicMock and directly assigns pipeline.pending_interaction, so it verifies the boundary hook but not that an interaction RPC is actually queued and propagated through the production path. A small follow-up using the real submission API or a minimal coordinator fake would make this regression test stronger.

[Bounty-hunter]

@tzhouam PTAL

[Gaohan123]

LGTM. Thanks. And please submit a PR to v0.32.0

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 44e0ea2bb706 refers to an earlier head; the pull request now points at 7ee1dc8b42c8.