[Perf][Cosmos3] Optimize I2V conditioning transfer (8.63% faster x-inference)

[david6666666]

Summary

#4340

This PR adds a small lossless Cosmos3 I2V input-transfer optimization found while analyzing a 4x B300 Cosmos3 I2V profile.

• Move the single I2V reference frame to the VAE device before expanding it into a full conditioning video, avoiding materializing [1, 3, T, H, W] on CPU and then issuing a large pageable H2D copy.
• Reuse cached VAE latent mean/std tensors for encode/decode normalization instead of recreating device tensors at each call site.
• Add a unit test that locks in the single-frame-before-expand behavior.

The original trace still shows GPU compute as the primary bottleneck (attention + GEMM), so this PR intentionally keeps the code change scoped to the low-risk CPU/H2D path identified in the I2V conditioning setup.

Validation

• python -m pytest tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
• python -m ruff check vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
• python -m ruff format --check vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
• python -m compileall vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

B300 I2V smoke benchmark

Setup: 4x B300, nvidia/Cosmos3-Nano, 1280x720, 189 frames, 35 steps, seed 1111, cfg_parallel_size=2, ulysses_degree=2, vae_patch_parallel_size=4, no torch profiler, lightweight diffusion pipeline profiler enabled.

Build	x-inference-time-s avg	stage_0_gen_ms avg	Pipeline forward avg	Output SHA
upstream/main @ 46e0b71	41.186s	34045.544ms	~30.646s	452d899a...
this PR @ ef478f3	34.351s	32587.751ms	~29.834s	452d899a...

The output MP4 hash is identical across upstream and patched runs for the fixed seed/input.

Note: ef478f38 has the same code tree as the previously benchmarked ef478f38; it only adds the DCO Signed-off-by: trailer.

Visual Evidence

GitHub-hosted comparison videos for the fixed seed/input smoke run. Both files are bit-identical (sha256: 452d899ab8246e33415d27f8c1557551800eff0c06d7da1dfd9fb7b479483d46), matching the PSNR/SSIM check above.

Build	Video	Notes
baseline 46e0b714	baseline_46e0b714_measured_2.mp4	upstream/main baseline output
optimized ef478f38	optimized_ef478f38_measured_2.mp4	PR output

Release asset page: https://github.com/david6666666/vllm-omni/releases/tag/pr-4467-cosmos3-i2v-videos## Reproduction Commands

Baseline and optimized runs used the same request parameters. Run the serve command from the corresponding worktree/commit, then send the curl request below. The benchmark numbers in this PR used no torch profiler and kept only the lightweight diffusion pipeline profiler enabled.

Baseline worktree/commit:

cd /home/zjy/code/david/worktree/vllm-omni-cosmos3-i2v-b300-upstream-verify
git checkout 46e0b714

Optimized worktree/commit:

cd /home/zjy/code/david/worktree/vllm-omni-cosmos3-i2v-b300-compute-vae-opt
git checkout ef478f38

Serve command used for the optimized run, port 8012:

CUDA_VISIBLE_DEVICES=4,5,6,7 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_ATTENTION_BACKEND=CUDNN_ATTN \
/home/zjy/code/david/.venv/bin/python -m vllm_omni.entrypoints.cli.main serve nvidia/Cosmos3-Nano \
  --omni \
  --host 127.0.0.1 \
  --port 8012 \
  --model-class-name Cosmos3OmniDiffusersPipeline \
  --no-guardrails \
  --cfg-parallel-size 2 \
  --ulysses-degree 2 \
  --vae-patch-parallel-size 4 \
  --enable-diffusion-pipeline-profiler

Serve command used for the baseline run, port 8013:

CUDA_VISIBLE_DEVICES=4,5,6,7 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_ATTENTION_BACKEND=CUDNN_ATTN \
/home/zjy/code/david/.venv/bin/python -m vllm_omni.entrypoints.cli.main serve nvidia/Cosmos3-Nano \
  --omni \
  --host 127.0.0.1 \
  --port 8013 \
  --model-class-name Cosmos3OmniDiffusersPipeline \
  --no-guardrails \
  --cfg-parallel-size 2 \
  --ulysses-degree 2 \
  --vae-patch-parallel-size 4 \
  --enable-diffusion-pipeline-profiler

