diff --git a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml index bd3828825d..bd2d4d1bf3 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml @@ -63,6 +63,9 @@ logical_axis_rules: [ # MoE Weights ['exp', ['context', 'expert']], ['embed_moe', ['fsdp']], + # wo's dims: FSDP on its mlp dim (see RoutedMoE.wo_kernel_axes) + ['mlp_moe_out', ['fsdp']], + ['embed_moe_out', []], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index ba872c9a4b..fca14624ed 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1523,12 +1523,13 @@ class Qwen3Next(BaseModel): "highest", description=( "Precision of the f32 matmuls that compose the GDN recurrent state across sequence-sharded context" - " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3." - " Only the XLA composition math changes; the Pallas kernels keep their own precision. It applies only" - " with use_gdn_kernel=True and sequence-sharded GDN context parallelism: gdn_cp_mode='seq', or 'auto'," - " which shards the sequence unless the GDN context parallelism is 2 and divides the number of key" - " heads. It has no effect with use_gdn_kernel=False (that path's context parallelism always uses" - " HIGHEST) or with head-sharded context parallelism." + " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3;" + " with 'high' the rank-local forward transition product and backward state gradient run as Pallas scans" + " (kernels/gdn/gdn_bwd/dh_scan.py) at the same 3 passes. Only the composition math changes; the Pallas GDN" + " kernels keep their own precision. It applies only with use_gdn_kernel=True and sequence-sharded GDN" + " context parallelism: gdn_cp_mode='seq', or 'auto', which shards the sequence unless the GDN context" + " parallelism is 2 and divides the number of key heads. It has no effect with use_gdn_kernel=False (that" + " path's context parallelism always uses HIGHEST) or with head-sharded context parallelism." ), ) diff --git a/src/maxtext/kernels/gdn/compute_gdn.py b/src/maxtext/kernels/gdn/compute_gdn.py index 36d8c7bf62..45a09d5eff 100644 --- a/src/maxtext/kernels/gdn/compute_gdn.py +++ b/src/maxtext/kernels/gdn/compute_gdn.py @@ -51,18 +51,24 @@ def invert_triangular_matrix( num_blocks = chunk // block_size def local_forward_sub(t_mat: jax.Array, b_mat: jax.Array) -> jax.Array: - x_list = [] - for i in range(block_size): - b_i = b_mat[:, i, :] - if i == 0: - x_i = b_i + # Right-looking substitution: once row j is final, subtract t[i, j] * x_j from every later row i. + # The rows stay in place as two sublane-aligned halves, so no step re-stacks the solved rows or + # moves a row of t from lanes to sublanes. Each row subtracts the same products as a left-looking + # solve, one at a time in column order. + half = block_size // 2 + x_lo, x_hi = b_mat[:, :half, :], b_mat[:, half:, :] + t_lo, t_hi = t_mat[:, :half, :], t_mat[:, half:, :] + rows_lo = jax.lax.broadcasted_iota(jnp.int32, (1, half, 1), 1) + rows_hi = jax.lax.broadcasted_iota(jnp.int32, (1, block_size - half, 1), 1) + for j in range(block_size - 1): + if j < half: + x_j = x_lo[:, j : j + 1, :] + x_lo = jnp.where(rows_lo > j, x_lo - t_lo[:, :, j : j + 1] * x_j, x_lo) + x_hi = x_hi - t_hi[:, :, j : j + 1] * x_j else: - stacked_x = jnp.stack(x_list, axis=1) - all_prev_t = t_mat[:, i, :i] - prev_sum = jnp.sum(all_prev_t[..., None] * stacked_x, axis=1) - x_i = b_i - prev_sum - x_list.append(x_i) - return jnp.stack(x_list, axis=1) + x_j = x_hi[:, j - half : j - half + 1, :] + x_hi = jnp.where(rows_hi > j - half, x_hi - t_hi[:, :, j : j + 1] * x_j, x_hi) + return jnp.concatenate([x_lo, x_hi], axis=1) x_blocks = [] iota_r = jax.lax.broadcasted_iota(jnp.int32, t.shape, 1) @@ -114,8 +120,8 @@ def chunked_gdn_per_seq( cfg: config.GDNConfig, seg_c_col: jax.Array | None = None, # [1, chunk, 1] seg_aux_col: jax.Array | None = None, # [1, chunk, 1] -) -> tuple[jax.Array, jax.Array, jax.Array]: - """Perform chunked GDN over input [num_heads, chunk, head_dim].""" +) -> tuple[jax.Array | None, jax.Array, jax.Array]: + """Perform chunked GDN over input [num_heads, chunk, head_dim]; out is None in PER_SEQ_STATES mode.""" # bf16 activations keep single-pass MXU matmuls even with an f32 compute dtype # (HIGHEST would triple the forward MXU cost); HIGHEST needs f32 activations and compute. dot_prec = ( @@ -236,23 +242,36 @@ def chunked_gdn_per_seq( # [num_v_heads, chunk, v_head_dim] u, w = jnp.split(merged_uw, [cfg.v_head_dim], axis=-1) - # [num_v_heads, chunk, kq_head_dim] - q_large_gating = q_repeat * gating_forward - # NOTE: Concatenate lhs with same rhs to leverage weight - # stationary architecture. - # [num_v_heads, 2 * chunk, kq_head_dim] - merged_w_q = jnp.concat([w, q_large_gating], axis=1) - # [num_v_heads, 2 * chunk, v_head_dim] - merged_ws_out_updated = jax.lax.dot( - merged_w_q, - state_prev, - dimension_numbers=(((2,), (1,)), ((0,), (0,))), - precision=dot_prec, - preferred_element_type=jnp.float32, - ) + states_only = cfg.mode == config.GDNMode.PER_SEQ_STATES + if states_only: + # No output: only the w rows of the merged matmul below. Each row's dot product is the same, + # so ws is bit-identical. + # [num_v_heads, chunk, v_head_dim] + ws = jax.lax.dot( + w, + state_prev, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + precision=dot_prec, + preferred_element_type=jnp.float32, + ) + else: + # [num_v_heads, chunk, kq_head_dim] + q_large_gating = q_repeat * gating_forward + # NOTE: Concatenate lhs with same rhs to leverage weight + # stationary architecture. + # [num_v_heads, 2 * chunk, kq_head_dim] + merged_w_q = jnp.concat([w, q_large_gating], axis=1) + # [num_v_heads, 2 * chunk, v_head_dim] + merged_ws_out_updated = jax.lax.dot( + merged_w_q, + state_prev, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + precision=dot_prec, + preferred_element_type=jnp.float32, + ) - # NOTE: Splitting along non sublane/lane dim is free. - ws, out_updated = jnp.split(merged_ws_out_updated, 2, axis=1) + # NOTE: Splitting along non sublane/lane dim is free. + ws, out_updated = jnp.split(merged_ws_out_updated, 2, axis=1) ws = ws.astype(cfg.dtypes.compute) # [num_v_heads, chunk, v_head_dim] @@ -273,6 +292,8 @@ def chunked_gdn_per_seq( # [num_v_heads, kq_head_dim, v_head_dim] state_updated = state_prev * gating_last state = state_updated + state_new + if states_only: + return None, state, t_inv # [num_kq_heads, chunk, chunk] out_qk = jax.lax.dot( @@ -314,8 +335,8 @@ def chunked_gdn( a_log: jax.Array, dt_bias: jax.Array, cfg: config.GDNConfig, -) -> tuple[jax.Array, jax.Array, jax.Array]: - """Perform chunked GDN over input [seq, num_heads, chunk, head_dim].""" +) -> tuple[jax.Array | None, jax.Array, jax.Array]: + """Perform chunked GDN over input [seq, num_heads, chunk, head_dim]; out is None in PER_SEQ_STATES mode.""" mask_dtype = get_mask_dtype(cfg.dtypes.compute) iota = jax.lax.broadcasted_iota(mask_dtype, (cfg.seq_tile_size, 1, cfg.chunk_size, 1), 2) @@ -373,10 +394,11 @@ def chunked_gdn( seg_c_col=seg_c_per_seq[idx] if cfg.has_seg_ids else None, seg_aux_col=seg_aux_per_seq[idx] if cfg.has_seg_ids else None, ) - out_list.append(out.swapaxes(0, 1)) + if out is not None: + out_list.append(out.swapaxes(0, 1)) state_list.append(state) t_inv_list.append(t_inv) - out = jnp.stack(out_list, axis=0) + out = jnp.stack(out_list, axis=0) if out_list else None state = jnp.stack(state_list, axis=0) t_inv = jnp.stack(t_inv_list, axis=0) return out, state, t_inv diff --git a/src/maxtext/kernels/gdn/config.py b/src/maxtext/kernels/gdn/config.py index 8ad3933295..b06926d24c 100644 --- a/src/maxtext/kernels/gdn/config.py +++ b/src/maxtext/kernels/gdn/config.py @@ -24,7 +24,8 @@ from jax.experimental.pallas import tpu as pltpu import jax.numpy as jnp -DEFAULT_VMEM_LIMIT_FACTOR: float = 0.90 +# The 64-v-head prefill kernel needs about 58 MiB of the 64 MiB VMEM on TPU7x. +DEFAULT_VMEM_LIMIT_FACTOR: float = 0.95 class GDNMode(enum.StrEnum): @@ -32,6 +33,8 @@ class GDNMode(enum.StrEnum): BATCHED = enum.auto() PER_SEQ = enum.auto() + # PER_SEQ that only carries the conv/recurrent states and t_inv: no output, no chunk states. + PER_SEQ_STATES = enum.auto() def get_seq_tile_size(self, tile_size: int) -> int: if self == GDNMode.BATCHED: @@ -82,6 +85,8 @@ class GDNConfig: num_buffers: int = 2 has_seg_ids: bool = False use_qk_norm_in_gdn: bool = True + # Prefill-only tiles are whole chunks: qkv is DMA'd in act_in dtype and upcast in VMEM. + is_prefill_only: bool = False @property def chunk_size(self) -> int: @@ -150,7 +155,13 @@ def get_scratch_shape_dict(self) -> dict[str, Any]: carry_conv_scratch = pltpu.VMEM(conv_shape, jnp.float32) carry_recurrent_scratch = pltpu.VMEM(recurrent_shape, jnp.float32) + # Compact fp32 copy of a prefill qkv chunk, read by Conv1D and the q/k/v loads. + qkv_scratch = None + if self.is_prefill_only: + qkv_scratch = pltpu.VMEM((self.seq_tile_size, self.chunk_size, 1, self.dim_size), jnp.float32) + return { "carry_conv_scratch_ref": carry_conv_scratch, "carry_recurrent_scratch_ref": carry_recurrent_scratch, + "qkv_scratch_ref": qkv_scratch, } diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py index 4ea358228b..e56aa53cd8 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/api.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py @@ -25,6 +25,7 @@ from .. import compute_conv1d as local_compute_conv1d from .. import wrapper as local_gdn_wrapper from . import cp_gdn +from . import dh_scan from .compute_conv1d_bwd import conv1d_silu_bwd from .compute_conv1d_bwd import conv1d_silu_fwd from .jax_compute_gdn_states import _compute_forward_conv_and_states @@ -130,6 +131,50 @@ def decoupled_conv1d_gdn_bwd_kernel( ) +def _fused_kernel_inputs( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, +) -> dict[str, jax.Array]: + """Lays out [batch, seq, ...] inputs for the fused Conv1D-GDN kernel: one equal-length prefill sequence per row.""" + num_seqs, seq_len, dim_size = qkv.shape + has_init = (conv_state is not None) or (recurrent_state is not None) + + if conv_state is None: + tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype) + elif conv_state.shape[0] == num_seqs: + tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0))) + else: + tokamax_conv_state = conv_state + + if recurrent_state is None: + tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) + elif recurrent_state.shape[0] == num_seqs: + tokamax_recurrent_state = jnp.pad(recurrent_state.astype(jnp.float32), ((1, 0), (0, 0), (0, 0), (0, 0))) + else: + tokamax_recurrent_state = recurrent_state.astype(jnp.float32) + + return { + "qkv": qkv.reshape(-1, dim_size), + "b": b.reshape(-1, b.shape[-1]), + "a": a.reshape(-1, a.shape[-1]), + "conv_state": tokamax_conv_state, + "recurrent_state": tokamax_recurrent_state, + "conv_weight": jnp.swapaxes(conv_weight, 0, 2), + "query_start_loc": jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32), + "state_indices": jnp.arange(1, num_seqs + 1, dtype=jnp.int32), + "distribution": jnp.array([0, 0, num_seqs], dtype=jnp.int32), + "seq_lens": jnp.full((num_seqs,), seq_len + (1 if has_init else 0), dtype=jnp.int32), + } + + def _run_local_gdn_decoupled_fwd( qkv: jax.Array, b: jax.Array, @@ -158,7 +203,7 @@ def _run_local_gdn_decoupled_fwd( Optional[jax.Array], ]: """Runs local GDN forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU.""" - batch_size, seq_len, dim_size = qkv.shape + batch_size, seq_len, _ = qkv.shape if seq_len % chunk_size != 0: raise ValueError( f"GDN forward kernel requires the local sequence length ({seq_len}) to be a multiple of" @@ -215,58 +260,27 @@ def _run_local_gdn_decoupled_fwd( ) return (out, states), t_inv, chunk_states - num_seqs = batch_size num_chunks = seq_len // chunk_size - - qkv_flat = qkv.reshape(-1, dim_size) - b_flat = b.reshape(-1, b.shape[-1]) - a_flat = a.reshape(-1, a.shape[-1]) - tokamax_conv_weight = jnp.swapaxes(conv_weight, 0, 2) - - has_init = (conv_state is not None) or (recurrent_state is not None) - query_start_loc = jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32) - state_indices = jnp.arange(1, num_seqs + 1, dtype=jnp.int32) - seq_lens = jnp.full((num_seqs,), seq_len + (1 if has_init else 0), dtype=jnp.int32) - distribution = jnp.array([0, 0, num_seqs], dtype=jnp.int32) - - if conv_state is None: - tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype) - elif conv_state.shape[0] == num_seqs: - tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0))) - else: - tokamax_conv_state = conv_state - - if recurrent_state is None: - tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) - elif recurrent_state.shape[0] == num_seqs: - tokamax_recurrent_state = jnp.pad(recurrent_state.astype(jnp.float32), ((1, 0), (0, 0), (0, 0), (0, 0))) - else: - tokamax_recurrent_state = recurrent_state.astype(jnp.float32) - ( core_attn_out_flat, (new_conv_state, new_recurrent_state), t_inv_raw, chunk_states_raw, ) = local_gdn_wrapper.fused_conv1d_gdn( - qkv=qkv_flat, - b=b_flat, - a=a_flat, - conv_state=tokamax_conv_state, - recurrent_state=tokamax_recurrent_state, - conv_weight=tokamax_conv_weight, + **_fused_kernel_inputs( + qkv, b, a, conv_weight, conv_state, recurrent_state, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size + ), conv_bias=conv_bias, a_log=a_log, dt_bias=dt_bias, - query_start_loc=query_start_loc, - state_indices=state_indices, - distribution=distribution, - seq_lens=seq_lens, n_kq=num_k_heads, n_v=num_v_heads, d_k=head_k_dim, d_v=head_v_dim, kernel_size=conv_kernel_size, + # seq_len is a multiple of chunk_size (checked above), so the prefill kernel writes every + # output row and the output buffer needs no zero fill. + zero_initialize_out=False, compute_precision=jnp.dtype(compute_dtype), mixed_tile_size=chunk_size, is_prefill_only=True, @@ -292,6 +306,90 @@ def _run_local_gdn_decoupled_fwd( ) +def _run_local_gdn_states_fwd( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, + segment_ids: Optional[jax.Array] = None, + conv_halo_seg: Optional[jax.Array] = None, + init_seg: Optional[jax.Array] = None, +) -> Tuple[Tuple[jax.Array, jax.Array], jax.Array]: + """_run_local_gdn_decoupled_fwd for callers that need only (states, t_inv), bit-identical to it. + + On TPU the kernel skips the output and the chunk states; on CPU XLA drops the unused pure-JAX outputs. + """ + if target_platform(qkv) == "cpu": + (_, states), t_inv, _ = _run_local_gdn_decoupled_fwd( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + segment_ids=segment_ids, + conv_halo_seg=conv_halo_seg, + init_seg=init_seg, + ) + return states, t_inv + batch_size, seq_len, _ = qkv.shape + if seq_len % chunk_size != 0: + raise ValueError( + f"GDN forward kernel requires the local sequence length ({seq_len}) to be a multiple of" + f" chunk_size ({chunk_size}); with sequence-sharded context parallelism this is seq_len / cp." + ) + fallback_reason = pallas_unsupported_reason(head_k_dim=head_k_dim, head_v_dim=head_v_dim, chunk_size=chunk_size) + if fallback_reason is not None: + raise ValueError(f"GDN Pallas TPU forward kernel does not support {fallback_reason}.") + + (new_conv_state, new_recurrent_state), t_inv_raw = local_gdn_wrapper.fused_conv1d_gdn_states( + **_fused_kernel_inputs( + qkv, b, a, conv_weight, conv_state, recurrent_state, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size + ), + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + n_kq=num_k_heads, + n_v=num_v_heads, + d_k=head_k_dim, + d_v=head_v_dim, + kernel_size=conv_kernel_size, + compute_precision=jnp.dtype(compute_dtype), + mixed_tile_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + segment_ids=segment_ids, + conv_halo_seg=conv_halo_seg, + init_seg=init_seg, + ) + t_inv = t_inv_raw.astype(jnp.float32).reshape(batch_size, seq_len // chunk_size, num_v_heads, chunk_size, chunk_size) + return (new_conv_state[1:].astype(qkv.dtype), new_recurrent_state[1:].astype(jnp.float32)), t_inv + + def _is_cp_active(cp_axis_name: str | tuple[str, ...] | None) -> bool: if cp_axis_name is None: return False @@ -377,8 +475,9 @@ def _run_cp_gdn_decoupled_fwd_impl( conv_halo = checkpoint_name(conv_halo, "gdn_cp_conv_halo") zero_rs = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) - # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv and S_ext_local - (_, states_pass1), t_inv, _ = _run_local_gdn_decoupled_fwd( + # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv and S_ext_local. Its output + # is not needed, so the kernel neither computes it nor writes it or the chunk states. + states_pass1, t_inv = _run_local_gdn_states_fwd( qkv, b, a, @@ -414,24 +513,43 @@ def _run_cp_gdn_decoupled_fwd_impl( conv_halo_seg=conv_halo_seg, output_dtype=jnp.float32, ) - m_local, s_ext_local = cp_gdn.compose_local_from_t_inv( - qkv_conv=k_conv, - b=b, - a=a, - a_log=a_log, - dt_bias=dt_bias, - t_inv=t_inv, - num_k_heads=num_k_heads, - num_v_heads=num_v_heads, - head_k_dim=head_k_dim, - head_v_dim=head_v_dim, - chunk_size=chunk_size, - use_qk_norm_in_gdn=use_qk_norm_in_gdn, - s_ext_pass1=states_pass1[1], - segment_ids=s_enc_local, - init_seg=init_seg, - precision=cp_matmul_precision, - ) + if target_platform(qkv) == "tpu" and cp_matmul_precision == jax.lax.Precision.HIGH: + # The same chunk transitions as compose_local_from_t_inv at the same 3-pass precision, multiplied in + # sequence in a Pallas scan instead of XLA's HBM-bound tree of f32 [chunks, heads, d_k, d_k] products. + m_local = dh_scan.gdn_cp_m_scan( + k_conv, + b, + a, + a_log, + dt_bias, + t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + segment_ids=s_enc_local, + init_seg=init_seg, + ) + s_ext_local = states_pass1[1].astype(jnp.float32) + else: + m_local, s_ext_local = cp_gdn.compose_local_from_t_inv( + qkv_conv=k_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + t_inv=t_inv, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + s_ext_pass1=states_pass1[1], + segment_ids=s_enc_local, + init_seg=init_seg, + precision=cp_matmul_precision, + ) m_local = checkpoint_name(m_local, "gdn_cp_m_local") h_init = recurrent_state.astype(jnp.float32) if recurrent_state is not None else zero_rs @@ -909,25 +1027,47 @@ def _gdn_decoupled_conv1d_bwd( if d_conv_state is not None else None ) - dm_local, ds_ext_local = cp_gdn.compose_bwd_local_from_t_inv( - qkv_conv=qkv_conv, - b=b, - a=a, - a_log=a_log, - dt_bias=dt_bias, - do=d_out, - t_inv=t_inv, - num_k_heads=num_k_heads, - num_v_heads=num_v_heads, - head_k_dim=head_k_dim, - head_v_dim=head_v_dim, - chunk_size=chunk_size, - use_qk_norm_in_gdn=use_qk_norm_in_gdn, - m_local_cached=m_local_fwd, - segment_ids=segment_ids, - init_seg=init_seg, - precision=cp_matmul_precision, - ) + if target_platform(qkv_conv) == "tpu" and m_local_fwd is not None and cp_matmul_precision == jax.lax.Precision.HIGH: + # ds_ext_local of compose_bwd_local_from_t_inv at the same 3-pass precision, as a reverse Pallas scan + # (Horner order) instead of XLA's HBM-bound tree. dm_local is the transposed forward transition, which + # compose_bwd_local_from_t_inv also returns when given m_local_cached. + ds_ext_local = dh_scan.gdn_cp_dh_scan( + qkv_conv, + b, + a, + a_log, + dt_bias, + d_out, + t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + segment_ids=segment_ids, + init_seg=init_seg, + ) + dm_local = jnp.swapaxes(m_local_fwd.astype(jnp.float32), -1, -2) + else: + dm_local, ds_ext_local = cp_gdn.compose_bwd_local_from_t_inv( + qkv_conv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=d_out, + t_inv=t_inv, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + m_local_cached=m_local_fwd, + segment_ids=segment_ids, + init_seg=init_seg, + precision=cp_matmul_precision, + ) dht_local, _ = cp_gdn.incoming_grad_state( dm_local, ds_ext_local, dht_final, cp_axis_name, precision=cp_matmul_precision ) diff --git a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py index 3d379e5cf7..f52565aa50 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py @@ -22,15 +22,23 @@ def make_bwd_block_specs( num_chunks: int, chunk_size: int, - dim_size: int, + qk_tile_size: int, + v_tile_size: int, num_v_heads: int, kq_head_dim: int, v_head_dim: int, + k_block_offset: int, + v_block_offset: int, padded_num_v_heads: int | None = None, has_dht: bool = False, has_dh0: bool = False, ) -> Tuple[list[pl.BlockSpec], list[pl.BlockSpec], int, int]: - """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward.""" + """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward. + + Head group g of q/k/v and dq/dk/dv is a column block of a token-major array (the k and v blocks of qkv_conv start + at block k_block_offset / v_block_offset), and do / chunk_states / t_inv are indexed by group, so the kernel reads + and writes them in place instead of through group-major copies. + """ if padded_num_v_heads is None: padded_num_v_heads = ((num_v_heads + 127) // 128) * 128 @@ -38,47 +46,57 @@ def rc(c): return num_chunks - 1 - c in_specs = [ - # 0: qkv_conv [B, G, C, chunk_size, dim_size] + # 0: q [B, C, chunk_size, dim_size], columns of group g pl.BlockSpec( - (None, None, None, chunk_size, dim_size), - lambda b, g, c: (b, g, rc(c), 0, 0), + (None, None, chunk_size, qk_tile_size), + lambda b, g, c: (b, rc(c), 0, g), + ), + # 1: k [B, C, chunk_size, dim_size], columns of group g + pl.BlockSpec( + (None, None, chunk_size, qk_tile_size), + lambda b, g, c: (b, rc(c), 0, k_block_offset + g), + ), + # 2: v [B, C, chunk_size, dim_size or v_size], columns of group g + pl.BlockSpec( + (None, None, chunk_size, v_tile_size), + lambda b, g, c: (b, rc(c), 0, v_block_offset + g), ), - # 1: b [B, G, C, chunk_size, padded_num_v_heads] + # 3: b [B, G, C, chunk_size, padded_num_v_heads] pl.BlockSpec( (None, None, None, chunk_size, padded_num_v_heads), lambda b, g, c: (b, g, rc(c), 0, 0), ), - # 2: a [B, G, C, chunk_size, padded_num_v_heads] + # 4: a [B, G, C, chunk_size, padded_num_v_heads] pl.BlockSpec( (None, None, None, chunk_size, padded_num_v_heads), lambda b, g, c: (b, g, rc(c), 0, 0), ), - # 3: do [B, G, C, chunk_size, num_v_heads, v_head_dim] + # 5: do [B, C, chunk_size, G, num_v_heads, v_head_dim] pl.BlockSpec( - (None, None, None, chunk_size, num_v_heads, v_head_dim), - lambda b, g, c: (b, g, rc(c), 0, 0, 0), + (None, None, chunk_size, None, num_v_heads, v_head_dim), + lambda b, g, c: (b, rc(c), 0, g, 0, 0), ), - # 4: chunk_states [B, G, C, num_v_heads, kq_head_dim, v_head_dim] + # 6: chunk_states [B, C, G, num_v_heads, kq_head_dim, v_head_dim] pl.BlockSpec( (None, None, None, num_v_heads, kq_head_dim, v_head_dim), - lambda b, g, c: (b, g, rc(c), 0, 0, 0), + lambda b, g, c: (b, rc(c), g, 0, 0, 0), ), - # 5: t_inv [B, G, C, num_v_heads, chunk_size, chunk_size] + # 7: t_inv [B, C, G, num_v_heads, chunk_size, chunk_size] pl.BlockSpec( (None, None, None, num_v_heads, chunk_size, chunk_size), - lambda b, g, c: (b, g, rc(c), 0, 0, 0), + lambda b, g, c: (b, rc(c), g, 0, 0, 0), ), - # 6: a_log [B, G, 1, padded_num_v_heads] + # 8: a_log [B, G, 1, padded_num_v_heads] pl.BlockSpec( (None, None, 1, padded_num_v_heads), lambda b, g, c: (b, g, 0, 0), ), - # 7: dt_bias [B, G, 1, padded_num_v_heads] + # 9: dt_bias [B, G, 1, padded_num_v_heads] pl.BlockSpec( (None, None, 1, padded_num_v_heads), lambda b, g, c: (b, g, 0, 0), ), - # 8: reset [B, G, C, 1, 128] + # 10: reset [B, G, C, 1, 128] pl.BlockSpec( (None, None, None, 1, 128), lambda b, g, c: (b, g, rc(c), 0, 0), @@ -86,7 +104,7 @@ def rc(c): ] if has_dht: in_specs.append( - # 9: dht [B, G, 1, num_v_heads, kq_head_dim, v_head_dim] + # 11: dht [B, G, 1, num_v_heads, kq_head_dim, v_head_dim] pl.BlockSpec( (None, None, 1, num_v_heads, kq_head_dim, v_head_dim), lambda b, g, c: (b, g, 0, 0, 0, 0), @@ -94,27 +112,37 @@ def rc(c): ) out_specs = [ - # 0: dy_conv [B, G, C, chunk_size, dim_size] + # 0: dq [B, C, chunk_size, q_size], columns of group g pl.BlockSpec( - (None, None, None, chunk_size, dim_size), - lambda b, g, c: (b, g, rc(c), 0, 0), + (None, None, chunk_size, qk_tile_size), + lambda b, g, c: (b, rc(c), 0, g), + ), + # 1: dk [B, C, chunk_size, k_size], columns of group g + pl.BlockSpec( + (None, None, chunk_size, qk_tile_size), + lambda b, g, c: (b, rc(c), 0, g), + ), + # 2: dv [B, C, chunk_size, v_size], columns of group g + pl.BlockSpec( + (None, None, chunk_size, v_tile_size), + lambda b, g, c: (b, rc(c), 0, g), ), - # 1: db [B, G, C, chunk_size, padded_num_v_heads] + # 3: db [B, G, C, chunk_size, padded_num_v_heads] pl.BlockSpec( (None, None, None, chunk_size, padded_num_v_heads), lambda b, g, c: (b, g, rc(c), 0, 0), ), - # 2: da [B, G, C, chunk_size, padded_num_v_heads] + # 4: da [B, G, C, chunk_size, padded_num_v_heads] pl.BlockSpec( (None, None, None, chunk_size, padded_num_v_heads), lambda b, g, c: (b, g, rc(c), 0, 0), ), - # 3: dal [B, G, C, 1, padded_num_v_heads] + # 5: dal [B, G, C, 1, padded_num_v_heads] pl.BlockSpec( (None, None, None, 1, padded_num_v_heads), lambda b, g, c: (b, g, rc(c), 0, 0), ), - # 4: ddt [B, G, C, 1, padded_num_v_heads] + # 6: ddt [B, G, C, 1, padded_num_v_heads] pl.BlockSpec( (None, None, None, 1, padded_num_v_heads), lambda b, g, c: (b, g, rc(c), 0, 0), @@ -122,7 +150,7 @@ def rc(c): ] if has_dh0: out_specs.append( - # 5: dh0 [B, G, 1, num_v_heads, kq_head_dim, v_head_dim] + # 7: dh0 [B, G, 1, num_v_heads, kq_head_dim, v_head_dim] pl.BlockSpec( (None, None, 1, num_v_heads, kq_head_dim, v_head_dim), lambda b, g, c: (b, g, 0, 0, 0, 0), diff --git a/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py b/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py new file mode 100644 index 0000000000..b4dba19cf4 --- /dev/null +++ b/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py @@ -0,0 +1,543 @@ +# 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 +# +# http://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. + +r"""Pallas scans that fold one sequence-sharded CP rank's GDN chunk transitions on Mosaic TPU. + +Per chunk $c$ of the rank, with $w_c = T_c^{-1} k^{\beta g}_c$, the forward state map is +$h \mapsto M_c h + S_c$ with $M_c = g^{last}_c I - (k^{bwd}_c)^\top w_c$, and the backward one is +$\mathrm{d}h \mapsto M_c^\top \mathrm{d}h + \mathrm{d}S_c$ with +$\mathrm{d}S_c = q_g^\top \mathrm{d}o - (k^{\beta g}_c)^\top T_c^{-\top} (\mathrm{attn}_c^\top \mathrm{d}o)$. + +- `gdn_cp_m_scan` returns $M_{N-1} \cdots M_0$, the `m_local` of `cp_gdn.compose_local_from_t_inv` + (its `s_ext_pass1` branch). +- `gdn_cp_dh_scan` returns $\mathrm{d}S_0 + M_0^\top(\mathrm{d}S_1 + M_1^\top(\cdots))$, the `ds_ext_local` of + `cp_gdn.compose_bwd_local_from_t_inv`, as a reverse scan in the form of the backward kernel's state recurrence. + +Both are the same terms as the `cp_gdn` balanced trees in sequential order, with every matmul at 3-pass bf16 +(`Precision.HIGH`, as `gdn_cp_matmul_precision=high` gives the trees) and the segment-masked gating cumsum at +`HIGHEST`. They read the head groups of qkv / do / t_inv in place, with the block layout of +`pallas_gdn_bwd_kernel`, and keep the running product in VMEM, so no [chunks, heads, d_k, d_k] f32 intermediate +reaches HBM. +""" + +import functools +from typing import Optional + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +from .. import compute_conv1d as local_compute_conv1d +from .pallas_mosaic_tpu_bwd import _gdn_matmul +from .runtime_utils import ensure_cpu_interpret_registered +from .runtime_utils import target_platform + +# Value heads per grid step. The dh-scan's per-chunk working set grows with the head tile; the M-scan is lighter +# and runs all 64 heads of the Qwen3.5-397B layer in one group. +_DH_SCAN_HEAD_TILE = 32 +_M_SCAN_HEAD_TILE = 64 + +_precise_matmul = functools.partial(_gdn_matmul, compute_dtype=jnp.bfloat16, precise=True) + + +def _pick_head_tile(num_v_heads: int, repeats: int, target: int) -> int: + """Largest divisor of num_v_heads up to target that holds whole GVA groups, as pallas_gdn_bwd_kernel picks it.""" + for candidate in range(min(num_v_heads, target), 0, -1): + if num_v_heads % candidate == 0 and candidate % repeats == 0: + return candidate + return repeats if num_v_heads % repeats == 0 else num_v_heads + + +def _grouped_gate_inputs( + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + segment_ids: Optional[jax.Array], + init_seg: Optional[jax.Array], + *, + num_chunks: int, + chunk_size: int, + num_groups: int, + tile_v_heads: int, +): + """Lays out b, a, a_log, dt_bias and the segment metadata per head group, as pallas_gdn_bwd_kernel does. + + Returns b / a as [B, G, C, chunk_size, padded], a_log / dt_bias as [B, G, 1, padded] and, per chunk, the active + segment entering it (lane 0 of [B, G, C, 1, 128]). With segment IDs, their signed encoding rides in b's spare lane + `tile_v_heads`, and b stays f32 then: bf16 rounds integers above 256, so distinct documents would share an ID. + """ + batch_size = b.shape[0] + seq_len = num_chunks * chunk_size + has_seg_ids = segment_ids is not None + padded = ((tile_v_heads + int(has_seg_ids) + 127) // 128) * 128 + lane_pad = ((0, 0),) * 4 + ((0, padded - tile_v_heads),) + + def per_group(x): + x = x.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_heads).transpose((0, 3, 1, 2, 4)) + return jnp.pad(x, lane_pad) + + def per_group_param(x): + x = jnp.pad(x.reshape(num_groups, tile_v_heads), ((0, 0), (0, padded - tile_v_heads))) + return jnp.broadcast_to(x[None, :, None, :], (batch_size, num_groups, 1, padded)) + + b_5d = per_group(b.astype(jnp.float32) if has_seg_ids else b) + a_5d = per_group(a) + if has_seg_ids: + s_enc = local_compute_conv1d.encode_segment_ids(segment_ids.reshape(batch_size, seq_len), init_seg=init_seg) + s_enc_3d = s_enc.reshape(batch_size, num_chunks, chunk_size) + b_5d = b_5d.at[:, :, :, :, tile_v_heads].set(s_enc_3d[:, None, :, :].astype(b_5d.dtype)) + if init_seg is not None: + init_active = jnp.abs(init_seg.astype(jnp.float32).reshape(batch_size, 1)) + else: + init_active = jnp.zeros((batch_size, 1), dtype=jnp.float32) + seg_prev = jnp.concatenate([init_active, jnp.abs(s_enc_3d[:, :-1, -1])], axis=1) + seg_prev_hbm = jnp.pad(seg_prev[:, :, None, None], ((0, 0), (0, 0), (0, 0), (0, 127))) + else: + seg_prev_hbm = jnp.zeros((batch_size, num_chunks, 1, 128), dtype=jnp.float32) + seg_prev_5d = jnp.broadcast_to(seg_prev_hbm[:, None], (batch_size, num_groups, num_chunks, 1, 128)) + return b_5d, a_5d, per_group_param(a_log), per_group_param(dt_bias), seg_prev_5d, padded + + +def _chunk_gates(b_ref, a_ref, a_log_ref, dt_bias_ref, seg_prev_ref, *, chunk_size, num_v_heads, has_seg_ids): + """Segment masks and decay gates of one chunk and head group, as cp_gdn computes them. + + Returns (valid [C, 1] or None, beta_h [H, C], cumsum_h [H, C], gating_forward [H, C, 1], gating_last [H, 1, 1], + gating_backward [H, C, 1], mask_causal [C, C]). + """ + b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32) + a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32) + a_log_val = a_log_ref[...][0, :num_v_heads].astype(jnp.float32) + dt_bias_val = dt_bias_ref[...][0, :num_v_heads].astype(jnp.float32) + beta = jax.nn.sigmoid(b_val) + log_g = -jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val) + mask_tril = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32)) + + valid_c = None + mask_cumsum = mask_tril + mask_causal = mask_tril + if has_seg_ids: + seg_c = b_ref[...][:, num_v_heads : num_v_heads + 1].astype(jnp.float32) + seg_prev = seg_prev_ref[...][0, 0].astype(jnp.float32) + valid_c = seg_c > 0.5 + active_c = jnp.abs(seg_c) + beta = jnp.where(valid_c, beta, 0.0) + log_g = jnp.where(valid_c, log_g, 0.0) + same_active = (jnp.abs(active_c - active_c.T) < 0.5) & (active_c > 0.5) + same_valid = (jnp.abs(seg_c - seg_c.T) < 0.5) & valid_c + mask_cumsum = mask_tril * same_active.astype(jnp.float32) + mask_causal = mask_tril * same_valid.astype(jnp.float32) + + # The cumsum feeds exp(); like cp_gdn it stays at HIGHEST (0/1 mask times f32: exact up to summation order). + cumsum_h = jnp.transpose(jnp.dot(mask_cumsum, log_g, precision=jax.lax.Precision.HIGHEST), (1, 0)) + gating_forward = jnp.exp(cumsum_h)[:, :, None] + gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None] + if has_seg_ids: + m_in = ((jnp.abs(seg_c - seg_prev) < 0.5) & (seg_prev > 0.5)).astype(jnp.float32) + m_out = ((jnp.abs(seg_c - active_c[-1:]) < 0.5) & (active_c[-1:] > 0.5)).astype(jnp.float32) + m_keep = ((jnp.abs(active_c[-1:] - seg_prev) < 0.5) & (seg_prev > 0.5)).astype(jnp.float32) + gating_forward = gating_forward * m_in[None, :, :] + gating_last = gating_last * m_keep[None, :, :] + bwd_diff = jnp.where(m_out.T > 0.5, cumsum_h[:, -1:] - cumsum_h, -1e4) + gating_backward = jnp.exp(bwd_diff)[:, :, None] * m_out[None, :, :] + else: + gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None] + beta_h = jnp.transpose(beta, (1, 0)) + return valid_c, beta_h, cumsum_h, gating_forward, gating_last, gating_backward, mask_causal + + +def _unit_rows(x: jax.Array, use_qk_norm: bool) -> jax.Array: + """The q / k l2 normalization of cp_gdn: x * rsqrt(sum(x**2) + 1e-6) over the head dim.""" + if not use_qk_norm: + return x + return x * jax.lax.rsqrt(jnp.sum(x**2, axis=-1, keepdims=True) + 1e-6) + + +def _m_scan_body( + k_ref, + b_ref, + a_ref, + t_inv_ref, + a_log_ref, + dt_bias_ref, + seg_prev_ref, + m_ref, + acc_scr, + *, + chunk_size: int, + num_kq_heads: int, + num_v_heads: int, + kq_head_dim: int, + num_chunks: int, + has_seg_ids: bool, + use_qk_norm: bool, +): + """Forward step per (batch, group, chunk): acc <- M_c acc = g_last acc - k_bwd^T (w (acc)), from acc = I.""" + c = pl.program_id(2) + + @pl.when(c == 0) + def _init(): + rows = jax.lax.broadcasted_iota(jnp.int32, (kq_head_dim, kq_head_dim), 0) + cols = jax.lax.broadcasted_iota(jnp.int32, (kq_head_dim, kq_head_dim), 1) + acc_scr[...] = jnp.broadcast_to((rows == cols).astype(jnp.float32)[None], acc_scr.shape) + + valid_c, beta_h, _, gating_forward, gating_last, gating_backward, _ = _chunk_gates( + b_ref, + a_ref, + a_log_ref, + dt_bias_ref, + seg_prev_ref, + chunk_size=chunk_size, + num_v_heads=num_v_heads, + has_seg_ids=has_seg_ids, + ) + k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) + if valid_c is not None: + k_orig = jnp.where(valid_c[:, :, None], k_orig, 0.0) + k_h = jnp.transpose(jnp.repeat(_unit_rows(k_orig, use_qk_norm), num_v_heads // num_kq_heads, axis=1), (1, 0, 2)) + k_beta_g = k_h * beta_h[:, :, None] * gating_forward + k_bwd = k_h * gating_backward + + acc = acc_scr[...] + w = _precise_matmul(t_inv_ref[...].astype(jnp.float32), k_beta_g) + acc_new = acc * gating_last - _precise_matmul(jnp.swapaxes(k_bwd, -1, -2), _precise_matmul(w, acc)) + acc_scr[...] = acc_new + + @pl.when(c == num_chunks - 1) + def _store(): + m_ref[0, ...] = acc_new + + +def _dh_scan_body( + q_ref, + k_ref, + b_ref, + a_ref, + do_ref, + t_inv_ref, + a_log_ref, + dt_bias_ref, + seg_prev_ref, + ds_ref, + acc_scr, + *, + chunk_size: int, + num_kq_heads: int, + num_v_heads: int, + kq_head_dim: int, + num_chunks: int, + has_seg_ids: bool, + use_qk_norm: bool, +): + """Reverse step per (batch, group, chunk): acc <- dS_c + M_c^T acc, from acc = 0. + + M_c^T acc + dS_c is evaluated as the backward kernel's state recurrence does it: + g_last acc + q_g^T do - k_beta_g^T (t_inv^T (attn^T do + k_bwd acc)). + """ + c = pl.program_id(2) + repeats = num_v_heads // num_kq_heads + + @pl.when(c == 0) + def _init(): + acc_scr[...] = jnp.zeros_like(acc_scr) + + valid_c, beta_h, cumsum_h, gating_forward, gating_last, gating_backward, mask_causal = _chunk_gates( + b_ref, + a_ref, + a_log_ref, + dt_bias_ref, + seg_prev_ref, + chunk_size=chunk_size, + num_v_heads=num_v_heads, + has_seg_ids=has_seg_ids, + ) + q_orig = q_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) + k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) + do_val = do_ref[...].astype(jnp.float32) + if valid_c is not None: + q_orig = jnp.where(valid_c[:, :, None], q_orig, 0.0) + k_orig = jnp.where(valid_c[:, :, None], k_orig, 0.0) + do_val = jnp.where(valid_c[:, :, None], do_val, 0.0) + q_scaled = _unit_rows(q_orig, use_qk_norm) * (1.0 / jnp.sqrt(kq_head_dim)) + k_scaled = _unit_rows(k_orig, use_qk_norm) + + q_h = jnp.transpose(jnp.repeat(q_scaled, repeats, axis=1), (1, 0, 2)) + k_h = jnp.transpose(jnp.repeat(k_scaled, repeats, axis=1), (1, 0, 2)) + do_h = jnp.transpose(do_val, (1, 0, 2)) + diff = cumsum_h[:, :, None] - cumsum_h[:, None, :] + g_mat_causal = jnp.exp(jnp.where(mask_causal[None] > 0.5, diff, -1e4)) * mask_causal[None] + q_g = q_h * gating_forward + k_beta_g = k_h * beta_h[:, :, None] * gating_forward + k_bwd = k_h * gating_backward + + # q k^T once per kq head, then repeated to its value heads: the same dot products as with repeated q and k. + qk = _precise_matmul(jnp.transpose(q_scaled, (1, 0, 2)), jnp.swapaxes(jnp.transpose(k_scaled, (1, 0, 2)), -1, -2)) + attn = jnp.repeat(qk, repeats, axis=0) * g_mat_causal + dv_attn = _precise_matmul(jnp.swapaxes(attn, -1, -2), do_h) + + acc = acc_scr[...] + t_inv_t = jnp.swapaxes(t_inv_ref[...].astype(jnp.float32), -1, -2) + d_v_beta = _precise_matmul(t_inv_t, dv_attn + _precise_matmul(k_bwd, acc)) + # q_g^T do - k_beta_g^T d_v_beta as one matmul over the stacked chunk rows. + acc_new = acc * gating_last + _precise_matmul( + jnp.swapaxes(jnp.concatenate([q_g, -k_beta_g], axis=1), -1, -2), + jnp.concatenate([do_h, d_v_beta], axis=1), + ) + acc_scr[...] = acc_new + + @pl.when(c == num_chunks - 1) + def _store(): + ds_ref[0, ...] = acc_new + + +def _run_scan(body, args, in_specs, out_shape, out_spec, acc_shape, *, grid, interpret, name): + """One pallas_call that pipelines `body` over grid=(B, G, C) with the running product in a VMEM scratch.""" + + def outer(*refs): + pltpu.emit_pipeline(body, grid=grid, in_specs=in_specs, out_specs=[out_spec])( + *refs[: len(args) + 1], scratches=tuple(refs[len(args) + 1 :]) + ) + + hbm = pltpu.MemorySpace.HBM + return pl.pallas_call( + outer, + grid=(), + out_shape=out_shape, + in_specs=[pl.BlockSpec(memory_space=hbm)] * len(args), + out_specs=pl.BlockSpec(memory_space=hbm), + scratch_shapes=[pltpu.VMEM(acc_shape, jnp.float32)], + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=int(0.9 * pltpu.get_tpu_info().vmem_capacity_bytes), + disable_bounds_checks=True, + ), + interpret=interpret, + name=name, + )(*args) + + +def _resolve_interpret(x: jax.Array, interpret: bool | pltpu.InterpretParams | None): + if interpret is None: + interpret = target_platform(x) == "cpu" + if interpret: + ensure_cpu_interpret_registered() + return interpret + + +def gdn_cp_m_scan( + k_conv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + t_inv: jax.Array, + *, + num_v_heads: int, + kq_head_dim: int, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, + segment_ids: Optional[jax.Array] = None, + init_seg: Optional[jax.Array] = None, + interpret: bool | pltpu.InterpretParams | None = None, +) -> jax.Array: + """M_local of `cp_gdn.compose_local_from_t_inv` (s_ext_pass1 branch), [B, num_v_heads, d_k, d_k] f32. + + Args: + k_conv: [B, S, num_kq_heads * d_k] conv output of the K channels, as passed to compose_local_from_t_inv. + b, a: [B, S, num_v_heads] beta and decay pre-activations. + a_log, dt_bias: [num_v_heads] decay parameters. + t_inv: [B, S / chunk_size, num_v_heads, chunk_size, chunk_size] per-chunk T^{-1} of the forward kernel. + num_v_heads, kq_head_dim, chunk_size, use_qk_norm_in_gdn: GDN layer geometry. + segment_ids, init_seg: Rank-local segment metadata, as passed to compose_local_from_t_inv. + interpret: Pallas interpret mode; None interprets on CPU. + """ + interpret = _resolve_interpret(k_conv, interpret) + batch_size, seq_len, k_size = k_conv.shape + num_chunks = seq_len // chunk_size + num_kq_heads = k_size // kq_head_dim + repeats = num_v_heads // num_kq_heads + tile_v_heads = _pick_head_tile(num_v_heads, repeats, _M_SCAN_HEAD_TILE) + num_groups = num_v_heads // tile_v_heads + tile_k_size = (tile_v_heads // repeats) * kq_head_dim + has_seg_ids = segment_ids is not None + + b_5d, a_5d, a_log_4d, dt_bias_4d, seg_prev_5d, padded = _grouped_gate_inputs( + b, + a, + a_log, + dt_bias, + segment_ids, + init_seg, + num_chunks=num_chunks, + chunk_size=chunk_size, + num_groups=num_groups, + tile_v_heads=tile_v_heads, + ) + args = [ + k_conv.reshape(batch_size, num_chunks, chunk_size, k_size), + b_5d, + a_5d, + t_inv.reshape(batch_size, num_chunks, num_groups, tile_v_heads, chunk_size, chunk_size).astype(jnp.float32), + a_log_4d, + dt_bias_4d, + seg_prev_5d, + ] + in_specs = [ + # k [B, C, chunk_size, k_size], columns of group g + pl.BlockSpec((None, None, chunk_size, tile_k_size), lambda b_, g, c: (b_, c, 0, g)), + # b, a [B, G, C, chunk_size, padded] + pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, c, 0, 0)), + pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, c, 0, 0)), + # t_inv [B, C, G, tile_v_heads, chunk_size, chunk_size] + pl.BlockSpec((None, None, None, tile_v_heads, chunk_size, chunk_size), lambda b_, g, c: (b_, c, g, 0, 0, 0)), + # a_log, dt_bias [B, G, 1, padded] + pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)), + pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)), + # segment entering the chunk [B, G, C, 1, 128] + pl.BlockSpec((None, None, None, 1, 128), lambda b_, g, c: (b_, g, c, 0, 0)), + ] + acc_shape = (tile_v_heads, kq_head_dim, kq_head_dim) + body = functools.partial( + _m_scan_body, + chunk_size=chunk_size, + num_kq_heads=tile_v_heads // repeats, + num_v_heads=tile_v_heads, + kq_head_dim=kq_head_dim, + num_chunks=num_chunks, + has_seg_ids=has_seg_ids, + use_qk_norm=use_qk_norm_in_gdn, + ) + m_6d = _run_scan( + body, + args, + in_specs, + jax.ShapeDtypeStruct((batch_size, num_groups, 1) + acc_shape, jnp.float32), + pl.BlockSpec((None, None, 1) + acc_shape, lambda b_, g, c: (b_, g, 0, 0, 0, 0)), + acc_shape, + grid=(batch_size, num_groups, num_chunks), + interpret=interpret, + name="gdn_cp_m_scan", + ) + return m_6d[:, :, 0].reshape(batch_size, num_v_heads, kq_head_dim, kq_head_dim) + + +def gdn_cp_dh_scan( + qkv_conv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + t_inv: jax.Array, + *, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, + segment_ids: Optional[jax.Array] = None, + init_seg: Optional[jax.Array] = None, + interpret: bool | pltpu.InterpretParams | None = None, +) -> jax.Array: + """dS_ext_local of `cp_gdn.compose_bwd_local_from_t_inv`, [B, num_v_heads, d_k, d_v] f32. + + Args: + qkv_conv: [B, S, 2 * num_kq_heads * d_k + num_v_heads * d_v] conv output (q, k, v channels). + b, a: [B, S, num_v_heads] beta and decay pre-activations. + a_log, dt_bias: [num_v_heads] decay parameters. + do: [B, S, num_v_heads, d_v] cotangent of the GDN output. + t_inv: [B, S / chunk_size, num_v_heads, chunk_size, chunk_size] per-chunk T^{-1} of the forward kernel. + num_v_heads, kq_head_dim, v_head_dim, chunk_size, use_qk_norm_in_gdn: GDN layer geometry. + segment_ids, init_seg: Rank-local segment metadata, as passed to compose_bwd_local_from_t_inv. + interpret: Pallas interpret mode; None interprets on CPU. + """ + interpret = _resolve_interpret(qkv_conv, interpret) + batch_size, seq_len, dim_size = qkv_conv.shape + num_chunks = seq_len // chunk_size + num_kq_heads = (dim_size - num_v_heads * v_head_dim) // (kq_head_dim * 2) + repeats = num_v_heads // num_kq_heads + tile_v_heads = _pick_head_tile(num_v_heads, repeats, _DH_SCAN_HEAD_TILE) + num_groups = num_v_heads // tile_v_heads + tile_q_size = (tile_v_heads // repeats) * kq_head_dim + k_block_offset = num_kq_heads * kq_head_dim // tile_q_size + has_seg_ids = segment_ids is not None + + b_5d, a_5d, a_log_4d, dt_bias_4d, seg_prev_5d, padded = _grouped_gate_inputs( + b, + a, + a_log, + dt_bias, + segment_ids, + init_seg, + num_chunks=num_chunks, + chunk_size=chunk_size, + num_groups=num_groups, + tile_v_heads=tile_v_heads, + ) + qkv_conv_4d = qkv_conv.reshape(batch_size, num_chunks, chunk_size, dim_size) + args = [ + qkv_conv_4d, + qkv_conv_4d, + b_5d, + a_5d, + do.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_heads, v_head_dim), + t_inv.reshape(batch_size, num_chunks, num_groups, tile_v_heads, chunk_size, chunk_size).astype(jnp.float32), + a_log_4d, + dt_bias_4d, + seg_prev_5d, + ] + + def rc(c): + return num_chunks - 1 - c + + in_specs = [ + # q, k [B, C, chunk_size, dim_size], columns of group g (k after the q blocks) + pl.BlockSpec((None, None, chunk_size, tile_q_size), lambda b_, g, c: (b_, rc(c), 0, g)), + pl.BlockSpec((None, None, chunk_size, tile_q_size), lambda b_, g, c: (b_, rc(c), 0, k_block_offset + g)), + # b, a [B, G, C, chunk_size, padded] + pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, rc(c), 0, 0)), + pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, rc(c), 0, 0)), + # do [B, C, chunk_size, G, tile_v_heads, d_v] + pl.BlockSpec((None, None, chunk_size, None, tile_v_heads, v_head_dim), lambda b_, g, c: (b_, rc(c), 0, g, 0, 0)), + # t_inv [B, C, G, tile_v_heads, chunk_size, chunk_size] + pl.BlockSpec((None, None, None, tile_v_heads, chunk_size, chunk_size), lambda b_, g, c: (b_, rc(c), g, 0, 0, 0)), + # a_log, dt_bias [B, G, 1, padded] + pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)), + pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)), + # segment entering the chunk [B, G, C, 1, 128] + pl.BlockSpec((None, None, None, 1, 128), lambda b_, g, c: (b_, g, rc(c), 0, 0)), + ] + acc_shape = (tile_v_heads, kq_head_dim, v_head_dim) + body = functools.partial( + _dh_scan_body, + chunk_size=chunk_size, + num_kq_heads=tile_v_heads // repeats, + num_v_heads=tile_v_heads, + kq_head_dim=kq_head_dim, + num_chunks=num_chunks, + has_seg_ids=has_seg_ids, + use_qk_norm=use_qk_norm_in_gdn, + ) + ds_6d = _run_scan( + body, + args, + in_specs, + jax.ShapeDtypeStruct((batch_size, num_groups, 1) + acc_shape, jnp.float32), + pl.BlockSpec((None, None, 1) + acc_shape, lambda b_, g, c: (b_, g, 0, 0, 0, 0)), + acc_shape, + grid=(batch_size, num_groups, num_chunks), + interpret=interpret, + name="gdn_cp_dh_scan", + ) + return ds_6d[:, :, 0].reshape(batch_size, num_v_heads, kq_head_dim, v_head_dim) diff --git a/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py b/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py index f4c490af8b..d57de46af2 100644 --- a/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py +++ b/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py @@ -139,13 +139,15 @@ def _bwd_gdn_pipeline_body( has_dh0: bool = False, ) -> None: """Inner kernel executed per (batch, group, chunk) by emit_pipeline with manual GDN backward.""" - if len(refs) == 17: + if len(refs) == 21: has_dht = True has_dh0 = True idx = 0 # pylint: disable=unbalanced-tuple-unpacking ( - qkv_conv_ref, + q_ref, + k_ref, + v_ref, b_ref, a_ref, do_ref, @@ -154,22 +156,24 @@ def _bwd_gdn_pipeline_body( a_log_ref, dt_bias_ref, reset_ref, - ) = refs[idx : idx + 9] - idx += 9 + ) = refs[idx : idx + 11] + idx += 11 if has_dht: dht_ref = refs[idx] idx += 1 else: dht_ref = None ( - dy_conv_ref, + dq_ref, + dk_ref, + dv_ref, d_b_ref, d_a_ref, d_a_log_ref, d_dt_bias_ref, - ) = refs[idx : idx + 5] + ) = refs[idx : idx + 7] # pylint: enable=unbalanced-tuple-unpacking - idx += 5 + idx += 7 if has_dh0: dh0_ref = refs[idx] idx += 1 @@ -198,17 +202,12 @@ def _init(): is_reset = reset_ref[...][0, 0] > 0.5 d_state = jnp.where(is_reset, 0.0, d_state_scr[...]) - compute_dtype = qkv_conv_ref.dtype - y_c = qkv_conv_ref[...] - - # Slice chunk inputs for this head group - q_orig = y_c[:, :q_size].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) - k_orig = y_c[:, q_size : q_size + k_size].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) - v = ( - y_c[:, q_size + k_size : q_size + k_size + v_size] - .reshape((chunk_size, num_v_heads, v_head_dim)) - .astype(jnp.float32) - ) + compute_dtype = q_ref.dtype + + # Chunk inputs for this head group + q_orig = q_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) + k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) + v = v_ref[...].reshape((chunk_size, num_v_heads, v_head_dim)).astype(jnp.float32) b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32) a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32) @@ -446,11 +445,10 @@ def _init(): d_q = jnp.where(valid_c[:, :, None], d_q, 0.0) d_k = jnp.where(valid_c[:, :, None], d_k, 0.0) - # Flatten gradients to chunk_size x dim_size and write to refs - d_q_flat = d_q.reshape(chunk_size, q_size).astype(dy_conv_ref.dtype) - d_k_flat = d_k.reshape(chunk_size, k_size).astype(dy_conv_ref.dtype) - d_v_flat = d_v_val.reshape(chunk_size, v_size).astype(dy_conv_ref.dtype) - dy_conv_ref[...] = jnp.concatenate([d_q_flat, d_k_flat, d_v_flat], axis=-1) + # Flatten gradients to chunk_size x q/k/v size and write to refs + dq_ref[...] = d_q.reshape(chunk_size, q_size).astype(dq_ref.dtype) + dk_ref[...] = d_k.reshape(chunk_size, k_size).astype(dk_ref.dtype) + dv_ref[...] = d_v_val.reshape(chunk_size, v_size).astype(dv_ref.dtype) if padded_num_v_heads > num_v_heads: d_b_ref[...] = jnp.pad(d_b_val, ((0, 0), (0, padded_num_v_heads - num_v_heads))).astype(d_b_ref.dtype) @@ -619,14 +617,15 @@ def pallas_gdn_bwd_kernel( padded_tile_v_heads = cfg.padded_num_v_heads - # 1. Prepare qkv_conv into (B, G, C, chunk_size, group_dim_size) - q = qkv_conv[:, :, :q_size] - k = qkv_conv[:, :, q_size : q_size + k_size] - v = qkv_conv[:, :, q_size + k_size :] - q_5d = q.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_q_size).transpose((0, 3, 1, 2, 4)) - k_5d = k.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_k_size).transpose((0, 3, 1, 2, 4)) - v_5d = v.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_size).transpose((0, 3, 1, 2, 4)) - qkv_conv_5d = jnp.concatenate([q_5d, k_5d, v_5d], axis=-1) + # 1. View qkv_conv as (B, C, chunk_size, dim_size); the kernel reads group g's q, k and v as column blocks + # (k after the num_groups q blocks, v after q + k). If q + k is not a whole number of v blocks, v is sliced out. + qkv_conv_4d = qkv_conv.reshape(batch_size, num_chunks, chunk_size, dim_size) + if (q_size + k_size) % tile_v_size == 0: + v_4d = qkv_conv_4d + v_block_offset = (q_size + k_size) // tile_v_size + else: + v_4d = qkv_conv_4d[..., q_size + k_size :] + v_block_offset = 0 # 2. Prepare b and a into (B, G, C, chunk_size, padded_tile_v_heads) # With sequence packing, the signed segment IDs are packed into the spare lane @@ -659,7 +658,7 @@ def pallas_gdn_bwd_kernel( ), ) - # 3. Prepare do into (B, G, C, chunk_size, tile_v_heads, v_head_dim) + # 3. View do as (B, C, chunk_size, G, tile_v_heads, v_head_dim); the kernel reads group g in place do_6d = do.reshape( batch_size, num_chunks, @@ -667,9 +666,9 @@ def pallas_gdn_bwd_kernel( num_groups, tile_v_heads, cfg.v_head_dim, - ).transpose((0, 3, 1, 2, 4, 5)) + ) - # 4. Prepare chunk_states into (B, G, C, tile_v_heads, kq_head_dim, v_head_dim) + # 4. View chunk_states as (B, C, G, tile_v_heads, kq_head_dim, v_head_dim) chunk_states_6d = chunk_states.reshape( batch_size, num_chunks, @@ -677,22 +676,18 @@ def pallas_gdn_bwd_kernel( tile_v_heads, cfg.kq_head_dim, cfg.v_head_dim, - ).swapaxes(1, 2) - - # 5. Prepare t_inv into (B, G, C, tile_v_heads, chunk_size, chunk_size) - t_inv_6d = ( - t_inv.reshape( - batch_size, - num_chunks, - num_groups, - tile_v_heads, - chunk_size, - chunk_size, - ) - .swapaxes(1, 2) - .astype(jnp.float32) ) + # 5. View t_inv as (B, C, G, tile_v_heads, chunk_size, chunk_size) + t_inv_6d = t_inv.reshape( + batch_size, + num_chunks, + num_groups, + tile_v_heads, + chunk_size, + chunk_size, + ).astype(jnp.float32) + # 6. Prepare a_log and dt_bias into (B, G, 1, padded_tile_v_heads) a_log_2d = a_log.reshape(num_groups, tile_v_heads) if padded_tile_v_heads > tile_v_heads: @@ -739,7 +734,9 @@ def pallas_gdn_bwd_kernel( reset_hbm_5d = jnp.broadcast_to(reset_hbm[:, None, :, :, :], (batch_size, num_groups, num_chunks, 1, 128)) in_args = [ - qkv_conv_5d, + qkv_conv_4d, + qkv_conv_4d, + v_4d, b_5d, a_5d, do_6d, @@ -767,20 +764,22 @@ def pallas_gdn_bwd_kernel( in_specs, out_specs, num_in, num_out = make_bwd_block_specs( num_chunks=num_chunks, chunk_size=chunk_size, - dim_size=group_dim_size, + qk_tile_size=tile_q_size, + v_tile_size=tile_v_size, num_v_heads=tile_v_heads, kq_head_dim=cfg.kq_head_dim, v_head_dim=cfg.v_head_dim, + k_block_offset=q_size // tile_k_size, + v_block_offset=v_block_offset, padded_num_v_heads=padded_tile_v_heads, has_dht=has_dht, has_dh0=has_dh0, ) out_shapes_list = [ - jax.ShapeDtypeStruct( - (batch_size, num_groups, num_chunks, chunk_size, group_dim_size), - qkv_conv.dtype, - ), + jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, q_size), qkv_conv.dtype), + jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, k_size), qkv_conv.dtype), + jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, v_size), qkv_conv.dtype), jax.ShapeDtypeStruct(b_5d.shape, b_5d.dtype), jax.ShapeDtypeStruct(a_5d.shape, a_5d.dtype), jax.ShapeDtypeStruct( @@ -849,7 +848,9 @@ def outer(*refs): if has_dh0: ( - dy_conv_chunks, + dq_4d, + dk_4d, + dv_4d, d_b_chunks, d_a_chunks, d_a_log_chunks, @@ -858,7 +859,9 @@ def outer(*refs): ) = pallas_out else: ( - dy_conv_chunks, + dq_4d, + dk_4d, + dv_4d, d_b_chunks, d_a_chunks, d_a_log_chunks, @@ -866,13 +869,10 @@ def outer(*refs): ) = pallas_out dh0_6d = None - # Reconstruct dy_conv: (B, G, C, S, group_dim_size) -> (B, seq_len, dim_size) - dq_5d = dy_conv_chunks[..., :tile_q_size] - dk_5d = dy_conv_chunks[..., tile_q_size : tile_q_size + tile_k_size] - dv_5d = dy_conv_chunks[..., tile_q_size + tile_k_size :] - dq = dq_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, q_size) - dk = dk_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, k_size) - dv = dv_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, v_size) + # Reconstruct dy_conv: (B, C, S, q/k/v size) -> (B, seq_len, dim_size) + dq = dq_4d.reshape(batch_size, seq_len, q_size) + dk = dk_4d.reshape(batch_size, seq_len, k_size) + dv = dv_4d.reshape(batch_size, seq_len, v_size) dy_conv_flat = jnp.concatenate([dq, dk, dv], axis=-1).astype(qkv_conv.dtype) # Reconstruct d_b and d_a: (B, G, C, S, padded_tile_v_heads) -> (B, seq_len, num_v_heads) diff --git a/src/maxtext/kernels/gdn/memory_ref.py b/src/maxtext/kernels/gdn/memory_ref.py index f29d79d81d..a1c4545ef5 100644 --- a/src/maxtext/kernels/gdn/memory_ref.py +++ b/src/maxtext/kernels/gdn/memory_ref.py @@ -281,6 +281,36 @@ def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): ).wait() +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class ChunkInBufferedRef(BaseBufferedRef): + """Input DMA reference that copies whole chunks of a [num_chunks, chunk, dim] array.""" + + def copy_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Initiates asynchronous copy of each tile's chunk into current input VMEM slot.""" + assert self.sem_recvs is not None + assert self.window_ref is not None + slot = self.current_copy_in_slot + sem = self.sem_recvs.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + chunk_idx = self.metadata_ref.get_record(p_id, idx).r_base // self.cfg.chunk_size + pltpu.make_async_copy(src_ref.at[chunk_idx], vmem_ref.at[idx], sem).start() + + def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Waits for asynchronous chunk copy to complete.""" + assert self.sem_recvs is not None + assert self.window_ref is not None + slot = self.current_wait_in_slot + sem = self.sem_recvs.at[slot] + vmem_ref = self.window_ref.at[slot] + + for idx in range(self.cfg.seq_tile_size): + pltpu.make_async_copy(vmem_ref.at[idx], vmem_ref.at[idx], sem).wait() + + @jax.tree_util.register_dataclass @dataclasses.dataclass(frozen=True, kw_only=True) class OutBufferedRef(BaseBufferedRef): @@ -496,15 +526,18 @@ def create_allocs( qkv_ref: jax.Array, b_ref: jax.Array, a_ref: jax.Array, - out_ref: jax.Array, + out_ref: jax.Array | None, conv_state_ref: jax.Array, recurrent_state_ref: jax.Array, cfg: config.GDNConfig, t_inv_ref: jax.Array | None = None, chunk_states_ref: jax.Array | None = None, ) -> tuple[Any, ...]: - """Creates all Pallas double-buffered allocations for GDN kernel.""" - qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.dim_size) + """Creates all Pallas double-buffered allocations for GDN kernel (no output buffers without out_ref).""" + if cfg.is_prefill_only: + qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, cfg.dim_size) + else: + qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.dim_size) ba_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.aligned_num_v_heads) out_shape = ( @@ -539,25 +572,31 @@ def create_allocs( cfg=cfg, metadata_ref=metadata_ref, ) - qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref) + if cfg.is_prefill_only: + qkv_alloc = ChunkInBufferedRef.input( + spec=qkv_spec, + dtype_or_type=qkv_ref, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + else: + qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref) b_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=b_ref) a_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=a_ref) - out_alloc = OutBufferedRef.output( - spec=block_spec_partial(block_shape=out_shape), - dtype_or_type=out_ref, - buffer_count=pipeline_mode.buffer_count, - use_lookahead=pipeline_mode.use_lookahead, - cfg=cfg, - metadata_ref=metadata_ref, - ) - - conv_spec = block_spec_partial(block_shape=conv_shape) - recurrent_spec = block_spec_partial(block_shape=recurrent_shape) + # A prefill call with one sequence copies the states in only before its first tile and out only after + # its last, so no copy is in flight while a tile uses the slot: one buffer suffices and frees VMEM. + state_pipeline_mode = pipeline_mode + if cfg.is_prefill_only and metadata_ref.s_idx_has_initial_state.shape[0] == 1: + state_pipeline_mode = pl.Buffered(buffer_count=1, use_lookahead=False) + conv_spec = block_spec_partial(block_shape=conv_shape, pipeline_mode=state_pipeline_mode) + recurrent_spec = block_spec_partial(block_shape=recurrent_shape, pipeline_mode=state_pipeline_mode) state_buffered_partial = functools.partial( StateBufferedRef.input_output, - buffer_count=pipeline_mode.buffer_count, - use_lookahead=pipeline_mode.use_lookahead, + buffer_count=state_pipeline_mode.buffer_count, + use_lookahead=state_pipeline_mode.use_lookahead, cfg=cfg, metadata_ref=metadata_ref, ) @@ -570,9 +609,19 @@ def create_allocs( a_alloc, conv_alloc, recurrent_alloc, - out_alloc, ] + if out_ref is not None: + out_alloc = OutBufferedRef.output( + spec=block_spec_partial(block_shape=out_shape), + dtype_or_type=out_ref, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + allocs.append(out_alloc) + if t_inv_ref is not None: t_inv_shape = ( cfg.seq_tile_size, diff --git a/src/maxtext/kernels/gdn/vmem_ldst.py b/src/maxtext/kernels/gdn/vmem_ldst.py index 7d604a69b1..3cc5225c48 100644 --- a/src/maxtext/kernels/gdn/vmem_ldst.py +++ b/src/maxtext/kernels/gdn/vmem_ldst.py @@ -134,6 +134,28 @@ def load_compact_to_large(vmem_ref: jax.Ref) -> jax.Array: return jnp.concat(vreg_list, axis=-1).reshape(new_shape) +def upcast_large_to_compact(large_ref: jax.Ref, compact_ref: jax.Ref) -> None: + """Upcast activations from large layout into a compact fp32 reference. + + The inverse of load_as_qkv_large: one strided store per 128-lane column puts + each token's lanes into its own row of the compact reference. + + Args: + large_ref: Activations in VMEM of shape [seq_tile_size, chunk_size, + dim_size]. + compact_ref: fp32 reference in VMEM of shape [seq_tile_size, chunk_size, 1, + dim_size]. + """ + + num_lanes = pltpu.get_tpu_info().num_lanes + lanes_per_col = compact_ref.shape[-1] // num_lanes + for idx in range(compact_ref.shape[0]): + compact_flat_ref = compact_ref.at[idx].reshape(-1, num_lanes) # pyrefly: ignore[missing-attribute] + for lane in range(lanes_per_col): + col = large_ref[idx, :, lane * num_lanes : (lane + 1) * num_lanes] + compact_flat_ref[lane::lanes_per_col] = col.astype(compact_ref.dtype) + + def load_and_select_states( metadata_ref: memory_ref.MetadataRef, p_id: jax.Array, diff --git a/src/maxtext/kernels/gdn/wrapper.py b/src/maxtext/kernels/gdn/wrapper.py index 7b951c7cee..65b52a8a89 100644 --- a/src/maxtext/kernels/gdn/wrapper.py +++ b/src/maxtext/kernels/gdn/wrapper.py @@ -95,14 +95,14 @@ def _pack_fwd_segment_metadata( def inner_kernel( # Inputs. - qkv_slot_ref: jax.Ref, # [seq, chunk, 1, dim_size] + qkv_slot_ref: jax.Ref, # [seq, chunk, 1, dim_size] ([seq, chunk, dim_size] for prefill-only) b_slot_ref: jax.Ref, # [seq, chunk, 1, num_v_heads] a_slot_ref: jax.Ref, # [seq, chunk, 1, num_v_heads] conv_state_slot_ref: jax.Ref, # [seq, prev_kernel_size, 1, dim_size] recurrent_slot_ref: jax.Ref, # [seq, num_v_heads, kq_head, v_head] - # Outputs. - out_slot_ref: jax.Array, # [seq * chunk, num_v_heads, v_head] - t_inv_slot_ref: jax.Array, # [seq, num_v_heads, chunk, chunk] + # Outputs, then scratches: + # out [seq * chunk, num_v_heads, v_head] (not in PER_SEQ_STATES), t_inv [seq, num_v_heads, chunk, chunk], + # chunk_states [seq, num_v_heads, kq_head, v_head] (PER_SEQ only). *args, cfg: config.GDNConfig, **kwargs, @@ -114,21 +114,28 @@ def inner_kernel( kernel invokes vmem_ldst to pre-processes data needed for compute and invokes compute_conv1d and compute_gdn for actual compute. """ - if cfg.mode == config.GDNMode.PER_SEQ: - chunk_states_slot_ref = args[0] - metadata_ref = args[1] - weights_ref = args[2] - carry_conv_scratch_ref = args[3] if len(args) > 3 else None - carry_recurrent_scratch_ref = args[4] if len(args) > 4 else None + if cfg.mode == config.GDNMode.PER_SEQ_STATES: + out_slot_ref = chunk_states_slot_ref = None + t_inv_slot_ref, *args = args + elif cfg.mode == config.GDNMode.PER_SEQ: + out_slot_ref, t_inv_slot_ref, chunk_states_slot_ref, *args = args else: + out_slot_ref, t_inv_slot_ref, *args = args chunk_states_slot_ref = None - metadata_ref = args[0] - weights_ref = args[1] - carry_conv_scratch_ref = args[2] if len(args) > 2 else None - carry_recurrent_scratch_ref = args[3] if len(args) > 3 else None + metadata_ref = args[0] + weights_ref = args[1] + carry_conv_scratch_ref = args[2] if len(args) > 2 else None + carry_recurrent_scratch_ref = args[3] if len(args) > 3 else None + qkv_scratch_ref = args[4] if len(args) > 4 else None p_id = pl.program_id(0) + # Prefill-only qkv arrives in act_in dtype and large layout. Upcast it into the + # compact fp32 scratch, which then serves as the qkv slot below. + if qkv_scratch_ref is not None: + vmem_ldst.upcast_large_to_compact(qkv_slot_ref, qkv_scratch_ref) + qkv_slot_ref = qkv_scratch_ref + # Prepare states. real_sizes, prev_conv, prev_recurrent = vmem_ldst.load_and_select_states( metadata_ref=metadata_ref, @@ -216,7 +223,8 @@ def inner_kernel( ) # Store output, recurrent, and t_inv to vmem. - out_slot_ref[...] = out.astype(out_slot_ref.dtype) + if out_slot_ref is not None: + out_slot_ref[...] = out.astype(out_slot_ref.dtype) recurrent_slot_ref[...] = new_recurrent_state.astype(recurrent_slot_ref.dtype) t_inv_slot_ref[...] = t_inv.astype(t_inv_slot_ref.dtype) if chunk_states_slot_ref is not None: @@ -236,19 +244,21 @@ def outer_kernel( recurrent_state_ref: jax.Array, _: jax.Array, weights_ref: memory_ref.WeightRefs, - # Outputs. - out_ref: jax.Array, - conv_state_out_ref: jax.Array, - recurrent_state_out_ref: jax.Array, - t_inv_ref: jax.Array, + # Outputs: out (not in PER_SEQ_STATES), conv_state, recurrent_state, t_inv, chunk_states (PER_SEQ only). + # The state outputs alias the state inputs, which the pipeline updates in place. *args, carry_conv_scratch_ref: jax.Array | None = None, carry_recurrent_scratch_ref: jax.Array | None = None, + qkv_scratch_ref: jax.Array | None = None, cfg: config.GDNConfig, **kwargs, ) -> None: """Setup memory allocations and emit pipeline for running inner_kernel.""" - del conv_state_out_ref, recurrent_state_out_ref + if cfg.mode == config.GDNMode.PER_SEQ_STATES: + out_ref = None + _, _, t_inv_ref, *args = args + else: + out_ref, _, _, t_inv_ref, *args = args chunk_states_ref = args[0] if (len(args) > 0 and cfg.mode == config.GDNMode.PER_SEQ) else None @@ -269,16 +279,11 @@ def outer_kernel( a_alloc = allocs[2] conv_alloc = allocs[3] recurrent_alloc = allocs[4] - out_alloc = allocs[5] - t_inv_alloc = allocs[6] - chunk_states_alloc = allocs[7] if len(allocs) > 7 else None + # [out,] t_inv[, chunk_states], in the order of the pipelined output refs below. + out_specs = [alloc.spec for alloc in allocs[5:]] num_tiles = metadata_ref.num_tiles[...] - out_specs = [out_alloc.spec, t_inv_alloc.spec] - if chunk_states_alloc is not None: - out_specs.append(chunk_states_alloc.spec) - pipeline_func = pltpu.emit_pipeline( body=functools.partial( inner_kernel, @@ -297,9 +302,7 @@ def outer_kernel( @pl.with_scoped(allocations=allocs) def _run(allocations): - out_args = [out_ref, t_inv_ref] - if chunk_states_ref is not None: - out_args.append(chunk_states_ref) + out_args = [ref for ref in (out_ref, t_inv_ref, chunk_states_ref) if ref is not None] pipeline_func( qkv_ref, b_ref, @@ -312,6 +315,7 @@ def _run(allocations): weights_ref, carry_conv_scratch_ref, carry_recurrent_scratch_ref, + qkv_scratch_ref, ), allocations=allocations, ) @@ -320,23 +324,7 @@ def _run(allocations): _run() -@jax.jit( - donate_argnames=("conv_state", "recurrent_state"), - static_argnames=( - "n_kq", - "n_v", - "d_k", - "d_v", - "kernel_size", - "decode_tile_size", - "mixed_tile_size", - "zero_initialize_out", - "compute_precision", - "is_prefill_only", - "use_qk_norm_in_gdn", - ), -) -def fused_conv1d_gdn( +def _fused_conv1d_gdn( qkv: jax.Array, # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size] b: jax.Array, # [batch_size, n_v] a: jax.Array, # [batch_size, n_v] @@ -365,15 +353,15 @@ def fused_conv1d_gdn( segment_ids: jax.Array | None = None, conv_halo_seg: jax.Array | None = None, init_seg: jax.Array | None = None, -) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array, jax.Array]: - """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states).""" + prefill_mode: config.GDNMode = config.GDNMode.PER_SEQ, +) -> tuple[jax.Array | None, tuple[jax.Array, jax.Array], jax.Array, jax.Array | None]: + """Implements fused_conv1d_gdn; prefill_mode=PER_SEQ_STATES returns no out and no chunk_states.""" act_in_dtype = qkv.dtype act_out_dtype = qkv.dtype conv_out_dtype = conv_state.dtype recurrent_out_dtype = recurrent_state.dtype assert a.dtype == b.dtype == qkv.dtype == act_in_dtype - qkv = qkv.astype(jnp.float32) b = b.astype(jnp.float32) a = a.astype(jnp.float32) conv_state = conv_state.astype(jnp.float32) @@ -441,7 +429,6 @@ def fused_conv1d_gdn( batch_padding_size = padded_batch_size - batch_size aligned_num_v_heads = tiling.align_to(n_v + extra_lanes, num_lanes) num_v_padding_size = aligned_num_v_heads - n_v - qkv = jnp.pad(qkv, ((0, batch_padding_size), (0, 0))) b = jnp.pad(b, ((0, batch_padding_size), (0, num_v_padding_size))) a = jnp.pad(a, ((0, batch_padding_size), (0, num_v_padding_size))) if has_seg_ids: @@ -456,7 +443,12 @@ def fused_conv1d_gdn( b = b.at[:batch_size, n_v].set(s_enc_flat) b = b.at[:batch_size, n_v + 1].set(seg_aux_flat) - qkv = qkv.reshape(padded_batch_size, 1, -1) + if is_prefill_only: + # Whole chunks, DMA'd in act_in dtype and upcast to fp32 in VMEM by the kernel. + qkv = qkv.reshape(batch_size // mixed_tile_size, mixed_tile_size, dim) + else: + qkv = jnp.pad(qkv.astype(jnp.float32), ((0, batch_padding_size), (0, 0))) + qkv = qkv.reshape(padded_batch_size, 1, -1) b = b.reshape(padded_batch_size, 1, -1) a = a.reshape(padded_batch_size, 1, -1) @@ -500,6 +492,7 @@ def call_kernel( v_head_dim=d_v, has_seg_ids=has_seg_ids, use_qk_norm_in_gdn=use_qk_norm_in_gdn, + is_prefill_only=is_prefill_only, dtypes=config.Dtypes( act_in=act_in_dtype, act_out=act_out_dtype, @@ -534,7 +527,7 @@ def call_kernel( input_output_aliases = {len(metadata_obj) + 3: 1, len(metadata_obj) + 4: 2} out_shape = cfg.get_out_shape() - if in_act is None and zero_initialize_out: + if in_act is None and zero_initialize_out and mode != config.GDNMode.PER_SEQ_STATES: in_act = jnp.zeros_like(out_shape) if in_act is not None: out_shape = in_act @@ -547,7 +540,12 @@ def call_kernel( cfg.dtypes.compute, ) - if mode == config.GDNMode.PER_SEQ: + if mode == config.GDNMode.PER_SEQ_STATES: + # No output, so the states move up to outputs 0 and 1. + input_output_aliases = {len(metadata_obj) + 3: 0, len(metadata_obj) + 4: 1} + out_shape_tuple = (in_conv_state, in_recurrent_state, t_inv_shape) + out_specs_tuple = (hbm_spec, hbm_spec, hbm_spec) + elif mode == config.GDNMode.PER_SEQ: chunk_states_shape = jax.ShapeDtypeStruct( (num_chunks, cfg.num_v_heads, cfg.kq_head_dim, cfg.v_head_dim), cfg.dtypes.compute, @@ -611,11 +609,15 @@ def call_kernel( out_conv_state = conv_state out_recurrent_state = recurrent_state - out_act, out_conv_state, out_recurrent_state, t_inv, chunk_states = call_kernel( - out_conv_state, out_recurrent_state, out_act, config.GDNMode.PER_SEQ - ) + if prefill_mode == config.GDNMode.PER_SEQ_STATES: + out_conv_state, out_recurrent_state, t_inv = call_kernel(out_conv_state, out_recurrent_state, None, prefill_mode) + out_act = chunk_states = None + else: + out_act, out_conv_state, out_recurrent_state, t_inv, chunk_states = call_kernel( + out_conv_state, out_recurrent_state, out_act, prefill_mode + ) + out_act = out_act.reshape(padded_batch_size, -1)[:batch_size] - out_act = out_act.reshape(padded_batch_size, -1)[:batch_size] out_conv_state = out_conv_state.astype(conv_out_dtype) out_conv_state = out_conv_state.reshape(conv_state_shape) if has_seg_ids and segment_ids is not None: @@ -630,3 +632,158 @@ def call_kernel( out_recurrent_state = out_recurrent_state.astype(recurrent_out_dtype) return out_act, (out_conv_state, out_recurrent_state), t_inv, chunk_states + + +@jax.jit( + donate_argnames=("conv_state", "recurrent_state"), + static_argnames=( + "n_kq", + "n_v", + "d_k", + "d_v", + "kernel_size", + "decode_tile_size", + "mixed_tile_size", + "zero_initialize_out", + "compute_precision", + "is_prefill_only", + "use_qk_norm_in_gdn", + ), +) +def fused_conv1d_gdn( + qkv: jax.Array, # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size] + b: jax.Array, # [batch_size, n_v] + a: jax.Array, # [batch_size, n_v] + conv_state: jax.Array, # [num_seqs + 1, kernel_size - 1, dim_size] + recurrent_state: jax.Array, # [num_seqs + 1, nv, dk, dv] + conv_weight: jax.Array, # [kernel_size - 1, dim_size] + conv_bias: jax.Array | None, # [dim_size] + a_log: jax.Array, # [n_v] + dt_bias: jax.Array, # [n_v] + query_start_loc: jax.Array, # [num_seqs + 1] + state_indices: jax.Array, # [num_seqs] + distribution: jax.Array, # [3] + seq_lens: jax.Array, # [num_seqs] + *, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + kernel_size: int, + zero_initialize_out: bool = True, + compute_precision: jnp.dtype = jnp.float32.dtype, + decode_tile_size: int | None = None, + mixed_tile_size: int | None = None, + is_prefill_only: bool = False, + use_qk_norm_in_gdn: bool = True, + segment_ids: jax.Array | None = None, + conv_halo_seg: jax.Array | None = None, + init_seg: jax.Array | None = None, +) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array, jax.Array]: + """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states).""" + return _fused_conv1d_gdn( + qkv, + b, + a, + conv_state, + recurrent_state, + conv_weight, + conv_bias, + a_log, + dt_bias, + query_start_loc, + state_indices, + distribution, + seq_lens, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + zero_initialize_out=zero_initialize_out, + compute_precision=compute_precision, + decode_tile_size=decode_tile_size, + mixed_tile_size=mixed_tile_size, + is_prefill_only=is_prefill_only, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + segment_ids=segment_ids, + conv_halo_seg=conv_halo_seg, + init_seg=init_seg, + ) + + +@jax.jit( + donate_argnames=("conv_state", "recurrent_state"), + static_argnames=( + "n_kq", + "n_v", + "d_k", + "d_v", + "kernel_size", + "mixed_tile_size", + "compute_precision", + "use_qk_norm_in_gdn", + ), +) +def fused_conv1d_gdn_states( + qkv: jax.Array, # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size] + b: jax.Array, # [batch_size, n_v] + a: jax.Array, # [batch_size, n_v] + conv_state: jax.Array, # [num_seqs + 1, kernel_size - 1, dim_size] + recurrent_state: jax.Array, # [num_seqs + 1, nv, dk, dv] + conv_weight: jax.Array, # [kernel_size - 1, dim_size] + conv_bias: jax.Array | None, # [dim_size] + a_log: jax.Array, # [n_v] + dt_bias: jax.Array, # [n_v] + query_start_loc: jax.Array, # [num_seqs + 1] + state_indices: jax.Array, # [num_seqs] + distribution: jax.Array, # [3] + seq_lens: jax.Array, # [num_seqs] + *, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + kernel_size: int, + mixed_tile_size: int, + compute_precision: jnp.dtype = jnp.float32.dtype, + use_qk_norm_in_gdn: bool = True, + segment_ids: jax.Array | None = None, + conv_halo_seg: jax.Array | None = None, + init_seg: jax.Array | None = None, +) -> tuple[tuple[jax.Array, jax.Array], jax.Array]: + """Prefill-only fused_conv1d_gdn that skips the output and chunk states, returning (states, t_inv). + + For callers that only carry the state forward, such as pass 1 of sequence-sharded context + parallelism: the kernel neither computes the output nor writes it or the chunk states to HBM. The + states and t_inv are bit-identical to fused_conv1d_gdn's. + """ + _, states, t_inv, _ = _fused_conv1d_gdn( + qkv, + b, + a, + conv_state, + recurrent_state, + conv_weight, + conv_bias, + a_log, + dt_bias, + query_start_loc, + state_indices, + distribution, + seq_lens, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + compute_precision=compute_precision, + mixed_tile_size=mixed_tile_size, + is_prefill_only=True, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + segment_ids=segment_ids, + conv_halo_seg=conv_halo_seg, + init_seg=init_seg, + prefill_mode=config.GDNMode.PER_SEQ_STATES, + ) + return states, t_inv diff --git a/src/maxtext/kernels/ragged/glu_rows.py b/src/maxtext/kernels/ragged/glu_rows.py new file mode 100644 index 0000000000..b69d422a10 --- /dev/null +++ b/src/maxtext/kernels/ragged/glu_rows.py @@ -0,0 +1,167 @@ +# 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. + +"""SiLU-GLU of the prefused gate/up GEMM output, and its gradient, over the valid rows of the expert buffer only. + +Only rows ``[0, limit)`` of the packed ``[buffer_size, 2 * mlp]`` buffer hold tokens. XLA's fusions stream the +whole buffer, and the backward writes the two half gradients and then concatenates them, one more pass. These +kernels visit the row blocks up to the one holding row ``limit - 1`` (later blocks map to it: no fetch, no +write, those output rows are left as they are), and the backward writes both halves of its block at once. +Consumers never read the rows left unwritten: gmm_v2 reads rows up to the sublane tile of row ``limit - 1``, +tgmm_v2 masks them. + +Both kernels run the op sequence of XLA's fusions for ``silu(gate) * up`` and its autodiff gradient (negate, +exp, + 1, divide, multiply) the way XLA runs those bf16 fusions on a TPU: in f32, rounding only the outputs. +""" + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +# Rows per grid step. A (1024, 1024) bf16 block is 2 MiB. A partial last block is fine: Pallas drops its +# out-of-range writes. +_BLOCK_ROWS = 1024 +# Rows per inner-loop tile. +_TILE_ROWS = 64 + + +def _sigmoid(x): + """XLA's expansion of logistic: 1 / (exp(-x) + 1).""" + return 1.0 / (jnp.exp(-x) + 1.0) + + +def _glu_fwd_tile(gate, up): + """`silu(gate) * up` as XLA's fusion computes it: in f32, rounded once to the output dtype.""" + g = gate.astype(jnp.float32) + return ((g * _sigmoid(g)) * up.astype(jnp.float32)).astype(gate.dtype) + + +def _glu_bwd_tile(gate, up, d_out): + """The autodiff gradient of `_glu_fwd_tile` as XLA's fusion computes it (f32, rounded once), -0 as +0.""" + g, u, d = gate.astype(jnp.float32), up.astype(jnp.float32), d_out.astype(jnp.float32) + sig = _sigmoid(g) + d_up_in = d * u + d_gate = d_up_in * sig + (g * d_up_in) * (sig * (1.0 - sig)) + d_up = (g * sig) * d + d_gate = jnp.where(d_gate == 0, 0.0, d_gate) + d_up = jnp.where(d_up == 0, 0.0, d_up) + return d_gate.astype(gate.dtype), d_up.astype(gate.dtype) + + +def kernel_supported(fused): + rows, width = fused.shape + return fused.dtype == jnp.bfloat16 and rows % 16 == 0 and width % 256 == 0 + + +def _last_block(limit): + """The row block holding row ``limit - 1`` (block 0 if there are no valid rows).""" + return jnp.maximum(pl.cdiv(limit, _BLOCK_ROWS) - 1, 0) + + +def _row_map(column_block): + def index_map(i, limit_ref): + return jnp.minimum(i, _last_block(limit_ref[0])), column_block + + return index_map + + +def _run(kernel, operands, in_specs, out_spec, out_shape, *, limit, flops, transcendentals, interpret, name): + rows = out_shape.shape[0] + bytes_accessed = sum(x.size * x.dtype.itemsize for x in operands) + out_shape.size * out_shape.dtype.itemsize + return pl.pallas_call( + kernel, + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=1, + grid=(pl.cdiv(rows, _BLOCK_ROWS),), + in_specs=in_specs, + out_specs=out_spec, + ), + out_shape=out_shape, + compiler_params=pltpu.CompilerParams(dimension_semantics=("arbitrary",), vmem_limit_bytes=48 * 1024 * 1024), + cost_estimate=pl.CostEstimate(flops=flops, transcendentals=transcendentals, bytes_accessed=bytes_accessed), + interpret=interpret, + name=name, + )(jnp.reshape(jnp.asarray(limit).astype(jnp.int32), (1,)), *operands) + + +def _visit_rows(i, limit, tile_fn): + """Runs `tile_fn(rows)` over the row tiles of block i if the block holds valid rows (or is block 0).""" + + @pl.when(i <= _last_block(limit)) + def _(): + def body(t, carry): + tile_fn(pl.ds(pl.multiple_of(t * _TILE_ROWS, _TILE_ROWS), _TILE_ROWS)) + return carry + + lax.fori_loop(0, _BLOCK_ROWS // _TILE_ROWS, body, 0) + + +def _glu_fwd_kernel(limit_ref, gate_ref, up_ref, out_ref): + def tile(rows): + out_ref[rows, :] = _glu_fwd_tile(gate_ref[rows, :], up_ref[rows, :]) + + _visit_rows(pl.program_id(0), limit_ref[0], tile) + + +def _glu_bwd_kernel(limit_ref, gate_ref, up_ref, d_ref, out_ref): + mlp = gate_ref.shape[1] + + def tile(rows): + d_gate, d_up = _glu_bwd_tile(gate_ref[rows, :], up_ref[rows, :], d_ref[rows, :]) + out_ref[rows, :mlp] = d_gate + out_ref[rows, mlp:] = d_up + + _visit_rows(pl.program_id(0), limit_ref[0], tile) + + +def glu_fwd_rows(fused, limit, *, interpret=False): + """``silu(gate) * up`` of rows ``[0, limit)`` of ``fused = [gate | up]``, ``[rows, 2 * mlp]`` -> ``[rows, mlp]``.""" + rows, width = fused.shape + mlp = width // 2 + return _run( + _glu_fwd_kernel, + (fused, fused), + [pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(0)), pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(1))], + pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(0)), + jax.ShapeDtypeStruct((rows, mlp), fused.dtype), + limit=limit, + flops=4 * rows * mlp, + transcendentals=2 * rows * mlp, + interpret=interpret, + name="glu_fwd_rows", + ) + + +def glu_bwd_rows(fused, d_out, limit, *, interpret=False): + """The ``[rows, 2 * mlp]`` gradient ``[d_gate | d_up]`` of `glu_fwd_rows` at rows ``[0, limit)``, -0 as +0.""" + rows, width = fused.shape + mlp = width // 2 + return _run( + _glu_bwd_kernel, + (fused, fused, d_out), + [ + pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(0)), + pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(1)), + pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(0)), + ], + pl.BlockSpec((_BLOCK_ROWS, width), _row_map(0)), + jax.ShapeDtypeStruct((rows, width), fused.dtype), + limit=limit, + flops=12 * rows * mlp, + transcendentals=2 * rows * mlp, + interpret=interpret, + name="glu_bwd_rows", + ) diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index dc563cef68..45f2a34bde 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -19,6 +19,7 @@ import jax.numpy as jnp from maxtext.kernels.ragged.ragged_gather import ragged_gather from maxtext.kernels.ragged.ragged_gather_reduce_v2 import ragged_gather_reduce +from maxtext.kernels.ragged.unsort_bwd_rows import unsort_bwd_rows def ring_ragged_sort( @@ -72,6 +73,8 @@ def ring_ragged_sort( only the tokens destined for local experts, padded with zeros elsewhere. - 1D tensor ``group_sizes_local`` tracking expert token counts. - 1D tensor ``topk_argsort_revert_indices`` for inverse routing. + - 1D tensor ``topk_argsort_indices``, the sort permutation (the inverse of + ``topk_argsort_revert_indices``), for the backward of :func:`ring_ragged_unsort`. """ @jax.custom_vjp @@ -88,17 +91,16 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): topk_indices_flat = topk_indices_local.flatten() # num_tokens_local x topk topk_argsort_indices = jnp.argsort(topk_indices_flat) # num_tokens_local x topk - token_indices = jnp.arange(num_tokens_local, dtype=jnp.int32).repeat(topk) # num_tokens_local x topk - token_indices_sorted = token_indices[topk_argsort_indices] # num_tokens_local x topk - group_sizes_local = jax.nn.one_hot(topk_indices_flat, num_experts, dtype=jnp.int32).sum(axis=0) # GLOBAL_NUM_EXPERTS topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices) # num_tokens_local x topk # Named for the `moe_route` remat key. Saving these integer results lets the backward replay # the token gather, whose output the expert GMMs need again, without replaying both argsorts. - token_indices_sorted = checkpoint_name(token_indices_sorted, "moe_route") + topk_argsort_indices = checkpoint_name(topk_argsort_indices, "moe_route") group_sizes_local = checkpoint_name(group_sizes_local, "moe_route") topk_argsort_revert_indices = checkpoint_name(topk_argsort_revert_indices, "moe_route") + # Flat slot j belongs to token j // topk. + token_indices_sorted = topk_argsort_indices // topk # num_tokens_local x topk shard_idx = jax.lax.axis_index(ep_name) local_num_experts = num_experts // ep_size @@ -145,7 +147,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): use_single_sparsecore=use_single_sparsecore, ) - out = (x, group_sizes_local, topk_argsort_revert_indices) + out = (x, group_sizes_local, topk_argsort_revert_indices, topk_argsort_indices) res = ( topk_argsort_revert_indices, @@ -173,7 +175,7 @@ def _ring_ragged_sort_bwd(res, g_out): local_buffer_size, _, ) = res - g_x, _, _ = g_out + g_x, _, _, _ = g_out # Restrict to the [start, end) source range via a validity bitmask. The # ragged kernel packs valid rows to the front of each row-partition and # only iterates over the populated prefix, so we hand it the mask directly @@ -236,6 +238,7 @@ def ring_ragged_unsort( sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, + topk_argsort_indices, topk, local_num_experts, ep_name, @@ -268,6 +271,8 @@ def ring_ragged_unsort( sorted_tokens_local: 2D ``[buffer_size, hidden]`` output tensor from local experts. group_sizes_local: 1D tensor tracking the token loads per expert. topk_argsort_revert_indices: 1D permutation restoring flat token positions. + topk_argsort_indices: 1D sort permutation from :func:`ring_ragged_sort`, the inverse of + ``topk_argsort_revert_indices``. Read by the backward only. topk: scalar ``int`` representing the routing top-k factor. local_num_experts: scalar ``int`` representing the count of experts hosted on this shard. ep_name: ``str`` identifying the expert parallel axis name. @@ -284,6 +289,7 @@ def _ring_ragged_unsort( sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, + topk_argsort_indices, topk_weights_flat, ): """Unsort and scatter activations.""" @@ -291,6 +297,7 @@ def _ring_ragged_unsort( sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, + topk_argsort_indices, topk_weights_flat, )[0] @@ -299,6 +306,7 @@ def _ring_ragged_unsort_fwd( sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, + topk_argsort_indices, topk_weights_flat, ): """Executes unsorting sending tokens back.""" @@ -361,6 +369,7 @@ def _ring_ragged_unsort_fwd( res = ( sorted_tokens_local, topk_argsort_revert_indices, + topk_argsort_indices, topk_weights_flat, shard_output_start, shard_output_end, @@ -388,6 +397,7 @@ def _ring_ragged_unsort_bwd(res, g_out): ( sorted_tokens_local, topk_argsort_revert_indices, + topk_argsort_indices, topk_weights_flat, shard_output_start, shard_output_end, @@ -396,9 +406,9 @@ def _ring_ragged_unsort_bwd(res, g_out): g_hidden_states_local = g_out n = topk_argsort_revert_indices.shape[0] - # Build the inverse permutation idx_inv such that idx_inv[j] = i + # The forward sort's permutation is the inverse permutation: idx_inv[j] = i # where revert[i] = j. - idx_inv = jnp.argsort(topk_argsort_revert_indices) + idx_inv = topk_argsort_indices # Handle the same two buffering modes for backward pass. # ragged_gather does the fan-out, by indexing into the un-expanded @@ -442,16 +452,14 @@ def _ring_ragged_unsort_bwd(res, g_out): bytes_accessed_override=gather_bytes_accessed_override, use_single_sparsecore=use_single_sparsecore, ) - # Mask out gradients for elements beyond the valid limit of the local buffer. + # Only rows [0, limit) of the local buffer hold tokens: weight their gradients and take their dots. limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size) - mask = jnp.arange(buffer_size) < limit - gathered = jnp.where(mask[:, None], gathered, 0.0) - grad_sorted_tokens = (gathered * sliced_weights[:, None]).astype(gathered.dtype) + grad_sorted_tokens, dot_local = unsort_bwd_rows(gathered, sorted_tokens_local, sliced_weights, limit) # Scatter the per-slot dot back to flat slot order; dropped slots stay zero. - dot_local = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local.astype(jnp.float32), axis=-1) + mask = jnp.arange(buffer_size) < limit slots = jnp.where(mask, sliced_idx_inv, n) grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop") - return grad_sorted_tokens, None, None, grad_topk_weights + return grad_sorted_tokens, None, None, None, grad_topk_weights _ring_ragged_unsort.defvjp(_ring_ragged_unsort_fwd, _ring_ragged_unsort_bwd) @@ -462,6 +470,7 @@ def _ring_ragged_unsort_bwd(res, g_out): sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, + topk_argsort_indices, topk_weights_flat, ) diff --git a/src/maxtext/kernels/ragged/unsort_bwd_rows.py b/src/maxtext/kernels/ragged/unsort_bwd_rows.py new file mode 100644 index 0000000000..148870b9cf --- /dev/null +++ b/src/maxtext/kernels/ragged/unsort_bwd_rows.py @@ -0,0 +1,179 @@ +# 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. + +"""The row-wise part of the ring unsort backward, on the TensorCore, over the valid rows of the buffer only. + +Only rows ``[0, limit)`` of the packed ``[buffer_size, hidden]`` expert buffer hold tokens; the rest is +ragged-buffer padding. XLA's fusion of the same three lines streams the whole buffer. This kernel visits the +row blocks up to the one holding row ``limit - 1``: the later blocks map to that block, so they are neither +fetched nor written, and those output rows are left as they are. No consumer reads them: the gmm_v2 / +tgmm_v2 kernels read rows up to the sublane tile of row ``limit - 1`` (tgmm masks them), and the weight-dot +scatter drops rows ``>= limit``. +""" + +import functools + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +# Rows per grid step: one (8, 128) tile of the lane-dense [rows / 128, 128] weights and dot. A partial last block +# is fine: Pallas drops its out-of-range writes, and its rows are past `limit`. +_BLOCK_ROWS = 1024 +# Hidden columns per grid step, at most. A (1024, 2048) bf16 block is 4 MiB; 3 double-buffered fit v7x VMEM. +_BLOCK_COLS = 2048 +_LANES = 128 + + +def unsort_bwd_rows_reference(gathered, sorted_tokens, weights, limit): + """XLA version: the weighted gradient and the f32 row dot of every row, zeroing rows ``>= limit``.""" + mask = jnp.arange(gathered.shape[0]) < limit + gathered = jnp.where(mask[:, None], gathered, 0.0) + grad = (gathered * weights[:, None]).astype(gathered.dtype) + dot = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens.astype(jnp.float32), axis=-1) + return grad, dot + + +def _block_cols(hidden): + return hidden if hidden <= _BLOCK_COLS else _BLOCK_COLS + + +def _kernel_supported(gathered, sorted_tokens, weights): + rows, hidden = gathered.shape + return ( + gathered.dtype == jnp.bfloat16 + and sorted_tokens.dtype == jnp.bfloat16 + and sorted_tokens.shape == gathered.shape + and weights.shape == (rows,) + and rows % _LANES == 0 + and hidden % _LANES == 0 + and hidden % _block_cols(hidden) == 0 + ) + + +def _last_block(limit): + """The row block holding row ``limit - 1`` (block 0 if there are no valid rows).""" + return jnp.maximum(pl.cdiv(limit, _BLOCK_ROWS) - 1, 0) + + +def _unsort_bwd_rows_kernel(limit_ref, g_ref, y_ref, w_ref, grad_ref, dot_ref, *, num_col_blocks): # pylint: disable=too-many-positional-arguments + """One (1024, cols) block: grad = bf16(f32(g) * w) and the partial f32 row dot, in 128 x 128 tiles.""" + i, j = pl.program_id(0), pl.program_id(1) + limit = limit_ref[0] + row0 = i * _BLOCK_ROWS + num_col_tiles = g_ref.shape[1] // _LANES + + def row_tile(masked, r, carry): + rows = pl.ds(pl.multiple_of(r * _LANES, _LANES), _LANES) + # Row a of the tile gets weight w[row0 + 128 r + a], broadcast along the lanes. + w_col = jnp.transpose(jnp.broadcast_to(w_ref[pl.ds(r, 1), :], (_LANES, _LANES))) + if masked: + row_ids = row0 + r * _LANES + lax.broadcasted_iota(jnp.int32, (_LANES, _LANES), 0) + valid = row_ids < limit + acc = jnp.zeros((_LANES, _LANES), jnp.float32) + for c in range(num_col_tiles): + cols = pl.ds(c * _LANES, _LANES) + g = g_ref[rows, cols].astype(jnp.float32) + if masked: + g = jnp.where(valid, g, 0.0) + grad_ref[rows, cols] = (g * w_col).astype(grad_ref.dtype) + acc = acc + g * y_ref[rows, cols].astype(jnp.float32) + # Lane-dense: entry a of the row sum over the transposed tile is row a's dot over this block's columns. + part = jnp.sum(jnp.transpose(acc), axis=0, keepdims=True) + if num_col_blocks > 1: + part = part + jnp.where(j == 0, 0.0, dot_ref[pl.ds(r, 1), :]) + dot_ref[pl.ds(r, 1), :] = part + return carry + + num_row_tiles = _BLOCK_ROWS // _LANES + + @pl.when(row0 + _BLOCK_ROWS <= limit) + def _full(): + lax.fori_loop(0, num_row_tiles, functools.partial(row_tile, False), 0) + + @pl.when(jnp.logical_and(row0 + _BLOCK_ROWS > limit, i == _last_block(limit))) + def _partial(): + lax.fori_loop(0, num_row_tiles, functools.partial(row_tile, True), 0) + + +def _unsort_bwd_rows_pallas(gathered, sorted_tokens, weights, limit, *, interpret=False): + """The Pallas TensorCore kernel; see `unsort_bwd_rows`.""" + rows, hidden = gathered.shape + block_cols = _block_cols(hidden) + num_row_blocks = pl.cdiv(rows, _BLOCK_ROWS) + num_col_blocks = hidden // block_cols + + def data_map(i, j, limit_ref): + last = _last_block(limit_ref[0]) + return jnp.minimum(i, last), jnp.where(i <= last, j, num_col_blocks - 1) + + def row_map(i, j, limit_ref): + del j + return jnp.minimum(i, _last_block(limit_ref[0])), 0 + + data_spec = pl.BlockSpec((_BLOCK_ROWS, block_cols), data_map) + lane_dense_spec = pl.BlockSpec((_BLOCK_ROWS // _LANES, _LANES), row_map) + grad, dot = pl.pallas_call( + functools.partial(_unsort_bwd_rows_kernel, num_col_blocks=num_col_blocks), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=1, + grid=(num_row_blocks, num_col_blocks), + in_specs=[data_spec, data_spec, lane_dense_spec], + out_specs=[data_spec, lane_dense_spec], + ), + out_shape=[ + jax.ShapeDtypeStruct((rows, hidden), gathered.dtype), + jax.ShapeDtypeStruct((rows // _LANES, _LANES), jnp.float32), + ], + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary"), + vmem_limit_bytes=48 * 1024 * 1024, + ), + cost_estimate=pl.CostEstimate( + flops=4 * rows * hidden, + transcendentals=0, + bytes_accessed=3 * rows * hidden * gathered.dtype.itemsize + 8 * rows, + ), + interpret=interpret, + name="unsort_bwd_rows", + )( + jnp.reshape(limit.astype(jnp.int32), (1,)), + gathered, + sorted_tokens, + jnp.reshape(weights.astype(jnp.float32), (rows // _LANES, _LANES)), + ) + return grad, jnp.reshape(dot, (rows,)) + + +def unsort_bwd_rows(gathered, sorted_tokens, weights, limit, *, interpret=None): + """The weighted gradient and the routing-weight row dot of the valid rows ``[0, limit)`` of the expert buffer. + + Args: + gathered: ``[rows, hidden]`` bf16 output gradient, gathered into the packed expert buffer order. + sorted_tokens: ``[rows, hidden]`` bf16 expert outputs in the same order. + weights: ``[rows]`` f32 routing weight of each buffer row. + limit: int32 scalar, the number of valid rows. + interpret: Pallas interpret mode (tests). None runs the kernel on a TPU target and the XLA version elsewhere. + + Returns: + ``(grad, dot)``: for rows ``< limit``, ``grad = bf16(f32(gathered) * weights)`` (bit-identical to the XLA + version) and ``dot = sum(f32(gathered) * f32(sorted_tokens), -1)`` (the summation order differs). Rows from + ``limit`` to the end of its 1024-row block are computed as zero rows; later rows are not written. + """ + use_kernel = interpret if interpret is not None else pltpu.is_tpu_device() + if not use_kernel or not _kernel_supported(gathered, sorted_tokens, weights): + return unsort_bwd_rows_reference(gathered, sorted_tokens, weights, limit) + return _unsort_bwd_rows_pallas(gathered, sorted_tokens, weights, jnp.asarray(limit), interpret=bool(interpret)) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index dda2f599eb..cf2bafa5c8 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -30,6 +30,7 @@ from jax.experimental.compute_on import compute_on from jax.experimental import xla_metadata from jax.experimental.layout import Layout, with_layout_constraint +from jax.experimental.pallas import tpu as pltpu import jax.numpy as jnp from jax.sharding import Mesh, NamedSharding from jax.sharding import PartitionSpec as P @@ -37,6 +38,7 @@ from maxtext.common.common_types import ShardMode from maxtext.kernels import megablox as mblx from maxtext.kernels import sort_activations +from maxtext.kernels.ragged import glu_rows from maxtext.kernels.ragged.ragged_sort import a2a_ragged_sort from maxtext.kernels.ragged.ragged_sort import a2a_ragged_unsort from maxtext.kernels.ragged.ragged_sort import ring_ragged_sort @@ -109,6 +111,9 @@ class RouteOutput: # log_moe_drops=True on the ring-of-experts ragged path only (None otherwise): the drop audit # `(dropped_assignments, max_load_ratio)`, see `RoutedMoE.permute`. Logging only. moe_audit: Optional[tuple[jax.Array, jax.Array]] = None + # Ring-of-experts ragged path only (None otherwise): the sort permutation, the inverse of + # `sorted_selected_experts` there. The unsort backward reads it instead of argsorting that. + topk_argsort_indices: Optional[jax.Array] = None def combine_moe_audits(audits): @@ -245,6 +250,23 @@ def _forward_optimization_barrier_bwd(_, cotangents): _forward_optimization_barrier.defvjp(_forward_optimization_barrier_fwd, _forward_optimization_barrier_bwd) +@jax.custom_vjp +def _backward_optimization_barrier(operands): + """The identity on the forward values and `jax.lax.optimization_barrier` on their cotangents.""" + return operands + + +def _backward_optimization_barrier_fwd(operands): + return operands, None + + +def _backward_optimization_barrier_bwd(_, cotangents): + return (jax.lax.optimization_barrier(cotangents),) + + +_backward_optimization_barrier.defvjp(_backward_optimization_barrier_fwd, _backward_optimization_barrier_bwd) + + def barrier_tie(x, dep, backward: bool = False): """Returns `(x, dep)` unchanged, tied together by `jax.lax.optimization_barrier`. @@ -301,6 +323,69 @@ def forward_row_major(x: jax.Array) -> jax.Array: return _forward_layout_constraint(x, tuple(range(x.ndim))) +def _split_gate_up(fused): + """The gate and up halves of the prefused gate/up GEMM output, with the remat names of `gmm_up`.""" + n = fused.shape[-1] // 2 + gate = adc.checkpoint_name(adc.checkpoint_name(fused[..., :n], "mlpwi_0"), "moe_mlpwi_0") + up = adc.checkpoint_name(fused[..., n:], "moe_mlpwi_1") + return gate, up + + +def fused_glu(fused, glu, limit=None): + """`glu(gate, up)` of the two halves of `fused`, the [..., 2 * mlp] output of the prefused gate/up GEMM. + + Autodiff transposes each half's slice into a zero-padded [..., 2 * mlp] gradient and adds the two, one more + full pass over the gradient. This backward computes the same two half gradients, with `jax.vjp` of `glu`, and + concatenates them, which XLA writes straight from the activation's backward fusion. + + `limit`, if given, says that only rows [0, limit) of the [rows, 2 * mlp] expert buffer `fused` hold tokens and + that `glu` is the bf16 SiLU GLU. On a TPU, Pallas kernels then compute the activation and its gradient on those + rows only, bit-identical to the XLA fusions, and write the gradient's two halves at once; the rows past the + 1024-row block of row limit - 1 are left unwritten in both results. + """ + if limit is not None and pltpu.is_tpu_device() and glu_rows.kernel_supported(fused): + return _silu_glu_rows(fused, limit) + return _fused_glu(fused, glu) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1,)) +def _fused_glu(fused, glu): + return glu(*_split_gate_up(fused)) + + +def _fused_glu_fwd(fused, glu): + gate, up = _split_gate_up(fused) + return glu(gate, up), (gate, up) + + +def _fused_glu_bwd(glu, residuals, cotangent): + _, glu_vjp = jax.vjp(glu, *residuals) + # The zero-padded add turns -0 into +0; so does this, so the gradient is bit-identical to autodiff's. + return (jnp.concatenate([jnp.where(d == 0, 0, d) for d in glu_vjp(cotangent)], axis=-1),) + + +_fused_glu.defvjp(_fused_glu_fwd, _fused_glu_bwd) + + +@jax.custom_vjp +def _silu_glu_rows(fused, limit): + return glu_rows.glu_fwd_rows(fused, limit) + + +def _silu_glu_rows_fwd(fused, limit): + # The kernels read both halves in place, so the residual is the whole buffer, under the names of both halves. + fused = adc.checkpoint_name(adc.checkpoint_name(adc.checkpoint_name(fused, "mlpwi_0"), "moe_mlpwi_0"), "moe_mlpwi_1") + return glu_rows.glu_fwd_rows(fused, limit), (fused, limit) + + +def _silu_glu_rows_bwd(residuals, cotangent): + fused, limit = residuals + return glu_rows.glu_bwd_rows(fused, cotangent, limit), None + + +_silu_glu_rows.defvjp(_silu_glu_rows_fwd, _silu_glu_rows_bwd) + + def get_batchsplit_init_kernel_axes(): return ( ("expert_only", "embed_moe", None), @@ -384,6 +469,37 @@ def fold_padding_into_forced_routing(forced_routed_experts, token_mask, num_expe return jnp.where(is_pad, jnp.asarray(PAD_TOKEN_ROUTED_EXPERT, base.dtype), base) +def _top_k_indices_where_needed(probs, needed, k): + """`jax.lax.top_k(probs, k)[1]` on the tokens where `needed` is set; other tokens may get any ids. + + Under router replay almost every token is forced or padding. When at most 256 tokens need the + router's own choice, top_k sorts only those rows (gathered, then scattered back) instead of every + row; otherwise it sorts every row. top_k is per row, so the needed rows are exact either way. + + Args: + probs: `[..., num_experts]` router probabilities. + needed: `[..., 1]` True for tokens whose ids are read. + k: top_k. + """ + max_rows = 256 + flat = probs.reshape(-1, probs.shape[-1]) + if flat.shape[0] <= max_rows: + return jax.lax.top_k(probs, k)[1] + needed = needed.reshape(-1) + + def compact(flat, needed): + # Fill entries repeat row 0, which only writes row 0's own top_k to row 0. + rows = jnp.nonzero(needed, size=max_rows)[0] + return jnp.zeros((flat.shape[0], k), jnp.int32).at[rows].set(jax.lax.top_k(flat[rows], k)[1]) + + def full(flat, needed): + del needed + return jax.lax.top_k(flat, k)[1] + + indices = jax.lax.cond(jnp.sum(needed) <= max_rows, compact, full, flat, needed) + return indices.reshape(probs.shape[:-1] + (k,)) + + def _top_2_in_group_sum(scores_grouped: jax.Array) -> jax.Array: """Sums the two largest scores along the last axis, accumulating in float32. @@ -667,6 +783,96 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. return output, pre_bias_logits +def _param_path(path) -> tuple: + """A params pytree path as the module path keys `nnx.iter_modules` yields, plus the param name. + + Drops the trailing `GetAttrKey("value")` of each Variable leaf. + """ + keys = [] + for k in path: + if isinstance(k, jax.tree_util.DictKey): + keys.append(k.key) + elif isinstance(k, jax.tree_util.SequenceKey): + keys.append(k.idx) + return tuple(keys) + + +def _gathered_kernels(layers: nnx.Module) -> dict: + """{kernel path within a scanned layer stack: its module}, for the kernels gathered once per stack.""" + kernels = {} + for path, module in nnx.iter_modules(layers): + if isinstance(module, GateLogit): + kernels[path + ("kernel",)] = module + elif isinstance(module, attentions.Attention) and hasattr(module, "query"): + # The backward pass otherwise gathers the query kernel twice per layer: for the rematerialized + # projection and for the input gradient. Under fused_qkv there is no `query` (it is `qkv_proj`). + kernels[path + ("query", "kernel")] = module.query + return kernels + + +def gather_scanned_kernels_once(layers: nnx.Module, params, mesh: Mesh, shard_mode): + """Gathers small FSDP-sharded kernels of a scanned layer stack once, before the scan, not in every layer. + + The MoE router kernel (embed x num_experts, f32 under float32_gate_logits) is small but FSDP-sharded, + so every layer all-gathers it on the critical path, synchronously on TPUs that run all-gathers on the + TensorCore, and the backward pass gathers it again. Replicating the stacked kernels up front is one + all-gather for all layers. `params` is the stack's Param state with the scan axis leading; pair with + `reduce_scatter_gathered_grads` inside the scan so each layer's gradient still lands sharded. + """ + gathered = _gathered_kernels(layers) + replicated = NamedSharding(mesh, P()) + + def replicate(path, x): + if _param_path(path) not in gathered: + return x + if shard_mode == ShardMode.EXPLICIT: + return jax.sharding.reshard(x, replicated) + return jax.lax.with_sharding_constraint(x, replicated) + + return jax.tree_util.tree_map_with_path(replicate, params) + + +def gathered_kernel_grad_shardings(layers: nnx.Module, mesh: Mesh, rules) -> dict: + """{kernel path within one layer: the kernel's own sharding} for the kernels `gather_scanned_kernels_once` gathers.""" + return { + path: NamedSharding(mesh, logical_to_mesh_axes(tuple(module.kernel_axes), mesh, rules)) + for path, module in _gathered_kernels(layers).items() + } + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1,)) +def _identity_with_grad_sharding(x, sharding): + """The identity on `x`; its cotangent is constrained to `sharding`.""" + return x + + +def _identity_with_grad_sharding_fwd(x, sharding): + del sharding + return x, None + + +def _identity_with_grad_sharding_bwd(sharding, _, cotangent): + return (jax.lax.with_sharding_constraint(cotangent, sharding),) + + +_identity_with_grad_sharding.defvjp(_identity_with_grad_sharding_fwd, _identity_with_grad_sharding_bwd) + + +def reduce_scatter_gathered_grads(params, shardings: dict): + """Identity on one layer's params whose backward puts the gathered kernels' gradients back on their sharding. + + A layer handed a slice of a replicated kernel otherwise gets a replicated gradient: an all-reduce per + layer, stacked at full size across the scan. Constraining it to the kernel's own sharding makes it a + reduce-scatter into a stacked shard, as when the layer gathered the kernel itself. + """ + + def apply(path, x): + keys = _param_path(path) + return _identity_with_grad_sharding(x, shardings[keys]) if keys in shardings else x + + return jax.tree_util.tree_map_with_path(apply, params) + + class RoutedMoE(nnx.Module): """Implements a routed MoE block.""" @@ -762,7 +968,14 @@ def __init__( self.wi_kernel_axes, self.wo_kernel_axes = get_batchsplit_init_kernel_axes() else: self.wi_kernel_axes = ("exp", "embed_moe", "mlp_moe") - self.wo_kernel_axes = ("exp", "mlp_moe", "embed_moe") + # wo is [experts, mlp, embed]. Mesh rules that define `mlp_moe_out`/`embed_moe_out` (cp-as-ep) shard its mlp + # dim over FSDP instead of its embed dim, so its gradient reduce-scatter scatters a middle dimension, like wi's, + # instead of being transposed to [embed, experts * mlp] and scattered on its major dimension (about 3x slower + # on the SparseCore for the same bytes). + if {"mlp_moe_out", "embed_moe_out"} <= {rule[0] for rule in self.config.logical_axis_rules}: + self.wo_kernel_axes = ("exp", "mlp_moe_out", "embed_moe_out") + else: + self.wo_kernel_axes = ("exp", "mlp_moe", "embed_moe") if self.config.attention in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"): # vLLM uses 'model' as the tensor parallelism axis name @@ -1136,7 +1349,6 @@ def get_topk( valid_token_mask = None if forced_routed_experts is not None: router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) - _, auto_indices = jax.lax.top_k(router_probs, self.num_experts_per_tok) slot_valid = valid_expert_mask(forced_routed_experts, self.num_experts) # A token has real forced routing only if its slots are within [0, num_experts). # Tokens with UNSET_ROUTED_EXPERT (-1) slots must fall back to auto_indices @@ -1144,7 +1356,6 @@ def get_topk( sorted_experts = jnp.sort(forced_routed_experts, axis=-1) has_duplicate = jnp.any(sorted_experts[..., 1:] == sorted_experts[..., :-1], axis=-1, keepdims=True) token_is_forced = jnp.logical_and(jnp.all(slot_valid, axis=-1, keepdims=True), jnp.logical_not(has_duplicate)) - top_k_indices = jnp.where(token_is_forced, forced_routed_experts, auto_indices) if self.config.moe_exclude_padding_from_routing: # A padding token (every slot PAD_TOKEN_ROUTED_EXPERT) is never forced (its slots are not # valid ids) and gets no valid slot, so its weights are zeroed below and its ids become -1. @@ -1152,6 +1363,11 @@ def get_topk( valid_token_mask = jnp.where(token_is_forced, slot_valid, jnp.logical_not(pad_token)) else: valid_token_mask = jnp.where(token_is_forced, slot_valid, jnp.ones_like(slot_valid, dtype=jnp.bool_)) + # Only tokens that are neither forced nor padding read auto_indices: every other id is replaced + # or masked below. + uses_auto = jnp.logical_and(jnp.logical_not(token_is_forced), jnp.any(valid_token_mask, axis=-1, keepdims=True)) + auto_indices = _top_k_indices_where_needed(router_probs, uses_auto, self.num_experts_per_tok) + top_k_indices = jnp.where(token_is_forced, forced_routed_experts, auto_indices) gather_indices = jnp.where(valid_token_mask, top_k_indices, 0) if self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: top_k_weights = jnp.take_along_axis(router_probs, gather_indices, axis=-1).astype(self.dtype) @@ -1420,7 +1636,7 @@ def permute( else: buffer_size = None - sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort( + sorted_inputs, group_size, sorted_selected_experts, topk_argsort_indices = ring_ragged_sort( inputs_2d, topk_indices_2d, self.config.num_experts, @@ -1437,6 +1653,7 @@ def permute( use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: + topk_argsort_indices = None flatten_selected_experts = jnp.ravel(selected_experts) # Must precede roll_to_expert_id: `(-1 - roll) % num_experts` wraps padding @@ -1542,6 +1759,7 @@ def permute( local_group_size, has_overflow, moe_audit, + topk_argsort_indices, ) def unpermute( @@ -1553,6 +1771,7 @@ def unpermute( sequence_length, use_custom_sort_vjp=True, group_sizes=None, + topk_argsort_indices=None, ): """Unpermute tokens to original order and combine weights.""" @@ -1566,6 +1785,7 @@ def unpermute( intermediate, group_sizes, sorted_selected_experts, + topk_argsort_indices, self.num_experts_per_tok, local_num_experts, self._expert_parallelism_name, @@ -2414,6 +2634,7 @@ def _ep_all_gather(z): local_group_sizes, has_overflow, moe_audit, + topk_argsort_indices, ) = self.permute( x, logits, @@ -2439,6 +2660,7 @@ def _ep_all_gather(z): local_group_sizes=local_group_sizes, has_overflow=has_overflow, moe_audit=moe_audit, + topk_argsort_indices=topk_argsort_indices, ), RouteMetadata( expert_shard_id=expert_shard_id, @@ -2472,6 +2694,7 @@ def ra2a_and_route( local_group_sizes, _, _, + _, ) = self.permute( x, logits, @@ -2678,6 +2901,10 @@ def gmm_up( # Weights are stored as (G,K,2N); w0/w1 are adjacent slices so XLA elides this concat. w_fused = jnp.concatenate([w0, w1], axis=-1) out = gmm_fn(x, w_fused, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes) + if not self.config.mlp_bias: + # Returned whole, as (out, None): `_down_projection` splits it in `fused_glu`, whose backward concatenates + # the halves' gradients, and the chunk pipeline's tie 4 holds the GEMM output, not a copy of its first half. + return out, None n = out.shape[-1] // 2 layer_w0, layer_w1 = out[:, :n], out[:, n:] if self.config.mlp_bias and w0_bias is not None and w1_bias is not None: @@ -2931,7 +3158,23 @@ def _up_projection(x, routing, route_metadata, w0, w1, w0_bias, w1_bias, wo_bias def _down_projection(output0, output1, gmm_fn, wo, wo_bias, routing, route_metadata): """The activation and the down projection, per routed token.""" - intermediate_layer = self.apply_ffn_activation(output0, output1) + if output1 is None: # output0 is the prefused gate/up GEMM output (see `gmm_up`) + # The ring's packed buffer holds its tokens in rows [0, valid_token_count); the SiLU GLU kernels skip the + # rest. apply_ffn_activation is the plain SiLU GLU unless the decoder block clips its inputs. + silu_glu = ( + self.config.mlp_activations[0] == "silu" + and self.dtype == jnp.bfloat16 + and self.config.decoder_block != ctypes.DecoderBlockType.GPT_OSS + and not ( + self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4) + and self.config.mlp_activations_limit > 0.0 + ) + ) + packed = self.config.use_ring_of_experts and output0.shape[0] < routing.sorted_selected_experts.shape[0] + limit = valid_token_count(output0, routing, route_metadata) if silu_glu and packed else None + intermediate_layer = fused_glu(output0, self.apply_ffn_activation, limit) + else: + intermediate_layer = self.apply_ffn_activation(output0, output1) wo_gather_axes, wo_tile_size = get_wo_gmm_params() intermediate_output = gmm_fn( intermediate_layer, @@ -2963,6 +3206,7 @@ def _ring_unpermute(intermediate_output, routing, batch_size, sequence_length): sequence_length=sequence_length, use_custom_sort_vjp=self.config.use_custom_sort_vjp, group_sizes=routing.group_sizes, + topk_argsort_indices=routing.topk_argsort_indices, ) return jnp.reshape( output, @@ -3166,6 +3410,11 @@ def _pipelined_route_and_compute(force_dropless): reads those barrier outputs, and XLA removes the reduce-scatter and its gather-reduce. The compiled backward has the collectives of the plain chunk loop. """ + nonlocal x, w0, w1, wo + # The backward hands the layer-input gradient on only once the expert-weight gradients exist, so each + # layer's tgmm runs under its own token reduce-scatters instead of at the end of the backward iteration, + # where every deferred layer's sorted-token buffers would still be live. + x, (w0, w1, wo) = _backward_optimization_barrier((x, (w0, w1, wo))) seq_len = x.shape[1] chunk = seq_len // n_chunks diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 890bd01179..de3d9be6a9 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1071,6 +1071,10 @@ def _apply_layers_sequentially( scan_axis = self.config.param_scan_axis if scan_axis != 0: params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0) if x.ndim > scan_axis else x, params) + gathered_grad_shardings = {} + if kv_caches_stacked is None and not self.config.parameter_memory_host_offload: + params = moe.gather_scanned_kernels_once(layers, params, self.mesh, self.config.shard_mode) + gathered_grad_shardings = moe.gathered_kernel_grad_shardings(layers, self.mesh, self.config.logical_axis_rules) layer_cls = layers.__class__ sig = inspect.signature(layer_cls.__call__) @@ -1105,6 +1109,9 @@ def layer_fn(carry, scanned_vars): kv_cache_layer = None forced_routed_experts_layer = None + if gathered_grad_shardings: + current_params = moe.reduce_scatter_gathered_grads(current_params, gathered_grad_shardings) + if self.config.parameter_memory_host_offload: current_params = jax.tree.map( lambda x: jax.device_put(x, max_utils.device_space()), diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 84c53cf96a..90d06c693f 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1473,6 +1473,10 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): # Reshape back to a single feature dimension for the final projection. # Shape from (B, S, H_v, D_v) -> (B, S, value_dim) gated_output = jnp.reshape(gated_output_reshaped, (batch, seq_len, -1), out_sharding=flat_sharding) + if model_mode == MODEL_MODE_TRAIN: + # Token-minor, as the backward's recomputed projection reads it: left to itself, layout assignment gives + # the forward's operand the norm's head-tiled layout and out_proj re-tiles it on every MXU fetch. + gated_output = moe._forward_layout_constraint(gated_output, (0, 2, 1)) # Final output shape: (B, S, E) output = self.out_proj(gated_output, out_sharding=out_sharding) diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 0f1d0ddc1b..e0fc9e6f18 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -1546,11 +1546,16 @@ def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None): keep the parameters' dtype; and the loss denominator to normalize them by. Under `retry_when_tokens_dropped`, followed by `has_moe_overflow`: whether any MoE layer's ragged buffer dropped tokens, in which case the gradients and denominator are zero and `fwd_bwd` - reruns the micro-batch once the flag reaches the host. + reruns the micro-batch once the flag reaches the host. Frozen parameters' gradients are zero. """ loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn def diff_wrapper(p, r, b): + if self._freeze_mask_fn is not None: + # Parameters `trainable_parameters_mask` freezes enter as constants. `_update_kernel` zeroes their + # gradients anyway, so the update is unchanged, but the backward pass no longer computes their weight + # gradients or reduce-scatters them (the MoE router gate's, every layer, every micro-batch). + p = jax.tree.map(lambda x, frozen: jax.lax.stop_gradient(x) if frozen else x, p, self._freeze_mask_fn(p)) mdl = nnx.merge(graphdef if graphdef is not None else self._model_graphdef, p, r, copy=True) if self._gen_model_input_fn is not None: # A gen_model_input_fn maps a payload to the loss fn's *keyword arguments* -- see diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index 674b7dc280..b562444b17 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -336,6 +336,8 @@ def __init__(self): payload = DummyPayload() t.compile(payload) t.fwd_bwd(payload) + # Not differentiated at all, rather than differentiated and zeroed in `update()`. + np.testing.assert_array_equal(np.asarray(t._accumulated_grads["router_gate"]), [0.0, 0.0]) t.update() np.testing.assert_allclose(np.asarray(t.model.router_gate.get_value()), [10.0, 10.0]) diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py index ee91c2b33e..d44e8f44b7 100644 --- a/tests/unit/gdn_bwd_pallas_test.py +++ b/tests/unit/gdn_bwd_pallas_test.py @@ -24,9 +24,12 @@ import jax import jax.numpy as jnp import numpy as np +import pytest from maxtext.configs import types as config_types +from maxtext.kernels.gdn import compute_gdn from maxtext.kernels.gdn import gdn_bwd_pallas +from maxtext.kernels.gdn.gdn_bwd import api as gdn_api from maxtext.models import qwen3 from maxtext.utils import maxtext_utils @@ -138,6 +141,28 @@ def test_chunk_forward_matches_jax_reference(self): np.testing.assert_allclose(out_emit, expected_out[0], rtol=1e-3, atol=1e-3) np.testing.assert_allclose(state_emit, expected_state[0], rtol=5e-3, atol=5e-3) + def test_invert_triangular_matrix_matches_float64_inverse(self): + """The forward kernels' blocked inverse of T = I + strictly-lower(decay * beta k k^T) is exact to f32 rounding. + + Chunks of 16 and less are one substitution block; 32 and 64 also substitute into rows that the + inter-block matmul has already updated. + """ + rng = np.random.default_rng(0) + num_heads, head_dim = 4, 128 + for chunk_size in (8, 16, 32, 64): + k = rng.standard_normal((num_heads, chunk_size, head_dim)) + k /= np.linalg.norm(k, axis=-1, keepdims=True) + beta = 1 / (1 + np.exp(-rng.standard_normal((num_heads, chunk_size, 1)))) + log_g = np.cumsum(-0.05 * np.abs(rng.standard_normal((num_heads, chunk_size))), axis=1) + decay = np.exp(log_g[:, :, None] - log_g[:, None, :]) * np.tril(np.ones((chunk_size, chunk_size)), -1) + t = (np.eye(chunk_size) + decay * np.einsum("hik,hjk->hij", k * beta, k)).astype(np.float32) + + t_inv = compute_gdn.invert_triangular_matrix(jnp.asarray(t), precision=jax.lax.Precision.HIGHEST) + + np.testing.assert_allclose( + np.asarray(t_inv, np.float64), np.linalg.inv(t.astype(np.float64)), rtol=0, atol=1e-6, err_msg=f"{chunk_size=}" + ) + def test_compute_forward_conv_and_states(self): batch_size = 1 chunk_size = 16 @@ -726,6 +751,107 @@ def test_run_local_gdn_decoupled_fwd_returns_cached_chunk_states(self): np.testing.assert_allclose(chunk_states, exp_chunk_states, rtol=tol, atol=tol) np.testing.assert_allclose(t_inv, exp_t_inv, rtol=tol, atol=tol) + def test_run_local_gdn_states_fwd_matches_full_forward(self): + """The states-only forward (CP pass 1) returns the full forward's states and t_inv bit for bit. + + On TPU this compares the PER_SEQ_STATES kernel, which skips the output, with the PER_SEQ kernel. + """ + is_tpu = jax.default_backend() == "tpu" + batch_size = 2 + chunk_size = 64 if is_tpu else 32 + seq_len = 4 * chunk_size + num_k_heads = 2 + num_v_heads = 4 + head_k_dim = 128 if is_tpu else 64 + head_v_dim = 128 if is_tpu else 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias = _init_fwd_inputs( + jax.random.PRNGKey(7), batch_size, seq_len, dim_size, num_v_heads, conv_kernel_size + ) + qkv, b, a = (x.astype(jnp.bfloat16) for x in (qkv, b, a)) + # Pass-1 inputs: the conv halo, a zero recurrent state, and packed rows with trailing padding. + conv_state = jax.random.normal(jax.random.PRNGKey(8), (batch_size, conv_kernel_size - 1, dim_size), jnp.bfloat16) + recurrent_state = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), jnp.float32) + seg = np.zeros((batch_size, seq_len), np.int32) + seg[0, : seq_len // 3] = 1 + seg[0, seq_len // 3 : seq_len - 10] = 2 + seg[1, : seq_len - 40] = 1 + segment_ids, conv_halo_seg, init_seg = gdn_api._local_segment_metadata(jnp.asarray(seg), True, conv_kernel_size) + args = (qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state) + kwargs = { + "num_k_heads": num_k_heads, + "num_v_heads": num_v_heads, + "head_k_dim": head_k_dim, + "head_v_dim": head_v_dim, + "conv_kernel_size": conv_kernel_size, + "chunk_size": chunk_size, + "use_qk_norm_in_gdn": True, + "compute_dtype": jnp.float32, + "segment_ids": segment_ids, + "conv_halo_seg": conv_halo_seg, + "init_seg": init_seg, + } + + (_, exp_states), exp_t_inv, _ = gdn_api._run_local_gdn_decoupled_fwd(*args, **kwargs) + states, t_inv = gdn_api._run_local_gdn_states_fwd(*args, **kwargs) + + for got, exp in zip((*states, t_inv), (*exp_states, exp_t_inv)): + self.assertEqual(got.dtype, exp.dtype) + np.testing.assert_array_equal(np.asarray(got, np.float32), np.asarray(exp, np.float32)) + + @pytest.mark.tpu_only + def test_prefill_bf16_qkv_matches_fp32_qkv_path(self): + """Prefill-only kernels upcast bf16 qkv chunks in VMEM; the general path upcasts in HBM. Same bits. + + Covers the full prefill kernel (PER_SEQ) and the states-only kernel (PER_SEQ_STATES). + """ + from maxtext.kernels.gdn import wrapper # pylint: disable=import-outside-toplevel + + batch_size, chunk_size, num_k_heads, num_v_heads, head_dim, conv_kernel_size = 2, 64, 2, 8, 128, 4 + seq_len = 4 * chunk_size + dim_size = num_k_heads * head_dim * 2 + num_v_heads * head_dim + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias = _init_fwd_inputs( + jax.random.PRNGKey(5), batch_size, seq_len, dim_size, num_v_heads, conv_kernel_size + ) + + def run(fn, **kwargs): + # Fresh states for every call: the wrapper donates them. + k1, k2 = jax.random.split(jax.random.PRNGKey(6)) + conv_state = jax.random.normal(k1, (batch_size + 1, conv_kernel_size - 1, dim_size), jnp.bfloat16) + recurrent_state = 0.1 * jax.random.normal(k2, (batch_size + 1, num_v_heads, head_dim, head_dim), jnp.float32) + return fn( + qkv=qkv.astype(jnp.bfloat16).reshape(-1, dim_size), + b=b.astype(jnp.bfloat16).reshape(-1, num_v_heads), + a=a.astype(jnp.bfloat16).reshape(-1, num_v_heads), + conv_state=conv_state, + recurrent_state=recurrent_state, + conv_weight=jnp.swapaxes(conv_weight, 0, 2), + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + query_start_loc=jnp.arange(0, (batch_size + 1) * seq_len, seq_len, dtype=jnp.int32), + state_indices=jnp.arange(1, batch_size + 1, dtype=jnp.int32), + distribution=jnp.array([0, 0, batch_size], dtype=jnp.int32), + seq_lens=jnp.full((batch_size,), seq_len + 1, dtype=jnp.int32), + n_kq=num_k_heads, + n_v=num_v_heads, + d_k=head_dim, + d_v=head_dim, + kernel_size=conv_kernel_size, + mixed_tile_size=chunk_size, + **kwargs, + ) + + out, (conv_out, rec_out), t_inv, chunk_states = run(wrapper.fused_conv1d_gdn, is_prefill_only=True) + (states_conv, states_rec), states_t_inv = run(wrapper.fused_conv1d_gdn_states) + ref_out, (ref_conv, ref_rec), ref_t_inv, ref_chunk_states = run(wrapper.fused_conv1d_gdn, is_prefill_only=False) + got = (out, conv_out, rec_out, t_inv, chunk_states, states_conv, states_rec, states_t_inv) + exp = (ref_out, ref_conv, ref_rec, ref_t_inv, ref_chunk_states, ref_conv, ref_rec, ref_t_inv) + for g, e in zip(got, exp): + np.testing.assert_array_equal(np.asarray(g), np.asarray(e)) + def test_decoupled_conv1d_gdn_kernel_gradient_with_initial_states(self): """Verifies custom VJP gradients when initial conv_state and recurrent_state are provided.""" is_tpu = jax.default_backend() == "tpu" @@ -879,6 +1005,101 @@ def test_gdn_kernel_bwd_multi_group_head_parallel(self): np.testing.assert_allclose(dal1, dal2, rtol=1e-3, atol=1e-3) np.testing.assert_allclose(ddt1, ddt2, rtol=1e-3, atol=1e-3) + @pytest.mark.tpu_only + def test_gdn_kernel_bwd_group_blocks_match_per_group_calls(self): + """Verifies the in-place group reads (q/k/v column blocks, do, chunk_states, t_inv) and dq/dk/dv writes. + + Every head group runs the same kernel body as a single-group call on that group's heads, so the 2-group call at + the Qwen3.5-397B per-device layer geometry (64 v / 16 kq heads, bf16) must match per-group calls bit for bit. + """ + batch_size = 1 + chunk_size = 64 + num_chunks = 8 + seq_len = num_chunks * chunk_size + num_k_heads = 16 + num_v_heads = 64 + head_tile = 32 + head_dim = 128 + q_size = num_k_heads * head_dim + dim_size = 2 * q_size + num_v_heads * head_dim + tile_q = q_size * head_tile // num_v_heads + tile_v = head_tile * head_dim + + k1, k2, k3, k4, k5, k6, k7, k8, k9 = jax.random.split(jax.random.PRNGKey(2468), 9) + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32).astype(jnp.bfloat16) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32).astype(jnp.bfloat16) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32).astype(jnp.bfloat16) + a_log = jax.random.normal(k4, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal(k6, (batch_size, seq_len, num_v_heads, head_dim), dtype=jnp.float32).astype(jnp.bfloat16) + chunk_states = 0.1 * jax.random.normal(k7, (batch_size, num_chunks, num_v_heads, head_dim, head_dim)) + t_inv = 0.1 * jax.random.normal(k8, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size)) + dht = 0.1 * jax.random.normal(k9, (batch_size, num_v_heads, head_dim, head_dim)) + # Two packed documents, the second starting mid-chunk, then padding. + segment_ids = jnp.zeros((batch_size, seq_len), jnp.int32) + segment_ids = segment_ids.at[:, : seq_len // 2 + 5].set(1).at[:, seq_len // 2 + 5 : seq_len - 16].set(2) + + run = functools.partial( + gdn_bwd_pallas.pallas_gdn_bwd_kernel, + kq_head_dim=head_dim, + v_head_dim=head_dim, + chunk_size=chunk_size, + head_tile=head_tile, + segment_ids=segment_ids, + return_dh0=True, + ) + full = run( + qkv_conv=qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + t_inv=t_inv, + num_v_heads=num_v_heads, + d_recurrent_state=dht, + ) + + per_group = [] + for g in range(num_v_heads // head_tile): + heads = slice(g * head_tile, (g + 1) * head_tile) + qkv_g = jnp.concatenate( + [ + qkv[..., g * tile_q : (g + 1) * tile_q], + qkv[..., q_size + g * tile_q : q_size + (g + 1) * tile_q], + qkv[..., 2 * q_size + g * tile_v : 2 * q_size + (g + 1) * tile_v], + ], + axis=-1, + ) + per_group.append( + run( + qkv_conv=qkv_g, + b=b[..., heads], + a=a[..., heads], + a_log=a_log[heads], + dt_bias=dt_bias[heads], + do=do[:, :, heads], + chunk_states=chunk_states[:, :, heads], + t_inv=t_inv[:, :, heads], + num_v_heads=head_tile, + d_recurrent_state=dht[:, heads], + ) + ) + + dy = [out[0] for out in per_group] + ref_dy = jnp.concatenate( + [d[..., :tile_q] for d in dy] + [d[..., tile_q : 2 * tile_q] for d in dy] + [d[..., 2 * tile_q :] for d in dy], + axis=-1, + ) + np.testing.assert_array_equal(full[0], ref_dy) + np.testing.assert_array_equal(full[1], jnp.concatenate([out[1] for out in per_group], axis=-1)) + np.testing.assert_array_equal(full[2], jnp.concatenate([out[2] for out in per_group], axis=-1)) + np.testing.assert_array_equal(full[5], jnp.concatenate([out[5] for out in per_group], axis=1)) + # d_a_log / d_dt_bias also sum per-chunk partials in XLA, whose order may depend on the number of groups. + np.testing.assert_allclose(full[3], jnp.concatenate([out[3] for out in per_group]), rtol=1e-6, atol=0) + np.testing.assert_allclose(full[4], jnp.concatenate([out[4] for out in per_group]), rtol=1e-6, atol=0) + def test_gdn_kernel_bwd_variable_length_segment_ids_reset(self): """Verifies segment_ids document boundaries reset carried state gradient to prevent leakage.""" is_tpu = jax.default_backend() == "tpu" @@ -1511,53 +1732,65 @@ def test_bwd_block_specs_and_symbolic_zero_dht_dh0_omission(self): in_00, out_00, n_in_00, n_out_00 = gdn_bwd_pallas.make_bwd_block_specs( num_chunks=2, chunk_size=16, - dim_size=128, + qk_tile_size=32, + v_tile_size=64, num_v_heads=4, kq_head_dim=16, v_head_dim=16, + k_block_offset=1, + v_block_offset=1, has_dht=False, has_dh0=False, ) - self.assertEqual(n_in_00, 9) - self.assertEqual(n_out_00, 5) - self.assertLen(in_00, 9) - self.assertLen(out_00, 5) + self.assertEqual(n_in_00, 11) + self.assertEqual(n_out_00, 7) + self.assertLen(in_00, 11) + self.assertLen(out_00, 7) _, _, n_in_10, n_out_10 = gdn_bwd_pallas.make_bwd_block_specs( num_chunks=2, chunk_size=16, - dim_size=128, + qk_tile_size=32, + v_tile_size=64, num_v_heads=4, kq_head_dim=16, v_head_dim=16, + k_block_offset=1, + v_block_offset=1, has_dht=True, has_dh0=False, ) - self.assertEqual((n_in_10, n_out_10), (10, 5)) + self.assertEqual((n_in_10, n_out_10), (12, 7)) _, _, n_in_01, n_out_01 = gdn_bwd_pallas.make_bwd_block_specs( num_chunks=2, chunk_size=16, - dim_size=128, + qk_tile_size=32, + v_tile_size=64, num_v_heads=4, kq_head_dim=16, v_head_dim=16, + k_block_offset=1, + v_block_offset=1, has_dht=False, has_dh0=True, ) - self.assertEqual((n_in_01, n_out_01), (9, 6)) + self.assertEqual((n_in_01, n_out_01), (11, 8)) _, _, n_in_11, n_out_11 = gdn_bwd_pallas.make_bwd_block_specs( num_chunks=2, chunk_size=16, - dim_size=128, + qk_tile_size=32, + v_tile_size=64, num_v_heads=4, kq_head_dim=16, v_head_dim=16, + k_block_offset=1, + v_block_offset=1, has_dht=True, has_dh0=True, ) - self.assertEqual((n_in_11, n_out_11), (10, 6)) + self.assertEqual((n_in_11, n_out_11), (12, 8)) is_tpu = jax.default_backend() == "tpu" batch_size = 1 diff --git a/tests/unit/gdn_cp_dh_scan_test.py b/tests/unit/gdn_cp_dh_scan_test.py new file mode 100644 index 0000000000..28857fa05b --- /dev/null +++ b/tests/unit/gdn_cp_dh_scan_test.py @@ -0,0 +1,189 @@ +# 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 +# +# http://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. +"""The Pallas CP composition scans (`dh_scan`) against the XLA `cp_gdn` trees they replace on TPU. + +Runs the kernel in Pallas interpret mode. XLA:CPU computes the tree's f32 matmuls in plain f32, so the comparison +bounds the scans' 3-pass bf16 matmul error and the reduction order, not TPU numerics. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.kernels.gdn.gdn_bwd import cp_gdn +from maxtext.kernels.gdn.gdn_bwd import dh_scan + +_NK, _NV, _D, _C, _NC = 2, 8, 128, 64, 4 + + +def _inputs(seed=0): + """One rank's K-slice conv output, gates and per-chunk T^{-1}, with weak decay.""" + keys = jax.random.split(jax.random.PRNGKey(seed), 6) + s = _NC * _C + k_conv = jax.random.normal(keys[0], (1, s, _NK * _D), jnp.float32) + b = jax.random.normal(keys[1], (1, s, _NV), jnp.float32).astype(jnp.bfloat16) + # Weak decay keeps the product of the 4 chunk transitions far from 0. + a = (jax.random.normal(keys[2], (1, s, _NV), jnp.float32) - 6.0).astype(jnp.bfloat16) + a_log = jax.random.normal(keys[3], (_NV,), jnp.float32) * 0.3 + dt_bias = jax.random.normal(keys[4], (_NV,), jnp.float32) * 0.3 + t_inv = jnp.tril(jax.random.normal(keys[5], (1, _NC, _NV, _C, _C), jnp.float32) * 0.05, -1) + jnp.eye(_C) + return k_conv, b, a, a_log, dt_bias, t_inv + + +def _tree_m(k_conv, b, a, a_log, dt_bias, t_inv, segment_ids, init_seg): + """m_local of the XLA tree (f32 on CPU).""" + m, _ = cp_gdn.compose_local_from_t_inv( + k_conv, + b, + a, + a_log, + dt_bias, + t_inv, + num_k_heads=_NK, + num_v_heads=_NV, + head_k_dim=_D, + head_v_dim=_D, + chunk_size=_C, + use_qk_norm_in_gdn=True, + s_ext_pass1=jnp.zeros((1, _NV, _D, _D), jnp.float32), + segment_ids=segment_ids, + init_seg=init_seg, + ) + return m + + +def _scan_m(k_conv, b, a, a_log, dt_bias, t_inv, segment_ids, init_seg): + """m_local of the Pallas scan, interpreted.""" + return dh_scan.gdn_cp_m_scan( + k_conv, + b, + a, + a_log, + dt_bias, + t_inv, + num_v_heads=_NV, + kq_head_dim=_D, + chunk_size=_C, + use_qk_norm_in_gdn=True, + segment_ids=segment_ids, + init_seg=init_seg, + interpret=True, + ) + + +def _rel_err(got, ref): + got, ref = np.asarray(got, np.float64), np.asarray(ref, np.float64) + return np.abs(got - ref).max() / np.abs(ref).max() + + +# A document continuing from the previous rank and spanning the whole rank (the long single-document case), and a +# packed row: that document ends inside chunk 1, two more follow with boundaries inside chunks, then padding. +_SINGLE = np.ones((1, _NC * _C), np.int32) +_PACKED = np.zeros((1, _NC * _C), np.int32) +_PACKED[0, :70], _PACKED[0, 70:200], _PACKED[0, 200:240] = 1, 2, 3 + + +@pytest.mark.parametrize("head_tile", [64, 4]) # one head group, two head groups +@pytest.mark.parametrize("layout", ["none", "single", "packed"]) +def test_m_scan_matches_compose_local_tree(layout, head_tile, monkeypatch): + monkeypatch.setattr(dh_scan, "_M_SCAN_HEAD_TILE", head_tile) + segment_ids = {"none": None, "single": jnp.asarray(_SINGLE), "packed": jnp.asarray(_PACKED)}[layout] + init_seg = None if segment_ids is None else jnp.ones((1,), jnp.float32) + args = _inputs() + ref = _tree_m(*args, segment_ids, init_seg) + got = _scan_m(*args, segment_ids, init_seg) + assert float(jnp.abs(ref).max()) > 1e-4, "degenerate reference: the transitions vanish" + assert _rel_err(got, ref) < 5e-5 + + +def test_m_scan_sees_chunk_order_and_segment_starts(): + """Negative controls: the comparison above can fail.""" + args = _inputs() + segment_ids = jnp.asarray(_SINGLE) + ref = _tree_m(*args, segment_ids, jnp.ones((1,), jnp.float32)) + k_conv, b, a, a_log, dt_bias, t_inv = args + reversed_chunks = _scan_m(k_conv, b, a, a_log, dt_bias, t_inv[:, ::-1], segment_ids, jnp.ones((1,), jnp.float32)) + assert _rel_err(reversed_chunks, ref) > 1e-3 + # Without the incoming document, the rank's first token starts a new one, and nothing from the incoming state + # reaches the rank's end state. + fresh = _scan_m(*args, segment_ids, jnp.zeros((1,), jnp.float32)) + assert float(jnp.abs(fresh).max()) == 0.0 + + +def _bwd_inputs(seed=1): + """One rank's full conv output (q, k, v), gates, output cotangent and per-chunk T^{-1}.""" + keys = jax.random.split(jax.random.PRNGKey(seed), 7) + s = _NC * _C + qkv_conv = jax.random.normal(keys[0], (1, s, 2 * _NK * _D + _NV * _D), jnp.float32).astype(jnp.bfloat16) + b = jax.random.normal(keys[1], (1, s, _NV), jnp.float32).astype(jnp.bfloat16) + a = (jax.random.normal(keys[2], (1, s, _NV), jnp.float32) - 2.0).astype(jnp.bfloat16) + a_log = jax.random.normal(keys[3], (_NV,), jnp.float32) * 0.3 + dt_bias = jax.random.normal(keys[4], (_NV,), jnp.float32) * 0.3 + do = jax.random.normal(keys[5], (1, s, _NV, _D), jnp.float32).astype(jnp.bfloat16) + t_inv = jnp.tril(jax.random.normal(keys[6], (1, _NC, _NV, _C, _C), jnp.float32) * 0.05, -1) + jnp.eye(_C) + return qkv_conv, b, a, a_log, dt_bias, do, t_inv + + +@pytest.mark.parametrize("head_tile", [32, 4]) # one head group, two head groups +@pytest.mark.parametrize("layout", ["none", "single", "packed"]) +def test_dh_scan_matches_compose_bwd_local_tree(layout, head_tile, monkeypatch): + monkeypatch.setattr(dh_scan, "_DH_SCAN_HEAD_TILE", head_tile) + segment_ids = {"none": None, "single": jnp.asarray(_SINGLE), "packed": jnp.asarray(_PACKED)}[layout] + init_seg = None if segment_ids is None else jnp.ones((1,), jnp.float32) + args = _bwd_inputs() + _, ref = cp_gdn.compose_bwd_local_from_t_inv( + *args, + num_k_heads=_NK, + num_v_heads=_NV, + head_k_dim=_D, + head_v_dim=_D, + chunk_size=_C, + use_qk_norm_in_gdn=True, + segment_ids=segment_ids, + init_seg=init_seg, + ) + got = dh_scan.gdn_cp_dh_scan( + *args, + num_v_heads=_NV, + kq_head_dim=_D, + v_head_dim=_D, + chunk_size=_C, + use_qk_norm_in_gdn=True, + segment_ids=segment_ids, + init_seg=init_seg, + interpret=True, + ) + assert float(jnp.abs(ref).max()) > 1e-4, "degenerate reference" + assert _rel_err(got, ref) < 5e-5 + # Negative control: a perturbed t_inv must show. + qkv_conv, b, a, a_log, dt_bias, do, t_inv = args + bad = dh_scan.gdn_cp_dh_scan( + qkv_conv, + b, + a, + a_log, + dt_bias, + do, + t_inv * 1.01, + num_v_heads=_NV, + kq_head_dim=_D, + v_head_dim=_D, + chunk_size=_C, + use_qk_norm_in_gdn=True, + segment_ids=segment_ids, + init_seg=init_seg, + interpret=True, + ) + assert _rel_err(bad, ref) > 1e-3 diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index fa4025072d..fa8f72cadc 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -14,6 +14,7 @@ """Mixture of Experts (MoE) tests.""" import functools +import re from types import SimpleNamespace import unittest from unittest import mock @@ -30,6 +31,7 @@ from jax.sharding import Mesh, PartitionSpec as P from maxtext.common.common_types import Config, DType from maxtext.configs import pyconfig +from maxtext.kernels.ragged import ragged_sort from maxtext.layers import linears from maxtext.layers import moe from maxtext.layers import nnx_wrappers @@ -267,6 +269,28 @@ def test_init(self): self.model.init({"params": self.rng, "dropout": self.rng}, x) +class TopKIndicesWhereNeededTest(unittest.TestCase): + + def test_matches_top_k_on_needed_tokens(self): + """Both branches give jax.lax.top_k's ids on the needed tokens: 256 compact, 257 sorts every row.""" + # Coarse logits make ties common; top_k breaks them by the lower index. + logits = jnp.round(jax.random.normal(jax.random.PRNGKey(0), (2, 512, 32)) * 4) / 4 + probs = jax.nn.softmax(logits, axis=-1) + expected = np.asarray(jax.lax.top_k(probs, 4)[1]) + top_k_where_needed = jax.jit(moe._top_k_indices_where_needed, static_argnums=2) # pylint: disable=protected-access + order = jax.random.permutation(jax.random.PRNGKey(1), 2 * 512) + for num_needed in (256, 257): + needed = jnp.zeros(2 * 512, jnp.bool_).at[order[:num_needed]].set(True) + got = np.asarray(top_k_where_needed(probs, needed.reshape(2, 512, 1), 4)).reshape(-1, 4) + needed = np.asarray(needed) + np.testing.assert_array_equal(got[needed], expected.reshape(-1, 4)[needed]) + if num_needed <= 256: + # The compact branch leaves the other tokens (but row 0) at 0. + np.testing.assert_array_equal(got[1:][~needed[1:]], 0) + else: + np.testing.assert_array_equal(got, expected.reshape(-1, 4)) + + class DeepSeekRoutingTest(unittest.TestCase): def setUp(self): @@ -3507,5 +3531,83 @@ def test_fp8_moe_with_weight_quant_config(self): self.assertEqual(model.wi_scale.shape, expected_wi_shape) +class FusedGluTest(unittest.TestCase): + """`fused_glu` matches autodiff of the split prefused gate/up output, bit for bit.""" + + def test_matches_autodiff_of_split_bitwise(self): + def glu(gate, up): + return (jax.nn.silu(gate) * up).astype(jnp.bfloat16) + + def split_then_glu(fused): + n = fused.shape[-1] // 2 + return glu(fused[:, :n], fused[:, n:]) + + k_fused, k_cotangent = jax.random.split(jax.random.key(0)) + fused = (3 * jax.random.normal(k_fused, (256, 2 * 128))).astype(jnp.bfloat16) + # Zero cotangent rows give -0 half gradients, which autodiff's zero-padded add turns into +0. + cotangent = jax.random.normal(k_cotangent, (256, 128)).astype(jnp.bfloat16).at[64:128].set(0) + bits = lambda a: np.asarray(a).view(np.uint16) + for wrap in (lambda f: f, jax.checkpoint): + want, want_vjp = jax.vjp(jax.jit(wrap(split_then_glu)), fused) + got, got_vjp = jax.vjp(jax.jit(wrap(lambda f: moe.fused_glu(f, glu))), fused) + np.testing.assert_array_equal(bits(got), bits(want)) + np.testing.assert_array_equal(bits(got_vjp(cotangent)[0]), bits(want_vjp(cotangent)[0])) + + +class RingRaggedUnsortGradTest(parameterized.TestCase): + """ring_ragged_unsort's custom backward against autodiff of the same forward written with plain gathers.""" + + @parameterized.named_parameters(("unbuffered", None), ("buffered", 8)) + def test_grad_matches_autodiff_without_argsort(self, buffer_size): + ep, num_tokens, hidden, topk, num_experts = 2, 8, 128, 2, 4 + local_num_experts = num_experts // ep + rng = np.random.default_rng(0) + # Small integers keep every product and sum exact, so the comparison is bitwise. + hidden_states = jnp.asarray(rng.integers(-4, 5, (ep, num_tokens, hidden)), jnp.bfloat16) + experts = jnp.asarray( + [[rng.choice(num_experts, topk, replace=False) for _ in range(num_tokens)] for _ in range(ep)], jnp.int32 + ) + weights = jnp.asarray(rng.integers(1, 4, (ep, num_tokens * topk)), jnp.float32) + cotangent = jnp.asarray(rng.integers(-4, 5, (ep, num_tokens, hidden)), jnp.float32) + fallbacks = {"enforce_gather_fallback": True, "enforce_gather_reduce_fallback": True} + + def sort(x, e): + return ragged_sort.ring_ragged_sort(x, e, num_experts, topk, "expert", ep, buffer_size=buffer_size, **fallbacks) + + # Each vmap lane is one expert shard. buffer_size=8 is the balanced load: the shards receive 9 and 6 slots, + # so one drops a slot and the other leaves buffer rows unused. + tokens, group_sizes, revert, perm = jax.vmap(sort, axis_name="expert")(hidden_states, experts) + + def unsort(tokens, weights, group_sizes, revert, perm): + return ragged_sort.ring_ragged_unsort( + tokens, group_sizes, revert, perm, topk, local_num_experts, "expert", weights, **fallbacks + ) + + def reference(tokens, weights, group_sizes, revert, perm): + del perm + offsets = jnp.cumulative_sum(group_sizes, include_initial=True) + shard = jax.lax.axis_index("expert") + start, end = offsets[shard * local_num_experts], offsets[(shard + 1) * local_num_experts] + if tokens.shape[0] < revert.shape[0]: # packed buffer: this shard's rows start at 0 + pos, valid = revert - start, (revert >= start) & (revert - start < jnp.minimum(end - start, tokens.shape[0])) + else: + pos, valid = revert, (revert >= start) & (revert < end) + rows = jnp.where(valid[:, None], tokens[jnp.where(valid, pos, 0)].astype(jnp.float32) * weights[:, None], 0.0) + return rows.reshape(num_tokens, topk, hidden).sum(axis=1) + + def grads(fn): + def loss(tokens, weights): + out = jax.vmap(fn, axis_name="expert")(tokens, weights, group_sizes, revert, perm) + return jnp.sum(out.astype(jnp.float32) * cotangent) + + return jax.grad(loss, argnums=(0, 1)) + + for actual, expected in zip(grads(unsort)(tokens, weights), grads(reference)(tokens, weights)): + np.testing.assert_array_equal(actual, expected) + # The backward reuses the sort's permutation instead of argsorting the revert indices. + jaxpr = str(jax.make_jaxpr(grads(unsort))(tokens, weights)) + self.assertIsNone(re.search(r"\bsort\[", jaxpr), "the unsort backward argsorts") + + if __name__ == "__main__": absltest.main() diff --git a/tests/unit/moe_valid_rows_kernels_test.py b/tests/unit/moe_valid_rows_kernels_test.py new file mode 100644 index 0000000000..2edc4891d8 --- /dev/null +++ b/tests/unit/moe_valid_rows_kernels_test.py @@ -0,0 +1,102 @@ +# 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 +# +# http://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. +"""The valid-rows MoE kernels (`unsort_bwd_rows`, `glu_rows`) against the XLA code they replace. + +On CPU the kernels run in Pallas interpret mode with the TPU7x chip description, over a buffer whose last +1024-row block is partial. On a TPU the SiLU-GLU kernels must match XLA's fusions bit for bit for every bf16 gate. +""" + +import jax +from jax._src.pallas.mosaic import tpu_info +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.kernels.ragged import glu_rows +from maxtext.kernels.ragged import unsort_bwd_rows as unsort +from maxtext.layers import moe + +_ROWS, _HIDDEN, _MLP = 2560, 256, 128 + + +@pytest.fixture(autouse=True) +def _tpu7x_interpret(): + if jax.default_backend() == "cpu": + tpu_info.registry["cpu"] = lambda: tpu_info.get_tpu_info_for_chip(tpu_info.ChipVersion.TPU_7X, 1) + tpu_info.get_tpu_info.cache_clear() + + +def _bf16(key, shape, scale=1.0): + return (jax.random.normal(key, shape, jnp.float32) * scale).astype(jnp.bfloat16) + + +def _bits(x, rows): + return np.asarray(x)[:rows].view(np.uint16) + + +def _silu_glu(gate, up): + return (jax.nn.silu(gate) * up).astype(jnp.bfloat16) + + +@pytest.mark.parametrize("limit", [0, 1000, 1700, _ROWS]) +def test_unsort_bwd_rows(limit): + """Weighted gradient bit-identical and row dot to f32 summation order on rows < limit, zeros up to the block end.""" + keys = jax.random.split(jax.random.PRNGKey(0), 3) + gathered, tokens = _bf16(keys[0], (_ROWS, _HIDDEN)), _bf16(keys[1], (_ROWS, _HIDDEN)) + weights = jax.random.uniform(keys[2], (_ROWS,), jnp.float32, -1.0, 1.0) + interpret = True if jax.default_backend() == "cpu" else None + grad, dot = unsort.unsort_bwd_rows(gathered, tokens, weights, jnp.int32(limit), interpret=interpret) + ref_grad, ref_dot = unsort.unsort_bwd_rows_reference(gathered, tokens, weights, jnp.int32(limit)) + block_end = min(_ROWS, max(-(-limit // 1024), 1) * 1024) + np.testing.assert_array_equal(_bits(grad, block_end), _bits(ref_grad, block_end)) + np.testing.assert_allclose(np.asarray(dot)[:limit], np.asarray(ref_dot)[:limit], rtol=0, atol=1e-4) + np.testing.assert_array_equal(np.asarray(dot)[limit:block_end], 0.0) + + +def _glu_inputs(rows, mlp): + """Every bf16 bit pattern as a gate value, against random up and output gradients.""" + gate = np.resize(np.arange(1 << 16, dtype=np.uint16).view(jnp.bfloat16), (rows, mlp)) + keys = jax.random.split(jax.random.PRNGKey(1), 2) + fused = jnp.concatenate([jnp.asarray(gate), _bf16(keys[0], (rows, mlp), 4.0)], axis=-1) + return fused, _bf16(keys[1], (rows, mlp), 4.0) + + +def test_glu_rows_interpret_rows_and_blocks(): + """CPU: the visited rows (partial last block) hold the per-tile formula; XLA:CPU may contract it differently.""" + if jax.default_backend() != "cpu": + pytest.skip("interpret-mode check") + fused, d_out = _glu_inputs(_ROWS, _MLP) + gate, up = fused[:, :_MLP], fused[:, _MLP:] + act = glu_rows.glu_fwd_rows(fused, jnp.int32(2100), interpret=True) + d_fused = glu_rows.glu_bwd_rows(fused, d_out, jnp.int32(2100), interpret=True) + ref_act = glu_rows._glu_fwd_tile(gate, up) # pylint: disable=protected-access + ref_d_fused = jnp.concatenate(glu_rows._glu_bwd_tile(gate, up, d_out), axis=-1) # pylint: disable=protected-access + for got, want in ((act, ref_act), (d_fused, ref_d_fused)): + np.testing.assert_allclose(np.asarray(got, np.float32), np.asarray(want, np.float32), rtol=2**-6, atol=1e-37) + + +@pytest.mark.tpu_only +def test_glu_rows_match_xla_fusions(): + """TPU: activation and gradient bit-identical to `fused_glu`'s XLA path, for every bf16 gate.""" + rows = 8192 + fused, d_out = _glu_inputs(rows, 1024) + limit = 6000 + block_end = 6144 + act = jax.jit(glu_rows.glu_fwd_rows)(fused, jnp.int32(limit)) + d_fused = jax.jit(glu_rows.glu_bwd_rows)(fused, d_out, jnp.int32(limit)) + # Jitted, so XLA fuses the reference as it does in the model (f32 inside the fusion, bf16 outputs). + ref_act = jax.jit(lambda f: moe.fused_glu(f, _silu_glu))(fused) + ref_d_fused = jax.jit(lambda f, d: jax.vjp(lambda x: moe.fused_glu(x, _silu_glu), f)[1](d)[0])(fused, d_out) + np.testing.assert_array_equal(_bits(act, block_end), _bits(ref_act, block_end)) + np.testing.assert_array_equal(_bits(d_fused, block_end), _bits(ref_d_fused, block_end))