feat(attention): fixed-m splash attention kernel with dynamic bounds and safety fallbacks by Perseus14 · Pull Request #477 · AI-Hypercomputer/maxdiffusion

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

Summary

Implements the single-device fixed-m splash attention kernel using Pallas on Cloud TPU.

Instead of tracking an online-softmax running max per KV block, eligible (head, Q-block) pairs
subtract a precomputed shift derived from a Cauchy-Schwarz bound on the logits, which lets the
numerator and denominator accumulate directly in FP32.

Dynamic eligibility bounding. The two constants are derived from the actual KV length rather
than pinned to one operating point:
$$C(N) = 127 - \lceil\log_2 N\rceil - \lceil\log_2 V_{\max}\rceil, \qquad W(N) = C(N) + 125, \qquad \text{gate} = \left\lfloor \frac{W(N)}{2} \right\rfloor$$
With the shift $m_i = \lceil U_i \rceil - C(N)$ built from the Cauchy-Schwarz bound
$U_i = \max_i \lVert q_i \rVert \max_j \lVert k_j \rVert$, every shifted exponent satisfies
$z_j - m_i \ge C(N) - (U_i + \lceil U_i \rceil) \ge -125 > -126$, so no term flushes to subnormal
and no accumulator overflows.
Virtual K-centering, opt-in. Centers logits via the register-level projection
$q_i^\top \bar{k}$ without ever writing $(K - \bar{k})$ back to HBM. This shrinks $U_i$ (and so
admits more heads); it does not loosen the gate — see below. On the ring path it activates
only when the caller supplies k_mean; see the next section for why that is load-bearing.
Ragged sequence handling. Dynamic tail slicing in last_compute_body_fixed lets physically
unpadded KV sequences run without out-of-bounds access.
Relayout optimization. Hoists the base-2 logit rescale $Q \cdot \log_2 e$ above the Ulysses
all-to-all so XLA fuses the scalar multiply into the upstream projection, saving 185 MB of HBM
traffic per layer.
Input validation & API hygiene. Explicitly rejects legacy 2D mk arrays ([H, num_q_blocks]) at entry with a clear ValueError, preserves keyword-only *, delimiter in make_custom_ring_attention, and derives head-dim padding dynamically from query.shape[-1].

This PR is additive: `main`'s $R > 1$ behaviour is unchanged

Centering and the Cauchy-Schwarz bound are a matched pair. If the kernel exponentiates centered
logits $q \cdot (k_j - \bar{k})$ while the caller's eligibility bound was built from raw $K$, the
bound caps the wrong quantity — $\lVert k \rVert$ does not bound $\lVert k - \bar{k} \rVert$ — and a
tile can clear the gate while the logit it was supposed to cap overflows fp32.
RingRawKeyBoundUnsoundTest pins that down in pure arithmetic: keys all of norm 100 give a raw bound
of 100 against a gate of 116, while the centered max logit is ~200 and the shifted exponent exceeds
$2^{128}$.

This revision makes centering opt-in: the kernel centers only when handed a k_mean.
Callers that supply one also supply centered norms; callers that do not keep the uncentered path,
which is sound against the two-sided $\lfloor W/2 \rfloor$ bound. main's $R > 1$ caller passes no
k_mean, so it is left running exactly as before and the restriction is gone. Net effect: this PR
adds the fixed-m kernel without changing any behaviour main already has, and merges safely on its
own. #478 then turns centering on by supplying both halves of the pair.

Why the gate is halved, and why `is_ring` does not change it

Review raised that get_fixed_m_constants accepts is_ring but ignores it, and that the gate is
half the legacy _FIXED_M_SAFE_BOUND = 213. Both observations are correct; the conclusion that the
centered path should get the full window $W$ is not.

The appealing argument is that K-centering forces the realized row max $M \ge 0$, so only one side
needs absorbing. It fails because the shift is built from $U$, not from $M$: the worst case stays
$C - (U + \lceil U \rceil)$ whether or not the keys are centered. Wiring is_ring up to return $W$
for the Ulysses path reproduces the exact bug the halving was introduced to fix — at $N = 4096$,
$W = 232$, so a query of norm $231.01$ against exactly-centered keys passes the gate and then
silently loses 11.8% of the softmax mass (0.882 vs 1.0) as its negative logits flush to zero.
test_adversarial_centered_keys_softmax_mass_loss and test_cpu_proof_invariant_bounds both pin
this down.

Measured cost of the strict gate. Instrumented over a full WAN 2.2 720p generation on v6e-8
(1280 metadata calls per gate, $N = 75{,}600$):

