[merceod]

Cosmos3-Edge t2v at 832x480x121, 20 steps, took 6.3 s per video on a B300 against SGLang's 4.3 s, and 9.5 videos per minute at c=4 against 14.3. nsys on one request put the denoise step at 221 ms of GPU for 542 kernels, most of it attention: the DiT declares the dense backend, FA3 has no kernel for sm_103, and AttentionManager.build degraded that to the paged backend's FA2 prefill. Four commits, each measured on its own (GPU 0, pinned, two sessions per arm, ten prompts, seed 0; the earlier run's SGLang and vLLM-Omni numbers for the same cells):

arm	c=1 p50	videos/min c=1 / 2 / 4	peak GB
main	6.27 s	9.5	37.8
+ dense pass on FlashInfer's ragged prefill where FA3 has no kernel	3.90 s	15.3	39.1
+ denoise step compiled in Inductor's deterministic mode	3.90 s	15.2	39.1
+ whole clips decoded on the shared channels-last Wan VAE decoder	3.69 s	16.0 / 15.4 / 14.0	38.2
+ the cross-request denoise step compiled	3.71 s	16.0 / 19.0 / 19.0	38.2
SGLang 0.5.22	4.32 s	14.0 / 14.3 / 14.3	14-17
vLLM-Omni 0.30	5.59 s	10.8 / 15.0 / 15.3

• attn: where FA3 has no kernel for the device (or does not import), the dense pass runs on FlashInfer's ragged prefill wrapper, same gathered prefix and concatenation, one wrapper per slot because its plan stages through a pinned buffer; backend: auto resolves to cuDNN or CUTLASS FMHA on Blackwell. The paged fallback stays for where FlashInfer cannot run. The plan's page lists are host tuples now, moved to the device when the prefix is gathered (once per layer on a gather miss), where before each step copied them from pageable memory and synced the stream. On the FlashInfer path the cu_seqlens stage through the wrapper's pinned buffer; the FA3 path still copies cu_q/cu_k from pageable memory each step. The denoise step goes from 221 to 92 ms of GPU.
• denoise_compile_mode: deterministic (default) picks kernel configs by rule so every server process runs the same kernels and one seed gives one video; default and level time kernels and differ per process. Sessions are bit-identical in the default mode.
• vae_decode_fast (default true): the non-windowed video walks decode through the shared Wan VAE decoder, channels-last and compiled per call, on a decoder-only copy of the VAE, as wan22 does: 1002 to 794 ms per clip at 48.3 dB against the stock decode. The streaming subclass keeps its own step. Images (one latent) keep the compiled stock vae.decode on the same copy, faster for a single frame: 10.3 / 21.6 ms per 640² / 1024² decode against main's 15.4 / 32.4.
• compile_denoise_batched (default true): the forward that batches two or more requests compiles too. Eager, its norms, rotary and residuals cost about 80 ms of GPU per extra row; c=2 goes from 7.3 to 6.0 s per pair.

Side arms, not kept: a 400-page KV pool gives the same speed at 24.6 GB peak instead of 38 (256 pages costs c=4), and the reasoner or the windowed decoder off saves 2-5 GB; those are config choices for a deployment rather than defaults.

Quality: each arm is 20.8-21.5 dB from main's videos, and any two attention kernels (FlashInfer cutlass vs cuDNN vs the FA2 fallback) differ by the same 21 dB on this model, while main against itself across fresh processes is 39 dB; every arm sits at 13.1 dB from the diffusers references, as main does. Frames 0, 60 and 120 of three prompts for main, the branch, the cuDNN variant and a fresh main show the same scenes and motion (a lecture hall, a walker between snowy firs, a bicycle against a tree). Sessions of one arm are bit-identical.

Tests: mstar/model/cosmos3/tests, test_cosmos3_edge_model.py, test_cosmos3_graph_switch.py, test_dense_attention.py: 252 passed, 23 skipped on CPU (new cases for the FlashInfer dense manager's plan and fallback order, the compile modes and their validation, the fast decode and its switches). e2e --tier full for cosmos3_edge, cosmos3_nano_ar and cosmos3_edge_droid on B300. ruff clean.

Not done: ROCm (the AITER dense path is unchanged, the FlashInfer branch is CUDA-only by construction); the i2v and image walks on the fast decode were inspected by eye after the review (frames 0/60/120 of three prompts, main vs branch, c=1 and c=4; image sheets at 640² and 1024²); a process that hosts both the whole-clip and the streaming decoder nodes loads two decoder copies, the same point the #860 review raised.

Since the review

B300, GPU 0 pinned, two sessions each, main (with #925) against this branch, and the earlier PR head bf913a8 in one session for reference:

cell	main s1 / s2	branch s1 / s2	PR head
t2v c=1 p50	6.67 / 6.61 s	4.07 / 4.06 s	3.68 s
t2v c=2 / c=4, videos/min	9.0 / 8.5	19.0 / 19.4	- / 19.5
i2v c=1 p50	6.56 / 6.53 s	3.98 / 3.94 s	-
mixed 832x480 and 480x832 at c=4 on a fresh server: p50 / p99, per min	27.7 / 87.9 s, 6.0; 27.9 / 45.8 s, 7.7	11.9 / 20.4 s, 17.6; 12.0 / 20.3 s, 17.6	16.6 / 56.1 s, 9.2
640² images c=1 p50 / c=4 per min	1.76-1.80 s / 49.7, 68.5	1.56-1.58 s / 80.7, 79.6	- / 32.7
1024² images c=1 p50 / c=4 per min	1.77-1.92 s / 44.8, 58.5	1.76-1.80 s / 52.4, 77.4	-
idle / peak GiB	38.2-38.8 / 40.2-41.3	36.5 / 36.5-37.1	37.3 / 37.5

• The dense backend holds the prefixed ragged wrapper that dense_or_flashinfer picks, the same pick the ragged resources make; the FlashInfer subclass, _plan_kernel, _varlen and kernel: object are gone. The wrapper is a PrefixedPrefillWrapper with FA3 and FlashInfer subclasses rather than the cross wrapper, since the cross wrapper is never causal and the dense plan can be. The FA3 path gets pinned cu_seqlens, so its per-step H2D no longer syncs.
• Its wrappers share one 512 MiB workspace taken at build, before capture planning reads free memory: peak 36.5-37.1 GiB against main's 40.2-41.3 (the PR head was 37.5).
• flashinfer_backend is checked at build with a tiny plan and run: a typo, fa3 on Blackwell or an unsupported kernel fails at boot with a ValueError naming the backend, the GPU and the FlashInfer version; auto, cudnn, cutlass and fa2 pass on B300.
• Deterministic compile is an engine compile mode: compile_kwargs("deterministic") and deterministic_inductor_options(), read by the Wan VAE decode, compile_blocks and Cosmos3's denoise_compile_mode, which now only maps the yaml knob onto it; a torch without the mode logs one warning.
• Instead of padding to buckets, the cross-request step compiles only its layer stack and final norm, once, at startup (NodeSubmodule.warmup_compiles, 3 to 7 s), with the pack length dynamic; patchify, the rotary tables and unpatchify stay eager per request. No batch size or mix of sizes compiles mid-traffic: the mixed c=4 case on a fresh server has p99 20.3-20.4 s (PR head 56.1, main 45.8-87.9), and 640² images at c=4 run 80/min (PR head 32.7, main 50-69). What still compiles per new size is the single-request step and the VAE decode, as on main.

Guards: branch against main 21.1 dB on t2v (the kernel band), both 13.1 dB from the diffusers refs; i2v 26.9 dB against main, 24.5 against the refs; single-request t2v bit-identical to the PR head's arm; sessions bit-identical on both arms; images 34.1 dB median against main at 640², 22.4 at 1024². Frames 0/60/120 of three prompts at c=1 and c=4 and the image sheets show the same scenes, no streaks.

One number moved the wrong way: single-request t2v is 4.06 s against the PR head's 3.68. The worker time is the same (3.13-3.15 s); the difference sits between the conductor's done and the first chunk (319 ms on the PR head, 477 ms after the main merge, 545 ms here), and main pays it too (6.27 s before, 6.67 now), so it came in with main, probably #903's output-path pool under an 8-core pin. Separate item.

Tests since the review: cosmos3 suites, test_cosmos3_edge_model, test_cosmos3_graph_switch, test_dense_attention, test_attn_backend_auto, test_accelerator_graph_compile_mode, test_wan_vae_compile_options, test_ragged_attention, test_node_fp32_precision: 336 passed, 78 skipped on CPU; test_dense_attention and test_attn_backend_auto on a B300 (FlashInfer parity against SDPA, the real cudnnn boot failure): 69 passed. Not run: Hopper and ROCm, SP for the batched layer stack.

[mstar-review]

Automated review

🟡 1 P2 on 0d7e2cb, light review

Description

• The description says the plan's page lists and cu_seqlens are host tuples now, moved to the device when the prefix is gathered. That holds only for the page lists. On the FA3/AITER path, DenseAttentionManager._plan_kernel still copies cu_q/cu_k to the device from pageable memory with torch.tensor(..., device=...) on every plan.
• The new FlashInfer dense branch runs on any CUDA GPU where the fa3-fwd wheel does not import, including Hopper and Ampere hosts without the wheel, which used to fall back to the paged backend. The description reports e2e runs on a B300 only. It should say whether that changed fallback was run on a non-Blackwell GPU, or note that it wasn't.
• The batched-step numbers come from a homogeneous 832x480x121 t2v workload. The description doesn't say whether a mixed workload was checked (several resolutions, i2v together with t2v), where compile_denoise_batched is likely to recompile most often.

Questions

• Has compile_denoise_batched been tried on a mixed workload? denoise_step_batched is compiled with dynamic=False, and it loops in Python over per-request token shapes and noisy-frame indexes. So I'd expect each new batch composition (sizes, resolutions, i2v vs t2v), not just each batch size, to trigger a compile on the serving path, and dynamo to fall back to eager once it hits its recompile limit. If those stalls show up in serving, the docs' line that 'the first batch of each size compiles on arrival' understates the cost, and the default may need to change.

What the review checked

• AttentionManager.build checks in order: FA3/AITER, then FlashInfer ragged, then paged. The CPU-device tests in test_attn_backend_auto.py and test_sliding_window_attention.py still get the paged fallback because _flashinfer_ragged_unavailable_reason rejects non-CUDA devices.
• dense.py imports ragged.wrappers at module level, but that module imports flashinfer lazily inside __init__, so FA3-only and ROCm hosts can still import dense.py.
• FlashInferDenseAttentionManager keeps one wrapper and workspace per (label, slot) and sets force_double_buffer. Plans are staged from CPU cu_seqlens, and plan() still refuses preplan and slot leases, so the path stays eager-only.
• The FlashInfer ragged wrapper uses the same default sm_scale (head_dim**-0.5) and the same bottom-right causal alignment as the FA3 call, and the new GPU test checks both against SDPA.
• Gathering the prefix pages moved into _prefix_kv, which caches per (request, label, generation, length), so the page-index H2D runs once per gather instead of once per plan.
• Only cosmos3 declares AttnBackend.DENSE, so the engine change doesn't affect other models (invariant 4).
• The fast whole-clip decode casts the node's own VAE to the decode dtype before switching it to channels-last. DecoderStep warms each new call signature, which covers the corrupted-first-execution issue the stock path handles with _warmed_decode_shapes. The streaming subclass opts out through _whole_clip_fast = False.
• denoise_compile_kwargs('level') reuses compile_kwargs(None), and level 2 there is max-autotune-no-cudagraphs, so no CUDA graph capture enters model code (invariant 1).
• check_compile runs in _load_config after model kwargs are applied, so compile_denoise=False passed from tests still validates.

_{Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.}

[dghosef]

running tests rn, some comments

[dghosef]

Can DenseAttentionManager just hold one of the existing ragged wrappers (picked by _build_dense_or_flashinfer) instead of a new subclass + _plan_kernel/_varlen + kernel: object, so there's one FA3 -> FlashInfer selection instead of two?

[merceod]

Done in 0386214: DenseAttentionManager takes the wrapper class dense_or_flashinfer picks (the same pick the ragged resources make) and holds one per (label, slot); the FlashInfer subclass, _plan_kernel, _varlen and kernel: object are gone. The wrapper is a PrefixedPrefillWrapper with FA3 and FlashInfer subclasses rather than the cross wrapper, since the cross wrapper is never causal and the dense plan can be.

[dghosef]

Can deterministic compile be a mode in the engine's compile_kwargs instead of a Cosmos3 knob? wan_vae.compile_decode and #926 each have their own version too. Also the hasattr fallback here is silent - should log, or drop it if every torch we support has the option.

[merceod]

Done in 443328c: compile_kwargs("deterministic") and deterministic_inductor_options() in the engine, read by wan_vae.compile_decode, compile_utils and Cosmos3's denoise_compile_mode, which now only maps the yaml knob onto them. A torch without the mode logs one warning instead of the silent hasattr.

[dghosef]

This is its own 512 MiB WorkspacePool per (label, slot), up to ~3 GiB, allocated after the KV pool is sized - can the wrappers share the node's workspace, since it's only scratch during run? (verifying on gpu rn)

[merceod]

Done in 0386214, one step short of the ask: the dense wrappers share one 512 MiB workspace the manager takes at build, before capture planning reads free memory, rather than borrowing the paged manager's, since resources hold no references to each other and sharing WorkspacePool across managers would change every model's FlashInfer buffers. Peak on the t2v cells 36.5-37.1 GiB against main's 40.2-41.3.

[dghosef]

denoise_step_action_batched stays eager, so either scope the comment/docs to video/image/sound or compile it too w/ a signature in can_batch (#860 does this w/ action_signature)?

[merceod]

Scoped in bf913a8: the config comment and the docs say video, image and sound, with the action step eager. Compiling it with a signature in can_batch can come with #860's action_signature.

[dghosef]

Tested on 2x H200 (no B300 available so couldn't test the Blackwell paths). t2v c=1 goes 6.90 -> 6.55 s and c=4 7.47 -> 9.62 videos/min, and Edge output is within main's own noise. A few things:

• denoise_step_batched compiles once per batch composition, so new mixes of sizes compile mid-traffic. Once everything's compiled it's a win (9.7 vs 7.7 videos/min w/ compile_denoise_batched: false), but on a fresh server w/ mixed 832x480/480x832 at c=4 it's worse than having it off (p99 47 s vs 32.5 s), and 640x640 images at c=4 dropped to ~50/min vs main's ~97. Once it hits the recompile limit it just warns and runs eager. Can we pad to fixed batch sizes and compile those at warmup? If not we should default it off
• A bad flashinfer_backend (e.g. cudnnn, or cutlass on Hopper) boots fine and then 500s every request. We should check it at startup and raise ValueError
• Verified the dense workspace is +0.7 GiB at the first request, so it lands after the KV pool is sized. Should share the node's workspace
• vae_decode_fast covers the image walk too, but it's only measured on video - for images it's ~36 ms vs ~24 ms w/ main's compiled vae.decode (+5% e2e at c=1), though the first image is ~10 s faster since it skips main's compile. Maybe keep vae.decode for T=1 and the fast path for video only?
• nit: the prefix gather still does torch.tensor(segment.pages) per layer per segment on a miss (56 H2Ds on the first step vs main's 2/step). Can we build it once per plan?

[merceod]

Thanks for the H200 run. Plan:

• batched step: pad the batch to the capture buckets and compile each (size bucket, rows) pair at warmup for the configured sizes, so no composition compiles mid-traffic; if that warmup is too long for a fresh server, the default flips to off.
• flashinfer_backend: checked at startup against the backends this GPU has, ValueError otherwise.
• workspace: the dense wrappers borrow the node's workspace during run instead of their own pools.
• vae_decode_fast: video only; T=1 keeps the compiled vae.decode.
• prefix gather: the page tensor built once per plan.
• DenseAttentionManager holds the ragged wrapper _build_dense_or_flashinfer picks, so there is one FA3 -> FlashInfer choice.
• deterministic compile becomes an engine compile option shared by the Wan VAE decode, Cosmos3 and LingBot-World: attention through the ragged resource, deterministic block compile, bf16 channels-last decode #926's blocks, with a log line where the torch has no deterministic mode.
• denoise_step_action_batched: the docs scope the batched step to video/image/sound; compiling it with a signature comes with Performance 3/7: image, video and world models #860's action_signature.

B300 numbers for each with the push.

[merceod]

Pushed bf913a8 for the small ones: the prefix page index is built once per gather (one H2D, not one per layer), images (one latent) keep the compiled stock vae.decode and the fast path takes clips only, and the docs scope the batched step to video/image/sound with the action step eager. The wrapper reuse, the shared workspace, the deterministic compile option, the batched-step warmup and the backend check follow with B300 numbers.

[merceod]

Pushed the rest (0386214, 443328c, 99b3374, plus two test fixes; main merged). The description has a 'Since the review' section with the B300 table: mixed-size c=4 on a fresh server p99 20 s (was 56 on the PR head, 46-88 on main), 640² images at c=4 80/min (was 33), peak memory 37 GiB vs main's 41, backend typos fail at boot, deterministic compile is an engine mode. One thing to flag: single-request t2v is 4.06 s now against the PR head's 3.68; it sits outside the graph (conductor done to first chunk) and main pays it too since #903, so it goes in its own item.

[dghosef]

lgtm