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:

Global Virtual K-Centering: Computes local key mean after Ulysses All-to-All on real unpadded tokens, reducing via cross-ring jax.lax.pmean when $R > 1$ for uniform head centering.
Cross-Ring Safety Fallback: Implements jax.lax.pmin reduction on v_ok_local. If any ring rank detects logit bound violation, all ranks safely transition to online softmax accumulation across the ring cycle.
Ring Relayout Fusion & Unpadded KV:
- Hoists $Q \cdot \log_2(e)$ before the Ulysses All-to-All in wrap_ulysses_ring_attention, fusing the operation into the Q producer.
- Slices ragged KV tail when $R=1$, eliminating dead sequence padding and saving 370 MB of HBM traffic per layer.
Config & Topology Guards: Enforces strict divisibility validation for context shards, head counts, and sequence lengths in attention_config_guards_test.py.

Verification

Unit Tests: 28/28 passed in ring_fixed_m_test.py (including two-rank adversarial centered keys 15/17 regression coverage) and 16/16 passed in attention_config_guards_test.py.
E2E Parity: Produces bit-identical video output to the baseline (4c18d4ce340e0da3a1d4c617926ad525) for tested generation while improving step time by -37.5 ms/step.
Linting: 100% compliant with pyink --pyink-indentation=2 --line-length=125 and ruff check.

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.

gemini-code-assist · 2026-09-13T06:12:16Z

-          query_proj = _unflatten_heads(query_proj, self.heads)
-          key_proj = _unflatten_heads(key_proj, self.heads)
+      if rotary_emb is not None and self.qk_norm and is_self_attention:
+        from maxdiffusion.kernels.fused_producers import fused_rmsnorm_rope


Inline imports inside frequently executed methods (like __call__ of an attention block) can introduce unnecessary overhead. It is highly recommended to move from maxdiffusion.kernels.fused_producers import fused_rmsnorm_rope to the top of the file.

gemini-code-assist · 2026-09-13T06:12:16Z

+        if self.qk_norm:
+          with self.conditional_named_scope("attn_q_norm"):
+            query_proj = self.norm_q(query_proj)
+          if not (not is_self_attention and cached_kv is not None and "text" in cached_kv):


The conditional expression contains a double negation (not (not is_self_attention ...)), which reduces readability. Simplifying it to use positive checks makes the logic much clearer and easier to maintain.

Suggested change

if not (not is_self_attention and cached_kv is not None and "text" in cached_kv):

if is_self_attention or cached_kv is None or "text" not in cached_kv:

