[Bugfix] Fix lost async diffusion outputs in request-level batching

[SamitHuang]

Purpose

Request-level batching (step_execution=false) stalls: requests hang until
_ASYNC_OUTPUT_TIMEOUT fires, and batches fail because OUTPUT_READY messages
are lost or futures orphaned. This PR fixes four independent defects on that
path, all still present on main.

This started as a drain-async-outputs-before-sleep change (the original scope
of this PR). Drain-before-sleep is retained as defensive hardening, but
end-to-end debugging showed it is not what caused the stalls — the four
fixes below are. The PR has been rescoped accordingly.

1. Per-request tensor views ship their entire batch storage (root cause).

_pack_tensor_if_large sized its pack/inline decision on the view:

if val.nelement() * val.element_size() > _SHM_TENSOR_THRESHOLD:

Batch outputs are split into per-request views of one batch tensor. Each view
is under the 1 MB threshold, so it stayed inline — but pickle serializes a
view's whole underlying storage, so a 32-request batch put the full batch
tensor on the wire once per request (~1.3 GB observed). That overflows the
MessageQueue inline buffer (24 MB) onto the ZMQ XPUB/SUB overflow path, and
PUB sockets silently drop messages under back-pressure — so OUTPUT_READY
vanished and every request in the batch timed out together.

The fix sizes the decision on max(view_bytes, storage_bytes). Packing still
copies only the view, so the SHM payload does not grow.

2. Batch split-map race in execute_batch.

The worker's background D2H/SHM thread can publish OUTPUT_READY before
execute_batch registers _batch_split_map[batch_id]. The pump then found no
map and cached the batch-level output under batch_id, while consumers waited
on per-request ids — hanging every request in the batch. execute_batch now
adopts an early-arriving output instead of registering a map nobody consumes.

3. result_mq has two concurrent writers.

MessageQueue.acquire_write() is single-writer: it reads current_idx, writes
that block, then advances the index. In async mode the background output thread
enqueues OUTPUT_READY while the main loop enqueues COMPUTE_DONE, so
unsynchronized writers can target the same ring block and clobber a message.
Writes are now serialized with a lock.

4. Single-request output future race during SHM unpacking.

In _result_pump, handling single-request OUTPUT_READY popped from _output_futures
under _futures_lock, then released the lock to perform unpack_diffusion_output_shm.
If the engine thread called wait_output_ready() during unpacking, it found
_completed_outputs empty and created a new pending future in _output_futures.
When _result_pump finished unpacking, it created its own future and stored it in
_completed_outputs, leaving the caller's future in _output_futures orphaned and
eternally pending. The fix unpacks SHM first, then atomically resolves _output_futures
or caches into _completed_outputs under _futures_lock.

Additionally, the engine now logs executor bookkeeping when the async wait times
out, so a future stall names the stage that lost the output instead of only
reporting the timeout.

Not a duplicate

Per-contribution policy, I checked open PRs for overlap:

• [Bugfix] Drain async diffusion outputs before worker shutdown #5758 (drain before worker shutdown) — different lifecycle event
  (shutdown vs sleep) and does not touch packing or the races.
• [codex] Fix diffusion worker timeout for large broadcast payloads #4845 raises the worker MessageQueue.dequeue timeout for payloads that
  take the overflow path. Materially different direction: this PR keeps those
  payloads off the overflow path entirely by packing them into SHM. The two are
  compatible and [codex] Fix diffusion worker timeout for large broadcast payloads #4845 corroborates that the overflow path is fragile.
• [Bugfix][Diffusion] Drop outputs for cancelled async waiters #5645 (drop outputs for cancelled waiters), [Bugfix] Make diffusion output wait timeout configurable #5725 (configurable output
  timeout), [Bugfix][XPU] Fix diffusion output corruption from async D2H copy #5782 / [XPU][Bugfix] Make Omni AR async output path device-agnostic #5569 (XPU async-output correctness), [Model] Cosmos3: Reduce video output transfer overhead #5712 (Cosmos3
  model-specific transfer) — all disjoint from these four defects.

I verified defects are still present on current main.

Note on the pickle allowlist

tools/pre_commit/check_pickle_imports.py gained one test-file entry. The
regression test measures the wire size of a diffusion output, and only pickle
reproduces the per-view storage duplication under test — torch.save
deduplicates shared storages, so it cannot express the bug. The test serializes
locally constructed tensors and loads no untrusted data.

Test Plan

vLLM-Omni Commit: 3fd9f30 (branch merged with upstream main @ 5be7541)

Unit tests (CPU), regression tests per defect:

pytest tests/diffusion/test_diffusion_ipc.py        tests/diffusion/test_result_pump.py        tests/diffusion/test_async_output_worker.py -q

End-to-end verification across multiple modalities:

1. Qwen-Image FlowGRPO LoRA training: request-level batching (step_execution=false, max_num_seqs=32, 4×GPU), test_freq=save_freq=1 exercising continuous rollout → validation → checkpoint (sleep/wake) cycles.
2. LTX-2.3 Text-to-Audio-Video FlowGRPO LoRA training: request-level async rollout (step_execution=false, 4×GPU, rollout_tp=2), joint video+audio generation with CLAP + ImageBind reward functions.

