Add heads_per_tile (multi-head-per-tile) support to Ulysses ring attention

[edgexyz]

Summary

This PR adds a heads_per_tile (mhpt — multi-head-per-tile) option to the tokamax / Ulysses ring splash-attention forward kernel (_splash_attention_forward_ring_raw). Instead of processing a single attention head per Pallas grid program, the kernel can now process heads_per_tile heads per program, amortizing per-tile overhead (grid launch, mask/scalar prefetch, VMEM pipelining) across several heads. This mirrors the existing mhpt fast-path that already exists for the non-ring custom splash kernel (ulysses_custom), extending the same optimization to the ulysses_ring / tokamax_ring attention modes used by Wan2.2 I2V on TPU v6e.

heads_per_tile is a pure tiling/scheduling choice: for a fixed input it must produce numerically identical results to heads_per_tile=1. The correctness evidence below shows it does (bit-identical output), and the perf motivation is that it lets a v6e-16 run tile 40 DiT heads into 20 (hpt=2) grid programs.

• Branch: ulysses-ring-vmem-passthrough — base main @ f62927f
• Commits (vs main):
  • 049210a — feat(splash-attention): heads_per_tile + configurable vmem_limit_bytes for ulysses ring attention. The feature itself; the block-index detail below is already correct here, and the mhpt invariance unit test ships in this same commit.
  • b47433e — docs(splash-ring): document mhpt kernel + dispatch design rationale.
  • 0d5d0f8 — test(splash-ring): make mhpt heads_per_tile test pass under multi-host.
• Files changed (branch vs main): 5 files, +372 / −25 (feature commit 049210a alone: 4 files, +311 / −25).

Motivation

The Wan2.2 DiT (num_attention_heads = 40) runs self/cross attention through splash attention. On the ring path each head was launched as its own grid program along grid dim 0. Batching multiple heads per program reduces launch/prefetch overhead and improves VMEM reuse. The non-ring custom kernel already had an mhpt path; this PR brings feature parity to the ring kernel so attention=ulysses_ring flash_block_sizes={... "heads_per_tile":N} is usable.

What changed

File	Change
src/maxdiffusion/kernels/splash_attention/splash_attention_kernel.py	New flash_attention_kernel_mhpt (per-tile inner loop over heads_per_tile heads with per-head m/l/o scratch); SplashConfig.heads_per_tile field + validation; _splash_attention_forward_ring_raw now selects the mhpt kernel, sizes the head-dim BlockSpec block to heads_per_tile, adjusts the grid to num_q_heads // heads_per_tile, and appends _hpt<N> to the kernel name.
src/maxdiffusion/max_utils.py	New TokamaxRingFlashBlockSizes hashable carrier and a branch in get_flash_block_sizes so ulysses_ring / tokamax_ring propagate heads_per_tile (plus block sizes) from config.flash_block_sizes.
src/maxdiffusion/models/attention_flax.py	convert_to_tokamax_splash_config now threads heads_per_tile (default 1) into the SplashConfig.

The three files above are the core mechanism. The rest of the branch: the mhpt invariance unit test in ring_attention_kernel_test.py (added in feature commit 049210a) and the assert_allclose_mcjax helper in splash_attention_test_utils.py (0d5d0f8) — see Correctness evidence; the in-code kernel/dispatch rationale docstring (b47433e); and a configurable vmem_limit_bytes bundled into the feature commit (out of scope for this summary).

Feature scope / guards

The mhpt ring path deliberately raises NotImplementedError outside the validated regime, so misconfigurations fail rather than silently miscomputing:

• static ring grids only (no dynamic/active-block grids),
• full MHA only (num_q_heads == num_kv_heads; MQA/GQA rejected),
• static FullMask only (no block/partial masks, no q_sequence, no mask_function),
• no attention sinks, no max-logit bounds, no logits soft-cap,
• HEAD_DIM_MINOR Q/K/V layouts only.

heads_per_tile=1 is the default everywhere and is completely unaffected (the mhpt kernel and block-size changes are only engaged when heads_per_tile > 1).

Key correctness detail — Pallas block-index vs. element-index

The subtlest part of heads_per_tile>1 is a Pallas block-index vs. element-index trap: getting it wrong produces broken output while still exiting cleanly (exit_status=0), so it is documented here as the crux of the correctness story. The kernel in this PR gets it right; the explanation below is the rationale (and a guard against reintroducing it).

In a Pallas BlockSpec, the index map returns a block index; Pallas multiplies it by the block size to get the element offset. The head-dim block size is heads_per_tile, so the index maps must not also multiply the program id — h is already the head-tile index:

# WRONG: head_index(h) = h * heads_per_tile used as the block index
# → element offset = (h * heads_per_tile) * heads_per_tile = h * heads_per_tile**2
q_index_map   = unravel(lambda h, i, j: from_head_minor((head_index(h), i, 0), q_layout))
out_index_map = unravel(lambda h, i, j: (head_index(h), i, 0))
prefix        = (_div(head_index(h), q_heads_per_kv_head),)   # k / v

