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.