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

test_accelerator_graph_matches_eager()	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).
test_anchor_encode_matches_full()	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.
test_batched_cfg_matches_sequential()	Running both guidance branches in one batched forward must match running them sequentially.
test_cache_once_matches_fused_exact()
test_cache_once_matches_fused_exact_t2v()
test_compile_vae_matches_eager()
test_compile_vae_matches_eager_t2v()
test_cross_request_batch_matches_individual()	Several requests denoised together in one batch must reproduce each request run alone.
test_dense_fa3_image_psnr()
test_dense_fa3_video_psnr()
test_engine_cache_path_image_psnr()
test_engine_cache_path_video_psnr()
test_windowed_kv_dense_matches_paged()	kv-mode denoise on the dense FA3 fast path — whose committed prefix mutates every window and is re-gathered on the stream's generation — against the pure paged FlashInfer backend.
test_windowed_kv_engine_release_and_psnr()	The served kv path on the real paged pool (production batched-CFG denoise + commit): pages of aged-out context are freed on the live request by the commit-time retention with exact token accounting, and window 0 matches the reference within the usual FlashInfer-vs-sdpa precision bar.
test_windowed_kv_matches_reference()	The served kv path — block-causal denoise over committed context, commit passes, page-floor release — must reproduce the hand-rolled reference bit-tightly on the sdpa resources (sequential guidance, the bit-exact regime), window by window.

mstar.model.cosmos3.tests.test_engine_cache.test_accelerator_graph_matches_eager()

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

mstar.model.cosmos3.tests.test_engine_cache.test_anchor_encode_matches_full()

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()

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()

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_cache_once_matches_fused_exact_t2v()

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_compile_vae_matches_eager()

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_compile_vae_matches_eager_t2v()

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_cross_request_batch_matches_individual()

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_dense_fa3_image_psnr()

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_dense_fa3_video_psnr()

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_engine_cache_path_image_psnr()

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_engine_cache_path_video_psnr()

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_windowed_kv_dense_matches_paged()

kv-mode denoise on the dense FA3 fast path — whose committed prefix mutates every window and is re-gathered on the stream’s generation — against the pure paged FlashInfer backend. Same machinery, different attention kernel: window 0 must agree at the usual kernel-precision bar; later windows compound the kernels’ rounding through the committed context (same autoregressive-feedback regime as the reference comparison) and get the corruption floor.

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_windowed_kv_engine_release_and_psnr()

The served kv path on the real paged pool (production batched-CFG denoise + commit): pages of aged-out context are freed on the live request by the commit-time retention with exact token accounting, and window 0 matches the reference within the usual FlashInfer-vs-sdpa precision bar. Later windows denoise against committed K/V that already carries the kernels’ rounding, so their trajectories legitimately diverge (autoregressive feedback, not an implementation error — implementation fidelity is the bit-exact sdpa check above); they get a corruption floor, not the precision bar.

Return type:

None

mstar.model.cosmos3.tests.test_engine_cache.test_windowed_kv_matches_reference()

The served kv path — block-causal denoise over committed context, commit passes, page-floor release — must reproduce the hand-rolled reference bit-tightly on the sdpa resources (sequential guidance, the bit-exact regime), window by window.

Return type:

None