gate	derivation	entries falling back	calls with any fallback
113	$\lfloor W/2 \rfloor$ (shipped)	0.73%	112 / 1280
201	relative mass loss $\le 2^{-10}$ incl. the $N$ multiplicity	0.45%	56 / 1280
218	relative mass loss $\le 2^{-10}$ per term	0.44%	56 / 1280
227	$W$	0.43%	54 / 1280

So the strict gate costs about 0.3 percentage points of extra fallback over the loosest alternative.

Verification

custom_splash_fixed_m_test.py: 23 passed on TPU (dynamic bounds, non-divisible sequences, legacy 2D mk rejection, adversarial centered-key regression, hybrid fallbacks).
ring_fixed_m_test.py: 36 passed, 1 skipped on TPU.
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:16Z

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

gemini-code-assist

Code Review

This pull request introduces dynamic fixed-m constants and safe bounds calculation based on KV sequence length, adds support for virtual K-centering, and enhances input validation and dtype safety checks in the Pallas flash attention kernel. It also expands the test suite to cover various edge cases, including per-Q-block fallback, batched isolation, and virtual K-centering. The review feedback focuses on optimizing the TPU kernel performance by replacing expensive dynamic integer division with optimized BlockSpec mapping and introducing a static use_k_centering flag to conditionally compile the centering logic at trace time.

syhuang22

A few things before this can go in though:

If this lands on its own, fixed-m gets turned off for R>1 (use_fixed_m = False in _ulysses_ring_custom_attention, it only comes back in #478). That breaks the recipe we ship today. Can we split the stack so each PR is safe to merge by itself?
Cutting the Ulysses gate from 213 to 113 feels more conservative than we need. With centered keys, the mass you can lose is bounded by 2^(ceil(U)-C-126), so a gate around C+116 already keeps the loss under 0.1%. Could you share fallback rates so we can pick the number?
mk[0] means something different now, but a 2D mk still gets silently broadcast. Can we just raise on that?
The V check never fired in my runs on WAN 2.2 (every call passed), and it only exists because C moved from 88 to 102. Is it worth the extra complexity?
Please put back the comments that explain the Mosaic cliff and k-smoothing. They save the next person a lot of pain.

syhuang22

Some line-level notes to go with my comment above.

syhuang22 · 2026-09-15T20:08:49Z

+        "(R = context_shards // ulysses_shards); falling back to online softmax. "
+        "Set ulysses_shards == ici_context_parallelism for R=1 to use fixed-m."
+    )
+    use_fixed_m = False


This turns off fixed-m for R>1, which is the ring2 x uly2 recipe we ship today. #478 turns it back on, so let's not land this one alone.

Fixed in 44dfe632 — removed this use_fixed_m = False override and made K-centering opt-in (k_mean is not None) on the ring path, so R > 1 continues to run fixed-m safely when PR #477 is merged on its own.

syhuang22 · 2026-09-15T20:08:49Z