Curl command used for the optimized measured run:

curl -sS -X POST http://127.0.0.1:8012/v1/videos/sync \
  -H "Accept: video/mp4" \
  -D /home/zjy/code/david/tmp/cosmos3_i2v_b300_perf_opt/run_20260616_0300/patched_i2v_h2d_cfg2_usp2_vae4/measured_2.headers.txt \
  -F model=nvidia/Cosmos3-Nano \
  -F prompt="The scene comes to life with smooth, natural motion." \
  -F negative_prompt="blurry, distorted, low quality" \
  -F size=1280x720 \
  -F num_frames=189 \
  -F fps=24 \
  -F num_inference_steps=35 \
  -F guidance_scale=6.0 \
  -F max_sequence_length=4096 \
  -F flow_shift=10.0 \
  -F extra_params='{"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
  -F seed=1111 \
  -F input_reference=@/home/zjy/.cache/huggingface/hub/models--benjiaiplayground--Cosmos3-Nano_fp8/snapshots/ee4be0843a7817e4f3840e5da91ec4afca55f75d/assets/example_i2v_input.jpg\;type=image/jpeg \
  -o /home/zjy/code/david/tmp/cosmos3_i2v_b300_perf_opt/run_20260616_0300/patched_i2v_h2d_cfg2_usp2_vae4/measured_2.mp4 \
  -w 'time_total=%{time_total}\n'

For the baseline measured run, use the same curl command with port 8013 and these output paths:

-D /home/zjy/code/david/tmp/cosmos3_i2v_b300_perf_opt/run_20260616_0300/upstream_46e0_cfg2_usp2_vae4/measured_2.headers.txt
-o /home/zjy/code/david/tmp/cosmos3_i2v_b300_perf_opt/run_20260616_0300/upstream_46e0_cfg2_usp2_vae4/measured_2.mp4

[david6666666]

Cosmos3 I2V B300 profile + optimization summary

Profile source: cosmos3_i2v_profile_20260615_205112 (analysis_all_ranks.txt, profiler_out_*.txt, server.log, headers, and trace_rank*.json.gz). Denominator is the summed 4-rank GPU trace span: 133.320s.

1. Bottleneck split from the original torch trace

Category	Total time	Span %	Notes
CPU bubble / GPU idle	18.825s	14.12%	GPU idle union gaps
Communication / NCCL	9.260s	6.95%	NCCL send/recv/all-to-all kernels
Memcpy / memset	0.924s	0.69%	Kept separate from model compute
Compute, excluding memcpy	104.311s	78.24%	non-NCCL GPU kernels
Compute, including memcpy	105.235s	78.93%	coarse reporting view

Rank-level split:

Rank	CPU bubble	Communication	Memcpy/set	Compute
rank0	14.15%	7.75%	1.38%	76.72%
rank1	14.19%	5.83%	0.44%	79.53%
rank2	13.82%	9.07%	0.48%	76.63%
rank3	14.32%	5.12%	0.46%	80.10%

Compute-heavy kernels remain the top-level bottleneck:

GPU busy subcategory	Total time	Span %	GPU busy %
Attention / CUDNN SDPA	53.516s	40.14%	46.74%
GEMM / matmul	32.978s	24.74%	28.80%
Elementwise	10.777s	8.08%	9.41%
VAE conv/media/layout kernels	3.912s	2.93%	3.42%
Norm / fused Triton	2.177s	1.63%	1.90%
Layout copy kernels	0.845s	0.63%	0.74%

2. Historical clean A/B runs, no torch profiler

These runs were used to decide where to spend implementation time. They are historical A/B numbers from the profiling branch, not all code changes in this PR.

