diff --git a/src/maxdiffusion/kernels/custom_splash_attention.py b/src/maxdiffusion/kernels/custom_splash_attention.py index 6bd3f3493..cfb354363 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 @@ -48,29 +49,96 @@ 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 -# 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 -# 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). -_FIXED_M_RING_SAFE_BOUND = _FIXED_M_SAFE_BOUND / 2.0 +# running max per KV block, eligible (head, Q-block) tiles subtract a +# precomputed shift built from a Cauchy-Schwarz bound on the logits +# (max_j q_i.k_j <= U = max_i ||q_i|| * max_j ||k_j||). With m = ceil(U) - C(N), +# the worst case over both logit signs is z - m >= C(N) - (U + ceil(U)), so the +# exponent stays >= -125 (above the f32 subnormal-flush floor 2^-126) when +# U <= floor((C(N) + 125) / 2). Tiles above that gate fall back to online +# softmax (the "sink" heads). `get_fixed_m_constants` derives C(N) and the gate +# from the KV length. +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, + 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_B = max_{i in B} ||q_i|| * max_j ||k_j|| be the Cauchy-Schwarz bound on query-key inner products + for Q-block B (per head). With base shift m_B = ceil(U_B) - C(N), the minimal shifted exponent stays >= -125.0 + (1 bit of margin above IEEE-754 normal floor -126.0) as long as U_B <= floor((C(N) + 125.0) / 2). + This bound is two-sided and guarantees no exponent underflows into subnormal range. + """ + + 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 + # With base shift m_B = ceil(U) - C(N), any shifted logit (q·k - m_B) stays in + # [-125, C(N)] as long as U <= floor((C(N) + 125) / 2), keeping exponents within FP32 normal range. + safe_window = recenter - fp32_min_normal_exp - 1.0 + 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 +157,22 @@ def _flash_attention_kernel( fuse_reciprocal: bool = True, use_fixed_m: bool = False, uniform_fixed_m: bool = False, + 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 +192,24 @@ 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 + # 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) + km = k_mean_ref[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) @@ -212,27 +297,38 @@ def compute_body_fixed(kv_compute_index, _): def last_compute_body_online(kv_compute_index): q = q_ref[...] slice_k_len = kv_seq_len % bkv_compute - slice_k = pl.ds(kv_compute_index * bkv_compute, slice_k_len) + aligned_k_len = ((slice_k_len + NUM_SUBLANES - 1) // NUM_SUBLANES) * NUM_SUBLANES + slice_k = pl.ds(kv_compute_index * bkv_compute, aligned_k_len) qk = lax.dot_general(k_ref[slice_k, :], q, NT_DIM_NUMBERS, preferred_element_type=float32) v_chunk = v_ref[slice_k, :] + if slice_k_len != aligned_k_len: + valid_k = jnp.arange(aligned_k_len)[:, None] < slice_k_len + qk = jnp.where(valid_k, qk, mask_value) + v_chunk = jnp.where(valid_k, v_chunk, 0) m_prev, l_prev, o_prev = _online_inner(qk, v_chunk, m_scratch_ref[...], l_scratch_ref[...], o_scratch_ref[:]) m_scratch_ref[...], l_scratch_ref[...] = m_prev, l_prev o_scratch_ref[:] = o_prev def last_compute_body_fixed(kv_compute_index): - # Ragged tail for the pinned fixed-m path: exact slice (padded keys are - # never touched -- with a pinned m their exp2(0 - m_fixed) would be huge - # garbage, so slicing, not masking, is load-bearing here). + # Ragged tail for the pinned fixed-m path: exact slice when aligned to + # NUM_SUBLANES (8); when unaligned, round the VMEM slice up to NUM_SUBLANES + # and mask padded rows to mask_value / 0 so exp2(mask_value - m_fixed) == 0. q = q_ref[...] slice_k_len = kv_seq_len % bkv_compute - slice_k = pl.ds(kv_compute_index * bkv_compute, slice_k_len) + aligned_k_len = ((slice_k_len + NUM_SUBLANES - 1) // NUM_SUBLANES) * NUM_SUBLANES + slice_k = pl.ds(kv_compute_index * bkv_compute, aligned_k_len) qk = lax.dot_general(k_ref[slice_k, :], q, NT_DIM_NUMBERS, preferred_element_type=float32) v_chunk = v_ref[slice_k, :] + if slice_k_len != aligned_k_len: + valid_k = jnp.arange(aligned_k_len)[:, None] < slice_k_len + qk = jnp.where(valid_k, qk, mask_value) + v_chunk = jnp.where(valid_k, v_chunk, 0) l_prev, o_prev = _fixed_inner(qk, v_chunk, m_scratch_ref[...], l_scratch_ref[...], o_scratch_ref[:]) 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: @@ -259,10 +355,16 @@ def body(): # Exactly ONE of these runs in the final KV block -- never both (see the # note on `uniform_fixed_m` above). # - # `_last_online` is the hybrid default: a fixed-m head arrives with - # m_scratch = ceil(bound) - C, and since that is an upper bound on every - # logit the online step's max leaves it unchanged and its rescale factor is - # exp2(0) = 1, so running the last block online is exact for fixed heads too. + # `_last_online` is the hybrid default: a fixed-m head arrives at the final + # KV block (`j == grid_width - 1`) with `m_scratch = ceil(bound) - C(N)`, + # which sits at most `C(N)` below the realized max `m_curr` (since + # `m_curr <= bound`). When + # `_online_inner` runs on that final block, `m_next = jnp.maximum(m_prev, m_curr)` + # updates the running max and rescales `l_prev` and `o_prev` by + # `alpha = exp2(m_prev - m_next) >= 2^-C(N)`. With the default V_max = 256, + # `C(N) <= 119`, so alpha stays above the FP32 subnormal floor (`2^-126`), + # preserving full precision while exporting the updated `m_next` via + # `m_ring_ref`. # `_last_fixed` keeps the bound pinned instead, which is what lets the ring's # accumulate merge assume every hop reports the identical m. def _last_online(): @@ -320,6 +422,110 @@ 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, +): + 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, + 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, + 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, + use_k_centering=use_k_centering, + ) + + def _flash_attention_kernel_mhpt( q_ref, k_ref, @@ -443,7 +649,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 +690,14 @@ def _splash_attention_forward( vmem_limit_bytes: int | None = None, use_fixed_m: bool = False, mk: jax.Array | None = None, + uniform_fixed_m: bool = False, + k_mean: jax.Array | None = None, + interpret: bool | None = None, ): + if interpret is None: + interpret = jax.default_backend() == "cpu" 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 +706,28 @@ 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) + + # 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 +741,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 +753,60 @@ 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, + ) + kernel_args = (mk, q, k, v) + else: + if k_mean.shape != (num_kv_heads, head_dim_qk): + raise ValueError(f"k_mean must have shape ({num_kv_heads}, {head_dim_qk}) indexed by KV head, got {k_mean.shape}") + if q_heads_per_kv_head > 1: + k_mean = jnp.repeat(k_mean, q_heads_per_kv_head, axis=0) + 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, + 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 +821,8 @@ def v_index_map(h, i, j, *_): vmem_limit_bytes=vmem_limit_bytes, ), out_shape=out_shapes, - )(mk, q, k, v) + interpret=interpret, + )(*kernel_args) return all_out[-1] @@ -578,6 +839,8 @@ def _splash_attention_forward_ring( use_fixed_m: bool = False, mk: jax.Array | None = None, uniform_fixed_m: bool = False, + k_mean: jax.Array | None = None, + interpret: bool | None = None, ): """Ring-specific forward path that returns pre-reciprocal fp32 accumulators. @@ -593,6 +856,9 @@ def _splash_attention_forward_ring( - `out` has shape `(num_q_heads, q_seq_len, head_dim_v)` (fp32, un-normalized), - `m` and `l` have shape `(num_q_heads, q_seq_len)` (fp32). """ + if interpret is None: + interpret = jax.default_backend() == "cpu" + num_q_heads, padded_q_seq_len, head_dim_qk = q.shape head_dim_v = v.shape[-1] bq, bkv = block_sizes.block_q, block_sizes.block_kv @@ -603,6 +869,8 @@ 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 def q_index_map(h, i, j, *_): @@ -617,11 +885,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 +905,80 @@ 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 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}") + + 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, + ) + kernel_args = (mk, q, k, v) + else: + if k_mean.shape != (num_kv_heads, head_dim_qk): + raise ValueError(f"k_mean must have shape ({num_kv_heads}, {head_dim_qk}) indexed by KV head, got {k_mean.shape}") + if q_heads_per_kv_head > 1: + k_mean = jnp.repeat(k_mean, q_heads_per_kv_head, axis=0) + 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, + 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 +993,8 @@ def v_index_map(h, i, j, *_): vmem_limit_bytes=vmem_limit_bytes, ), out_shape=out_shapes, - )(mk, q, k, v) + interpret=interpret, + )(*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 +1023,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 +1101,17 @@ 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, + interpret: bool | None = None, ): - def _splash_attention(q, k, v, mk=None): + 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." + ) + + 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 +1139,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, + k_mean=k_mean, + interpret=interpret, ) 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..f45a2b310 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) idx = lax.axis_index(ring_axis) exp_fn = jnp.exp2 if use_base2_exp else jnp.exp @@ -859,6 +863,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 = True, + 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 +892,60 @@ 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: Optional per-KV-head key mean, shape (num_kv_heads, head_dim_qk). + When supplied (together with norms computed on the centered keys), logits + are virtually centered; when None the kernel uses raw keys. + uniform_fixed_m: True forces the fixed-m accumulate path and bypasses the + eligibility gates and `v_ok`; False forces the per-hop LSE path; None + (default) dispatches on `all_fixed_global`. 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) + local_recenter, per_shard_bound = custom_splash.get_fixed_m_constants(orig_kv_seq_len) if bidirectional: if perm is not None or (ring_size is not None and ring_size != axis_size): raise ValueError( @@ -914,68 +967,132 @@ 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: if the caller supplies `k_mean` (and matching centered + # norms), logits are virtually centered; otherwise keys are raw. Either way + # eligibility uses the two-sided floor(W/2) bound, so no row-max >= 0 + # assumption is needed. + # 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). The global gate uses + # floor(W(N_total)/2); when all_fixed_global holds, every hop evaluates the + # identical m_fixed, enabling direct FP32 (o_sum, l_sum) accumulation. + # In the hybrid (LSE-merge) branch each hop is gated on its own, against the + # two-sided bound for the local shard length, floor(W(N_local)/2). + # All Cauchy-Schwarz gating is computed in squared-norm space + # (|q|^2 * R_k^2 <= floor(W/2)^2), with no square roots in the gates. 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 @@ -1005,7 +1122,9 @@ def _accumulate_scan(_): vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=True, mk=mk_arr, - # This branch only runs under `all_fixed_global`, so the kernel is + k_mean=k_mean, + # This branch runs under `all_fixed_global` (or when the caller + # forces it with `uniform_fixed_m=True`), 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 # KV block instead of two (a two-body last block is what triggers @@ -1036,11 +1155,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, @@ -1054,6 +1177,7 @@ def fixed_body(carry, hop, is_last_hop): vmem_limit_bytes=vmem_limit_bytes, 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 +1209,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 +1261,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 +1274,51 @@ 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 <= floor(W/2)^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). With the default + `uniform_fixed_m=None`, pass `v_ok=False` to force the online fallback. + + `uniform_fixed_m`: True forces the fixed-m accumulate path and bypasses the + eligibility gates and `v_ok`; False forces the per-hop LSE path; None + (default) dispatches on `all_fixed_global`. """ + 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 +1335,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 3db542586..2d6b161d2 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -846,6 +846,76 @@ 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], 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) + + # The kernel's grid is ceil(q_len / block_q); callers pad Q to a multiple of + # block_q first, so floor == ceil here. Fail loudly if that contract breaks + # rather than silently dropping the ragged tail's gating metadata. + if q_len % block_q != 0: + raise ValueError( + f"_compute_fixed_m_metadata expects query padded to a multiple of block_q, got q_len={q_len}, block_q={block_q}." + ) + 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, @@ -863,10 +933,12 @@ 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, spatiotemporal_config: Optional[dict] = None, spatiotemporal_shape: Optional[Tuple[int, int, int]] = None, + kv_heads: int | None = None, ) -> jax.Array: """Ulysses sequence-parallel attention. @@ -879,7 +951,7 @@ def _ulysses_attention( 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: @@ -888,8 +960,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, " @@ -931,27 +1001,40 @@ 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) + + 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) + kv_pad_size = 1 if actual_kv_seq_len % 8 == 0 else bkv + key, _, key_seq_len = _pad_data_for_flash(raw_key, heads, kv_pad_size) + value, _, _ = _pad_data_for_flash(raw_value, heads, kv_pad_size) 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, @@ -960,24 +1043,51 @@ 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, + ) + 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, + ) + + 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 @@ -1419,10 +1529,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)) @@ -1444,27 +1552,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) + 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) + 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, @@ -1477,23 +1609,49 @@ 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, + ) + 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, ) + + 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. @@ -1511,6 +1669,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) @@ -1784,6 +1944,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..0d16b5839 100644 --- a/src/maxdiffusion/tests/custom_splash_fixed_m_test.py +++ b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py @@ -17,8 +17,8 @@ """Unit tests for the fixed-m path of the custom splash attention kernel. The fixed-m optimization replaces the online-softmax running max with a -precomputed per-query Cauchy-Schwarz bound for eligible heads, falling back to -online softmax for "sink" heads whose bound exceeds the no-flush gate. These +precomputed Cauchy-Schwarz bound for eligible (head, Q-block) tiles, falling +back to online softmax for "sink" tiles whose bound exceeds the no-flush gate. These tests check that, mirroring the production calling convention, the kernel: * matches an f32 softmax reference for both online and fixed-m modes, @@ -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) + 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,734 @@ 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) + 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): + """Per-sample mk rows from `_compute_fixed_m_metadata` stay independent under vmap. + + An outlier in sample 0 disqualifies only sample 0's tile in `mk`; sample 1 + stays fully eligible. Note that production's `all_fixed` is reduced over + the whole batch, so the outlier still routes the batch to the hybrid + kernel, where each sample's tiles follow their own `mk`. + """ + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + 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) + + # Keys are already centered above, so no k_mean is passed. + mk_arr, all_fixed = _compute_fixed_m_metadata(q_in, k_in, block_q=self.block_sizes.block_q) + self.assertEqual(mk_arr.shape, (2, 2, self.num_heads, self.seq_len // self.block_sizes.block_q)) + + # Sample 0: only head 0, block 1 is disqualified. + self.assertEqual(float(mk_arr[0, 1, 0, 1]), 0.0) + self.assertEqual(float(mk_arr[0, 1, 0, 0]), 1.0) + self.assertTrue(bool(jnp.all(mk_arr[0, 1, 1:] > 0.5))) + # Sample 1: every head and block stays eligible. + self.assertTrue(bool(jnp.all(mk_arr[1, 1] > 0.5))) + # The batch-wide predicate is False, so production would run the hybrid kernel. + self.assertFalse(bool(all_fixed)) + + 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): + """Output stays continuous as the Cauchy-Schwarz bound sweeps across the fixed-m gate. + + Cases well below the gate must run fixed-m and cases well above must fall + back; the cases within bf16 rounding of the gate may land on either side. + """ + 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) + 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, mk = self._run_kernel_per_q_block(q, k, v) + if target_bound <= safe_bound - 2.0: + self.assertTrue(bool(jnp.all(mk[1] > 0.5)), f"expected fixed-m at bound={target_bound}") + if target_bound >= safe_bound + 2.0: + self.assertTrue(bool(jnp.all(mk[1] == 0.0)), f"expected fallback at bound={target_bound}") + 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] + headroom_bits = math.ceil(math.log2(custom_splash.DEFAULT_MAX_V_BOUND)) # |V| <= 256 -> 8 bits + + 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) + # 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 the default |V| headroom: ceil(log2(N)) + C(N) + headroom_bits <= 127 + max_accum_bits = math.ceil(math.log2(n)) + self.assertLessEqual( + max_accum_bits + recenter + headroom_bits, + 127.0, + f"Overflow violation on Ulysses: max bits={max_accum_bits + recenter + 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) + # 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: ceil(log2(N_total)) + C_ring + headroom_bits <= 127 + ring_max_bits = math.ceil(math.log2(n_total)) + self.assertLessEqual( + ring_max_bits + ring_recenter + headroom_bits, + 127.0, + f"Overflow violation on Ring (R={ring_size}): max bits={ring_max_bits + ring_recenter + 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) + + 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_anti_aligned_uncentered_keys_heavily_negative_logits(self): + """Fixed-m stays exact when every logit is heavily negative and keys are NOT centered. + + This is the two-sided worst case the floor(W/2) gate exists for: q and k + are anti-aligned, so every base-2 logit sits near -U (row max ~ -109) while + U stays just under the gate. With m = ceil(U) - C(N) the shifted exponents + land near -117, close to but above the -126 flush floor. No k_mean is + passed, so nothing re-centers the logits. + """ + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + shape = (self.num_heads, self.seq_len, self.head_dim) + # Kernel-domain inputs (already base-2 scaled); all values are exact in bf16. + q_sign = jnp.where(jax.random.bernoulli(jax.random.PRNGKey(7), 0.5, shape[:2]), 0.5, -0.5) + q_in = jnp.zeros(shape, jnp.float32).at[:, :, 0].set(10.5).at[:, :, 1].set(q_sign).astype(jnp.bfloat16) + c = jax.random.uniform(jax.random.PRNGKey(8), shape[:2], jnp.float32, -2.0, 2.0) + k_in = jnp.zeros(shape, jnp.float32).at[:, :, 0].set(-10.5).at[:, :, 1].set(c).astype(jnp.bfloat16) + v = jax.random.normal(jax.random.PRNGKey(9), shape, jnp.bfloat16) + + bq = self.block_sizes.block_q + mk_arr, all_fixed = _compute_fixed_m_metadata(q_in[None], k_in[None], block_q=bq) + _, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len) + qf, kf = q_in.astype(jnp.float32), k_in.astype(jnp.float32) + logits = jnp.einsum("hsd,htd->hst", qf, kf) + u = float(jnp.sqrt((qf**2).sum(-1).max() * (kf**2).sum(-1).max())) + self.assertGreater(u, safe_bound - 5.0) # near the gate ... + self.assertTrue(bool(all_fixed)) # ... but admitted + self.assertLess(float(logits.max()), -100.0) # every logit heavily negative + + 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, + ) + out_fixed = jnp.swapaxes(kernel(q_in, k_in, v, mk_arr[0]), 1, 2).astype(jnp.float32) + ref = jax.nn.softmax(logits * math.log(2.0), axis=-1) @ v.astype(jnp.float32) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_fixed)))) + self.assertLess(float(jnp.max(jnp.abs(out_fixed - ref))), 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) + 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) + 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) + # 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) + 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), the negative-logit terms + (1/17 ~ 5.9% of the softmax mass) flush to zero in FP32 (S - m_base < -126), so the output + collapses to 1.0 instead of 15/17 ~ 0.882 (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 safely falls back to online softmax: the output is + # 15/17 (~0.8824) rather than the flushed 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 output must be ~15/17 (~0.8824), got {float(out[0, 0, 0]):.6f}", + ) + + +class FixedMAttentionFlaxIntegrationTest(unittest.TestCase): + """End-to-end fixed-m through the attention_flax entry points. + + The tests above drive the kernel directly. These go through the production + wrappers (`_ulysses_attention` and the R=1 branch of + `_ulysses_ring_custom_attention`), which build the metadata, pick the + uniform/hybrid kernel with `lax.cond` and call `make_splash_mha`. They run on + a 1-device mesh, so on CPU the Pallas kernel runs in interpret mode. + """ + + batch = 1 + seq_len = 256 + num_heads = 2 + head_dim = 128 + + def setUp(self): + super().setUp() + from flax.linen import partitioning as nn_partitioning + import numpy as np + from jax.sharding import Mesh + from maxdiffusion.models import attention_flax + + self.af = attention_flax + self.nn_partitioning = nn_partitioning + self.mesh = Mesh(np.array(jax.devices()[:1]).reshape(1, 1, 1, 1), ("data", "fsdp", "context", "tensor")) + self.rules = ( + (attention_flax.BATCH, "data"), + (attention_flax.SELF_ATTN_HEAD, None), + (attention_flax.SELF_ATTN_Q_LENGTH, "context"), + (attention_flax.SELF_ATTN_KV_LENGTH, "context"), + (attention_flax.D_KV, None), + ) + self.axis_names_q = ( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_Q_LENGTH, + attention_flax.D_KV, + ) + self.axis_names_kv = ( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_KV_LENGTH, + attention_flax.D_KV, + ) + self.block_sizes = {"block_q": 128, "block_kv": 128, "block_kv_compute": 128, "block_kv_compute_in": 128} + + b, s, h, d = self.batch, self.seq_len, self.num_heads, self.head_dim + kq, kk, kv = jax.random.split(jax.random.PRNGKey(0), 3) + self.q = jax.random.normal(kq, (b, s, h * d), jnp.float32).astype(jnp.bfloat16) + # The wrappers do not apply 1/sqrt(d); fold it into K like the Wan caller does. + self.k = (jax.random.normal(kk, (b, s, h * d), jnp.float32) * d**-0.5).astype(jnp.bfloat16) + self.v = jax.random.normal(kv, (b, s, h * d), jnp.float32).astype(jnp.bfloat16) + + def _reference(self) -> jax.Array: + """Plain fp32 softmax attention on the same bf16 inputs, shape (b, s, h*d).""" + b, s, h, d = self.batch, self.seq_len, self.num_heads, self.head_dim + q, k, v = (x.astype(jnp.float32).reshape(b, s, h, d) for x in (self.q, self.k, self.v)) + probs = jax.nn.softmax(jnp.einsum("bshd,bthd->bhst", q, k), axis=-1) + return jnp.einsum("bhst,bthd->bshd", probs, v).reshape(b, s, h * d) + + def _assert_inputs_take_fixed_m(self): + """Guards against the test silently exercising only the online fallback.""" + b, s, h, d = self.batch, self.seq_len, self.num_heads, self.head_dim + q = (self.q.reshape(b, s, h, d).transpose(0, 2, 1, 3) * _LOG2E).astype(jnp.bfloat16) + k = self.k.reshape(b, s, h, d).transpose(0, 2, 1, 3) + k_mean = jnp.mean(k.astype(jnp.float32), axis=2) + _, all_fixed = self.af._compute_fixed_m_metadata(q, k, block_q=128, k_mean=k_mean, value=self.v) + self.assertTrue(bool(all_fixed)) + + def _assert_close_to_reference(self, out: jax.Array): + self.assertEqual(out.shape, self.q.shape) + out = out.astype(jnp.float32) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + # Measured max|err| is ~5e-3 on CPU interpret mode, mostly the bf16 rounding + # of an O(1) output (half an ulp at [1, 2) is ~3.9e-3); 1e-2 leaves headroom. + self.assertLess(float(jnp.max(jnp.abs(out - self._reference()))), 1e-2) + + def test_ulysses_attention_fixed_m_matches_reference(self): + self._assert_inputs_take_fixed_m() + for per_q_block in (True, False): + with self.subTest(per_q_block=per_q_block): + with self.mesh, self.nn_partitioning.axis_rules(self.rules): + out = self.af._ulysses_attention( + self.q, + self.k, + self.v, + heads=self.num_heads, + mesh=self.mesh, + axis_names_q=self.axis_names_q, + axis_names_kv=self.axis_names_kv, + flash_block_sizes=self.block_sizes, + dtype=jnp.bfloat16, + use_custom_kernel=True, + use_fixed_m=True, + per_q_block=per_q_block, + ) + self._assert_close_to_reference(out) + + def test_ulysses_ring_custom_r1_fixed_m_matches_reference(self): + self._assert_inputs_take_fixed_m() + with self.mesh, self.nn_partitioning.axis_rules(self.rules): + out = self.af._ulysses_ring_custom_attention( + self.q, + self.k, + self.v, + heads=self.num_heads, + mesh=self.mesh, + axis_names_q=self.axis_names_q, + axis_names_kv=self.axis_names_kv, + flash_block_sizes=self.block_sizes, + dtype=jnp.bfloat16, + ulysses_shards=1, + use_fixed_m=True, + ) + self._assert_close_to_reference(out) + 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..01e9d8e13 100644 --- a/src/maxdiffusion/tests/ring_fixed_m_test.py +++ b/src/maxdiffusion/tests/ring_fixed_m_test.py @@ -16,10 +16,11 @@ """Unit tests for the fixed-m path of the custom RING attention. -The ring path gates fixed-m PER (head, K-shard) against the halved -un-smoothed bound, rotates each K shard's max row norm alongside K/V, and -merges the per-hop partials in LSE space (invariant to fixed-m's bound -overshoot). These tests check, against an f32 dense-softmax reference: +The ring path gates fixed-m globally against floor(W(N_total)/2) (accumulate +merge when every (head, shard) passes) and otherwise PER (head, K-shard) +against floor(W(N_local)/2), merging the per-hop partials in LSE space +(invariant to fixed-m's bound overshoot). These tests check, against an f32 +dense-softmax reference: * the untouched online ring path (regression guard), * fixed-m with every (head, shard) eligible, @@ -53,8 +54,6 @@ class RingFixedMTest(unittest.TestCase): def setUp(self): super().setUp() - if jax.default_backend() != "tpu": - self.skipTest("Only supported on TPUs.") if len(jax.devices()) < _RING_SIZE: self.skipTest(f"Requires {_RING_SIZE} devices.") self.scale = 1.0 / math.sqrt(self.head_dim) @@ -94,9 +93,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 +111,28 @@ 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() + recenter, _ = custom_splash.get_fixed_m_constants(self.shard_len * _RING_SIZE) + dtype_safe = custom_splash.fixed_m_dtype_is_safe(ql.dtype, 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,26 +142,40 @@ 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) return _body(q_in, k_in, v) + def _shard_max_sq_norms(self, x): + """(heads, ring_size) max squared row norm of each rank's shard.""" + xf = x.astype(jnp.float32) + sq = (xf * xf).sum(-1).reshape(self.num_heads, _RING_SIZE, self.shard_len) + return sq.max(axis=-1) + def _gate_per_shard(self, q_in, k_in): - """(heads, ring_size) eligibility against the halved un-smoothed bound.""" - qf = q_in.astype(jnp.float32) - kf = k_in.astype(jnp.float32) - qn = jnp.sqrt((qf * qf).sum(-1)) # (heads, total) - kn = jnp.sqrt((kf * kf).sum(-1)) - gates = [] - for r in range(_RING_SIZE): - rows = slice(r * self.shard_len, (r + 1) * self.shard_len) - # Stationary q max is per-RANK, but for the gate check we use the global - # q max: it upper-bounds every rank's local max, so "eligible globally" - # implies eligible on every rank. - bound = qn.max(axis=1) * kn[:, rows].max(axis=1) - gates.append(bound <= custom_splash._FIXED_M_RING_SAFE_BOUND) - return jnp.stack(gates, axis=1) + """(heads, q_rank, k_shard) per-hop eligibility, as the kernel's LSE path computes it. + + Each rank gates its stationary local q against every K shard with the + two-sided bound for the local shard length, floor(W(shard_len)/2). + """ + _, per_shard_bound = custom_splash.get_fixed_m_constants(self.shard_len) + qn_sq = self._shard_max_sq_norms(q_in) # (heads, q_rank) + kn_sq = self._shard_max_sq_norms(k_in) # (heads, k_shard) + return qn_sq[:, :, None] * kn_sq[:, None, :] <= per_shard_bound**2 + + def _gate_global(self, q_in, k_in): + """(heads,) global eligibility; all True means the kernel takes the accumulate merge.""" + _, global_bound = custom_splash.get_fixed_m_constants(self.shard_len * _RING_SIZE) + qn_sq = self._shard_max_sq_norms(q_in).max(axis=1) + kn_sq = self._shard_max_sq_norms(k_in).max(axis=1) + return qn_sq * kn_sq <= global_bound**2 def _run_and_compare(self, q, k, v, use_fixed_m): q_in, k_in = self._scaled_inputs(q, k) @@ -153,21 +186,28 @@ def _run_and_compare(self, q, k, v, use_fixed_m): def _gate(self, q, k): return self._gate_per_shard(*self._scaled_inputs(q, k)) + def _global_gate(self, q, k): + return self._gate_global(*self._scaled_inputs(q, k)) + def test_online_ring_matches_reference(self): q, k, v = self._random_qkv() self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=False), 2e-2) def test_fixed_m_all_eligible_matches_reference(self): q, k, v = self._random_qkv() + self.assertTrue(bool(jnp.all(self._global_gate(q, k)))) # accumulate merge self.assertTrue(bool(jnp.all(self._gate(q, k)))) self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=True), 2e-2) def test_sink_head_falls_back_everywhere(self): total = self.shard_len * _RING_SIZE q, k, v = self._random_qkv(q_gain=(0, slice(0, total), 40.0)) + global_gate = self._global_gate(q, k) + self.assertFalse(bool(global_gate[0])) # forces the per-hop LSE path + self.assertTrue(bool(jnp.all(global_gate[1:]))) gate = self._gate(q, k) - self.assertFalse(bool(jnp.any(gate[0]))) # head 0 online on every shard - self.assertTrue(bool(jnp.all(gate[1:]))) + self.assertFalse(bool(jnp.any(gate[0]))) # head 0 online on every (rank, shard) + self.assertTrue(bool(jnp.all(gate[1:]))) # other heads fixed on every hop self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=True), 2e-2) def test_fixed_m_accumulate_ragged_tail(self): @@ -176,6 +216,7 @@ def test_fixed_m_accumulate_ragged_tail(self): # covering the pinned fixed-m path's exact-slice tail handling. self.block_sizes = custom_splash._BlockSizes(block_q=1024, block_kv=768, block_kv_compute=384, block_kv_compute_in=384) q, k, v = self._random_qkv() + self.assertTrue(bool(jnp.all(self._global_gate(q, k)))) self.assertTrue(bool(jnp.all(self._gate(q, k)))) self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=True), 2e-2) @@ -183,11 +224,215 @@ def test_mixed_fixed_online_across_shards(self): # Amplify head 0's keys on shard 1 only: head 0 is fixed on shard 0 but # online on shard 1 -- the mixed-partial merge the LSE space exists for. q, k, v = self._random_qkv(k_gain=(0, slice(self.shard_len, self.shard_len * _RING_SIZE), 40.0)) + self.assertFalse(bool(self._global_gate(q, k)[0])) # forces the per-hop LSE path gate = self._gate(q, k) - self.assertTrue(bool(gate[0, 0])) - self.assertFalse(bool(gate[0, 1])) + self.assertTrue(bool(jnp.all(gate[0, :, 0]))) # every rank: fixed on shard 0 + self.assertFalse(bool(jnp.any(gate[0, :, 1]))) # every rank: online on shard 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 produce the same output. + + The norm representation is declared via `fixed_m_norms_squared`; it can't + be inferred from magnitude. 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. + + 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 = C(4096) = 107, far beyond float16's ceiling. + """ + recenter, _ = custom_splash.get_fixed_m_constants(4096) + 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) + + 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) + + 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) + + 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()