# RIGHT: h is already the head-*tile* index; the block size is heads_per_tile,
# so the block index is simply h → element offset = h * heads_per_tile
q_index_map   = unravel(lambda h, i, j: from_head_minor((h, i, 0), q_layout))
out_index_map = unravel(lambda h, i, j: (h, i, 0))
prefix        = (_div(h, q_heads_per_kv_head),)               # k / v

Failure mode if done wrong (Wan DiT, num_q_heads=40, heads_per_tile=2, grid dim 0 = 40/2 = 20 programs): program h would write output heads starting at h*2*2 = 4h instead of 2h. Only heads {0,1,4,5,8,9,…,36,37} plus the clamped tail {38,39} would ever be written — 18 of 40 heads would stay at their uninitialized/zero value, with reads scrambled the same way. The pipeline still runs to completion (exit_status=0) but the decoded video is effectively static — which is why this is called out explicitly.

The index maps align with the already-correct reference mhpt kernel in src/maxdiffusion/kernels/custom_splash_attention.py (the ulysses_custom path), whose q_index_map returns (h, i, 0). The q_heads_per_kv_head != 1 guard matches that file's assert num_q_heads == num_kv_heads (mhpt requires no GQA).

Correctness evidence

Verified on a live TPU v6e-16 (ici_context_parallelism=4 ici_tensor_parallelism=4, ulysses_shards=4), Wan2.2-I2V-A14B, running the same generate config three times and changing only the attention setting / code:

Run	attention	heads_per_tile	output md5sum	mp4 size	4-worker exit
Baseline	ulysses_ring	1	c8a0ff4182551867cb0bb75c768c6d72	496,148 B	all 0
This PR	ulysses_ring	2	c8a0ff4182551867cb0bb75c768c6d72	496,148 B	all 0
Wrong block-index (contrast)	ulysses_ring	2	b31f17bf126bfc41844b092e176c741e	5,128 B	all 0

The heads_per_tile=2 output is byte-for-byte identical (same md5) to the heads_per_tile=1 baseline. Because mhpt is a pure tiling optimization, bit-identical output against the single-head baseline is the strongest correctness signal available — it proves the tiling changes scheduling only, not results. The third row is a development-time contrast: with the wrong block-index (see Key correctness detail) the run still exits 0 but produces a differing md5 and a 5 KB, ≈40 B/frame (i.e. static) video — the failure mode the exit code alone hid.

Additional cross-checks performed during the investigation:

• A mask_padding_tokens=False variant of the wrong-index run stayed 5 KB, ruling out segment-id / padding masking as the cause and isolating it to head addressing.
• The XLA trace kernel name carried the expected _hpt2 suffix, confirming the run actually exercised the ring mhpt kernel (not the ulysses_custom path).

Automated unit test — heads_per_tile invariance (multi-host)

RingAttentionHeadsPerTileTest.test_heads_per_tile_matches_single_head (src/maxdiffusion/kernels/splash_attention/ring_attention_kernel_test.py) locks in the invariance this PR relies on: for fixed random MHA inputs it runs the ring kernel with heads_per_tile ∈ {2, 4} and asserts the output matches the heads_per_tile=1 baseline (rtol = atol = 5e-3; ring_size=2, num_heads=4, seq_len=2048, HEAD_DIM=128, bf16, static FullMask).

Run it — on all workers at once (the v6e-16 slice only initializes when every host starts together; a lone host hangs in jax.devices()):

# fill in for your slice:
TPU_NAME=<your v6e-16 node>   ZONE=<its zone>
MAXDIFFUSION_HOME=<path to maxdiffusion checkout>   PYTHON=<tpu venv python>

gcloud compute tpus tpu-vm ssh "$TPU_NAME" --zone="$ZONE" --worker=all \
  --command="cd $MAXDIFFUSION_HOME && \
    $PYTHON -m pytest \
      src/maxdiffusion/kernels/splash_attention/ring_attention_kernel_test.py \
      -k test_heads_per_tile_matches_single_head -q"
# each worker prints: "2 passed, 16 deselected"  (exit 0)

Verified on a live TPU v6e-16, launched on all 4 workers simultaneously (the slice cannot initialize on a single host):

host (process_index 0–3)	result
all 4 hosts	2 passed, 16 deselected · exit 0 · ≈11 s each

Both parameter cases (hpt=2, hpt=4) pass on every host. Making the test green under multi-controller (multi-host) JAX required one fix (commit 0d5d0f8 on ulysses-ring-vmem-passthrough): the test builds its mesh from jax.devices()[:ring_size] (all on process 0), so the stock np.testing.assert_allclose — which pulls the sharded jax.Array back to host — succeeded only on the owning process and raised RuntimeError: ... spans non-addressable devices on the other three, failing the test on all hosts but one. A new SplashAttentionTestCase.assert_allclose_mcjax helper evaluates the comparison on-device, reads the scalar verdict only on the owner, and broadcasts it to every process via multihost_utils.broadcast_one_to_all (one matching collective per process → no deadlock). The test therefore passes on all hosts while still genuinely running the kernel on TPU — no skipTest.

