Repository navigation
Commit 07d7e72
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
- src/maxdiffusion
- kernels
- splash_attention
- models
- tests
0 commit comments