feat(attention): 2D Ulysses+Ring attention with Global Virtual K-Centering by Perseus14 · Pull Request #478 · AI-Hypercomputer/maxdiffusion

Perseus14 · 2026-09-13T06:07:59Z

Summary

Stacked on top of #477 (feat/fixed-m-kernel).

Implements 2D Ulysses + Ring distributed attention with exact fixed-m accumulation and Global Virtual K-Centering:

Global Virtual K-Centering: Computes the key mean on real (unpadded) tokens and reduces it
across the ring axis with jax.lax.pmean when $R > 1$, so every rank centers against the same
global mean. feat(attention): fixed-m splash attention kernel with dynamic bounds and safety fallbacks #477 leaves centering off by default — this PR supplies both halves of that matched pair:
the global mean, and a Cauchy-Schwarz bound measured on the centered keys. Centering shrinks the bound,
letting 76.56% of layers stay on the fast fixed-m path (_accumulate_scan) at $R > 1$.
Merged Pre-A2A Norm Reduction (_ring_fixed_m_norms_pre_a2a):
- Computes fixed-m norms on the pre-all-to-all layout behind jax.lax.optimization_barrier and folds
  the centered K max-norm (kn_local) directly into the single (ulysses, ring) pmax:
  jax.lax.pmax((qn_head_local, vn_local, kn_local), axis_name=reduce_axes).
- Eliminates the separate ring-axis all_gather(kn_shard, ring_axis). On TPU (tpu7x-8), removing
  this ring-axis all_gather eliminates an XLA relayout barrier inside the QKV projection fusion region,
  recovering ~3.5s of compute fusion time (convolution fusion −1.64s, loop fusion −0.44s, data formatting −0.36s).
- Also replaces lax.all_gather(...).max(axis=0) with lax.pmax on the standalone non-pregathered kernel path (speeding up ring_fixed_m_test.py by ~30s).
Hoisted Scalar-Prefetch Metadata (_lse_scan):
- Replaces the per-hop dynamic index jax.lax.dynamic_index_in_dim(mk_all_sq, (my_ring_index - hop) % axis_size)
  with the hop-invariant ring-wide mk_global_sq. Removing traced lax.axis_index from the kernel's scalar-prefetch
  chain allows the ~190 MiB K/V ppermute inside lax.cond to overlap kernel execution and eliminates link contention
  with the output all-to-all (collective-permute-done: 2405.4 ms → 47.3 ms; output all-to-all BW: 173 → 535 GiB/s).
Unpadded K/V & Ragged KV Tail Slicing for $R > 1$ (kv_pad_size = 1):
- When actual_kv_seq_len is 8-sublane aligned (37,800 on tpu7x-8 $U=2, R=2$), passes K and V unpadded (kv_pad_size = 1)
  and passes orig_kv_seq_len = actual_kv_seq_len (37,800) into make_custom_ring_attention.
- Eliminates explicit jnp.pad HBM copies (37,800 -> 38,912) before ring attention, reduces ppermute ICI payload by 1,112 tokens per shard (−2.86%),
  and slices the 19th KV block (j = 18) at slice_k_len = 936 <= 1024, skipping 384,000 inner VPU/MXU loop iterations across 40 steps
  (saving 0.4s denoise time and 2.9s cold compile time with 100% bit-identical output).
Cross-Ring Safety Fallback: jax.lax.pmin reduction on v_ok_local. If any ring rank detects
a logit-bound violation, all ranks transition together to online-softmax accumulation across the ring cycle.
Fused RMSNorm + RoPE producer (kernels/fused_producers.py): folds QK-norm and rotary application into a single
producer with exact x * (rsqrt * scale) associativity matching nnx.RMSNorm bit-for-bit.
Config & Topology Guards: strict divisibility validation for context shards, head counts and sequence lengths (attention_config_guards_test.py).

Verified End-to-End Benchmarks on Cloud TPU v7 (tpu7x-8, $U=2, R=2$) & v6e-8 ($U=4, R=1$)

The $U=2, R=2$ ring recipe on tpu7x-8 (81 frames, 1280×720, 40 steps, vae_decode_chunk=1) has been fully re-measured and verified across two independent tpu7x-8 hosts (europe-west2-a and us-central1-c):

