Diffusion Perf Opt

vLLM Omni Performance Optimization

Use this skill to run a disciplined optimization loop for vLLM Omni diffusion workloads. Keep two ideas separate: real performance baselines are collected with low overhead, while torch profiler traces are diagnostic artifacts and may distort latency.

First Questions

Before proposing changes, ask for the optimization scene if it is not already known:

• GPU model, card count, topology, and whether NVLink is present.
• Model and pipeline, for example Wan2.2 I2V A14B.
• User workload: resolution, frames, steps, batch/concurrency, CFG scales, prompt/image inputs.
• Three runnable commands for each model/workload, or scripts that generate them:
  • Server startup command: exact vllm serve command, environment variables, model path, port, parallelism flags, profiler flags, and precision/compile settings.
  • User request command: exact single-request client command (for example curl or Python client) with fixed prompt/media/seed/size/steps. Use this to validate correctness and collect per-request stage timings.
  • User benchmark command: exact repeatable benchmark command or script with warmup count, measured iteration count, concurrency/batch policy, output directory, and summary format.
• Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
• Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
• Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
• Quality validation method: exact-output comparison when feasible, or metrics such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker checks, baseline self-run variance, and reviewable output artifacts.

Workflow

1. Freeze the measurement protocol and commands.

   • Before analyzing parallel strategies or traces, establish the three commands for the target model/workload: server startup, single user request, and user benchmark.
   • Prefer checked-in scripts or one-off shell scripts over free-form commands in chat. The scripts should make workload variables explicit, including model path, port/base URL, prompt/media inputs, resolution, frames, steps, seed, warmup, measured iterations, concurrency, and output path.
   • Measurement reports must preserve the concrete process, not just the final numbers. For each tested configuration, write down the server startup command, the user request command, the user benchmark command or polling loop, the final server response metrics, and the client-side timing/HTTP result.
   • Present all duration and latency values in milliseconds (ms) in tables and summaries. If an API returns seconds, convert to ms before reporting. Keep raw units only inside quoted raw logs or source snippets.
   • Keep profiler modes separate in the command set:
     • Baseline/benchmark commands should avoid torch profiler and stack collection.
     • Diagnostic commands may enable torch profiler after a baseline identifies a bottleneck.
   • If any of the three commands is missing, ask for it or create a proposed script before proceeding.
   • Make clear which command is authoritative for correctness validation and which command is authoritative for performance numbers.

2. Collect a real baseline.

   • Disable PyTorch profiler and stack collection.
   • Disable or fix torch.compile state so A/B is fair.
   • Avoid --enforce-eager for production-speed baselines unless eager is the target.
   • Prefer --log-stats and --enable-diffusion-pipeline-profiler for low-overhead stage timing. PR 3069 has the relevant metrics/log-stats changes; if local code does not include them, fetch or cherry-pick the minimal metrics changes rather than merging unrelated PR drift.
   • Run warmup requests and exclude first-request lazy init.

3. Model the parallel strategy before testing.

   • Estimate compute and communication for self-attention, cross-attention, FFN, CFG, VAE encode/decode, and HSDP.
   • Select a small candidate matrix rather than exhaustively testing everything.
   • Typical 2-card candidates: USP=2, CFG=2, USP=1/HSDP on-off, VAE parallel on-off if memory allows.
   • Typical 4-card candidates: USP=4, CFG=2 x USP=2, USP=2 x HSDP, VAE parallel on-off.
   • Typical 8-card candidates for official CFG-enabled video diffusion models:
     • Primary candidate: CFG=2 x USP=4 with VAE patch parallel across all 8 ranks.
     • Compare against: CFG=1 x USP=8 to test whether larger sequence parallel groups beat CFG branch parallelism.
     • Isolate VAE parallelism: keep the DiT strategy fixed and compare VAE patch parallel on/off or different VAE patch sizes.
     • If long-video diffuse remains dominant and Ulysses all-to-all is suspected, test a hybrid sequence strategy such as CFG=2 x USP=2 x Ring=2.
     • Test HSDP off only if the model fits without it; treat HSDP primarily as a memory strategy until A/B proves otherwise.
   • Prefer USP/SP for long video token sequences; prefer CFG parallel when CFG doubles transformer forwards and sequence length is modest.
   • Convert workload shape to latent/patch token counts before choosing candidates. For video models, record frames, latent frames, latent HxW, patch size, approximate token count, and which stages should scale with the token count.