Reproduce (end-to-end video correctness)

Wan2.2-I2V-A14B generate_wan.py on a v6e-16, run once with heads_per_tile=1
(baseline) and once with =2, everything else identical, then compare the produced
mp4s. Launch on --worker=all so all 4 hosts start together (single-host hangs):

# fill in for your slice:
TPU_NAME=<your v6e-16 node>   ZONE=<its zone>
MAXDIFFUSION_HOME=<path to maxdiffusion checkout>   PYTHON=<tpu venv python>   OUT=<output dir>

CFG=src/maxdiffusion/configs/base_wan_i2v_27b.yml
COMMON="height=576 width=1024 num_inference_steps=1 per_device_batch_size=0.0625 \
guidance_scale_low=1.0 guidance_scale_high=1.0 \
ici_context_parallelism=4 ici_tensor_parallelism=4"
BLK() { echo "{\"block_q\":1024,\"block_kv\":1024,\"block_kv_compute\":256,\"block_q_dkv\":1024,\"block_kv_dkv\":1024,\"block_kv_dkv_compute\":256,\"block_q_dq\":1024,\"block_kv_dq\":1024,\"heads_per_tile\":$1}"; }

for N in 1 2; do
  gcloud compute tpus tpu-vm ssh "$TPU_NAME" --zone="$ZONE" --worker=all \
    --command="cd $MAXDIFFUSION_HOME && \
      $PYTHON src/maxdiffusion/generate_wan.py $CFG $COMMON \
        attention=ulysses_ring flash_block_sizes='$(BLK $N)' \
        run_name=hpt$N output_dir=$OUT/hpt$N"
done

# worker-0's mp4 from each run must be byte-identical:
md5sum "$OUT"/hpt1/**/wan_output_0_0.mp4 "$OUT"/hpt2/**/wan_output_0_0.mp4

Smoke config used for the evidence above: 576x1024, num_inference_steps=1,
per_device_batch_size=0.0625, guidance_scale_low=guidance_scale_high=1.0, the
flash_block_sizes shown (with heads_per_tile = 1 or 2). Cold compile 213–423 s, warm
generation ≈15.6 s/step, all 4 workers exit_status=0. (Internally these runs were driven
by the run_generate_wan_profiler.sh launcher, which wraps exactly the generate_wan.py
invocation above and fans it out to every worker.)

Scope, limitations & follow-ups

• Forward only. This is the inference/generate forward kernel; the mhpt fast-path is not wired into the backward kernel (fine for generation, which is all ulysses_ring serves here).
• MHA + static FullMask only (enforced by the guards above).
• Unit test (mhpt invariance). test_heads_per_tile_matches_single_head asserts mhpt(N) == mhpt(1) for N ∈ {2, 4}, guarding against exactly the block-index class of error in Key correctness detail. It ships in the feature commit (049210a) and was made multi-host-safe in 0d5d0f8 (the new assert_allclose_mcjax helper); it passes on a live TPU v6e-16 across all 4 hosts (see Automated unit test) — i.e. it runs on the real cluster, not just interpret=True.

[Perseus14]

@eltsai @csgoogle Could you please review this?

[csgoogle]

Do you think it's better to rename the CustomFlashBlockSizes and reuse that for all code paths

[edgexyz]

Both carriers existed for the same reason — the upstream jax splash_attention_kernel.BlockSizes is frozen + slotted and can't hold the extra fields (block_kv_compute_in for custom; block_*_dkv / heads_per_tile / vmem_limit_bytes for ring), and a plain dict isn't hashable.

Merged TokamaxRingFlashBlockSizes into CustomFlashBlockSizes (field union) and dropped the separate class; the ring path now returns CustomFlashBlockSizes too. Both consumers read by field name via getattr, so no consumer changes were needed — only the two returns in get_flash_block_sizes, which keep their own populate logic (custom uses lenient .get(...); ring requires the dkv keys and forces use_fused_bwd_kernel=True). Done in 0cd23dd.

[csgoogle]

Isn't it better to do all gather and verify it?

[edgexyz]

The test now builds the mesh from all devices (ring_size = jax.device_count()), so the outputs are sharded across every process; both results are then gathered with multihost_utils.process_allgather(..., tiled=True) and compared with the standard assert_allclose. The bespoke assert_allclose_mcjax helper is deleted.

Verified on a live v6e-16 (4 hosts x 4 chips, launched with --worker=all): all 4 hosts report 2 passed (hpt=2 and hpt=4 vs the hpt=1 baseline) in ~34s each. Done in 758cb1f.

[csgoogle]

LGTM, please squish all the commits

The merge-base changed after approval.

[csgoogle]

to make code more readable, you can move this to a method

[csgoogle]

To make tests more robust you can also add tests cases for negative scenarios like MQA, GQA which you mentioned. Also minor nit, keep names more descriptive, like instead of run use may be run_multi_head_kernel