From 069e7fb0f1b505d1c215f3c6bc1d8a9fb67edf17 Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Wed, 23 Sep 2026 12:05:46 +0530 Subject: [PATCH] perf(wan2.2): Optimize attention kernels, rematerialization, and Ulysses parallelism for TPU v7x - Implement custom Pallas splash & ring attention forward and fused backward kernels - Optimize rematerialization policy with custom checkpointing rules - Implement Ulysses sequence parallelism with 2-phase all-to-all communications - Eliminate device-to-host sync stall via deterministic host-side expert routing - Support multi-host evaluation metric gathering with multihost_utils.process_allgather - Specify sequential dimension semantics in Pallas backward kernels for safe accumulation - Attach pl.CostEstimate to the splash forward and fused backward pallas_calls so XLA's latency-hiding scheduler can overlap collectives (FSDP reduce-scatter) with attention - Address PR review comments: slice row-0 LSE to cut saved residual memory by 8x, warn on unaliased dQ scratch buffer allocation, document inference-only paths, pass residual_checkpoint_name in tokamax_ring_custom, revert default block_kv_compute_in to None in max_utils, restore architectural comments in ring attention, and add unit test for saved remat residuals. --- .../kernels/custom_splash_attention.py | 1021 ++++++++++++++++- .../splash_attention/ring_attention_kernel.py | 261 ++++- src/maxdiffusion/max_utils.py | 35 +- src/maxdiffusion/models/attention_flax.py | 68 +- .../tests/custom_splash_backward_test.py | 558 +++++++++ src/maxdiffusion/trainers/wan_trainer_2_2.py | 21 +- 6 files changed, 1895 insertions(+), 69 deletions(-) create mode 100644 src/maxdiffusion/tests/custom_splash_backward_test.py diff --git a/src/maxdiffusion/kernels/custom_splash_attention.py b/src/maxdiffusion/kernels/custom_splash_attention.py index 6bd3f3493..f85e520b5 100644 --- a/src/maxdiffusion/kernels/custom_splash_attention.py +++ b/src/maxdiffusion/kernels/custom_splash_attention.py @@ -17,6 +17,7 @@ """Custom Pallas flash attention kernel for TPU.""" import functools +import warnings import jax import jax.numpy as jnp @@ -31,8 +32,26 @@ NT_DIM_NUMBERS = (((1,), (1,)), ((), ())) +LN2 = float(np.log(2.0)) + + class _BlockSizes: - __slots__ = ("block_q", "block_kv", "block_kv_compute", "block_kv_compute_in") + __slots__ = ( + "block_q", + "block_kv", + "block_kv_compute", + "block_kv_compute_in", + "block_q_dkv", + "block_kv_dkv", + "block_kv_dkv_compute", + "block_kv_dkv_compute_in", + "block_q_dq", + "block_kv_dq", + "block_kv_dq_compute", + "block_kv_dq_compute_in", + "use_fused_bwd_kernel", + "dq_reduction_steps", + ) def __init__( self, @@ -40,11 +59,57 @@ def __init__( block_kv: int, block_kv_compute: int | None = None, block_kv_compute_in: int = 256, + block_q_dkv: int | None = None, + block_kv_dkv: int | None = None, + block_kv_dkv_compute: int | None = None, + block_kv_dkv_compute_in: int | None = None, + block_q_dq: int | None = None, + block_kv_dq: int | None = None, + block_kv_dq_compute: int | None = None, + block_kv_dq_compute_in: int | None = None, + use_fused_bwd_kernel: bool = True, + dq_reduction_steps: int | None = 3, ): self.block_q = block_q self.block_kv = block_kv self.block_kv_compute = block_kv_compute if block_kv_compute is not None else block_kv self.block_kv_compute_in = block_kv_compute_in + self.block_q_dkv = block_q_dkv if block_q_dkv is not None else block_q + self.block_kv_dkv = block_kv_dkv if block_kv_dkv is not None else block_kv + self.block_kv_dkv_compute = block_kv_dkv_compute if block_kv_dkv_compute is not None else self.block_kv_dkv + self.block_kv_dkv_compute_in = block_kv_dkv_compute_in if block_kv_dkv_compute_in is not None else block_kv_compute_in + self.block_q_dq = block_q_dq if block_q_dq is not None else block_q + self.block_kv_dq = block_kv_dq if block_kv_dq is not None else block_kv + self.block_kv_dq_compute = block_kv_dq_compute if block_kv_dq_compute is not None else self.block_kv_dq + self.block_kv_dq_compute_in = block_kv_dq_compute_in if block_kv_dq_compute_in is not None else block_kv_compute_in + self.use_fused_bwd_kernel = use_fused_bwd_kernel + self.dq_reduction_steps = dq_reduction_steps + + def _as_tuple(self): + return ( + self.block_q, + self.block_kv, + self.block_kv_compute, + self.block_kv_compute_in, + self.block_q_dkv, + self.block_kv_dkv, + self.block_kv_dkv_compute, + self.block_kv_dkv_compute_in, + self.block_q_dq, + self.block_kv_dq, + self.block_kv_dq_compute, + self.block_kv_dq_compute_in, + self.use_fused_bwd_kernel, + self.dq_reduction_steps, + ) + + def __eq__(self, other): + if not isinstance(other, _BlockSizes): + return False + return self._as_tuple() == other._as_tuple() + + def __hash__(self): + return hash(self._as_tuple()) # Fixed-m softmax-bound constants. Instead of tracking the online-softmax @@ -66,6 +131,27 @@ def __init__( _FIXED_M_RING_SAFE_BOUND = _FIXED_M_SAFE_BOUND / 2.0 +def _attention_cost_estimate( + num_heads: int, + q_len: int, + kv_len: int, + head_dim_qk: int, + head_dim_v: int, + num_matmuls_qk: int, + num_matmuls_v: int, + bytes_accessed: int, +) -> pl.CostEstimate: + """Cost of a splash kernel for XLA's scheduler (and XProf FLOP accounting). + + Without an estimate the latency-hiding scheduler treats these custom calls as + free, so it cannot overlap the Ulysses all-to-alls of one head chunk with the + attention kernel of another and leaves every all-to-all exposed. + """ + scores = num_heads * q_len * kv_len + flops = 2 * scores * (num_matmuls_qk * head_dim_qk + num_matmuls_v * head_dim_v) + return pl.CostEstimate(flops=int(flops), transcendentals=int(scores), bytes_accessed=int(bytes_accessed)) + + def _flash_attention_kernel( mk_ref, q_ref, @@ -89,6 +175,7 @@ def _flash_attention_kernel( fuse_reciprocal: bool = True, use_fixed_m: bool = False, uniform_fixed_m: bool = False, + save_lse: bool = False, ): float32 = jnp.float32 head_dim_v_repeats, rem = divmod(head_dim_v, NUM_SUBLANES) @@ -315,7 +402,11 @@ def end(): # can merge shard contributions and normalize only once at the very end. o_ref[...] = o_scratch_ref[...].astype(o_ref.dtype) if l_ring_ref is not None: - l_ring_ref[...] = l.astype(l_ring_ref.dtype) + if save_lse: + log = jnp.log2 if use_base2_exp else jnp.log + l_ring_ref[...] = (m_scratch_ref[...] + log(l)).astype(l_ring_ref.dtype) + else: + l_ring_ref[...] = l.astype(l_ring_ref.dtype) if m_ring_ref is not None: m_ring_ref[...] = m_scratch_ref[...].astype(m_ring_ref.dtype) @@ -483,6 +574,7 @@ def _splash_attention_forward( vmem_limit_bytes: int | None = None, use_fixed_m: bool = False, mk: jax.Array | None = None, + save_residuals: bool = False, ): num_q_heads, padded_q_seq_len, head_dim_qk = q.shape head_dim_v = v.shape[-1] @@ -513,6 +605,15 @@ def k_index_map(h, i, j, *_): def v_index_map(h, i, j, *_): return (h // q_heads_per_kv_head, j, 0) + grid_width = (actual_kv_seq_len + bkv - 1) // bkv + grid_height = (actual_q_seq_len + bq - 1) // bq + active_q_len = grid_height * bq + active_kv_len = grid_width * bkv + + q_in = jnp.pad(q, ((0, 0), (0, active_q_len - q.shape[1]), (0, 0))) if q.shape[1] < active_q_len else q + k_in = jnp.pad(k, ((0, 0), (0, active_kv_len - k.shape[1]), (0, 0))) if k.shape[1] < active_kv_len else k + v_in = jnp.pad(v, ((0, 0), (0, active_kv_len - v.shape[1]), (0, 0))) if v.shape[1] < active_kv_len else v + in_specs = [ pl.BlockSpec((None, bq, head_dim_qk), q_index_map), pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), @@ -522,7 +623,7 @@ def v_index_map(h, i, j, *_): jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), q.dtype), + jax.ShapeDtypeStruct((num_q_heads, head_dim_v, active_q_len), q.dtype), ] out_specs = [ pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), @@ -530,9 +631,26 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), pl.BlockSpec((None, head_dim_v, bq), out_index_map), ] - grid_width = (actual_kv_seq_len + bkv - 1) // bkv - grid_height = (actual_q_seq_len + bq - 1) // bq + if save_residuals: + out_shapes.append(jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, active_q_len), jnp.float32)) + out_specs.append(pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map)) + grid = (num_q_heads, grid_height, grid_width) + itemsize = q.dtype.itemsize + cost_estimate = _attention_cost_estimate( + num_q_heads, + actual_q_seq_len, + actual_kv_seq_len, + head_dim_qk, + head_dim_v, + num_matmuls_qk=1, + num_matmuls_v=1, + bytes_accessed=itemsize + * ( + num_q_heads * active_q_len * (head_dim_qk + head_dim_v) + + grid_height * num_kv_heads * active_kv_len * (head_dim_qk + head_dim_v) + ), + ) all_out = pl.pallas_call( functools.partial( @@ -546,6 +664,7 @@ def v_index_map(h, i, j, *_): kv_seq_len=actual_kv_seq_len, use_base2_exp=use_base2_exp, use_fixed_m=use_fixed_m, + save_lse=save_residuals, ), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=1, @@ -561,8 +680,12 @@ def v_index_map(h, i, j, *_): vmem_limit_bytes=vmem_limit_bytes, ), out_shape=out_shapes, - )(mk, q, k, v) - return all_out[-1] + cost_estimate=cost_estimate, + )(mk, q_in, k_in, v_in) + out = all_out[3][:, :, :actual_q_seq_len] + if save_residuals: + return out, all_out[4][:, 0, :actual_q_seq_len] + return out def _splash_attention_forward_ring( @@ -617,6 +740,15 @@ def k_index_map(h, i, j, *_): def v_index_map(h, i, j, *_): return (h // q_heads_per_kv_head, j, 0) + grid_width = (actual_kv_seq_len + bkv - 1) // bkv + grid_height = (actual_q_seq_len + bq - 1) // bq + active_q_len = grid_height * bq + active_kv_len = grid_width * bkv + + q_in = jnp.pad(q, ((0, 0), (0, active_q_len - q.shape[1]), (0, 0))) if q.shape[1] < active_q_len else q + k_in = jnp.pad(k, ((0, 0), (0, active_kv_len - k.shape[1]), (0, 0))) if k.shape[1] < active_kv_len else k + v_in = jnp.pad(v, ((0, 0), (0, active_kv_len - v.shape[1]), (0, 0))) if v.shape[1] < active_kv_len else v + in_specs = [ pl.BlockSpec((None, bq, head_dim_qk), q_index_map), pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), @@ -626,9 +758,9 @@ def v_index_map(h, i, j, *_): jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), jax.ShapeDtypeStruct((head_dim_v, bq), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, head_dim_v, actual_q_seq_len), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), - jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, actual_q_seq_len), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, head_dim_v, active_q_len), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, active_q_len), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, NUM_SUBLANES, active_q_len), jnp.float32), ] out_specs = [ pl.BlockSpec((NUM_SUBLANES, bq), lambda *_: (0, 0)), @@ -638,8 +770,6 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), pl.BlockSpec((None, NUM_SUBLANES, bq), out_index_map), ] - grid_width = (actual_kv_seq_len + bkv - 1) // bkv - grid_height = (actual_q_seq_len + bq - 1) // bq grid = (num_q_heads, grid_height, grid_width) # Scalar-prefetch operand carrying per-head fixed-m data (same convention as @@ -678,10 +808,10 @@ def v_index_map(h, i, j, *_): vmem_limit_bytes=vmem_limit_bytes, ), out_shape=out_shapes, - )(mk, q, k, v) - out = jnp.swapaxes(all_out[3], 1, 2) # (h, head_dim_v, s) -> (h, s, head_dim_v) - l = all_out[4][:, 0, :] # (h, s) - m = all_out[5][:, 0, :] # (h, s) + )(mk, q_in, k_in, v_in) + out = jnp.swapaxes(all_out[3][:, :, :actual_q_seq_len], 1, 2) # (h, head_dim_v, s) -> (h, s, head_dim_v) + l = all_out[4][:, 0, :actual_q_seq_len] # (h, s) + m = all_out[5][:, 0, :actual_q_seq_len] # (h, s) return out, m, l @@ -774,6 +904,832 @@ def out_index_map(h, i, j, *_): return all_out[-1] +def _compute_p_and_ds(qk_slice, v_slice, do, lse, di, exp_fn): + """Computes attention weights p and gradient ds for a KV chunk.""" + p_curr = exp_fn(qk_slice - lse) + dp_curr = lax.dot_general( + v_slice, + do.astype(v_slice.dtype), + (((1,), (0,)), ((), ())), + preferred_element_type=jnp.float32, + ) + ds_curr = p_curr * (dp_curr - di) + return p_curr, ds_curr + + +def _compute_dq_tile(ds_curr, k_slice): + """Computes dQ gradient tile for a given dS and K chunk.""" + return lax.dot_general( + ds_curr.astype(k_slice.dtype), + k_slice, + (((0,), (0,)), ((), ())), + preferred_element_type=jnp.float32, + ) + + +def _accumulate_dk_dv(sub_slice, p_curr, ds_curr, q, do, dk_scratch_ref, dv_scratch_ref): + """Accumulates gradient contributions for dK and dV tiles into scratch references.""" + float32 = jnp.float32 + dv_curr = lax.dot_general( + p_curr.astype(do.dtype), + do, + NT_DIM_NUMBERS, + preferred_element_type=float32, + ) + dv_scratch_ref[sub_slice, :] = dv_scratch_ref[sub_slice, :] + dv_curr + + dk_curr = lax.dot_general( + ds_curr.astype(q.dtype), + q, + (((1,), (0,)), ((), ())), + preferred_element_type=float32, + ) + dk_scratch_ref[sub_slice, :] = dk_scratch_ref[sub_slice, :] + dk_curr + + +def _run_kv_compute_loop( + j, + grid_width: int, + bkv: int, + bkv_compute: int, + kv_seq_len: int, + compute_body, + last_compute_body, +): + """Executes the loop over KV chunks with exact tail handling.""" + assert bkv % bkv_compute == 0 + + @pl.when(j != grid_width - 1) + def body(): + lax.fori_loop(0, (bkv // bkv_compute), compute_body, None, unroll=True) + + @pl.when(j == grid_width - 1) + def last_body(): + if kv_seq_len % bkv == 0: + iter_num = bkv // bkv_compute + lax.fori_loop(0, iter_num, compute_body, None, unroll=True) + else: + remain_kv_seq_len = kv_seq_len % bkv + iter_num = (remain_kv_seq_len + bkv_compute - 1) // bkv_compute + if remain_kv_seq_len % bkv_compute == 0: + lax.fori_loop(0, iter_num, compute_body, None, unroll=True) + else: + lax.fori_loop(0, iter_num - 1, compute_body, None, unroll=True) + last_compute_body(iter_num - 1) + + +def _flash_attention_dq_kernel( + q_ref, + k_ref, + v_ref, + do_ref, + lse_ref, + di_ref, + dq_scratch_ref, + dq_ref, + *, + grid_width: int, + bkv: int, + bkv_compute: int, + bkv_compute_in: int, + kv_seq_len: int, + use_base2_exp: bool = True, +): + float32 = jnp.float32 + _, _, j = pl.program_id(0), pl.program_id(1), pl.program_id(2) + exp = jnp.exp2 if use_base2_exp else jnp.exp + + @pl.when(j == 0) + def init(): + dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref) + + def _dq_inner(qk, k_chunk, v_chunk, do, lse, di, dq_prev): + step = bkv_compute_in + for idx in range(0, qk.shape[0], step): + sub_len = min(step, qk.shape[0] - idx) + qk_slice = qk[idx : idx + sub_len] + v_slice = v_chunk[idx : idx + sub_len] + k_slice = k_chunk[idx : idx + sub_len] + + _, ds_curr = _compute_p_and_ds(qk_slice, v_slice, do, lse, di, exp) + dq_prev = dq_prev + _compute_dq_tile(ds_curr, k_slice) + return dq_prev + + def compute_body(kv_compute_index, _): + q = q_ref[...] + do = do_ref[...] + lse = lse_ref[0:1, :] + di = di_ref[0:1, :] + base_offset = kv_compute_index * bkv_compute + slice_k = pl.ds(base_offset, bkv_compute) + k_chunk = k_ref[slice_k, :] + v_chunk = v_ref[slice_k, :] + qk = lax.dot_general(k_chunk, q, NT_DIM_NUMBERS, preferred_element_type=float32) + dq_scratch_ref[...] = _dq_inner(qk, k_chunk, v_chunk, do, lse, di, dq_scratch_ref[...]) + + def last_compute_body(kv_compute_index): + q = q_ref[...] + do = do_ref[...] + lse = lse_ref[0:1, :] + di = di_ref[0:1, :] + slice_k_len = kv_seq_len % bkv_compute + slice_k = pl.ds(kv_compute_index * bkv_compute, slice_k_len) + k_chunk = k_ref[slice_k, :] + v_chunk = v_ref[slice_k, :] + qk = lax.dot_general(k_chunk, q, NT_DIM_NUMBERS, preferred_element_type=float32) + dq_scratch_ref[...] = _dq_inner(qk, k_chunk, v_chunk, do, lse, di, dq_scratch_ref[...]) + + _run_kv_compute_loop(j, grid_width, bkv, bkv_compute, kv_seq_len, compute_body, last_compute_body) + + @pl.when(j == grid_width - 1) + def end(): + if use_base2_exp: + dq_ref[...] = (dq_scratch_ref[...] * LN2).astype(dq_ref.dtype) + else: + dq_ref[...] = dq_scratch_ref[...].astype(dq_ref.dtype) + + +def _flash_attention_dkv_kernel( + q_ref, + k_ref, + v_ref, + do_ref, + lse_ref, + di_ref, + dk_scratch_ref, + dv_scratch_ref, + dk_ref, + dv_ref, + *, + grid_width: int, + total_q_steps: int, + bkv: int, + bkv_compute: int, + bkv_compute_in: int, + kv_seq_len: int, + use_base2_exp: bool = True, +): + float32 = jnp.float32 + _, j, step_q = pl.program_id(0), pl.program_id(1), pl.program_id(2) + exp = jnp.exp2 if use_base2_exp else jnp.exp + + @pl.when(step_q == 0) + def init(): + dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref) + dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref) + + def _dkv_inner(base_offset, qk, q, v_chunk, do, lse, di): + step = bkv_compute_in + for idx in range(0, qk.shape[0], step): + sub_len = min(step, qk.shape[0] - idx) + sub_slice = pl.ds(base_offset + idx, sub_len) + qk_slice = qk[idx : idx + sub_len] + v_slice = v_chunk[idx : idx + sub_len] + + p_curr, ds_curr = _compute_p_and_ds(qk_slice, v_slice, do, lse, di, exp) + _accumulate_dk_dv(sub_slice, p_curr, ds_curr, q, do, dk_scratch_ref, dv_scratch_ref) + + def compute_body(kv_compute_index, _): + q = q_ref[...] + do = do_ref[...] + lse = lse_ref[0:1, :] + di = di_ref[0:1, :] + base_offset = kv_compute_index * bkv_compute + slice_k = pl.ds(base_offset, bkv_compute) + k_chunk = k_ref[slice_k, :] + v_chunk = v_ref[slice_k, :] + qk = lax.dot_general(k_chunk, q, NT_DIM_NUMBERS, preferred_element_type=float32) + _dkv_inner(base_offset, qk, q, v_chunk, do, lse, di) + + def last_compute_body(kv_compute_index): + q = q_ref[...] + do = do_ref[...] + lse = lse_ref[0:1, :] + di = di_ref[0:1, :] + base_offset = kv_compute_index * bkv_compute + slice_k_len = kv_seq_len % bkv_compute + slice_k = pl.ds(base_offset, slice_k_len) + k_chunk = k_ref[slice_k, :] + v_chunk = v_ref[slice_k, :] + qk = lax.dot_general(k_chunk, q, NT_DIM_NUMBERS, preferred_element_type=float32) + _dkv_inner(base_offset, qk, q, v_chunk, do, lse, di) + + _run_kv_compute_loop(j, grid_width, bkv, bkv_compute, kv_seq_len, compute_body, last_compute_body) + + @pl.when(step_q == total_q_steps - 1) + def end(): + if use_base2_exp: + dk_ref[...] = (dk_scratch_ref[...] * LN2).astype(dk_ref.dtype) + else: + dk_ref[...] = dk_scratch_ref[...].astype(dk_ref.dtype) + dv_ref[...] = dv_scratch_ref[...].astype(dv_ref.dtype) + + +def _flash_attention_bwd_fused_kernel( + q_ref, + k_ref, + v_ref, + do_ref, + lse_ref, + di_ref, + dq_alias_ref, + dq_ref, + dk_ref, + dv_ref, + dq_scratch_ref, + dk_scratch_ref, + dv_scratch_ref, + *, + grid_width: int, + grid_height: int, + q_heads_per_kv_head: int, + bkv: int, + bkv_compute: int, + bkv_compute_in: int, + kv_seq_len: int, + use_base2_exp: bool = True, + use_dq_aliasing: bool = False, +): + """Fused backward kernel iterating over KV outer (grid_width) and Q inner (num_q_heads, grid_height). + + Computes dQ, dK, and dV in a single pass without masks, ignoring KV tail padding via + slice bounds and accumulating dK/dV in VMEM across Q steps and KV head groups. + """ + float32 = jnp.float32 + j, h_q, i = pl.program_id(0), pl.program_id(1), pl.program_id(2) + exp = jnp.exp2 if use_base2_exp else jnp.exp + + q_head_in_group = lax.rem(h_q, q_heads_per_kv_head) + should_init_dkv = jnp.logical_and(i == 0, q_head_in_group == 0) + should_write_dkv = jnp.logical_and(i == grid_height - 1, q_head_in_group == q_heads_per_kv_head - 1) + + @pl.when(should_init_dkv) + def init_dkv(): + dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref) + dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref) + + dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref) + + def _bwd_fused_inner(base_offset, qk, k_chunk, v_chunk, q, do, lse, di): + step = bkv_compute_in + dq_acc = dq_scratch_ref[...] + for idx in range(0, qk.shape[0], step): + sub_len = min(step, qk.shape[0] - idx) + sub_slice = pl.ds(base_offset + idx, sub_len) + qk_slice = qk[idx : idx + sub_len] + k_slice = k_chunk[idx : idx + sub_len] + v_slice = v_chunk[idx : idx + sub_len] + + p_curr, ds_curr = _compute_p_and_ds(qk_slice, v_slice, do, lse, di, exp) + _accumulate_dk_dv(sub_slice, p_curr, ds_curr, q, do, dk_scratch_ref, dv_scratch_ref) + dq_acc = dq_acc + _compute_dq_tile(ds_curr, k_slice) + dq_scratch_ref[...] = dq_acc + + def compute_body(kv_compute_index, _): + q = q_ref[...] + do = do_ref[...] + lse = lse_ref[0:1, :] + di = di_ref[0:1, :] + base_offset = kv_compute_index * bkv_compute + slice_k = pl.ds(base_offset, bkv_compute) + k_chunk = k_ref[slice_k, :] + v_chunk = v_ref[slice_k, :] + qk = lax.dot_general(k_chunk, q, NT_DIM_NUMBERS, preferred_element_type=float32) + _bwd_fused_inner(base_offset, qk, k_chunk, v_chunk, q, do, lse, di) + + def last_compute_body(kv_compute_index): + q = q_ref[...] + do = do_ref[...] + lse = lse_ref[0:1, :] + di = di_ref[0:1, :] + base_offset = kv_compute_index * bkv_compute + slice_k_len = kv_seq_len % bkv_compute + slice_k = pl.ds(base_offset, slice_k_len) + k_chunk = k_ref[slice_k, :] + v_chunk = v_ref[slice_k, :] + qk = lax.dot_general(k_chunk, q, NT_DIM_NUMBERS, preferred_element_type=float32) + _bwd_fused_inner(base_offset, qk, k_chunk, v_chunk, q, do, lse, di) + + _run_kv_compute_loop(j, grid_width, bkv, bkv_compute, kv_seq_len, compute_body, last_compute_body) + + if use_base2_exp: + dq_val = (dq_scratch_ref[...] * LN2).astype(dq_ref.dtype) + else: + dq_val = dq_scratch_ref[...].astype(dq_ref.dtype) + + if use_dq_aliasing: + dq_ref[...] = dq_alias_ref[...] + dq_val + else: + dq_ref[...] = dq_val + + @pl.when(should_write_dkv) + def end_dkv(): + if use_base2_exp: + dk_ref[...] = (dk_scratch_ref[...] * LN2).astype(dk_ref.dtype) + else: + dk_ref[...] = dk_scratch_ref[...].astype(dk_ref.dtype) + dv_ref[...] = dv_scratch_ref[...].astype(dv_ref.dtype) + + +def _splash_attention_bwd_fused( + q: jax.Array, + k: jax.Array, + v: jax.Array, + do: jax.Array, + lse: jax.Array, + di: jax.Array, + block_sizes: _BlockSizes, + actual_q_seq_len: int, + actual_kv_seq_len: int, + padded_q_seq_len: int, + padded_kv_seq_len: int, + use_base2_exp: bool = True, + use_experimental_scheduler: bool = False, + vmem_limit_bytes: int | None = None, +): + num_q_heads, _, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + num_kv_heads = k.shape[0] + q_heads_per_kv_head = num_q_heads // num_kv_heads + + bq = block_sizes.block_q_dkv + bkv = block_sizes.block_kv_dkv + bkv_compute = block_sizes.block_kv_dkv_compute + bkv_compute_in = block_sizes.block_kv_dkv_compute_in + + grid_width = (actual_kv_seq_len + bkv - 1) // bkv + grid_height = (actual_q_seq_len + bq - 1) // bq + active_q_len = grid_height * bq + active_kv_len = grid_width * bkv + + if actual_q_seq_len < active_q_len: + pad_q = active_q_len - actual_q_seq_len + q_bwd = jnp.pad(q[:, :actual_q_seq_len, :], ((0, 0), (0, pad_q), (0, 0))) + do_bwd = jnp.pad(do[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + lse_bwd = jnp.pad(lse[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + di_bwd = jnp.pad(di[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + elif q.shape[1] > active_q_len: + q_bwd = q[:, :active_q_len, :] + do_bwd = do[:, :, :active_q_len] + lse_bwd = lse[:, :, :active_q_len] + di_bwd = di[:, :, :active_q_len] + else: + q_bwd, do_bwd, lse_bwd, di_bwd = q, do, lse, di + + k_bwd = jnp.pad(k, ((0, 0), (0, active_kv_len - k.shape[1]), (0, 0))) if k.shape[1] < active_kv_len else k + v_bwd = jnp.pad(v, ((0, 0), (0, active_kv_len - v.shape[1]), (0, 0))) if v.shape[1] < active_kv_len else v + + use_dq_aliasing = block_sizes.dq_reduction_steps == 3 and grid_width > 3 + + if use_dq_aliasing: + + def dq_index_map_aliased(j, h_q, i, *_): + return (j % 3, h_q, i, 0) + + dq_spec = pl.BlockSpec((None, None, bq, head_dim_qk), dq_index_map_aliased) + dq_alias_spec = dq_spec + dq_dtype = jnp.float32 + dq_shape = jax.ShapeDtypeStruct((3, num_q_heads, active_q_len, head_dim_qk), dq_dtype) + dq_init = jnp.zeros((3, num_q_heads, active_q_len, head_dim_qk), dtype=dq_dtype) + else: + if grid_width > 3: + warnings.warn( + f"Custom splash backward running without dQ aliasing (dq_reduction_steps={block_sizes.dq_reduction_steps}). " + f"Allocating unaliased dQ scratch shape ({grid_width}, {num_q_heads}, {active_q_len}, {head_dim_qk}) in HBM, " + "which may consume significant memory on long sequences." + ) + + def dq_index_map_standard(j, h_q, i, *_): + return (j, h_q, i, 0) + + dq_spec = pl.BlockSpec((None, None, bq, head_dim_qk), dq_index_map_standard) + dq_alias_spec = None + dq_dtype = q.dtype if grid_width == 1 else jnp.float32 + dq_shape = jax.ShapeDtypeStruct((grid_width, num_q_heads, active_q_len, head_dim_qk), dq_dtype) + dq_init = None + + in_specs = [ + pl.BlockSpec((None, bq, head_dim_qk), lambda j, h_q, i, *_: (h_q, i, 0)), + pl.BlockSpec((None, bkv, head_dim_qk), lambda j, h_q, i, *_: (h_q // q_heads_per_kv_head, j, 0)), + pl.BlockSpec((None, bkv, head_dim_v), lambda j, h_q, i, *_: (h_q // q_heads_per_kv_head, j, 0)), + pl.BlockSpec((None, head_dim_v, bq), lambda j, h_q, i, *_: (h_q, 0, i)), + pl.BlockSpec((None, NUM_SUBLANES, bq), lambda j, h_q, i, *_: (h_q, 0, i)), + pl.BlockSpec((None, NUM_SUBLANES, bq), lambda j, h_q, i, *_: (h_q, 0, i)), + dq_alias_spec, + ] + + out_shapes = [ + dq_shape, + jax.ShapeDtypeStruct((num_kv_heads, active_kv_len, head_dim_qk), k.dtype), + jax.ShapeDtypeStruct((num_kv_heads, active_kv_len, head_dim_v), v.dtype), + jax.ShapeDtypeStruct((bq, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct((bkv, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct((bkv, head_dim_v), jnp.float32), + ] + out_specs = [ + dq_spec, + pl.BlockSpec((None, bkv, head_dim_qk), lambda j, h_q, i, *_: (h_q // q_heads_per_kv_head, j, 0)), + pl.BlockSpec((None, bkv, head_dim_v), lambda j, h_q, i, *_: (h_q // q_heads_per_kv_head, j, 0)), + pl.BlockSpec((bq, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((bkv, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((bkv, head_dim_v), lambda *_: (0, 0)), + ] + grid = (grid_width, num_q_heads, grid_height) + input_output_aliases = {6: 0} if use_dq_aliasing else {} + itemsize = q.dtype.itemsize + cost_estimate = _attention_cost_estimate( + num_q_heads, + actual_q_seq_len, + actual_kv_seq_len, + head_dim_qk, + head_dim_v, + num_matmuls_qk=3, # QK^T recompute, dK, dQ + num_matmuls_v=2, # dP, dV + bytes_accessed=( + grid_width * num_q_heads * active_q_len * ((head_dim_qk + head_dim_v) * itemsize + 2 * NUM_SUBLANES * 4) + + 2 * grid_width * num_q_heads * active_q_len * head_dim_qk * 4 # f32 dQ partials read+write + + 2 * num_kv_heads * active_kv_len * (head_dim_qk + head_dim_v) * itemsize # K, V in; dK, dV out + ), + ) + + all_out = pl.pallas_call( + functools.partial( + _flash_attention_bwd_fused_kernel, + grid_width=grid_width, + grid_height=grid_height, + q_heads_per_kv_head=q_heads_per_kv_head, + bkv=bkv, + bkv_compute=bkv_compute, + bkv_compute_in=bkv_compute_in, + kv_seq_len=actual_kv_seq_len, + use_base2_exp=use_base2_exp, + use_dq_aliasing=use_dq_aliasing, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=0, + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": use_experimental_scheduler}, + disable_bounds_checks=True, + skip_device_barrier=True, + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=out_shapes, + input_output_aliases=input_output_aliases, + cost_estimate=cost_estimate, + )(q_bwd, k_bwd, v_bwd, do_bwd, lse_bwd, di_bwd, dq_init) + + dq_unreduced, dk, dv = all_out[0], all_out[1], all_out[2] + if grid_width == 1: + dq = dq_unreduced[0].astype(q.dtype) + else: + dq = dq_unreduced.sum(axis=0).astype(q.dtype) + + if active_q_len > padded_q_seq_len: + dq = dq[:, :padded_q_seq_len, :] + elif active_q_len < padded_q_seq_len: + dq = jnp.pad(dq, ((0, 0), (0, padded_q_seq_len - active_q_len), (0, 0))) + + if active_kv_len > padded_kv_seq_len: + dk = dk[:, :padded_kv_seq_len, :] + dv = dv[:, :padded_kv_seq_len, :] + elif active_kv_len < padded_kv_seq_len: + pad_kv = padded_kv_seq_len - active_kv_len + dk = jnp.pad(dk, ((0, 0), (0, pad_kv), (0, 0))) + dv = jnp.pad(dv, ((0, 0), (0, pad_kv), (0, 0))) + + return dq, dk, dv + + +def _splash_attention_backward( + q: jax.Array, + k: jax.Array, + v: jax.Array, + o: jax.Array, + lse: jax.Array, + do: jax.Array, + block_sizes: _BlockSizes, + q_seq_len: int | None = None, + kv_seq_len: int | None = None, + use_base2_exp: bool = True, + use_experimental_scheduler: bool = False, + vmem_limit_bytes: int | None = None, + di: jax.Array | None = None, +): + num_q_heads, padded_q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + num_kv_heads = k.shape[0] + padded_kv_seq_len = k.shape[1] + + actual_q_seq_len = q_seq_len if q_seq_len is not None else padded_q_seq_len + actual_kv_seq_len = kv_seq_len if kv_seq_len is not None else padded_kv_seq_len + q_heads_per_kv_head = num_q_heads // num_kv_heads + + if di is None: + di_vec = jnp.sum(o.astype(jnp.float32) * do.astype(jnp.float32), axis=1) + di = jnp.broadcast_to(di_vec[:, None, :], (num_q_heads, NUM_SUBLANES, o.shape[2])) + + if lse.ndim == 2: + lse = jnp.broadcast_to(lse[:, None, :], (num_q_heads, NUM_SUBLANES, lse.shape[1])) + + if block_sizes.use_fused_bwd_kernel: + return _splash_attention_bwd_fused( + q=q, + k=k, + v=v, + do=do, + lse=lse, + di=di, + block_sizes=block_sizes, + actual_q_seq_len=actual_q_seq_len, + actual_kv_seq_len=actual_kv_seq_len, + padded_q_seq_len=padded_q_seq_len, + padded_kv_seq_len=padded_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + ) + + # --- 1. Compute dQ --- + bq_dq, bkv_dq = block_sizes.block_q_dq, block_sizes.block_kv_dq + bkv_dq_compute = block_sizes.block_kv_dq_compute + bkv_dq_compute_in = block_sizes.block_kv_dq_compute_in + grid_width_dq = (actual_kv_seq_len + bkv_dq - 1) // bkv_dq + grid_height_dq = (actual_q_seq_len + bq_dq - 1) // bq_dq + active_q_len_dq = grid_height_dq * bq_dq + active_kv_len_dq = grid_width_dq * bkv_dq + + if actual_q_seq_len < active_q_len_dq: + pad_q = active_q_len_dq - actual_q_seq_len + q_dq = jnp.pad(q[:, :actual_q_seq_len, :], ((0, 0), (0, pad_q), (0, 0))) + do_dq = jnp.pad(do[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + lse_dq = jnp.pad(lse[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + di_dq = jnp.pad(di[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + elif q.shape[1] > active_q_len_dq: + q_dq = q[:, :active_q_len_dq, :] + do_dq = do[:, :, :active_q_len_dq] + lse_dq = lse[:, :, :active_q_len_dq] + di_dq = di[:, :, :active_q_len_dq] + else: + q_dq, do_dq, lse_dq, di_dq = q, do, lse, di + + k_dq = jnp.pad(k, ((0, 0), (0, active_kv_len_dq - k.shape[1]), (0, 0))) if k.shape[1] < active_kv_len_dq else k + v_dq = jnp.pad(v, ((0, 0), (0, active_kv_len_dq - v.shape[1]), (0, 0))) if v.shape[1] < active_kv_len_dq else v + + dq_in_specs = [ + pl.BlockSpec((None, bq_dq, head_dim_qk), lambda h, i, j, *_: (h, i, 0)), + pl.BlockSpec((None, bkv_dq, head_dim_qk), lambda h, i, j, *_: (h // q_heads_per_kv_head, j, 0)), + pl.BlockSpec((None, bkv_dq, head_dim_v), lambda h, i, j, *_: (h // q_heads_per_kv_head, j, 0)), + pl.BlockSpec((None, head_dim_v, bq_dq), lambda h, i, j, *_: (h, 0, i)), + pl.BlockSpec((None, NUM_SUBLANES, bq_dq), lambda h, i, j, *_: (h, 0, i)), + pl.BlockSpec((None, NUM_SUBLANES, bq_dq), lambda h, i, j, *_: (h, 0, i)), + ] + dq_out_shapes = [ + jax.ShapeDtypeStruct((bq_dq, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, active_q_len_dq, head_dim_qk), q.dtype), + ] + dq_out_specs = [ + pl.BlockSpec((bq_dq, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((None, bq_dq, head_dim_qk), lambda h, i, j, *_: (h, i, 0)), + ] + dq_grid = (num_q_heads, grid_height_dq, grid_width_dq) + + _, dq = pl.pallas_call( + functools.partial( + _flash_attention_dq_kernel, + grid_width=grid_width_dq, + bkv=bkv_dq, + bkv_compute=bkv_dq_compute, + bkv_compute_in=bkv_dq_compute_in, + kv_seq_len=actual_kv_seq_len, + use_base2_exp=use_base2_exp, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=0, + in_specs=dq_in_specs, + out_specs=dq_out_specs, + grid=dq_grid, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary"), + flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": use_experimental_scheduler}, + disable_bounds_checks=True, + skip_device_barrier=True, + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=dq_out_shapes, + )(q_dq, k_dq, v_dq, do_dq, lse_dq, di_dq) + + if active_q_len_dq > padded_q_seq_len: + dq = dq[:, :padded_q_seq_len, :] + elif active_q_len_dq < padded_q_seq_len: + dq = jnp.pad(dq, ((0, 0), (0, padded_q_seq_len - active_q_len_dq), (0, 0))) + + # --- 2. Compute dK, dV --- + bq_dkv, bkv_dkv = block_sizes.block_q_dkv, block_sizes.block_kv_dkv + bkv_dkv_compute = block_sizes.block_kv_dkv_compute + bkv_dkv_compute_in = block_sizes.block_kv_dkv_compute_in + grid_width_dkv = (actual_kv_seq_len + bkv_dkv - 1) // bkv_dkv + grid_height_dkv = (actual_q_seq_len + bq_dkv - 1) // bq_dkv + active_q_len_dkv = grid_height_dkv * bq_dkv + active_kv_len_dkv = grid_width_dkv * bkv_dkv + total_q_steps = q_heads_per_kv_head * grid_height_dkv + + if actual_q_seq_len < active_q_len_dkv: + pad_q = active_q_len_dkv - actual_q_seq_len + q_dkv = jnp.pad(q[:, :actual_q_seq_len, :], ((0, 0), (0, pad_q), (0, 0))) + do_dkv = jnp.pad(do[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + lse_dkv = jnp.pad(lse[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + di_dkv = jnp.pad(di[:, :, :actual_q_seq_len], ((0, 0), (0, 0), (0, pad_q))) + elif q.shape[1] > active_q_len_dkv: + q_dkv = q[:, :active_q_len_dkv, :] + do_dkv = do[:, :, :active_q_len_dkv] + lse_dkv = lse[:, :, :active_q_len_dkv] + di_dkv = di[:, :, :active_q_len_dkv] + else: + q_dkv, do_dkv, lse_dkv, di_dkv = q, do, lse, di + + k_dkv = jnp.pad(k, ((0, 0), (0, active_kv_len_dkv - k.shape[1]), (0, 0))) if k.shape[1] < active_kv_len_dkv else k + v_dkv = jnp.pad(v, ((0, 0), (0, active_kv_len_dkv - v.shape[1]), (0, 0))) if v.shape[1] < active_kv_len_dkv else v + + def q_step_map(h_kv, j, step_q, *_): + h_q = h_kv * q_heads_per_kv_head + (step_q // grid_height_dkv) + i = step_q % grid_height_dkv + return (h_q, i, 0) + + def do_step_map(h_kv, j, step_q, *_): + h_q = h_kv * q_heads_per_kv_head + (step_q // grid_height_dkv) + i = step_q % grid_height_dkv + return (h_q, 0, i) + + dkv_in_specs = [ + pl.BlockSpec((None, bq_dkv, head_dim_qk), q_step_map), + pl.BlockSpec((None, bkv_dkv, head_dim_qk), lambda h_kv, j, step_q, *_: (h_kv, j, 0)), + pl.BlockSpec((None, bkv_dkv, head_dim_v), lambda h_kv, j, step_q, *_: (h_kv, j, 0)), + pl.BlockSpec((None, head_dim_v, bq_dkv), do_step_map), + pl.BlockSpec((None, NUM_SUBLANES, bq_dkv), do_step_map), + pl.BlockSpec((None, NUM_SUBLANES, bq_dkv), do_step_map), + ] + dkv_out_shapes = [ + jax.ShapeDtypeStruct((bkv_dkv, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct((bkv_dkv, head_dim_v), jnp.float32), + jax.ShapeDtypeStruct((num_kv_heads, active_kv_len_dkv, head_dim_qk), k.dtype), + jax.ShapeDtypeStruct((num_kv_heads, active_kv_len_dkv, head_dim_v), v.dtype), + ] + dkv_out_specs = [ + pl.BlockSpec((bkv_dkv, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((bkv_dkv, head_dim_v), lambda *_: (0, 0)), + pl.BlockSpec((None, bkv_dkv, head_dim_qk), lambda h_kv, j, step_q, *_: (h_kv, j, 0)), + pl.BlockSpec((None, bkv_dkv, head_dim_v), lambda h_kv, j, step_q, *_: (h_kv, j, 0)), + ] + dkv_grid = (num_kv_heads, grid_width_dkv, total_q_steps) + + _, _, dk, dv = pl.pallas_call( + functools.partial( + _flash_attention_dkv_kernel, + grid_width=grid_width_dkv, + total_q_steps=total_q_steps, + bkv=bkv_dkv, + bkv_compute=bkv_dkv_compute, + bkv_compute_in=bkv_dkv_compute_in, + kv_seq_len=actual_kv_seq_len, + use_base2_exp=use_base2_exp, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=0, + in_specs=dkv_in_specs, + out_specs=dkv_out_specs, + grid=dkv_grid, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary"), + flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": use_experimental_scheduler}, + disable_bounds_checks=True, + skip_device_barrier=True, + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=dkv_out_shapes, + )(q_dkv, k_dkv, v_dkv, do_dkv, lse_dkv, di_dkv) + + if active_kv_len_dkv > padded_kv_seq_len: + dk = dk[:, :padded_kv_seq_len, :] + dv = dv[:, :padded_kv_seq_len, :] + elif active_kv_len_dkv < padded_kv_seq_len: + pad_kv = padded_kv_seq_len - active_kv_len_dkv + dk = jnp.pad(dk, ((0, 0), (0, pad_kv), (0, 0))) + dv = jnp.pad(dv, ((0, 0), (0, pad_kv), (0, 0))) + + return dq, dk, dv + + +@functools.partial( + jax.custom_vjp, + nondiff_argnames=( + "block_sizes", + "q_seq_len", + "kv_seq_len", + "use_base2_exp", + "use_experimental_scheduler", + "vmem_limit_bytes", + "residual_checkpoint_name", + ), +) +def _splash_attention_custom( + q: jax.Array, + k: jax.Array, + v: jax.Array, + block_sizes: _BlockSizes, + q_seq_len: int | None = None, + kv_seq_len: int | None = None, + use_base2_exp: bool = True, + use_experimental_scheduler: bool = False, + vmem_limit_bytes: int | None = None, + residual_checkpoint_name: str | None = None, +): + out = _splash_attention_forward( + q, + k, + v, + block_sizes, + q_seq_len=q_seq_len, + kv_seq_len=kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + save_residuals=False, + ) + if residual_checkpoint_name is not None: + out = jax.ad_checkpoint.checkpoint_name(out, name=residual_checkpoint_name) + return out + + +def _splash_attention_fwd( + q: jax.Array, + k: jax.Array, + v: jax.Array, + block_sizes: _BlockSizes, + q_seq_len: int | None = None, + kv_seq_len: int | None = None, + use_base2_exp: bool = True, + use_experimental_scheduler: bool = False, + vmem_limit_bytes: int | None = None, + residual_checkpoint_name: str | None = None, +): + out, lse = _splash_attention_forward( + q, + k, + v, + block_sizes, + q_seq_len=q_seq_len, + kv_seq_len=kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + save_residuals=True, + ) + if residual_checkpoint_name is not None: + out = jax.ad_checkpoint.checkpoint_name(out, name=residual_checkpoint_name) + lse = jax.ad_checkpoint.checkpoint_name(lse, name=residual_checkpoint_name) + return out, (q, k, v, out, lse) + + +def _splash_attention_bwd( + block_sizes: _BlockSizes, + q_seq_len: int | None, + kv_seq_len: int | None, + use_base2_exp: bool, + use_experimental_scheduler: bool, + vmem_limit_bytes: int | None, + residual_checkpoint_name: str | None, + residuals, + do: jax.Array, +): + del residual_checkpoint_name + q, k, v, out, lse = residuals + dq, dk, dv = _splash_attention_backward( + q, + k, + v, + out, + lse, + do, + block_sizes, + q_seq_len=q_seq_len, + kv_seq_len=kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + ) + return dq, dk, dv + + +_splash_attention_custom.defvjp(_splash_attention_fwd, _splash_attention_bwd) + + def make_splash_mha( block_sizes: _BlockSizes, orig_q_seq_len: int | None = None, @@ -783,8 +1739,20 @@ def make_splash_mha( use_experimental_scheduler: bool = False, vmem_limit_bytes: int | None = None, use_fixed_m: bool = False, + residual_checkpoint_name: str | None = None, ): + """Builds a splash multi-head attention callable. + + Note: + - Standard dense training uses `_splash_attention_custom` (which defines custom VJP + forward and backward rules with activation checkpointing). + - `heads_per_tile > 1`, `use_fixed_m=True`, and `mk is not None` are inference-only + optimizations. They execute forward-only kernels without custom backward passes. + """ + def _splash_attention(q, k, v, mk=None): + # Note: heads_per_tile > 1 and fixed-m/mk paths are forward-only (inference-only); + # they do not define a custom backward VJP rule. if heads_per_tile > 1: if use_fixed_m: raise NotImplementedError("fixed-m is not supported with heads_per_tile > 1") @@ -800,18 +1768,31 @@ def _splash_attention(q, k, v, mk=None): use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, ) - return _splash_attention_forward( + if use_fixed_m or mk is not None: + return _splash_attention_forward( + q, + k, + v, + block_sizes, + q_seq_len=orig_q_seq_len, + kv_seq_len=orig_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=use_fixed_m, + mk=mk, + ) + return _splash_attention_custom( q, k, v, - block_sizes, + block_sizes=block_sizes, q_seq_len=orig_q_seq_len, kv_seq_len=orig_kv_seq_len, use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, - use_fixed_m=use_fixed_m, - mk=mk, + residual_checkpoint_name=residual_checkpoint_name, ) return _splash_attention diff --git a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py index bc49c5af7..7dd8cb7b1 100644 --- a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py +++ b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py @@ -19,6 +19,7 @@ from typing import Any import jax +from jax import ad_checkpoint from jax import lax from jax import tree_util import jax.numpy as jnp @@ -354,6 +355,10 @@ def _ring_attention_fwd( ring_axis=ring_axis, rotate_segment_ids=rotate_segment_ids, ) + if config is not None and config.residual_checkpoint_name is not None: + out = ad_checkpoint.checkpoint_name(out, name=config.residual_checkpoint_name) + if logsumexp is not None: + logsumexp = ad_checkpoint.checkpoint_name(logsumexp, name=config.residual_checkpoint_name) residuals = (q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info) return out, residuals @@ -431,6 +436,8 @@ def _ring_attention_custom( ring_axis=ring_axis, rotate_segment_ids=rotate_segment_ids, ) + if config is not None and config.residual_checkpoint_name is not None: + out = ad_checkpoint.checkpoint_name(out, name=config.residual_checkpoint_name) return out @@ -859,8 +866,9 @@ def _custom_ring_attention_forward( bidirectional: bool = False, use_fixed_m: bool = False, fixed_m_norms: tuple[jax.Array, jax.Array] | None = None, -) -> jax.Array: - """Forward-only ring attention using the custom dense splash kernel. + save_residuals: bool = False, +) -> jax.Array | tuple[jax.Array, jax.Array]: + """Ring attention using the custom dense splash kernel. Args: q: Query shard, shape `(num_q_heads, q_seq_len, head_dim_qk)`. Stationary @@ -870,7 +878,6 @@ def _custom_ring_attention_forward( the ring axis. v: Value shard, shape `(num_kv_heads, kv_seq_len, head_dim_v)`. Rotated. block_sizes: Custom-kernel block sizes (block_q / block_kv / block_kv_compute). - bkv_compute_in: Inner VPU register-tiling step for the custom kernel. orig_q_seq_len: Un-padded local query length (grid bound). orig_kv_seq_len: Un-padded local key/value length (grid bound). Assumed equal across all shards (uniform per-shard padding), matching the @@ -888,9 +895,12 @@ def _custom_ring_attention_forward( perm: Explicit `ppermute` permutation. Defaults to a full-axis +1 rotation. For the hybrid split, pass a perm that rotates K/V *within each ring sub-group only* (built by the caller from the U x R factorization). + save_residuals: If True, returns `(out, lse)` where `lse` is the global + logsumexp of shape `(num_q_heads, orig_q_seq_len)`. Returns: - Normalized attention output, shape `(num_q_heads, q_seq_len, head_dim_v)`. + Normalized attention output of shape `(num_q_heads, q_seq_len, head_dim_v)`, + or `(out, lse)` if `save_residuals=True`. """ axis_size = lax.axis_size(ring_axis) if bidirectional: @@ -901,6 +911,8 @@ def _custom_ring_attention_forward( ) if use_fixed_m: raise NotImplementedError("fixed-m is not yet supported on the bidirectional ring path.") + if save_residuals: + raise NotImplementedError("save_residuals is not yet supported on the bidirectional ring path.") return _custom_bidirectional_ring_forward( q, k, @@ -922,6 +934,7 @@ def _custom_ring_attention_forward( shift = partial(lax.ppermute, axis_name=ring_axis, perm=perm) exp_fn = jnp.exp2 if use_base2_exp else jnp.exp + log_fn = jnp.log2 if use_base2_exp else jnp.log num_q_heads = q.shape[0] head_dim_v = v.shape[-1] @@ -936,9 +949,10 @@ def _custom_ring_attention_forward( # invariant to the kernel's m convention, so overshoot cancels exactly. # The K-shard norms rotate WITH K/V (a (heads,)-sized ppermute) instead of # being re-reduced per hop, which would stall the kernel's scalar prefetch. + if save_residuals: + raise NotImplementedError("save_residuals is not supported with use_fixed_m.") if fixed_m_norms is None: raise ValueError("use_fixed_m on the ring path requires fixed_m_norms=(qn_max, mk_h).") - log_fn = jnp.log2 if use_base2_exp else jnp.log qn_max, mk_h_init = fixed_m_norms tiny = jnp.finfo(jnp.float32).tiny # Finite (not -inf) init: the first merge computes exp(init - lse_new) = 0.0 @@ -1128,9 +1142,215 @@ def _lse_scan(_): l_inv = jnp.where(l_final == 0.0, 0.0, 1.0 / l_final) out = (o_final * l_inv[..., None]).astype(q.dtype) + if save_residuals: + lse = m_final + log_fn(jnp.maximum(l_final, jnp.finfo(jnp.float32).tiny)) + lse = jnp.where(l_final == 0.0, mask_value, lse) + return out, lse return out +def _custom_ring_attention_backward( + q: jax.Array, + k: jax.Array, + v: jax.Array, + out: jax.Array, + lse: jax.Array, + do: jax.Array, + *, + block_sizes: "custom_splash._BlockSizes", + orig_q_seq_len: int, + orig_kv_seq_len: int, + use_base2_exp: bool, + use_experimental_scheduler: bool, + vmem_limit_bytes: int | None, + ring_axis: str, + ring_size: int | None = None, + perm: list[tuple[int, int]] | tuple[tuple[int, int], ...] | None = None, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Backward ring attention using the custom dense splash backward kernel.""" + axis_size = lax.axis_size(ring_axis) + if ring_size is None: + ring_size = axis_size + if perm is None: + perm = [(i, (i + 1) % axis_size) for i in range(axis_size)] + + shift = partial(lax.ppermute, axis_name=ring_axis, perm=perm) + + di_vec = jnp.sum(out.astype(jnp.float32) * do.astype(jnp.float32), axis=-1) + di_expanded = jnp.broadcast_to(di_vec[:, None, :], (q.shape[0], custom_splash.NUM_SUBLANES, orig_q_seq_len)) + out_swapped = jnp.swapaxes(out, 1, 2) + do_swapped = jnp.swapaxes(do, 1, 2) + lse_expanded = jnp.broadcast_to(lse[:, None, :], (q.shape[0], custom_splash.NUM_SUBLANES, orig_q_seq_len)) + + dq_accum = jnp.zeros_like(q, dtype=jnp.float32) + dk_accum = jnp.zeros_like(k, dtype=jnp.float32) + dv_accum = jnp.zeros_like(v, dtype=jnp.float32) + + k_current, v_current = k, v + for hop in range(ring_size): + if hop != ring_size - 1: + k_next = shift(k_current) + v_next = shift(v_current) + + dq_i, dk_i, dv_i = custom_splash._splash_attention_backward( # pylint: disable=protected-access + q, + k_current, + v_current, + out_swapped, + lse_expanded, + do_swapped, + block_sizes, + q_seq_len=orig_q_seq_len, + kv_seq_len=orig_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + di=di_expanded, + ) + dq_accum = dq_accum + dq_i.astype(jnp.float32) + dk_accum = dk_accum + dk_i.astype(jnp.float32) + dv_accum = dv_accum + dv_i.astype(jnp.float32) + + if ring_size > 1: + dk_accum = shift(dk_accum) + dv_accum = shift(dv_accum) + if hop != ring_size - 1: + k_current, v_current = k_next, v_next + + return dq_accum.astype(q.dtype), dk_accum.astype(k.dtype), dv_accum.astype(v.dtype) + + +@partial( + jax.custom_vjp, + nondiff_argnames=( + "block_sizes", + "orig_q_seq_len", + "orig_kv_seq_len", + "use_base2_exp", + "use_experimental_scheduler", + "vmem_limit_bytes", + "mask_value", + "ring_axis", + "ring_size", + "perm", + "residual_checkpoint_name", + ), +) +def _custom_ring_attention_custom( + q: jax.Array, + k: jax.Array, + v: jax.Array, + block_sizes: "custom_splash._BlockSizes", + orig_q_seq_len: int, + orig_kv_seq_len: int, + use_base2_exp: bool, + use_experimental_scheduler: bool, + vmem_limit_bytes: int | None, + mask_value: float, + ring_axis: str, + ring_size: int | None, + perm: tuple[tuple[int, int], ...] | None, + residual_checkpoint_name: str | None = None, +) -> jax.Array: + out = _custom_ring_attention_forward( + q, + k, + v, + block_sizes=block_sizes, + orig_q_seq_len=orig_q_seq_len, + orig_kv_seq_len=orig_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + mask_value=mask_value, + ring_axis=ring_axis, + ring_size=ring_size, + perm=list(perm) if perm is not None else None, + save_residuals=False, + ) + if residual_checkpoint_name is not None: + out = ad_checkpoint.checkpoint_name(out, name=residual_checkpoint_name) + return out + + +def _custom_ring_attention_fwd( + q: jax.Array, + k: jax.Array, + v: jax.Array, + block_sizes: "custom_splash._BlockSizes", + orig_q_seq_len: int, + orig_kv_seq_len: int, + use_base2_exp: bool, + use_experimental_scheduler: bool, + vmem_limit_bytes: int | None, + mask_value: float, + ring_axis: str, + ring_size: int | None, + perm: tuple[tuple[int, int], ...] | None, + residual_checkpoint_name: str | None = None, +): + out, lse = _custom_ring_attention_forward( + q, + k, + v, + block_sizes=block_sizes, + orig_q_seq_len=orig_q_seq_len, + orig_kv_seq_len=orig_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + mask_value=mask_value, + ring_axis=ring_axis, + ring_size=ring_size, + perm=list(perm) if perm is not None else None, + save_residuals=True, + ) + if residual_checkpoint_name is not None: + out = ad_checkpoint.checkpoint_name(out, name=residual_checkpoint_name) + lse = ad_checkpoint.checkpoint_name(lse, name=residual_checkpoint_name) + return out, (q, k, v, out, lse) + + +def _custom_ring_attention_bwd( + block_sizes: "custom_splash._BlockSizes", + orig_q_seq_len: int, + orig_kv_seq_len: int, + use_base2_exp: bool, + use_experimental_scheduler: bool, + vmem_limit_bytes: int | None, + mask_value: float, + ring_axis: str, + ring_size: int | None, + perm: tuple[tuple[int, int], ...] | None, + residual_checkpoint_name: str | None, + residuals, + do: jax.Array, +): + del mask_value, residual_checkpoint_name + q, k, v, out, lse = residuals + dq, dk, dv = _custom_ring_attention_backward( + q, + k, + v, + out, + lse, + do, + block_sizes=block_sizes, + orig_q_seq_len=orig_q_seq_len, + orig_kv_seq_len=orig_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=ring_size, + perm=list(perm) if perm is not None else None, + ) + return dq, dk, dv + + +_custom_ring_attention_custom.defvjp(_custom_ring_attention_fwd, _custom_ring_attention_bwd) + + def make_custom_ring_attention( *, block_sizes: "custom_splash._BlockSizes", @@ -1146,8 +1366,9 @@ def make_custom_ring_attention( bidirectional: bool = False, use_fixed_m: bool = False, fixed_m_norms: tuple[jax.Array, jax.Array] | None = None, + residual_checkpoint_name: str | None = None, ): - """Builds a forward-only ring-attention callable around the custom kernel. + """Builds a ring-attention callable around the custom kernel (supports forward & backward). The returned function takes a single (un-batched) `(q, k, v)` triple of shape `(num_heads, seq, head_dim)` and is meant to be `jax.vmap`-ped over the batch @@ -1162,9 +1383,29 @@ def make_custom_ring_attention( one hop at a time) for a NON-wrapping ring axis, avoiding the diameter-length wrap hop. Requires `perm=None` and the full real ring axis (no sub-group). """ + perm_tuple = tuple(tuple(x) for x in perm) if perm is not None else None def _ring(q, k, v): - return _custom_ring_attention_forward( + if use_fixed_m or bidirectional: + return _custom_ring_attention_forward( + q, + k, + v, + block_sizes=block_sizes, + orig_q_seq_len=orig_q_seq_len, + orig_kv_seq_len=orig_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + mask_value=mask_value, + ring_axis=ring_axis, + ring_size=ring_size, + perm=perm, + bidirectional=bidirectional, + use_fixed_m=use_fixed_m, + fixed_m_norms=fixed_m_norms, + ) + return _custom_ring_attention_custom( q, k, v, @@ -1177,10 +1418,8 @@ def _ring(q, k, v): mask_value=mask_value, ring_axis=ring_axis, ring_size=ring_size, - perm=perm, - bidirectional=bidirectional, - use_fixed_m=use_fixed_m, - fixed_m_norms=fixed_m_norms, + perm=perm_tuple, + residual_checkpoint_name=residual_checkpoint_name, ) return _ring diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index 171fb91ea..21d1cc7b0 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -314,7 +314,9 @@ def download_blobs(source_gcs_folder, local_destination): if len(file_split[-1]) <= 0: continue download_to_filename = os.path.join(directory, file_split[-1]) - if not os.path.isfile(download_to_filename): + if not os.path.isfile(download_to_filename) or ( + blob.size is not None and os.path.getsize(download_to_filename) != blob.size + ): blob.download_to_filename(download_to_filename) return os.path.join(local_destination, prefix_name) @@ -823,6 +825,16 @@ class CustomFlashBlockSizes: block_kv: int | None = None block_kv_compute: int | None = None block_kv_compute_in: int | None = None + block_q_dkv: int | None = None + block_kv_dkv: int | None = None + block_kv_dkv_compute: int | None = None + block_kv_dkv_compute_in: int | None = None + block_q_dq: int | None = None + block_kv_dq: int | None = None + block_kv_dq_compute: int | None = None + block_kv_dq_compute_in: int | None = None + use_fused_bwd_kernel: bool = True + dq_reduction_steps: int | None = 3 heads_per_tile: int | None = None vmem_limit_bytes: int | None = None @@ -882,16 +894,27 @@ def get_flash_block_sizes(config): block_kv=user_block_sizes.get("block_kv"), block_kv_compute=user_block_sizes.get("block_kv_compute"), block_kv_compute_in=user_block_sizes.get("block_kv_compute_in"), + block_q_dkv=user_block_sizes.get("block_q_dkv", user_block_sizes.get("block_q")), + block_kv_dkv=user_block_sizes.get("block_kv_dkv", user_block_sizes.get("block_kv")), + block_kv_dkv_compute=user_block_sizes.get("block_kv_dkv_compute", user_block_sizes.get("block_kv_compute")), + block_kv_dkv_compute_in=user_block_sizes.get( + "block_kv_dkv_compute_in", user_block_sizes.get("block_kv_compute_in") + ), + block_q_dq=user_block_sizes.get("block_q_dq", user_block_sizes.get("block_q")), + block_kv_dq=user_block_sizes.get("block_kv_dq", user_block_sizes.get("block_kv")), + block_kv_dq_compute=user_block_sizes.get("block_kv_dq_compute", user_block_sizes.get("block_kv_compute")), + block_kv_dq_compute_in=user_block_sizes.get("block_kv_dq_compute_in", user_block_sizes.get("block_kv_compute_in")), + use_fused_bwd_kernel=bool(user_block_sizes.get("use_fused_bwd_kernel", True)), + dq_reduction_steps=user_block_sizes.get("dq_reduction_steps", 3), heads_per_tile=user_block_sizes.get("heads_per_tile"), vmem_limit_bytes=user_block_sizes.get("vmem_limit_bytes"), ) if attention_is_tokamax: max_logging.log( - "Tokamax kernel specified, Note: Tokamax only supports fused backward kernel." - "Hence following flash block properties specified will be ignored:" - f"block_q: {user_block_sizes['block_q']}," - f"block_q_dq: {user_block_sizes.get('block_q_dq')}," - f"block_kv_dq: {user_block_sizes.get('block_kv_dq')}," + "Tokamax kernel specified, Note: Tokamax only supports fused backward kernel. " + "Hence following flash block properties specified will be ignored: " + f"block_q_dq: {user_block_sizes.get('block_q_dq')}, " + f"block_kv_dq: {user_block_sizes.get('block_kv_dq')}, " f"use_fused_bwd_kernel: {user_block_sizes.get('use_fused_bwd_kernel')}" ) use_fused_bwd = True if attention_is_tokamax else bool(user_block_sizes.get("use_fused_bwd_kernel", False)) diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 3db542586..6358b46b3 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -72,16 +72,16 @@ def _coerce_tokamax_block_sizes(block_sizes): # Tokamax requires fused bwd; convert if needed. - if getattr(block_sizes, "use_fused_bwd_kernel", False): + if isinstance(block_sizes, splash_attention_kernel.BlockSizes) and getattr(block_sizes, "use_fused_bwd_kernel", False): return block_sizes # Fall back if some fields are missing. bq = block_sizes.block_q - bkv = getattr(block_sizes, "block_kv", bq) - bkv_compute = getattr(block_sizes, "block_kv_compute", bkv) - bq_dkv = getattr(block_sizes, "block_q_dkv", bq) - bkv_dkv = getattr(block_sizes, "block_kv_dkv", bkv) - bkv_dkv_compute = getattr(block_sizes, "block_kv_dkv_compute", bkv_compute) + bkv = getattr(block_sizes, "block_kv", None) or bq + bkv_compute = getattr(block_sizes, "block_kv_compute", None) or bkv + bq_dkv = getattr(block_sizes, "block_q_dkv", None) or bq + bkv_dkv = getattr(block_sizes, "block_kv_dkv", None) or bkv + bkv_dkv_compute = getattr(block_sizes, "block_kv_dkv_compute", None) or bkv_compute return splash_attention_kernel.BlockSizes( block_q=bq, block_kv=bkv, @@ -409,6 +409,36 @@ def _extract_custom_block_sizes(flash_block_sizes): return bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes +def _build_custom_splash_block_sizes(flash_block_sizes, bq, bkv, bkv_compute, bkv_compute_in): + """Constructs custom_splash._BlockSizes preserving backward tile settings.""" + + def _get(attr, default=None): + if flash_block_sizes is None: + return default + if isinstance(flash_block_sizes, dict): + val = flash_block_sizes.get(attr, None) + else: + val = getattr(flash_block_sizes, attr, None) + return default if val is None else val + + return custom_splash._BlockSizes( + block_q=bq, + block_kv=bkv, + block_kv_compute=bkv_compute, + block_kv_compute_in=bkv_compute_in, + block_q_dkv=_get("block_q_dkv", bq), + block_kv_dkv=_get("block_kv_dkv", bkv), + block_kv_dkv_compute=_get("block_kv_dkv_compute", bkv_compute), + block_kv_dkv_compute_in=_get("block_kv_dkv_compute_in", bkv_compute_in), + block_q_dq=_get("block_q_dq", bq), + block_kv_dq=_get("block_kv_dq", bkv), + block_kv_dq_compute=_get("block_kv_dq_compute", bkv_compute), + block_kv_dq_compute_in=_get("block_kv_dq_compute_in", bkv_compute_in), + use_fused_bwd_kernel=bool(_get("use_fused_bwd_kernel", True)), + dq_reduction_steps=_get("dq_reduction_steps", 3), + ) + + def _build_padding_segment_ids( query_seq_len: int, q_padded_len: int, @@ -644,12 +674,7 @@ def wrap_flash_attention(query, key, value, attention_mask): key_local, _, key_seq_len = _pad_data_for_flash(key, heads, bkv) value_local, _, _ = _pad_data_for_flash(value, heads, bkv) - bsizes = custom_splash._BlockSizes( - block_q=bq, - block_kv=bkv, - block_kv_compute=bkv_compute, - block_kv_compute_in=bkv_compute_in, - ) + bsizes = _build_custom_splash_block_sizes(flash_block_sizes, bq, bkv, bkv_compute, bkv_compute_in) ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( block_sizes=bsizes, orig_q_seq_len=query_seq_len, @@ -658,6 +683,7 @@ def wrap_flash_attention(query, key, value, attention_mask): use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, ring_axis="context", + residual_checkpoint_name=residual_checkpoint_name, ) vmapped_ring = jax.vmap(ring_kernel, in_axes=(0, 0, 0)) attention_output = vmapped_ring(query_local, key_local, value_local) @@ -953,12 +979,7 @@ def wrap_ulysses_attention(query, key, value, attention_mask): fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) - bsizes = custom_splash._BlockSizes( - block_q=bq, - block_kv=bkv, - block_kv_compute=bkv_compute, - block_kv_compute_in=bkv_compute_in, - ) + bsizes = _build_custom_splash_block_sizes(flash_block_sizes, bq, bkv, bkv_compute, bkv_compute_in) splash_kernel = custom_splash.make_splash_mha( block_sizes=bsizes, @@ -969,6 +990,7 @@ def wrap_ulysses_attention(query, key, value, attention_mask): use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=use_fixed_m, + residual_checkpoint_name=residual_checkpoint_name, ) if use_fixed_m: @@ -1466,12 +1488,7 @@ def wrap_ulysses_ring_attention(query, key, value): fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) - bsizes = custom_splash._BlockSizes( - block_q=bq, - block_kv=bkv, - block_kv_compute=bkv_compute, - block_kv_compute_in=bkv_compute_in, - ) + bsizes = _build_custom_splash_block_sizes(flash_block_sizes, bq, bkv, bkv_compute, bkv_compute_in) 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). @@ -1486,6 +1503,7 @@ def wrap_ulysses_ring_attention(query, key, value): use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=use_fixed_m, + residual_checkpoint_name=residual_checkpoint_name, ) if use_fixed_m: attention_output = jnp.swapaxes( @@ -1511,6 +1529,7 @@ def wrap_ulysses_ring_attention(query, key, value): bidirectional=bidirectional, use_fixed_m=use_fixed_m, fixed_m_norms=fixed_m_norms, + residual_checkpoint_name=residual_checkpoint_name, ) 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) @@ -1915,6 +1934,7 @@ def tokamax_ring_custom_kernel(q, k, v, context): context["dtype"], attention_kernel="tokamax_ring_custom", mask_padding_tokens=context["mask_padding_tokens"], + residual_checkpoint_name=context.get("residual_checkpoint_name"), attention_mask=context["attention_mask"], use_base2_exp=context.get("use_base2_exp", True), use_experimental_scheduler=context.get("use_experimental_scheduler", False), diff --git a/src/maxdiffusion/tests/custom_splash_backward_test.py b/src/maxdiffusion/tests/custom_splash_backward_test.py new file mode 100644 index 000000000..460bc70dc --- /dev/null +++ b/src/maxdiffusion/tests/custom_splash_backward_test.py @@ -0,0 +1,558 @@ +""" +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. +""" + +"""Unit tests for custom splash attention and custom ring attention backward pass.""" + +import contextlib +import functools +import io +import math +import unittest + +import jax +from jax import ad_checkpoint +import jax.numpy as jnp +import numpy as np + +from maxdiffusion.kernels import custom_splash_attention as custom_splash +from maxdiffusion.kernels.splash_attention import ring_attention_kernel + +_LOG2E = math.log2(math.e) +_LN2 = math.log(2.0) + + +def _reference_attention(q, k, v, actual_q_len, actual_kv_len, use_base2_exp): + """Reference attention in fp32 returning shape (Hq, Dv, actual_q_len).""" + q_f = q[:, :actual_q_len, :].astype(jnp.float32) + k_f = k[:, :actual_kv_len, :].astype(jnp.float32) + v_f = v[:, :actual_kv_len, :].astype(jnp.float32) + hq = q_f.shape[0] + hkv = k_f.shape[0] + q_per_kv = hq // hkv + k_f = jnp.repeat(k_f, q_per_kv, axis=0) + v_f = jnp.repeat(v_f, q_per_kv, axis=0) + logits = jnp.einsum("hsd,htd->hst", q_f, k_f) + if use_base2_exp: + logits = logits * _LN2 + probs = jax.nn.softmax(logits, axis=-1) + out = jnp.einsum("hst,htd->hds", probs, v_f) + return out.astype(q.dtype) + + +class CustomSplashBackwardTest(unittest.TestCase): + """Tests custom splash attention backward pass and ring attention backward pass.""" + + def setUp(self): + super().setUp() + if self._testMethodName != "test_remat_saves_attention_residuals" and jax.default_backend() != "tpu": + self.skipTest("Custom Pallas splash kernel requires TPU.") + self.block_sizes = custom_splash._BlockSizes( + block_q=1024, + block_kv=1024, + block_kv_compute=512, + block_kv_compute_in=256, + ) + + def test_remat_saves_attention_residuals(self): + """Verifies that ad_checkpoint saves out and lse under remat without recomputing them.""" + bsizes = custom_splash._BlockSizes( + block_q=128, + block_kv=128, + block_kv_compute=128, + block_kv_compute_in=128, + use_fused_bwd_kernel=True, + ) + attn_fn = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=128, + orig_kv_seq_len=128, + residual_checkpoint_name="attn_output", + ) + + def model_fn(q, k, v): + out = attn_fn(q, k, v) + return jnp.sum(out) + + chk_fn = jax.checkpoint( + model_fn, + policy=ad_checkpoint.checkpoint_policies.save_only_these_names("attn_output"), + ) + + q = jax.ShapeDtypeStruct((2, 128, 64), dtype=jnp.bfloat16) + k = jax.ShapeDtypeStruct((2, 128, 64), dtype=jnp.bfloat16) + v = jax.ShapeDtypeStruct((2, 128, 64), dtype=jnp.bfloat16) + + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + ad_checkpoint.print_saved_residuals(chk_fn, q, k, v) + output = buf.getvalue() + + self.assertIn("attn_output", output) + # Verify that both the bf16 output tensor and fp32 lse tensor are saved + self.assertIn("bf16[2,64,128]", output) + self.assertIn("f32[2,128]", output) + + def test_single_device_bwd_exact_multiples(self): + hq, hkv, sq, skv, d = 4, 4, 2048, 2048, 128 + scale = 1.0 / math.sqrt(d) + k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(0), 4) + q = (jax.random.normal(k1, (hq, sq, d), jnp.bfloat16) * scale * _LOG2E).astype(jnp.bfloat16) + k = jax.random.normal(k2, (hkv, skv, d), jnp.bfloat16) * scale + v = jax.random.normal(k3, (hkv, skv, d), jnp.bfloat16) * scale + do = jax.random.normal(k4, (hq, d, sq), jnp.bfloat16) + + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=sq, + orig_kv_seq_len=skv, + use_base2_exp=True, + ) + + out_kern, vjp_kern = jax.vjp(kernel, q, k, v) + dq_kern, dk_kern, dv_kern = vjp_kern(do) + + out_ref, vjp_ref = jax.vjp( + lambda q_, k_, v_: _reference_attention(q_, k_, v_, sq, skv, True), + q, + k, + v, + ) + dq_ref, dk_ref, dv_ref = vjp_ref(do) + + self.assertLess(float(jnp.max(jnp.abs(out_kern.astype(jnp.float32) - out_ref.astype(jnp.float32)))), 1e-3) + self.assertLess(float(jnp.max(jnp.abs(dq_kern.astype(jnp.float32) - dq_ref.astype(jnp.float32)))), 1e-3) + self.assertLess(float(jnp.max(jnp.abs(dk_kern.astype(jnp.float32) - dk_ref.astype(jnp.float32)))), 1e-3) + self.assertLess(float(jnp.max(jnp.abs(dv_kern.astype(jnp.float32) - dv_ref.astype(jnp.float32)))), 2e-3) + + def test_single_device_bwd_ragged_gqa_and_vmapped(self): + batch, hq, hkv, sq, skv, d = 2, 4, 2, 2048, 2048, 128 + actual_q_len, actual_kv_len = 1500, 1600 + scale = 1.0 / math.sqrt(d) + k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(1), 4) + q = (jax.random.normal(k1, (batch, hq, sq, d), jnp.bfloat16) * scale * _LOG2E).astype(jnp.bfloat16) + k = jax.random.normal(k2, (batch, hkv, skv, d), jnp.bfloat16) * scale + v = jax.random.normal(k3, (batch, hkv, skv, d), jnp.bfloat16) * scale + do = jax.random.normal(k4, (batch, hq, d, actual_q_len), jnp.bfloat16) + + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=actual_q_len, + orig_kv_seq_len=actual_kv_len, + use_base2_exp=True, + ) + vmapped_kernel = jax.vmap(kernel, in_axes=(0, 0, 0)) + + out_kern, vjp_kern = jax.vjp(vmapped_kernel, q, k, v) + dq_kern, dk_kern, dv_kern = vjp_kern(do) + + vmapped_ref = jax.vmap( + lambda q_, k_, v_: _reference_attention(q_, k_, v_, actual_q_len, actual_kv_len, True), + in_axes=(0, 0, 0), + ) + out_ref, vjp_ref = jax.vjp(vmapped_ref, q, k, v) + dq_ref, dk_ref, dv_ref = vjp_ref(do) + + self.assertLess(float(jnp.max(jnp.abs(out_kern.astype(jnp.float32) - out_ref.astype(jnp.float32)))), 1e-3) + self.assertLess( + float( + jnp.max( + jnp.abs(dq_kern[:, :, :actual_q_len].astype(jnp.float32) - dq_ref[:, :, :actual_q_len].astype(jnp.float32)) + ) + ), + 1e-3, + ) + self.assertEqual(float(jnp.max(jnp.abs(dq_kern[:, :, actual_q_len:].astype(jnp.float32)))), 0.0) + self.assertLess( + float( + jnp.max( + jnp.abs(dk_kern[:, :, :actual_kv_len].astype(jnp.float32) - dk_ref[:, :, :actual_kv_len].astype(jnp.float32)) + ) + ), + 1e-3, + ) + self.assertEqual(float(jnp.max(jnp.abs(dk_kern[:, :, actual_kv_len:].astype(jnp.float32)))), 0.0) + self.assertLess( + float( + jnp.max( + jnp.abs(dv_kern[:, :, :actual_kv_len].astype(jnp.float32) - dv_ref[:, :, :actual_kv_len].astype(jnp.float32)) + ) + ), + 2e-3, + ) + self.assertEqual(float(jnp.max(jnp.abs(dv_kern[:, :, actual_kv_len:].astype(jnp.float32)))), 0.0) + + def test_single_device_bwd_fused_aliasing_and_unfused(self): + hq, hkv, sq, skv, d = 4, 2, 2048, 4096, 128 + actual_q_len, actual_kv_len = 1500, 3500 # grid_width = ceil(3500/512) = 7 > 3 + scale = 1.0 / math.sqrt(d) + k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(42), 4) + q = (jax.random.normal(k1, (hq, sq, d), jnp.bfloat16) * scale * _LOG2E).astype(jnp.bfloat16) + k = jax.random.normal(k2, (hkv, skv, d), jnp.bfloat16) * scale + v = jax.random.normal(k3, (hkv, skv, d), jnp.bfloat16) * scale + do = jax.random.normal(k4, (hq, d, actual_q_len), jnp.bfloat16) + + out_ref, vjp_ref = jax.vjp( + lambda q_, k_, v_: _reference_attention(q_, k_, v_, actual_q_len, actual_kv_len, True), + q, + k, + v, + ) + dq_ref, dk_ref, dv_ref = vjp_ref(do) + + for use_fused, dq_red in [(True, 3), (True, None), (False, None)]: + bs = custom_splash._BlockSizes( + block_q=512, + block_kv=512, + block_kv_compute=256, + block_kv_compute_in=256, + use_fused_bwd_kernel=use_fused, + dq_reduction_steps=dq_red, + ) + kernel = custom_splash.make_splash_mha( + block_sizes=bs, + orig_q_seq_len=actual_q_len, + orig_kv_seq_len=actual_kv_len, + use_base2_exp=True, + ) + out_k, vjp_k = jax.vjp(kernel, q, k, v) + dq_k, dk_k, dv_k = vjp_k(do) + self.assertLess(float(jnp.max(jnp.abs(out_k.astype(jnp.float32) - out_ref.astype(jnp.float32)))), 1e-3) + self.assertLess( + float(jnp.max(jnp.abs(dq_k[:, :actual_q_len].astype(jnp.float32) - dq_ref[:, :actual_q_len].astype(jnp.float32)))), + 1e-2, + ) + self.assertEqual(float(jnp.max(jnp.abs(dq_k[:, actual_q_len:].astype(jnp.float32)))), 0.0) + self.assertLess( + float( + jnp.max(jnp.abs(dk_k[:, :actual_kv_len].astype(jnp.float32) - dk_ref[:, :actual_kv_len].astype(jnp.float32))) + ), + 1e-2, + ) + self.assertEqual(float(jnp.max(jnp.abs(dk_k[:, actual_kv_len:].astype(jnp.float32)))), 0.0) + self.assertLess( + float( + jnp.max(jnp.abs(dv_k[:, :actual_kv_len].astype(jnp.float32) - dv_ref[:, :actual_kv_len].astype(jnp.float32))) + ), + 1e-2, + ) + self.assertEqual(float(jnp.max(jnp.abs(dv_k[:, actual_kv_len:].astype(jnp.float32)))), 0.0) + + def test_ring_attention_bwd_ragged_gqa_vmapped(self): + devices = jax.devices() + if len(devices) < 4: + self.skipTest("Requires at least 4 TPU devices for ring attention test.") + ring_size = 4 + mesh = jax.sharding.Mesh(np.array(devices[:ring_size]), ("context",)) + batch, hq, hkv, sq_local, skv_local, d = 2, 4, 2, 1024, 1024, 128 + actual_q_local, actual_kv_local = 800, 900 + bsizes = custom_splash._BlockSizes( + block_q=512, + block_kv=512, + block_kv_compute=256, + block_kv_compute_in=256, + ) + + scale = 1.0 / math.sqrt(d) + k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(2), 4) + q_global = (jax.random.normal(k1, (ring_size, batch, hq, sq_local, d), jnp.bfloat16) * scale * _LOG2E).astype( + jnp.bfloat16 + ) + k_global = jax.random.normal(k2, (ring_size, batch, hkv, skv_local, d), jnp.bfloat16) * scale + v_global = jax.random.normal(k3, (ring_size, batch, hkv, skv_local, d), jnp.bfloat16) * scale + do_global = jax.random.normal(k4, (ring_size, batch, hq, actual_q_local, d), jnp.bfloat16) + + def _ref_single_batch(qg, kg, vg): + q_list = [qg[r, :, :actual_q_local, :].astype(jnp.float32) for r in range(ring_size)] + k_list = [kg[r, :, :actual_kv_local, :].astype(jnp.float32) for r in range(ring_size)] + v_list = [vg[r, :, :actual_kv_local, :].astype(jnp.float32) for r in range(ring_size)] + q_cat = jnp.concatenate(q_list, axis=1) + k_cat = jnp.concatenate(k_list, axis=1) + v_cat = jnp.concatenate(v_list, axis=1) + k_cat = jnp.repeat(k_cat, hq // hkv, axis=0) + v_cat = jnp.repeat(v_cat, hq // hkv, axis=0) + logits = jnp.einsum("hsd,htd->hst", q_cat, k_cat) * _LN2 + probs = jax.nn.softmax(logits, axis=-1) + out_cat = jnp.einsum("hst,htd->hsd", probs, v_cat) + return jnp.stack( + [out_cat[:, r * actual_q_local : (r + 1) * actual_q_local, :] for r in range(ring_size)], + axis=0, + ).astype(qg.dtype) + + ref_fn = jax.vmap(_ref_single_batch, in_axes=(1, 1, 1), out_axes=1) + out_ref, vjp_ref = jax.vjp(ref_fn, q_global, k_global, v_global) + dq_ref, dk_ref, dv_ref = vjp_ref(do_global) + + for use_fused in [True, False]: + bsizes = custom_splash._BlockSizes( + block_q=512, + block_kv=512, + block_kv_compute=256, + block_kv_compute_in=256, + use_fused_bwd_kernel=use_fused, + ) + ring_kernel = ring_attention_kernel.make_custom_ring_attention( + block_sizes=bsizes, + orig_q_seq_len=actual_q_local, + orig_kv_seq_len=actual_kv_local, + use_base2_exp=True, + ring_axis="context", + ) + vmapped_ring = jax.vmap(ring_kernel, in_axes=(0, 0, 0)) + + p = jax.sharding.PartitionSpec("context", None, None, None, None) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(p, p, p, p), + out_specs=(p, p, p, p), + check_vma=False, + ) + def run_ring(q_sh, k_sh, v_sh, do_sh): + q_s, k_s, v_s, do_s = q_sh[0], k_sh[0], v_sh[0], do_sh[0] + out_s, vjp_s = jax.vjp(vmapped_ring, q_s, k_s, v_s) + dq_s, dk_s, dv_s = vjp_s(do_s) + return out_s[None], dq_s[None], dk_s[None], dv_s[None] + + out_ring, dq_ring, dk_ring, dv_ring = run_ring(q_global, k_global, v_global, do_global) + + self.assertLess(float(jnp.max(jnp.abs(out_ring.astype(jnp.float32) - out_ref.astype(jnp.float32)))), 1e-3) + self.assertLess( + float( + jnp.max( + jnp.abs( + dq_ring[:, :, :, :actual_q_local].astype(jnp.float32) + - dq_ref[:, :, :, :actual_q_local].astype(jnp.float32) + ) + ) + ), + 1e-3, + ) + self.assertEqual(float(jnp.max(jnp.abs(dq_ring[:, :, :, actual_q_local:].astype(jnp.float32)))), 0.0) + self.assertLess( + float( + jnp.max( + jnp.abs( + dk_ring[:, :, :, :actual_kv_local].astype(jnp.float32) + - dk_ref[:, :, :, :actual_kv_local].astype(jnp.float32) + ) + ) + ), + 1e-3, + ) + self.assertEqual(float(jnp.max(jnp.abs(dk_ring[:, :, :, actual_kv_local:].astype(jnp.float32)))), 0.0) + self.assertLess( + float( + jnp.max( + jnp.abs( + dv_ring[:, :, :, :actual_kv_local].astype(jnp.float32) + - dv_ref[:, :, :, :actual_kv_local].astype(jnp.float32) + ) + ) + ), + 2e-3, + ) + self.assertEqual(float(jnp.max(jnp.abs(dv_ring[:, :, :, actual_kv_local:].astype(jnp.float32)))), 0.0) + + def test_single_device_bwd_unpadded_non_divisible_seqlens(self): + """Tests directly passing unpadded Q/K/V whose sequence lengths are not divisible by block size.""" + hq, hkv, sq, skv, d = 4, 2, 1350, 1777, 128 + scale = 1.0 / math.sqrt(d) + k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(101), 4) + q = (jax.random.normal(k1, (hq, sq, d), jnp.bfloat16) * scale * _LOG2E).astype(jnp.bfloat16) + k = jax.random.normal(k2, (hkv, skv, d), jnp.bfloat16) * scale + v = jax.random.normal(k3, (hkv, skv, d), jnp.bfloat16) * scale + do = jax.random.normal(k4, (hq, d, sq), jnp.bfloat16) + + out_ref, vjp_ref = jax.vjp( + lambda q_, k_, v_: _reference_attention(q_, k_, v_, sq, skv, True), + q, + k, + v, + ) + dq_ref, dk_ref, dv_ref = vjp_ref(do) + + for use_fused, dq_red in [(True, 3), (True, None), (False, None)]: + bs = custom_splash._BlockSizes( + block_q=512, + block_kv=512, + block_kv_compute=256, + block_kv_compute_in=128, + use_fused_bwd_kernel=use_fused, + dq_reduction_steps=dq_red, + ) + kernel = custom_splash.make_splash_mha( + block_sizes=bs, + orig_q_seq_len=sq, + orig_kv_seq_len=skv, + use_base2_exp=True, + ) + out_k, vjp_k = jax.vjp(kernel, q, k, v) + dq_k, dk_k, dv_k = vjp_k(do) + + self.assertEqual(out_k.shape, (hq, d, sq)) + self.assertEqual(dq_k.shape, (hq, sq, d)) + self.assertEqual(dk_k.shape, (hkv, skv, d)) + self.assertEqual(dv_k.shape, (hkv, skv, d)) + + self.assertLess(float(jnp.max(jnp.abs(out_k.astype(jnp.float32) - out_ref.astype(jnp.float32)))), 1e-3) + self.assertLess(float(jnp.max(jnp.abs(dq_k.astype(jnp.float32) - dq_ref.astype(jnp.float32)))), 1e-2) + self.assertLess(float(jnp.max(jnp.abs(dk_k.astype(jnp.float32) - dk_ref.astype(jnp.float32)))), 1e-2) + self.assertLess(float(jnp.max(jnp.abs(dv_k.astype(jnp.float32) - dv_ref.astype(jnp.float32)))), 1e-2) + + def test_single_device_bwd_padded_with_garbage_in_kv_tail(self): + """Verifies that non-divisible orig_kv_seq_len ignores extreme garbage values in padded KV tail.""" + hq, hkv, padded_sq, padded_skv, d = 4, 2, 2048, 2048, 128 + actual_q_len, actual_kv_len = 1350, 1777 + scale = 1.0 / math.sqrt(d) + k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(102), 4) + q = (jax.random.normal(k1, (hq, padded_sq, d), jnp.bfloat16) * scale * _LOG2E).astype(jnp.bfloat16) + k = jax.random.normal(k2, (hkv, padded_skv, d), jnp.bfloat16) * scale + v = jax.random.normal(k3, (hkv, padded_skv, d), jnp.bfloat16) * scale + # Inject huge garbage numbers into the KV padding tail [actual_kv_len:] + k = k.at[:, actual_kv_len:, :].set(1000.0) + v = v.at[:, actual_kv_len:, :].set(1000.0) + do = jax.random.normal(k4, (hq, d, actual_q_len), jnp.bfloat16) + + out_ref, vjp_ref = jax.vjp( + lambda q_, k_, v_: _reference_attention(q_, k_, v_, actual_q_len, actual_kv_len, True), + q, + k, + v, + ) + dq_ref, dk_ref, dv_ref = vjp_ref(do) + + for use_fused in [True, False]: + bs = custom_splash._BlockSizes( + block_q=512, + block_kv=512, + block_kv_compute=256, + block_kv_compute_in=128, + use_fused_bwd_kernel=use_fused, + ) + kernel = custom_splash.make_splash_mha( + block_sizes=bs, + orig_q_seq_len=actual_q_len, + orig_kv_seq_len=actual_kv_len, + use_base2_exp=True, + ) + out_k, vjp_k = jax.vjp(kernel, q, k, v) + dq_k, dk_k, dv_k = vjp_k(do) + + self.assertLess(float(jnp.max(jnp.abs(out_k.astype(jnp.float32) - out_ref.astype(jnp.float32)))), 1e-3) + self.assertLess( + float(jnp.max(jnp.abs(dq_k[:, :actual_q_len].astype(jnp.float32) - dq_ref[:, :actual_q_len].astype(jnp.float32)))), + 1e-2, + ) + self.assertEqual(float(jnp.max(jnp.abs(dq_k[:, actual_q_len:].astype(jnp.float32)))), 0.0) + self.assertLess( + float( + jnp.max(jnp.abs(dk_k[:, :actual_kv_len].astype(jnp.float32) - dk_ref[:, :actual_kv_len].astype(jnp.float32))) + ), + 1e-2, + ) + self.assertEqual(float(jnp.max(jnp.abs(dk_k[:, actual_kv_len:].astype(jnp.float32)))), 0.0) + self.assertLess( + float( + jnp.max(jnp.abs(dv_k[:, :actual_kv_len].astype(jnp.float32) - dv_ref[:, :actual_kv_len].astype(jnp.float32))) + ), + 1e-2, + ) + self.assertEqual(float(jnp.max(jnp.abs(dv_k[:, actual_kv_len:].astype(jnp.float32)))), 0.0) + + def test_ring_attention_bwd_unpadded_non_divisible_seqlens(self): + """Tests 4-device ring attention with unpadded local sequence lengths not divisible by block sizes.""" + devices = jax.devices() + if len(devices) < 4: + self.skipTest("Requires at least 4 TPU devices for ring attention test.") + ring_size = 4 + mesh = jax.sharding.Mesh(np.array(devices[:ring_size]), ("context",)) + batch, hq, hkv, sq_local, skv_local, d = 2, 4, 2, 650, 733, 128 + + scale = 1.0 / math.sqrt(d) + k1, k2, k3, k4 = jax.random.split(jax.random.PRNGKey(103), 4) + q_global = (jax.random.normal(k1, (ring_size, batch, hq, sq_local, d), jnp.bfloat16) * scale * _LOG2E).astype( + jnp.bfloat16 + ) + k_global = jax.random.normal(k2, (ring_size, batch, hkv, skv_local, d), jnp.bfloat16) * scale + v_global = jax.random.normal(k3, (ring_size, batch, hkv, skv_local, d), jnp.bfloat16) * scale + do_global = jax.random.normal(k4, (ring_size, batch, hq, sq_local, d), jnp.bfloat16) + + def _ref_single_batch(qg, kg, vg): + q_cat = jnp.concatenate([qg[r].astype(jnp.float32) for r in range(ring_size)], axis=1) + k_cat = jnp.concatenate([kg[r].astype(jnp.float32) for r in range(ring_size)], axis=1) + v_cat = jnp.concatenate([vg[r].astype(jnp.float32) for r in range(ring_size)], axis=1) + k_cat = jnp.repeat(k_cat, hq // hkv, axis=0) + v_cat = jnp.repeat(v_cat, hq // hkv, axis=0) + logits = jnp.einsum("hsd,htd->hst", q_cat, k_cat) * _LN2 + probs = jax.nn.softmax(logits, axis=-1) + out_cat = jnp.einsum("hst,htd->hsd", probs, v_cat) + return jnp.stack( + [out_cat[:, r * sq_local : (r + 1) * sq_local, :] for r in range(ring_size)], + axis=0, + ).astype(qg.dtype) + + ref_fn = jax.vmap(_ref_single_batch, in_axes=(1, 1, 1), out_axes=1) + out_ref, vjp_ref = jax.vjp(ref_fn, q_global, k_global, v_global) + dq_ref, dk_ref, dv_ref = vjp_ref(do_global) + + for use_fused in [True, False]: + bsizes = custom_splash._BlockSizes( + block_q=512, + block_kv=512, + block_kv_compute=256, + block_kv_compute_in=128, + use_fused_bwd_kernel=use_fused, + ) + ring_kernel = ring_attention_kernel.make_custom_ring_attention( + block_sizes=bsizes, + orig_q_seq_len=sq_local, + orig_kv_seq_len=skv_local, + use_base2_exp=True, + ring_axis="context", + ) + vmapped_ring = jax.vmap(ring_kernel, in_axes=(0, 0, 0)) + + p = jax.sharding.PartitionSpec("context", None, None, None, None) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(p, p, p, p), + out_specs=(p, p, p, p), + check_vma=False, + ) + def run_ring(q_sh, k_sh, v_sh, do_sh): + q_s, k_s, v_s, do_s = q_sh[0], k_sh[0], v_sh[0], do_sh[0] + out_s, vjp_s = jax.vjp(vmapped_ring, q_s, k_s, v_s) + dq_s, dk_s, dv_s = vjp_s(do_s) + return out_s[None], dq_s[None], dk_s[None], dv_s[None] + + out_ring, dq_ring, dk_ring, dv_ring = run_ring(q_global, k_global, v_global, do_global) + + self.assertEqual(out_ring.shape, (ring_size, batch, hq, sq_local, d)) + self.assertEqual(dq_ring.shape, (ring_size, batch, hq, sq_local, d)) + self.assertEqual(dk_ring.shape, (ring_size, batch, hkv, skv_local, d)) + self.assertEqual(dv_ring.shape, (ring_size, batch, hkv, skv_local, d)) + + self.assertLess(float(jnp.max(jnp.abs(out_ring.astype(jnp.float32) - out_ref.astype(jnp.float32)))), 1e-3) + self.assertLess(float(jnp.max(jnp.abs(dq_ring.astype(jnp.float32) - dq_ref.astype(jnp.float32)))), 1e-3) + self.assertLess(float(jnp.max(jnp.abs(dk_ring.astype(jnp.float32) - dk_ref.astype(jnp.float32)))), 1e-3) + self.assertLess(float(jnp.max(jnp.abs(dv_ring.astype(jnp.float32) - dv_ref.astype(jnp.float32)))), 2e-3) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/trainers/wan_trainer_2_2.py b/src/maxdiffusion/trainers/wan_trainer_2_2.py index 13650a0d9..6ea105114 100644 --- a/src/maxdiffusion/trainers/wan_trainer_2_2.py +++ b/src/maxdiffusion/trainers/wan_trainer_2_2.py @@ -16,6 +16,7 @@ import functools import datetime +import random import threading from concurrent.futures import ThreadPoolExecutor import numpy as np @@ -24,9 +25,9 @@ from flax.linen import partitioning as nn_partitioning import jax.numpy as jnp import jax +from jax.experimental import multihost_utils from jax.sharding import PartitionSpec as P import jaxopt -from jax.experimental import multihost_utils from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2 from maxdiffusion.schedulers import FlaxFlowMatchScheduler @@ -387,11 +388,12 @@ def training_loop_2_2( start_step_time = datetime.datetime.now() next_batch_future = executor.submit(load_next_batch, train_data_iterator, example_batch, self.config) - # Determine active expert on host and synchronize with device - step_rng, cond_rng = jax.random.split(step_rng, num=2) + # Determine active expert deterministically on host without device-to-host sync stall + host_rng = random.Random(self.config.seed + step_int) shift = getattr(pipeline.scheduler.config, "shift", 1.0) or 1.0 u_boundary = compute_unshifted_boundary(self.config.boundary_ratio, shift) - is_high_noise = bool(jax.random.uniform(cond_rng) >= u_boundary) + is_high_noise = host_rng.random() >= u_boundary + is_high_noise_dev = jnp.bool_(is_high_noise) if is_high_noise: active_lr = float(learning_rate_scheduler_high(host_step_high)) else: @@ -403,7 +405,7 @@ def training_loop_2_2( nn_partitioning.axis_rules(self.config.logical_axis_rules), ): state_low, state_high, scheduler_state, train_metric, _ = p_train_step( - state_low, state_high, example_batch, step_rng, scheduler_state, jnp.bool_(is_high_noise) + state_low, state_high, example_batch, step_rng, scheduler_state, is_high_noise_dev ) train_metric["scalar"]["learning/loss"].block_until_ready() last_step_completion = datetime.datetime.now() @@ -522,7 +524,7 @@ def eval_2_2(self, mesh, eval_rng_key, step, p_eval_step, state_low, state_high, writer.add_scalar("learning/eval_loss", final_eval_loss, step) -def train_step_2_2(state_low, state_high, data, rng, scheduler_state, is_high_noise, scheduler, config): +def train_step_2_2(state_low, state_high, data, rng, scheduler_state, is_high_noise=None, scheduler=None, config=None): """Wan 2.2 joint dual-expert training step. Expert Routing: @@ -540,8 +542,6 @@ def train_step_2_2(state_low, state_high, data, rng, scheduler_state, is_high_no applying Flow Match time shift S(u) so that the resulting timesteps cover [0, boundary_ratio * num_train_timesteps], matching the low-noise regime used during inference. """ - _, new_rng, timestep_rng, dropout_rng, noise_rng = jax.random.split(rng, num=5) - data = {k: v[: config.global_batch_size_to_train_on] for k, v in data.items()} bsz = data["latents"].shape[0] @@ -549,6 +549,11 @@ def train_step_2_2(state_low, state_high, data, rng, scheduler_state, is_high_no shift = getattr(scheduler.config, "shift", 1.0) or 1.0 u_boundary = compute_unshifted_boundary(config.boundary_ratio, shift) + if is_high_noise is None: + _, new_rng, timestep_rng, dropout_rng, cond_rng, noise_rng = jax.random.split(rng, num=6) + is_high_noise = jax.random.uniform(cond_rng) >= u_boundary + else: + _, new_rng, timestep_rng, dropout_rng, noise_rng = jax.random.split(rng, num=5) def compute_loss_high(high_params, s_high): u = jax.random.uniform(timestep_rng, (bsz,), minval=u_boundary, maxval=1.0) t_shifted = (u * shift) / (1.0 + (shift - 1.0) * u)