[rahul-steiger-nv]

Contributor intent (required human answers)

1. What outcome do you want to achieve?

Lower the VRAM required for Cosmos3 generation

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

The lifetime of unused GPU buffers

3. How will this work move that variable?

Release buffers once they are no longer needed, reducing peak memory usage

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

People who want to run Cosmos3 models with less VRAM.

Purpose

Reduce Cosmos3 memory by releasing obsolete conditioning, VAE cache/chunk references and delivered engine outputs while preserving consumer ownership. Internal channels_last decoder padding probes are skipped; lossless, encoder and conv_out checks remain.

Distributed tiling now gathers only to rank 0, preallocates temporal tile output, and releases consumed tile/send/cache buffers. HCCL retains guarded all_gather for compatibility. Temporal history and gathered views remain valid through decoding/merging. Reverted the unrelated decode.py copyright change.

Test Plan

Cosmos3-Nano BF16, GB300 node nvl72d037-T03, 1280×720, 189 frames, 35 steps, seed 42, guidance 6, flow shift 10; FA4, regional compilation, channels-last VAE, layerwise DiT offload. ARM64 Enroot v0.32.0rc1, PyTorch 2.13.0+cu130. Variants ran sequentially on the same node, each with 1 warmup + 3 measured requests. Allocated-memory/probe diagnostics ran separately.

Test Result

Fresh full-PR comparison

Baseline ad3a730e8 → head 437ecb192. All results below are from the fresh rerun, with identical primed compiler kernels per topology and unchanged kernel-choice audits. The independent-autotuning trial was excluded because it selected different kernels.

Single GPU, untiled attribution (CFG/Ulysses/VAE parallelism 1). Lifetime/buffer-only: 1287a65c5; probe-only: 7eed6f2be. Cold means the first full-size request with cold VAE verdicts and a primed compiler cache.

Variant	Allocated GiB, cold / steady	Reserved MiB, cold / steady	Median (range), s
Baseline	15.452 / 14.227	19,604 / 19,604	77.136 (77.098–77.198)
Lifetime/buffer fixes only	14.618 / 12.405	19,378 / 19,378	77.071 (77.060–77.135)
Probe skip only	13.252 / 14.225	17,104 / 17,554	77.143 (77.077–77.160)
Full PR	12.403 / 12.404	16,652 / 16,652	77.131 (77.064–77.138)

Lifetime/buffer fixes save 1.822 GiB steady allocated memory. Probe skipping mainly reduces first-use allocations and retained reservation; its steady allocated saving is only 1.85 MiB. The non-probe variant includes all current buffer fixes; distributed tile changes are inactive in this single-GPU case.

Four GPUs, tiled: CFG2 × Ulysses2, VAE tile parallelism 4, 256×256 tiles.

Measurement	Baseline	Full PR
Steady peak reserved, reporting worker	13,522 MiB	11,000 MiB
Rank 0 request peak allocated	10.355 GiB	9.221 GiB
Rank 1 request peak allocated	9.425 GiB	7.181 GiB
Ranks 2/3 request peak allocated	9.613 GiB	7.337 GiB
Median time (range), s	24.159 (24.100–24.189)	24.089 (24.079–24.101)

Reserved memory falls 2,522 MiB (18.65%). Repeating the baseline confirmed 13,522 MiB (median 24.199 s, range 24.191–24.242). No meaningful speedup is claimed. Non-root VAE allocated peaks fall to approximately 5.24 GiB; diagnostics verify actual tiled execution and zero non-root receive buffers, versus 2,080,899,072 bytes per rank before.

All compared raw BF16 VAE and returned float32 output hashes match; maximum error, MAE and RMSE are zero on this workload. Reserved memory is the reporting worker's PyTorch allocator pool, not total device VRAM; allocated peaks come from separate diagnostics.

Validation: 135 integrated CPU tests passed; final gather/backend suite 18 passed; real 2/4-GPU gather smoke passed. Targeted GPU tests: 16 passed, two FP16 posterior cases fail identically on baseline. CI selectors cover the added CPU and distributed tests; hosted execution is not claimed. Existing mypy errors remain; restored legacy copyright is the requested header-check exception.

AI assistance: Codex assisted with code, tests, documentation, and validation.

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[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-26T18:03:32.592933Z We couldn’t complete this review. Comment @codex review to try again.	caca3c6	Draft marked ready

ℹ️ 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.

[vllm-omni-review-bot]

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

Module owners: @david6666666 @Isotr0py @princepride

Routing: @david6666666 via module of the changed files, CODEOWNERS; @Isotr0py via module of the changed files; @princepride via module of the changed files

@rahul-steiger-nv, 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]

Omni ReviewBot: three questions on the performance claim

@rahul-steiger-nv this PR reads as a performance or value claim:

• claim: Separate diagnostics show output-worker peak allocated memory falling 10.125 → 9.221 GiB (8.93%).

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.

[princepride]

Thanks! The release points look safe to me, and keeping the probe on conv_out makes sense given its known padding mismatch. A few suggestions:

1. Bigger win in distributed decode: gather_tensors uses all_gather, so every rank allocates world_size padded tile buffers, and non-zero ranks immediately discard them. Switching to dist.gather(..., dst=0) would leave non-zero ranks holding only the send buffer, which should save more than the dels here.
2. Attribution: could you add a quick ablation (lifetime fixes only vs. probe skip only)? The probe skip is the only change that relaxes a numerical guarantee, and it only affects the first decode per shape. Peak reserved is sticky after warmup, so it's hard to tell how much of the 19.14 → 16.26 GiB is steady-state.
3. tile_exec still holds every temporal chunk plus the torch.cat result (~2× tile). Preallocating the output and copying in each chunk, like decode_frames now does, would avoid that.
4. Nit: the same del cache, payload isn't applied in _run_norm_act_cached_conv or _run_conv_out_channels_last.
5. Nit: the decode.py SPDX header change looks unrelated; maybe revert it, as in [Bugfix] Avoid retaining tensors in Quack FP8 fallback logging #9262.

[rahul-steiger-nv]

@princepride Thank you for the review. I updated the PR including the description. Would it be possible for you take another look at it?

[princepride]

LGTM