Skip to content

Commit 07d7e72

Browse files
committed
feat(attention): fixed-m splash attention kernel with dynamic bounds 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.
1 parent 3b6ee39 commit 07d7e72

5 files changed

Lines changed: 1749 additions & 219 deletions

File tree

0 commit comments

Comments
 (0)