…ering Implements 2D Ulysses + Ring attention with exact fixed-m accumulation: - Global Virtual K-Centering via cross-ring pmean(k) - Cross-ring v_ok reduction via jax.lax.pmin preventing unsafe fixed-m hops - Per-Q-block fixed-m accumulation and pre-gathered mk norms - Full GQA support across arbitrary Ulysses and Ring shard configurations - Multi-device hardware regression test suite including adversarial V tests Config guards ------------- Two config mistakes were previously accepted in silence and produced results that looked fine but did not mean what the config said. 1. Degenerate ring. The ring variants derive R = context_shards / ulysses_shards. Setting U == CP yields R=1: no KV is rotated, the cross-ring mean for virtual K-centering is skipped, and the output is mathematically identical to the corresponding non-ring ulysses_custom* kernel. This is easy to hit (CP=4, ulysses_shards=4) and the run is indistinguishable from a real ring result unless you diff the output bytes against the non-ring kernel. _warn_if_ring_is_degenerate now logs once per distinct config and, where one exists, names the largest ulysses_shards that divides both the mesh and the head counts to give a genuine ring. 2. Ignored ulysses_shards. The non-ring _ulysses_attention path always shards over the full context axis, so its Ulysses degree is implicitly CP. A user-supplied ulysses_shards that disagreed was silently dropped; _validate_implicit_ulysses_degree now raises instead. -1 (unset) and a value already equal to CP are still accepted, so existing configs are unaffected. attention_config_guards_test covers CP in {1,2,4,8,16,32,64} against Wan's 40 heads, which is not a power of two and so exercises the divisibility logic properly. Also registers ulysses_ring_custom_fixed_m_per_q_block in the tile-size grid search kernel set; it was the only ring kernel missing, so the optional search would have tuned it via the non-ring code path. Relayout optimisations ---------------------- Three bit-identical wins in _ulysses_attention, each measured separately. 1. Unpadded V into the fixed-m metadata reduction (-1.3% end to end). _compute_fixed_m_metadata was reading the padded V, putting a 193MB pad plus its reduction on the critical path, chained behind the V all-to-all and leaving that collective fully exposed. The reduction feeds `all_fixed`, which gates the whole kernel through a lax.cond, so everything upstream of it is on the critical path by construction. The padding is zeros and the statistic is a max of squares, so the padded rows cannot change the result. 2. Hoist `query * LOG2E` above the all-to-alls (-12.96 ms/step). A scalar elementwise multiply commutes exactly with an all-to-all, which is pure data movement. Applied after the collective it sat between the collective and the kernel and XLA wrapped it in relayout copies; applied before, it fuses into the producer of Q and its 185MB round-trip disappears. The use_custom_kernel guard matters: the non-custom path must not see a pre-scaled Q, and because every change here is bit-identical by design, an output hash would not catch that mistake. 3. Stop padding K/V to the block size (-23.22 ms/step). The fixed-m kernel *slices* the ragged KV tail rather than masking it -- see last_compute_body_fixed in custom_splash_attention.py, where slice_k_len derives from the unpadded kv_seq_len -- so it never reads a padded K/V row. Materialising the pad cost 2 x 185MB of HBM traffic per layer for nothing. Passing flash_block_size=1 makes only the sequence pad a no-op while preserving the head_dim->128 pad, the reshape and the (tensor, kv_size, seq_len) return contract, so configs with head_dim < 128 are unaffected. Note the asymmetry: the same trick applied to Q is a large regression (+65.8 ms/step measured) and is deliberately not done. The Q pad is load-bearing for _compute_fixed_m_metadata, which needs q_len to divide block_q exactly. custom_splash_unpadded_test guards item 3 directly: it runs the kernel twice with identical logical inputs, once padded and once not, and asserts exact equality. Coverage includes the fixed-m, hybrid and online paths, bkv_compute < bkv ragged tails, and repeated invocations that re-poison the padding region with NaN between runs, so a kernel that started reading the tail would fail rather than quietly return stale values. Measured on Wan2.2-T2V-A14B, 40 steps, 720x1280, 81 frames, CP=4/DP=2 on v6e-8. Counterbalanced ABBA design, n=4 pairs, paired t intervals: items 2+3 stacked -36.00 ms/step, 95% CI [-42.5, -29.5] Denoise 139.8s -> 136.456s overall, byte-identical video throughout (md5 513c32812f6daab717ffd780574360e2).

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

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

Perseus14 force-pushed the feat/fixed-m-kernel branch from 57a4676 to e2b4cf8 Compare September 13, 2026 06:29

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

Perseus14 force-pushed the feat/fixed-m-kernel branch from e2b4cf8 to 304dd39 Compare September 13, 2026 06:49

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

Perseus14 force-pushed the feat/fixed-m-kernel branch from 304dd39 to c1c87a2 Compare September 13, 2026 07:10

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

Perseus14 force-pushed the feat/fixed-m-kernel branch from c1c87a2 to c557008 Compare September 13, 2026 07:40

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/fixed-m-kernel branch from c557008 to 49c3b6c Compare September 13, 2026 16:21

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/fixed-m-kernel branch from 49c3b6c to 9731972 Compare September 13, 2026 19:09

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/fixed-m-kernel branch from 9731972 to 5317aca Compare September 13, 2026 20:16

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/fixed-m-kernel branch from 5317aca to 4c47d80 Compare September 14, 2026 05:57

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

Perseus14 force-pushed the feat/fixed-m-kernel branch from 4c47d80 to 3b9d7aa Compare September 14, 2026 06:08

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 from 080bcca to cf7d6e0 Compare September 14, 2026 07:49

Perseus14 force-pushed the feat/ring-attention branch from cf7d6e0 to b5ebff4 Compare September 14, 2026 09:39

Navigation Menu

Uh oh!

feat(attention): 2D Ulysses+Ring attention with Global Virtual K-Centering - #478

feat(attention): 2D Ulysses+Ring attention with Global Virtual K-Centering#478
Perseus14 wants to merge 1 commit into
feat/fixed-m-kernelfrom
feat/ring-attention

Perseus14 commented Sep 13, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 13, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Sep 13, 2026

Uh oh!

gemini-code-assist Bot Sep 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

	if not (not is_self_attention and cached_kv is not None and "text" in cached_kv):
	if is_self_attention or cached_kv is None or "text" not in cached_kv:

Uh oh!

Conversation

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

Uh oh!

Summary

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!

gemini-code-assist Bot Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Perseus14 commented Sep 13, 2026 •

edited

Loading