4. Search for the best parallel configuration.

   • Start with a small matrix that answers one question per comparison:
     • CFG parallel vs larger USP: compare CFG=2 x USP=world/2 against CFG=1 x USP=world for CFG-enabled workloads.
     • VAE patch parallel value: compare the best DiT strategy with VAE patch parallel enabled and disabled.
     • HSDP cost: compare HSDP on/off only when both configurations fit memory.
     • Ulysses vs Ring: test a Ring/Ulysses hybrid only after long-sequence diffuse is confirmed dominant.
   • For each configuration, create or record a stable config id, for example A_cfg2_usp4_vaepp8_hsdp_tiling.
   • For each config id, capture:
     • exact server command and environment variables,
     • observed distributed setup from logs, such as SP groups, CFG groups, HSDP shard/replicate sizes, VAE patch size,
     • exact request command for every scenario,
     • final server response metrics and client-side elapsed time,
     • output artifact paths and any failed/empty responses.
   • Use one warmup request per scenario, then at least three measured repeats for the shortlisted config. For exploratory matrix pruning, one measured request is acceptable only if the margin is large; label it as one-shot.
   • Compare configurations by stage, not only end-to-end latency. A configuration can improve diffuse while hurting vae.decode; record both effects.
   • Select the best config only after checking the target metric, dominant stages, memory headroom, and output correctness.

5. Run targeted A/B tests.

   • Change one variable per test.
   • Keep model, input, seed, request parameters, GPU placement, and warmup policy fixed.
   • Record latency, stage timings, memory, output quality, and logs.
   • Report comparison tables in ms. Include at least end-to-end client time, server inference_time, server stage generation time if available, vae.encode, diffuse, vae.decode, and peak memory.

6. Enforce the quality and precision gate for every optimization.

   • Do not mark an optimization ready until quality is checked on the same workload: model, prompt/media, seed, shape, steps, scheduler, dtype, backend, parallelism, and output encoding.
   • For math-preserving changes, require exact/near-exact agreement or show that differences are within baseline self-run variance.
   • For precision, quantization, approximate attention, sparse/custom kernels, or backend changes, include reviewable artifacts and metrics such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
   • If quality tolerance is not stated, default to quality-preserving behavior. Failed or inconclusive quality validation blocks a ready-to-merge claim.

7. Collect diagnostic trace only after narrowing hypotheses.

   • Use torch profiler for a small number of requests.
   • Run two separate diagnostic traces instead of mixing concerns:
     • Operator/shape trace: enable torch_profiler_record_shapes=True and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
     • Host-stack trace: enable torch_profiler_with_stack=True and normally keep shape collection disabled. Use this to map CPU/Python host gaps, synchronization points, scheduler paths, and request handling overhead.
   • Keep both trace commands and reports separate from baseline/benchmark commands. Torch profiler latency is diagnostic only and must not be used as the final latency claim.
   • Prefer profiling only the narrowed dominant scenario, for example the highest-resolution/video-length case where diffuse dominates.
   • If profiler endpoints are available, run one warmup request first, then call /start_profile, run one profiled request, and call /stop_profile. This keeps model initialization and warmup out of the diagnostic trace.
   • Start by analyzing rank 0 only. Expand to more ranks only if rank 0 suggests imbalance, unclear GPU idle/free bubbles, high NCCL wait, server timing mismatch, CFG branch imbalance, or USP group stragglers.
   • Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
   • Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.

8. Analyze host, communication, and operators.

   • Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
   • Separate real GPU idle from profiler overhead such as CUPTI Command Buffer Full.
   • Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
   • Rank operator work by total CUDA time and by repeated small-kernel launch count.

9. Produce an optimization plan.

   • Classify candidates as P0/P1/P2.
   • For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
   • Include a concrete quality/precision validation gate for each candidate.
   • Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
   • End the plan with a user-facing candidate selection table. The assistant should not automatically choose a risky optimization just because it is technically possible. Present the options clearly and let the user decide which item is worth implementing based on latency target, engineering budget, memory headroom, and quality tolerance.
   • Organize the plan by optimization layer, not by a flat list of ideas:
     • host/runtime optimization,
     • measurement/benchmark reliability,
     • parallelism and communication,
     • VAE encode/decode and media pre/post processing,
     • operator fusion and layout cleanup,
     • attention main-path optimization,
     • algorithmic, precision, or approximation changes.
   • For each layer, tie every candidate back to evidence from baseline metrics, diagnostic trace, source code, or output quality requirements. Do not list generic optimizations without a trace or workload reason.

Priority Rules

• P0: low risk, likely useful, or required for trustworthy measurement. Examples: real baseline, warmup, targeted parallel A/B, disabling avoidable empty_cache, scheduler coefficient caching.
• P1: meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
• P2: high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.

