[merceod]

Follows #963 (merged as bf70108). Measured against #963's head f109674 on B300, then rebased onto main 897eca5 with no change to the patch.

#963 skips FlashInfer's plan when a dense slot's layout hasn't changed, in the dense manager only. This moves the check into _RaggedPrefillBase._plan, so every manager on that base gets it: ragged self and cross attention on FlashInfer and FA3, and the dense manager, which loses its own copy.

What changed

• _RaggedPrefillBase._plan (attn/ragged/wrappers.py). A wrapper keeps its last schedule until it's planned again and run only reads it, so a layout it already holds returns before any work. The layout it compares against is kept on the wrapper. Wrappers are per (label, slot) and per (bucket, slot, label), so each slot's check is against that slot's own schedule, as Wan2.2 step: one-pass adaLN per frame, fused q/k norm+RoPE, no host syncs; Cosmos3: keep the dense attention plan #963 described, never against the previous step.
• A device layout always plans. Comparing one would read it on the host, which is a stream sync.
• A plan that raises leaves nothing to reuse. The remembered layout is cleared before the kernel plan, since a failed plan may already have overwritten the schedule.
• A dropped pre-plan is staged again (_RaggedManagerBase.clear_preplan -> forget_plan). Its H2Ds went on the plan stream and the step that runs instead has no event to wait on, so skipping would read the schedule with nothing ordering it against those copies. This comes from reading the stage order, not from a failure I saw. The block-causal wrapper plans every step and has nothing to forget.
• DenseAttentionManager._build_plan drops _planned and always calls wrapper.plan. It now builds its two pinned layouts per step instead of only on a change, about 6-9 µs on this box.
• The TODO: cache the latest plan state in _plan_layouts is gone.

Performance

This doesn't move a served number, and the gate shows that. Captured steps plan on the plan thread, off the critical path, and a miss there is 0.2-0.3 ms. Where a skip costs real time is the eager FlashInfer path, and Cosmos3 already had that from #963. So this is about having one check instead of two, and every ragged model getting it, more than about speed.

What a skip saves, host time around plan(), median of 200 (B300, FlashInfer 0.7.2.post1):

path	miss	skip
captured, auto	0.22-0.30 ms	3-6 µs
captured, cuDNN	0.09 ms	3 µs
eager, auto	0.75 ms one row, 3.0 ms eight rows	3-7 µs
eager, cuDNN	0.08 ms	3 µs

Gate on 1x B300 (sm_103, 1100 W cap), torch 2.12.1+cu132, FlashInfer 0.7.2.post1, -O2, MSTAR_SHM_ARENA=1, server pinned to 8 cores and client to 4, fresh servers alternated with #963's head, one warm Inductor cache, no compaction stalls or recompile lines in any counted session. Step is DiT ms per step: (compute span − other nodes) / forwards. Wan runs configs/wan22.yaml plus compile_dit_sizes: ["832x480x33"]; Cosmos3-Edge runs the stock config. Both video models at 20 steps.

Three models needed something beyond their stock config to reach this code at all, the same in both arms:

• FLUX.2-klein and Z-Image ship attention_backend: sdpa, which plans nothing, so they run flashinfer here. Graph mode then needs MSTAR_WORKSPACE_BUFFER_MB=2048: FlashInfer's plan wants 866 MB at bs=8 against the 512 MB default.
• LTX-2.5 runs on one GPU with flashinfer_backend: cudnn. Its plan wants 153 MB on B300 against the 128 MB it declares, and that value has no YAML override.

	c=1 step, ms: #963	this PR	c=1 p50, s: #963	this PR	c=4 or 8 p50, s: #963	this PR	c=4 or 8 per min: #963	this PR
Wan	62.4 / 62.3	62.3 / 62.4	2.14 / 2.14	2.14 / 2.17	4.08 / 4.11	4.11 / 4.11	33.0 / 33.1	33.1 / 33.1
Cosmos3-Edge	108.4 / 108.9	108.3 / 108.7	3.65 / 3.65	3.62 / 3.64	11.78 / 11.92	11.94 / 11.88	19.7 / 19.7	19.6 / 19.7
LTX-2.5	242.7 / 241.4	244.2 / 243.3	3.13 / 3.08	3.16 / 3.10	9.26 / 8.87	9.43 / 9.50	25.0 / 25.2	24.9 / 24.8
FLUX.2-klein	46.3 / 45.4 / 45.6	46.0 / 45.9 / 45.3	see below	see below	1.88 / 1.95 / 1.79	1.99 / 1.97 / 1.84	252.9 / 246.9 / 254.9	244.2 / 253.2 / 249.6
Z-Image	62.9 / 63.0 / 62.7	62.9 / 62.7 / 62.8	0.76 / 0.74 / 0.74	0.75 / 0.75 / 0.74	4.56 / 4.57 / 4.46	4.21 / 4.48 / 4.35	107.6 / 100.4 / 103.4	105.9 / 102.3 / 103.1