+  # This guarantees that negative logits never flush to zero in normal FP32, preventing silent
+  # loss of significant softmax probability mass even when keys are mean-centered.
+  # Both ring and non-ring paths strictly adhere to this two-sided bound.
+  safe_bound = float(int(safe_window // 2))


This halves the Ulysses gate (213 -> 113). With centered keys the row mean is 0, so Jensen caps the lost mass at 2^(ceil(U)-C-126). Something like C+116 keeps it under 2^-10 without halving. What do fallback rates look like at 113?

On Wan 2.2 720p/81f (40 steps), per_q_block=True + K-centering shrinks the realized Cauchy-Schwarz product $U = |q| \cdot |k - \bar{k}|$ by $>2.2\times$, so 99.4% of Q-blocks pass gate 113 on 1D Ulysses (U=4), and 82.0% of calls pass on Ring R=2 (higher than main's 70.0% at gate 213).

The reason we keep floor(W/2) = 113 even with centered keys is that $m_B = \lceil U_{\text{block}} \rceil - C$ is shared across all rows in a Q-block (or head), not computed per individual row. If a block has an outlier row with $U_{\text{block}} \approx C + 116$ and a quiet row with $|q_i| \approx 0$ ($z_{i,j} \approx 0$), shifting the quiet row by $m_B$ gives exponents $0 - m_B \approx -116$. If $U_{\text{block}}$ approached $C + 126$, quiet rows in that block would underflow to zero denominator ($0/0 = \text{NaN}$) despite centered keys. floor(W/2) prevents both negative-logit underflow on high-norm rows and zero-logit underflow on quiet rows in the same block.

syhuang22 · 2026-09-15T20:08:49Z

+
+def get_fixed_m_constants(
+    kv_seq_len: int,
+    is_ring: bool = False,


is_ring isn't used. Drop it?

Kept is_ring: bool = False in the signature (with del is_ring and an explicit note at L166–188) so the ~20 callers across attention_flax.py, ring_attention_kernel.py, and the test suite self-document whether they are passing the distributed ring sequence length (shard_len * R) vs local shard length. Happy to remove the keyword arg across those call sites if you prefer a stricter signature!

syhuang22 · 2026-09-15T20:08:49Z

+    raise ValueError("`mk` metadata array is required when `use_fixed_m=True`.")
+  if mk is None:
+    mk = jnp.zeros((2, num_q_heads, grid_height), jnp.float32)
+  elif mk.ndim == 2:


mk[0] used to be max||k||, now it's m. Broadcasting an old 2D mk would silently read one as the other. Can we just raise? (same at L797)

Good catch — updated both _splash_attention_forward (L780) and _splash_attention_forward_ring (L981) to raise an explicit ValueError when mk.ndim == 2, and added test_legacy_2d_mk_raises_value_error in custom_splash_fixed_m_test.py.

syhuang22 · 2026-09-15T20:08:49Z

-  # block degrades the instruction schedule of the WHOLE grid -- measured 3x
-  # slower end to end, which is the cliff the design doc's D3 warns about.
-  # Keeping this one flag rather than two makes that combination unspellable.
  fixed_only = use_fixed_m and uniform_fixed_m


Can we keep the comment that was here? It's the only thing explaining why pinning and uniform_fixed_m must go together (two-body last block = Mosaic cliff, ~3x slower).

Restored in custom_splash_attention.py:L239-L250!

syhuang22 · 2026-09-15T20:08:50Z

+  v_ok = 1.0 if dtype_safe else 0.0
+  if dtype_safe and value is not None:
+    v_max_sq = (value.astype(jnp.float32) ** 2).max()
+    v_ok = (v_max_sq <= (v_max_bound**2)).astype(jnp.float32)


This never tripped in my WAN 2.2 runs. It's only needed because C went 88 -> 102, which buys gate 106 -> 113. Worth it?

Yes — raising $C(N)$ from 88 to 102 gives us +7 on the safe bound (106 -> 113), which meaningfully improves fixed-m hit rates. On the ring path (_ring_fixed_m_norms_pre_a2a), vn_local is fused into the single (ulysses, ring) pmax collective alongside qn and kn, so it adds zero extra collectives and <0.05 ms/step while guaranteeing fail-closed protection against FP32 numerator overflow (inf/NaN) if a LoRA or fine-tuned checkpoint produces |V| > 256.

syhuang22 · 2026-09-15T20:08:50Z

-        # precondition that keeps the fixed-m Cauchy-Schwarz bound flush-free.
-        key = key - jnp.mean(key, axis=2, keepdims=True)
+        k_mean = jnp.mean(real_key.astype(jnp.float32), axis=2)
+        if k_mean.shape[-1] < 128:


Nit: hard-coded 128, can we use the padded head dim? (same at L1556)

Done! Moved _pad_data_for_flash before the k_mean pad in _ulysses_attention (attention_flax.py:L1071) and replaced hard-coded 128 with query.shape[-1] (matching q.shape[-1] in ring_attention_kernel.py:L945).

syhuang22 · 2026-09-15T20:08:50Z

      use_base2_exp=context.get("use_base2_exp", True),
      use_experimental_scheduler=context.get("use_experimental_scheduler", False),
      use_fixed_m=True,
+      per_q_block=False,


Heads-up: ulysses_custom_fixed_m goes from per-row m to per-head m here. Fine for safety, but worth mentioning in the description.

Good callout! Updated the PR description to explicitly note that ulysses_custom_fixed_m uses per-head $m_B$ metadata (per_q_block=False), while ulysses_custom_fixed_m_per_q_block provides finer per-Q-block granularity.

syhuang22 · 2026-09-15T20:08:50Z

@@ -1132,7 +1255,6 @@ def _lse_scan(_):


 def make_custom_ring_attention(


Was removing the keyword-only * on purpose? With this many bool args it's a nice guard.

Restored *, after orig_kv_seq_len: int in make_custom_ring_attention (ring_attention_kernel.py:L1308). All callers already use keyword arguments.

syhuang22 · 2026-09-15T20:08:50Z

-    enhance_prompt = (
-        prompt_enhancement_words_threshold > 0 and prompt_word_count < prompt_enhancement_words_threshold
-    )
+    enhance_prompt = prompt_enhancement_words_threshold > 0 and prompt_word_count < prompt_enhancement_words_threshold


Unrelated formatting change, mind dropping it?

Dropped in 44dfe632 — generate_ltx_video.py is completely untouched now.

…and safety fallbacks Implements exact fixed-m splash attention in Pallas on TPU: - Dynamic C(N) headroom constants guaranteeing FP32 accumulator safety - fixed_m_dtype_is_safe checks rejecting FP16/FP8 exponent overflow - Value bound validation (|V| <= 256) with safe online softmax fallback - Unit tests covering all boundary conditions, dtypes, and scale factors Explicit metadata contracts on the fixed-m ring path ---------------------------------------------------- Three implicit contracts are made explicit. Each failed silently rather than loudly, and one of them produced non-finite output on TPU. 1. Norm representation is declared, not inferred. The gate previously guessed whether `fixed_m_norms` were squared with `(qn.max() * mk.max()) < 1000.0`. Magnitude cannot answer that question: legacy unsquared norms of (1000, 2) have a true bound of 2000, but read as already-squared they yield sqrt(2000) ~= 44.7 -- a ~45x under-estimate that admits fixed-m where it must fall back, and overflows. Replaced by `fixed_m_norms_squared` (default True, matching every in-tree caller); the test harness is migrated to squared norms. 2. The V-safety predicate is required, not assumed. Unlike the Cauchy-Schwarz norm bounds, the V-magnitude and dtype verdict is not re-derivable from a single hop's Q/K, so the kernel cannot reconstruct it. Omission previously meant "safe", which let fixed-m run on inputs it cannot represent: float16 with Q=K=0 and V=1 returns inf instead of 1.0. The ring path now raises unless `v_ok` is passed, mirroring the existing `fixed_m_recenter` rule, and `_ulysses_ring_custom_attention` computes it -- dtype safety plus |V| <= DEFAULT_MAX_V_BOUND, reduced with pmin over BOTH internal axes, since after the all-to-all neither axis alone observes the whole V and the fixed-m branch must be taken uniformly by every ppermute participant. 3. Norm shape is validated against per_q_block. This is the defect behind the `test_sink_head_falls_back_everywhere` TPU failure. Both gates compute `qn * mk[:, None]`, so a (num_heads,) array supplied while per_q_block=True does not raise -- it broadcasts to (num_heads, num_heads), pairing head j's query norm with head h's key norm. A sink head then inherits a small bound from an unrelated head, is wrongly marked eligible, and the kernel evaluates exp2(large_logit - small_m) -> inf. The kernel now rejects the mismatch, and the test declares per_q_block=False to match the per-head norms it supplies, as the production ring caller already did. Regression coverage: six backend-independent contract tests (both omissions raise, v_ok=False is accepted, mis-shaped norms are rejected, correctly shaped per-Q-block norms are accepted, fp16 is rejected while bf16/fp32 pass) plus two TPU tests (the two declared norm representations must agree, and an explicit unsafe verdict must force a finite fallback). Verified on v6e-8: ring_fixed_m_test 13 passed; attention_test, custom_splash_fixed_m_test, attention_block_sizes_test and ring_fixed_m_test together 58 passed.

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

Perseus14 mentioned this pull request Sep 13, 2026

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

Open

gemini-code-assist Bot reviewed Sep 13, 2026

View reviewed changes

Comment thread src/maxdiffusion/kernels/custom_splash_attention.py

Comment thread src/maxdiffusion/kernels/custom_splash_attention.py

Comment thread src/maxdiffusion/kernels/custom_splash_attention.py Outdated

Comment thread src/maxdiffusion/kernels/custom_splash_attention.py Outdated

Perseus14 force-pushed the feat/fixed-m-kernel branch 10 times, most recently from 3b9d7aa to cb1e4d1 Compare September 15, 2026 06:37

syhuang22 self-requested a review September 15, 2026 20:04

syhuang22 reviewed Sep 15, 2026

View reviewed changes

Perseus14 force-pushed the feat/fixed-m-kernel branch 4 times, most recently from 44dfe63 to fc91941 Compare September 17, 2026 12:31

syhuang22 previously approved these changes Sep 17, 2026

View reviewed changes

github-actions Bot added the pull ready label Sep 17, 2026

Perseus14 dismissed syhuang22’s stale review via f4413b1 September 17, 2026 18:41

Perseus14 force-pushed the feat/fixed-m-kernel branch 2 times, most recently from f4413b1 to d57775a Compare September 17, 2026 18:52

Perseus14 force-pushed the feat/fixed-m-kernel branch from d57775a to 678e806 Compare September 17, 2026 18:57

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

Perseus14 self-assigned this Sep 17, 2026

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

		@@ -1132,7 +1255,6 @@ def _lse_scan(_):


		def make_custom_ring_attention(

Conversation

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

Uh oh!

Summary

This PR is additive: main's $R &gt; 1$ behaviour is unchanged

Why the gate is halved, and why is_ring does not change it

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!

Uh oh!

Uh oh!

syhuang22 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Perseus14 commented Sep 13, 2026 •

edited

Loading

This PR is additive: `main`'s $R > 1$ behaviour is unchanged

Why the gate is halved, and why `is_ring` does not change it

syhuang22 left a comment •

edited

Loading