Config	Measured x-inference-time-s	Avg	Readout
cfg_parallel_size=2, ulysses_degree=2, vae_patch_parallel_size=4 baseline	58.499, 61.159, 52.858	57.505s	Current profiled config on the older baseline
Same config, eager path	39.761, 39.947	39.854s	Avoided some compile/profile effects
cfg=2, usp=2, vaepp=1	37.775, 37.861	37.818s	VAEPP=4 was not clearly helping in this setup
cfg=2, usp=2, vaepp=2	40.748, 40.992	40.870s	Slower than VAEPP=1 here
cfg=1, usp=4, vaepp=4	36.934, 37.023	36.979s	Best lossless parallel config in this sweep
cfg=2, usp=2, vaepp=4, fp8	33.029, 32.960	32.995s	Fastest measured, but quality/numerical-risk category

3. This PR: patch-only comparison against latest upstream/main

Setup: 4x B300, nvidia/Cosmos3-Nano, 1280x720, 189 frames, 35 steps, seed 1111, cfg_parallel_size=2, ulysses_degree=2, vae_patch_parallel_size=4, no torch profiler, lightweight stage profiler enabled. Base is upstream/main@46e0b714; PR head is ef478f38.

Build	x-inference-time-s	Avg	stage_0_gen_ms	Avg	Pipeline forward avg	Output SHA
upstream/main @ 46e0b714	44.877, 37.495	41.186s	34489.015, 33602.072	34045.544ms	~30.646s	452d899a...
this PR @ ef478f38	34.443, 34.259	34.351s	32675.854, 32499.647	32587.751ms	~29.834s	452d899a...
Delta	-10.434, -3.236	-6.835s / -16.6%	-1813.161, -1102.425	-1457.793ms / -4.3%	~-0.812s / -2.6%	identical

The identical MP4 SHA across upstream and patched runs confirms this change is lossless for the fixed seed/input smoke case.

4. Optimization points and status

Optimization point	Evidence	PR status	Result / next action
Move I2V single reference frame before full-video expand	Host-stack run showed many H2D copies; original code materialized [1, 3, T, H, W] on CPU before .to(device)	Implemented	Avoids large pageable H2D for the repeated conditioning video; unit test locks in shape order
Cache VAE latents_mean/std tensors	Encode/decode recreated small device tensors at multiple call sites	Implemented	Removes repeated tensor construction/H2D in VAE normalization path
VAE decode output broadcast	Upstream already calls tiled_decode(..., broadcast_result=False) in the distributed VAE path	Confirmed, no new PR diff	Keep as-is; no duplicate change needed
CFG/USP/VAEPP sweep	A/B showed cfg=1, usp=4, vaepp=4 and vaepp=1 are worth follow-up	Not in this PR	Should be a separate config/default discussion because rank communication balance changes
Attention/GEMM optimization	Attention + GEMM are ~64.88% of span / ~75.54% of GPU busy	Not in this PR	Next code-level work: attention backend/layout conversion audit and GEMM path inspection
CPU bubble attribution	Bubble is 14.12%, but the first stack run did not capture useful stacks	Partially informed this PR	Re-run short host-stack profile with idle-gap attribution to find remaining sync/log/prepost issues

5. Overall readout

The main bottleneck is still GPU compute, especially CUDNN SDPA + GEMM. Communication is visible and rank-imbalanced, but at ~7% span it is not the first bottleneck. This PR lands the lowest-risk lossless fix found during the CPU/H2D audit: avoid full-video CPU materialization for Cosmos3 I2V conditioning and cache VAE normalization tensors. The remaining larger wins should come from attention/GEMM layout/backend work and a separate parallelism sweep.

[david6666666]

Losslessness check + optimized profile comparison

Clarifying the optimization status first:

Item	In this PR?	Lossless?	Why
Move the I2V reference frame to VAE device before expand(...).contiguous()	Yes	Yes	It changes where the repeated conditioning video is materialized, not the tensor values used by VAE encode.
Cache VAE latents_mean/std tensors	Yes	Yes	It reuses the same config constants on the same device/dtype instead of recreating them per call.
VAE decode broadcast behavior	No new diff here	Yes / no quality change	This was only confirmed as already handled upstream via broadcast_result=False; this PR does not change output semantics there.

