mstar.model.cosmos3.tests.test_engine_cache#
GPU parity for the cache-once engine path of the Cosmos3 generator.
The understanding tower runs once and writes its per-layer K/V; the generation
tower then runs each denoise step re-reading that frozen K/V (the text tokens get
no timestep embedding, so their K/V is denoise-step independent — caching it once
is exact). This checks the Cosmos3DiTSubmodule prefill + denoise loop against
the fused Cosmos3Pipeline that runs the whole transformer every step, for both
image (single frame) and video (multi-frame, fps-modulated mRoPE) generation.
- Two GPU-gated checks per mode (need
COSMOS3_NANO_DIR+ CUDA; skipped otherwise): with an in-process sdpa cache (same attention kernel as the fused pipeline), the cache-once output is bit-for-bit identical;
with the engine’s FlashInfer paged cache (the served path), the decoded output matches the fused pipeline within PSNR >= 30 (FlashInfer-vs-sdpa precision).
Run: COSMOS3_NANO_DIR=<snap> python3 test_engine_cache.py
Functions
Image-to-video only consumes latent frame 0, and the Wan VAE encodes it as a standalone causal anchor, so the encoder node's single-frame i2v encode must give a bit-identical frame 0 to the repeat-padded full-clip encode it runs for action policy / forward-dynamics — at a fraction of the cost. |
|
Running both guidance branches in one batched forward must match running them sequentially. |
|
Several requests denoised together in one batch must reproduce each request run alone. |
|
The captured-graph denoise step is the served path's accelerator: both guidance branches run in one captured forward (~2x faster than the eager step). |
|
- mstar.model.cosmos3.tests.test_engine_cache.test_anchor_encode_matches_full()[source]#
Image-to-video only consumes latent frame 0, and the Wan VAE encodes it as a standalone causal anchor, so the encoder node’s single-frame i2v encode must give a bit-identical frame 0 to the repeat-padded full-clip encode it runs for action policy / forward-dynamics — at a fraction of the cost.
- Return type:
None
- mstar.model.cosmos3.tests.test_engine_cache.test_batched_cfg_matches_sequential()[source]#
Running both guidance branches in one batched forward must match running them sequentially. The two paths differ only in bf16 GEMM rounding (a batched matmul tiles differently), so compare the decoded images by PSNR.
- Return type:
None
- mstar.model.cosmos3.tests.test_engine_cache.test_cache_once_matches_fused_exact()[source]#
- Return type:
None
- mstar.model.cosmos3.tests.test_engine_cache.test_cache_once_matches_fused_exact_t2v()[source]#
- Return type:
None
- mstar.model.cosmos3.tests.test_engine_cache.test_compile_vae_matches_eager()[source]#
- Return type:
None
- mstar.model.cosmos3.tests.test_engine_cache.test_compile_vae_matches_eager_t2v()[source]#
- Return type:
None
- mstar.model.cosmos3.tests.test_engine_cache.test_cross_request_batch_matches_individual()[source]#
Several requests denoised together in one batch must reproduce each request run alone. Distinct prompts are decoded and compared to the fused pipeline: batching must (a) keep each request isolated — its own image far closer than any other request’s — and (b) not lose quality versus the bs=1 path (per-prompt fidelity varies with the FlashInfer kernel, so the bar is relative to bs=1, not an absolute PSNR).
- Return type:
None
- mstar.model.cosmos3.tests.test_engine_cache.test_cuda_graph_matches_eager()[source]#
The captured-graph denoise step is the served path’s accelerator: both guidance branches run in one captured forward (~2x faster than the eager step). Each captured forward matches eager to within bf16 (the first step differs by ~one ULP); the multistep solver amplifies that into a small latent spread, but the decoded image is unchanged — so gate the decoded image against the fused pipeline, the same bar the eager engine path meets.
- Return type:
None