Image concurrency is 8, video and LTX 4. LTX-2.5's c=4 p50 reads about 4% higher here, with throughput down 1%; two servers per arm can't separate that from its own spread (8.87-9.26 s on #963 alone), and a plan there costs 0.09 ms. Wan and Cosmos3 are the controls here: Wan declares no attention resource, and Cosmos3 already had the check. Both arms ran at the same clocks: c=1 at 1485-1552 MHz and 1049-1057 W on Cosmos3, 1522-1755 MHz and 854-903 W on Wan, 2025-2032 MHz on klein, about 1965 MHz on Z-Image, about 1610 MHz on LTX-2.5.

FLUX.2-klein's c=1 p50 reads worse and isn't. The DiT step is the same, but across servers p50 went 0.39-0.43 s on #963 and 0.42-0.47 s here, in both orders of the pair. That gap is in the text encoder, which is bimodal, about 40 or 100 ms per request regardless of prompt, and on this box's powersave E-cores that's a core at 3.2 GHz or at 800 MHz. Two checks:

FLUX.2-klein c=1	p50, s	slow text encodes
this branch with the skip disabled, 2 servers	0.42 / 0.44	9 / 10 of 20
this branch, same round, 2 servers	0.42 / 0.42	6 / 7 of 20
#963, server cores kept busy at SCHED_IDLE	0.342 / 0.342	0 / 0 of 20
this branch, server cores kept busy at SCHED_IDLE	0.340 / 0.343	0 / 0 of 20

Turning the skip off doesn't close it, and with the governor out of the way the arms are equal. So it's how often a request lands on a parked core, which varies by server, and not this change.

Correctness

Same-seed outputs, against #963, with #963 against itself as the yardstick:

cell	#963 vs itself	this PR vs #963
all five models, c=1	byte-identical	byte-identical: Wan 10/10, Cosmos3-Edge 10/10, klein 20/20, Z-Image 20/20, LTX-2.5 8/8
Wan c=4	20/20	20/20
Cosmos3-Edge c=4	13-20/20, the rest 18-26 dB (4 servers)	13-14/20, the rest 19-20 dB
klein c=8	22/48, the rest 59.67 dB, lowest 57.77	21-27/48, the rest 59.6 dB, lowest 56.98
Z-Image c=8	3/48, the rest 47.33 dB, lowest 22.94	6-9/48, the rest 48.3-55.3 dB, lowest 22.88
LTX-2.5 c=4	2/16, the rest 25.26 dB, lowest 19.51	1-2/16, the rest 26.2-27.5 dB, lowest 18.82

At concurrency a row's output depends on which rows share its forward, so servers differ with or without this change. That batch makeup is the cause is my reading of those numbers; I didn't log it per step. The lowest figures for klein and LTX-2.5 sit just under #963's own, by 0.8 and 0.7 dB, over a handful of files in a spread that wide.

Tests

• ruff check .: clean.
• CPU Core (python -m pytest test/modular test/distributed test/chatterbox test/kokoro test/omnivoice test/sampling_test mstar/model/*/tests benchmark/worker_phases, CUDA hidden, -p xdist -n 8): Wan2.2 step: one-pass adaLN per frame, fused q/k norm+RoPE, no host syncs; Cosmos3: keep the dense attention plan #963's head 50 failed / 8809 passed / 801 skipped; this branch 51 failed / 8845 passed / 801 skipped. 49 failures are shared by test id; the rest are the eval family below. The five test_eval_compare* cases that flipped between trees (they run fake servers under timeouts) all pass on each tree run alone on an idle box. test_piecewise_config_signature and test_tp_async_lockstep_sim ran serially on both and are counted above. test_flex_attention_resource, test_waypoint_components and test_waypoint_dit crash in native code on this box's CPU on both trees and aren't counted.
• New: test_ragged_plan_cache.py, 38 tests on mocked FlashInfer and FA3 wrappers, every manager on this base: reuse, the miss, each slot owning its schedule, captured wrappers, promoted and dropped pre-plans, a plan that raises, a device layout, and a dropped block-causal pre-plan. Mutating the skip, the clear before planning, forget_plan, the device guard and causal in the key each fails it. test_dense_attention.py loses its copy of the reuse test.
• GPU, 1x B300, on 35f012f, real FlashInfer kernels: test_ragged_plan_cache, test_ragged_attention, test_ragged_dense_attention, test_ragged_attention_cpu, test_dense_attention, test_attn_backend_auto, test_cross_attention_plan: 184 passed, 25 skipped, 3 failed. The same 3 fail on Wan2.2 step: one-pass adaLN per frame, fused q/k norm+RoPE, no host syncs; Cosmos3: keep the dense attention plan #963's head (test_piecewise_* in test_ragged_attention, an error of 0.03125 against a 0.03 limit).

Not in this PR

• RaggedBlockCausalWrapper isn't on this base and still plans every step.
• workspace_mb has no YAML override, and on B300 graph-mode FlashInfer wants more workspace than some of these models allow (klein 866 MB at bs=8 against the 512 MB default, LTX-2.5 153 MB against its declared 128 MB). So stock LTX-2.5 can't boot on flashinfer here. That's worth its own fix.

[NSagan271]

LGTM, pending some live testing

[Gaurav-Shah05]

I pushed the two small points on top: the device-layout guard is gone, since every caller builds cu_seqlens on the host, and a minor comment edit.

I tested it with configs/cosmos3_edge.yaml one 832x480, 121-frame, 20-step video at a time, two server sessions per arm, against main at 897eca5:

	main	branch
DiT step	226.3 ms	225.8 ms

Every video is identical between the arms.

[NSagan271]

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human.

Live test: #1076 at 7386585, merged onto main ad0a21e, H200

Verdict: no behaviour change and no regressions on any model I ran. Seeded outputs are as identical between main and the PR as between two main servers, abort storms are clean, and throughput is level, including the Qwen3.5 0.8B CPU-overhead check.

Every model ran in the order main, PR, PR, main, with a fresh server each time and the Rust runtime.

GPU tests (1× H200)

PR tree 191 passed, 25 skipped. Main tree 154 passed, 25 skipped; that's a different set, since the PR adds tests. CPU suites are covered by CI and weren't run here.

Correctness: does skipping the re-plan change outputs?

model (ragged attention user)	result
BAGEL, cfg_parallel	Every phase (solo, concurrent, storm, solo again) was byte-identical across all 4 servers. Repeating a solo request on one server matched 3/6 on every server alike, since the first request takes a different path.
Cosmos3-Edge	Solo, and solo again after the storms, 8/8 identical across all servers. Concurrent 4–6/6 vs solo on both arms (batching nondeterminism).
FLUX.2-klein, attention_backend: flashinfer	Solo vs solo again 6/6 on every server. Across servers 3 of 6 images differ at ~60 dB on both arms, and main's second server was byte-identical to the PR's second server: cross-process noise.
Qwen3-Omni colocated (encoders on ragged attention)	Speech 10/10 on every server. Cross-server text agreement 14–18/20 on both arms alike. Storms clean.
Qwen3.5 0.8B (multimodal) and Qwen3-ASR	Greedy text identical across all servers. Storms clean.

No request failed on any server.

Performance, p50 latency (s), main servers → PR servers

bench	main	PR
Omni a2s	1.59, 1.74	1.86, 1.64
Omni a2t	0.50, 0.48	0.51, 0.46
Omni i2t	2.56, 2.75	2.73, 2.64
Qwen3-ASR	0.20, 0.21	0.20, 0.20
BAGEL t2i	52.9, 53.1	53.4, 53.0

Qwen3.5 0.8B: worker CPU overhead (text tok/s; 128–384 output tokens with ignore-EOS)

	main	PR
t2t c=8	3181, 3094	3115, 3147
t2t c=32	7516, 7320	7596, 7445
t2t c=128	13227, 12973	12328, 13750
i2t c=16	5233, 5058	5339, 5155
i2t c=64	10242, 12358	12365, 12677

Every pair is inside the run-to-run spread. 640/640 requests at c=128 succeeded on every server, with no errors in any server log.

Not run

• FLUX.2-klein throughput: the benchmark runner has no klein model choice. The seeded klein correctness run above covers the attention path.
• Z-Image: no weights on this machine.