Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
3b9d7aa to
cb1e4d1
Compare
There was a problem hiding this comment.
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 = Falsein_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 2Dmkstill 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
left a comment
syhuang22
left a comment
There was a problem hiding this comment.
Some line-level notes to go with my comment above.
| "(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 |
There was a problem hiding this comment.
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.
| # 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)) |
There was a problem hiding this comment.
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?
|
|
||
| def get_fixed_m_constants( | ||
| kv_seq_len: int, | ||
| is_ring: bool = False, |
There was a problem hiding this comment.
is_ring isn't used. Drop it?
| 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: |
There was a problem hiding this comment.
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)
| # 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 |
There was a problem hiding this comment.
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).
| 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) |
There was a problem hiding this comment.
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?
| # 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: |
There was a problem hiding this comment.
Nit: hard-coded 128, can we use the padded head dim? (same at L1556)
| 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, |
There was a problem hiding this comment.
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.
| @@ -1132,7 +1255,6 @@ def _lse_scan(_): | |||
|
|
|||
|
|
|||
| def make_custom_ring_attention( | |||
There was a problem hiding this comment.
Was removing the keyword-only * on purpose? With this many bool args it's a nice guard.
| 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 |
There was a problem hiding this comment.
Unrelated formatting change, mind dropping it?
3067e91 to
970a828
Compare
…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.
970a828 to
44dfe63
Compare
2 participants
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.
than pinned to one operating point:
With the shift
and no accumulator overflows.
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.last_compute_body_fixedlets physicallyunpadded KV sequences run without out-of-bounds access.
all-to-all so XLA fuses the scalar multiply into the upstream projection, saving 185 MB of HBM
traffic per layer.
This PR is additive:$R > 1$ behaviour is unchanged
main'sCentering and the Cauchy-Schwarz bound are a matched pair. If the kernel exponentiates centered$q \cdot (k_j - \bar{k})$ while the caller's eligibility bound was built from raw $K$ , the$\lVert k \rVert$ does not bound $\lVert k - \bar{k} \rVert$ — and a
$2^{128}$ .
logits
bound caps the wrong quantity —
tile can clear the gate while the logit it was supposed to cap overflows fp32.
RingRawKeyBoundUnsoundTestpins that down in pure arithmetic: keys all of norm 100 give a raw boundof 100 against a gate of 116, while the centered max logit is ~200 and the shifted exponent exceeds
An earlier revision of this PR handled that by having the ring kernel compute$R > 1$ to stop the now-mismatched caller
$R > 1$ ) and left the disabled code path carrying a latent bug.
k_meanitself andcenter unconditionally, then disabling fixed-m for
norms from being used. Reviewers correctly pointed out that this made the PR a regression against
main(which does run fixed-m atThis revision instead makes centering opt-in: the kernel centers only when handed a$\lfloor W/2 \rfloor$ bound. $R > 1$ caller passes no
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
main'sk_mean, so it is left running exactly as before and the restriction is gone. Net effect: this PRadds the fixed-m kernel without changing any behaviour
mainalready has, and merges safely on itsown. #478 then turns centering on by supplying both halves of the pair.
Why the gate is halved, and why
is_ringdoes not change itReview raised that$W$ is not, and we verified that the hard way.
get_fixed_m_constantsacceptsis_ringbut ignores it, and that the gate ishalf the legacy
_FIXED_M_SAFE_BOUND = 213. Both observations are correct; the conclusion that thecentered path should get the full window
The appealing argument is that K-centering forces the realized row max$M \ge 0$ , so only one side$U$ , not from $M$ : the worst case stays
$C - (U + \lceil U \rceil)$ whether or not the keys are centered. Wiring $W$ $N = 4096$ ,
$W = 232$ , so a query of norm $231.01$ against exactly-centered keys passes the gate and then
$N$ they pass — the ring passes the whole
needs absorbing. It fails because the shift is built from
is_ringup to returnfor the Ulysses path reproduces exactly the bug the halving was introduced to fix — at
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_lossandtest_cpu_proof_invariant_boundsboth pinthis down, and both fail under the loosened gate. The parameter is retained (documented, with the
derivation) because the two call sites differ in the
distributed length — and because it marks the seam where a future kernel that shifts by the realized
row max could legitimately diverge.
Measured cost of the strict gate. Instrumented over a full WAN 2.2 720p generation on v6e-8$N = 75{,}600$ ):
(1280 metadata calls per gate,
So the strict gate costs about 0.3 percentage points of extra fallback over the loosest alternative.
The largest observed bound is 1134.58 — far above every candidate gate — so genuine sink heads exist
that fall back regardless, and no choice of gate removes the online-softmax path.
Verification
custom_splash_fixed_m_test.py: 22 passed on v6e-8 (dynamic bounds, non-divisible sequences, theadversarial centered-key 15/17 regression, hybrid fallbacks).
ring_fixed_m_test.py: 16 passed on v6e-8.(
Only interpret mode is supported on CPU backend); they need a TPU runner.pyink --pyink-indentation=2 --line-length=125andruff checkclean.