Repository navigation
[Hardware][Ascend] Reclaim denoise-loop allocator blocks before streaming VAE decode - #9313
Conversation
…ming VAE decode WanStreamingDecoder.decode_chunk runs right after each AR block's denoise steps, and the denoise loop's cached allocator blocks stay resident under the decode's per-frame allocations. No existing empty_cache call site covers this path. Reclaim them through current_omni_platform.empty_cache() behind a non-CPU device guard; the lazy import keeps the module's device-less contract tests free of the platform plugin machinery. Measured on 8x Ascend 910B3, lingbot-world-v2-14b stepwise streaming (TP2+USP4, 480x832, 117-frame session): per-card HBM peak drops from 65,521/65,536 MB to 61,273 MB, and post-session residency from ~63.6 GiB to ~52-53 GiB. Outputs stay byte-identical for a fixed seed; e2e and inter-chunk latency are unchanged. Signed-off-by: dongyang1006 <40831215@qq.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
This PR matches CODEOWNERS paths: *****. Code owners: @Gaohan123 @hsliuustc0106 @ywang96 Routing: @Gaohan123 via CODEOWNERS; @hsliuustc0106 via CODEOWNERS; @ywang96 via CODEOWNERS @dongyang1006, 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. |
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
@dongyang1006 — please tie the memory/latency claim to the rebased head e911c2b3647d78e1ba2c1afcc767f70b77678cc5.
- Supply a matched baseline/head run on the stated 8×910B3 setup, with the complete camera-action input and commands, equal warmups/repetitions, raw per-card memory samples and latency/output results. The current measurements explicitly use
131a773...; one without-reclaim run versus four candidate runs does not establish “zero measured cost.” Distinguish sampled device residency from an allocator peak. - Include sequential-session measurements for the residency-compounding claim, and a focused regression that reaches the new non-CPU reclaim branch. The unchanged CPU tests bypass this branch. Acceptance: the intended reclaim executes without changing streamed frame/chunk completeness or fixed-input output, including successive sessions.
This is an evidence gap, not a demonstrated new code defect. Inspection was source-only; I did not run accelerator tests. The disclosed CUDA/ROCm/XPU coverage gap remains.
hsliuustc0106
added
NPU
Cursor
1. What outcome do you want to achieve?
Keep stepwise streaming sessions off the physical HBM cap: turn a generation
peak that rides 15 MB under the card limit into one with real headroom, and
stop unreclaimed allocator residency from compounding on long-lived servers.
2. What is the highest-leverage variable for that outcome?
Per-card HBM peak during generation, and post-session allocator residency —
on memory-tight streaming deployments these decide whether a session survives
any transient allocation at all.
3. How will this work move that variable?
Reclaim the denoise loop's cached allocator blocks with
current_omni_platform.empty_cache()at the top ofdecode_chunk, so thestreaming VAE decode's per-frame buffers allocate into freed space instead of
on top of resident denoise fragments — ~4.2 GiB of headroom recovered at zero
measured cost.
4. Whose customer or user needs does this address, and what are those needs?
Operators running long-lived stepwise streaming servers on memory-tight Ascend
deployments (sequential sessions, shared devices). They need sessions that do
not ride the physical cap and do not leak ~10 GiB of residency per session.
Purpose
WanStreamingDecoder.decode_chunkruns once per committed AR block, rightafter the four DMD denoise steps of that block. The denoise loop leaves cached
allocator blocks behind, and the streaming decode then allocates its
per-frame buffers on the same device with those fragments still resident.
Nothing on this path reclaims them: the existing
current_omni_platform. empty_cache()call sites indiffusion_worker.pycover the error-cleanup,paged-KV init, memory-profiling, sleep, and pipeline-hot-swap paths only —
none of them run between a chunk's denoise and its decode.
On memory-tight streaming deployments this is the difference between riding
the physical cap and having real headroom. Measured on 8x Ascend 910B3
(64 GiB HBM per card), lingbot-world-v2-14b stepwise streaming
(
lingbot_world_v2_stepwise.yaml, TP2+USP4, 480x832, 117-frame /11-chunk session, seed 42):
Sessions complete either way on this workload, but running 15 MB under the
physical cap leaves no room for any transient allocation, and the unreclaimed
residency compounds the risk for long-lived servers (sequential sessions,
shared devices). The reclaim recovers ~4.2 GiB of headroom at zero measured
cost: e2e and inter-chunk latency are unchanged within noise
(64.3–64.5 s / ~6.6 s per chunk over four runs with the reclaim; a single
without-reclaim run was slightly slower at 68.7 s), and outputs stay
byte-identical for a fixed seed (md5-equal MP4s with/without the change).
Design notes:
current_omni_platform.empty_cache()— the repo-blessed platformabstraction already used by the diffusion worker — instead of
torch.cuda.empty_cache()(banned by thecheck-torch-cuda-callpre-commit gate) or
torch.accelerator.empty_cache()(fails on NPU:RuntimeError: Allocator for npu is not a DeviceAllocatorwithtorch_npu 2.10).
only
torchis imported so the contract can be exercised without a device,a checkpoint, or the distributed VAE stack. Keeping the platform import
inside the non-CPU branch preserves that property for the device-less
contract tests (they never take the branch) at a per-chunk call frequency
where the cached import is negligible.
decode_chunkis also the decode entry point for thecausvid_wanasyncstreaming path (
async_stream_decode.py), which interleaves denoise anddecode on the worker's consumer threads; the reclaim runs there too, on the
same non-CPU guard.
main@131a7731a44e2a64feac165580ce99d36beb94ae(file identical to thev0.30.0tag). This branch is rebased onto current
main, where [Hardware][Ascend] Add opt-in LingBot SP8 output paths #8610 addedoutput-format validations at the top of
decode_chunk; the reclaim now sitsahead of those validations as the method's first statement. The reclaim
logic itself is unchanged, and
tests/diffusion/ar_diffusion/test_streaming_decode.pyis byte-identicalbetween the two revisions.
Test Plan
vLLM Version: 0.30.0 (vllm-ascend 0.19.1rc2.dev, torch 2.10.0 +
torch_npu 2.10.0.post4, CANN 9.1.0)
vLLM-Omni Commit: measured on
main@131a7731a44e2a64feac165580ce99d36beb94ae(=v0.30.0file state); branchrebased onto current
main— see the rebase note under Purpose.npu-smi infopolled every 3 s(8 cards), run once with and once without the change.
Test Result
test_streaming_decode.py: 20 passed (CPU, with and without the change).byte-identical output to the without-reclaim run (same seed, md5-equal).
residency ~63.6 -> ~52-53 GiB).
through the platform abstraction, and the CPU contract tests cover the
guarded path's syntax/import shape.
AI assistance: the diff, measurements, and this description were drafted with
an AI coding agent (opencode / GLM); I reviewed every line, validated on
8x910B3, and take responsibility for the claims above. Validation commands
and hardware are listed in the Test Plan.