Test Result

Unit tests:

49 passed, 16 warnings in 44.15s

All unit tests pass cleanly, including the new atomic resolution and IPC storage regression tests.

End-to-end, before vs after:

• Before:
  • Qwen-Image died at checkpoint+validation cycle with TimeoutError in step_streaming; all requests in the batch timed out simultaneously.
  • LTX-2.3 single-request async rollout hung in Step 1 with TimeoutError due to orphaned futures in _output_futures.
• After:
  • Qwen-Image: 10/10 training steps, 10/10 validation cycles, 10/10 checkpoints, 0 timeouts, 0 errors. Mean generation time improved from ~130 s to ~102 s per step.
  • LTX-2.3: 10/10 training steps, 2/2 validation cycles, 0 timeouts, 0 errors, with expected reward progression (ImageBind: 0.1915 → 0.2066).

ruff check and ruff format --check clean on all touched files; pre-commit hooks pass.

Disclosure

AI assistance (Cursor) was used for the investigation, the patch, and the tests.
A human submitter has reviewed every changed line, ran the commands above, and
can defend the change end-to-end.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/index.md, docs/design/module/diffusion/offloader.md, docs/design/module/diffusion/diffusion_runtime.md.

Module owners: @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi @yuanheng-zhao @lishunyang12 @fhfuih

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

[SamitHuang]

Tracking issue: #6024

[Copilot]

Pull request overview

Fixes a request-mode race where diffusion outputs may still be undergoing async D2H/shared-memory packing when the worker enters sleep and releases device memory, causing missing OUTPUT_READY messages and engine-side timeouts.

Changes:

• Track in-flight async diffusion outputs via a pending counter + condition variable.
• Add drain_async_outputs() and invoke it before sleep-related RPCs / sleep tasks so memory isn’t released while packing is active.
• Add unit tests covering draining behavior and ensuring non-sleep RPCs don’t drain.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
vllm_omni/diffusion/worker/diffusion_worker.py	Adds async-output draining and hooks it into sleep entry points to prevent freeing tensors during background packing.
tests/diffusion/test_async_output_worker.py	Adds unit tests validating pending tracking, drain semantics, and correct integration with sleep RPC handling.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[fhfuih]

Overall LGTM. This PR fixes the conflict between the async output channel and the free-up action of a sleep task.

[yuanheng-zhao]

Wonder why we have a manual clamp here. It should be matching pairs I think

[SamitHuang]

it's a defensive safeguard for shutdown and test teardown paths, ensures the counter never underflows to negative values or desyncs the wait_for predicate

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

[SamitHuang]

The storage-aware packing change fixes the shared-view amplification reported here, but I think there is a longer-term transport issue worth discussing.

SHM selection is still made per value, while MessageQueue chooses its overflow path based on the aggregate serialized message size. With this PR's packing logic, I measured 32 independent 900 KB tensors: none was packed because each value is below the 1 MB threshold, but the resulting message was 28,813,783 bytes, above the 24 MiB ring-buffer limit. That still reaches the ZMQ XPUB/SUB overflow path and therefore leaves the same silent-drop failure mode reachable through a different payload shape.

I don't think this should block the immediate fix, which correctly reduces the shared-view case from 128,007,220 bytes to 3,604 bytes in my focused check. Could maintainers discuss the desired long-term contract for oversized aggregate messages? Some options are:

• make packing aggregate-size-aware;
• provide chunking/reliable backpressure for the large-message path; or
• fail fast with an actionable size error instead of allowing a silent PUB/SUB drop.

A MessageQueue integration regression at batch size 32 would also help: assert that a large request-level batch either remains within the ring-buffer limit or is delivered reliably through the overflow path.

Agree:

1. Near-Term (vllm-omni IPC):
   We can make pack_diffusion_output_shm aggregate-aware when processing BatchRunnerOutput. For example, estimating the aggregate un-packed tensor footprint and dynamically lowering the per-tensor packing threshold (or greedily packing tensors in descending size order until the remaining inline payload is safely within ~12–16 MiB). This guarantees that diffusion batch messages never hit the MessageQueue overflow path regardless of request count or tensor dimensions.

