Skip to content

cosmos3_video_usage.md

Latest commit

 

History

History
129 lines (98 loc) · 5.31 KB

cosmos3_video_usage.md

File metadata and controls

129 lines (98 loc) · 5.31 KB
Raw

Cosmos3-Nano text2video (FP8) — Usage & Benchmark

Kernelized, CUDA-graphed Cosmos3-Nano text2video denoise, packaged as a first-class FlashRT model (config="cosmos3_video"). Run it through the standard flash_rt.load_model(...) API.

This is a self-contained two-tower MoT denoise backbone driven for the video path: the gen tower is the all-noisy vision latent, the head is llm2vae, and the output is unpatchified to a [1, 48, T, H, W] vision latent. Conditioning (text / VAE encode) is upstream and consumed from the official reference dump; VAE decode to pixels is the downstream step.

1. Precision & speed (RTX 5090 / sm120)

Quantization choice matters: the video latent is far more quant-sensitive than the AV action head, so only FP8 (E4M3) is near-lossless; NVFP4 GEMM and int8-sage attention degrade it and are off by default.

480p / 49 frames / 10-step UniPC (denoise loop; VAE decode listed separately):

path denoise VAE decode E2E latent cos notes
official Cosmos3-Nano ~53.5 s ~2.3 s ~55.8 s — (ref) eager, uncached
FlashRT bf16 4.4 s 2.3 s 6.7 s 0.9989 CUDA graph + static text-KV cache
FlashRT fp8 2.5 s 2.3 s 4.8 s 0.986 near-lossless
FlashRT fp8 + TeaCache(3,5,7) + fp4-VAE ~1.8 s 1.26 s ~3.0 s 0.981 default-recommended

→ ~18× faster end-to-end than the official pipeline, near-lossless.

Decoded-frame quality (256p, PSNR of decoded frames vs the official video):

denoise precision frame PSNR
bf16 41.4 dB
fp8 34.2 dB
fp4 22.8 dB (lossy — not recommended)

How to read these numbers. The official baseline runs eager (no CUDA graph) and recomputes the static text tower every step, so most of the ~18× is FlashRT's CUDA graph + static text-KV cache (same precision); the quantization + TeaCache kernels contribute ~2.2× on top of a graphed bf16 baseline. Both are real; the ~18× is the end-user speedup switching from the official pipeline.

How it gets there: FP8 weights/activations (fp8_gemm_descale_bf16out) · static text K/V cache (text tower computed once) · fused qk-norm+rope · TeaCache training-free step caching · one CUDA graph per compute step · near-lossless fp4 Wan-VAE conv decode.

2. Prerequisites

  • GPU: RTX 5090 (sm120). FlashRT sm120 image (flash_rt_kernels.so + flash_rt_fa2.so).
  • Build the model-local kernels once on the target GPU (isolated; does not rebuild flash_rt_kernels.so):
cd flash_rt/models/cosmos3_video/kernels && python3 setup.py build_ext --inplace
  • Required files (paths are env/arg-driven — no host paths baked in):
    • Cosmos3 flat-format weights .safetensors (the base text2video model, converted from the public diffusers transformer).
    • The official reference dump tensors.safetensors (conditioning: text/VAE-encode tokens, rope tables, initial latent, timestep embeds).

3. Quickstart

python3 examples/cosmos3_video_quickstart.py \
    --checkpoint <cosmos3 flat weights .safetensors> \
    --ref <.../tensors.safetensors> \
    --teacache-skip 3,5,7

Expected (RTX 5090, 480p/49f/10-step):

[cosmos3_video] denoise 1809.1 ms  quant=fp8  teacache_skip=[3,5,7]
[cosmos3_video] latent (1, 48, 13, 30, 52)
[cosmos3_video] latent cos 0.98125  rel_l2 19.408%  (vs official reference)

Programmatic (all config via typed parameters — no environment knobs):

import flash_rt
model = flash_rt.load_model("<weights>", config="cosmos3_video",
                            hardware="rtx_sm120", use_fp8=True)
model.set_prompt(ref="<.../tensors.safetensors>")   # conditioning
out = model.infer(teacache_skip="3,5,7", shift=10.0,
                  compare_ref=True, return_metadata=True)
# out["latent"]     -> [1,48,T,H,W] denoised vision latent
# out["latency_ms"] -> denoise loop wall time
# out["cos"]        -> latent cosine vs official once/final_vision_latent
#
# model.infer() with no return_metadata returns the latent tensor directly.

4. Configuration (typed parameters, default-safe)

where parameter default effect
load_model(...) use_fp8 True True → FP8 (near-lossless) · False → bf16 (reference)
set_prompt(...) ref — official reference dump (conditioning); required
infer(...) teacache_skip "" step-cache skip steps, e.g. "3,5,7" (cos 0.99) or "2,4,6,8" (faster)
infer(...) shift 10.0 UniPC shift
infer(...) compare_ref False also return cos / rel_l2 vs the official latent

The fp4 path and per-projection bf16 overrides remain available as constructor arguments on the pipeline class for experiments (lossy for video; not exposed via load_model). TeaCache is training-free step caching (it reuses the cached velocity on skip steps), not step distillation.

5. Scope

  • The denoise policy is packaged here; VAE decode to frames is downstream (the Wan2.2 VAE + its fp4/fp8 conv acceleration are applied to the user's VAE).
  • Conditioning (text encode, VAE encode) is upstream, consumed from the reference dump passed to set_prompt(ref=...).
  • Additive & isolated: the production flash_rt_kernels.so and its CMake are untouched; the model-local kernels are an isolated extension.