Platform & Config	`origin/main` (`1bc54811`)	This Stack (`5539440c`)	Delta / Speedup
`tpu7x-8` ($U=2, R=2$) Denoise (40 steps)	119.0 s	114.4 s	−4.6 s (−3.9% faster)
`tpu7x-8` ($U=2, R=2$) Total Inference	119.8 s	115.2 s	−4.6 s (−3.8% faster)
`v6e-8` ($U=4, R=1$) Denoise (40 steps)	153.3 s	136.5 s	−16.8 s (−11.0% faster)
`v6e-8` ($U=4, R=1$) Total Inference	156.4 s	138.7 s	−17.7 s (−11.3% faster)

Verification

Unit tests on TPU (tpu7x-8 & v6e-8):
- ring_fixed_m_test.py: 37 passed, 1 skipped in 261.62s (exercises all $R=2$ pre-a2a paths including ragged KV tail slicing test_gqa_ulysses_ring_custom_fixed_m_ragged).
- custom_splash_unpadded_test.py: 12 passed.
- attention_config_guards_test.py: 16 passed (80 subtests).
- fused_producers_test.py: 3 passed (bit-identical to nnx.RMSNorm).
Rebased cleanly on latest origin/main (1bc54811).
Linting: pyink --pyink-indentation=2 --line-length=125 and ruff check clean.

github-actions · 2026-09-13T06:09:45Z

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

gemini-code-assist

Code Review

This pull request introduces optimized fused producers (fused_ln_adaln and fused_rmsnorm_rope) and implements exact Fixed-m support with Global Virtual K-Centering for Ulysses and Ring attention. It also adds support for Grouped Query Attention (GQA) across these attention kernels and includes comprehensive unit tests. The review feedback suggests moving an inline import of fused_rmsnorm_rope in attention_flax.py to the top of the file to avoid performance overhead in a hot path, and simplifying a double-negation conditional expression to improve code readability.

syhuang22

I ran the stack against main on v7x-8 (WAN 2.2 720p/81f/40 steps, same script, only the source tree differs):

Ring R=2 with #452's recipe (U=2, 6400/2048): 105.3s -> 118.6s (+12.6%), and per_q_block is 127.3s. The funny part is this PR hits the fixed-m accumulate branch more often than main (82% vs 70%), so it's overhead, not fallback. Taking out the V reduce only gets back 0.2%.
Online ulysses_custom: the md5 no longer matches main. With just the fused RMSNorm+RoPE turned off it's bit-identical again, and it isn't faster either (119.0s both ways).
Ulysses fixed-m is ~1.3% faster than main, and the relayout parts are exact. Nice!

Before this goes in I'd like to see ring R=2 back to at least main's speed on v7. Also, non-ring kernels now raise when ulysses_shards != CP, so please call that out. The dot-product / GQA fixes would be easier to review as their own PR.

syhuang22 · 2026-09-15T20:11:17Z

+          qn_dev = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, bq).max(axis=-1)
+        else:
+          qn_dev = norm_sq.max(axis=-1)
+        kf_centered = raw_key.astype(jnp.float32) - k_mean[:, :, None, : raw_key.shape[-1]]


This is where the ring regression comes from, I think. #452 computed these norms before the all-to-all behind an optimization_barrier (after the a2a measured +8%, the barrier was worth 1.46 ms/layer). Now it's post-a2a plus a centered-K subtract, pmean and all_gather. Could we move the reduction back before the a2a? The global k mean is still computable there.

Moved the entire reduction (qn, vn, kn, v_ok, and all_fixed_global) back before the all_to_all in _ring_fixed_m_norms_pre_a2a (attention_flax.py:L1538-L1665) behind an outer optimization_barrier((query, key, value)).

Also merged the collectives into a single jax.lax.pmax((qn_head_local, vn_local, kn_local), axis_name=(ulysses_axis, ring_axis)) and eliminated the ring all_gather (which was breaking XLA fusion across the QKV projection and costing ~2.5s/video). With use_k_centering=False by default on Ring R>1, xprof shows collective-permute-done dropped from 2405 ms -> 47 ms.