2. Long-Term (vLLM MessageQueue transport) (rely on vllm upstream upgrade)
   The fallback to ZMQ XPUB/SUB for oversized messages remains a fragile path across the framework (as also observed in [codex] Fix diffusion worker timeout for large broadcast payloads #4845). Replacing the overflow path with SHM-backed chunking or a reliable socket mechanism with backpressure / fail-fast size checks in upstream vLLM would eliminate the silent-drop failure mode fundamentally.

3. Regression Test:
   Add a synthetic 32-request test case to tests/diffusion/test_diffusion_ipc.py to assert that multi-request batch outputs are safely bounded below the ring-buffer limit.

[hsliuustc0106]

does this PR solve #5793 ?

[hsliuustc0106]

how is this related to ASYNC_OUTPUT_TIMEOUT?

[SamitHuang]

They operate at two different layers of the lifecycle and protect against different failure modes:

[ Worker Process ]                                    [ Engine / Client Process ]
Worker finishes forward pass
   │
   ├──> Enqueues D2H/SHM task to background thread
   │    and sends COMPUTE_DONE ──────────────────────> Engine waits for OUTPUT_READY
   │                                                   (deadline: _ASYNC_OUTPUT_TIMEOUT = 30.0s)
   │
Trainer/Engine invokes sleep/wake or memory release
   │
Worker calls drain_async_outputs()
(deadline: _ASYNC_OUTPUT_DRAIN_TIMEOUT_S = 10.0s)
   │
   ├── Background thread finishes D2H/SHM (~10-50 ms)
   └── Worker safely releases GPU memory

Parameter	Location	Scope & Role
_ASYNC_OUTPUT_TIMEOUT (30.0s)	Engine (diffusion_engine.py)	Total deadline for the engine/client to receive OUTPUT_READY (D2H transfer + SHM write + queue
delivery) after compute finishes.
_ASYNC_OUTPUT_DRAIN_TIMEOUT_S (10.0s)	Worker (diffusion_worker.py)	Barrier timeout when a worker is asked to sleep or release device memory. It waits for
already-enqueued background D2H/SHM tasks to finish so GPU tensors aren't freed while being read.

Why _ASYNC_OUTPUT_DRAIN_TIMEOUT_S is set to 10.0s:

1. Draining is fast: Draining only waits for in-flight D2H memory copies and POSIX SHM writes, which typically complete in under 50–100 ms.
2. Preventing timeout inversion: If drain timeout were 60.0s (larger than the engine's 30.0s), a wedged worker background thread could hold the sleep RPC for 60s while the
   engine already timed out at 30s, causing desynchronized aborts across processes. Keeping drain timeout at 10.0s guarantees worker memory barriers complete well within the
   engine's 30.0s wait window.

[SamitHuang]

does this PR solve #5793 ?

no, #5793 (ResultPump crash on cancelled futures) is addressed directly by #5983, for safe cancellation handling
This PR fixes the cases of normal requests were erroneously stalling or timing out in the first place:

• tensor view storage amplification over ZMQ
• batch split-map race
• result_mq single-writer corruption
• and atomic future resolution in the pump

[hsliuustc0106]

does this PR solve #5793 ?

no, #5793 (ResultPump crash on cancelled futures) is addressed directly by #5983, for safe cancellation handling This PR fixes the cases of normal requests were erroneously stalling or timing out in the first place:

• tensor view storage amplification over ZMQ
• batch split-map race
• result_mq single-writer corruption
• and atomic future resolution in the pump

for video generation, I think we need to optimize tensor transmission from engine/server output to client. Also, we should be more concise about what's the meaning of the global env vars and can we avoid them

[SamitHuang]

for video generation, I think we need to optimize tensor transmission from engine/server output to client. Also, we should be more concise about what's the meaning of the global env vars and can we avoid them

1. On the video payload transmisstion

Right, for video models (e.g., LTX-2.3, MiniMax-H3), the video payloads (uncompressed raw video tensors) are much larger than image payload. When the Engine/API server sends final outputs over HTTP API to an external user/client, returning raw video arrays (e.g. JSON/base64 of 1.2 GB) over the network is inefficient.

A better solution is either compressed video streaming (for web clients) or zero-copy shared-memory pointers (for co-located trainers like verl-omni).

There are 2 hops in the transmission pipeline:
• Hop 1: GPU Worker Process $\rightarrow$ Engine / Scheduler Process (Internal IPC)
• Hop 2: Engine / Server $\rightarrow$ External Client (HTTP / API Serving)

This PR is a targeted bugfix solving async output drops, worker-to-engine IPC buffer overflow, and pump thread race conditions, and has fixed and optimized Hop 1 via POSIX Shared Memory. For Hop 2 optimization (either compress the video or use zero-copy SHM), I think we should do it in another followup PR.

2. On the global env var

This PR doesn't add new env var. _ASYNC_OUTPUT_DRAIN_TIMEOUT_S = 10.0 is purely an internal Python constant in diffusion_worker.py. I have updated the comment and docs to explain it

[hsliuustc0106]

This PR is a targeted bugfix solving async output drops, worker-to-engine IPC buffer overflow, and pump thread race conditions, and has fixed and optimized Hop 1 via POSIX Shared Memory. For Hop 2 optimization (either compress the video or use zero-copy SHM), I think we should do it in another followup PR.

Please mark it with a new issue and we need to proceed to optimize the transmission. BTW, I remember we have done similar optimization to the WAN model before but it seems that optimization is not repo-wide:)

[SamitHuang]

Please mark it with a new issue and we need to proceed to optimize the transmission. BTW, I remember we have done similar optimization to the WAN model before but it seems that optimization is not repo-wide:)

track here: #6212

[hsliuustc0106]

lgtm