Every implemented optimization needs A/B validation and a passing quality/precision gate, including math-preserving changes such as padding trim, layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics when available, keep reviewable artifacts, and compare against baseline self-run variance for precision, quantization, approximate kernels, or backend changes. If quality validation fails or is inconclusive, do not present the optimization as ready to merge.

Optimization Layers

After baseline, parallel-search, and diagnostic traces, summarize optimization opportunities by layer. This is the core of the performance analysis: the goal is to connect evidence to a scoped implementation and a validation plan.

Host and Runtime Optimization

Purpose: remove CPU/Python stalls, synchronization points, allocator overhead, and request-path overhead that leave GPU lanes empty.

Evidence to look for:

• High idle_pct or large GAP blocks in trace_analyzer.py.
• Host-stack trace lines such as torch.cuda.empty_cache, cudaStreamSynchronize, cudaDeviceSynchronize, Python locks, scheduler waits, image/video preprocessing, or repeated small allocation paths.
• Difference between client wall-clock time and server inference_time_s.

Typical candidates:

• Make avoidable torch.cuda.empty_cache() optional or guard it by memory headroom.
• Cache scheduler coefficients, timesteps, masks, or other tiny repeated CPU computations when the request shape/steps are fixed.
• Remove avoidable host-device synchronizations and blocking logging/stat calls.
• Move expensive preprocessing out of the critical path or cache fixed prompt, image, and transform work for benchmark scenarios.
• Ensure benchmark scripts record client-side elapsed time, HTTP status, output path, and server response metrics.

Priority guidance:

• Usually P0 when the change is measurement reliability or an obvious removable synchronization.
• Usually P1 when it changes scheduling, memory lifetime, or request execution order.

Validation:

• Re-run non-profiler baseline with same workload and seed.
• Confirm peak memory headroom if disabling cache cleanup.
• Confirm generated output exists and quality/seed behavior is unchanged.

Parallelism and Communication Optimization

Purpose: choose the right decomposition for CFG branches, sequence tokens, model weights, VAE tiles, and rank topology.

Evidence to look for:

• Baseline A/B across CFG, USP/SP, Ring, HSDP/FSDP, and VAE patch parallelism.
• Stage timing shifts: diffuse, vae.encode, vae.decode, and server end-to-end.
• NCCL kernel time from trace, not only user_annotation time.
• Rank imbalance across SP group ranks or CFG branch ranks.
• Memory headroom and OOM risk.

Typical candidates:

• CFG=2 x USP=world/2 versus CFG=1 x USP=world for CFG-enabled models.
• VAE patch parallel on/off or patch size tuning.
• HSDP/FSDP on/off only if both configurations fit memory.
• Ulysses versus Ulysses+Ring only after long-sequence diffuse is confirmed dominant and all-to-all is suspected.
• Rank mapping/topology changes if all-rank traces show stragglers or NCCL wait.
• Buffer reuse or preallocation for FSDP/HSDP all-gather paths.

Priority guidance:

• P0/P1 for configuration-only changes with strong measured wins.
• P1 for buffer reuse or rank mapping changes.
• P2 for invasive distributed algorithm changes.

Validation:

• Measure by stage and memory, not only end-to-end.
• Use one-variable A/B with identical prompt/media/seed/shape/steps.
• When communication is suspected, compare rank0-3 in one USP group and rank0 versus rank4 across CFG branches for CFG=2 x USP=4.

VAE and Media Pipeline Optimization

Purpose: reduce encode/decode, tiling, split/gather, and media conversion time.

Evidence to look for:

• Large vae.encode or vae.decode in low-overhead stage timings.
• Host-stack gaps in VAE tile split, gather, merge, broadcast, or image/video transforms.
• VAE kernels or cuDNN convolution in operator trace.
• Whether every rank needs the final decoded tensor.

Typical candidates:

• Keep VAE patch parallel enabled when it has clear measured benefit.
• Reduce VAE gather/broadcast to only ranks that need the final media output.
• Reuse tile metadata, split buffers, or gather buffers.
• Evaluate bf16/autocast behavior for VAE only with visual quality checks.
• Avoid redundant image conversion, resize, or tensor construction in repeated benchmark runs.

Priority guidance:

• P0/P1 if VAE is a large share of the target workload or if a host gap is obvious and low risk.
• Lower priority when diffuse dominates and VAE is already patch-parallelized.

Validation:

• Compare vae.encode, vae.decode, server end-to-end, and peak memory.
• Check output video integrity, artifacts, flicker, and seed stability.