No quality-related knobs are included in this PR. The historical FP8/config A/B numbers in the previous comment were used for prioritization only; FP8 and other quality/numerical-risk changes are not part of this patch.

Baseline commit vs optimized output quality

Baseline: upstream/main@46e0b714
Optimized: this PR ef478f38
Request: nvidia/Cosmos3-Nano, 1280x720, 189 frames, 24 fps, 35 steps, seed 1111, same reference image and prompt.

Metric	Baseline measured_2	Optimized measured_2	Result
Resolution	1280x720	1280x720	Match
Frames	189	189	Match
FPS	24	24	Match
Duration	7.875s	7.875s	Match
MP4 SHA256	452d899ab8246e33415d27f8c1557551800eff0c06d7da1dfd9fb7b479483d46	452d899ab8246e33415d27f8c1557551800eff0c06d7da1dfd9fb7b479483d46	Bit-identical
PSNR	average:inf min:inf max:inf	Compared to baseline	Identical frames
SSIM	All:1.000000 (inf)	Compared to baseline	Identical frames

Commands used:

ffmpeg -hide_banner -nostats \
  -i upstream_46e0_cfg2_usp2_vae4/measured_2.mp4 \
  -i patched_i2v_h2d_cfg2_usp2_vae4/measured_2.mp4 \
  -lavfi psnr -f null -

ffmpeg -hide_banner -nostats \
  -i upstream_46e0_cfg2_usp2_vae4/measured_2.mp4 \
  -i patched_i2v_h2d_cfg2_usp2_vae4/measured_2.mp4 \
  -lavfi ssim -f null -

Optimized profile comparison

This comparison uses the clean no-torch-profiler run with the lightweight diffusion pipeline profiler enabled. I am using this view because the original torch profiler was useful for attribution, but it also perturbs end-to-end latency. The low-level original trace still says compute is the main bottleneck; this PR targets the I2V CPU/H2D setup path.

Build	Run	x-inference-time-s	stage_0_gen_ms	Pipeline forward	VAE encode	Diffuse	VAE decode	MP4 encode
baseline 46e0b714	measured_1	44.877s	34489.015ms	~30.651s	~1.712s	~24.120s	~3.919s per-rank avg	10342.99ms
baseline 46e0b714	measured_2	37.495s	33602.072ms	~30.641s	~1.681s	~24.167s	~3.919s per-rank avg	3851.51ms
optimized ef478f38	measured_1	34.443s	32675.854ms	~29.824s	~1.691s	~24.111s	~3.936s	~1725ms
optimized ef478f38	measured_2	34.259s	32499.647ms	~29.844s	~1.68-1.69s	~24.136s	~3.934s	~1718ms

Averaged view:

Metric	Baseline avg	Optimized avg	Delta
x-inference-time-s	41.186s	34.351s	-6.835s / -16.6%
stage_0_gen_ms	34045.544ms	32587.751ms	-1457.793ms / -4.3%
Pipeline forward	~30.646s	~29.834s	~-0.812s / -2.6%
Diffuse	~24.143s	~24.123s	roughly unchanged
VAE decode	~3.919s per-rank avg	~3.935s	roughly unchanged
Output quality	SHA-identical, PSNR=inf, SSIM=1.0	SHA-identical, PSNR=inf, SSIM=1.0	no quality change

Readout: the optimized profile does not claim to solve the dominant attention/GEMM compute bottleneck. The measurable improvement is in the surrounding I2V conditioning / host-transfer path and stage overhead while keeping the generated video bit-identical to the baseline commit.

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

[hsliuustc0106]

I2V conditioning transfer optimization looks correct. Moving single frame to GPU with pin_memory before expanding avoids large pageable H2D copies.

[david6666666]

@lishunyang12 @RuixiangMa ptal thx

[hsliuustc0106]

lgtm