diff --git a/src/maxdiffusion/kernels/custom_splash_attention.py b/src/maxdiffusion/kernels/custom_splash_attention.py index 6bd3f3493..3b3603cba 100644 --- a/src/maxdiffusion/kernels/custom_splash_attention.py +++ b/src/maxdiffusion/kernels/custom_splash_attention.py @@ -17,6 +17,7 @@ """Custom Pallas flash attention kernel for TPU.""" import functools +import math import jax import jax.numpy as jnp @@ -50,27 +51,152 @@ def __init__( # Fixed-m softmax-bound constants. Instead of tracking the online-softmax # running max per KV block, eligible heads subtract a precomputed per-query # upper bound on the logits (Cauchy-Schwarz: max_j q_i.k_j <= ||q_i|| * -# max_j||k_j||). _FIXED_M_RECENTER (C) shifts the exp2 exponents up so the +# max_j||k_j||). The recenter constant C shifts the exp2 exponents up so the # largest surviving term stays above the f32 subnormal-flush floor 2^-126: # with k-smoothing the per-row max is >= 0, so the max term has exponent -# >= -ceil(bound) + C, which stays > -126 while ceil(bound) <= -# _FIXED_M_SAFE_BOUND (= C + 126 - 1 of margin). Heads whose worst-case bound -# exceeds the gate fall back to online softmax (the "sink" heads). -_FIXED_M_RECENTER = 88.0 -_FIXED_M_SAFE_BOUND = 213.0 +# >= -ceil(bound) + C, which stays > -126 while ceil(bound) <= the safe +# bound (= C + 126 - 1 of margin). Heads whose worst-case bound exceeds the +# gate fall back to online softmax (the "sink" heads). +# # Ring-path gate: the ring processes UN-smoothed K shards (no ring rank holds # the full K to compute a mean, and a per-shard mean would shift each hop's # logits differently, breaking the cross-shard merge). Without k-smoothing the # per-row max logit has no >=0 guarantee, so the safe bound halves (calibrated # for ring_size=2, matching DiffusionServing's ring gate). +# +# These fixed values are superseded by `get_fixed_m_constants`, which derives +# C(N) and the gate from the actual KV length rather than from a single +# hard-coded operating point. They are retained as the reference point the +# dynamic derivation is calibrated against, and are exercised by the tests. +# +# Note what the pair encodes: 213 == C + 125 == W is the full no-flush window +# assuming the shift tracks the realized row max. It does not -- the shift is +# built from the Cauchy-Schwarz bound -- so W admits inputs that silently flush +# their negative logits, and `get_fixed_m_constants` returns floor(W/2) for +# every caller. The halved ring constant here is the value that generalised; +# 213 is kept only as the historical calibration point. + +_FIXED_M_RECENTER = 88.0 +_FIXED_M_SAFE_BOUND = 213.0 _FIXED_M_RING_SAFE_BOUND = _FIXED_M_SAFE_BOUND / 2.0 +FP32_OUTPUT_HEADROOM_BITS = 8.0 # Assumes default activation |V| <= 2**FP32_OUTPUT_HEADROOM_BITS = 256.0 +DEFAULT_MAX_V_BOUND = 256.0 -def _flash_attention_kernel( + +def fixed_m_dtype_is_safe(dtype, recenter: float) -> bool: + """Whether `dtype` can hold the fixed-m softmax weights without overflowing. + + Fixed-m deliberately parks the un-normalized weights at up to `2**recenter`, + a range derived against FP32's exponent (see `get_fixed_m_constants`). The + kernel then narrows them to the activation dtype for the S@V matmul + (`s_curr.astype(q_ref.dtype)`), so a dtype with a *smaller exponent range* + silently overflows to inf even though the FP32 bound analysis passed. + + bfloat16 and float32 both have 8-bit exponents (maxexp 128) and are safe for + every C(N) this module produces. float16 has a 5-bit exponent (maxexp 16) and + is not: at N=4096 with |V| <= 256, C(N) = 107 and 2**107 is far beyond + float16's 65504 ceiling. The fp8 formats fail for the same reason. + + This is deliberately expressed in terms of the exponent range rather than an + allowlist so narrower formats are rejected automatically. + + Args: + dtype: Activation dtype the kernel will narrow the weights to. + recenter: The fixed-m constant C(N) from `get_fixed_m_constants`. + + Returns: + True if `2**recenter` is representable in `dtype`. + """ + return float(jnp.finfo(jnp.dtype(dtype)).maxexp) > float(recenter) + + +def get_fixed_m_constants( + kv_seq_len: int, + is_ring: bool = False, + v_max_bound: float = DEFAULT_MAX_V_BOUND, +) -> tuple[float, float]: + """Computes dynamic fixed-m constants C(N) and safe bounds based on KV sequence length. + + Mathematical Derivations: + 1. Overflow Ceiling: + For a given upper bound on value activation magnitude |V| <= V_max: + output_headroom_bits = ceil(log2(V_max)). + The ceiling constant C(N) = 127.0 - ceil(log2(N)) - output_headroom_bits guarantees that: + - Denominator accumulator: l = sum_j 2^{z_j - m} <= N * 2^C(N) <= 2^{127 - headroom} < 2^{128} + - Numerator accumulator: |o_d| = |sum_j V_{j,d} 2^{z_j - m}| <= V_max * N * 2^C(N) <= 2^{127} < 2^{128} + preventing IEEE-754 FP32 overflow for all activations |V| <= V_max. + + 2. Subnormal Underflow Floor (Cauchy-Schwarz Proof): + Let U_i = max_i ||q_i|| * max_j ||k_j|| be the Cauchy-Schwarz bound on query-key inner products. + By Cauchy-Schwarz inequality, for all tokens j: + z_j = Q_i . K_j >= -||Q_i|| * ||K_j|| >= -U_i. + With the fixed-m base shift defined as m_i = ceil(U_i) - C(N): + z_j - m_i >= -U_i - (ceil(U_i) - C(N)) = C(N) - (U_i + ceil(U_i)). + To guarantee that no term underflows into the subnormal range (requiring minimal shifted exponent >= -125.0, + providing 1 bit of margin above IEEE-754 normal floor -126.0): + U + ceil(U) <= W(N) = C(N) + 125.0 => U <= floor(W(N) / 2). + If U_i <= floor(W(N) / 2), it is mathematically impossible to underflow below -125.0. + + This bound is two-sided for *both* the ring and the Ulysses path. K-centering + (which only the Ulysses/centered path applies) guarantees the realized row max + M >= 0, but m_i above is built from the Cauchy-Schwarz bound U rather than from + M, so the -U side is still exposed and the full window W would admit inputs that + flush their negative logits. Centering pays off by shrinking U itself, since mk + is then measured on centered keys. See `is_ring` below. + """ + + if kv_seq_len is None or kv_seq_len <= 0: + raise ValueError(f"kv_seq_len must be a positive integer to compute dynamic fixed-m constants, got {kv_seq_len=}") + if v_max_bound <= 0.0: + raise ValueError(f"v_max_bound must be a positive float, got {v_max_bound=}") + + fp32_max_exp = 128.0 + fp32_min_normal_exp = -126.0 + output_headroom_bits = float(max(0, math.ceil(math.log2(float(v_max_bound))))) + + max_accumulation_bits = float(math.ceil(math.log2(float(kv_seq_len)))) + + # C(N) = 127.0 - max_accumulation_bits - output_headroom_bits + recenter = fp32_max_exp - max_accumulation_bits - output_headroom_bits - 1.0 + + # Safe window W(N) = C(N) - (-126.0) - 1.0 = C(N) + 125.0 + safe_window = recenter - fp32_min_normal_exp - 1.0 + + # `is_ring` does NOT select a different bound, and that is deliberate. + # + # It is tempting to give the centered Ulysses path the full window W: k-centering + # forces the row max M >= 0, so "only one side needs absorbing". That reasoning is + # wrong here because the shift is not the realized row max. `_compute_fixed_m_metadata` + # sets m_base = ceil(U) - C from the *Cauchy-Schwarz* bound U, so the worst-case + # shifted exponent is C - (U + ceil(U)) whether or not the keys are centered, and + # only U <= floor(W/2) keeps that above -125. Centering helps by shrinking U itself + # (mk is measured on centered keys), not by making the inequality one-sided. + # + # This was tried and measured: at N=4096 the full window is W=232, and a query of + # norm 231.01 against exactly-centered keys then passes the gate and silently loses + # ~11.8% of the softmax mass (0.882 vs 1.0) as the negative logits flush to zero. + # `test_adversarial_centered_keys_softmax_mass_loss` and + # `test_cpu_proof_invariant_bounds` both pin this down. The legacy pair + # (_FIXED_M_SAFE_BOUND = 213 = W, _FIXED_M_RING_SAFE_BOUND = W/2) is where the loose + # value came from; the halving is the fix, not a ring-specific tightening. + # + # The parameter is kept because the two call sites differ in the *N* they pass -- the + # ring passes the whole distributed length (per-shard length x R), Ulysses passes its + # own -- so it documents which derivation the caller believes it is in, and leaves a + # seam if a future kernel ever shifts by the realized row max. + del is_ring + safe_bound = float(int(safe_window // 2)) + + return recenter, safe_bound + + +def _flash_attention_kernel_impl( mk_ref, q_ref, k_ref, v_ref, + k_mean_ref, m_scratch_ref, l_scratch_ref, o_scratch_ref, @@ -89,13 +215,24 @@ def _flash_attention_kernel( fuse_reciprocal: bool = True, use_fixed_m: bool = False, uniform_fixed_m: bool = False, + fixed_m_recenter: float | None = None, + q_heads_per_kv_head: int = 1, + use_k_centering: bool = False, ): + """Pallas Mosaic TPU flash attention kernel with fixed-m support. + + Scalar Prefetch Multiplexing: + `mk_ref` is a multiplexed scalar prefetch buffer of shape `(2, num_heads, num_q_blocks)` + passing both the precomputed block fixed-m base shift and discrete predicate in a single scalar memory slot: + - `mk_ref[0, h, i]`: Precomputed block shift m_B = ceil(max_i ||q_i|| * max_j ||k_j||) - C. + - `mk_ref[1, h, i]`: Gating eligibility predicate (1.0 for fixed-m, 0.0 for online). + """ float32 = jnp.float32 head_dim_v_repeats, rem = divmod(head_dim_v, NUM_SUBLANES) if rem != 0: raise NotImplementedError(f"{head_dim_v=} should be a multiple of {NUM_SUBLANES}") - h, _, j = pl.program_id(0), pl.program_id(1), pl.program_id(2) + h, i, j = pl.program_id(0), pl.program_id(1), pl.program_id(2) exp = jnp.exp2 if use_base2_exp else jnp.exp sv_dims = (((0,), (0,)), ((), ())) @@ -115,16 +252,28 @@ def _flash_attention_kernel( if uniform_fixed_m and not use_fixed_m: raise ValueError("uniform_fixed_m requires use_fixed_m.") - # Per-head dispatch: heads inside the no-flush window run fixed-m, the rest - # keep online softmax. Branch once per head (body level), never per step. - is_fixed = (mk_ref[1, h] > 0.5) if (use_fixed_m and not fixed_only) else False + if use_fixed_m and fixed_m_recenter is None: + raise ValueError("fixed_m_recenter must be specified when use_fixed_m=True.") + + # Per-(head, Q-block) dispatch: heads / Q-blocks inside the no-flush window run + # fixed-m, the rest keep online softmax. + if use_fixed_m and not fixed_only: + is_fixed = mk_ref[1, h, i] > 0.5 + else: + is_fixed = False def _write_fixed_m(): - # Per-query Cauchy-Schwarz bound m_i = ceil(||q_i|| * max_j||k_j||) - C. - qf = q_ref[...].astype(float32) - qn = jnp.sqrt((qf * qf).sum(axis=1))[None, :] # (1, bq) per-query norm - bound = qn * mk_ref[0, h] - m_fixed = jnp.ceil(bound) - _FIXED_M_RECENTER + # Precomputed block bound m_B = ceil(max_i ||q_i|| * max_j ||k_j||) - C. + # Virtual K-centering applies the row-specific projection: m_i = m_B + q_i^T \bar{k}. + m_base = mk_ref[0, h, i] + if use_k_centering and k_mean_ref is not None: + qf = q_ref[...].astype(float32) + kv_h = h // q_heads_per_kv_head if q_heads_per_kv_head > 1 else h + km = k_mean_ref[kv_h, :].astype(float32) + mu = (qf * km[None, :]).sum(axis=1)[None, :] + m_fixed = m_base + mu + else: + m_fixed = m_base m_scratch_ref[...] = jnp.broadcast_to(m_fixed, m_scratch_ref.shape) @pl.when(j == 0) @@ -232,7 +381,8 @@ def last_compute_body_fixed(kv_compute_index): l_scratch_ref[...] = l_prev o_scratch_ref[:] = o_prev - assert bkv % bkv_compute == 0 + if bkv % bkv_compute != 0: + raise ValueError(f"block_kv ({bkv}) must be divisible by block_kv_compute ({bkv_compute})") if fixed_only: @@ -320,6 +470,118 @@ def end(): m_ring_ref[...] = m_scratch_ref[...].astype(m_ring_ref.dtype) +def _flash_attention_kernel( + mk_ref, + q_ref, + k_ref, + v_ref, + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + o_ref, + l_ring_ref=None, + m_ring_ref=None, + *, + mask_value: float, + grid_width: int, + bkv: int, + bkv_compute: int, + bkv_compute_in: int, + head_dim_v: int, + kv_seq_len: int, + use_base2_exp: bool = True, + fuse_reciprocal: bool = True, + use_fixed_m: bool = False, + uniform_fixed_m: bool = False, + fixed_m_recenter: float | None = None, + q_heads_per_kv_head: int = 1, +): + return _flash_attention_kernel_impl( + mk_ref, + q_ref, + k_ref, + v_ref, + None, + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + o_ref, + l_ring_ref=l_ring_ref, + m_ring_ref=m_ring_ref, + mask_value=mask_value, + grid_width=grid_width, + bkv=bkv, + bkv_compute=bkv_compute, + bkv_compute_in=bkv_compute_in, + head_dim_v=head_dim_v, + kv_seq_len=kv_seq_len, + use_base2_exp=use_base2_exp, + fuse_reciprocal=fuse_reciprocal, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=fixed_m_recenter, + q_heads_per_kv_head=q_heads_per_kv_head, + use_k_centering=False, + ) + + +def _flash_attention_kernel_kcentered( + mk_ref, + q_ref, + k_ref, + v_ref, + k_mean_ref, + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + o_ref, + l_ring_ref=None, + m_ring_ref=None, + *, + mask_value: float, + grid_width: int, + bkv: int, + bkv_compute: int, + bkv_compute_in: int, + head_dim_v: int, + kv_seq_len: int, + use_base2_exp: bool = True, + fuse_reciprocal: bool = True, + use_fixed_m: bool = False, + uniform_fixed_m: bool = False, + fixed_m_recenter: float | None = None, + q_heads_per_kv_head: int = 1, + use_k_centering: bool = True, +): + return _flash_attention_kernel_impl( + mk_ref, + q_ref, + k_ref, + v_ref, + k_mean_ref, + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + o_ref, + l_ring_ref=l_ring_ref, + m_ring_ref=m_ring_ref, + mask_value=mask_value, + grid_width=grid_width, + bkv=bkv, + bkv_compute=bkv_compute, + bkv_compute_in=bkv_compute_in, + head_dim_v=head_dim_v, + kv_seq_len=kv_seq_len, + use_base2_exp=use_base2_exp, + fuse_reciprocal=fuse_reciprocal, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=fixed_m_recenter, + q_heads_per_kv_head=q_heads_per_kv_head, + use_k_centering=use_k_centering, + ) + + def _flash_attention_kernel_mhpt( q_ref, k_ref, @@ -443,7 +705,8 @@ def last_compute_body(kv_compute_index): l_scratch_ref[h_local] = l_prev o_scratch_ref[h_local] = o_prev - assert bkv % bkv_compute == 0 + if bkv % bkv_compute != 0: + raise ValueError(f"block_kv ({bkv}) must be divisible by block_kv_compute ({bkv_compute})") @pl.when(j != grid_width - 1) def body(): @@ -483,14 +746,12 @@ def _splash_attention_forward( vmem_limit_bytes: int | None = None, use_fixed_m: bool = False, mk: jax.Array | None = None, + fixed_m_recenter: float | None = None, + uniform_fixed_m: bool = False, + k_mean: jax.Array | None = None, ): num_q_heads, padded_q_seq_len, head_dim_qk = q.shape head_dim_v = v.shape[-1] - # Scalar-prefetch operand carrying per-head fixed-m data: - # mk[0, h] = max_j||k_j|| (Cauchy-Schwarz factor), mk[1, h] = eligibility. - # A dummy is supplied for online callers; the kernel ignores it. - if mk is None: - mk = jnp.zeros((2, num_q_heads), jnp.float32) bq, bkv = block_sizes.block_q, block_sizes.block_kv bkv_compute = block_sizes.block_kv_compute bkv_compute_in = block_sizes.block_kv_compute_in @@ -499,7 +760,31 @@ def _splash_attention_forward( actual_q_seq_len = q_seq_len if q_seq_len is not None else padded_q_seq_len actual_kv_seq_len = kv_seq_len if kv_seq_len is not None else padded_kv_seq_len + if num_q_heads % num_kv_heads != 0: + raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA.") q_heads_per_kv_head = num_q_heads // num_kv_heads + grid_width = (actual_kv_seq_len + bkv - 1) // bkv + grid_height = (actual_q_seq_len + bq - 1) // bq + grid = (num_q_heads, grid_height, grid_width) + + if use_fixed_m and fixed_m_recenter is None: + fixed_m_recenter, _ = get_fixed_m_constants(actual_kv_seq_len, is_ring=False) + + # Scalar-prefetch operand carrying per-head / per-Q-block fixed-m data: + # mk[0, h, i] = m_B (precomputed block fixed-m base shift), mk[1, h, i] = eligibility. + # A dummy is supplied for online callers; the kernel ignores it. + if use_fixed_m and mk is None: + 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: + raise ValueError( + "2D `mk` arrays (2, num_q_heads) are not supported: `mk[0]` now stores the precomputed " + "base shift m_B rather than legacy max||k||. Pass a 3D (2, num_q_heads, num_q_blocks) array." + ) + + if mk.shape[0] != 2 or mk.shape[1] != num_q_heads or mk.shape[2] != grid_height: + raise ValueError(f"mk must have shape (2, {num_q_heads}, {grid_height}), got {mk.shape}") def q_index_map(h, i, j, *_): return (h, i, 0) @@ -513,11 +798,6 @@ def k_index_map(h, i, j, *_): def v_index_map(h, i, j, *_): return (h // q_heads_per_kv_head, j, 0) - in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), - pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), - pl.BlockSpec((None, bkv, head_dim_v), v_index_map), - ] out_shapes = [ jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), @@ -530,23 +810,64 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), pl.BlockSpec((None, head_dim_v, bq), out_index_map), ] - grid_width = (actual_kv_seq_len + bkv - 1) // bkv - grid_height = (actual_q_seq_len + bq - 1) // bq - grid = (num_q_heads, grid_height, grid_width) + + if k_mean is None: + in_specs = [ + pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), + pl.BlockSpec((None, bkv, head_dim_v), v_index_map), + ] + kernel_fn = functools.partial( + _flash_attention_kernel, + mask_value=DEFAULT_MASK_VALUE, + grid_width=grid_width, + bkv=bkv, + bkv_compute=bkv_compute, + bkv_compute_in=bkv_compute_in, + head_dim_v=head_dim_v, + kv_seq_len=actual_kv_seq_len, + use_base2_exp=use_base2_exp, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=fixed_m_recenter, + q_heads_per_kv_head=q_heads_per_kv_head, + ) + kernel_args = (mk, q, k, v) + else: + if k_mean.shape[0] == num_q_heads and num_q_heads != num_kv_heads: + k_mean = k_mean[::q_heads_per_kv_head] + if k_mean.shape[0] < num_kv_heads or k_mean.shape[1] != head_dim_qk: + raise ValueError(f"k_mean must have shape (>={num_kv_heads}, {head_dim_qk}), got {k_mean.shape}") + pad_h = (NUM_SUBLANES - (k_mean.shape[0] % NUM_SUBLANES)) % NUM_SUBLANES + if pad_h > 0: + k_mean = jnp.pad(k_mean, ((0, pad_h), (0, 0))) + + in_specs = [ + pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), + pl.BlockSpec((None, bkv, head_dim_v), v_index_map), + pl.BlockSpec((k_mean.shape[0], head_dim_qk), lambda *_: (0, 0)), + ] + kernel_fn = functools.partial( + _flash_attention_kernel_kcentered, + mask_value=DEFAULT_MASK_VALUE, + grid_width=grid_width, + bkv=bkv, + bkv_compute=bkv_compute, + bkv_compute_in=bkv_compute_in, + head_dim_v=head_dim_v, + kv_seq_len=actual_kv_seq_len, + use_base2_exp=use_base2_exp, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=fixed_m_recenter, + q_heads_per_kv_head=q_heads_per_kv_head, + use_k_centering=True, + ) + kernel_args = (mk, q, k, v, k_mean) all_out = pl.pallas_call( - functools.partial( - _flash_attention_kernel, - mask_value=DEFAULT_MASK_VALUE, - grid_width=grid_width, - bkv=bkv, - bkv_compute=bkv_compute, - bkv_compute_in=bkv_compute_in, - head_dim_v=head_dim_v, - kv_seq_len=actual_kv_seq_len, - use_base2_exp=use_base2_exp, - use_fixed_m=use_fixed_m, - ), + kernel_fn, grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=1, in_specs=in_specs, @@ -561,7 +882,7 @@ def v_index_map(h, i, j, *_): vmem_limit_bytes=vmem_limit_bytes, ), out_shape=out_shapes, - )(mk, q, k, v) + )(*kernel_args) return all_out[-1] @@ -578,6 +899,8 @@ def _splash_attention_forward_ring( use_fixed_m: bool = False, mk: jax.Array | None = None, uniform_fixed_m: bool = False, + fixed_m_recenter: float | None = None, + k_mean: jax.Array | None = None, ): """Ring-specific forward path that returns pre-reciprocal fp32 accumulators. @@ -603,8 +926,16 @@ def _splash_attention_forward_ring( actual_q_seq_len = q_seq_len if q_seq_len is not None else padded_q_seq_len actual_kv_seq_len = kv_seq_len if kv_seq_len is not None else padded_kv_seq_len + if num_q_heads % num_kv_heads != 0: + raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA.") q_heads_per_kv_head = num_q_heads // num_kv_heads + if use_fixed_m and fixed_m_recenter is None: + fixed_m_recenter, _ = get_fixed_m_constants(actual_kv_seq_len, is_ring=True) + + if use_fixed_m and mk is None: + raise ValueError("`mk` metadata array is required when `use_fixed_m=True`.") + def q_index_map(h, i, j, *_): return (h, i, 0) @@ -617,11 +948,6 @@ def k_index_map(h, i, j, *_): def v_index_map(h, i, j, *_): return (h // q_heads_per_kv_head, j, 0) - in_specs = [ - pl.BlockSpec((None, bq, head_dim_qk), q_index_map), - pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), - pl.BlockSpec((None, bkv, head_dim_v), v_index_map), - ] out_shapes = [ jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), @@ -642,28 +968,81 @@ def v_index_map(h, i, j, *_): grid_height = (actual_q_seq_len + bq - 1) // bq grid = (num_q_heads, grid_height, grid_width) - # Scalar-prefetch operand carrying per-head fixed-m data (same convention as - # `_splash_attention_forward`): mk[0, h] = max_j||k_j|| over ALL ring shards - # (the caller all-reduces this over the ring axis), mk[1, h] = eligibility. + # Scalar-prefetch operand carrying per-head / per-Q-block fixed-m data: + # mk[0, h, i] = m_B, the precomputed block fixed-m base shift derived from + # the bound max_i||q_i|| * max_j||k_j|| taken over ALL ring shards (the + # caller all-reduces that norm over the ring axis before forming m_B). + # mk[1, h, i] = eligibility. # A dummy is supplied for online callers; the kernel ignores it. + if use_fixed_m and mk is None: + raise ValueError("`mk` metadata array is required when `use_fixed_m=True`.") if mk is None: - mk = jnp.zeros((2, num_q_heads), jnp.float32) + mk = jnp.zeros((2, num_q_heads, grid_height), jnp.float32) + elif mk.ndim == 2: + raise ValueError( + "2D `mk` arrays (2, num_q_heads) are not supported: `mk[0]` now stores the precomputed " + "base shift m_B rather than legacy max||k||. Pass a 3D (2, num_q_heads, num_q_blocks) array." + ) + + if k_mean is None: + in_specs = [ + pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), + pl.BlockSpec((None, bkv, head_dim_v), v_index_map), + ] + kernel_fn = functools.partial( + _flash_attention_kernel, + mask_value=DEFAULT_MASK_VALUE, + grid_width=grid_width, + bkv=bkv, + bkv_compute=bkv_compute, + bkv_compute_in=bkv_compute_in, + head_dim_v=head_dim_v, + kv_seq_len=actual_kv_seq_len, + use_base2_exp=use_base2_exp, + fuse_reciprocal=False, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=fixed_m_recenter, + q_heads_per_kv_head=q_heads_per_kv_head, + ) + kernel_args = (mk, q, k, v) + else: + if k_mean.shape[0] == num_q_heads and num_q_heads != num_kv_heads: + k_mean = k_mean[::q_heads_per_kv_head] + if k_mean.shape[0] < num_kv_heads or k_mean.shape[1] != head_dim_qk: + raise ValueError(f"k_mean must have shape (>={num_kv_heads}, {head_dim_qk}), got {k_mean.shape}") + pad_h = (NUM_SUBLANES - (k_mean.shape[0] % NUM_SUBLANES)) % NUM_SUBLANES + if pad_h > 0: + k_mean = jnp.pad(k_mean, ((0, pad_h), (0, 0))) + + in_specs = [ + pl.BlockSpec((None, bq, head_dim_qk), q_index_map), + pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), + pl.BlockSpec((None, bkv, head_dim_v), v_index_map), + pl.BlockSpec((k_mean.shape[0], head_dim_qk), lambda *_: (0, 0)), + ] + kernel_fn = functools.partial( + _flash_attention_kernel_kcentered, + mask_value=DEFAULT_MASK_VALUE, + grid_width=grid_width, + bkv=bkv, + bkv_compute=bkv_compute, + bkv_compute_in=bkv_compute_in, + head_dim_v=head_dim_v, + kv_seq_len=actual_kv_seq_len, + use_base2_exp=use_base2_exp, + fuse_reciprocal=False, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=fixed_m_recenter, + q_heads_per_kv_head=q_heads_per_kv_head, + use_k_centering=True, + ) + kernel_args = (mk, q, k, v, k_mean) all_out = pl.pallas_call( - functools.partial( - _flash_attention_kernel, - mask_value=DEFAULT_MASK_VALUE, - grid_width=grid_width, - bkv=bkv, - bkv_compute=bkv_compute, - bkv_compute_in=bkv_compute_in, - head_dim_v=head_dim_v, - kv_seq_len=actual_kv_seq_len, - use_base2_exp=use_base2_exp, - fuse_reciprocal=False, - use_fixed_m=use_fixed_m, - uniform_fixed_m=uniform_fixed_m, - ), + kernel_fn, grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=1, in_specs=in_specs, @@ -678,7 +1057,7 @@ def v_index_map(h, i, j, *_): vmem_limit_bytes=vmem_limit_bytes, ), out_shape=out_shapes, - )(mk, q, k, v) + )(*kernel_args) out = jnp.swapaxes(all_out[3], 1, 2) # (h, head_dim_v, s) -> (h, s, head_dim_v) l = all_out[4][:, 0, :] # (h, s) m = all_out[5][:, 0, :] # (h, s) @@ -707,8 +1086,10 @@ def _splash_attention_forward_mhpt( actual_kv_seq_len = kv_seq_len if kv_seq_len is not None else k.shape[1] hpt = heads_per_tile - assert num_q_heads % hpt == 0, f"num_heads {num_q_heads} must be divisible by heads_per_tile {hpt}" - assert num_q_heads == num_kv_heads, "MHPT currently requires num_q_heads == num_kv_heads (no GQA)" + if num_q_heads % hpt != 0: + raise ValueError(f"num_heads {num_q_heads} must be divisible by heads_per_tile {hpt}") + if num_q_heads != num_kv_heads: + raise ValueError(f"MHPT currently requires num_q_heads == num_kv_heads (no GQA), got {num_q_heads=} vs {num_kv_heads=}") def q_index_map(h, i, j, *_): return (h, i, 0) @@ -783,8 +1164,23 @@ def make_splash_mha( use_experimental_scheduler: bool = False, vmem_limit_bytes: int | None = None, use_fixed_m: bool = False, + uniform_fixed_m: bool = False, + fixed_m_recenter: float | None = None, ): - def _splash_attention(q, k, v, mk=None): + if use_fixed_m: + if not use_base2_exp: + raise NotImplementedError( + "fixed-m softmax bounds are derived strictly for base-2 exponents. Please set use_base2_exp=True." + ) + if fixed_m_recenter is None: + fixed_m_recenter, _ = get_fixed_m_constants(orig_kv_seq_len, is_ring=False) + recenter = fixed_m_recenter + else: + recenter = None + + def _splash_attention(q, k, v, mk=None, k_mean=None): + if use_fixed_m and mk is None: + raise ValueError("`mk` metadata array is required when `use_fixed_m=True`.") if heads_per_tile > 1: if use_fixed_m: raise NotImplementedError("fixed-m is not supported with heads_per_tile > 1") @@ -812,6 +1208,9 @@ def _splash_attention(q, k, v, mk=None): vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=use_fixed_m, mk=mk, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=recenter, + k_mean=k_mean, ) return _splash_attention diff --git a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py index bc49c5af7..67f735651 100644 --- a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py +++ b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py @@ -168,7 +168,8 @@ def body( unroll=True, ) # type: ignore[arg-type] # Final normalization - assert l_final.dtype == jnp.float32 + if l_final.dtype != jnp.float32: + raise TypeError(f"l_final must have dtype float32, got {l_final.dtype}") l_inv = jnp.where(l_final == 0.0, 0.0, 1.0 / l_final) out = (o_final * l_inv[..., None]).astype(q.dtype) # Final logsumexp for residuals @@ -697,7 +698,8 @@ def make_ring_attention( is_dkv=True, return_dynamic_grid=config.dq_reduction_steps == 3, ) - assert (mask_function_fwd is None) == (mask_function_dkv is None) + if (mask_function_fwd is None) != (mask_function_dkv is None): + raise ValueError("mask_function_fwd and mask_function_dkv must both be None or both be provided") dkv_mask_sparsity = float(np.mean(dkv_mask_info.block_mask != 0)) dkv_mask_info = tree_util.tree_map(jnp.array, dkv_mask_info) @@ -760,6 +762,8 @@ def _custom_bidirectional_ring_forward( axis (no sub-group perm). """ axis_size = lax.axis_size(ring_axis) + effective_kv_seq_len = orig_kv_seq_len * axis_size + recenter, ring_safe_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=True) idx = lax.axis_index(ring_axis) exp_fn = jnp.exp2 if use_base2_exp else jnp.exp @@ -774,6 +778,7 @@ def _attn(kc, vc): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + fixed_m_recenter=recenter, ) return o.astype(jnp.float32), m.astype(jnp.float32), l.astype(jnp.float32) @@ -859,6 +864,12 @@ def _custom_ring_attention_forward( bidirectional: bool = False, use_fixed_m: bool = False, fixed_m_norms: tuple[jax.Array, jax.Array] | None = None, + fixed_m_norms_squared: bool = True, + per_q_block: bool = False, + pregathered_mk: bool = False, + k_mean: jax.Array | None = None, + uniform_fixed_m: bool | None = None, + v_ok: jax.Array | bool | None = None, ) -> jax.Array: """Forward-only ring attention using the custom dense splash kernel. @@ -882,17 +893,57 @@ def _custom_ring_attention_forward( mask_value: Initial running-max value for the online softmax. ring_axis: Name of the mesh axis to rotate K/V over (e.g. "context"). ring_size: Number of ring steps to scan over. Defaults to the full size of - `ring_axis`. For a hybrid Ulysses+Ring (USP) split this is the ring - sub-group size R (< full axis size), so each device only rotates within its - ring sub-group. + `ring_axis`. For fixed-m, ring_size must equal the size of ring_axis (2D + Ulysses+Ring should use a dedicated ring mesh axis). For online ring on a + flattened axis, this is the ring sub-group size R (< full axis size). perm: Explicit `ppermute` permutation. Defaults to a full-axis +1 rotation. - For the hybrid split, pass a perm that rotates K/V *within each ring - sub-group only* (built by the caller from the U x R factorization). + For fixed-m, the canonical ring permutation is required. For the online + hybrid split on a flattened axis, pass a perm that rotates K/V within each + ring sub-group only. + k_mean: Per-head global key mean vector (num_kv_heads, head_dim_qk) for + Global Virtual K-Centering. When provided or automatically reduced across + the ring axis, logit centering guarantees row max >= 0 across all hops. Returns: Normalized attention output, shape `(num_q_heads, q_seq_len, head_dim_v)`. """ axis_size = lax.axis_size(ring_axis) + effective_ring_size = ring_size if ring_size is not None else axis_size + effective_kv_seq_len = orig_kv_seq_len * effective_ring_size + + num_q_heads = q.shape[0] + num_kv_heads = k.shape[0] + head_dim_v = v.shape[-1] + + if use_fixed_m and not use_base2_exp: + raise NotImplementedError( + "fixed-m softmax bounds are derived strictly for base-2 exponents. Please set use_base2_exp=True." + ) + + # Virtual K-centering is OPT-IN: it happens only when the caller supplies + # `k_mean`. This matters for correctness, not taste. The caller's + # Cauchy-Schwarz metadata bounds whichever keys it measured, so a kernel that + # centers on its own would be bounding `q . (k_j - k_mean)` with a norm taken + # over raw `k` -- and `||k||` does not bound `||k - k_mean||`. Centering + # unconditionally therefore silently invalidates every caller that has not + # been taught to centre its norms too. Callers that want centering pass + # `k_mean` and centered norms together; callers that do not get the + # uncentered path, which is sound against the two-sided `floor(W/2)` bound. + + if use_fixed_m and num_q_heads != num_kv_heads: + if num_q_heads % num_kv_heads != 0: + raise ValueError( + f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA ring fixed-m." + ) + q_heads_per_kv_head = num_q_heads // num_kv_heads + if k_mean is not None and k_mean.shape[0] == num_kv_heads: + k_mean = jnp.repeat(k_mean, q_heads_per_kv_head, axis=0) + + if use_fixed_m and k_mean is not None and k_mean.shape[-1] < q.shape[-1]: + k_mean = jnp.pad(k_mean, ((0, 0), (0, q.shape[-1] - k_mean.shape[-1]))) + + global_recenter, global_centered_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=False) + local_recenter, per_shard_bound = custom_splash.get_fixed_m_constants(orig_kv_seq_len, is_ring=True) if bidirectional: if perm is not None or (ring_size is not None and ring_size != axis_size): raise ValueError( @@ -914,68 +965,131 @@ def _custom_ring_attention_forward( mask_value=mask_value, ring_axis=ring_axis, ) + if use_fixed_m and ring_size is not None and ring_size != axis_size: + raise NotImplementedError( + f"fixed-m ring attention requires ring_size == ring axis size (got ring_size={ring_size}, axis_size={axis_size}); " + "use a dedicated ring mesh axis for 2D Ulysses+Ring." + ) if ring_size is None: ring_size = axis_size + canonical_perm = [(i, (i + 1) % axis_size) for i in range(axis_size)] + if use_fixed_m and perm is not None and perm != canonical_perm: + raise NotImplementedError( + "fixed-m ring attention currently requires the canonical ring permutation " + f"[(i, (i + 1) % axis_size)], got perm={perm}." + ) if perm is None: - perm = [(i, (i + 1) % axis_size) for i in range(axis_size)] + perm = canonical_perm shift = partial(lax.ppermute, axis_name=ring_axis, perm=perm) exp_fn = jnp.exp2 if use_base2_exp else jnp.exp - num_q_heads = q.shape[0] - head_dim_v = v.shape[-1] - if use_fixed_m: - # Fixed-m ring: each hop gates PER (head, K-shard) against the halved - # un-smoothed bound, so a head can be fixed on one shard and online on - # another. A fixed hop returns m = the Cauchy-Schwarz upper bound (not the - # rowmax); the naive (m, l) merge below would then flush the other hop's - # partial (exp(m_other - m_bound) underflows once the overshoot exceeds - # the f32 window). Merge in LSE space instead: lse = m + log(l) is - # invariant to the kernel's m convention, so overshoot cancels exactly. - # The K-shard norms rotate WITH K/V (a (heads,)-sized ppermute) instead of - # being re-reduced per hop, which would stall the kernel's scalar prefetch. + # Fixed-m ring: with Global Virtual K-Centering, the keys are centered against + # the global mean across all ring ranks (k_mean = pmean(mean(k), ring_axis)). + # Because logit centering guarantees max_j (q^T (k_j - k_mean)) >= 0 globally, + # the centered Cauchy-Schwarz bound applies across the entire distributed sequence. + # We gather each rank's squared K-shard norms once before the scan: mk_all_sq (R, heads), + # and form mk_global_sq = mk_all_sq.max(axis=0). When all_fixed_global holds, every hop + # evaluates identical m_fixed, enabling direct FP32 (o_sum, l_sum) accumulation. + # In the hybrid fallback branch, individual hops do NOT have a zero-mean guarantee + # in isolation, so per-hop eligibility uses the conservative two-sided per_shard_bound. + # All Cauchy-Schwarz gating is computed in exact squared-norm space (|q|^2 * R_k^2 <= W^2) + # eliminating square roots from all query block and hop gating checks. if fixed_m_norms is None: - raise ValueError("use_fixed_m on the ring path requires fixed_m_norms=(qn_max, mk_h).") + raise ValueError("use_fixed_m on the ring path requires fixed_m_norms=(qn_max_sq, mk_h_sq).") + # The V-magnitude / dtype safety verdict is NOT re-derivable from Q/K norms, + # so the kernel cannot reconstruct it and must not assume it. Treating an + # omitted predicate as permission silently re-enables fixed-m for inputs it + # cannot represent -- e.g. float16 with Q=K=0 and V=1 overflows to inf. Fail + # closed: require the caller to state the verdict explicitly. + if v_ok is None: + raise ValueError( + "use_fixed_m on the ring path requires an explicit `v_ok` predicate " + "(the cross-ring-reduced V-magnitude and dtype safety verdict). Pass " + "v_ok=False to force the online fallback if you have not computed it." + ) log_fn = jnp.log2 if use_base2_exp else jnp.log - qn_max, mk_h_init = fixed_m_norms + qn_norm, mk_norm = fixed_m_norms + # Norm representation is declared by the caller, never inferred. Magnitude + # cannot identify whether norms are squared: legacy unsquared norms of + # (1000, 2) have a true bound of 2000, but any magnitude test that reads + # that product as already-squared yields sqrt(2000) ~= 44.7 and wrongly + # admits fixed-m, which overflows. See fixed_m_norms_squared in the + # make_custom_ring_attention docstring. + if not fixed_m_norms_squared: + qn_max_sq, mk_h_init_sq = qn_norm**2, mk_norm**2 + else: + qn_max_sq, mk_h_init_sq = qn_norm, mk_norm + if num_q_heads != num_kv_heads and mk_h_init_sq.shape[-1] == num_kv_heads: + q_heads_per_kv_head = num_q_heads // num_kv_heads + mk_h_init_sq = jnp.repeat(mk_h_init_sq, q_heads_per_kv_head, axis=-1) tiny = jnp.finfo(jnp.float32).tiny # Finite (not -inf) init: the first merge computes exp(init - lse_new) = 0.0 # exactly; a -inf init meeting an empty partial would produce inf - inf = NaN. lse_init = -1e30 - # Every rank's K-shard norms, gathered ONCE before the scan: (R, heads). - # Rotating mk alongside K/V instead (a third per-hop ppermute feeding the - # kernel's scalar prefetch) serialized the K/V rotation against the kernel - # (trace: collective-permute-done 0.004s -> 0.467s per window); a local - # index into a pre-gathered array keeps the per-hop gate collective-free. - # A caller holding the full table already (e.g. a static weight-derived - # bound, identical on every rank) passes it as (ring_size, heads) and - # skips the gather -- an all_gather of a constant is NOT folded by XLA - # and would still occupy the async-collective machinery every call. - if mk_h_init.ndim == 2: - mk_all = mk_h_init + # Every rank's squared K-shard norms, gathered ONCE before the scan: (R, heads). + # A pre-gathered array keeps the per-hop gate collective-free and avoids + # serializing a third ppermute alongside K/V transfers. + if pregathered_mk or (mk_h_init_sq.ndim == 2 and mk_h_init_sq.shape[0] == axis_size): + mk_all_sq = mk_h_init_sq else: - mk_all = lax.all_gather(mk_h_init, ring_axis) # (axis_size, heads) + mk_all_sq = lax.all_gather(mk_h_init_sq, ring_axis) # (axis_size, heads) my_ring_index = lax.axis_index(ring_axis) - # GLOBAL bound = max over every shard's mk. When ALL local heads pass the - # gate at this single bound, every hop's pinned m is IDENTICAL (it depends - # only on the stationary q rows and the global bound), so hop partials - # combine by PURE ACCUMULATION: o += o_hop, l += l_hop, one normalize at - # the end -- no per-hop LSE math or [H,S,D] divides. The predicate is - # device-uniform ALONG THE RING (the caller pmaxes qn over the ring axis - # and mk_all is a gathered table), so every ppermute participant takes the - # same lax.cond branch; ulysses ranks may diverge freely (no ulysses - # collective lives inside the branches). - mk_global = mk_all.max(axis=0) # (heads,) - all_fixed_global = jnp.all( - qn_max * mk_global <= custom_splash._FIXED_M_RING_SAFE_BOUND # pylint: disable=protected-access - ) + num_q_blocks = (orig_q_seq_len + block_sizes.block_q - 1) // block_sizes.block_q + mk_global_sq = mk_all_sq.max(axis=0) # (heads,) + + # Validate the query-norm rank against `per_q_block`. Both gates below + # multiply qn by `mk[:, None]`, so a (heads,) array supplied while + # per_q_block=True does NOT raise -- it broadcasts to (heads, heads) and + # silently pairs head j's query norm with head h's key norm. A sink head + # then inherits a small bound from some other head, is marked eligible, and + # the kernel evaluates exp2(large_logit - small_m) -> inf. Shape is part of + # the contract, so check it rather than let NumPy guess. + expected_qn_shape = (num_q_heads, num_q_blocks) if per_q_block else (num_q_heads,) + if qn_max_sq.shape != expected_qn_shape: + raise ValueError( + f"fixed_m_norms[0] must have shape {expected_qn_shape} for " + f"per_q_block={per_q_block} (num_q_heads={num_q_heads}, " + f"num_q_blocks={num_q_blocks}), got {qn_max_sq.shape}. A (num_heads,) " + "array with per_q_block=True would broadcast to (heads, heads) and " + "mix head norms together." + ) + if mk_global_sq.shape != (num_q_heads,): + raise ValueError(f"fixed_m_norms[1] must reduce to shape ({num_q_heads},) per rank, got {mk_global_sq.shape}.") + + global_centered_bound_sq = global_centered_bound**2 + per_shard_bound_sq = per_shard_bound**2 + + # Global V-magnitude / dtype safety verdict. Unlike the Cauchy-Schwarz + # norm bounds this is NOT re-derivable from a single hop's Q/K, so it has + # to be carried in and applied to every eligibility decision below -- + # including the per-hop ones in `fixed_body`. + v_gate = True if v_ok is None else v_ok + + if not per_q_block: + bound_sq_1d = qn_max_sq * mk_global_sq + all_fixed_local = jnp.all(bound_sq_1d <= global_centered_bound_sq) & v_gate + all_fixed_global = lax.pmin(all_fixed_local, ring_axis) + m_base_1d = jnp.ceil(jnp.sqrt(bound_sq_1d)) - global_recenter + m_base_expanded = jnp.broadcast_to(m_base_1d[:, None], (num_q_heads, num_q_blocks)) + fixed_ok_expanded = jnp.ones_like(m_base_expanded) + mk_arr = jnp.stack([m_base_expanded, fixed_ok_expanded], axis=0) + qn_blocks_sq = jnp.broadcast_to(qn_max_sq[:, None], (num_q_heads, num_q_blocks)) + else: + qn_blocks_sq = qn_max_sq + bound_blocks_sq = qn_blocks_sq * mk_global_sq[:, None] + fixed_ok_local = bound_blocks_sq <= global_centered_bound_sq # pylint: disable=protected-access + all_fixed_local = jnp.all(fixed_ok_local) & v_gate + all_fixed_global = lax.pmin(all_fixed_local, ring_axis) + m_base = jnp.ceil(jnp.sqrt(bound_blocks_sq)) - global_recenter + fixed_ok_expanded = jnp.ones_like(m_base) + mk_arr = jnp.stack([m_base, fixed_ok_expanded], axis=0) # (2, heads, num_q_blocks) def _accumulate_scan(_): - mk_arr = jnp.stack([mk_global, jnp.ones_like(mk_global)]) o_sum = jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32) l_sum = jnp.zeros((num_q_heads, orig_q_seq_len), jnp.float32) k_current, v_current = k, v @@ -1003,8 +1117,10 @@ def _accumulate_scan(_): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + fixed_m_recenter=global_recenter, use_fixed_m=True, mk=mk_arr, + k_mean=k_mean, # This branch only runs under `all_fixed_global`, so the kernel is # told at compile time that every head is fixed: no per-head scalar # dispatch, and -- load-bearing -- a single body in the ragged last @@ -1036,11 +1152,15 @@ def fixed_body(carry, hop, is_last_hop): # perm src i -> dst i+1: after `hop` shifts this rank holds the K shard # of ring rank (my_index - hop) mod R; its norms come from the local table. - mk_h = jax.lax.dynamic_index_in_dim(mk_all, (my_ring_index - hop) % axis_size, keepdims=False) - fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_RING_SAFE_BOUND).astype( # pylint: disable=protected-access - jnp.float32 - ) - mk_arr = jnp.stack([mk_h, fixed_ok]) + mk_h_sq = jax.lax.dynamic_index_in_dim(mk_all_sq, (my_ring_index - hop) % axis_size, keepdims=False) + bound_hop_sq = qn_blocks_sq * mk_h_sq[:, None] + # `v_gate` is load-bearing here. The Cauchy-Schwarz term is per-hop, but + # V-magnitude and dtype safety are global; recomputing eligibility from + # Q/K norms alone would re-enable fixed-m on this hop even when the + # caller's global V check already rejected it, overflowing to inf. + fixed_ok = ((bound_hop_sq <= per_shard_bound_sq) & v_gate).astype(jnp.float32) # pylint: disable=protected-access + m_base_hop = jnp.ceil(jnp.sqrt(bound_hop_sq)) - local_recenter + mk_arr = jnp.stack([m_base_hop, fixed_ok], axis=0) o_curr, m_curr, l_curr = custom_splash._splash_attention_forward_ring( # pylint: disable=protected-access q, @@ -1052,8 +1172,10 @@ def fixed_body(carry, hop, is_last_hop): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + fixed_m_recenter=local_recenter, use_fixed_m=True, mk=mk_arr, + k_mean=k_mean, ) m_curr = m_curr.astype(jnp.float32) l_curr = l_curr.astype(jnp.float32) @@ -1085,7 +1207,12 @@ def _lse_scan(_): carry, _ = fixed_body(carry, hop, hop == ring_size - 1) return carry[0].astype(q.dtype) - return lax.cond(all_fixed_global, _accumulate_scan, _lse_scan, None) + if uniform_fixed_m is True: + return _accumulate_scan(None) + elif uniform_fixed_m is False: + return _lse_scan(None) + else: + return lax.cond(all_fixed_global, _accumulate_scan, _lse_scan, None) o_init = jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32) l_init = jnp.zeros((num_q_heads, orig_q_seq_len), jnp.float32) @@ -1132,7 +1259,6 @@ def _lse_scan(_): def make_custom_ring_attention( - *, block_sizes: "custom_splash._BlockSizes", orig_q_seq_len: int, orig_kv_seq_len: int, @@ -1146,24 +1272,47 @@ def make_custom_ring_attention( bidirectional: bool = False, use_fixed_m: bool = False, fixed_m_norms: tuple[jax.Array, jax.Array] | None = None, + fixed_m_norms_squared: bool = True, + per_q_block: bool = True, + pregathered_mk: bool = False, + k_mean: jax.Array | None = None, + uniform_fixed_m: bool | None = None, + v_ok: jax.Array | bool | None = None, ): """Builds a forward-only ring-attention callable around the custom kernel. The returned function takes a single (un-batched) `(q, k, v)` triple of shape - `(num_heads, seq, head_dim)` and is meant to be `jax.vmap`-ped over the batch - axis inside the attention `shard_map` (the `ppermute` rotates over `ring_axis`, - which is a mesh axis and independent of the vmap batch axis). - - `ring_size` / `perm` let a caller restrict the rotation to a ring sub-group of - the axis (for the hybrid Ulysses+Ring / USP split); when omitted the rotation - covers the whole `ring_axis`. - - `bidirectional=True` selects the wrap-free schedule (streams K/V both directions - one hop at a time) for a NON-wrapping ring axis, avoiding the diameter-length - wrap hop. Requires `perm=None` and the full real ring axis (no sub-group). + `(num_heads, seq, head_dim)` and optional per-batch `fixed_m_norms=(qn_max_sq, mk_h_sq)` + and `k_mean` to be `jax.vmap`-ped over the batch axis inside the attention `shard_map`. + + `fixed_m_norms_squared` declares the representation of `fixed_m_norms`. The + kernel gates in squared-norm space (`|q|^2 * R_k^2 <= W^2`), which is the + default and what every in-tree caller supplies. Callers holding legacy + unsquared norms must say so by passing False; the representation is never + inferred. It cannot be: magnitude does not distinguish the two. Unsquared + norms of (1000, 2) have a true bound of 2000, and any magnitude test that + reads that product as already-squared gets sqrt(2000) ~= 44.7 -- a ~45x + under-estimate that wrongly admits fixed-m and overflows to inf. + + `v_ok` is a global (already cross-ring-reduced) scalar predicate asserting that + the value magnitudes and activation dtype are safe for fixed-m. It is closed + over rather than passed per call, since it is invariant across the batch. It is + **required** when `use_fixed_m=True`: unlike the Cauchy-Schwarz norm bounds it + is not re-derivable from a single hop's Q/K, so the kernel cannot reconstruct + it, and defaulting it to "safe" silently re-enables fixed-m on inputs that + overflow (float16 with Q=K=0, V=1 yields inf). Pass `v_ok=False` to force the + online fallback. """ + if use_fixed_m and not use_base2_exp: + raise NotImplementedError( + "fixed-m softmax bounds are derived strictly for base-2 exponents. Please set use_base2_exp=True." + ) + default_fixed_m_norms = fixed_m_norms + default_k_mean = k_mean - def _ring(q, k, v): + def _ring(q, k, v, fixed_m_norms=None, k_mean=None): + norms = fixed_m_norms if fixed_m_norms is not None else default_fixed_m_norms + km = k_mean if k_mean is not None else default_k_mean return _custom_ring_attention_forward( q, k, @@ -1180,7 +1329,13 @@ def _ring(q, k, v): perm=perm, bidirectional=bidirectional, use_fixed_m=use_fixed_m, - fixed_m_norms=fixed_m_norms, + fixed_m_norms=norms, + fixed_m_norms_squared=fixed_m_norms_squared, + per_q_block=per_q_block, + pregathered_mk=pregathered_mk, + k_mean=km, + uniform_fixed_m=uniform_fixed_m, + v_ok=v_ok, ) return _ring diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 7b2ba0df7..7992f4ecd 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -840,6 +840,69 @@ def ring_scan_body(carry, _): # --------------------------------------------------------------------------- +def _compute_fixed_m_metadata( + query: jax.Array, + key: jax.Array, + block_q: int, + safe_bound: float | None = None, + recenter: float | None = None, + per_q_block: bool = True, + k_mean: jax.Array | None = None, + value: jax.Array | None = None, + v_max_bound: float = 256.0, +) -> tuple[jax.Array, jax.Array]: + """Computes Cauchy-Schwarz norm bounds and per-Q-block (or per-head) fixed-m metadata.""" + batch_size, num_q_heads, q_len, _ = query.shape + num_kv_heads = key.shape[1] + if safe_bound is None or recenter is None: + rec, bnd = custom_splash.get_fixed_m_constants(key.shape[2], is_ring=False, v_max_bound=v_max_bound) + if safe_bound is None: + safe_bound = bnd + if recenter is None: + recenter = rec + safe_bound_sq = safe_bound**2 + if k_mean is not None: + centered_k = key.astype(jnp.float32) - k_mean[:, :, None, : key.shape[-1]] + mk_h_sq = (centered_k**2).sum(axis=-1).max(axis=-1) + else: + mk_h_sq = (key.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_kv_heads) + + if num_q_heads != num_kv_heads: + if num_q_heads % num_kv_heads != 0: + raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA fixed-m.") + q_heads_per_kv_head = num_q_heads // num_kv_heads + mk_h_sq = jnp.repeat(mk_h_sq, q_heads_per_kv_head, axis=1) # (batch, num_q_heads) + + dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, recenter) + 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) + + num_q_blocks = q_len // block_q + if per_q_block: + norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) # (batch, num_q_heads, q_len) + qn_max_sq = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, block_q).max( + axis=-1 + ) # (batch, num_q_heads, num_q_blocks) + bound_sq = qn_max_sq * mk_h_sq[:, :, None] + fixed_ok = (bound_sq <= safe_bound_sq).astype(jnp.float32) * v_ok + m_base = jnp.ceil(jnp.sqrt(bound_sq)) - recenter + mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, num_q_heads, num_q_blocks) + all_fixed = jnp.all(fixed_ok > 0.5) + else: + qn_max_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_q_heads) + bound_sq_1d = qn_max_sq * mk_h_sq + fixed_ok_1d = (bound_sq_1d <= safe_bound_sq).astype(jnp.float32) * v_ok + m_base_1d = jnp.ceil(jnp.sqrt(bound_sq_1d)) - recenter + m_base_expanded = jnp.broadcast_to(m_base_1d[:, :, None], (batch_size, num_q_heads, num_q_blocks)) + fixed_ok_expanded = jnp.broadcast_to(fixed_ok_1d[:, :, None], (batch_size, num_q_heads, num_q_blocks)) + mk_arr = jnp.stack([m_base_expanded, fixed_ok_expanded], axis=1) # (batch, 2, num_q_heads, num_q_blocks) + all_fixed = jnp.all(fixed_ok_1d > 0.5) + + return mk_arr, all_fixed + + def _ulysses_attention( query: jax.Array, key: jax.Array, @@ -857,21 +920,17 @@ def _ulysses_attention( use_base2_exp: bool = True, use_experimental_scheduler: bool = False, use_fixed_m: bool = False, + per_q_block: bool = True, ulysses_attention_chunks: int = 1, preserve_asymmetric_block_sizes: bool = False, + kv_heads: int | None = None, ) -> jax.Array: - """Ulysses sequence-parallel attention. - - Tensors arrive sequence-sharded on the context axis. Inside a shard_map the - all-to-all collectives trade sequence shards for head shards, run local - splash attention on the full sequence with a subset of heads, then - all-to-all back. - """ + """Ulysses sequence-parallel attention.""" axis_name = CONTEXT num_shards = mesh.shape[axis_name] query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_shards) - key, _ = _reshape_data_for_flash(key, heads, num_shards) + key, orig_kv_seq_len = _reshape_data_for_flash(key, heads, num_shards) value, _ = _reshape_data_for_flash(value, heads, num_shards) attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) if attention_mask is not None and use_custom_kernel: @@ -880,8 +939,6 @@ def _ulysses_attention( "(it only handles padding via orig_seq_len); got a non-None attention_mask." ) num_heads = query.shape[1] - # Ulysses only redistributes existing heads across the context mesh; unlike - # the earlier draft, we fail fast instead of padding synthetic heads. if num_heads % num_shards != 0: raise ValueError( "Ulysses attention requires the number of heads to be divisible by the context shard count, " @@ -923,27 +980,39 @@ def wrap_ulysses_attention(query, key, value, attention_mask): if use_base2_exp: query = query * LOG2E + raw_key = key + raw_query = query + raw_value = value + context_q_seq_len = raw_query.shape[2] + actual_kv_seq_len = orig_kv_seq_len + + real_key = raw_key[:, :, :actual_kv_seq_len, :] + + recenter, safe_bound = custom_splash.get_fixed_m_constants(actual_kv_seq_len, is_ring=False) + + k_mean = None if use_fixed_m: - # k-smoothing (output-invariant): subtracting the per-row key mean - # forces every logit row to have mean 0, hence row-max >= 0 — the - # 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: + k_mean = jnp.pad(k_mean, ((0, 0), (0, 0), (0, 128 - k_mean.shape[-1]))) - query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, bq) - key, _, key_seq_len = _pad_data_for_flash(key, heads, bkv) - value, _, _ = _pad_data_for_flash(value, heads, bkv) + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) + key, _, key_seq_len = _pad_data_for_flash(raw_key, heads, bkv) + value, _, _ = _pad_data_for_flash(raw_value, heads, bkv) mk_arr = None + all_fixed = None if use_fixed_m: - # Per-(local-)head Cauchy-Schwarz inputs over the (batch, seq) slice; - # padded rows have zero norm and never raise the max. mk[0] feeds the - # in-kernel per-query bound, mk[1] flags heads within the no-flush gate. - qf = query.astype(jnp.float32) - kf = key.astype(jnp.float32) - qn_max = jnp.sqrt((qf * qf).sum(-1)).max(axis=(0, 2)) # (local_heads,) - mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=(0, 2)) # (local_heads,) - fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) - mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) + mk_arr, all_fixed = _compute_fixed_m_metadata( + query, + real_key, + block_q=bq, + safe_bound=safe_bound, + recenter=recenter, + per_q_block=per_q_block, + k_mean=k_mean, + value=value, + ) bsizes = custom_splash._BlockSizes( block_q=bq, @@ -952,24 +1021,53 @@ def wrap_ulysses_attention(query, key, value, attention_mask): block_kv_compute_in=bkv_compute_in, ) - splash_kernel = custom_splash.make_splash_mha( - block_sizes=bsizes, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_len, - heads_per_tile=heads_per_tile, - use_base2_exp=use_base2_exp, - use_experimental_scheduler=use_experimental_scheduler, - vmem_limit_bytes=vmem_limit_bytes, - use_fixed_m=use_fixed_m, - ) - if use_fixed_m: - vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None)) - attention_output = vmapped_splash(query, key, value, mk_arr) + splash_kernel_uniform = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_len, + heads_per_tile=heads_per_tile, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=True, + uniform_fixed_m=True, + fixed_m_recenter=recenter, + ) + splash_kernel_hybrid = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_len, + heads_per_tile=heads_per_tile, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=True, + uniform_fixed_m=False, + fixed_m_recenter=recenter, + ) + + def _run_uniform(q, k, v, m, km): + return jax.vmap(splash_kernel_uniform, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + + def _run_hybrid(q, k, v, m, km): + return jax.vmap(splash_kernel_hybrid, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + + raw_out = jax.lax.cond(all_fixed, _run_uniform, _run_hybrid, query, key, value, mk_arr, k_mean) else: + splash_kernel = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_len, + heads_per_tile=heads_per_tile, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=False, + ) vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0)) - attention_output = vmapped_splash(query, key, value) - attention_output = jnp.swapaxes(attention_output, 2, 3) + raw_out = vmapped_splash(query, key, value) + attention_output = jnp.swapaxes(raw_out, 2, 3) attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) else: # Run the same local splash kernel as standard TPU flash attention, but now @@ -1411,10 +1509,8 @@ def wrap_ulysses_ring_attention(query, key, value): # is still globally sharded, so the reduction becomes a per-layer # all-reduce over the context axis: measured WORSE (+54 ms per forward). query, key = jax.lax.optimization_barrier((query, key)) - qn_local = _max_row_norm_per_head(query) - kn_local = _max_row_norm_per_head(key) - if use_base2_exp: - qn_local = qn_local * LOG2E + qn_local = (_max_row_norm_per_head(query) * (LOG2E if use_base2_exp else 1.0)) ** 2 + kn_local = _max_row_norm_per_head(key) ** 2 # The accumulate-vs-LSE lax.cond predicate must be uniform along the RING # axis (every ppermute participant takes the same branch). qn_all = jax.lax.pmax(qn_local, (ring_axis, ulysses_axis)) @@ -1436,27 +1532,51 @@ def wrap_ulysses_ring_attention(query, key, value): if use_base2_exp: query = query * LOG2E + k_mean = None if use_fixed_m and num_ring_shards == 1: - # K-smoothing precondition for fixed-m (R=1 / pure-ulysses semantics, - # same as _ulysses_attention). The R>1 ring path deliberately does NOT - # smooth: no ring rank holds the full K to compute a mean, and a per- - # shard mean would shift each hop's logits differently, breaking the - # cross-shard merge; it gates on the un-smoothed halved bound instead. - kbar = jnp.mean(key, axis=2, keepdims=True) - key = key - kbar + k_mean = jnp.mean(key.astype(jnp.float32), axis=2) + if k_mean.shape[-1] < 128: + k_mean = jnp.pad(k_mean, ((0, 0), (0, 0), (0, 128 - k_mean.shape[-1]))) query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, bq) key, _, key_seq_len = _pad_data_for_flash(key, heads, bkv) value, _, _ = _pad_data_for_flash(value, heads, bkv) + v_ok = None + if use_fixed_m and num_ring_shards > 1: + # V-magnitude and dtype safety are properties of the WHOLE distributed + # problem, not of any single hop, and unlike the Cauchy-Schwarz norm + # bounds they are not re-derivable from a hop's Q/K. The kernel therefore + # cannot reconstruct this verdict, and omitting it let fixed-m run on + # inputs it cannot represent -- float16 with Q=K=0 and V=1 overflows to + # inf instead of returning 1.0. + effective_kv_seq_len = key_seq_len * num_ring_shards + global_recenter, _ = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=False) + dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, global_recenter) + # The sequence pad is zeros, and a max of squares is unchanged by zeros, + # so reading the padded V here is exact. + v_max_sq = (value.astype(jnp.float32) ** 2).max() + v_ok_local = (v_max_sq <= (custom_splash.DEFAULT_MAX_V_BOUND**2)) & dtype_safe + # Reduce over BOTH internal axes: after the all-to-all each device holds a + # head slice of one ring chunk, so neither axis alone sees the whole V. + # The fixed-m branch must also be taken uniformly by every participant of + # the ppermute, which a pmin over both axes guarantees. + v_ok = jax.lax.pmin(v_ok_local, (ring_axis, ulysses_axis)) + mk_arr = None + all_fixed = None if use_fixed_m and num_ring_shards == 1: - qf = query.astype(jnp.float32) - kf = key.astype(jnp.float32) - qn_max = jnp.sqrt((qf * qf).sum(-1)).max(axis=(0, 2)) # (local_heads,) - mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=(0, 2)) # (local_heads,) local - fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) - mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) + recenter, safe_bound = custom_splash.get_fixed_m_constants(key_seq_len, is_ring=False) + mk_arr, all_fixed = _compute_fixed_m_metadata( + query, + key[:, :, :key_seq_len, :], + block_q=bq, + safe_bound=safe_bound, + recenter=recenter, + per_q_block=False, + k_mean=k_mean, + value=value, + ) bsizes = custom_splash._BlockSizes( block_q=bq, @@ -1469,23 +1589,51 @@ def wrap_ulysses_ring_attention(query, key, value): # splash kernel (fuse_reciprocal, no fp32 online-softmax residual windows). # Same math as the 1-step ring, and it fits BQ=8448 where the ring kernel # OOMs (its 3x residual windows). make_splash_mha returns [H, D, S]. - splash_kernel = custom_splash.make_splash_mha( - block_sizes=bsizes, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_len, - heads_per_tile=heads_per_tile, - use_base2_exp=use_base2_exp, - use_experimental_scheduler=use_experimental_scheduler, - vmem_limit_bytes=vmem_limit_bytes, - use_fixed_m=use_fixed_m, - ) if use_fixed_m: - attention_output = jnp.swapaxes( - jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), - 2, - 3, + splash_kernel_uniform = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=query_seq_len, + orig_kv_seq_len=key_seq_len, + heads_per_tile=heads_per_tile, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=True, + uniform_fixed_m=True, + fixed_m_recenter=recenter, + ) + splash_kernel_hybrid = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=query_seq_len, + orig_kv_seq_len=key_seq_len, + heads_per_tile=heads_per_tile, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=True, + uniform_fixed_m=False, + fixed_m_recenter=recenter, ) + + def _run_uniform(q, k, v, m, km): + return jax.vmap(splash_kernel_uniform, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + + def _run_hybrid(q, k, v, m, km): + return jax.vmap(splash_kernel_hybrid, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + + raw_out = jax.lax.cond(all_fixed, _run_uniform, _run_hybrid, query, key, value, mk_arr, k_mean) + attention_output = jnp.swapaxes(raw_out, 2, 3) else: + splash_kernel = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=query_seq_len, + orig_kv_seq_len=key_seq_len, + heads_per_tile=heads_per_tile, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=False, + ) attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value), 2, 3) else: # (2b) Ring (full ppermute over the cross-chip ring axis) with the custom kernel. @@ -1503,6 +1651,8 @@ def wrap_ulysses_ring_attention(query, key, value): bidirectional=bidirectional, use_fixed_m=use_fixed_m, fixed_m_norms=fixed_m_norms, + v_ok=v_ok, + per_q_block=False, ) attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0))(query, key, value) attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) @@ -1774,6 +1924,32 @@ def ulysses_custom_fixed_m_kernel(q, k, v, context): 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, + ulysses_attention_chunks=context.get("ulysses_attention_chunks", 1), + ) + + +@register_kernel("ulysses_custom_fixed_m_per_q_block") +def ulysses_custom_fixed_m_per_q_block_kernel(q, k, v, context): + return _ulysses_attention( + q, + k * context["scale"], + v, + context["heads"], + context["mesh"], + context["axis_names_q"], + context["axis_names_kv"], + context["flash_block_sizes"], + context["dtype"], + mask_padding_tokens=context["mask_padding_tokens"], + residual_checkpoint_name=context["residual_checkpoint_name"], + attention_mask=context["attention_mask"], + use_custom_kernel=True, + 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=True, + ulysses_attention_chunks=context.get("ulysses_attention_chunks", 1), ) diff --git a/src/maxdiffusion/tests/custom_splash_fixed_m_test.py b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py index 0f80a8306..749bdfdbc 100644 --- a/src/maxdiffusion/tests/custom_splash_fixed_m_test.py +++ b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py @@ -87,8 +87,15 @@ def _run_kernel(self, q: jax.Array, k: jax.Array, v: jax.Array, use_fixed_m: boo k_in = k_in - jnp.mean(k_in, axis=1, keepdims=True) qn = jnp.sqrt((q_in.astype(jnp.float32) ** 2).sum(-1)).max(axis=1) mk_h = jnp.sqrt((k_in.astype(jnp.float32) ** 2).sum(-1)).max(axis=1) - eligible = (qn * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) - mk = jnp.stack([mk_h, eligible]) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + bound = qn * mk_h + eligible = (bound <= safe_bound).astype(jnp.float32) + m_base = jnp.ceil(bound) - recenter + num_q_blocks = self.seq_len // self.block_sizes.block_q + mk = jnp.stack([ + jnp.broadcast_to(m_base[:, None], (self.num_heads, num_q_blocks)), + jnp.broadcast_to(eligible[:, None], (self.num_heads, num_q_blocks)), + ]) kernel = custom_splash.make_splash_mha( block_sizes=self.block_sizes, orig_q_seq_len=self.seq_len, @@ -121,14 +128,575 @@ def test_fixed_m_matches_reference(self): fixed, _ = self._run_kernel(q, k, v, use_fixed_m=True) self.assertLess(float(jnp.max(jnp.abs(fixed - self._reference(q, k, v)))), 2e-2) + def _run_kernel_per_q_block( + self, q: jax.Array, k: jax.Array, v: jax.Array, uniform_fixed_m: bool = False + ) -> tuple[jax.Array, jax.Array]: + """Runs the custom kernel with 3D per-Q-block mk inputs.""" + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in = k * self.scale + k_in = k_in - jnp.mean(k_in, axis=1, keepdims=True) + + bq = self.block_sizes.block_q + num_q_blocks = self.seq_len // bq + qf = q_in.astype(jnp.float32) + kf = k_in.astype(jnp.float32) + qf_blocks = qf.reshape(self.num_heads, num_q_blocks, bq, self.head_dim) + qn_max = jnp.sqrt((qf_blocks * qf_blocks).sum(-1)).max(axis=-1) # (heads, num_q_blocks) + mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=1) # (heads,) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + bound = qn_max * mk_h[:, None] + fixed_ok = (bound <= safe_bound).astype(jnp.float32) + m_base = jnp.ceil(bound) - recenter + mk = jnp.stack([m_base, fixed_ok], axis=0) # (2, heads, num_q_blocks) + + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=uniform_fixed_m, + ) + out = kernel(q_in, k_in, v, mk) + out = jnp.swapaxes(out, 1, 2) + return out.astype(jnp.float32), mk + def test_sink_head_falls_back_to_online(self): """An out-of-gate head is flagged ineligible and stays finite (no flush).""" q, k, v = self._random_qkv(q_gain=6.0, k_gain=6.0) fixed, mk = self._run_kernel(q, k, v, use_fixed_m=True) - self.assertEqual(float(mk[1][0]), 0.0) # head 0 is a sink -> ineligible + self.assertTrue(bool(jnp.all(mk[1][0] == 0.0))) # head 0 is a sink -> ineligible self.assertTrue(bool(jnp.all(mk[1][1:] > 0.5))) # the rest stay eligible self.assertTrue(bool(jnp.all(jnp.isfinite(fixed)))) + def test_per_q_block_sink_fallback(self): + """Per-Q-block eligibility keeps normal Q-blocks fixed while sinking outlier blocks.""" + q, k, v = self._random_qkv(k_gain=2.0) + # Amplify only Q-block 1 of Head 0 (bq = 2048, so indices 2048:4096) + q = q.at[0, 2048:].multiply(10.0) + + fixed, mk = self._run_kernel_per_q_block(q, k, v) + # Head 0, Block 0 should be eligible (1.0) + self.assertEqual(float(mk[1, 0, 0]), 1.0) + # Head 0, Block 1 should be ineligible (0.0) due to amplified Q outlier + self.assertEqual(float(mk[1, 0, 1]), 0.0) + # All other heads should be eligible across both blocks + self.assertTrue(bool(jnp.all(mk[1, 1:, :] > 0.5))) + self.assertTrue(bool(jnp.all(jnp.isfinite(fixed)))) + # Check numerical agreement against online kernel running on the same centered inputs + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in_centered = (k * self.scale) - jnp.mean(k * self.scale, axis=1, keepdims=True) + kernel_online = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=False, + ) + online_centered = jnp.swapaxes(kernel_online(q_in, k_in_centered, v), 1, 2).astype(jnp.float32) + self.assertLess(float(jnp.max(jnp.abs(fixed - online_centered))), 1e-2) + + def test_batched_fixed_m_isolation(self): + """Batch-isolated gating ensures outliers in one sample do not contaminate other samples.""" + q0, k0, v0 = self._random_qkv(k_gain=2.0) + # Sample 0 has an outlier in Q-block 1 of Head 0 + q0 = q0.at[0, 2048:].multiply(10.0) + + # Sample 1 is completely clean + q1, k1, v1 = self._random_qkv(k_gain=1.0) + + q = jnp.stack([q0, q1], axis=0) # (2, heads, seq, dim) + k = jnp.stack([k0, k1], axis=0) + v = jnp.stack([v0, v1], axis=0) + + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in = k * self.scale + k_in = k_in - jnp.mean(k_in, axis=2, keepdims=True) + + bq = self.block_sizes.block_q + num_q_blocks = self.seq_len // bq + qf = q_in.astype(jnp.float32) + kf = k_in.astype(jnp.float32) + qf_blocks = qf.reshape(2, self.num_heads, num_q_blocks, bq, self.head_dim) + qn_max_sq = (qf_blocks * qf_blocks).sum(-1).max(axis=-1) # (2, heads, num_q_blocks) + mk_h_sq = (kf * kf).sum(-1).max(axis=-1) # (2, heads) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + bound_sq = qn_max_sq * mk_h_sq[:, :, None] + fixed_ok = (bound_sq <= (safe_bound**2)).astype(jnp.float32) + m_base = jnp.ceil(jnp.sqrt(bound_sq)) - recenter + mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (2, 2, heads, num_q_blocks) + + # Verify Sample 0 has Head 0 Block 1 disqualified (0.0) + self.assertEqual(float(mk_arr[0, 1, 0, 1]), 0.0) + # Verify Sample 1 has ALL heads and ALL blocks eligible (1.0) - zero contamination! + self.assertTrue(bool(jnp.all(mk_arr[1, 1] > 0.5))) + + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=False, + ) + vmapped_kernel = jax.vmap(kernel, in_axes=(0, 0, 0, 0)) + out = vmapped_kernel(q_in, k_in, v, mk_arr) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + + def test_uniform_fixed_matches_hybrid(self): + """Uniform-fixed kernel matches hybrid kernel and f32 reference when all eligible.""" + q, k, v = self._random_qkv() + hybrid_out, mk = self._run_kernel_per_q_block(q, k, v, uniform_fixed_m=False) + uniform_out, _ = self._run_kernel_per_q_block(q, k, v, uniform_fixed_m=True) + ref = self._reference(q, k, v) + + self.assertTrue(bool(jnp.all(mk[1] > 0.5))) + self.assertLess(float(jnp.max(jnp.abs(uniform_out - hybrid_out))), 5e-3) + self.assertLess(float(jnp.max(jnp.abs(uniform_out - ref))), 2e-2) + + def test_missing_mk_raises_value_error(self): + """When use_fixed_m=True, omitting mk raises an immediate ValueError.""" + q, k, v = self._random_qkv() + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + ) + with self.assertRaises(ValueError): + kernel(q, k, v, mk=None) + + def test_legacy_2d_mk_raises_value_error(self): + """Passing a legacy 2D mk array (2, heads) raises ValueError rather than misinterpreting max||k|| as m_B.""" + q, k, v = self._random_qkv() + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + ) + legacy_mk = jnp.zeros((2, self.num_heads), dtype=jnp.float32) + with self.assertRaises(ValueError): + kernel(q, k, v, mk=legacy_mk) + + def test_phase_transition_boundary_continuity(self): + """Verifies seamless output continuity between fixed-m and online mode across the dynamic safe bound threshold.""" + q_base, k_base, v = self._random_qkv() + q_normed = q_base / jnp.sqrt((q_base.astype(jnp.float32) ** 2).sum(-1, keepdims=True)) + k_normed = k_base / jnp.sqrt((k_base.astype(jnp.float32) ** 2).sum(-1, keepdims=True)) + + _, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + test_bounds = [ + safe_bound - 2.0, + safe_bound - 0.5, + safe_bound - 0.01, + safe_bound, + safe_bound + 0.01, + safe_bound + 0.5, + safe_bound + 2.0, + ] + for target_bound in test_bounds: + factor = math.sqrt(target_bound / _LOG2E / self.scale) + q = (q_normed * factor).astype(jnp.bfloat16) + k = (k_normed * factor).astype(jnp.bfloat16) + + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in = k * self.scale + k_in = k_in - jnp.mean(k_in, axis=1, keepdims=True) + + out_gated, _ = self._run_kernel_per_q_block(q, k, v) + kernel_online = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=False, + ) + out_online = jnp.swapaxes(kernel_online(q_in, k_in, v), 1, 2).astype(jnp.float32) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_gated)))) + diff = float(jnp.max(jnp.abs(out_gated - out_online))) + self.assertLess(diff, 2e-2, f"Discontinuity at bound={target_bound}, diff={diff}") + + def test_cpu_proof_invariant_bounds(self): + """Verifies that mathematical underflow and overflow invariants hold across sequence lengths.""" + # Test sequence lengths across short, medium, and production Wan2.2 dimensions + test_lengths = [1, 2, 512, 1024, 16384, 75600, 151200] + + for n in test_lengths: + # 1. Pure Ulysses (Two-sided bound): M in [-U, U] + recenter, safe_bound = custom_splash.get_fixed_m_constants(n, is_ring=False) + # Minimal shifted exponent at worst-case extremum M = -U + exponent_centered = -safe_bound - (math.ceil(safe_bound) - recenter) + self.assertGreaterEqual( + exponent_centered, + -125.0, + f"Underflow violation on Ulysses: {exponent_centered=} for N={n}", + ) + # Non-overflow check with explicit 8-bit FP32 output headroom: log2(N) + C(N) + FP32_OUTPUT_HEADROOM_BITS <= 127 + max_accum_bits = math.ceil(math.log2(n)) + self.assertLessEqual( + max_accum_bits + recenter + custom_splash.FP32_OUTPUT_HEADROOM_BITS, + 127.0, + f"Overflow violation on Ulysses: max bits={max_accum_bits + recenter + custom_splash.FP32_OUTPUT_HEADROOM_BITS} for N={n}", + ) + + # 2. Ring Attention (Uncentered across R hops): M >= -U + for ring_size in [2, 4, 8]: + n_total = n * ring_size + ring_recenter, ring_safe_bound = custom_splash.get_fixed_m_constants(n_total, is_ring=True) + # Minimal shifted exponent at worst-case extremum M = -U + exponent_ring = -ring_safe_bound - (math.ceil(ring_safe_bound) - ring_recenter) + self.assertGreaterEqual( + exponent_ring, + -125.0, + f"Underflow violation on Ring (R={ring_size}): {exponent_ring=} for N={n}", + ) + # Ring direct accumulation non-overflow check: log2(N_total) + C_ring + FP32_OUTPUT_HEADROOM_BITS <= 127 + ring_max_bits = math.ceil(math.log2(n_total)) + self.assertLessEqual( + ring_max_bits + ring_recenter + custom_splash.FP32_OUTPUT_HEADROOM_BITS, + 127.0, + f"Overflow violation on Ring (R={ring_size}): max bits={ring_max_bits + ring_recenter + custom_splash.FP32_OUTPUT_HEADROOM_BITS} for N={n}", + ) + + def test_non_divisible_sequence_context_padding_fixed_m(self): + """Verifies that non-divisible sequences (e.g. S=1001 padded for 8 shards) are correctly masked without zero-padding pollution.""" + seq_len = 1001 + context_shards = 8 + rem = seq_len % context_shards + padded_seq_len = seq_len + (context_shards - rem) # 1008 + heads = 4 + dim = 64 + bq = 512 + + q_raw = jax.random.normal(jax.random.PRNGKey(101), (heads, seq_len, dim), jnp.bfloat16) + k_raw = jax.random.normal(jax.random.PRNGKey(102), (heads, seq_len, dim), jnp.bfloat16) + v_raw = jax.random.normal(jax.random.PRNGKey(103), (heads, seq_len, dim), jnp.bfloat16) + + # Reference dense attention on true unpadded inputs + ref_out = self._reference(q_raw, k_raw, v_raw) + + # Pad inputs as _reshape_data_for_flash would for context sharding + q_pad = jnp.pad(q_raw, ((0, 0), (0, padded_seq_len - seq_len), (0, 0))) + k_pad = jnp.pad(k_raw, ((0, 0), (0, padded_seq_len - seq_len), (0, 0))) + v_pad = jnp.pad(v_raw, ((0, 0), (0, padded_seq_len - seq_len), (0, 0))) + + # Compute unpadded K centering and metadata + k_mean = jnp.mean(k_raw.astype(jnp.float32) * self.scale, axis=1) # (heads, dim) + recenter, safe_bound = custom_splash.get_fixed_m_constants(seq_len, is_ring=False) + + q_in = (q_pad * _LOG2E).astype(jnp.bfloat16) + k_in = (k_pad * self.scale).astype(jnp.bfloat16) + + num_q_blocks = math.ceil(padded_seq_len / bq) + # Pad to systolic block_q boundary + pad_bq = num_q_blocks * bq + q_in_padded = jnp.pad(q_in, ((0, 0), (0, pad_bq - padded_seq_len), (0, 0))) + k_in_padded = jnp.pad(k_in, ((0, 0), (0, pad_bq - padded_seq_len), (0, 0))) + v_in_padded = jnp.pad(v_pad, ((0, 0), (0, pad_bq - padded_seq_len), (0, 0))) + + # Metadata computed on real keys + k_centered = (k_raw.astype(jnp.float32) * self.scale) - k_mean[:, None, :] + mk_h = jnp.sqrt((k_centered**2).sum(-1)).max(axis=-1) + qf_blocks = q_in_padded.astype(jnp.float32).reshape(heads, num_q_blocks, bq, dim) + qn_max = jnp.sqrt((qf_blocks * qf_blocks).sum(-1)).max(axis=-1) + bound = qn_max * mk_h[:, None] + fixed_ok = (bound <= safe_bound).astype(jnp.float32) + m_base = jnp.ceil(bound) - recenter + mk = jnp.stack([m_base, fixed_ok], axis=0) + + block_sizes = custom_splash._BlockSizes(block_q=bq, block_kv=bq, block_kv_compute=bq, block_kv_compute_in=bq) + kernel = custom_splash.make_splash_mha( + block_sizes=block_sizes, + orig_q_seq_len=padded_seq_len, + orig_kv_seq_len=seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=True, + ) + out = jnp.swapaxes(kernel(q_in_padded, k_in_padded, v_in_padded, mk, k_mean), 1, 2).astype(jnp.float32) + out_sliced = out[:, :seq_len, :] + + diff = float(jnp.max(jnp.abs(out_sliced - ref_out))) + self.assertTrue(bool(jnp.all(jnp.isfinite(out_sliced)))) + self.assertLess( + diff, + 2e-2, + f"Non-divisible sequence output diverged from reference: {diff=}", + ) + + def test_pathological_keys_extreme_negative_logits(self): + """Verifies stability when logits are heavily negative and close to underflow.""" + q, k, v = self._random_qkv() + # Shift keys far into negative space so dot products are mostly negative + k_pathological = k - 30.0 + out_fixed, _ = self._run_kernel_per_q_block(q, k_pathological, v) + out_online, _ = self._run_kernel(q, k_pathological, v, use_fixed_m=False) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_fixed)))) + diff = float(jnp.max(jnp.abs(out_fixed - out_online))) + self.assertLess(diff, 2e-2) + + def test_extreme_dynamic_range_inputs(self): + """Verifies that norm computation and gating remain robust with wide dynamic ranges.""" + shape = (self.num_heads, self.seq_len, self.head_dim) + scales = jnp.array([1e-3, 0.1, 0.5, 1.0, 1.5])[:, None, None] + q = (jax.random.normal(jax.random.PRNGKey(42), shape, jnp.bfloat16) * scales).astype(jnp.bfloat16) + k = (jax.random.normal(jax.random.PRNGKey(43), shape, jnp.bfloat16) * scales).astype(jnp.bfloat16) + v = jax.random.normal(jax.random.PRNGKey(44), shape, jnp.bfloat16) + + out, mk = self._run_kernel_per_q_block(q, k, v) + ref = self._reference(q, k, v) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + self.assertLess(float(jnp.max(jnp.abs(out - ref))), 3e-2) + + def test_virtual_k_centering_matches_explicit(self): + """Virtual K-centering with raw keys matches explicit K-centering numerically.""" + q, k, v = self._random_qkv() + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in_raw = (k * self.scale).astype(jnp.bfloat16) + k_mean = jnp.mean(k_in_raw.astype(jnp.float32), axis=1) + + k_in_centered = k_in_raw.astype(jnp.float32) - k_mean[:, None, :] + mk_h_sq = (k_in_centered**2).sum(axis=-1).max(axis=1) + mk_h = jnp.sqrt(mk_h_sq) + + bq = self.block_sizes.block_q + num_q_blocks = self.seq_len // bq + qf = q_in.astype(jnp.float32) + qf_blocks = qf.reshape(self.num_heads, num_q_blocks, bq, self.head_dim) + qn_max = jnp.sqrt((qf_blocks * qf_blocks).sum(-1)).max(axis=-1) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + bound = qn_max * mk_h[:, None] + fixed_ok = (bound <= safe_bound).astype(jnp.float32) + m_base = jnp.ceil(bound) - recenter + mk = jnp.stack([m_base, fixed_ok], axis=0) + + # Virtual K-centering with raw uncentered keys + kernel_virtual = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=True, + ) + out_virtual = jnp.swapaxes(kernel_virtual(q_in, k_in_raw, v, mk, k_mean), 1, 2).astype(jnp.float32) + + # Explicit centering with centered keys + k_centered_bf16 = k_in_centered.astype(jnp.bfloat16) + kernel_explicit = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=True, + ) + out_explicit = jnp.swapaxes(kernel_explicit(q_in, k_centered_bf16, v, mk), 1, 2).astype(jnp.float32) + + ref = self._reference(q, k, v) + diff_virtual_explicit = float(jnp.max(jnp.abs(out_virtual - out_explicit))) + diff_virtual_ref = float(jnp.max(jnp.abs(out_virtual - ref))) + + self.assertLess(diff_virtual_explicit, 2e-3) + self.assertLess(diff_virtual_ref, 2e-2) + self.assertTrue(bool(jnp.all(jnp.isfinite(out_virtual)))) + + def test_virtual_k_centering_per_q_block_hybrid_fallback(self): + """Exercises Virtual K-Centering + Per-Q-Block Hybrid dispatch with mixed fixed/online tiles.""" + q, k, v = self._random_qkv() + bq = self.block_sizes.block_q + num_q_blocks = self.seq_len // bq + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in_raw = (k * self.scale).astype(jnp.bfloat16) + k_mean = jnp.mean(k_in_raw.astype(jnp.float32), axis=1) + + k_in_centered = k_in_raw.astype(jnp.float32) - k_mean[:, None, :] + mk_h_sq = (k_in_centered**2).sum(axis=-1).max(axis=1) + mk_h = jnp.sqrt(mk_h_sq) + + # Test hybrid dispatch where Head 0 Block 0 is Fixed-M and Block 1 is Online Fallback + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + qf_blocks = q_in.astype(jnp.float32).reshape(self.num_heads, num_q_blocks, bq, self.head_dim) + qn_max = jnp.sqrt((qf_blocks * qf_blocks).sum(-1)).max(axis=-1) + bound = qn_max * mk_h[:, None] + m_base = jnp.ceil(bound) - recenter + fixed_ok = jnp.ones((self.num_heads, num_q_blocks), dtype=jnp.float32).at[0, 1].set(0.0) + mk = jnp.stack([m_base, fixed_ok], axis=0) + + # Verify Block 0 is fixed (1.0), Block 1 is online fallback (0.0) on Head 0 + self.assertEqual(float(mk[1, 0, 0]), 1.0) + self.assertEqual(float(mk[1, 0, 1]), 0.0) + + # Hybrid kernel with raw uncentered keys + k_mean + kernel_hybrid = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=False, + ) + out_hybrid = jnp.swapaxes(kernel_hybrid(q_in, k_in_raw, v, mk, k_mean), 1, 2).astype(jnp.float32) + + # Dense f32 reference + ref = self._reference(q, k, v) + diff = float(jnp.max(jnp.abs(out_hybrid - ref))) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_hybrid)))) + self.assertLess(diff, 2e-2, f"Hybrid virtual K output diverged from reference: diff={diff}") + + +class FixedMDtypeSafetyTest(unittest.TestCase): + """P2 regression: dtypes that cannot represent 2**C(N) must not use fixed-m. + + Fixed-m parks the un-normalized softmax weights at up to 2**C(N), a range + derived against FP32's exponent. The kernel narrows them to the activation + dtype for the S@V matmul, so a dtype with a smaller exponent range overflows + to inf even when the FP32 bound analysis passes. + + Backend-agnostic on purpose: this gate is pure Python/jnp, so it should be + enforced in CI even where no TPU is attached. + """ + + def test_float16_is_rejected(self): + recenter, _ = custom_splash.get_fixed_m_constants(4096, is_ring=False) + # C(4096) with |V| <= 256 is 107; float16 tops out at 2**16. + self.assertGreater(recenter, 16.0) + self.assertFalse(custom_splash.fixed_m_dtype_is_safe(jnp.float16, recenter)) + + def test_bfloat16_and_float32_are_accepted(self): + recenter, _ = custom_splash.get_fixed_m_constants(4096, is_ring=False) + self.assertTrue(custom_splash.fixed_m_dtype_is_safe(jnp.bfloat16, recenter)) + self.assertTrue(custom_splash.fixed_m_dtype_is_safe(jnp.float32, recenter)) + + def test_gate_tracks_recenter_not_a_hardcoded_allowlist(self): + """A small enough C(N) is representable even in float16.""" + self.assertTrue(custom_splash.fixed_m_dtype_is_safe(jnp.float16, 4.0)) + self.assertFalse(custom_splash.fixed_m_dtype_is_safe(jnp.float16, 200.0)) + + +class FixedMMetadataSafetyTest(unittest.TestCase): + """Backend-independent regression tests for fixed-m metadata gating.""" + + def test_adversarial_v_magnitude_safely_disqualifies_fixed_m(self): + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + batch = 1 + num_heads = 4 + seq_len = 4096 + dim = 64 + bq = 512 + + q = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + k = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + v_overflow = jnp.full((batch, num_heads, seq_len, dim), 512.0, dtype=jnp.bfloat16) + + # With adversarial V=512 (> 256 default bound), fixed_ok must be 0.0, safely falling back to online + mk_arr, all_fixed = _compute_fixed_m_metadata(q, k, block_q=bq, value=v_overflow) + self.assertFalse(bool(all_fixed)) + self.assertTrue(bool(jnp.all(mk_arr[:, 1] == 0.0))) + + # With normal V <= 256, fixed_ok should remain 1.0 (all eligible) + v_normal = jnp.full((batch, num_heads, seq_len, dim), 1.0, dtype=jnp.bfloat16) + mk_arr_normal, all_fixed_normal = _compute_fixed_m_metadata(q, k, block_q=bq, value=v_normal) + self.assertTrue(bool(all_fixed_normal)) + self.assertTrue(bool(jnp.all(mk_arr_normal[:, 1] == 1.0))) + + def test_float16_query_disqualifies_fixed_m_metadata(self): + """fp16, N=4096, Q=K=0, |V|=1 must report all_fixed=False and fixed_ok=0.""" + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + batch, num_heads, seq_len, dim, bq = 1, 2, 4096, 128, 512 + q = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.float16) + k = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.float16) + v = jnp.full((batch, num_heads, seq_len, dim), 1.0, dtype=jnp.float16) + + mk_arr, all_fixed = _compute_fixed_m_metadata(q, k, block_q=bq, value=v) + self.assertFalse(bool(all_fixed), "fp16 must not be eligible for fixed-m") + self.assertTrue(bool(jnp.all(mk_arr[:, 1] == 0.0))) + + def test_bfloat16_same_case_remains_eligible(self): + """Control: the identical case in bf16 must still take the fast path.""" + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + batch, num_heads, seq_len, dim, bq = 1, 2, 4096, 128, 512 + q = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + k = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + v = jnp.full((batch, num_heads, seq_len, dim), 1.0, dtype=jnp.bfloat16) + + _, all_fixed = _compute_fixed_m_metadata(q, k, block_q=bq, value=v) + self.assertTrue(bool(all_fixed)) + + def test_adversarial_centered_keys_softmax_mass_loss(self): + """Adversarial regression: centered keys with large query norm must NOT be eligible for fixed-m. + + If admitted under a loose one-sided bound (e.g. safe_bound ~ 232), negative logits flush + to zero in FP32 (S - m_base < -126), losing significant softmax probability mass (0.882 vs 1.0, + a silent ~0.118 error). The tightened two-sided bound (safe_bound = safe_window // 2) rejects + this input, ensuring safe fallback to online softmax. + """ + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + seq_len = 4096 + dim = 128 + num_heads = 1 + batch = 1 + bq = 512 + + # Query: [231, 2, 0, ...] + q = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + q = q.at[:, :, :, 0].set(231.0).at[:, :, :, 1].set(2.0) + + # Half keys: [0, 1, 0, ...], half keys: [0, -1, 0, ...] (centered, mean = 0) + k = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + k = k.at[:, :, : seq_len // 2, 1].set(1.0).at[:, :, seq_len // 2 :, 1].set(-1.0) + + # Values: half +1, half -1 + v = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + v = v.at[:, :, : seq_len // 2, :].set(1.0).at[:, :, seq_len // 2 :, :].set(-1.0) + + k_mean = jnp.mean(k.astype(jnp.float32), axis=2) + mk_arr, all_fixed = _compute_fixed_m_metadata(q, k, block_q=bq, k_mean=k_mean, value=v) + + # Must be flagged ineligible for fixed-m under the tightened two-sided bound + self.assertFalse( + bool(all_fixed), + "Adversarial centered input must not be eligible for fixed-m", + ) + self.assertTrue( + bool(jnp.all(mk_arr[:, 1] == 0.0)), + "fixed_ok predicate must be 0.0 across all blocks", + ) + + # Verify that the kernel execution safely falls back to online softmax, preserving + # the negative logit probability mass near 15/17 (~0.8824) rather than flushing to 1.0. + block_sizes = custom_splash._BlockSizes(block_q=bq, block_kv=1024, block_kv_compute=512, block_kv_compute_in=256) + out = custom_splash._splash_attention_forward( + q[0], + k[0], + v[0], + block_sizes=block_sizes, + q_seq_len=seq_len, + kv_seq_len=seq_len, + use_base2_exp=True, + use_fixed_m=True, + mk=mk_arr[0], + k_mean=k_mean[0], + ) + # Transpose from (heads, dim, seq_len) -> (heads, seq_len, dim) + out = jnp.swapaxes(out, 1, 2) + expected = 15.0 / 17.0 + self.assertLess( + float(jnp.max(jnp.abs(out.astype(jnp.float32) - expected))), + 5e-3, + f"Fallback online softmax must preserve negative logit probability mass near 15/17 (~0.8824), got {float(out[0, 0, 0]):.6f}", + ) + if __name__ == "__main__": unittest.main() diff --git a/src/maxdiffusion/tests/ring_fixed_m_test.py b/src/maxdiffusion/tests/ring_fixed_m_test.py index 0f4cc6876..6f29f8a4e 100644 --- a/src/maxdiffusion/tests/ring_fixed_m_test.py +++ b/src/maxdiffusion/tests/ring_fixed_m_test.py @@ -94,9 +94,13 @@ def _reference(self, q_in, k_in, v): logits = jnp.einsum("hqd,hkd->hqk", qf, kf) # LOG2E & scale pre-folded return jax.nn.softmax(logits * math.log(2.0), axis=-1) @ vf - def _run_ring(self, q_in, k_in, v, use_fixed_m): + def _run_ring(self, q_in, k_in, v, use_fixed_m, norms_squared: bool = True, v_ok_override=None): """Runs the custom ring under shard_map with per-rank fixed_m_norms - from the LOCAL q / initial K shard.""" + from the LOCAL q / initial K shard. + + `norms_squared` selects which representation to hand the kernel. Both are + valid so long as they are *declared*; the kernel never infers them. + """ spec = jax.sharding.PartitionSpec(None, _RING_AXIS, None) @functools.partial( @@ -108,12 +112,27 @@ def _run_ring(self, q_in, k_in, v, use_fixed_m): ) def _body(ql, kl, vl): fixed_m_norms = None + v_ok = None if use_fixed_m: qf = ql.astype(jnp.float32) kf = kl.astype(jnp.float32) - qn_max = jnp.sqrt((qf * qf).sum(-1)).max(axis=1) # (heads,) - mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=1) # (heads,) local shard - fixed_m_norms = (qn_max, mk_h) + # Squared norms are the kernel's declared default contract. sqrt is + # monotonic, so max-then-square and square-then-max agree exactly. + qn_max_sq = (qf * qf).sum(-1).max(axis=1) # (heads,) + mk_h_sq = (kf * kf).sum(-1).max(axis=1) # (heads,) local shard + if norms_squared: + fixed_m_norms = (qn_max_sq, mk_h_sq) + else: + fixed_m_norms = (jnp.sqrt(qn_max_sq), jnp.sqrt(mk_h_sq)) + if v_ok_override is None: + # The V/dtype safety verdict the production caller computes. It is + # global, so it is reduced across the ring before use. + v_max_sq = (vl.astype(jnp.float32) ** 2).max() + dtype_safe = custom_splash.fixed_m_dtype_is_safe(ql.dtype, custom_splash._FIXED_M_RECENTER) + v_ok_local = (v_max_sq <= (custom_splash.DEFAULT_MAX_V_BOUND**2)) & dtype_safe + v_ok = jax.lax.pmin(v_ok_local, axis_name=_RING_AXIS) + else: + v_ok = v_ok_override ring = ring_attention_kernel.make_custom_ring_attention( block_sizes=self.block_sizes, orig_q_seq_len=self.shard_len, @@ -123,6 +142,12 @@ def _body(ql, kl, vl): ring_size=_RING_SIZE, use_fixed_m=use_fixed_m, fixed_m_norms=fixed_m_norms, + fixed_m_norms_squared=norms_squared, + v_ok=v_ok, + # These norms are per-head, not per-Q-block, which is what the + # production ring caller supplies. Declaring it keeps the (heads,) + # array from broadcasting against mk[:, None] into (heads, heads). + per_q_block=False, ) return ring(ql, kl, vl) @@ -188,6 +213,213 @@ def test_mixed_fixed_online_across_shards(self): self.assertFalse(bool(gate[0, 1])) self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=True), 2e-2) + def test_declared_unsquared_norms_match_squared(self): + """The two declared norm representations must agree exactly. + + Regression test for the removed magnitude heuristic. Previously the kernel + guessed whether norms were squared by testing `qn.max() * mk.max() < 1000`, + which silently mis-classifies whenever the product straddles that constant + -- reading unsquared norms as squared under-estimates the bound by up to + the square root of its own magnitude, admits fixed-m where it must fall + back, and overflows to inf. With the representation declared rather than + inferred, both spellings of the same inputs must produce the same output. + """ + q, k, v = self._random_qkv(q_gain=(0, slice(0, self.shard_len * _RING_SIZE), 40.0)) + q_in, k_in = self._scaled_inputs(q, k) + out_sq = self._run_ring(q_in, k_in, v, use_fixed_m=True, norms_squared=True).astype(jnp.float32) + out_unsq = self._run_ring(q_in, k_in, v, use_fixed_m=True, norms_squared=False).astype(jnp.float32) + self.assertTrue(bool(jnp.all(jnp.isfinite(out_sq)))) + self.assertTrue(bool(jnp.all(jnp.isfinite(out_unsq)))) + self.assertLess(float(jnp.max(jnp.abs(out_sq - out_unsq))), 2e-2) + + def test_v_ok_false_forces_finite_output(self): + """An explicit unsafe verdict must force the online fallback. + + This is the shape of the FP16 / Q=K=0 / V=1 overflow: when the safety + predicate says no, fixed-m must not run, whatever the Q/K norms imply. + """ + q, k, v = self._random_qkv() + q_in, k_in = self._scaled_inputs(q, k) + out = self._run_ring(q_in, k_in, v, use_fixed_m=True, v_ok_override=False).astype(jnp.float32) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + self.assertLess(float(jnp.max(jnp.abs(out - self._reference(q_in, k_in, v)))), 2e-2) + + +class RingFixedMContractTest(unittest.TestCase): + """Backend-independent checks on the fixed-m ring API contract. + + These assert on errors raised during tracing, so they need neither a TPU nor + a real Pallas lowering and run everywhere CI does. + """ + + def _make(self, **kwargs): + kwargs.setdefault("per_q_block", False) + return ring_attention_kernel.make_custom_ring_attention( + block_sizes=custom_splash._BlockSizes(block_q=128, block_kv=128, block_kv_compute=128, block_kv_compute_in=128), + orig_q_seq_len=128, + orig_kv_seq_len=128, + use_base2_exp=True, + ring_axis=_RING_AXIS, + ring_size=1, + **kwargs, + ) + + def _trace(self, ring, num_heads: int = 1): + """Traces the ring callable under a 1-device mesh; never reaches the device.""" + mesh = jax.sharding.Mesh(np.asarray(jax.devices()[:1]), (_RING_AXIS,)) + spec = jax.sharding.PartitionSpec(None, _RING_AXIS, None) + shape = (num_heads, 128, 128) + + @functools.partial(jax.shard_map, mesh=mesh, in_specs=(spec, spec, spec), out_specs=spec, check_vma=False) + def _body(q, k, v): + return ring(q, k, v) + + zeros = jnp.zeros(shape, jnp.bfloat16) + return jax.eval_shape(_body, zeros, zeros, zeros) + + def test_fixed_m_requires_explicit_v_ok(self): + """Omitting the safety predicate must fail loudly, not default to 'safe'. + + The kernel cannot re-derive the V-magnitude / dtype verdict from a single + hop's Q/K, so treating omission as permission let fixed-m run on inputs it + cannot represent (float16 with Q=K=0 and V=1 returned inf instead of 1.0). + """ + norms = (jnp.ones((1,), jnp.float32), jnp.ones((1,), jnp.float32)) + ring = self._make(use_fixed_m=True, fixed_m_norms=norms) + with self.assertRaises(ValueError) as ctx: + self._trace(ring) + self.assertIn("v_ok", str(ctx.exception)) + + def test_fixed_m_requires_norms(self): + ring = self._make(use_fixed_m=True, v_ok=True) + with self.assertRaises(ValueError) as ctx: + self._trace(ring) + self.assertIn("fixed_m_norms", str(ctx.exception)) + + def test_explicit_v_ok_false_is_accepted(self): + """v_ok=False is a valid answer and must not trip the 'omitted' check.""" + norms = (jnp.ones((1,), jnp.float32), jnp.ones((1,), jnp.float32)) + ring = self._make(use_fixed_m=True, fixed_m_norms=norms, v_ok=False) + self._trace(ring) # must not raise + + def test_per_head_norms_with_per_q_block_are_rejected(self): + """A (heads,) query norm under per_q_block=True must not broadcast. + + This is the defect behind the sink-head CI failure. Both eligibility gates + compute `qn * mk[:, None]`, so a (heads,) array does not raise under + per_q_block=True -- it broadcasts to (heads, heads), pairing head j's query + norm with head h's key norm. A sink head inherits a small bound from an + unrelated head, is wrongly marked fixed-eligible, and the kernel then + evaluates exp2(large_logit - small_m), which overflows to inf. + """ + norms = (jnp.ones((4,), jnp.float32), jnp.ones((4,), jnp.float32)) + ring = self._make(use_fixed_m=True, fixed_m_norms=norms, v_ok=True, per_q_block=True) + with self.assertRaises(ValueError) as ctx: + self._trace(ring, num_heads=4) + self.assertIn("per_q_block", str(ctx.exception)) + + def test_per_q_block_norms_with_correct_shape_are_accepted(self): + """The properly shaped (heads, num_q_blocks) array must pass.""" + # orig_q_seq_len=128 and block_q=128 give exactly one Q block. + norms = (jnp.ones((4, 1), jnp.float32), jnp.ones((4,), jnp.float32)) + ring = self._make(use_fixed_m=True, fixed_m_norms=norms, v_ok=True, per_q_block=True) + self._trace(ring, num_heads=4) # must not raise + + def test_fp16_is_rejected_by_dtype_safety(self): + """The dtype half of the safety predicate must reject narrow exponents. + + float16 has a 5-bit exponent (maxexp 16); fixed-m parks weights at + 2**recenter with recenter ~= 88, which is far beyond float16's ceiling. + """ + recenter, _ = custom_splash.get_fixed_m_constants(4096, is_ring=False) + self.assertFalse(custom_splash.fixed_m_dtype_is_safe(jnp.float16, recenter)) + self.assertTrue(custom_splash.fixed_m_dtype_is_safe(jnp.bfloat16, recenter)) + self.assertTrue(custom_splash.fixed_m_dtype_is_safe(jnp.float32, recenter)) + + +class RingRawKeyBoundUnsoundTest(unittest.TestCase): + """Why virtual K-centering on the ring path has to be opt-in. + + Centering and the Cauchy-Schwarz bound are a matched pair. If a kernel + exponentiates centered logits `q . (k_j - k_mean)` while the caller's + eligibility bound was built from the raw, uncentered `k`, the bound caps the + wrong quantity: `||k||` does not bound `||k - k_mean||`, so a tile can clear + the gate while the centered logit it is supposed to cap overflows fp32. + + That is why `make_custom_ring_attention` centers only when the caller passes + `k_mean` -- a kernel that centered on its own would silently invalidate every + caller that had not also been taught to centre its norms. Callers that want + centering supply `k_mean` and centered norms together; callers that do not + keep the uncentered path, which is sound against the two-sided `floor(W/2)` + bound. + + These assertions are pure arithmetic -- no TPU, no kernel -- so they pin the + failure mode itself rather than one kernel's symptom of it. + """ + + total_kv = 4096 + head_dim = 128 + + def _adversarial_keys(self): + """One key at +100, the rest at -100, on a single active dimension. + + Every key has the same norm (100), so the raw bound is small, but the + population mean sits at ~-99.95 and the lone positive key is ~200 away + from it. + """ + k = jnp.full((self.total_kv,), -100.0, dtype=jnp.float32) + k = k.at[0].set(100.0) + keys = jnp.zeros((self.total_kv, self.head_dim), dtype=jnp.float32).at[:, 0].set(k) + query = jnp.zeros((self.head_dim,), dtype=jnp.float32).at[0].set(1.0) + return query, keys + + def test_raw_bound_admits_a_tile_the_centered_bound_rejects(self): + query, keys = self._adversarial_keys() + _, safe_bound = custom_splash.get_fixed_m_constants(self.total_kv, is_ring=True) + + q_norm = float(jnp.linalg.norm(query)) + raw_bound = q_norm * float(jnp.linalg.norm(keys, axis=-1).max()) + + k_mean = keys.mean(axis=0) + centered_bound = q_norm * float(jnp.linalg.norm(keys - k_mean, axis=-1).max()) + + # The raw bound clears the gate ... + self.assertLessEqual(raw_bound, safe_bound) + # ... but the quantity the kernel actually exponentiates does not. + self.assertGreater(centered_bound, safe_bound) + # The gap is the whole bug: ~100 vs ~200 against a 116 ceiling. + self.assertGreater(centered_bound, 1.9 * raw_bound) + + def test_centered_logit_overflows_fp32_under_the_raw_bound(self): + """With `m` taken from the raw bound, the shifted exponent leaves fp32.""" + query, keys = self._adversarial_keys() + recenter, _ = custom_splash.get_fixed_m_constants(self.total_kv, is_ring=True) + + k_mean = keys.mean(axis=0) + max_centered_logit = float(((keys - k_mean) @ query).max()) + fixed_m = float(jnp.linalg.norm(query)) * float(jnp.linalg.norm(keys, axis=-1).max()) + + # fixed-m parks the max weight at 2**recenter, so the realised exponent is + # (z - m) + recenter. fp32 tops out at 2**128. + shifted_exponent = max_centered_logit - fixed_m + recenter + self.assertGreater(shifted_exponent, 128.0) + + def test_centering_restores_a_sound_bound(self): + """Bounding the centered keys is what makes the gate honest again.""" + query, keys = self._adversarial_keys() + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.total_kv, is_ring=True) + + k_mean = keys.mean(axis=0) + centered = keys - k_mean + centered_bound = float(jnp.linalg.norm(query)) * float(jnp.linalg.norm(centered, axis=-1).max()) + + # Correctly rejected, so this tile takes the online-softmax path. + self.assertGreater(centered_bound, safe_bound) + # And had it been admitted, the bound would genuinely cap the logit. + max_centered_logit = float((centered @ query).max()) + self.assertLessEqual(max_centered_logit, centered_bound + 1e-3) + self.assertLessEqual(max_centered_logit - centered_bound + recenter, 128.0) + if __name__ == "__main__": unittest.main()