diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml index bf8e1c740..85d5f8b59 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -83,16 +83,26 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, tokamax_ring_custom, ulysses, ulysses_custom, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_bidir +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, tokamax_ring_custom, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block, ulysses_ring_custom_bidir # # Best 2D-ring / USP (Ulysses x ring) configs for WAN2.2-T2V-A14B (720x1280, 81 frames) # Set attention=ulysses_ring_custom and ulysses_shards=U (ring degree R=CP/U): # CP4 (v7x-8): ulysses_shards=2 (R=2), BQ=9472 # CP8 (v7x-8): ulysses_shards=4 (R=2), BQ=9472 # CP16 (v7x-16): ulysses_shards=8 (R=2), BQ=9472 +# +# WARNING: ulysses_shards only has an effect on the *ring* variants. +# - Setting U == CP gives R=1, which is a DEGENERATE ring: no KV is rotated +# and the result is mathematically identical to the non-ring +# ulysses_custom* kernel. Such a run must not be reported as a ring result. +# The attention layer logs a warning when this happens. +# - The non-ring ulysses/ulysses_custom* kernels always use the full context +# axis, so their Ulysses degree is fixed at CP. Passing a different +# ulysses_shards to them now raises instead of being silently ignored. use_base2_exp: True use_experimental_scheduler: True -# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. +use_k_centering: False +# For attention=ulysses_ring*, hidden Ulysses shard count; ring shards are context / this. ulysses_shards: -1 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. # For communication-compute overlap to be effective, enable the following XLA flags: diff --git a/src/maxdiffusion/kernels/fused_producers.py b/src/maxdiffusion/kernels/fused_producers.py new file mode 100644 index 000000000..cccd6da66 --- /dev/null +++ b/src/maxdiffusion/kernels/fused_producers.py @@ -0,0 +1,109 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Optimized fused producers for Wan Attention.""" + +from typing import Tuple + +import jax +import jax.numpy as jnp + + +def fused_rmsnorm_rope( + raw_q: jax.Array, + raw_k: jax.Array, + q_norm_scale: jax.Array, + k_norm_scale: jax.Array, + freqs_cis: jax.Array, + q_heads: int = 40, + kv_heads: int | None = None, + dim_head: int = 128, + eps: float = 1e-6, + heads: int | None = None, +) -> Tuple[jax.Array, jax.Array]: + """Fusion-friendly FP32 RMSNorm + BF16 RoPE + Head Transposition producer. + + Performs FP32 RMSNorm normalization for maximum stability, casts to input dtype + (e.g. BF16), and applies RoPE rotation and head transposition in BF16 precision, + avoiding excess FP32 VPU/VMEM cycles on long sequence lengths. Fully supports GQA + where q_heads != kv_heads. + + Args: + raw_q: Raw query projection of shape [B, Sq, Dq] (where Dq = q_heads * dim_head). + raw_k: Raw key projection of shape [B, Sk, Dk] (where Dk = kv_heads * dim_head). + q_norm_scale: RMSNorm scale parameter for query of shape [Dq]. + k_norm_scale: RMSNorm scale parameter for key of shape [Dk]. + freqs_cis: Complex rotary embedding tensor of shape [1, 1, S, dim_head // 2]. + q_heads: Number of query attention heads. + kv_heads: Number of key/value attention heads (defaults to q_heads for MHA). + dim_head: Dimension of each attention head. + eps: Epsilon for RMSNorm numerical stability. + heads: Deprecated alias for q_heads. + + Returns: + Transposed and RoPE-rotated (q_out, k_out) of shapes [B, q_heads, Sq, dim_head] + and [B, kv_heads, Sk, dim_head]. + """ + if heads is not None: + q_heads = heads + kv_heads = q_heads if kv_heads is None else kv_heads + B, Sq, Dq = raw_q.shape + _, Sk, Dk = raw_k.shape + + if Dq != q_heads * dim_head: + raise ValueError(f"raw_q feature dim ({Dq}) must equal q_heads ({q_heads}) * dim_head ({dim_head})") + if Dk != kv_heads * dim_head: + raise ValueError(f"raw_k feature dim ({Dk}) must equal kv_heads ({kv_heads}) * dim_head ({dim_head})") + + # 1. FP32 RMSNorm for stability, then cast directly to target activation dtype. + # + # Association matters: Flax's `_normalize` computes `mul = rsqrt(var + eps)`, + # then `mul *= scale`, then `y = x * mul` -- i.e. x * (rsqrt * scale). Folding + # left-to-right as (x * rsqrt) * scale rounds differently and makes this path + # drift from `nnx.RMSNorm` bit-for-bit. Keep the parenthesisation below in + # step with Flax so the fused producer stays a pure fusion, not a numerical + # change. + q_fp32 = raw_q.astype(jnp.float32) + q_rms = jax.lax.rsqrt(jnp.mean(jnp.square(q_fp32), axis=-1, keepdims=True) + eps) + q_norm = (q_fp32 * (q_rms * q_norm_scale.astype(jnp.float32))).astype(raw_q.dtype) + + k_fp32 = raw_k.astype(jnp.float32) + k_rms = jax.lax.rsqrt(jnp.mean(jnp.square(k_fp32), axis=-1, keepdims=True) + eps) + k_norm = (k_fp32 * (k_rms * k_norm_scale.astype(jnp.float32))).astype(raw_k.dtype) + + # 2. Reshape and transpose to [B, heads, S, dim_head] + q_h = q_norm.reshape(B, Sq, q_heads, dim_head).transpose(0, 2, 1, 3) + k_h = k_norm.reshape(B, Sk, kv_heads, dim_head).transpose(0, 2, 1, 3) + + # 3. Direct RoPE with freqs_cis [1, 1, S, dim_head // 2] in input dtype + cos = jnp.real(freqs_cis).astype(raw_q.dtype) + sin = jnp.imag(freqs_cis).astype(raw_q.dtype) + cos_q, sin_q = cos[:, :, :Sq, :], sin[:, :, :Sq, :] + cos_k, sin_k = cos[:, :, :Sk, :], sin[:, :, :Sk, :] + + q_pairs = q_h.reshape(B, q_heads, Sq, -1, 2) + q_0, q_1 = q_pairs[..., 0], q_pairs[..., 1] + q_out_0 = q_0 * cos_q - q_1 * sin_q + q_out_1 = q_0 * sin_q + q_1 * cos_q + q_out = jnp.stack([q_out_0, q_out_1], axis=-1).reshape(B, q_heads, Sq, dim_head) + + k_pairs = k_h.reshape(B, kv_heads, Sk, -1, 2) + k_0, k_1 = k_pairs[..., 0], k_pairs[..., 1] + k_out_0 = k_0 * cos_k - k_1 * sin_k + k_out_1 = k_0 * sin_k + k_1 * cos_k + k_out = jnp.stack([k_out_0, k_out_1], axis=-1).reshape(B, kv_heads, Sk, dim_head) + + return q_out, k_out diff --git a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py index 67f735651..4f301ebb1 100644 --- a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py +++ b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py @@ -870,6 +870,7 @@ def _custom_ring_attention_forward( k_mean: jax.Array | None = None, uniform_fixed_m: bool | None = None, v_ok: jax.Array | bool | None = None, + all_fixed_global: jax.Array | bool | None = None, ) -> jax.Array: """Forward-only ring attention using the custom dense splash kernel. @@ -936,13 +937,16 @@ def _custom_ring_attention_forward( 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 k_mean is not None and k_mean.shape[0] == num_q_heads: + k_mean = k_mean[::q_heads_per_kv_head] - 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]))) + if use_fixed_m and k_mean is not None: + pad_h = (custom_splash.NUM_SUBLANES - (k_mean.shape[0] % custom_splash.NUM_SUBLANES)) % custom_splash.NUM_SUBLANES + pad_d = max(0, q.shape[-1] - k_mean.shape[-1]) + if pad_h > 0 or pad_d > 0: + k_mean = jnp.pad(k_mean, ((0, pad_h), (0, pad_d))) - global_recenter, global_centered_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=False) + global_recenter, global_centered_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=True) local_recenter, per_shard_bound = custom_splash.get_fixed_m_constants(orig_kv_seq_len, is_ring=True) if bidirectional: if perm is not None or (ring_size is not None and ring_size != axis_size): @@ -990,13 +994,14 @@ def _custom_ring_attention_forward( # the global mean across all ring ranks (k_mean = pmean(mean(k), ring_axis)). # Because logit centering guarantees max_j (q^T (k_j - k_mean)) >= 0 globally, # the centered Cauchy-Schwarz bound applies across the entire distributed sequence. - # We gather each rank's squared K-shard norms once before the scan: mk_all_sq (R, heads), - # and form mk_global_sq = mk_all_sq.max(axis=0). When all_fixed_global holds, every hop - # evaluates identical m_fixed, enabling direct FP32 (o_sum, l_sum) accumulation. - # In the hybrid fallback branch, individual hops do NOT have a zero-mean guarantee - # in isolation, so per-hop eligibility uses the conservative two-sided per_shard_bound. - # All Cauchy-Schwarz gating is computed in exact squared-norm space (|q|^2 * R_k^2 <= W^2) - # eliminating square roots from all query block and hop gating checks. + # Ring-wide max squared K-shard norm: mk_global_sq (heads,). When all_fixed_global + # holds, every hop evaluates identical m_fixed, enabling direct FP32 (o_sum, l_sum) + # accumulation. In the hybrid fallback branch, individual hops do NOT have a + # zero-mean guarantee in isolation, so eligibility uses the conservative two-sided + # per_shard_bound evaluated against mk_global_sq (see `_lse_scan` for why the + # per-hop index was removed). All Cauchy-Schwarz gating is computed in exact + # squared-norm space (|q|^2 * R_k^2 <= W^2) eliminating square roots from all + # query block and hop gating checks. if fixed_m_norms is None: raise ValueError("use_fixed_m on the ring path requires fixed_m_norms=(qn_max_sq, mk_h_sq).") # The V-magnitude / dtype safety verdict is NOT re-derivable from Q/K norms, @@ -1030,17 +1035,18 @@ def _custom_ring_attention_forward( # exactly; a -inf init meeting an empty partial would produce inf - inf = NaN. lse_init = -1e30 - # 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 + # Ring-wide max K norm: (heads,). When `mk_h_init_sq` is passed as a 2D + # `(R, heads)` array, collapse the hop axis directly. When 1D `(heads,)`, + # either it is already ring-reduced (`pregathered_mk=True`) or reduce across + # `ring_axis` via `pmax` rather than `all_gather(...).max(axis=0)`. + if mk_h_init_sq.ndim == 2: + mk_global_sq = mk_h_init_sq.max(axis=0) + elif pregathered_mk: + mk_global_sq = mk_h_init_sq else: - mk_all_sq = lax.all_gather(mk_h_init_sq, ring_axis) # (axis_size, heads) - my_ring_index = lax.axis_index(ring_axis) + mk_global_sq = lax.pmax(mk_h_init_sq, ring_axis) 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 @@ -1072,24 +1078,28 @@ def _custom_ring_attention_forward( 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)) + if uniform_fixed_m is None and all_fixed_global is None: + all_fixed_local = jnp.all(bound_sq_1d <= global_centered_bound_sq) & v_gate + all_fixed_global = lax.pmin(all_fixed_local, ring_axis) 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) + if uniform_fixed_m is None and all_fixed_global is 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) def _accumulate_scan(_): + if not per_q_block: + 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) + else: + 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) 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 @@ -1139,7 +1149,7 @@ def _accumulate_scan(_): # the conditional, ~half of fixed-m's whole kernel win). return (o_sum * l_inv[..., None]).astype(q.dtype) - def fixed_body(carry, hop, is_last_hop): + def fixed_body(carry, is_last_hop, mk_arr): o_run, lse_run, k_current, v_current = carry # Prefetch the next shard while computing on this one. The last hop skips # it: nothing consumes the rotated shard, and the collective would still @@ -1150,18 +1160,6 @@ def fixed_body(carry, hop, is_last_hop): k_next = shift(k_current) v_next = shift(v_current) - # 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_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, k_current, @@ -1194,17 +1192,63 @@ def fixed_body(carry, hop, is_last_hop): o_new = (w_run[..., None] * o_run + w_curr[..., None] * o_norm) / denom[..., None] return (o_new, lse_new + log_fn(denom), k_next, v_next), None - fixed_init = ( - jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32), - jnp.full((num_q_heads, orig_q_seq_len), lse_init, jnp.float32), - k, - v, - ) - def _lse_scan(_): - carry = fixed_init + # Precompute the scalar-prefetch metadata BEFORE the ring loop so no VPU + # compute or SMEM scalar-prefetch barrier sits between `shift(k_current)` + # (`collective-permute-start`) and `_splash_attention_forward_ring`. + # + # The bound uses the ring-wide `mk_global_sq` rather than indexing + # `mk_all_sq` at `(my_ring_index - hop) % axis_size`: + # + # * Correctness. Taking the max over hops can only enlarge the + # Cauchy-Schwarz bound, so `fixed_ok` is never set where the per-hop + # bound would have cleared it -- the gate gets strictly more + # conservative, never less. Under the pregathered path (see + # `_ring_fixed_m_norms_pre_a2a`) every hop already carries the + # ring-wide max, so this is exact rather than merely safe. + # * Performance. Worth ~5.7s per 40-step 720p denoise on a v7x-8 + # (120.6s -> 114.9s), and it is the collectives, confirmed by trace: + # + # category main before after + # collective-permute-done 38.0 2405.4 47.3 ms + # all-to-all 2135.7 3920.5 2088.7 ms + # + # `my_ring_index` is a traced `lax.axis_index`, so indexing with it + # put the ring index on the kernel's scalar-prefetch operand. The + # kernel could not be issued until that resolved, and the ~190 MiB + # K/V `ppermute` issued just before it -- which exists precisely to + # be hidden behind that kernel -- was left fully exposed, running at + # 43 GiB/s. It then contended with the output all-to-all, which fell + # from 506 GiB/s to 173 GiB/s on a byte-identical transfer. Removing + # the index restores both: the rotation hides again and the output + # a2a returns to 535 GiB/s. + # + # Why this costs so much when the branch runs only ~23% of the time: + # the `lax.cond` boundary. XLA's latency-hiding scheduler will not + # move a collective across a conditional, so a ppermute inside a + # branch has only that branch's own body to hide behind -- and one + # scalar-prefetch dependency is enough to consume all of it. + # Consistent with that, pinning `uniform_fixed_m=False`, which takes + # the `elif` below and emits no cond at all, independently recovers + # most of the same time (measured 116.0s) despite running this branch + # 100% of the time. + # + # All hops therefore share one metadata array; it is still materialised + # per hop below to keep `fixed_body`'s signature unchanged. + bound_hop_sq = qn_blocks_sq * mk_global_sq[:, None] + 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_uniform = jnp.stack([m_base_hop, fixed_ok], axis=0) + mk_arr_hops = [mk_arr_uniform] * ring_size + + carry = ( + jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32), + jnp.full((num_q_heads, orig_q_seq_len), lse_init, jnp.float32), + k, + v, + ) for hop in range(ring_size): - carry, _ = fixed_body(carry, hop, hop == ring_size - 1) + carry, _ = fixed_body(carry, hop == ring_size - 1, mk_arr_hops[hop]) return carry[0].astype(q.dtype) if uniform_fixed_m is True: @@ -1262,6 +1306,7 @@ def make_custom_ring_attention( block_sizes: "custom_splash._BlockSizes", orig_q_seq_len: int, orig_kv_seq_len: int, + *, use_base2_exp: bool = True, use_experimental_scheduler: bool = False, vmem_limit_bytes: int | None = None, @@ -1278,6 +1323,7 @@ def make_custom_ring_attention( k_mean: jax.Array | None = None, uniform_fixed_m: bool | None = None, v_ok: jax.Array | bool | None = None, + all_fixed_global: jax.Array | bool | None = None, ): """Builds a forward-only ring-attention callable around the custom kernel. @@ -1336,6 +1382,7 @@ def _ring(q, k, v, fixed_m_norms=None, k_mean=None): k_mean=km, uniform_fixed_m=uniform_fixed_m, v_ok=v_ok, + all_fixed_global=all_fixed_global, ) return _ring diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 7992f4ecd..67ac0648b 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -28,6 +28,7 @@ from maxdiffusion.kernels.splash_attention import splash_attention_kernel as tokamax_splash_attention_kernel from maxdiffusion.kernels.splash_attention import ring_attention_kernel as tokamax_ring_attention_kernel from maxdiffusion.kernels.splash_attention import base as tokamax_splash_base +from maxdiffusion.kernels.fused_producers import fused_rmsnorm_rope from einops import rearrange from .. import common_types, max_logging from maxdiffusion.tpu_utils import get_tpu_type, TpuType @@ -68,6 +69,18 @@ def _coerce_tokamax_block_sizes(block_sizes): + if isinstance(block_sizes, dict): + return splash_attention_kernel.BlockSizes( + block_q=block_sizes.get("block_q", 512), + block_kv=block_sizes.get("block_kv", 512), + block_kv_compute=block_sizes.get("block_kv_compute", 512), + block_q_dkv=block_sizes.get("block_q_dkv", 512), + block_kv_dkv=block_sizes.get("block_kv_dkv", 512), + block_kv_dkv_compute=block_sizes.get("block_kv_dkv_compute", 512), + block_q_dq=block_sizes.get("block_q_dq", None), + block_kv_dq=block_sizes.get("block_kv_dq", None), + use_fused_bwd_kernel=block_sizes.get("use_fused_bwd_kernel", False), + ) # Tokamax requires fused bwd; convert if needed. if getattr(block_sizes, "use_fused_bwd_kernel", False): return block_sizes @@ -179,6 +192,95 @@ def _replace_mesh_axis_names(axis_names, old_axis: str, new_axes: tuple[str, ... return jax.sharding.PartitionSpec(*(_replace_mesh_axis(axis_name, old_axis, new_axes) for axis_name in axis_names)) +# Attention kernels are traced once per layer per transformer, so an +# unconditional log would repeat dozens of times per run and be ignored. +_WARNED_ONCE: set[str] = set() + + +def _warn_once(key: str, message: str) -> None: + """Logs `message` the first time `key` is seen in this process.""" + if key in _WARNED_ONCE: + return + _WARNED_ONCE.add(key) + max_logging.log(message) + + +def _validate_implicit_ulysses_degree(requested_ulysses_shards: int, context_shards: int, kernel_name: str) -> None: + """Rejects a `ulysses_shards` request the non-ring Ulysses path cannot honour. + + The non-ring kernels always shard heads across the *entire* context mesh + axis, so their Ulysses degree is implicitly `context_shards`. Silently + ignoring a different explicit request has previously caused benchmarks to + believe they were measuring U=2 while actually measuring U=4. + """ + if requested_ulysses_shards is None or requested_ulysses_shards <= 0: + return # Unset: the implicit degree is what the caller wants. + if requested_ulysses_shards == context_shards: + return # Explicit request agrees with what this path will do. + raise ValueError( + f"attention='{kernel_name}' cannot honour ulysses_shards={requested_ulysses_shards}: " + f"the non-ring Ulysses path always splits heads across the full context mesh axis, " + f"so its Ulysses degree is fixed at context_shards={context_shards}. " + f"Either set ulysses_shards={context_shards} (or leave it unset), or switch to a ring " + f"variant such as 'ulysses_ring_custom_fixed_m_per_q_block', which accepts " + f"ulysses_shards=U and forms a ring of degree R=context_shards/U." + ) + + +def _largest_ulysses_shards_for_real_ring(context_shards: int, heads: int | None = None, kv_heads: int | None = None): + """Largest Ulysses degree that still leaves a real ring (R > 1), or None if impossible. + + A usable Ulysses degree U must divide the context shard count *and* both head + counts, mirroring the constraints the ring path itself enforces. Returning the + largest such U below `context_shards` yields the smallest ring degree R > 1, + which is normally the cheapest real ring for a given mesh. + """ + for candidate in range(context_shards - 1, 0, -1): + if context_shards % candidate != 0: + continue + if heads is not None and heads % candidate != 0: + continue + if kv_heads is not None and kv_heads % candidate != 0: + continue + return candidate + return None + + +def _warn_if_ring_is_degenerate( + num_ring_shards: int, + num_ulysses_shards: int, + context_shards: int, + heads: int | None = None, + kv_heads: int | None = None, +) -> None: + """Warns when a ring variant collapses to R=1 and is really running plain Ulysses.""" + if num_ring_shards != 1: + return + + if context_shards <= 1: + advice = ( + "There is only one context shard, so no ring is possible on this mesh; " + "increase ici_context_parallelism to use a ring." + ) + else: + suggestion = _largest_ulysses_shards_for_real_ring(context_shards, heads, kv_heads) + if suggestion is None: + advice = ( + f"No ulysses_shards below context_shards={context_shards} divides both the mesh and the " + f"head counts (heads={heads}, kv_heads={kv_heads}), so this mesh cannot form a real ring." + ) + else: + advice = f"For a real ring set ulysses_shards={suggestion} (ring degree R={context_shards // suggestion})." + + _warn_once( + f"degenerate_ring:{context_shards}:{num_ulysses_shards}", + f"[attention] Ring degree R=1 (context_shards={context_shards} / ulysses_shards={num_ulysses_shards}). " + f"This ring variant is degenerate: no KV is rotated, the cross-ring mean for virtual K-centering " + f"is skipped, and the result is mathematically identical to the corresponding non-ring " + f"'ulysses_custom*' kernel. Do NOT report this as a ring-attention result. " + advice, + ) + + def _create_internal_ulysses_ring_mesh( mesh: Mesh, ring_shards: int, @@ -446,7 +548,10 @@ def _build_padding_segment_ids( kv_mask_for_batch = jnp.concatenate( [ kv_mask_for_batch, - jnp.zeros((attention_mask.shape[0], kv_padded_len - key_seq_len), jnp.int32), + jnp.zeros( + (attention_mask.shape[0], kv_padded_len - key_seq_len), + jnp.int32, + ), ], axis=1, ) @@ -563,6 +668,11 @@ def _run_chunked_ulysses_attention( Returns: The concatenated attention output tensor. """ + if query.shape[1] != key.shape[1] and ulysses_attention_chunks > 1: + raise NotImplementedError( + f"GQA (query heads {query.shape[1]} != key heads {key.shape[1]}) with " + f"ulysses_attention_chunks={ulysses_attention_chunks} > 1 is not supported." + ) head_chunk_ranges = _ulysses_head_chunk_ranges(num_heads, ulysses_shards, ulysses_attention_chunks) if len(head_chunk_ranges) > 1: chunk_outputs = [ @@ -840,69 +950,6 @@ def ring_scan_body(carry, _): # --------------------------------------------------------------------------- -def _compute_fixed_m_metadata( - query: jax.Array, - key: jax.Array, - block_q: int, - safe_bound: float | None = None, - recenter: float | None = None, - per_q_block: bool = True, - k_mean: jax.Array | None = None, - value: jax.Array | None = None, - v_max_bound: float = 256.0, -) -> tuple[jax.Array, jax.Array]: - """Computes Cauchy-Schwarz norm bounds and per-Q-block (or per-head) fixed-m metadata.""" - batch_size, num_q_heads, q_len, _ = query.shape - num_kv_heads = key.shape[1] - if safe_bound is None or recenter is None: - rec, bnd = custom_splash.get_fixed_m_constants(key.shape[2], is_ring=False, v_max_bound=v_max_bound) - if safe_bound is None: - safe_bound = bnd - if recenter is None: - recenter = rec - safe_bound_sq = safe_bound**2 - if k_mean is not None: - centered_k = key.astype(jnp.float32) - k_mean[:, :, None, : key.shape[-1]] - mk_h_sq = (centered_k**2).sum(axis=-1).max(axis=-1) - else: - mk_h_sq = (key.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_kv_heads) - - if num_q_heads != num_kv_heads: - if num_q_heads % num_kv_heads != 0: - raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA fixed-m.") - q_heads_per_kv_head = num_q_heads // num_kv_heads - mk_h_sq = jnp.repeat(mk_h_sq, q_heads_per_kv_head, axis=1) # (batch, num_q_heads) - - dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, recenter) - v_ok = 1.0 if dtype_safe else 0.0 - if dtype_safe and value is not None: - v_max_sq = (value.astype(jnp.float32) ** 2).max() - v_ok = (v_max_sq <= (v_max_bound**2)).astype(jnp.float32) - - num_q_blocks = q_len // block_q - if per_q_block: - norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) # (batch, num_q_heads, q_len) - qn_max_sq = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, block_q).max( - axis=-1 - ) # (batch, num_q_heads, num_q_blocks) - bound_sq = qn_max_sq * mk_h_sq[:, :, None] - fixed_ok = (bound_sq <= safe_bound_sq).astype(jnp.float32) * v_ok - m_base = jnp.ceil(jnp.sqrt(bound_sq)) - recenter - mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, num_q_heads, num_q_blocks) - all_fixed = jnp.all(fixed_ok > 0.5) - else: - qn_max_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_q_heads) - bound_sq_1d = qn_max_sq * mk_h_sq - fixed_ok_1d = (bound_sq_1d <= safe_bound_sq).astype(jnp.float32) * v_ok - m_base_1d = jnp.ceil(jnp.sqrt(bound_sq_1d)) - recenter - m_base_expanded = jnp.broadcast_to(m_base_1d[:, :, None], (batch_size, num_q_heads, num_q_blocks)) - fixed_ok_expanded = jnp.broadcast_to(fixed_ok_1d[:, :, None], (batch_size, num_q_heads, num_q_blocks)) - mk_arr = jnp.stack([m_base_expanded, fixed_ok_expanded], axis=1) # (batch, 2, num_q_heads, num_q_blocks) - all_fixed = jnp.all(fixed_ok_1d > 0.5) - - return mk_arr, all_fixed - - def _ulysses_attention( query: jax.Array, key: jax.Array, @@ -920,30 +967,53 @@ def _ulysses_attention( use_base2_exp: bool = True, use_experimental_scheduler: bool = False, use_fixed_m: bool = False, - per_q_block: bool = True, ulysses_attention_chunks: int = 1, preserve_asymmetric_block_sizes: bool = False, - kv_heads: int | None = None, + per_q_block: bool = True, + kv_heads: Optional[int] = None, + ulysses_shards: int = -1, + kernel_name: str = "ulysses_custom", ) -> jax.Array: - """Ulysses sequence-parallel attention.""" + """Ulysses sequence-parallel attention. + + Tensors arrive sequence-sharded on the context axis. Inside a shard_map the + all-to-all collectives trade sequence shards for head shards, run local + splash attention on the full sequence with a subset of heads, then + all-to-all back. + + The Ulysses degree of this path is implicitly the full context mesh axis; an + explicit `ulysses_shards` that disagrees is rejected rather than ignored. + """ axis_name = CONTEXT num_shards = mesh.shape[axis_name] + _validate_implicit_ulysses_degree(ulysses_shards, num_shards, kernel_name) + if kv_heads is None: + kv_heads = heads query, orig_q_seq_len = _reshape_data_for_flash(query, 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) + key, orig_kv_seq_len = _reshape_data_for_flash(key, kv_heads, num_shards) + value, _ = _reshape_data_for_flash(value, kv_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: raise NotImplementedError( "The custom dense splash kernel (use_custom_kernel) does not support attention_mask " "(it only handles padding via orig_seq_len); got a non-None attention_mask." ) - num_heads = query.shape[1] - if num_heads % num_shards != 0: + num_q_heads = query.shape[1] + num_kv_heads = key.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_q_heads % num_shards != 0: raise ValueError( - "Ulysses attention requires the number of heads to be divisible by the context shard count, " - f"got heads={num_heads} and context_shards={num_shards}." + "Ulysses attention requires the number of query heads to be divisible by the context shard count, " + f"got q_heads={num_q_heads} and context_shards={num_shards}." ) + if num_kv_heads % num_shards != 0: + raise ValueError( + "Ulysses attention requires the number of KV heads to be divisible by the context shard count, " + f"got kv_heads={num_kv_heads} and context_shards={num_shards}." + ) + num_heads = num_q_heads if not use_custom_kernel: block_sizes = _select_flash_block_sizes( @@ -961,6 +1031,13 @@ def _ulysses_attention( mask_needs_ulysses_gather = _mesh_axis_in_spec(kv_axis_names[2], axis_name) def wrap_ulysses_attention(query, key, value, attention_mask): + # Apply the base-2 rescale of Q *before* the all-to-all. A scalar elementwise + # multiply commutes exactly with the collective (which is pure data movement), + # so this is bit-identical. Done after the a2a it sat between the collective + # and the kernel and XLA wrapped it in relayout copies; done before, it fuses + # into the producer of Q and its 185MB round-trip disappears. + if use_custom_kernel and use_base2_exp: + query = query * LOG2E # Swap sharding: each device gives up a slice of heads and gathers # a slice of sequence, so the local kernel sees the full sequence. query = jax.lax.all_to_all(query, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) @@ -975,11 +1052,16 @@ def wrap_ulysses_attention(query, key, value, attention_mask): "The custom dense splash kernel (use_custom_kernel) does not support attention_mask " "(it only handles padding via orig_seq_len); got a non-None attention_mask." ) - bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes) - - if use_base2_exp: - query = query * LOG2E + ( + bq, + bkv, + bkv_compute, + bkv_compute_in, + heads_per_tile, + vmem_limit_bytes, + ) = _extract_custom_block_sizes(flash_block_sizes) + # NOTE: the base-2 rescale of Q is applied before the all-to-all above. raw_key = key raw_query = query raw_value = value @@ -990,15 +1072,27 @@ def wrap_ulysses_attention(query, key, value, attention_mask): recenter, safe_bound = custom_splash.get_fixed_m_constants(actual_kv_seq_len, is_ring=False) + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) k_mean = None if use_fixed_m: + # Virtual k-centering (output-invariant): project q^T \bar{k} inside the + # kernel registers without writing back / materializing (K - \bar{k}) in HBM. + # Computed strictly on real (unpadded) tokens. 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(raw_query, heads, bq) - key, _, key_seq_len = _pad_data_for_flash(raw_key, heads, bkv) - value, _, _ = _pad_data_for_flash(raw_value, heads, bkv) + pad_h = (custom_splash.NUM_SUBLANES - (k_mean.shape[1] % custom_splash.NUM_SUBLANES)) % custom_splash.NUM_SUBLANES + pad_d = max(0, query.shape[-1] - k_mean.shape[-1]) + if pad_h > 0 or pad_d > 0: + k_mean = jnp.pad(k_mean, ((0, 0), (0, pad_h), (0, pad_d))) + # K/V are passed with NO sequence padding. The fixed-m kernel slices the + # ragged KV tail (`last_compute_body_fixed` in custom_splash_attention.py) + # using slice lengths derived from the unpadded `orig_kv_seq_len`, so it + # never reads a padded K/V row; materialising the pad cost 2 x 185MB of HBM + # traffic per layer for nothing. Passing flash_block_size=1 makes only the + # sequence pad a no-op -- the head_dim->128 pad, the reshape and the + # returned (tensor, kv_size, seq_len) contract are all preserved, so + # configs with head_dim < 128 still behave exactly as before. + key, _, key_seq_len = _pad_data_for_flash(raw_key, heads, 1) + value, _, _ = _pad_data_for_flash(raw_value, heads, 1) mk_arr = None all_fixed = None @@ -1011,7 +1105,12 @@ def wrap_ulysses_attention(query, key, value, attention_mask): recenter=recenter, per_q_block=per_q_block, k_mean=k_mean, - value=value, + # Use the unpadded V: `all_fixed` gates the whole kernel through a + # lax.cond, so anything feeding it sits on the critical path. Reading + # the padded copy chained a 193MB pad + reduction behind the V + # all-to-all and left that collective fully exposed. The padding is + # zeros and the check is a max of squares, so this is output-invariant. + value=raw_value, ) bsizes = custom_splash._BlockSizes( @@ -1032,7 +1131,6 @@ def wrap_ulysses_attention(query, key, value, attention_mask): vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=True, uniform_fixed_m=True, - fixed_m_recenter=recenter, ) splash_kernel_hybrid = custom_splash.make_splash_mha( block_sizes=bsizes, @@ -1044,7 +1142,6 @@ def wrap_ulysses_attention(query, key, value, attention_mask): vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=True, uniform_fixed_m=False, - fixed_m_recenter=recenter, ) def _run_uniform(q, k, v, m, km): @@ -1053,7 +1150,16 @@ def _run_uniform(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 = 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, @@ -1066,9 +1172,18 @@ def _run_hybrid(q, k, v, m, km): use_fixed_m=False, ) vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0)) - 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) + attention_output = vmapped_splash(query, key, value) + attention_output = attention_output[:, :, :kv_size, :context_q_seq_len].astype(query.dtype) + # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. + # Sequence axis is at index 3, heads axis is at index 1. + attention_output = jax.lax.all_to_all( + attention_output, + axis_name=axis_name, + split_axis=3, + concat_axis=1, + tiled=True, + ) + return attention_output else: # Run the same local splash kernel as standard TPU flash attention, but now # on full-sequence / fewer-heads tensors produced by the all-to-all above. @@ -1108,15 +1223,15 @@ def _run_hybrid(q, k, v, m, km): attention_output = vmapped_splash(query, key, value, segment_ids) attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) - # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. - attention_output = jax.lax.all_to_all( - attention_output, - axis_name=axis_name, - split_axis=2, - concat_axis=1, - tiled=True, - ) - return attention_output + # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. + attention_output = jax.lax.all_to_all( + attention_output, + axis_name=axis_name, + split_axis=2, + concat_axis=1, + tiled=True, + ) + return attention_output devices_in_batch_sharding = mesh.shape["data"] * (mesh.shape["fsdp"] if "fsdp" in mesh.shape else 1) if not (query.shape[0] / devices_in_batch_sharding).is_integer(): @@ -1141,7 +1256,11 @@ def _run_hybrid(q, k, v, m, km): # Folding batch into heads destroys the one-mask-per-example association. # Keep the optimization for the common unmasked path only. fold_batch = ( - attention_mask is None and batch > 1 and devices_in_batch_sharding == 1 and (batch * num_heads) % num_shards == 0 + attention_mask is None + and batch > 1 + and devices_in_batch_sharding == 1 + and num_q_heads == num_kv_heads + and (batch * num_heads) % num_shards == 0 ) if fold_batch: query = query.reshape(1, batch * num_heads, *query.shape[2:]) @@ -1151,12 +1270,18 @@ def _run_hybrid(q, k, v, m, km): else: effective_num_heads = num_heads + out_q_axis_names = ( + jax.sharding.PartitionSpec(q_axis_names[0], q_axis_names[1], q_axis_names[3], q_axis_names[2]) + if use_custom_kernel + else q_axis_names + ) + if attention_mask is None: sharded_ulysses_attention = jax.shard_map( lambda q, k, v: wrap_ulysses_attention(q, k, v, None), mesh=mesh, in_specs=(q_axis_names, kv_axis_names, kv_axis_names), - out_specs=q_axis_names, + out_specs=out_q_axis_names, check_vma=False, ) @@ -1168,7 +1293,7 @@ def run_ulysses_attention(q, k, v): wrap_ulysses_attention, mesh=mesh, in_specs=(q_axis_names, kv_axis_names, kv_axis_names, mask_axis_names), - out_specs=q_axis_names, + out_specs=out_q_axis_names, check_vma=False, ) @@ -1185,10 +1310,19 @@ def run_ulysses_attention(q, k, v): run_ulysses_attention, ) - if fold_batch: - x = x.reshape(batch, num_heads, *x.shape[2:]) - x = x[:, :, :orig_q_seq_len, :] - x = _reshape_heads_to_head_dim(x) + if use_custom_kernel: + if fold_batch: + x = x.reshape(batch, num_heads, *x.shape[2:]) + x = x[:, :, :, :orig_q_seq_len] + b, h, d, s = x.shape + x = jnp.transpose(x, (0, 3, 1, 2)).reshape(b, -1, h * d) + axis_names = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD)) + x = jax.lax.with_sharding_constraint(x, axis_names) + else: + if fold_batch: + x = x.reshape(batch, num_heads, *x.shape[2:]) + x = x[:, :, :orig_q_seq_len, :] + x = _reshape_heads_to_head_dim(x) return x @@ -1213,6 +1347,7 @@ def _ulysses_ring_attention( ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, preserve_asymmetric_block_sizes: bool = False, + kv_heads: int | None = None, ) -> jax.Array: """2D context-parallel attention using a private Ulysses x ring mesh. @@ -1221,6 +1356,8 @@ def _ulysses_ring_attention( Ulysses all-to-all over the hidden Ulysses axis, and rotates K/V over the hidden ring axis. """ + if kv_heads is None: + kv_heads = heads context_axis = CONTEXT if context_axis not in mesh.shape: @@ -1237,10 +1374,22 @@ def _ulysses_ring_attention( ) if heads % num_ulysses_shards != 0: raise ValueError( - "Ulysses ring attention requires the number of heads to be divisible by the requested Ulysses shard count, " + "Ulysses ring attention requires the number of query heads to be divisible by the requested Ulysses shard count, " f"got heads={heads} and ulysses_shards={num_ulysses_shards}." ) + if kv_heads % num_ulysses_shards != 0: + raise ValueError( + "Ulysses ring attention requires the number of KV heads to be divisible by the requested Ulysses shard count, " + f"got kv_heads={kv_heads} and ulysses_shards={num_ulysses_shards}." + ) num_ring_shards = num_context_shards // num_ulysses_shards + _warn_if_ring_is_degenerate( + num_ring_shards, + num_ulysses_shards, + num_context_shards, + heads=heads, + kv_heads=kv_heads, + ) internal_mesh = _create_internal_ulysses_ring_mesh( mesh, ring_shards=num_ring_shards, @@ -1252,8 +1401,8 @@ def _ulysses_ring_attention( num_sequence_shards = num_context_shards query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_sequence_shards) - key, _ = _reshape_data_for_flash(key, heads, num_sequence_shards) - value, _ = _reshape_data_for_flash(value, heads, num_sequence_shards) + key, _ = _reshape_data_for_flash(key, kv_heads, num_sequence_shards) + value, _ = _reshape_data_for_flash(value, kv_heads, num_sequence_shards) attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) num_heads = query.shape[1] @@ -1296,8 +1445,8 @@ def wrap_ulysses_ring_attention(query, key, value, attention_mask): block_q = max(*block_q_sizes) query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, block_q) block_kv = max(*block_kv_sizes) - key, _, key_seq_len = _pad_data_for_flash(key, heads, block_kv) - value, _, _ = _pad_data_for_flash(value, heads, block_kv) + key, _, key_seq_len = _pad_data_for_flash(key, kv_heads, block_kv) + value, _, _ = _pad_data_for_flash(value, kv_heads, block_kv) q_padded_len = query.shape[2] kv_padded_len = key.shape[2] @@ -1358,7 +1507,11 @@ def wrap_ulysses_ring_attention(query, key, value, attention_mask): sharded_ulysses_ring_attention = jax.shard_map( lambda q, k, v: wrap_ulysses_ring_attention(q, k, v, None), mesh=internal_mesh, - in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names), + in_specs=( + internal_q_axis_names, + internal_kv_axis_names, + internal_kv_axis_names, + ), out_specs=internal_q_axis_names, check_vma=False, ) @@ -1370,7 +1523,12 @@ def run_ulysses_ring_attention(q, k, v): sharded_ulysses_ring_attention = jax.shard_map( wrap_ulysses_ring_attention, mesh=internal_mesh, - in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names, internal_mask_axis_names), + in_specs=( + internal_q_axis_names, + internal_kv_axis_names, + internal_kv_axis_names, + internal_mask_axis_names, + ), out_specs=internal_q_axis_names, check_vma=False, ) @@ -1395,10 +1553,236 @@ def run_ulysses_ring_attention(q, k, v): def _max_row_norm_per_head(x: jax.Array) -> jax.Array: - """Largest row L2 norm per head of a `[B, H, S, D]` activation.""" - row_sq = jnp.square(x).sum(axis=-1, dtype=jnp.float32) - # 1.01 keeps the result an upper bound despite bf16 mantissa loss. - return jnp.sqrt(row_sq.max(axis=(0, 2))) * 1.01 + """Largest row L2 norm estimate per head of a `[..., H, S, D]` activation in FP32 with conservative exponent margin, preserving batch dims.""" + row_sq = (x.astype(jnp.float32) ** 2).sum(axis=-1) + return jnp.sqrt(row_sq.max(axis=-1)) + + +def _slice_own_ulysses_heads(x: jax.Array, ulysses_axis: str, num_ulysses_shards: int, axis: int) -> jax.Array: + """Slices an all-heads array down to the heads this rank owns after the a2a. + + `all_to_all(split_axis=1, concat_axis=2, tiled=True)` hands rank `r` the head + block `[r * H/U, (r+1) * H/U)`, so the same static block size with a + rank-dependent offset recovers exactly the heads the rank now holds. + """ + heads_per_dev = x.shape[axis] // num_ulysses_shards + start = jax.lax.axis_index(ulysses_axis) * heads_per_dev + return jax.lax.dynamic_slice_in_dim(x, start, heads_per_dev, axis=axis) + + +def _ring_fixed_m_norms_pre_a2a( + query: jax.Array, + key: jax.Array, + value: jax.Array, + *, + ulysses_axis: str, + ring_axis: str, + num_ulysses_shards: int, + num_ring_shards: int, + block_q: int, + per_q_block: bool, + use_k_centering: bool = False, +): + """Computes all R>1 fixed-m norms and global eligibility predicates *pre* a2a. + + Inputs are the shard-local activations as they arrive from the QKV + projections: `[B, H_all, S/(U*R), D]` -- every head, a 1/U slice of this ring + shard's sequence. The equivalent post-a2a arrays are `[B, H_all/U, S/R, D]`: + the same elements, redistributed. Both forms therefore admit the same + reductions, but doing them here (including `v_max_sq`, `v_ok`, and + `all_fixed_global`) is materially cheaper: + + * Every reduction reads the projection's natural output layout. After the + all-to-all the arrays carry the collective's layout, and XLA inserts + relayout copies to feed post-a2a reductions. + * All reductions and cross-chip `pmin`/`pmax` collectives become + independent of `all_to_all(query, key, value)`, allowing XLA's + latency-hiding scheduler to overlap them with the all-to-all instead of + serialising `a2a -> post-a2a V reduce -> ring pmin -> ring pmin -> lax.cond`. + + Note: callers must apply `jax.lax.optimization_barrier((query, key, value))` + in the outer scope so both this function and the subsequent `all_to_all` + consume the exact same barriered tensors. + + Returns `(key_out, qn_dev, mk_all_sq_dev, v_ok, all_fixed_global)` sliced + to the heads this Ulysses rank owns and ready for immediate `jax.lax.cond` + dispatch post-a2a. + """ + reduce_axes = (ulysses_axis, ring_axis) + key_f32 = key.astype(jnp.float32) + + q_norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) + qn_head_local = q_norm_sq.max(axis=-1) + vn_local = (value.astype(jnp.float32) ** 2).max(axis=(2, 3)) + + if use_k_centering: + # Optional K-centering: computes global mean and subtracts before a2a. + k_mean_all = jax.lax.pmean(jnp.mean(key_f32, axis=2), axis_name=reduce_axes) + centered_f32 = key_f32 - k_mean_all[:, :, None, :] + key_out = centered_f32.astype(key.dtype) + kn_local = jnp.sum(centered_f32**2, axis=-1).max(axis=-1) + else: + # High-performance uncentered path: key is completely untouched, so all-to-all + # starts immediately in parallel with norm reductions, eliminating the pmean + # collective, 194MB/layer HBM subtraction, and collective serialization. + key_out = key + kn_local = jnp.sum(key_f32**2, axis=-1).max(axis=-1) + + # Global Q/V/K max norms in a SINGLE (ulysses, ring) pmax. + # + # This used to be two collectives -- a pmax of (Q, V) over (ulysses, ring), + # then `all_gather(pmax(kn_local, ulysses), ring)` to give every hop its own + # K shard's norm. The gather is what made this expensive, and not for the + # reason one would guess: its payload is only `heads` floats. Profiling shows + # the collective time is unchanged by removing it; what changes is *compute*. + # On TPU the ring-axis gather forces a relayout of the small norm arrays that + # sits in the middle of the QKV projection's fusion region, and XLA then fails + # to fuse across it. Removing it recovered ~2.5 s per 40-step denoise + # (convolution fusion -1.64 s, loop fusion -0.44 s, data formatting -0.36 s). + qn_head_global, vn_global, kn_global = jax.lax.pmax((qn_head_local, vn_local, kn_local), axis_name=reduce_axes) + + if not per_q_block: + qn_all = qn_head_global + else: + batch, num_q_heads, local_seq = q_norm_sq.shape + post_a2a_seq = local_seq * num_ulysses_shards + num_q_blocks = -(-post_a2a_seq // block_q) + padded_seq = num_q_blocks * block_q + scattered = jnp.zeros((batch, num_q_heads, padded_seq), q_norm_sq.dtype) + seq_offset = jax.lax.axis_index(ulysses_axis) * local_seq + scattered = jax.lax.dynamic_update_slice_in_dim(scattered, q_norm_sq, seq_offset, axis=2) + qn_local = scattered.reshape(batch, num_q_heads, num_q_blocks, block_q).max(axis=-1) + qn_all = jax.lax.pmax(qn_local, ulysses_axis) + + # `mk_all_sq` keeps its (batch, ring, heads) shape so the kernel's per-hop + # interface is unchanged, but every hop now carries the ring-wide max rather + # than that hop's own shard norm. Using the max over hops can only enlarge the + # Cauchy-Schwarz bound, so fixed-m eligibility becomes *more* conservative, + # never less -- and it is what `mk_global_sq` below already collapsed it to, + # so `all_fixed_global` is bit-identical to the per-hop version. + mk_all_sq = jnp.broadcast_to(kn_global[:, None, :], (kn_global.shape[0], num_ring_shards, kn_global.shape[1])) + + # Slice down to the heads owned by this Ulysses rank post-a2a. + qn_dev = _slice_own_ulysses_heads(qn_all, ulysses_axis, num_ulysses_shards, axis=1) + qn_head_global_dev = _slice_own_ulysses_heads(qn_head_global, ulysses_axis, num_ulysses_shards, axis=1) + mk_all_sq_dev = _slice_own_ulysses_heads(mk_all_sq, ulysses_axis, num_ulysses_shards, axis=2) + vn_dev = _slice_own_ulysses_heads(vn_global, ulysses_axis, num_ulysses_shards, axis=1) + + num_q_heads_dev = qn_dev.shape[1] + num_kv_heads_dev = mk_all_sq_dev.shape[2] + if num_q_heads_dev != num_kv_heads_dev: + if num_q_heads_dev % num_kv_heads_dev != 0: + raise ValueError( + f"num_q_heads ({num_q_heads_dev}) must be divisible by num_kv_heads ({num_kv_heads_dev}) for GQA ring fixed-m." + ) + q_heads_per_kv_head = num_q_heads_dev // num_kv_heads_dev + mk_all_sq_dev = jnp.repeat(mk_all_sq_dev, q_heads_per_kv_head, axis=2) + + # Evaluate global V safety and Cauchy-Schwarz fixed-m eligibility pre-a2a. + # Because qn_head_global_dev, vn_dev, and mk_global_sq are already reduced over + # ring_axis, both v_ok and all_fixed_global are bit-identical across all ring + # ranks with zero post-all_gather collectives. + effective_kv_seq_len = key.shape[2] * num_ulysses_shards * num_ring_shards + global_recenter, global_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=True) + global_bound_sq = global_bound**2 + dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, global_recenter) + + v_max_sq = vn_dev.max() + v_ok = (v_max_sq <= (custom_splash.DEFAULT_MAX_V_BOUND**2)) & dtype_safe + mk_global_sq = mk_all_sq_dev.max(axis=1) + + bound_head_sq = qn_head_global_dev * mk_global_sq + all_fixed_global = jnp.all(bound_head_sq <= global_bound_sq) & v_ok + return key_out, qn_dev, mk_all_sq_dev, v_ok, all_fixed_global + + +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. + + Args: + query: Padded query activation, shape `(batch, local_heads, padded_q_len, head_dim)`. + key: Key activation (raw unpadded or padded), shape `(batch, local_heads, kv_len, head_dim)`. + block_q: Query tile block size. + safe_bound: Maximum safe norm product threshold before falling back to online softmax. + recenter: Fixed-m dynamic recenter constant C(N). + per_q_block: If True, evaluates gating independently per query tile. If False, + evaluates monolithic gating per head. + k_mean: Optional mean key vector for Virtual K-centering, shape `(batch, local_heads, head_dim)`. + value: Optional value activation, shape `(batch, local_heads, kv_len, head_dim_v)`, used to + verify that |V| <= v_max_bound to guarantee against FP32 overflow. + v_max_bound: Maximum safe value magnitude (default 256.0). + + Returns: + mk_arr: Gating metadata array of shape `(batch, 2, local_heads, num_q_blocks)` + multiplexing precomputed block base shifts and binary fixed-m gating predicates into a single + Pallas scalar prefetch memory slot: + - `mk_arr[:, 0, h, i]`: Precomputed block base shift m_B = ceil(max_i ||q_i|| * max_j ||k_j||) - C. + - `mk_arr[:, 1, h, i]`: Discrete eligibility predicate (1.0 for fixed-m, 0.0 for online). + all_fixed: Boolean scalar indicating if all elements are eligible for uniform fixed-m. + """ + batch_size, num_q_heads, q_len, _ = query.shape + num_kv_heads = key.shape[1] + if safe_bound is None or recenter is None: + rec, bnd = custom_splash.get_fixed_m_constants(key.shape[2], is_ring=False, v_max_bound=v_max_bound) + if safe_bound is None: + safe_bound = bnd + if recenter is None: + recenter = rec + safe_bound_sq = safe_bound**2 + if k_mean is not None: + centered_k = key.astype(jnp.float32) - k_mean[:, :num_kv_heads, 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) + + # Fixed-m weights reach 2**recenter before being narrowed to the activation + # dtype for the S@V matmul. If that dtype's exponent range cannot hold them + # (fp16, fp8), the FP32 bound analysis is irrelevant -- the narrowing itself + # overflows to inf -- so disqualify every head up front. + dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, recenter) + v_ok = 1.0 if dtype_safe else 0.0 + if dtype_safe and value is not None: + v_max_sq = (value.astype(jnp.float32) ** 2).max() + v_ok = (v_max_sq <= (v_max_bound**2)).astype(jnp.float32) + + num_q_blocks = q_len // block_q + if per_q_block: + norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) # (batch, num_q_heads, q_len) + qn_max_sq = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, block_q).max( + axis=-1 + ) # (batch, num_q_heads, num_q_blocks) + bound_sq = qn_max_sq * mk_h_sq[:, :, None] + fixed_ok = (bound_sq <= safe_bound_sq).astype(jnp.float32) * v_ok + m_base = jnp.ceil(jnp.sqrt(bound_sq)) - recenter + mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, num_q_heads, num_q_blocks) + all_fixed = jnp.all(fixed_ok > 0.5) + else: + qn_max_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_q_heads) + bound_sq_1d = qn_max_sq * mk_h_sq + fixed_ok_1d = (bound_sq_1d <= safe_bound_sq).astype(jnp.float32) * v_ok + m_base_1d = jnp.ceil(jnp.sqrt(bound_sq_1d)) - recenter + m_base_expanded = jnp.broadcast_to(m_base_1d[:, :, None], (batch_size, num_q_heads, num_q_blocks)) + fixed_ok_expanded = jnp.broadcast_to(fixed_ok_1d[:, :, None], (batch_size, num_q_heads, num_q_blocks)) + mk_arr = jnp.stack([m_base_expanded, fixed_ok_expanded], axis=1) # (batch, 2, num_q_heads, num_q_blocks) + all_fixed = jnp.all(fixed_ok_1d > 0.5) + + return mk_arr, all_fixed def _ulysses_ring_custom_attention( @@ -1420,28 +1804,16 @@ def _ulysses_ring_custom_attention( bidirectional: bool = False, use_fixed_m: bool = False, ulysses_attention_chunks: int = 1, + per_q_block: bool = True, + kv_heads: int | None = None, + use_k_centering: bool = False, ) -> jax.Array: - """Hybrid Ulysses + Ring (USP) with the CUSTOM splash kernel on main's mesh. + """2D USP attention (Ulysses + Ring) using custom splash kernel with exact Fixed-m support.""" + if kv_heads is None: + kv_heads = heads - Uses origin/main's explicit internal `(ring, ulysses)` mesh - (`_create_internal_ulysses_ring_mesh`, commit c104db51) instead of single-axis - collective sub-groups: the public `context` axis is reshaped with the Ulysses - axis innermost, so the Ulysses all-to-all stays INTRA-chip and the ring rotates - ACROSS chips. The per-shard attention is our custom splash kernel - (`make_custom_ring_attention`), not the tokamax_ring kernel main uses. - - 1. all-to-all over the (intra-chip) Ulysses axis: trade sequence for heads; - 2. ring (full ppermute) over the (cross-chip) ring axis, online-softmax merge; - 3. all-to-all back to restore the sequence-sharded / full-heads layout. - - U = ulysses_shards (from config); R = context // U. U=context -> pure - Ulysses, U=1 -> pure Ring (all on the same custom kernel). - """ if attention_mask is not None: - raise NotImplementedError( - "ulysses_ring_custom does not support attention_mask (the custom splash kernels only " - "handle padding via orig_seq_len); got a non-None attention_mask." - ) + raise NotImplementedError("ulysses_ring_custom does not support attention_mask.") axis_name = "context" num_context_shards = mesh.shape[axis_name] num_ulysses_shards = ulysses_shards @@ -1454,13 +1826,34 @@ def _ulysses_ring_custom_attention( ) num_ring_shards = num_context_shards // num_ulysses_shards + # Virtual K-centering is switched ON for the ring path here. The kernel only + # centers when it is handed a `k_mean` (it stays uncentered otherwise, which is + # what the preceding change ships), so enabling it means supplying both halves + # of the matched pair: the global mean itself, and a Cauchy-Schwarz bound + # derived from the centered keys rather than the raw ones (see kf_centered + # below). Centering shrinks the bound, which is what lets more heads stay on + # the fixed-m path at R > 1. + _warn_if_ring_is_degenerate( + num_ring_shards, + num_ulysses_shards, + num_context_shards, + heads=heads, + kv_heads=kv_heads, + ) + query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards) - key, _ = _reshape_data_for_flash(key, heads, num_context_shards) - value, _ = _reshape_data_for_flash(value, heads, num_context_shards) + key, orig_kv_seq_len = _reshape_data_for_flash(key, kv_heads, num_context_shards) + value, _ = _reshape_data_for_flash(value, kv_heads, num_context_shards) num_heads = query.shape[1] if num_heads % num_ulysses_shards != 0: - raise ValueError(f"Ulysses+Ring requires heads divisible by U={num_ulysses_shards}, got heads={num_heads}.") - + raise ValueError(f"Ulysses+Ring requires query heads divisible by U={num_ulysses_shards}, got heads={num_heads}.") + if kv_heads % num_ulysses_shards != 0: + raise ValueError(f"Ulysses+Ring requires KV heads divisible by U={num_ulysses_shards}, got kv_heads={kv_heads}.") + if num_ring_shards > 1 and orig_kv_seq_len % num_context_shards != 0: + raise ValueError( + f"2D Ulysses+Ring attention requires sequence length to be divisible by context_shards={num_context_shards}, " + f"got orig_kv_seq_len={orig_kv_seq_len}." + ) ( bq, bkv, @@ -1469,13 +1862,10 @@ def _ulysses_ring_custom_attention( heads_per_tile, vmem_limit_bytes, ) = _extract_custom_block_sizes(flash_block_sizes) - if heads_per_tile > 1: - raise NotImplementedError("ulysses_ring_custom currently supports heads_per_tile == 1 only.") - + if heads_per_tile > 1 and num_ring_shards > 1: + raise NotImplementedError("heads_per_tile > 1 is not supported for multi-shard ring attention.") internal_mesh = _create_internal_ulysses_ring_mesh(mesh, num_ring_shards, num_ulysses_shards) - ring_axis = INTERNAL_RING_AXIS - ulysses_axis = INTERNAL_ULYSSES_AXIS - + ring_axis, ulysses_axis = INTERNAL_RING_AXIS, INTERNAL_ULYSSES_AXIS q_axis_names = nn.logical_to_mesh_axes(axis_names_q) kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv) internal_q_axis_names = _replace_mesh_axis_names(q_axis_names, axis_name, (ring_axis, ulysses_axis)) @@ -1493,179 +1883,182 @@ def _ulysses_ring_custom_attention( check_vma=False, ) def wrap_ulysses_ring_attention(query, key, value): - fixed_m_norms = None + # Apply the base-2 rescale of Q *before* the all-to-all. A scalar elementwise + # multiply commutes exactly with the collective (which is pure data movement), + # so this is bit-identical. Done after the a2a it sat between the collective + # and the kernel and XLA wrapped it in relayout copies; done before, it fuses + # into the producer of Q and its 185MB round-trip disappears. + if use_base2_exp: + query = query * LOG2E + + # (0) R>1 fixed-m reductions and global eligibility predicates, computed + # entirely on the pre-a2a layout so zero reductions or collectives sit + # between `all_to_all` and `jax.lax.cond`. + qn_dev, mk_all_sq, v_ok, all_fixed_global = None, None, None, None if use_fixed_m and num_ring_shards > 1: - # Fixed-m's Cauchy-Schwarz inputs, reduced on the PRE-a2a activation so - # the reduction overlaps the all_to_all instead of stalling the first - # ring step (taking them after the a2a measured +8% end to end). - # - # The barrier is load-bearing: the norms are a second consumer of these - # activations, and without it XLA duplicates the producer chain into the - # norm fusion instead of materializing once -- worth 1.46 ms/layer, the - # difference between fixed-m breaking even and winning. - # - # Reducing them further upstream (on the flat [B, S, H*D] form, where - # head_dim is contiguous) is exact and looks cheaper, but there the array - # 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) * (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)) - mk_all = jax.lax.pmax(kn_local, ulysses_axis) - heads_per_dev = qn_all.shape[0] // num_ulysses_shards - start_head = jax.lax.axis_index(ulysses_axis) * heads_per_dev - fixed_m_norms = ( - jax.lax.dynamic_slice_in_dim(qn_all, start_head, heads_per_dev), - jax.lax.dynamic_slice_in_dim(mk_all, start_head, heads_per_dev), + query, key, value = jax.lax.optimization_barrier((query, key, value)) + ( + key, + qn_dev, + mk_all_sq, + v_ok, + all_fixed_global, + ) = _ring_fixed_m_norms_pre_a2a( + query, + key, + value, + ulysses_axis=ulysses_axis, + ring_axis=ring_axis, + num_ulysses_shards=num_ulysses_shards, + num_ring_shards=num_ring_shards, + block_q=bq, + per_q_block=per_q_block, + use_k_centering=use_k_centering, ) - # (1) Ulysses all-to-all over the (intra-chip) ulysses axis: heads -> sequence, - # so each device holds the full ring-chunk sequence with heads/U heads. + # (1) Ulysses All-to-All: heads -> sequence a2a = functools.partial(jax.lax.all_to_all, axis_name=ulysses_axis, tiled=True) query = a2a(query, split_axis=1, concat_axis=2) key = a2a(key, split_axis=1, concat_axis=2) value = a2a(value, split_axis=1, concat_axis=2) - if use_base2_exp: - query = query * LOG2E - - k_mean = None - if use_fixed_m and num_ring_shards == 1: - 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) + # NOTE: the base-2 rescale of Q is applied before the all-to-all above. + 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 if num_ring_shards == 1 else raw_key.shape[2] + + if use_fixed_m and num_ring_shards == 1 and use_k_centering: + # Optional K-centering: Center key directly in JAX so Pallas kernel runs with pristine 4 operands + kbar = jnp.mean( + raw_key[:, :, :actual_kv_seq_len, :].astype(jnp.float32), + axis=2, + keepdims=True, + ) + raw_key = (raw_key.astype(jnp.float32) - kbar).astype(raw_key.dtype) + real_key = raw_key[:, :, :actual_kv_seq_len, :] + else: + real_key = raw_key[:, :, :actual_kv_seq_len, :] - v_ok = None - if use_fixed_m and num_ring_shards > 1: - # V-magnitude and dtype safety are properties of the WHOLE distributed - # problem, not of any single hop, and unlike the Cauchy-Schwarz norm - # bounds they are not re-derivable from a hop's Q/K. The kernel therefore - # cannot reconstruct this verdict, and omitting it let fixed-m run on - # inputs it cannot represent -- float16 with Q=K=0 and V=1 overflows to - # inf instead of returning 1.0. - effective_kv_seq_len = key_seq_len * num_ring_shards - global_recenter, _ = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=False) - dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, global_recenter) - # The sequence pad is zeros, and a max of squares is unchanged by zeros, - # so reading the padded V here is exact. - v_max_sq = (value.astype(jnp.float32) ** 2).max() - v_ok_local = (v_max_sq <= (custom_splash.DEFAULT_MAX_V_BOUND**2)) & dtype_safe - # Reduce over BOTH internal axes: after the all-to-all each device holds a - # head slice of one ring chunk, so neither axis alone sees the whole V. - # The fixed-m branch must also be taken uniformly by every participant of - # the ppermute, which a pmin over both axes guarantees. - v_ok = jax.lax.pmin(v_ok_local, (ring_axis, ulysses_axis)) - - mk_arr = None - all_fixed = None + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) + # When num_ring_shards == 1 or actual_kv_seq_len is aligned to 8 sublanes, K/V + # are passed with NO sequence padding. The kernel slices the ragged KV tail + # using slice lengths derived from actual_kv_seq_len, avoiding sequence pad + # HBM copies and redundant ppermute/MXU compute on padded keys. + kv_pad_size = 1 if (num_ring_shards == 1 or actual_kv_seq_len % 8 == 0) else bkv + key, _, key_seq_len = _pad_data_for_flash(raw_key, kv_heads, kv_pad_size) + value, _, _ = _pad_data_for_flash(raw_value, kv_heads, kv_pad_size) + ring_kv_seq_len = actual_kv_seq_len if actual_kv_seq_len % 8 == 0 else key_seq_len + + mk_arr, all_fixed = None, None if use_fixed_m and num_ring_shards == 1: - recenter, safe_bound = custom_splash.get_fixed_m_constants(key_seq_len, is_ring=False) + recenter, safe_bound = custom_splash.get_fixed_m_constants(actual_kv_seq_len, is_ring=False) mk_arr, all_fixed = _compute_fixed_m_metadata( query, - key[:, :, :key_seq_len, :], - block_q=bq, + real_key, + bq, safe_bound=safe_bound, recenter=recenter, - per_q_block=False, - k_mean=k_mean, - value=value, + per_q_block=per_q_block, + k_mean=None, + value=raw_value, ) - bsizes = custom_splash._BlockSizes( - block_q=bq, - block_kv=bkv, - block_kv_compute=bkv_compute, - block_kv_compute_in=bkv_compute_in, - ) + bsizes = custom_splash._BlockSizes(bq, bkv, bkv_compute, bkv_compute_in) + + # (2a) R=1: Dedicated single-device splash kernel with fixed-m or online softmax if num_ring_shards == 1: - # (2a) R=1: the ring is trivial (no rotation) -> use the lighter dedicated - # 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]. if use_fixed_m: 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, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_len, heads_per_tile=heads_per_tile, use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=True, uniform_fixed_m=True, - fixed_m_recenter=recenter, ) splash_kernel_hybrid = custom_splash.make_splash_mha( block_sizes=bsizes, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_len, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_len, heads_per_tile=heads_per_tile, use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=True, uniform_fixed_m=False, - fixed_m_recenter=recenter, ) - def _run_uniform(q, k, v, m, km): - return jax.vmap(splash_kernel_uniform, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + def _run_uniform(q, k, v, m): + return jax.vmap(splash_kernel_uniform, in_axes=(0, 0, 0, 0))(q, k, v, m) - 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) + def _run_hybrid(q, k, v, m): + return jax.vmap(splash_kernel_hybrid, in_axes=(0, 0, 0, 0))(q, k, v, m) - 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) + raw_out = jax.lax.cond(all_fixed, _run_uniform, _run_hybrid, query, key, value, mk_arr) 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, + 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, ) - attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value), 2, 3) + raw_out = jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value) + attention_output = jnp.swapaxes(raw_out, 2, 3) + + # (2b) Ring: Cross-chip ppermute schedule with custom ring kernel else: - # (2b) Ring (full ppermute over the cross-chip ring axis) with the custom kernel. - # bidirectional=True -> wrap-free schedule (streams K/V both directions one hop - # at a time), for a non-wrapping ring axis. Selected by attention=ulysses_ring_custom_bidir. - ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( - block_sizes=bsizes, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_len, - use_base2_exp=use_base2_exp, - use_experimental_scheduler=use_experimental_scheduler, - vmem_limit_bytes=vmem_limit_bytes, - ring_axis=ring_axis, - ring_size=num_ring_shards, - 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) + if use_fixed_m: + ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( + block_sizes=bsizes, + orig_q_seq_len=query_seq_len, + orig_kv_seq_len=ring_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + ring_axis=ring_axis, + ring_size=num_ring_shards, + bidirectional=bidirectional, + use_fixed_m=True, + per_q_block=per_q_block, + pregathered_mk=True, + v_ok=v_ok, + all_fixed_global=all_fixed_global, + ) + attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0, (0, 0)))(query, key, value, (qn_dev, mk_all_sq)) + else: + ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( + block_sizes=bsizes, + orig_q_seq_len=query_seq_len, + orig_kv_seq_len=ring_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + ring_axis=ring_axis, + ring_size=num_ring_shards, + bidirectional=bidirectional, + use_fixed_m=False, + ) + attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0))(query, key, value) - # (3) Ulysses all-to-all back: sequence -> heads, restoring the layout. - attention_output = a2a(attention_output, split_axis=2, concat_axis=1) - return attention_output + attention_output = attention_output[:, :, :context_q_seq_len, :kv_size].astype(query.dtype) + + # (3) Ulysses All-to-All back: sequence -> heads + return a2a(attention_output, split_axis=2, concat_axis=1) x = _run_chunked_ulysses_attention( query, key, value, - num_heads, + heads, num_ulysses_shards, ulysses_attention_chunks, wrap_ulysses_ring_attention, @@ -1688,17 +2081,48 @@ def _apply_attention_dot( float32_qk_product: bool, use_memory_efficient_attention: bool, attention_mask: Array = None, + kv_heads: int | None = None, ): """Apply Attention.""" + effective_kv_heads = kv_heads if kv_heads is not None else heads if split_head_dim: - b = key.shape[0] - query_states = jnp.reshape(query, (b, -1, heads, dim_head)) - key_states = jnp.reshape(key, (b, -1, heads, dim_head)) - value_states = jnp.reshape(value, (b, -1, heads, dim_head)) + + def _to_bshd(x: Array, n_heads: int) -> Array: + """Normalise to [B, S, H, D]. + + Callers that apply rotary embeddings hand us [B, H, S, D] (see + `_unflatten_heads`), while the flat path supplies [B, S, H*D]. Only the + latter can be reshaped into [B, S, H, D]; reinterpreting [B, H, S, D] + that way keeps the shape legal but interleaves heads with tokens, so it + corrupts the output silently. Transpose the 4-D case instead. + """ + if x.ndim == 4: + return jnp.swapaxes(x, 1, 2) + return jnp.reshape(x, (x.shape[0], -1, n_heads, dim_head)) + + query_states = _to_bshd(query, heads) + key_states = _to_bshd(key, effective_kv_heads) + value_states = _to_bshd(value, effective_kv_heads) + if heads != effective_kv_heads: + num_repeats = heads // effective_kv_heads + key_states = jnp.repeat(key_states, num_repeats, axis=2) + value_states = jnp.repeat(value_states, num_repeats, axis=2) else: query_states = _reshape_heads_to_batch_dim(query, heads) - key_states = _reshape_heads_to_batch_dim(key, heads) - value_states = _reshape_heads_to_batch_dim(value, heads) + key_states = _reshape_heads_to_batch_dim(key, effective_kv_heads) + value_states = _reshape_heads_to_batch_dim(value, effective_kv_heads) + if heads != effective_kv_heads: + num_repeats = heads // effective_kv_heads + b = query.shape[0] + s_k = key_states.shape[1] + key_states = jnp.repeat(key_states.reshape(b, effective_kv_heads, s_k, -1), num_repeats, axis=1).reshape( + b * heads, s_k, -1 + ) + value_states = jnp.repeat( + value_states.reshape(b, effective_kv_heads, s_k, -1), + num_repeats, + axis=1, + ).reshape(b * heads, s_k, -1) if float32_qk_product: query_states = query_states.astype(jnp.float32) @@ -1804,6 +2228,7 @@ def dot_product_kernel(q, k, v, context): context["float32_qk_product"], context["use_memory_efficient_attention"], context["attention_mask"], + kv_heads=context.get("kv_heads", None), ) @@ -1826,6 +2251,9 @@ def ulysses_custom_kernel(q, k, v, context): use_base2_exp=context.get("use_base2_exp", True), use_experimental_scheduler=context.get("use_experimental_scheduler", False), ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), + ulysses_shards=context.get("ulysses_shards", -1), + kernel_name="ulysses_custom", ) @@ -1848,16 +2276,40 @@ def ulysses_ring_custom_kernel(q, k, v, context): use_base2_exp=context.get("use_base2_exp", True), use_experimental_scheduler=context.get("use_experimental_scheduler", False), ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), ) @register_kernel("ulysses_ring_custom_fixed_m") def ulysses_ring_custom_fixed_m_kernel(q, k, v, context): - """fixed-m variant of ulysses_ring_custom: the per-shard custom splash kernel - uses the Cauchy-Schwarz fixed-m softmax bound (no in-kernel running-max - rescale). max||k|| and the K-smoothing mean are taken LOCALLY per ring shard - (no per-layer ring collective); the outer ring online-softmax merge still - re-normalizes across shards, so per-shard bounds stay correct.""" + """fixed-m variant of ulysses_ring_custom with monolithic per-head gating.""" + return _ulysses_ring_custom_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"], + ulysses_shards=context["ulysses_shards"], + 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), + kv_heads=context.get("kv_heads", None), + use_k_centering=context.get("use_k_centering", False), + ) + + +@register_kernel("ulysses_ring_custom_fixed_m_per_q_block") +def ulysses_ring_custom_fixed_m_per_q_block_kernel(q, k, v, context): + """fixed-m variant of ulysses_ring_custom with per-Q-block gating.""" return _ulysses_ring_custom_attention( q, k * context["scale"], @@ -1875,7 +2327,10 @@ def ulysses_ring_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=True, ulysses_attention_chunks=context.get("ulysses_attention_chunks", 1), + kv_heads=context.get("kv_heads", None), + use_k_centering=context.get("use_k_centering", False), ) @@ -1902,6 +2357,7 @@ def ulysses_ring_custom_bidir_kernel(q, k, v, context): use_experimental_scheduler=context.get("use_experimental_scheduler", False), bidirectional=True, ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), ) @@ -1925,7 +2381,10 @@ def ulysses_custom_fixed_m_kernel(q, k, v, context): 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), + ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), + ulysses_shards=context.get("ulysses_shards", -1), + kernel_name="ulysses_custom_fixed_m", ) @@ -1949,7 +2408,10 @@ def ulysses_custom_fixed_m_per_q_block_kernel(q, k, v, context): 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), + ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), + ulysses_shards=context.get("ulysses_shards", -1), + kernel_name="ulysses_custom_fixed_m_per_q_block", ) @@ -1970,6 +2432,9 @@ def ulysses_kernel(q, k, v, context): attention_mask=context["attention_mask"], ulysses_attention_chunks=context["ulysses_attention_chunks"], preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), + kv_heads=context.get("kv_heads", None), + ulysses_shards=context.get("ulysses_shards", -1), + kernel_name="ulysses", ) @@ -1993,6 +2458,7 @@ def ulysses_ring_kernel(q, k, v, context): ulysses_shards=context["ulysses_shards"], ulysses_attention_chunks=context["ulysses_attention_chunks"], preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), + kv_heads=context.get("kv_heads", None), ) @@ -2118,6 +2584,8 @@ def _apply_attention( ulysses_attention_chunks: int = 1, is_causal: bool = False, preserve_asymmetric_block_sizes: bool = False, + kv_heads: Optional[int] = None, + use_k_centering: bool = False, ): """Routes to different attention kernels using a module-level registry.""" @@ -2133,7 +2601,12 @@ def _apply_attention( "ulysses", "ulysses_custom", "ulysses_custom_fixed_m", + "ulysses_custom_fixed_m_per_q_block", "ulysses_ring", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_fixed_m_per_q_block", + "ulysses_ring_custom_bidir", ]: can_use_flash_attention = ( query.shape[seq_len_idx] >= flash_min_seq_length @@ -2165,6 +2638,7 @@ def _apply_attention( context = { "heads": heads, + "kv_heads": kv_heads, "mesh": mesh, "axis_names_q": axis_names_q, "axis_names_kv": axis_names_kv, @@ -2185,6 +2659,7 @@ def _apply_attention( "dpa_layer": dpa_layer, "is_causal": is_causal, "preserve_asymmetric_block_sizes": preserve_asymmetric_block_sizes, + "use_k_centering": use_k_centering, } # Module-level Registry lookup @@ -2421,12 +2896,15 @@ def __init__( use_experimental_scheduler: bool = False, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + kv_heads: Optional[int] = None, + use_k_centering: bool = False, ): self.dpa_layer = None self.use_base2_exp = use_base2_exp self.use_experimental_scheduler = use_experimental_scheduler self.ulysses_shards = ulysses_shards self.ulysses_attention_chunks = ulysses_attention_chunks + self.use_k_centering = use_k_centering if attention_kernel == "cudnn_flash_te": from transformer_engine.jax.flax.transformer import DotProductAttention # pytype: disable=import-error @@ -2451,6 +2929,7 @@ def __init__( self.mesh = mesh self.scale = scale self.heads = heads + self.kv_heads = kv_heads self.dim_head = dim_head self.attention_kernel = attention_kernel self.use_memory_efficient_attention = use_memory_efficient_attention @@ -2499,6 +2978,8 @@ def apply_attention( ulysses_shards=(self.ulysses_shards if hasattr(self, "ulysses_shards") else -1), ulysses_attention_chunks=(self.ulysses_attention_chunks if hasattr(self, "ulysses_attention_chunks") else 1), preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes, + kv_heads=self.kv_heads, + use_k_centering=self.use_k_centering if hasattr(self, "use_k_centering") else False, ) @@ -2522,6 +3003,8 @@ class AttentionOp(nn.Module): ulysses_shards: int = -1 ulysses_attention_chunks: int = 1 is_causal: bool = False + kv_heads: Optional[int] = None + use_k_centering: bool = False def setup(self): self.dpa_layer = None @@ -2580,6 +3063,8 @@ def apply_attention( ulysses_attention_chunks=self.ulysses_attention_chunks, is_causal=self.is_causal, preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes, + kv_heads=self.kv_heads, + use_k_centering=self.use_k_centering if hasattr(self, "use_k_centering") else False, ) @@ -2623,6 +3108,7 @@ def __init__( "use_experimental_scheduler": False, "ulysses_shards": -1, "ulysses_attention_chunks": 1, + "use_k_centering": False, **(attention_config or {}), } @@ -2638,6 +3124,8 @@ def __init__( self.value_axis_names = value_axis_names self.out_axis_names = out_axis_names self.enable_jax_named_scopes = enable_jax_named_scopes + self.is_self_attention = is_self_attention + self.eps = eps cross_attention_remapped_to_flash = not is_self_attention and attention_kernel in ( "tokamax_ring", @@ -2645,9 +3133,11 @@ def __init__( "ulysses_ring", "ulysses_ring_custom", "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_fixed_m_per_q_block", "ulysses_ring_custom_bidir", "ulysses_custom", "ulysses_custom_fixed_m", + "ulysses_custom_fixed_m_per_q_block", ) cross_attention_uses_local_kv = not is_self_attention and ( cross_attention_remapped_to_flash or attention_kernel in ("flash", "tokamax_flash", "cudnn_flash_te") @@ -2668,7 +3158,12 @@ def __init__( elif attention_kernel in ("tokamax_ring", "tokamax_ring_custom", "ulysses_ring") and not is_self_attention: attention_kernel = "tokamax_flash" # do not use ring attention for cross attention elif ( - attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir", "ulysses_ring_custom_fixed_m") + attention_kernel + in ( + "ulysses_ring_custom", + "ulysses_ring_custom_bidir", + "ulysses_ring_custom_fixed_m", + ) and not is_self_attention ): attention_kernel = "ulysses_custom" # plain ulysses (no ring) for cross attention @@ -2676,6 +3171,7 @@ def __init__( self.image_seq_len = image_seq_len # New for I2V tpu_type = get_tpu_type() self.alignment = 256 if tpu_type in [TpuType.TPU_V6_LITE, TpuType.TPU_7X] else 128 + self.precision = precision self.attention_op = NNXAttentionOp( mesh=mesh, @@ -2698,6 +3194,7 @@ def __init__( use_experimental_scheduler=attention_config["use_experimental_scheduler"], ulysses_shards=attention_config["ulysses_shards"], ulysses_attention_chunks=attention_config["ulysses_attention_chunks"], + use_k_centering=attention_config["use_k_centering"], ) # None axes corresponds to the stacked weights across all blocks # because of the use of nnx.vmap and nnx.scan. @@ -2850,9 +3347,9 @@ def _apply_rope(self, xq: jax.Array, xk: jax.Array, freqs_cis: jax.Array) -> Tup xk_out_0 = xk_0 * cos - xk_1 * sin xk_out_1 = xk_0 * sin + xk_1 * cos - # 5. Stack and reshape back to original - xq_out = jnp.stack([xq_out_0, xq_out_1], axis=-1).reshape(xq.shape) - xk_out = jnp.stack([xk_out_0, xk_out_1], axis=-1).reshape(xk.shape) + # 5. Concatenate along last axis instead of stack + reshape to prevent layout fragmentation + xq_out = jnp.concatenate([xq_out_0[..., None], xq_out_1[..., None]], axis=-1).reshape(xq.shape) + xk_out = jnp.concatenate([xk_out_0[..., None], xk_out_1[..., None]], axis=-1).reshape(xk.shape) return xq_out, xk_out @@ -2874,7 +3371,10 @@ def __call__( if encoder_hidden_states is not None: encoder_hidden_states = nn.with_logical_constraint(encoder_hidden_states, (BATCH, LENGTH, HEAD)) dtype = hidden_states.dtype - is_self_attention = encoder_hidden_states is None + if encoder_hidden_states is not None and encoder_hidden_states is not hidden_states: + is_self_attention = False + else: + is_self_attention = getattr(self, "is_self_attention", True) if encoder_hidden_states is None: encoder_hidden_states = hidden_states @@ -2889,11 +3389,12 @@ def __call__( with jax.named_scope("query_proj"): query_proj = self.query(hidden_states) - if self.qk_norm: - with self.conditional_named_scope("attn_q_norm"): - query_proj = self.norm_q(query_proj) - - if not is_self_attention and cached_kv is not None and "text" in cached_kv: + if is_self_attention: + with jax.named_scope("key_proj"): + key_proj = self.key(hidden_states) + with jax.named_scope("value_proj"): + value_proj = self.value(hidden_states) + elif cached_kv is not None and "text" in cached_kv: key_proj, value_proj = cached_kv["text"] else: with jax.named_scope("key_proj"): @@ -2901,17 +3402,36 @@ def __call__( with jax.named_scope("value_proj"): value_proj = self.value(encoder_hidden_states) - if self.qk_norm: - with self.conditional_named_scope("attn_k_norm"): - key_proj = self.norm_k(key_proj) - - if rotary_emb is not None: - with self.conditional_named_scope("attn_rope"): - query_proj = _unflatten_heads(query_proj, self.heads) - key_proj = _unflatten_heads(key_proj, self.heads) + if rotary_emb is not None and self.qk_norm and is_self_attention: + with self.conditional_named_scope("fused_rmsnorm_rope"): + q_scale = self.norm_q.scale[...] + k_scale = self.norm_k.scale[...] + query_proj, key_proj = fused_rmsnorm_rope( + query_proj, + key_proj, + q_scale, + k_scale, + rotary_emb, + q_heads=self.heads, + dim_head=self.dim_head, + eps=self.eps, + ) value_proj = _unflatten_heads(value_proj, self.heads) - # output of _unflatten_heads Batch, heads, seq_len, head_dim - query_proj, key_proj = self._apply_rope(query_proj, key_proj, rotary_emb) + else: + if self.qk_norm: + with self.conditional_named_scope("attn_q_norm"): + query_proj = self.norm_q(query_proj) + if is_self_attention or cached_kv is None or "text" not in cached_kv: + with self.conditional_named_scope("attn_k_norm"): + key_proj = self.norm_k(key_proj) + + if rotary_emb is not None: + with self.conditional_named_scope("attn_rope"): + query_proj = _unflatten_heads(query_proj, self.heads) + key_proj = _unflatten_heads(key_proj, self.heads) + value_proj = _unflatten_heads(value_proj, self.heads) + # output of _unflatten_heads Batch, heads, seq_len, head_dim + query_proj, key_proj = self._apply_rope(query_proj, key_proj, rotary_emb) query_proj = checkpoint_name(query_proj, "query_proj") key_proj = checkpoint_name(key_proj, "key_proj") diff --git a/src/maxdiffusion/tests/attention_config_guards_test.py b/src/maxdiffusion/tests/attention_config_guards_test.py new file mode 100644 index 000000000..351732f30 --- /dev/null +++ b/src/maxdiffusion/tests/attention_config_guards_test.py @@ -0,0 +1,207 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""CPU checks for the attention-config guards across TPU topologies. + +Guard 1: a non-ring ulysses kernel must REJECT a ulysses_shards it cannot honour. +Guard 2: a ring kernel must WARN when it degenerates to R=1, and the remedy it + suggests must be valid for the actual mesh and head count. +""" +import unittest +from unittest import mock + +from maxdiffusion.models import attention_flax + +# Context-parallel degrees reachable on real slices: v5e/v6e-8 (CP up to 8), +# v6e-16 / v7x-16 (CP up to 16), and larger multi-host slices. +TOPOLOGIES = [1, 2, 4, 8, 16, 32, 64] + +# Wan 2.2 T2V-A14B has 40 attention heads, which is NOT a power of two -- this +# is precisely why a naive "context_shards // 2" suggestion is unsafe. +WAN_HEADS = 40 + + +class ImplicitUlyssesDegreeGuardTest(unittest.TestCase): + + def test_unset_is_allowed_on_every_topology(self): + for cp in TOPOLOGIES: + with self.subTest(cp=cp): + attention_flax._validate_implicit_ulysses_degree(-1, cp, "ulysses_custom") + attention_flax._validate_implicit_ulysses_degree(0, cp, "ulysses_custom") + attention_flax._validate_implicit_ulysses_degree(None, cp, "ulysses_custom") + + def test_matching_request_is_allowed_on_every_topology(self): + for cp in TOPOLOGIES: + with self.subTest(cp=cp): + attention_flax._validate_implicit_ulysses_degree(cp, cp, "ulysses_custom") + + def test_mismatched_request_is_rejected_on_every_topology(self): + for cp in TOPOLOGIES: + for requested in {1, 2, cp // 2, cp * 2} - {cp, 0}: + if requested <= 0: + continue + with self.subTest(cp=cp, requested=requested): + with self.assertRaises(ValueError) as ctx: + attention_flax._validate_implicit_ulysses_degree(requested, cp, "ulysses_custom_fixed_m_per_q_block") + msg = str(ctx.exception) + self.assertIn(f"ulysses_shards={requested}", msg) + self.assertIn(f"context_shards={cp}", msg) + + def test_rejection_names_the_offending_kernel(self): + with self.assertRaises(ValueError) as ctx: + attention_flax._validate_implicit_ulysses_degree(2, 4, "ulysses_custom_fixed_m_per_q_block") + self.assertIn("ulysses_custom_fixed_m_per_q_block", str(ctx.exception)) + + +class RealRingSuggestionTest(unittest.TestCase): + + def test_suggestion_is_valid_for_every_topology(self): + """Whatever U we suggest must satisfy every constraint the ring enforces.""" + for cp in TOPOLOGIES: + with self.subTest(cp=cp, heads=WAN_HEADS): + u = attention_flax._largest_ulysses_shards_for_real_ring(cp, WAN_HEADS, WAN_HEADS) + if u is None: + # Only legitimate when no divisor below cp works. + self.assertTrue( + all(cp % c != 0 or WAN_HEADS % c != 0 for c in range(1, cp)), + f"returned None for cp={cp} despite a valid candidate existing", + ) + continue + self.assertLess(u, cp) + self.assertEqual(cp % u, 0, "suggested U must divide the context shard count") + self.assertEqual(WAN_HEADS % u, 0, "suggested U must divide the head count") + self.assertGreater(cp // u, 1, "suggested U must leave a real ring R>1") + + def test_no_suggestion_for_single_shard(self): + self.assertIsNone(attention_flax._largest_ulysses_shards_for_real_ring(1, WAN_HEADS, WAN_HEADS)) + + def test_prefers_smallest_real_ring(self): + # cp=8, heads=40 -> U=4 gives R=2, the cheapest real ring. + self.assertEqual(attention_flax._largest_ulysses_shards_for_real_ring(8, 40, 40), 4) + # cp=16, heads=40 -> U=16 and U=8 both divide 16, but only 8 divides 40. + self.assertEqual(attention_flax._largest_ulysses_shards_for_real_ring(16, 40, 40), 8) + # cp=32, heads=40 -> largest common divisor below 32 is 8, giving R=4. + self.assertEqual(attention_flax._largest_ulysses_shards_for_real_ring(32, 40, 40), 8) + + def test_respects_asymmetric_kv_heads(self): + # GQA-style: 40 query heads but 8 KV heads restricts U to divisors of 8. + u = attention_flax._largest_ulysses_shards_for_real_ring(16, 40, 8) + self.assertEqual(8 % u, 0) + self.assertEqual(40 % u, 0) + self.assertEqual(16 % u, 0) + + +class DegenerateRingWarningTest(unittest.TestCase): + + def setUp(self): + attention_flax._WARNED_ONCE.clear() + + def test_warns_on_every_topology_when_degenerate(self): + for cp in TOPOLOGIES: + with self.subTest(cp=cp): + attention_flax._WARNED_ONCE.clear() + with mock.patch.object(attention_flax.max_logging, "log") as log: + attention_flax._warn_if_ring_is_degenerate(1, cp, cp, heads=WAN_HEADS, kv_heads=WAN_HEADS) + log.assert_called_once() + msg = log.call_args[0][0] + self.assertIn("R=1", msg) + self.assertIn("Do NOT report this as a ring-attention result", msg) + + def test_suggested_remedy_in_message_is_actionable(self): + for cp in [2, 4, 8, 16, 32]: + with self.subTest(cp=cp): + attention_flax._WARNED_ONCE.clear() + with mock.patch.object(attention_flax.max_logging, "log") as log: + attention_flax._warn_if_ring_is_degenerate(1, cp, cp, heads=WAN_HEADS, kv_heads=WAN_HEADS) + msg = log.call_args[0][0] + expected_u = attention_flax._largest_ulysses_shards_for_real_ring(cp, WAN_HEADS, WAN_HEADS) + self.assertIn(f"ulysses_shards={expected_u}", msg) + self.assertIn(f"R={cp // expected_u}", msg) + + def test_single_shard_mesh_does_not_suggest_zero(self): + """Regression: context_shards//2 would have advised the impossible U=0.""" + with mock.patch.object(attention_flax.max_logging, "log") as log: + attention_flax._warn_if_ring_is_degenerate(1, 1, 1, heads=WAN_HEADS, kv_heads=WAN_HEADS) + msg = log.call_args[0][0] + self.assertNotIn("ulysses_shards=0", msg) + self.assertIn("only one context shard", msg) + + def test_silent_for_real_ring_on_every_topology(self): + for cp in [2, 4, 8, 16, 32]: + u = attention_flax._largest_ulysses_shards_for_real_ring(cp, WAN_HEADS, WAN_HEADS) + with self.subTest(cp=cp, u=u): + attention_flax._WARNED_ONCE.clear() + with mock.patch.object(attention_flax.max_logging, "log") as log: + attention_flax._warn_if_ring_is_degenerate(cp // u, u, cp, heads=WAN_HEADS, kv_heads=WAN_HEADS) + log.assert_not_called() + + def test_warns_only_once_per_configuration(self): + with mock.patch.object(attention_flax.max_logging, "log") as log: + for _ in range(80): # ~40 layers x 2 transformers + attention_flax._warn_if_ring_is_degenerate(1, 4, 4, heads=WAN_HEADS, kv_heads=WAN_HEADS) + log.assert_called_once() + + def test_distinct_configurations_each_warn(self): + with mock.patch.object(attention_flax.max_logging, "log") as log: + attention_flax._warn_if_ring_is_degenerate(1, 4, 4, heads=WAN_HEADS, kv_heads=WAN_HEADS) + attention_flax._warn_if_ring_is_degenerate(1, 8, 8, heads=WAN_HEADS, kv_heads=WAN_HEADS) + self.assertEqual(log.call_count, 2) + + +class KernelClassificationDriftTest(unittest.TestCase): + """Every registered ulysses-ring kernel must be classified for tile sizing. + + `local_tiled_seq_len` silently falls through to `return full_seq` for any + attention name it does not recognise, which yields a mesh-independent (and + therefore wrong) tile length. A newly registered ring kernel must not be able + to slip through that default. + """ + + def test_all_registered_ulysses_ring_kernels_are_classified(self): + from maxdiffusion.utils import tile_size_grid_search as tsgs + + registered = {name for name in attention_flax.KERNEL_REGISTRY if name.startswith("ulysses_ring")} + self.assertTrue(registered, "expected at least one registered ulysses_ring kernel") + missing = registered - tsgs.ULYSSES_RING_ATTENTION_KERNELS + self.assertEqual( + missing, + set(), + f"these ulysses_ring kernels are missing from ULYSSES_RING_ATTENTION_KERNELS and would " + f"get a mesh-independent tile length: {sorted(missing)}", + ) + + def test_classified_kernels_scale_with_topology(self): + from maxdiffusion.utils.tile_size_grid_search import local_tiled_seq_len + + full_seq = 6144 + for attention in sorted(attention_flax.KERNEL_REGISTRY): + if not attention.startswith("ulysses_ring"): + continue + for cp in [2, 4, 8, 16]: + u = attention_flax._largest_ulysses_shards_for_real_ring(cp, WAN_HEADS, WAN_HEADS) + with self.subTest(attention=attention, cp=cp, u=u): + local = local_tiled_seq_len(full_seq, attention, context_shards=cp, ulysses_shards=u) + # Ulysses gathers u chunks of the context-local sequence. + self.assertEqual(local, (full_seq // cp) * u) + self.assertLess( + local, + full_seq, + "a sharded mesh must tile less than the full sequence", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/custom_splash_unpadded_test.py b/src/maxdiffusion/tests/custom_splash_unpadded_test.py new file mode 100644 index 000000000..f783d3fd5 --- /dev/null +++ b/src/maxdiffusion/tests/custom_splash_unpadded_test.py @@ -0,0 +1,407 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""C3 gate: may the custom splash kernel be handed PHYSICALLY UNPADDED q/k/v? + +Context +------- +Production calls `_pad_data_for_flash` on q, k and v, materialising a full +185 MiB copy of each purely to round the sequence dimension up to a multiple of +the Pallas block size. Static analysis says most of that is unnecessary: + + * K/V: the kernel never *masks* the KV tail, it *slices* it + (`slice_k_len = kv_seq_len % bkv_compute`, using the UNPADDED length -- see + `last_compute_body_fixed` / `last_compute_body_online` in + custom_splash_attention.py and the load-bearing comment above them). No + compute path reads a padded K or V row, so their contents are irrelevant. + + * Q: rows are independent (the running max is over KV, never across the `bq` + row axis), and the kernel's OUTPUT is already ragged today -- `out_shape`'s + last dim is `actual_q_seq_len` against a `bq`-wide BlockSpec -- so Pallas is + already clipping a non-divisible last block in this very kernel. + +The one thing static analysis CANNOT settle is whether Pallas clips the ragged +last-block *input* DMA, or issues an unclipped read past the end of the array. +`compiler_params` sets `disable_bounds_checks=True`. If the DMA is unclipped, +passing an unpadded array is a silent out-of-bounds HBM read. + +**That is the only question these tests exist to answer.** Everything else about +C3 is already proven on paper. + +Why the existing coverage does not answer it +-------------------------------------------- +`custom_splash_fixed_m_test.test_non_divisible_sequence_context_padding_fixed_m` +looks like it covers this, but it pads its inputs to a block boundary before the +call (`q_in_padded = jnp.pad(...)`) and only passes a ragged *logical* length. +That is exactly today's regime. No existing test ever hands the kernel a +physically non-block-aligned array. + +Test design +----------- +Each case runs the kernel twice with IDENTICAL LOGICAL INPUTS: + reference -- physically padded to a block boundary (today's behaviour) + candidate -- physically unpadded (what C3 proposes) +and asserts the outputs are **bit-identical**. This is pure data movement: the +grid, the block sizes and every slice length are computed from the unpadded +logical length and are therefore identical between the two runs. The same +arithmetic happens in the same order on the same values, so anything other than +exact equality means memory outside the array was read. + +K/V-unpadded and Q-unpadded are separate cases so the two halves of C3 can be +gated independently. + +Two failure modes are guarded against explicitly: + 1. Passing by luck on a freshly-zeroed buffer -- see `_dirty_device_memory`, + and the repeat-run case. + 2. A harness bug making both sides equally wrong -- every case also checks the + padded reference against a dense f32 softmax reference. + +> A NOTE ON THE LIMITS OF THIS TEST. `_dirty_device_memory` is a heuristic. JAX +> gives no control over HBM placement, so we cannot *guarantee* the bytes past +> the end of the array are NaN. A PASS is therefore strong but not absolute +> evidence; a FAILURE is conclusive. Treat a pass as "no evidence of an +> unclipped DMA under adversarial conditions", not as a proof of clipping. +""" + +import gc +import math +import unittest + +import jax +import jax.numpy as jnp + +from maxdiffusion.kernels import custom_splash_attention as custom_splash + +_LOG2E = math.log2(math.e) + + +def _dirty_device_memory(num_buffers: int = 8, mb_each: int = 64) -> None: + """Fills and releases device buffers of NaN to poison the allocator pool. + + The failure mode this defends against: an unclipped DMA reads whatever + happens to sit past the end of the array. On a fresh device that is often + zero, which is exactly the padding value the kernel would have seen anyway -- + so the bug would produce a correct answer and the test would pass for the + wrong reason. + + By allocating NaN buffers and then dropping them, subsequent allocations are + likely (not guaranteed) to be served from memory containing NaN. NaN is a + stronger probe than a large finite value: a large finite value could be + squashed back to something finite by a downstream mask or a saturating + operation, whereas NaN propagates through every arithmetic path in the + softmax and cannot be masked away once it enters an accumulation. + """ + elems = mb_each * 1024 * 1024 // 4 + junk = [] + for _ in range(num_buffers): + junk.append(jax.block_until_ready(jnp.full((elems,), jnp.nan, dtype=jnp.float32))) + del junk + gc.collect() + + +class CustomSplashUnpaddedInputTest(unittest.TestCase): + """Bit-identity of the kernel when fed physically unpadded q/k/v.""" + + heads = 4 + head_dim = 64 + + # Deliberately non-divisible by `bq` in BOTH dimensions. + # q_len = 1008, bq = 512 -> grid_height = 2, last block covers [512, 1024) + # against an array of only 1008 rows. + # kv_len = 1001, bkv = 512 -> grid_width = 2, tail slice length 489. + q_len = 1008 + kv_len = 1001 + bq = 512 + + def setUp(self): + super().setUp() + if jax.default_backend() == "cpu": + self.skipTest("Pallas splash kernel requires TPU.") + self.scale = 1.0 / math.sqrt(self.head_dim) + self.grid_height = math.ceil(self.q_len / self.bq) + self.q_padded_len = self.grid_height * self.bq + self.kv_padded_len = math.ceil(self.kv_len / self.bq) * self.bq + + # ---------------------------------------------------------------- helpers + + def _inputs(self): + """Returns logical (unpadded) bf16 q, k, v in the kernel's own convention.""" + q = jax.random.normal( + jax.random.PRNGKey(11), + (self.heads, self.q_len, self.head_dim), + jnp.bfloat16, + ) + k = jax.random.normal( + jax.random.PRNGKey(12), + (self.heads, self.kv_len, self.head_dim), + jnp.bfloat16, + ) + v = jax.random.normal( + jax.random.PRNGKey(13), + (self.heads, self.kv_len, self.head_dim), + jnp.bfloat16, + ) + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in = (k * self.scale).astype(jnp.bfloat16) + return q_in, k_in, v, q, k + + def _pad_seq(self, x, target): + if x.shape[1] == target: + return x + return jnp.pad(x, ((0, 0), (0, target - x.shape[1]), (0, 0))) + + def _metadata(self, q_in, k_in): + """Builds `mk` and `k_mean` from UNPADDED inputs. + + This is the C3b norm-vector trick: the per-row norms are computed on the + unpadded query and then the *norm vector* is zero-padded to the block grid, + instead of zero-padding the (far larger) activation. Zero rows contribute 0 + to a max over non-negative values, so the resulting `mk` is identical to the + one derived from a zero-padded activation -- bit-identical, not merely + close. + """ + k_mean = jnp.mean(k_in.astype(jnp.float32), axis=1) # (heads, dim) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.kv_len, is_ring=False) + + k_centered = k_in.astype(jnp.float32) - k_mean[:, None, :] + mk_h = jnp.sqrt((k_centered**2).sum(-1)).max(axis=-1) # (heads,) + + row_norm_sq = (q_in.astype(jnp.float32) ** 2).sum(-1) # (heads, q_len) + pad = self.q_padded_len - row_norm_sq.shape[1] + if pad: + row_norm_sq = jnp.pad(row_norm_sq, ((0, 0), (0, pad))) + qn_max = jnp.sqrt(row_norm_sq.reshape(self.heads, self.grid_height, self.bq).max(axis=-1)) + + bound = qn_max * mk_h[:, None] + mk = jnp.stack( + [jnp.ceil(bound) - recenter, (bound <= safe_bound).astype(jnp.float32)], + axis=0, + ) + return mk, k_mean + + def _run( + self, + q_in, + k_in, + v_in, + mk, + k_mean, + *, + use_fixed_m, + uniform_fixed_m, + bkv_compute=None, + ): + """Invokes the kernel. Block sizes derive only from logical lengths.""" + # `uniform_fixed_m` is a sub-mode of fixed-m; the kernel rejects the + # combination outright ("uniform_fixed_m requires use_fixed_m"). Asserting + # here turns a misconfigured case into a loud harness failure instead of a + # case that aborts at kernel construction and never touches the device -- + # which would look like a kernel finding but is really a test bug. + assert not (uniform_fixed_m and not use_fixed_m), "uniform_fixed_m requires use_fixed_m" + bkv_compute = bkv_compute or self.bq + block_sizes = custom_splash._BlockSizes( + block_q=self.bq, + block_kv=self.bq, + block_kv_compute=bkv_compute, + block_kv_compute_in=bkv_compute, + ) + kernel = custom_splash.make_splash_mha( + block_sizes=block_sizes, + orig_q_seq_len=self.q_len, + orig_kv_seq_len=self.kv_len, + use_base2_exp=True, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + ) + if use_fixed_m: + out = kernel(q_in, k_in, v_in, mk, k_mean) + else: + out = kernel(q_in, k_in, v_in) + return jnp.swapaxes(out, 1, 2).astype(jnp.float32) # (heads, seq, dim) + + def _dense_reference(self, q, k, v): + qf, kf, vf = (x.astype(jnp.float32) for x in (q, k, v)) + logits = jnp.einsum("hsd,htd->hst", qf, kf) * self.scale + return jnp.einsum("hst,htd->hsd", jax.nn.softmax(logits, axis=-1), vf) + + def _assert_reference_is_sane(self, reference, q, k, v): + """Guards against a harness bug that would make both sides equally wrong.""" + self.assertTrue(bool(jnp.all(jnp.isfinite(reference))), "padded reference is not finite") + self.assertGreater(float(jnp.max(jnp.abs(reference))), 0.0, "padded reference is all zeros") + dense = self._dense_reference(q, k, v) + diff = float(jnp.max(jnp.abs(reference[:, : self.q_len] - dense))) + self.assertLess(diff, 5e-2, f"padded reference disagrees with dense softmax: {diff=}") + + def _compare( + self, + *, + unpad_q, + unpad_kv, + use_fixed_m=True, + uniform_fixed_m=None, + bkv_compute=None, + dirty=True, + ): + """Core assertion: unpadded inputs reproduce padded inputs bit-for-bit. + + `uniform_fixed_m` defaults to tracking `use_fixed_m` rather than to a bare + `True`, because `uniform_fixed_m=True` with `use_fixed_m=False` is rejected + by the kernel at construction time. Defaulting it to `True` made the online + cases abort before reaching the device -- they looked like failures of the + kernel when they were failures of this harness. + """ + if uniform_fixed_m is None: + uniform_fixed_m = use_fixed_m + q_in, k_in, v, q_raw, k_raw = self._inputs() + mk, k_mean = self._metadata(q_in, k_in) + + reference = self._run( + self._pad_seq(q_in, self.q_padded_len), + self._pad_seq(k_in, self.kv_padded_len), + self._pad_seq(v, self.kv_padded_len), + mk, + k_mean, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + bkv_compute=bkv_compute, + ) + jax.block_until_ready(reference) + self._assert_reference_is_sane(reference, q_raw, k_raw, v) + + if dirty: + _dirty_device_memory() + + candidate = self._run( + q_in if unpad_q else self._pad_seq(q_in, self.q_padded_len), + k_in if unpad_kv else self._pad_seq(k_in, self.kv_padded_len), + v if unpad_kv else self._pad_seq(v, self.kv_padded_len), + mk, + k_mean, + use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + bkv_compute=bkv_compute, + ) + jax.block_until_ready(candidate) + + self.assertEqual(reference.shape, candidate.shape) + self.assertTrue( + bool(jnp.all(jnp.isfinite(candidate))), + "unpadded run produced non-finite values", + ) + self.assertTrue( + bool(jnp.array_equal(reference, candidate)), + "unpadded inputs changed the result -- the kernel is reading past the end " + f"of the array (max abs delta {float(jnp.max(jnp.abs(reference - candidate)))})", + ) + return reference, candidate + + # ------------------------------------------------------------ C3a: K and V + + def test_kv_unpadded_fixed_m(self): + """C3a gate. Padded K/V rows are sliced away, so removing them must be a no-op.""" + self._compare(unpad_q=False, unpad_kv=True) + + def test_kv_unpadded_hybrid(self): + """C3a on the hybrid kernel, whose tail goes through `_last_online`.""" + self._compare(unpad_q=False, unpad_kv=True, uniform_fixed_m=False) + + def test_kv_unpadded_online(self): + """C3a on the pure online path -- the simplest probe of the DMA question.""" + self._compare(unpad_q=False, unpad_kv=True, use_fixed_m=False) + + def test_kv_unpadded_multi_iteration_tail(self): + """C3a with bkv_compute < bkv, exercising the fori_loop + ragged remainder.""" + self._compare(unpad_q=False, unpad_kv=True, bkv_compute=self.bq // 2) + + def test_kv_unpadded_online_multi_iteration_tail(self): + """Same, on the online path. + + The online path tracks a *running* max instead of a pinned one, so + `_last_online` is genuinely different code from `_last_fixed` even though + both slice the tail the same way. It gets its own ragged-remainder case + rather than inheriting confidence from the fixed-m result. + """ + self._compare(unpad_q=False, unpad_kv=True, use_fixed_m=False, bkv_compute=self.bq // 2) + + # ---------------------------------------------------------------- C3b: Q + + def test_q_unpadded_fixed_m(self): + """C3b gate. Relies on the norm-vector padding in `_metadata`.""" + self._compare(unpad_q=True, unpad_kv=False) + + def test_q_unpadded_hybrid(self): + self._compare(unpad_q=True, unpad_kv=False, uniform_fixed_m=False) + + def test_q_unpadded_online(self): + self._compare(unpad_q=True, unpad_kv=False, use_fixed_m=False) + + # ------------------------------------------------------------ both halves + + def test_all_unpadded_fixed_m(self): + self._compare(unpad_q=True, unpad_kv=True) + + def test_all_unpadded_hybrid(self): + self._compare(unpad_q=True, unpad_kv=True, uniform_fixed_m=False) + + # ------------------------------------------------------- adversarial cases + + def test_all_unpadded_is_stable_across_repeats(self): + """A single cold run can pass by luck on zeroed memory; N runs cannot. + + Between repeats the allocator pool is re-poisoned with NaN. If the kernel + were reading past the end of the array, the value it reads would change + from run to run and at least one repeat would diverge. + """ + first = None + for i in range(4): + _dirty_device_memory(num_buffers=4) + _, candidate = self._compare(unpad_q=True, unpad_kv=True, dirty=False) + if first is None: + first = candidate + else: + self.assertTrue( + bool(jnp.array_equal(first, candidate)), + f"unpadded run {i} differs from run 0 -- result depends on memory " + "outside the array, which is the signature of an unclipped DMA", + ) + + def test_all_unpadded_survives_interleaved_kernel_invocation(self): + """Runs a differently-shaped kernel first so VMEM holds real, non-zero data. + + `_dirty_device_memory` poisons HBM; this poisons the VMEM scratch buffers + that a ragged block DMA would only partially overwrite. The interleaved call + uses a large-magnitude value tensor so any leakage is numerically obvious + rather than lost in rounding. + """ + q_in, k_in, v, _, _ = self._inputs() + mk, k_mean = self._metadata(q_in, k_in) + loud = (jnp.ones_like(v) * 64).astype(v.dtype) + jax.block_until_ready( + self._run( + self._pad_seq(q_in, self.q_padded_len), + self._pad_seq(k_in, self.kv_padded_len), + self._pad_seq(loud, self.kv_padded_len), + mk, + k_mean, + use_fixed_m=True, + uniform_fixed_m=True, + ) + ) + self._compare(unpad_q=True, unpad_kv=True, dirty=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/dot_fallback_layout_test.py b/src/maxdiffusion/tests/dot_fallback_layout_test.py new file mode 100644 index 000000000..c79e71a83 --- /dev/null +++ b/src/maxdiffusion/tests/dot_fallback_layout_test.py @@ -0,0 +1,165 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Layout regression tests for the short-sequence dot-product fallback. + +Sequences below `flash_min_seq_length` bypass the flash/ulysses kernels and +run `_apply_attention_dot`. Callers that apply rotary embeddings hand the +dispatcher `[B, H, S, D]` -- the dispatcher says so itself, reading the +sequence length from axis 2 when `ndim == 4` -- but the `split_head_dim` path +reshaped those tensors as if they were `[B, S, H*D]`. + +Both layouts hold the same number of elements, so the reshape succeeded and +returned a wrong answer with no error. That silent corruption is what these +tests exist to prevent. They are pure JAX and run on CPU. +""" + +import math +import unittest + +import jax +import jax.numpy as jnp +import numpy as np + +from maxdiffusion.models.attention_flax import _apply_attention_dot + + +def _call_dot(query, key, value, heads, dim_head, kv_heads=None): + """Runs the fallback and returns [B, H, S, D]. + + `_apply_attention_dot` emits the flat `[B, S, H*D]` form, so unflatten it + here rather than at each call site. + """ + out = _apply_attention_dot( + query=query, + key=key, + value=value, + dtype=jnp.float32, + heads=heads, + dim_head=dim_head, + scale=1.0 / math.sqrt(dim_head), + split_head_dim=True, + float32_qk_product=True, + use_memory_efficient_attention=False, + attention_mask=None, + kv_heads=kv_heads, + ) + batch, seq, _ = out.shape + return jnp.swapaxes(out.reshape(batch, seq, heads, dim_head), 1, 2) + + +def _reference_attention(query, key, value, scale): + """Dense f32 attention on explicit [B, H, S, D] inputs.""" + q, k, v = (x.astype(jnp.float32) for x in (query, key, value)) + logits = jnp.einsum("bhqd,bhkd->bhqk", q, k) * scale + return jnp.einsum("bhqk,bhkd->bhqd", jax.nn.softmax(logits, axis=-1), v) + + +class DotFallbackLayoutTest(unittest.TestCase): + """`_apply_attention_dot` must transpose, not reshape, 4-D inputs.""" + + def test_zero_logits_return_per_head_token_means(self): + """The reviewer's counterexample, reproduced exactly. + + One active head dimension, three tokens, two heads. Q = K = 0 makes every + logit zero, so softmax is uniform and each head's output is the mean of + its values over tokens: + + head 0 values [1, 2, 3] -> 2 + head 1 values [10, 20, 30] -> 20 + + Reinterpreting [B, H, S, D] as [B, S, H, D] instead walks the buffer + [1, 2, 3, 10, 20, 30] as three token-pairs (1,2), (3,10), (20,30), + producing [8, 14]. + """ + heads, seq, dim_head = 2, 3, 1 + shape = (1, heads, seq, dim_head) + query = jnp.zeros(shape, jnp.float32) + key = jnp.zeros(shape, jnp.float32) + value = jnp.array([[[[1.0], [2.0], [3.0]], [[10.0], [20.0], [30.0]]]], dtype=jnp.float32) + self.assertEqual(value.shape, shape) + + out = np.asarray(_call_dot(query, key, value, heads, dim_head)) + + np.testing.assert_allclose(out[0, 0], 2.0, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(out[0, 1], 20.0, rtol=1e-5, atol=1e-5) + # The exact wrong answer the reshape produced. + self.assertFalse(np.allclose(out[0, 0], 8.0)) + self.assertFalse(np.allclose(out[0, 1], 14.0)) + + def test_matches_reference_on_random_4d_inputs(self): + heads, seq, dim_head = 4, 8, 16 + shape = (2, heads, seq, dim_head) + query = jax.random.normal(jax.random.PRNGKey(0), shape, jnp.float32) + key = jax.random.normal(jax.random.PRNGKey(1), shape, jnp.float32) + value = jax.random.normal(jax.random.PRNGKey(2), shape, jnp.float32) + + out = np.asarray(_call_dot(query, key, value, heads, dim_head)) + expected = np.asarray(_reference_attention(query, key, value, 1.0 / math.sqrt(dim_head))) + np.testing.assert_allclose(out, expected, rtol=1e-4, atol=1e-4) + + def test_three_dim_inputs_agree_with_four_dim(self): + """The flat [B, S, H*D] contract must keep working, and agree.""" + heads, seq, dim_head = 4, 8, 16 + batch = 2 + shape = (batch, heads, seq, dim_head) + query = jax.random.normal(jax.random.PRNGKey(3), shape, jnp.float32) + key = jax.random.normal(jax.random.PRNGKey(4), shape, jnp.float32) + value = jax.random.normal(jax.random.PRNGKey(5), shape, jnp.float32) + + def flatten(x): + return jnp.swapaxes(x, 1, 2).reshape(batch, seq, heads * dim_head) + + out_4d = np.asarray(_call_dot(query, key, value, heads, dim_head)) + out_3d = np.asarray(_call_dot(flatten(query), flatten(key), flatten(value), heads, dim_head)) + np.testing.assert_allclose(out_4d, out_3d, rtol=1e-5, atol=1e-5) + + def test_gqa_repeat_still_applies_on_4d(self): + """Head repetition must happen after the transpose, on the head axis.""" + heads, kv_heads, seq, dim_head = 4, 2, 8, 16 + q = jax.random.normal(jax.random.PRNGKey(6), (1, heads, seq, dim_head), jnp.float32) + k = jax.random.normal(jax.random.PRNGKey(7), (1, kv_heads, seq, dim_head), jnp.float32) + v = jax.random.normal(jax.random.PRNGKey(8), (1, kv_heads, seq, dim_head), jnp.float32) + + out = np.asarray(_call_dot(q, k, v, heads, dim_head, kv_heads=kv_heads)) + expected = np.asarray( + _reference_attention( + q, + jnp.repeat(k, heads // kv_heads, axis=1), + jnp.repeat(v, heads // kv_heads, axis=1), + 1.0 / math.sqrt(dim_head), + ) + ) + np.testing.assert_allclose(out, expected, rtol=1e-4, atol=1e-4) + + +class DispatcherLayoutContractTest(unittest.TestCase): + """The threshold check and the dot path must agree on where S lives.""" + + def test_seq_len_axis_for_4d_is_the_third_axis(self): + """`_apply_attention` reads S from axis 2 for 4-D, i.e. [B, H, S, D]. + + That is the convention `_apply_attention_dot` now honours by transposing. + If the dispatcher ever moves to [B, S, H, D], the transpose becomes wrong + and this assertion should be revisited alongside it. + """ + query = jnp.zeros((1, 4, 128, 8), jnp.float32) # B, H, S, D + seq_len_idx = 2 if query.ndim == 4 else 1 + self.assertEqual(query.shape[seq_len_idx], 128) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/fused_producers_test.py b/src/maxdiffusion/tests/fused_producers_test.py new file mode 100644 index 000000000..2b0a9e78f --- /dev/null +++ b/src/maxdiffusion/tests/fused_producers_test.py @@ -0,0 +1,120 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Numerical-equivalence tests for the fused attention producers. + +These run on CPU; they pin numerical contracts, not kernel performance. +""" + +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx + +from maxdiffusion.kernels.fused_producers import fused_rmsnorm_rope + + +def _reference_rmsnorm(x, scale, eps=1e-6): + """Mirrors flax.nnx.RMSNorm's association: x * (rsqrt(var + eps) * scale). + + Flax's `_normalize` builds `mul = rsqrt(var + eps)`, folds the scale into it + with `mul *= scale`, and only then applies `y *= mul`. Reassociating this as + `(x * rsqrt) * scale` rounds differently, so the order is load-bearing. + """ + var = jnp.mean(jnp.square(x.astype(jnp.float32)), axis=-1, keepdims=True) + return x.astype(jnp.float32) * (jax.lax.rsqrt(var + eps) * scale.astype(jnp.float32)) + + +class FusedRmsNormAssociationTest(unittest.TestCase): + """The fused producer must be a pure fusion, never a numerical change.""" + + def test_matches_flax_rmsnorm_bit_for_bit(self): + dim = 128 + key = jax.random.PRNGKey(0) + k1, k2 = jax.random.split(key) + x = jax.random.normal(k1, (2, 64, dim), jnp.float32) + scale = jax.random.normal(k2, (dim,), jnp.float32) + + layer = nnx.RMSNorm( + dim, + epsilon=1e-6, + dtype=jnp.float32, + param_dtype=jnp.float32, + rngs=nnx.Rngs(0), + ) + layer.scale.value = scale + + np.testing.assert_array_equal( + np.asarray(_reference_rmsnorm(x, scale)), + np.asarray(layer(x)), + err_msg="Reference helper must reproduce nnx.RMSNorm exactly.", + ) + + def test_left_to_right_association_is_not_equivalent(self): + """Guards the reason this test exists: the two orders genuinely differ.""" + dim = 128 + key = jax.random.PRNGKey(1) + k1, k2 = jax.random.split(key) + x = jax.random.normal(k1, (2, 64, dim), jnp.float32) + scale = jax.random.normal(k2, (dim,), jnp.float32) + + rsqrt = jax.lax.rsqrt(jnp.mean(jnp.square(x), axis=-1, keepdims=True) + 1e-6) + folded = x * (rsqrt * scale) # Flax order + left_to_right = (x * rsqrt) * scale + + self.assertFalse( + bool(jnp.all(folded == left_to_right)), + "If these ever become identical the association guard above is vacuous.", + ) + + def test_fused_producer_q_matches_flax_rmsnorm(self): + """End-to-end: the q path of the fused producer must match nnx.RMSNorm.""" + b, seq, q_heads, dim_head = 1, 8, 2, 8 + d_model = q_heads * dim_head + key = jax.random.PRNGKey(2) + k1, k2, k3 = jax.random.split(key, 3) + + raw_q = jax.random.normal(k1, (b, seq, d_model), jnp.float32) + raw_k = jax.random.normal(k2, (b, seq, d_model), jnp.float32) + q_scale = jax.random.normal(k3, (d_model,), jnp.float32) + k_scale = jnp.ones((d_model,), jnp.float32) + + # Identity rotation isolates the RMSNorm from the RoPE. + freqs_cis = jnp.ones((1, 1, seq, dim_head // 2), jnp.complex64) + + q_out, _ = fused_rmsnorm_rope( + raw_q, + raw_k, + q_scale, + k_scale, + freqs_cis, + q_heads=q_heads, + kv_heads=q_heads, + dim_head=dim_head, + ) + + expected = _reference_rmsnorm(raw_q, q_scale).reshape(b, seq, q_heads, dim_head).transpose(0, 2, 1, 3) + np.testing.assert_array_equal( + np.asarray(q_out.astype(jnp.float32)), + np.asarray(expected.astype(raw_q.dtype).astype(jnp.float32)), + err_msg="Fused RMSNorm+RoPE must be bit-identical to nnx.RMSNorm under an identity rotation.", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/wan/wan_vace_transformer_test.py b/src/maxdiffusion/tests/wan/wan_vace_transformer_test.py index ce5fdf2db..adb4788aa 100644 --- a/src/maxdiffusion/tests/wan/wan_vace_transformer_test.py +++ b/src/maxdiffusion/tests/wan/wan_vace_transformer_test.py @@ -120,13 +120,25 @@ def test_wan_vace_block_returns_the_correct_shape(self): apply_input_projection=True, apply_output_projection=True, ) + + @nnx.jit + def forward(m, h, eh, ch, temb, rot): + return m( + hidden_states=h, + encoder_hidden_states=eh, + control_hidden_states=ch, + temb=temb, + rotary_emb=rot, + ) + with mesh: - conditioning_states, control_hidden_states = wan_vace_block( - hidden_states=dummy_hidden_states, - encoder_hidden_states=dummy_encoder_hidden_states, - control_hidden_states=dummy_control_hidden_states, - temb=dummy_temb, - rotary_emb=dummy_rotary_emb, + conditioning_states, control_hidden_states = forward( + wan_vace_block, + dummy_hidden_states, + dummy_encoder_hidden_states, + dummy_control_hidden_states, + dummy_temb, + dummy_rotary_emb, ) assert conditioning_states.shape == dummy_hidden_states.shape assert control_hidden_states.shape == dummy_hidden_states.shape diff --git a/src/maxdiffusion/utils/tile_size_grid_search.py b/src/maxdiffusion/utils/tile_size_grid_search.py index 3b7afb90b..b72f1495e 100644 --- a/src/maxdiffusion/utils/tile_size_grid_search.py +++ b/src/maxdiffusion/utils/tile_size_grid_search.py @@ -46,6 +46,7 @@ "ulysses_ring", "ulysses_ring_custom", "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_fixed_m_per_q_block", "ulysses_ring_custom_bidir", })