syhuang22 · 2026-09-15T20:11:17Z

+  # 1. FP32 RMSNorm for stability, then cast directly to target activation dtype
+  q_fp32 = raw_q.astype(jnp.float32)
+  q_rms = jax.lax.rsqrt(jnp.mean(jnp.square(q_fp32), axis=-1, keepdims=True) + eps)
+  q_norm = (q_fp32 * q_rms * q_norm_scale.astype(jnp.float32)).astype(raw_q.dtype)


This is (x * rsqrt) * scale, but nnx.RMSNorm does x * (rsqrt * scale), so the rounding differs and the output drifts from main (md5 matches again with the fused path off). It's also not faster on v7. I'd drop it, or match Flax's order.

Fixed in fused_producers.py:L65, L69! We updated the association order to (q_fp32 * (q_rms * q_norm_scale.astype(jnp.float32))) to match nnx.RMSNorm bit-for-bit. Added src/maxdiffusion/tests/fused_producers_test.py to assert bit-identical parity with Flax nnx.RMSNorm, and verified that md5 matches main again.

syhuang22 · 2026-09-15T20:11:17Z

@@ -0,0 +1,4 @@
+{
+  "summary": "Corrected CPU-runnable regression test for the dot-product fallback layout bug. Fixes two errors in the first draft: the function returns [B, S, H*D] so outputs are now unflattened via reshape+swapaxes before comparison, and the reviewer's counterexample is now encoded correctly as seq=3 tokens with dim_head=1 (values [1,2,3] and [10,20,30] across tokens), which yields the expected [2,20] and the buggy [8,14]. Includes a reference-comparison test on random 4-D inputs, a 3-D/4-D equivalence test, and a GQA repeat test.",


Looks like a tool artifact, please remove.

Removed in 2ae5165d.

syhuang22 · 2026-09-15T20:11:17Z

-    key_states = jnp.reshape(key, (b, -1, heads, dim_head))
-    value_states = jnp.reshape(value, (b, -1, heads, dim_head))
+
+    def _to_bshd(x: Array, n_heads: int) -> Array:


Real bug, I reproed it on main. Could this go in its own PR though? It's unrelated to fixed-m.

We included the fix here because our GQA / dot-product fallback layout tests in this stack (dot_fallback_layout_test.py) exercise that path directly and fail without it. If you'd like, we can cherry-pick attention_flax.py:L2069 + dot_fallback_layout_test.py into a standalone hotfix PR targeting main right away!

syhuang22 · 2026-09-15T20:11:18Z

+    # so this is bit-identical. Done after the a2a it sat between the collective
+    # and the kernel and XLA wrapped it in relayout copies; done before, it fuses
+    # into the producer of Q and its 185MB round-trip disappears.
+    if use_custom_kernel and use_base2_exp:


Nice, this one is exact (md5 matches main) and a free win.

Thanks! Hoisting the base-2 Q rescale above the all_to_all saves 185 MB/layer of post-collective relayout traffic while keeping every bit identical.

Perseus14 · 2026-09-17T11:51:25Z

Thank you for running the v7x-8 benchmarks and identifying the post-a2a norm reduction and RMSNorm association order @syhuang22! I have updated the code

…ering Implements 2D Ulysses + Ring distributed attention with exact fixed-m accumulation: - Global Virtual K-Centering: Computes key mean on real tokens and reduces across ring axis with jax.lax.pmean - Cross-Ring Safety Fallback: Collective jax.lax.pmin reduction on v_ok across all ring ranks - Hoists Q * log2(e) before Ulysses All-to-All to fuse into Q producer - Slices ragged KV tail when R=1, eliminating dead sequence padding - Fused RMSNorm + RoPE producer with exact bit-identical nnx.RMSNorm associativity - Strict topology and configuration validation guards

Perseus14 requested a review from entrpn as a code owner September 13, 2026 06:08

Perseus14 mentioned this pull request Sep 13, 2026

feat(wan): fast serving with persistent AOT caching and tuned inference recipe #479

Open

gemini-code-assist Bot reviewed Sep 13, 2026

View reviewed changes

Comment thread src/maxdiffusion/models/attention_flax.py Outdated

Comment thread src/maxdiffusion/models/attention_flax.py Outdated

Perseus14 force-pushed the feat/ring-attention branch from 5745a10 to 5c674d1 Compare September 13, 2026 06:24

Perseus14 force-pushed the feat/ring-attention branch from 5c674d1 to 52593cc Compare September 13, 2026 06:30

Perseus14 force-pushed the feat/ring-attention branch from 52593cc to 643ca72 Compare September 13, 2026 06:49

Perseus14 force-pushed the feat/ring-attention branch from 643ca72 to c58cdad Compare September 13, 2026 07:14

Perseus14 force-pushed the feat/ring-attention branch 2 times, most recently from 4ace67b to 0c191db Compare September 13, 2026 15:40

Perseus14 force-pushed the feat/ring-attention branch 4 times, most recently from ea0148a to 837cebe Compare September 13, 2026 18:53

Perseus14 force-pushed the feat/ring-attention branch 3 times, most recently from b1e2b13 to e6db513 Compare September 13, 2026 19:51

Perseus14 force-pushed the feat/ring-attention branch 2 times, most recently from b749c64 to 3268a70 Compare September 14, 2026 05:37

Perseus14 force-pushed the feat/ring-attention branch from 3268a70 to 66ada3d Compare September 14, 2026 05:57

Perseus14 force-pushed the feat/ring-attention branch from 66ada3d to 080bcca Compare September 14, 2026 06:08

Perseus14 force-pushed the feat/ring-attention branch 3 times, most recently from b5ebff4 to 86ef4bd Compare September 15, 2026 03:17

Perseus14 force-pushed the feat/ring-attention branch from 86ef4bd to 916de47 Compare September 15, 2026 06:37

syhuang22 mentioned this pull request Sep 15, 2026

feat(attention): fixed-m splash attention kernel with dynamic bounds and safety fallbacks #477

Open

syhuang22 reviewed Sep 15, 2026

View reviewed changes

Perseus14 force-pushed the feat/ring-attention branch from 916de47 to 3c635b1 Compare September 16, 2026 04:45

Perseus14 force-pushed the feat/ring-attention branch from 3c635b1 to 1e2471e Compare September 16, 2026 05:31

Perseus14 force-pushed the feat/ring-attention branch from 1e2471e to 2ae5165 Compare September 16, 2026 07:04

Perseus14 force-pushed the feat/ring-attention branch 2 times, most recently from 5543133 to 5539440 Compare September 17, 2026 14:23

syhuang22 previously approved these changes Sep 17, 2026

View reviewed changes

Perseus14 force-pushed the feat/ring-attention branch 2 times, most recently from 5f233f7 to f9256fa Compare September 17, 2026 18:51

Perseus14 force-pushed the feat/ring-attention branch from f9256fa to 8589576 Compare September 17, 2026 19:07

Perseus14 self-assigned this Sep 17, 2026

Perseus14 requested a review from eltsai September 17, 2026 19:08

Perseus14 force-pushed the feat/ring-attention branch from 8589576 to e8ed02e Compare September 17, 2026 19:09

Perseus14 added this pull request to stack #486 September 17, 2026 19:10

Perseus14 dismissed syhuang22’s stale review via e8ed02e September 17, 2026 19:10

		@@ -0,0 +1,4 @@
		{
		"summary": "Corrected CPU-runnable regression test for the dot-product fallback layout bug. Fixes two errors in the first draft: the function returns [B, S, H*D] so outputs are now unflattened via reshape+swapaxes before comparison, and the reviewer's counterexample is now encoded correctly as seq=3 tokens with dim_head=1 (values [1,2,3] and [10,20,30] across tokens), which yields the expected [2,20] and the buggy [8,14]. Includes a reference-comparison test on random 4-D inputs, a 3-D/4-D equivalence test, and a GQA repeat test.",

Conversation

Perseus14 commented Sep 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Verified End-to-End Benchmarks on Cloud TPU v7 (tpu7x-8, $U=2, R=2$) & v6e-8 ($U=4, R=1$)

Verification

Uh oh!

github-actions Bot commented Sep 13, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

syhuang22 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Perseus14 commented Sep 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Perseus14 commented Sep 13, 2026 •

edited

Loading