diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 69d8dc0f60..c5f2f118fe 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -123,6 +123,10 @@ grad_dtype: "float32" # activation dtypes. dtype: "bfloat16" + +# GDN precision configuration +gdn_state_dtype: "float32" +gdn_decay_dtype: "float32" # used to configure quantization in the transformer layers, defaults to null implying bf16. # possible alternative settings are as follows: # 'int8' for dynamic range quantization using 8-bits @@ -1357,6 +1361,8 @@ gdn_chunk_size: 64 use_qk_norm_in_gdn: true # The ratio of dimension to apply ROPE on partial_rotary_factor: 1.0 +# Whether to use GDN Pallas kernel +use_gdn_kernel: false use_tokamax_splash: false # Setting this flag will use a non-pallas implementation. diff --git a/src/maxtext/configs/models/qwen3.5-tiny.yml b/src/maxtext/configs/models/qwen3.5-tiny.yml new file mode 100644 index 0000000000..422219db26 --- /dev/null +++ b/src/maxtext/configs/models/qwen3.5-tiny.yml @@ -0,0 +1,72 @@ +# 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. + +# Tiny version of Qwen3.5 for fast testing and execution. + +decoder_block: "qwen3_5" + +# Core Architectural Parameters +base_emb_dim: 256 +base_num_decoder_layers: 2 +base_num_query_heads: 4 +base_num_kv_heads: 2 +head_dim: 256 +vocab_size: 248320 +normalization_layer_epsilon: 1.0e-6 + +# MoE Specific Parameters +# Set base_mlp_dim to match base_moe_mlp_dim to pass validation for fully MoE models. +base_mlp_dim: 256 +base_moe_mlp_dim: 256 +num_experts: 8 +shared_experts: 1 +num_experts_per_tok: 2 +norm_topk_prob: true + +# GatedDeltaNet Specific Parameters for Linear Attention (GDN) +inhomogeneous_layer_cycle_interval: 2 +gdn_conv_kernel_dim: 4 +gdn_key_head_dim: 128 +gdn_value_head_dim: 128 +gdn_num_key_heads: 2 +gdn_num_value_heads: 4 +gdn_chunk_size: 64 + +# RoPE Settings +rope_max_timescale: 10000000 +partial_rotary_factor: 0.25 + +# General Model Settings +enable_dropout: false + +# Vision Encoder Configuration (need to set use_multimodal=true) +vision_encoder_block: "qwen3_5" +# Based on Qwen3.5 MoE Vision Model Config +image_size_for_vit: 768 +hidden_size_for_vit: 1152 +intermediate_size_for_vit: 4304 +num_attention_heads_for_vit: 16 +num_hidden_layers_for_vit: 27 +num_channels_for_vit: 3 +patch_size_for_vit: 16 +temporal_patch_size_for_vit: 2 +spatial_merge_size_for_vit: 2 +out_hidden_size_for_vit: 256 # Projects to decoder emb_dim (256) +num_position_embeddings_for_vit: 2304 +deepstack_visual_indexes_for_vit: [] # No deepstack for Qwen3.5 VL +rope_theta_for_vit: 10000 + +# MRoPE Settings (Multi-dimensional RoPE for multimodal) +use_mrope: true +mrope_section: [11, 11, 10] diff --git a/src/maxtext/configs/pyconfig.py b/src/maxtext/configs/pyconfig.py index f784f27a64..b5bd17ec0d 100644 --- a/src/maxtext/configs/pyconfig.py +++ b/src/maxtext/configs/pyconfig.py @@ -337,6 +337,8 @@ def __init__(self, pydantic_config: types.MaxTextConfig): final_dict["dtype"] = jnp.dtype(final_dict["dtype"]) final_dict["grad_dtype"] = jnp.dtype(final_dict["grad_dtype"]) final_dict["weight_dtype"] = jnp.dtype(final_dict["weight_dtype"]) + final_dict["gdn_state_dtype"] = jnp.dtype(final_dict.get("gdn_state_dtype", "float32")) + final_dict["gdn_decay_dtype"] = jnp.dtype(final_dict.get("gdn_decay_dtype", "float32")) final_dict["mu_dtype"] = ( final_dict["weight_dtype"] if not final_dict["mu_dtype"] else jnp.dtype(final_dict["mu_dtype"]) ) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index d35d6b9be0..a754351fa4 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -284,6 +284,7 @@ class ProfilerType(str, Enum): "qwen3-custom-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "qwen3.5-tiny", "gpt3-175b", "gpt3-22b", "gpt3-6b", @@ -480,6 +481,13 @@ class DataTypes(BaseModel): description="If True, sets activations to float32 before the nonlinearity.", ) dtype_mm: str = Field("float32", description="Data type for multimodal model's vision encoder") + gdn_state_dtype: DType = Field( + DType.FLOAT32, description="The data type for GDN recurrent states." + ) + gdn_decay_dtype: DType = Field( + DType.FLOAT32, + description="The data type for GDN decay parameters (A_log, dt_bias).", + ) class Quantization(BaseModel): @@ -1154,6 +1162,10 @@ class Qwen3Next(BaseModel): description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.", ) partial_rotary_factor: float = Field(1.0, description="The ratio of dimension to apply ROPE on") + use_gdn_kernel: bool = Field( + False, + description="Whether to use GDN Pallas kernel.", + ) # ---------------------------------------------------------------------------- @@ -4319,6 +4331,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "qwen3-vl-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "qwen3.5-tiny", "maxtext-omni-gemma3-qwen3", "cosmos3-nano-reasoner", "cosmos3-super-reasoner", diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 1826ef576b..4f3c70f19b 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -1090,7 +1090,7 @@ def validate_inputs( assert group_offset.shape == (1,) - size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype) + size_lhs_sublane = max(pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype), 16) size_lhs_sublane = min(size_lhs_sublane, size_m) if fuse_act is not None: num_lanes = pltpu.get_tpu_info().num_lanes diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py index a5f3f6338d..2bf524a5e4 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py @@ -221,6 +221,9 @@ def make_tgmm_configs( size_lhs_sublane = min(size_lhs_sublane, size_m) size_rhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(rhs.dtype) size_rhs_sublane = min(size_rhs_sublane, size_m) + common_sublane = min(size_lhs_sublane, size_rhs_sublane) + size_lhs_sublane = common_sublane + size_rhs_sublane = common_sublane assert size_lhs_sublane == size_rhs_sublane, ( f"size_lhs_sublane should be the same as size_rhs_sublane {lhs.dtype=}," f" {rhs.dtype=}" ) diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index a9bacf692f..962f1d595c 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -734,6 +734,7 @@ def _apply_embedding( "qwen3-vl-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "qwen3.5-tiny", "maxtext-omni-gemma3-qwen3", ]: y = mm_utils.merge_mm_embeddings( @@ -754,6 +755,7 @@ def _apply_embedding( "qwen3-vl-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "qwen3.5-tiny", ]: y = mm_utils.merge_mm_embeddings( text_embeddings=y, diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index c1d567eea3..80b0a5d533 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1478,6 +1478,7 @@ def _apply_embedding( "qwen3-vl-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "qwen3.5-tiny", "maxtext-omni-gemma3-qwen3", "cosmos3-nano-reasoner", "cosmos3-super-reasoner", @@ -1499,6 +1500,7 @@ def _apply_embedding( "qwen3-vl-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "qwen3.5-tiny", "cosmos3-nano-reasoner", "cosmos3-super-reasoner", }: diff --git a/src/maxtext/models/kernels/gdn/__init__.py b/src/maxtext/models/kernels/gdn/__init__.py new file mode 100644 index 0000000000..8c28c13e0b --- /dev/null +++ b/src/maxtext/models/kernels/gdn/__init__.py @@ -0,0 +1,42 @@ +# 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. +# ============================================================================== + +"""Local GDN forward kernel package with triangular inverse matrix caching.""" + +from . import compute_conv1d +from . import compute_gdn +from . import config +from . import gdn_bwd_pallas +from . import memory_ref +from . import metadata +from . import tiling +from . import vmem_ldst +from . import wrapper +from .gdn_bwd_pallas import gdn_fused_conv1d +from .gdn_bwd_pallas import pallas_gdn_bwd_kernel + +__all__ = [ + "compute_conv1d", + "compute_gdn", + "config", + "gdn_bwd_pallas", + "gdn_fused_conv1d", + "memory_ref", + "metadata", + "pallas_gdn_bwd_kernel", + "tiling", + "vmem_ldst", + "wrapper", +] diff --git a/src/maxtext/models/kernels/gdn/compute_conv1d.py b/src/maxtext/models/kernels/gdn/compute_conv1d.py new file mode 100644 index 0000000000..76152641b5 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/compute_conv1d.py @@ -0,0 +1,75 @@ +# 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. +# ============================================================================== + +"""In-VMEM causal depthwise Conv1D computation.""" + +import jax +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + except (ImportError, ModuleNotFoundError): + from . import config + + +def causal_conv1d( + real_sizes: jax.Array, # [seq] + lhs: jax.Array, # [seq, chunk, q, dim_size] + conv_weight: jax.Array, # [prev_kernel_size, 1, dim_size] + conv_bias: jax.Array | None, # [dim_size] + cfg: config.GDNConfig, +) -> tuple[jax.Array, jax.Array]: + """Perform causal Conv1D. Returns Conv1D output and convolution states.""" + + assert lhs.ndim == 4 + + out_list = [] + + for c_idx in range(cfg.chunk_size): + out = jnp.zeros((cfg.seq_tile_size, 1, cfg.dim_size), jnp.float32) + + end_idx = c_idx + cfg.prev_kernel_size + start_idx = 1 + end_idx - cfg.kernel_size + for k in range(cfg.kernel_size): + lhs_curr = lhs[:, start_idx + k] + out += lhs_curr * conv_weight[k : k + 1] + + if conv_bias is not None: + out += conv_bias.reshape(1, 1, -1) + + out_list.append(out) + + # Last prev_kernel_size elements needs to be returned as conv_state. However, + # real_sizes may be smaller than chunk_size. Therefore, slicing last + # prev_kernel_size elements does not gurantee numeric correctness. Instead, + # kernel iterate each rows and perform masking to fetch correct values. + # NOTE: lhs[:, : prev_kernel_size] can be skipped since they were loaded from + # previous conv states. + new_conv_state = lhs[:, 1 : cfg.kernel_size] + real_sizes = real_sizes.reshape(-1, 1, 1, 1) + # NOTE: Even though for loop is invoked twice, since they are static loops, + # compiler will perform loop fusion. + for c_idx in range(2, cfg.chunk_size + 1): + row_end = c_idx + cfg.prev_kernel_size + new_conv_state = jnp.where( + c_idx == real_sizes, + lhs[:, c_idx:row_end], + new_conv_state, + ) + + return jnp.stack(out_list, axis=1), new_conv_state diff --git a/src/maxtext/models/kernels/gdn/compute_gdn.py b/src/maxtext/models/kernels/gdn/compute_gdn.py new file mode 100644 index 0000000000..f97cd949f9 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/compute_gdn.py @@ -0,0 +1,476 @@ +# 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. +# ============================================================================== + +"""Core GDN forward computation with triangular inverse matrix caching.""" + +import jax +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + except (ImportError, ModuleNotFoundError): + from . import config + + +def l2_norm(x: jax.Array, eps: float = 1e-6) -> jax.Array: + norm = jnp.sqrt(jnp.sum(x * x, axis=-1, keepdims=True, dtype=x.dtype) + eps) + return x / norm + + +def get_mask_dtype(dtype: jnp.dtype) -> jnp.dtype: + match jnp.dtype(dtype).itemsize: + case 4: + return jnp.int32 + case 2: + return jnp.int16 + case _: + raise ValueError(f"Unsupported dtype: {dtype}") + + +# NOTE: Fork of recurrent_scan_v2.py but applied various optimizations. +def invert_triangular_matrix(t: jax.Array, block_size: int = 16) -> jax.Array: + """Compute invert matrix of a given triauglar matrix.""" + + # NOTE: if chunk_size=1, compiler will perform DCE. + out_dtype = t.dtype + chunk = t.shape[-1] + block_size = min(block_size, chunk) + 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 + 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_blocks = [] + iota_r = jax.lax.broadcasted_iota(jnp.int32, t.shape, 1) + iota_c = jax.lax.broadcasted_iota(jnp.int32, t.shape, 2) + identity_mask = jnp.where(iota_r == iota_c, 1.0, 0.0) + for i in range(num_blocks): + start, end = i * block_size, (i + 1) * block_size + e_block = identity_mask[:, start:end, :] + + if i == 0: + target_b = e_block + else: + interaction_t = t[:, start:end, :start] + solved_x = jnp.concatenate(x_blocks, axis=1) + prev_sum = jax.lax.dot( + interaction_t, + solved_x, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ) + target_b = e_block - prev_sum + + # NOTE: Utilize fp32 to minimize cost of sublane rolling. + local_t = t[:, start:end, start:end].astype(jnp.float32) + x_block = local_forward_sub(local_t, target_b) + x_blocks.append(x_block.astype(out_dtype)) + + return jnp.concatenate(x_blocks, axis=1) + + +def fused_transpose_broadcast( + x: jax.Array, src_dim: int, dst_dim: int +) -> jax.Array: + """Perform 1D transpose where results are broadcasted along src_dim.""" + assert x.shape[dst_dim] == 1 + + dtype = x.dtype + mask_dtype = get_mask_dtype(dtype) + mask_shape = list(x.shape) + mask_size = mask_shape[src_dim] + mask_shape[dst_dim] = mask_size + src_mask = jax.lax.broadcasted_iota(mask_dtype, mask_shape, src_dim) + dst_mask = jax.lax.broadcasted_iota(mask_dtype, mask_shape, dst_dim) + mask = src_mask == dst_mask + return jnp.where(mask, x, 0).sum(axis=src_dim, keepdims=True, dtype=dtype) + + +def chunked_gdn_per_seq( + q_large: jax.Array, # [num_kq_heads, chunk, kq_head_dim] + k_large: jax.Array, # [num_kq_heads, chunk, kq_head_dim] + v_large: jax.Array, # [num_v_heads, chunk, v_head_dim] + gating_log: jax.Array, # [1, 1, num_v_heads] + beta: jax.Array, # [1, 1, num_v_heads] + state_prev: jax.Array, # [num_v_heads, kq_head_dim, v_head_dim] + cfg: config.GDNConfig, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Perform chunked GDN over input [num_heads, chunk, head_dim].""" + + # NOTE: Repeat along non lane/sublane dim is free. + q_repeat = jnp.repeat(q_large, cfg.v_per_kq_head, axis=0) + k_repeat = jnp.repeat(k_large, cfg.v_per_kq_head, axis=0) + + # Compute cumulative sum of decay. + # [1, 1, num_v_heads] + g_cum_sum_list = [gating_log[:, :1]] + for row in range(1, cfg.chunk_size): + g_cum_sum_list.append(g_cum_sum_list[-1] + gating_log[:, row : row + 1]) + # [1, chunk, num_v_heads] + g_cum_sum_log = jnp.concat(g_cum_sum_list, axis=1) + + # [num_v_heads, chunk, 1] + g_cum_sum_log = fused_transpose_broadcast(g_cum_sum_log, src_dim=2, dst_dim=0) + g_cum_sum_log = g_cum_sum_log[: cfg.num_v_heads] + beta = fused_transpose_broadcast(beta, src_dim=2, dst_dim=0) + beta_large = beta[: cfg.num_v_heads] + + # [num_v_heads, 1, chunk] + g_cum_sum_log_t = fused_transpose_broadcast( + g_cum_sum_log, src_dim=1, dst_dim=2 + ) + # [num_v_heads, chunk, chunk] + g_cum_sum_diff_log = g_cum_sum_log - g_cum_sum_log_t + gating_map = jnp.exp(g_cum_sum_diff_log) + # [num_v_heads, chunk, 1] + gating_backward = jnp.exp(-g_cum_sum_diff_log[..., -1:]) + # [num_v_heads, chunk, 1] + gating_forward = jnp.exp(g_cum_sum_log) + # [num_v_heads, 1, 1] + gating_last = gating_forward[:, -1:] + + mask_dtype = get_mask_dtype(cfg.dtypes.compute) + iota_r = jax.lax.broadcasted_iota(mask_dtype, gating_map.shape, 1) + iota_c = jax.lax.broadcasted_iota(mask_dtype, gating_map.shape, 2) + identity_mask = iota_r == iota_c + strictly_lower_mask = iota_r > iota_c + lower_mask = iota_r >= iota_c + # [num_v_heads, chunk, chunk] + gating_map_masked = jnp.where(strictly_lower_mask, gating_map, 0) + + # [num_v_heads, chunk, kq_head_dim] + k_beta_repeat = k_repeat * beta_large + + # [num_v_heads, chunk, chunk] + beta_k_k_t = jax.lax.dot( + k_beta_repeat, + k_repeat, + dimension_numbers=(((2,), (2,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfg.dtypes.compute) + gating_beta_k_k_t = gating_map_masked * beta_k_k_t + t = jnp.where(identity_mask, 1, gating_beta_k_k_t) + + # [num_v_heads, chunk, chunk] + t_inv = invert_triangular_matrix(t) + + # [num_v_heads, chunk, v_head_dim] + v_beta_large = v_large * beta_large + # [num_v_heads, chunk, kv_head_dim] + k_beta_gating = k_beta_repeat * gating_forward + # NOTE: If v_head_dim < mxu size, concatenating them will help increase mxu + # utilization. Also, if v_head_dim is multiple of lane size, concat / split + # along lane dim is free - making this optimization strictly beneficial. + # [num_v_heads, chunk, v_head_dim + kq_head_dim] + merged_v_k = jnp.concat([v_beta_large, k_beta_gating], axis=-1) + merged_uw = jax.lax.dot( + t_inv, + merged_v_k, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfg.dtypes.compute) + + # [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,))), + 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) + ws = ws.astype(cfg.dtypes.compute) + + # [num_v_heads, chunk, v_head_dim] + u_ws = u - ws + + # [num_v_heads, chunk, kq_head_dim] + k_repeat_gating = k_repeat * gating_backward + + # [num_v_heads, kq_head_dim, v_head_dim] + state_new = jax.lax.dot( + k_repeat_gating, + u_ws, + dimension_numbers=(((1,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ) + + # [num_v_heads, kq_head_dim, v_head_dim] + state_updated = state_prev * gating_last + state = state_updated + state_new + + # [num_kq_heads, chunk, chunk] + out_qk = jax.lax.dot( + q_large, + k_large, + dimension_numbers=(((2,), (2,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfg.dtypes.compute) + # NOTE: must perform repeat after matmul to reduce required compute. + # [num_v_heads, chunk, chunk] + out_qk = jnp.repeat(out_qk, cfg.v_per_kq_head, axis=0) + out_qk *= gating_map + out_qk = jnp.where(lower_mask, out_qk, 0) + + # [num_v_heads, chunk, v_head_dim] + out_new = jax.lax.dot( + out_qk, + u_ws, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ) + out = out_updated + out_new + + return out, state, t_inv + + +def chunked_gdn( + real_sizes: jax.Array, + q_large: jax.Array, + k_large: jax.Array, + v_large: jax.Array, + b_large: jax.Array, + a_large: jax.Array, + state_prev: jax.Array, + 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].""" + + 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 + ) + mask = iota < real_sizes.reshape(-1, 1, 1, 1).astype(mask_dtype) + + # [seqs, num_kq_heads, chunk, kq_head_dim] + q_large = jnp.where(mask, q_large.astype(cfg.dtypes.compute), 0) + k_large = jnp.where(mask, k_large.astype(cfg.dtypes.compute), 0) + # [seqs, num_v_heads, chunk, v_head_dim] + v_large = jnp.where(mask, v_large.astype(cfg.dtypes.compute), 0) + + b_large = b_large.astype(cfg.dtypes.compute) + a_large = a_large.astype(cfg.dtypes.compute) + + a_log = a_log.reshape(1, 1, 1, -1).astype(cfg.dtypes.compute) + dt_bias = dt_bias.reshape(1, 1, 1, -1).astype(cfg.dtypes.compute) + + # NOTE: Any element-wise computations should occur before repeat. + q_large = l2_norm(q_large) + q_scale = cfg.kq_head_dim**-0.5 + q_large *= q_scale + k_large = l2_norm(k_large) + + # [seqs, 1, chunk, num_v_heads] + beta = jax.nn.sigmoid(b_large) + gating_log = -jnp.exp(a_log) * jax.nn.softplus(a_large + dt_bias) + + beta = jnp.where(mask, beta, 0) + # NOTE: Masked gating_log will evaluate to jnp.exp(0)=1. gating (decay) must + # be masked to 1 since it signifies that strength of state from previous row + # will be 1 (i.e., no decay) if current row is invalid. + gating_log = jnp.where(mask, gating_log, 0) + + out_list = [] + state_list = [] + t_inv_list = [] + for idx in range(cfg.seq_tile_size): + out, state, t_inv = chunked_gdn_per_seq( + q_large[idx], + k_large[idx], + v_large[idx], + gating_log[idx], + beta[idx], + state_prev[idx], + cfg, + ) + out_list.append(out.swapaxes(0, 1)) + state_list.append(state) + t_inv_list.append(t_inv) + out = jnp.stack(out_list, axis=0) + state = jnp.stack(state_list, axis=0) + t_inv = jnp.stack(t_inv_list, axis=0) + return out, state, t_inv + + +def recurrent_gdn_per_seq( + q_compact: jax.Array, # [num_kq_heads, chunk, 1, kq_head_dim] + k_compact: jax.Array, # [num_kq_heads, chunk, 1, kq_head_dim] + k_compact_t: jax.Array, # [num_kq_heads, chunk, kq_head_dim, 1] + v_compact: jax.Array, # [num_v_heads, chunk, 1, v_head_dim] + gating_log: jax.Array, # [num_v_heads, chunk, 1, 1] + beta: jax.Array, # [num_v_heads, chunk, 1, 1] + state: jax.Array, # [num_v_heads, kq_head_dim, v_head_dim] + cfgs: config.GDNConfig, +) -> tuple[jax.Array, jax.Array]: + """Perform recurrent GDN over input [num_heads, chunk, 1, head_dim].""" + + out_list = [] + for c_idx in range(cfgs.chunk_size): + # [num_v_heads, 1, kq_head_dim] + q_curr = q_compact[:, c_idx] + q_curr = jnp.repeat(q_curr, cfgs.v_per_kq_head, axis=0) + k_curr = k_compact[:, c_idx] + k_curr = jnp.repeat(k_curr, cfgs.v_per_kq_head, axis=0) + + # [num_v_heads, 1, v_head_dim] + v_curr = v_compact[:, c_idx] + + # [num_v_heads, kq_head_dim, 1] + k_curr_t = k_compact_t[:, c_idx] + k_curr_t = jnp.repeat(k_curr_t, cfgs.v_per_kq_head, axis=0) + + # [num_v_heads, 1, 1] + beta_curr = beta[:, c_idx] + gating_curr = gating_log[:, c_idx] + + # [num_v_heads, kq_head_dim, v_head_dim] + state_updated = state * gating_curr + + # [num_v_heads, 1, v_head_dim] + v_updated = jax.lax.dot( + k_curr, + state_updated, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfgs.dtypes.compute) + + # [num_v_heads, 1, v_head_dim] + v_diff = v_curr - v_updated + v_new = beta_curr * v_diff + + # [num_v_heads, kq_head_dim, v_head_dim] + # NOTE: Multiplication with k_curr_t needs to be deferred as much as + # possible as it expands the dimension size by kq_head_dim. + state_new = k_curr_t * v_new + # [num_v_heads, kq_head_dim, v_head_dim] + state = state_updated + state_new + + # [num_v_heads, 1, v_head_dim] + out = jax.lax.dot( + q_curr, + state, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfgs.dtypes.compute) + + out_list.append(out[:, 0, :]) + + return jnp.stack(out_list, axis=0), state + + +def recurrent_gdn( + real_sizes: jax.Array, + q_compact: jax.Array, + k_compact: jax.Array, + v_compact: jax.Array, + b_compact: jax.Array, + a_compact: jax.Array, + state_prev: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + cfg: config.GDNConfig, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Perform recurrent GDN over input [seq, num_heads, chunk, 1, head_dim].""" + + mask_dtype = get_mask_dtype(cfg.dtypes.compute) + iota = jax.lax.broadcasted_iota( + mask_dtype, (cfg.seq_tile_size, 1, cfg.chunk_size, 1, 1), 2 + ) + mask = iota < real_sizes.reshape(-1, 1, 1, 1, 1).astype(mask_dtype) + + # [seqs, num_kq_heads, chunk, 1, kq_head_dim] + q_compact = jnp.where(mask, q_compact.astype(cfg.dtypes.compute), 0) + k_compact = jnp.where(mask, k_compact.astype(cfg.dtypes.compute), 0) + # [seqs, num_v_heads, chunk, 1, v_head_dim] + v_compact = jnp.where(mask, v_compact.astype(cfg.dtypes.compute), 0) + + b_compact = b_compact.astype(cfg.dtypes.compute) + a_compact = a_compact.astype(cfg.dtypes.compute) + + a_log = a_log.reshape(1, 1, 1, 1, -1).astype(cfg.dtypes.compute) + dt_bias = dt_bias.reshape(1, 1, 1, 1, -1).astype(cfg.dtypes.compute) + + # [seqs, num_kq_heads, chunk, 1, kq_head_dim] + q_compact = l2_norm(q_compact) + q_scale = cfg.kq_head_dim**-0.5 + q_compact *= q_scale + k_compact = l2_norm(k_compact) + k_compact_t = fused_transpose_broadcast(k_compact, src_dim=4, dst_dim=3) + + beta = jax.nn.sigmoid(b_compact) + gating_log = -jnp.exp(a_log) * jax.nn.softplus(a_compact + dt_bias) + + beta = jnp.where(mask, beta, 0) + # NOTE: Masked gating_log will evaluate to jnp.exp(0)=1. gating (decay) must + # be masked to 1 since it signifies that strength of state from previous row + # will be 1 (i.e., no decay) if current row is invalid. + gating_log = jnp.where(mask, gating_log, 0) + gating_log = jnp.exp(gating_log) + + beta = fused_transpose_broadcast(beta, src_dim=4, dst_dim=1) + beta = beta[:, : cfg.num_v_heads] + gating_log = fused_transpose_broadcast(gating_log, src_dim=4, dst_dim=1) + gating_log = gating_log[:, : cfg.num_v_heads] + + out_list = [] + new_state_list = [] + + for idx in range(cfg.seq_tile_size): + out, state = recurrent_gdn_per_seq( + q_compact[idx], + k_compact[idx], + k_compact_t[idx], + v_compact[idx], + gating_log[idx], + beta[idx], + state_prev[idx], + cfg, + ) + out_list.append(out) + new_state_list.append(state) + + out = jnp.stack(out_list, axis=0) + new_recurrent_state = jnp.stack(new_state_list, axis=0) + t_inv = jnp.ones( + (cfg.seq_tile_size, cfg.num_v_heads, 1, 1), dtype=cfg.dtypes.compute + ) + + return out, new_recurrent_state, t_inv diff --git a/src/maxtext/models/kernels/gdn/config.py b/src/maxtext/models/kernels/gdn/config.py new file mode 100644 index 0000000000..74be137655 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/config.py @@ -0,0 +1,150 @@ +# 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. +# ============================================================================== + +"""GDN Configuration dataclass defining GDN tiling, dtypes, and kernel configurations.""" + +import dataclasses +import enum +from typing import Any + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + + +DEFAULT_VMEM_LIMIT_FACTOR: float = 0.90 + + +class GDNMode(enum.StrEnum): + BATCHED = enum.auto() + PER_SEQ = enum.auto() + + def get_seq_tile_size(self, tile_size: int) -> int: + if self == GDNMode.BATCHED: + return tile_size + return 1 + + def get_chunk_size(self, tile_size: int) -> int: + if self == GDNMode.BATCHED: + return 1 + return tile_size + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class Dtypes: + act_in: jnp.dtype + act_out: jnp.dtype + compute: jnp.dtype + recurrent_state: jnp.dtype + conv_state: jnp.dtype + + +def get_vmem_limit_bytes( + vmem_limit_factor: float = DEFAULT_VMEM_LIMIT_FACTOR, +) -> int: + """Returns the maximum allowable VMEM capacity budget in bytes.""" + tpu_info = pltpu.get_tpu_info() + return int(vmem_limit_factor * tpu_info.vmem_capacity_bytes) + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class GDNConfig: + mode: GDNMode + dtypes: Dtypes + batch_size: int + dim_size: int + kernel_size: int + tile_size: int + num_kq_heads: int + num_v_heads: int + kq_head_dim: int + v_head_dim: int + num_buffers: int = 2 + + @property + def chunk_size(self) -> int: + return self.mode.get_chunk_size(self.tile_size) + + @property + def seq_tile_size(self) -> int: + return self.mode.get_seq_tile_size(self.tile_size) + + @property + def prev_kernel_size(self) -> int: + return self.kernel_size - 1 + + @property + def v_dim_size(self) -> int: + return self.num_v_heads * self.v_head_dim + + @property + def kq_dim_size(self) -> int: + return self.num_kq_heads * self.kq_head_dim + + @property + def v_per_kq_head(self) -> int: + return self.num_v_heads // self.num_kq_heads + + @property + def aligned_num_v_heads(self) -> int: + tpu_info = pltpu.get_tpu_info() + num_lanes = tpu_info.num_lanes + return pl.cdiv(self.num_v_heads, num_lanes) * num_lanes + + def get_kernel_name(self) -> str: + return ( + f"fused_conv1d_gdn_{self.mode.value}_b{self.seq_tile_size}" + f"_c{self.chunk_size}" + ) + + def get_metadata(self) -> dict[str, str | int | float]: + cfgs_dict = dataclasses.asdict(self) + ret = {} + for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict): + key = jax.tree_util.keystr(path, simple=True, separator=".") + if not isinstance(val, str | int | float): + val = str(val) + ret[key] = val + return ret + + def get_out_shape(self) -> jax.ShapeDtypeStruct: + return jax.ShapeDtypeStruct( + (self.batch_size, self.num_v_heads, self.v_head_dim), + self.dtypes.act_out, + ) + + def get_scratch_shape_dict(self) -> dict[str, Any]: + conv_shape = (self.seq_tile_size, self.prev_kernel_size, 1, self.dim_size) + recurrent_shape = ( + self.seq_tile_size, + self.num_v_heads, + self.kq_head_dim, + self.v_head_dim, + ) + + carry_conv_scratch = carry_recurrent_scratch = None + # NOTE: Currently, batched mode only supports case where 1 seq = 1 tile. + # Therefore, inter tile carry is not needed. + if self.mode != GDNMode.BATCHED: + carry_conv_scratch = pltpu.VMEM(conv_shape, jnp.float32) + carry_recurrent_scratch = pltpu.VMEM(recurrent_shape, jnp.float32) + + return dict( + carry_conv_scratch_ref=carry_conv_scratch, + carry_recurrent_scratch_ref=carry_recurrent_scratch, + ) diff --git a/src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py b/src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py new file mode 100644 index 0000000000..6a0658ce1a --- /dev/null +++ b/src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py @@ -0,0 +1,1936 @@ +# 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. + +"""Canonical Gated Delta Net (GDN) backward pass using Pallas emit_pipeline. + +Decoupled GDN Architecture: +1. Conv1D forward is paired with SiLU in pure JAX (`conv1d_silu_fwd`), and GDN + forward caches the triangular inverse matrices (t_inv) in residuals while + running dense systolic matmuls on the TPU MXU. +2. Pallas GDN Backward: `pallas_gdn_bwd_kernel` executes the 40+ GDN + adjoint matrix recurrences via `pltpu.emit_pipeline` with vectorized head + processing and zero HBM intermediate spills. +3. Decoupled Conv1D Backward: `conv1d_silu_bwd` executes immediately after + Pallas via `lax.conv_general_dilated` in native JAX. + +Why Decoupled is Optimal: +Decoupling Conv1D backward into native JAX immediately after Pallas avoids +binding Conv1D gradient live ranges with GDN adjoint matrix state buffers, +eliminates vector register spills, keeps peak VMEM well within fast VMEM +boundaries, and maximizes training throughput. +""" + +import functools +import math +from typing import Any, Optional, Tuple + +import jax +from jax.ad_checkpoint import checkpoint_name +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +try: + from maxtext.layers import normalizations +except ImportError: + from maxtext.src.maxtext.layers import normalizations + +try: + import jax.experimental.xla_metadata + if not hasattr(jax.experimental.xla_metadata, "must_fuse_call"): + jax.experimental.xla_metadata.must_fuse_call = ( + lambda *args, **kwargs: (lambda fn: fn) + ) +except Exception: + pass + + +try: + from maxtext.models.kernels.gdn import compute_gdn as local_compute_gdn + from maxtext.models.kernels.gdn import wrapper as local_gdn_wrapper +except ImportError: + try: + from maxtext.src.maxtext.models.kernels.gdn import compute_gdn as local_compute_gdn + from maxtext.src.maxtext.models.kernels.gdn import wrapper as local_gdn_wrapper + except ImportError: + from . import compute_gdn as local_compute_gdn + from . import wrapper as local_gdn_wrapper + + +# ============================================================================== +# SECTION 1: CPU Interpret & Runtime Helpers +# ============================================================================== + + +def ensure_cpu_interpret_registered() -> None: + """Ensures Pallas CPU interpretation registers TPU hardware info without top-level import side-effects.""" + try: + from jax._src.pallas.mosaic import tpu_info as _ti # noqa: E402 + + if "cpu" not in _ti.registry: + _ti.registry["cpu"] = lambda: _ti.get_tpu_info_for_chip( + _ti.ChipVersion.TPU_V6E, 1 + ) + try: + _ti.get_tpu_info.cache_clear() + except Exception: + pass + except Exception: # pragma: no cover + pass + + try: + from jax._src.pallas.mosaic import pipeline as _pl_pipeline # noqa: E402 + + if not getattr(_pl_pipeline, "_is_cpu_safe_cbs_patched", False): + _orig_cbs = getattr( + _pl_pipeline, + "_original_create_bounded_slice", + _pl_pipeline._create_bounded_slice, + ) + _pl_pipeline._original_create_bounded_slice = _orig_cbs + + def _cpu_safe_create_bounded_slice( + slice_start, + slice_size, + block_size, + dim_size, + tiling=None, + *args, + **kwargs, + ): + if isinstance(slice_size, int) and ( + tiling is None or slice_size % tiling == 0 + ): + return pl.ds(slice_start, slice_size) + return _orig_cbs( + slice_start, + slice_size, + block_size, + dim_size, + tiling, + *args, + **kwargs, + ) + + _pl_pipeline._create_bounded_slice = _cpu_safe_create_bounded_slice + _pl_pipeline._is_cpu_safe_cbs_patched = True + except Exception: # pragma: no cover + pass + + +@jax.custom_vjp +def invert_triangular_matrix(t: jax.Array) -> jax.Array: + """Computes inverse of unit lower-triangular matrix using Tokamax block forward substitution.""" + return local_compute_gdn.invert_triangular_matrix(t, block_size=16) + + +def _invert_triangular_matrix_fwd(t: jax.Array): + t_inv = invert_triangular_matrix(t) + return t_inv, t_inv + + +def _invert_triangular_matrix_bwd(res, g): + t_inv = res + grad_t = jnp.tril(-(t_inv.mT @ g @ t_inv.mT), k=-1) + return (grad_t,) + + +invert_triangular_matrix.defvjp( + _invert_triangular_matrix_fwd, + _invert_triangular_matrix_bwd, +) + + +def chunk_forward( + q: jax.Array, + k: jax.Array, + v: jax.Array, + b_val: jax.Array, + a_val: jax.Array, + a_log_val: jax.Array, + dt_bias_val: jax.Array, + state_prev: jax.Array, + *, + kq_head_dim: int, + repeats: int, + chunk_size: int, + use_qk_norm_in_gdn: bool = False, +) -> Tuple[jax.Array, jax.Array]: + """Computes one chunk forward pass for GDN v3 with WY delta rule.""" + out, state_new, _ = chunk_forward_with_tinv( + q=q, + k=k, + v=v, + b_val=b_val, + a_val=a_val, + a_log_val=a_log_val, + dt_bias_val=dt_bias_val, + state_prev=state_prev, + kq_head_dim=kq_head_dim, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + return out, state_new + + +def chunk_forward_with_tinv( + q: jax.Array, + k: jax.Array, + v: jax.Array, + b_val: jax.Array, + a_val: jax.Array, + a_log_val: jax.Array, + dt_bias_val: jax.Array, + state_prev: jax.Array, + *, + kq_head_dim: int, + repeats: int, + chunk_size: int, + use_qk_norm_in_gdn: bool = False, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Computes one chunk forward pass for GDN v3 and returns (out, state_new, t_inv).""" + q = q.astype(jnp.float32) + k = k.astype(jnp.float32) + v = v.astype(jnp.float32) + if use_qk_norm_in_gdn: + q = normalizations.l2norm(q, dim=-1, eps=1e-6) + k = normalizations.l2norm(k, dim=-1, eps=1e-6) + scale = 1.0 / jnp.sqrt(kq_head_dim) + q = q * scale + b_val = b_val.astype(jnp.float32) + a_val = a_val.astype(jnp.float32) + a_log_val = a_log_val.astype(jnp.float32) + dt_bias_val = dt_bias_val.astype(jnp.float32) + state_prev = state_prev.astype(jnp.float32) + q_rep = jnp.repeat(q, repeats, axis=1) + k_rep = jnp.repeat(k, repeats, axis=1) + + beta = jax.nn.sigmoid(b_val) + + # EXACT GDN v3 gating formula + log_g = -jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val) + + # Fast MXU cumsum replacement + mask_cumsum = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=log_g.dtype)) + cumsum_log_g = jnp.dot(mask_cumsum, log_g) + + # Transpose to head-first: (H, C, D) + q_h = jnp.transpose(q_rep, (1, 0, 2)) + k_h = jnp.transpose(k_rep, (1, 0, 2)) + v_h = jnp.transpose(v, (1, 0, 2)) + beta_h = jnp.transpose(beta, (1, 0)) + cumsum_h = jnp.transpose(cumsum_log_g, (1, 0)) + + diff = cumsum_h[:, :, None] - cumsum_h[:, None, :] + mask_strict = jnp.tril( + jnp.ones((chunk_size, chunk_size), dtype=diff.dtype), k=-1 + ) + safe_diff_strict = jnp.where(mask_strict[None, :, :] == 1.0, diff, -1e4) + g_mat_strict = jnp.exp(safe_diff_strict) * mask_strict[None, :, :] + + mask_causal = jnp.tril( + jnp.ones((chunk_size, chunk_size), dtype=diff.dtype), k=0 + ) + safe_diff_causal = jnp.where(mask_causal[None, :, :] == 1.0, diff, -1e4) + g_mat_causal = jnp.exp(safe_diff_causal) * mask_causal[None, :, :] + + gating_forward = jnp.exp(cumsum_h)[:, :, None] + gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None] + gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None] + + # WY Representation: T = unit lower-triangular Gram matrix + k_beta = k_h * beta_h[:, :, None] + S = jnp.matmul(k_beta, jnp.swapaxes(k_h, -1, -2)) * g_mat_strict + identity_mask = jnp.eye(chunk_size, dtype=S.dtype)[None, :, :] + t = jnp.where(identity_mask == 1.0, 1.0, S) + A = invert_triangular_matrix(t) + + v_beta = v_h * beta_h[:, :, None] + k_beta_g = k_beta * gating_forward + u = jnp.matmul(A, v_beta) + w = jnp.matmul(A, k_beta_g) + + # Delta error subtraction against recurrent state + ws = jnp.matmul(w, state_prev) + v_new = u - ws + + # Output: cross-chunk state read + intra-chunk attention with v_new + q_g = q_h * gating_forward + out_cross = jnp.matmul(q_g, state_prev) + + attn = jnp.matmul(q_h, jnp.swapaxes(k_h, -1, -2)) * g_mat_causal + out_intra = jnp.matmul(attn, v_new) + + out = out_cross + out_intra + out = jnp.transpose(out, (1, 0, 2)) + + # State update: decayed previous state + rank-1 update from chunk with v_new + state_prev_decayed = state_prev * gating_last + k_scaled = k_h * gating_backward + state_new_intra = jnp.matmul(jnp.swapaxes(k_scaled, -1, -2), v_new) + state_new = state_prev_decayed + state_new_intra + + return out, state_new, A + + +def chunk_state_forward_with_cached_tinv( + k: jax.Array, + v: jax.Array, + b_val: jax.Array, + a_val: jax.Array, + a_log_val: jax.Array, + dt_bias_val: jax.Array, + state_prev: jax.Array, + t_inv: jax.Array, + *, + repeats: int, + chunk_size: int, + use_qk_norm_in_gdn: bool = False, +) -> jax.Array: + """Computes next recurrent state for one chunk using cached t_inv without triangular inversion.""" + v = v.astype(jnp.float32) + k = k.astype(jnp.float32) + if use_qk_norm_in_gdn: + k = normalizations.l2norm(k, dim=-1, eps=1e-6) + k_rep = jnp.repeat(k, repeats, axis=1) + + b_val = b_val.astype(jnp.float32) + a_val = a_val.astype(jnp.float32) + a_log_val = a_log_val.astype(jnp.float32) + dt_bias_val = dt_bias_val.astype(jnp.float32) + state_prev = state_prev.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_cumsum = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=log_g.dtype)) + cumsum_log_g = jnp.dot(mask_cumsum, log_g) + + k_h = jnp.transpose(k_rep, (1, 0, 2)) + v_h = jnp.transpose(v, (1, 0, 2)) + beta_h = jnp.transpose(beta, (1, 0)) + cumsum_h = jnp.transpose(cumsum_log_g, (1, 0)) + + gating_forward = jnp.exp(cumsum_h)[:, :, None] + gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None] + gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None] + + A = t_inv.astype(jnp.float32) + k_beta = k_h * beta_h[:, :, None] + v_beta = v_h * beta_h[:, :, None] + k_beta_g = k_beta * gating_forward + + u = jnp.matmul(A, v_beta) + w = jnp.matmul(A, k_beta_g) + + ws = jnp.matmul(w, state_prev) + v_new = u - ws + + state_prev_decayed = state_prev * gating_last + k_scaled = k_h * gating_backward + state_new_intra = jnp.matmul(jnp.swapaxes(k_scaled, -1, -2), v_new) + state_new = state_prev_decayed + state_new_intra + + return state_new + + +# ============================================================================== +# SECTION 2: Conv1D + SiLU Forward & Backward Primitives (Pure JAX) +# ============================================================================== + + +def conv1d_silu_fwd( + qkv: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + kernel_size: int, +) -> Tuple[jax.Array, jax.Array]: + """Forward Conv1D + SiLU returning (conv_out, qkv_conv).""" + batch, seq_len, dim_size = qkv.shape + if conv_weight.ndim == 3: + conv_weight_3d = conv_weight.astype(jnp.float32) + else: + conv_weight_3d = conv_weight[:, None, :].astype(jnp.float32) + + conv_input = jnp.pad( + qkv.astype(jnp.float32), ((0, 0), (kernel_size - 1, 0), (0, 0)) + ) + conv_out = sum( + conv_input[:, k : k + seq_len, :] * conv_weight_3d[k, 0, :] + for k in range(kernel_size) + ) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(jnp.float32) + qkv_conv = jax.nn.silu(conv_out) + return conv_out, qkv_conv.astype(qkv.dtype) + + +def conv1d_silu_bwd( + qkv: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + dy: jax.Array, + kernel_size: int, +) -> Tuple[jax.Array, jax.Array, Optional[jax.Array]]: + """Dedicated Conv1D + SiLU backward pass using JAX primitives.""" + batch, seq_len, dim_size = qkv.shape + if conv_weight.ndim == 3: + conv_weight_3d = conv_weight.astype(jnp.float32) + else: + conv_weight_3d = conv_weight[:, None, :].astype(jnp.float32) + + # 1. Forward pass: z = conv1d(x) + b + conv_input = jnp.pad( + qkv.astype(jnp.float32), ((0, 0), (kernel_size - 1, 0), (0, 0)) + ) + conv_out = sum( + conv_input[:, k : k + seq_len, :] * conv_weight_3d[k, 0, :] + for k in range(kernel_size) + ) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(jnp.float32) + z = conv_out + + # 2. Adjoint: dz = dy * SiLU'(z) + sig_z = jax.nn.sigmoid(z) + silu_prime = sig_z * (1.0 + z * (1.0 - sig_z)) + dz = dy.astype(jnp.float32) * silu_prime + + # 3. Parameter gradients: + # Bias gradient: db = sum(dz) + if conv_bias is not None: + db = jnp.sum(dz, axis=(0, 1)).astype(conv_bias.dtype) + if conv_bias.ndim != 1: + db = db.reshape(conv_bias.shape) + else: + db = None + + # Weight gradient: dw[k] = sum_{b,t} dz[b,t] * conv_input[b, t+k] + dw_rows = [] + for k in range(kernel_size): + x_k = conv_input[:, k : k + seq_len, :] + dw_rows.append(jnp.sum(dz * x_k, axis=(0, 1))) + dw = jnp.stack(dw_rows, axis=0) + if conv_weight.ndim == 3: + dw = dw[:, None, :].astype(conv_weight.dtype) + else: + dw = dw.astype(conv_weight.dtype) + + # 4. Input gradient: dx = transposed convolution of dz with reversed w + dz_pad = jnp.pad(dz, ((0, 0), (0, kernel_size - 1), (0, 0))) + w_rev = conv_weight_3d[::-1] + dx = sum( + dz_pad[:, k : k + seq_len, :] * w_rev[k, 0, :] + for k in range(kernel_size) + ) + dx = dx.astype(qkv.dtype) + + return dx, dw, db + + +# ============================================================================== +# SECTION 3: Pallas GDN Backward Pipeline Kernel (Mosaic TPU) +# ============================================================================== + + +def make_bwd_block_specs( + batch_size: int, + num_chunks: int, + chunk_size: int, + dim_size: int, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + padded_num_v_heads: int | None = None, + g: Any = None, + kernel_size: int = 4, + pad_len: int = 8, +) -> Tuple[list[pl.BlockSpec], list[pl.BlockSpec], int, int]: + """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward.""" + del batch_size, kernel_size, pad_len, g + if padded_num_v_heads is None: + padded_num_v_heads = ((num_v_heads + 127) // 128) * 128 + rc = lambda c: num_chunks - 1 - c + in_specs = [ + pl.BlockSpec( + (None, None, chunk_size, dim_size), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, num_v_heads, v_head_dim), + lambda b, c: (b, rc(c), 0, 0, 0), + ), + pl.BlockSpec( + (None, None, num_v_heads, kq_head_dim, v_head_dim), + lambda b, c: (b, rc(c), 0, 0, 0), + ), + pl.BlockSpec( + (None, None, num_v_heads, chunk_size, chunk_size), + lambda b, c: (b, rc(c), 0, 0, 0), + ), + pl.BlockSpec( + (None, 1, padded_num_v_heads), + lambda b, c: (b, 0, 0), + ), + pl.BlockSpec( + (None, 1, padded_num_v_heads), + lambda b, c: (b, 0, 0), + ), + pl.BlockSpec( + (None, None, 1, 128), + lambda b, c: (b, rc(c), 0, 0), + ), + ] + out_specs = [ + pl.BlockSpec( + (None, None, chunk_size, dim_size), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, 1, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, 1, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + ] + return in_specs, out_specs, len(in_specs), len(out_specs) + + +def _bwd_gdn_pipeline_body( + qkv_conv_ref: Any, + b_ref: Any, + a_ref: Any, + do_ref: Any, + chunk_states_ref: Any, + t_inv_ref: Any, + a_log_ref: Any, + dt_bias_ref: Any, + reset_ref: Any, + dy_conv_ref: Any, + d_b_ref: Any, + d_a_ref: Any, + d_a_log_ref: Any, + d_dt_bias_ref: Any, + d_state_scr: Any, + *, + chunk_size: int, + dim_size: int, + num_kq_heads: int, + num_v_heads: int, + padded_num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + use_qk_norm_in_gdn: bool, + kernel_size: int = 4, + pad_len: int = 8, +) -> None: + """Inner kernel executed per (batch, chunk) by emit_pipeline with manual GDN backward.""" + del kernel_size, pad_len + c = pl.program_id(1) + repeats = num_v_heads // num_kq_heads + q_size = num_kq_heads * kq_head_dim + k_size = num_kq_heads * kq_head_dim + v_size = num_v_heads * v_head_dim + + @pl.when(c == 0) + def _init(): + d_state_scr[...] = jnp.zeros( + (num_v_heads, kq_head_dim, v_head_dim), dtype=jnp.float32 + ) + + is_reset = reset_ref[...][0, 0] > 0.5 + d_state = jnp.where(is_reset, 0.0, d_state_scr[...]) + 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) + ) + + b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32) + a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32) + do_val = do_ref[...].astype(jnp.float32) + state_prev = chunk_states_ref[...].astype(jnp.float32) + t_inv_val = t_inv_ref[...].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) + + scale = 1.0 / jnp.sqrt(kq_head_dim) + mask_cumsum = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32)) + mask_strict = jnp.tril( + jnp.ones((chunk_size, chunk_size), dtype=jnp.float32), k=-1 + ) + mask_causal = jnp.tril( + jnp.ones((chunk_size, chunk_size), dtype=jnp.float32), k=0 + ) + + if use_qk_norm_in_gdn: + norm_q = normalizations.l2norm(q_orig, dim=-1, eps=1e-6) + norm_k = normalizations.l2norm(k_orig, dim=-1, eps=1e-6) + q_scaled = norm_q * scale + k_scaled_val = norm_k + else: + q_scaled = q_orig * scale + k_scaled_val = k_orig + + q_rep = jnp.repeat(q_scaled, repeats, axis=1) + k_rep = jnp.repeat(k_scaled_val, repeats, axis=1) + beta = jax.nn.sigmoid(b_val) + + sp_input = a_val + dt_bias_val + sp_val = jax.nn.softplus(sp_input) + exp_a_log = jnp.exp(a_log_val) + log_g = -exp_a_log * sp_val + + cumsum_log_g = jnp.dot(mask_cumsum, log_g) + + q_h = jnp.transpose(q_rep, (1, 0, 2)) + k_h = jnp.transpose(k_rep, (1, 0, 2)) + v_h = jnp.transpose(v, (1, 0, 2)) + beta_h = jnp.transpose(beta, (1, 0)) + cumsum_h = jnp.transpose(cumsum_log_g, (1, 0)) + do_h = jnp.transpose(do_val, (1, 0, 2)) + + diff = cumsum_h[:, :, None] - cumsum_h[:, None, :] + safe_diff_strict = jnp.where(mask_strict[None, :, :] == 1.0, diff, -1e4) + g_mat_strict = jnp.exp(safe_diff_strict) * mask_strict[None, :, :] + + safe_diff_causal = jnp.where(mask_causal[None, :, :] == 1.0, diff, -1e4) + g_mat_causal = jnp.exp(safe_diff_causal) * mask_causal[None, :, :] + + gating_forward = jnp.exp(cumsum_h)[:, :, None] + gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None] + gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None] + + k_beta = k_h * beta_h[:, :, None] + k_h_T = jnp.swapaxes(k_h, -1, -2) + S_unmasked = jnp.matmul(k_beta, k_h_T) + + A = t_inv_val + + v_beta = v_h * beta_h[:, :, None] + k_beta_g = k_beta * gating_forward + + u = jnp.matmul(A, v_beta) + w = jnp.matmul(A, k_beta_g) + + ws = jnp.matmul(w, state_prev) + v_new = u - ws + + q_g = q_h * gating_forward + attn_unmasked = jnp.matmul(q_h, k_h_T) + attn = attn_unmasked * g_mat_causal + + k_scaled_bwd = k_h * gating_backward + + dv_attn = jnp.matmul(jnp.swapaxes(attn, -1, -2), do_h) + + dv_new = dv_attn + jnp.matmul(k_scaled_bwd, d_state) + d_attn = jnp.matmul(do_h, jnp.swapaxes(v_new, -1, -2)) + + du = dv_new + dw = -jnp.matmul(dv_new, jnp.swapaxes(state_prev, -1, -2)) + + d_state_prev = ( + d_state * gating_last + + jnp.matmul(jnp.swapaxes(q_g, -1, -2), do_h) + - jnp.matmul(jnp.swapaxes(w, -1, -2), dv_new) + ) + + A_T = jnp.swapaxes(A, -1, -2) + d_v_beta = jnp.matmul(A_T, du) + d_k_beta_g = jnp.matmul(A_T, dw) + + dA = jnp.matmul(du, jnp.swapaxes(v_beta, -1, -2)) + jnp.matmul( + dw, jnp.swapaxes(k_beta_g, -1, -2) + ) + + dS = jnp.tril(-jnp.matmul(jnp.matmul(A_T, dA), A_T), k=-1) + + d_S_unmasked = dS * g_mat_strict + d_k_beta_from_S = jnp.matmul(d_S_unmasked, k_h) + d_k_h_from_S = jnp.matmul(jnp.swapaxes(d_S_unmasked, -1, -2), k_beta) + + d_k_beta = d_k_beta_g * gating_forward + d_k_beta_from_S + d_beta_h = jnp.sum(d_v_beta * v_h, axis=-1) + jnp.sum( + d_k_beta * k_h, axis=-1 + ) + d_v_h = d_v_beta * beta_h[:, :, None] + + d_attn_unmasked = d_attn * g_mat_causal + d_q_h_from_attn = jnp.matmul(d_attn_unmasked, k_h) + d_k_h_from_attn = jnp.matmul(jnp.swapaxes(d_attn_unmasked, -1, -2), q_h) + + d_q_g = jnp.matmul(do_h, jnp.swapaxes(state_prev, -1, -2)) + d_q_h_from_q_g = d_q_g * gating_forward + + d_k_scaled = jnp.matmul(v_new, jnp.swapaxes(d_state, -1, -2)) + d_k_h_from_k_scaled = d_k_scaled * gating_backward + + d_q_h = d_q_h_from_q_g + d_q_h_from_attn + d_k_h = ( + d_k_h_from_attn + + d_k_h_from_S + + d_k_beta * beta_h[:, :, None] + + d_k_h_from_k_scaled + ) + + # Gating adjoints + d_gating_forward = jnp.sum(d_q_g * q_h, axis=-1) + jnp.sum( + d_k_beta_g * k_beta, axis=-1 + ) + d_cumsum_from_fwd = d_gating_forward * gating_forward[:, :, 0] + + d_gating_last = jnp.sum(d_state * state_prev, axis=(-1, -2)) + d_cumsum_last = d_gating_last * jnp.exp(cumsum_h[:, -1]) + + d_gating_backward = jnp.sum(d_k_scaled * k_h, axis=-1) + d_diff_bwd = d_gating_backward * gating_backward[:, :, 0] + + d_g_strict = dS * S_unmasked + d_g_causal = d_attn * attn_unmasked + d_diff = (d_g_strict * g_mat_strict) + (d_g_causal * g_mat_causal) + d_cumsum_from_diff = jnp.sum(d_diff, axis=2) - jnp.sum(d_diff, axis=1) + + d_cumsum_h = d_cumsum_from_fwd + d_cumsum_from_diff - d_diff_bwd + last_col_addition = (d_cumsum_last + jnp.sum(d_diff_bwd, axis=-1))[:, None] + d_cumsum_h = d_cumsum_h + jnp.pad( + last_col_addition, ((0, 0), (chunk_size - 1, 0)) + ) + + d_cumsum_log_g = jnp.transpose(d_cumsum_h, (1, 0)) + d_log_g = jnp.dot(mask_cumsum.T, d_cumsum_log_g) + + sig_sp = jax.nn.sigmoid(sp_input) + d_a_val = d_log_g * (-exp_a_log * sig_sp) + d_a_log_val = jnp.sum(d_log_g * (-exp_a_log * sp_val), axis=0) + d_dt_bias_val = jnp.sum(d_a_val, axis=0) + + d_b_val = (jnp.transpose(d_beta_h, (1, 0))) * beta * (1.0 - beta) + + d_v_val = jnp.transpose(d_v_h, (1, 0, 2)) + d_q_rep = jnp.transpose(d_q_h, (1, 0, 2)) + d_k_rep = jnp.transpose(d_k_h, (1, 0, 2)) + + d_q_proj = jnp.sum( + d_q_rep.reshape(chunk_size, num_kq_heads, repeats, kq_head_dim), + axis=2, + ) + d_k_proj = jnp.sum( + d_k_rep.reshape(chunk_size, num_kq_heads, repeats, kq_head_dim), + axis=2, + ) + + if use_qk_norm_in_gdn: + d_q_scaled = d_q_proj * scale + r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-12) + q_unit = q_orig / r_q + d_q = ( + d_q_scaled + - q_unit * jnp.sum(d_q_scaled * q_unit, axis=-1, keepdims=True) + ) / r_q + + r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-12) + k_unit = k_orig / r_k + d_k = ( + d_k_proj + - k_unit * jnp.sum(d_k_proj * k_unit, axis=-1, keepdims=True) + ) / r_k + else: + d_q = d_q_proj * scale + d_k = d_k_proj + + # 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) + + 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) + d_a_ref[...] = jnp.pad( + d_a_val, ((0, 0), (0, padded_num_v_heads - num_v_heads)) + ).astype(d_a_ref.dtype) + d_a_log_ref[...] = jnp.pad( + d_a_log_val[None, :], ((0, 0), (0, padded_num_v_heads - num_v_heads)) + ).astype(d_a_log_ref.dtype) + d_dt_bias_ref[...] = jnp.pad( + d_dt_bias_val[None, :], ((0, 0), (0, padded_num_v_heads - num_v_heads)) + ).astype(d_dt_bias_ref.dtype) + else: + d_b_ref[...] = d_b_val.astype(d_b_ref.dtype) + d_a_ref[...] = d_a_val.astype(d_a_ref.dtype) + d_a_log_ref[...] = d_a_log_val[None, :].astype(d_a_log_ref.dtype) + d_dt_bias_ref[...] = d_dt_bias_val[None, :].astype(d_dt_bias_ref.dtype) + + d_state_scr[...] = d_state_prev.astype(d_state_scr.dtype) + + +def _pallas_gdn_bwd_kernel_single_group( + qkv_conv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + chunk_states: jax.Array, + t_inv: jax.Array, + *, + num_v_heads: int, + num_kq_heads: int, + kq_head_dim: int, + v_head_dim: int, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, + vmem_limit_mb: Optional[int] = None, + segment_ids: Optional[jax.Array] = None, + interpret: bool | pltpu.InterpretParams | None = None, + name: str = "gdn_bwd_kernel", +) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Executes single head-group Pallas emit_pipeline kernel.""" + batch_size, seq_len, group_dim_size = qkv_conv.shape + num_chunks = seq_len // chunk_size + padded_num_v_heads = ((num_v_heads + 127) // 128) * 128 + + # Reshape inputs into chunked tensors for emit_pipeline + qkv_conv_4d = qkv_conv.reshape( + batch_size, num_chunks, chunk_size, group_dim_size + ) + + b_4d = b.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + if padded_num_v_heads > num_v_heads: + b_4d = jnp.pad( + b_4d, + ((0, 0), (0, 0), (0, 0), (0, padded_num_v_heads - num_v_heads)), + ) + + a_4d = a.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + if padded_num_v_heads > num_v_heads: + a_4d = jnp.pad( + a_4d, + ((0, 0), (0, 0), (0, 0), (0, padded_num_v_heads - num_v_heads)), + ) + + do_5d = do.reshape( + batch_size, num_chunks, chunk_size, num_v_heads, v_head_dim + ) + + if a_log.ndim == 1: + a_log_3d = jnp.broadcast_to( + a_log[None, None, :], (batch_size, 1, num_v_heads) + ) + elif a_log.ndim == 2: + a_log_3d = a_log[:, None, :] + else: + a_log_3d = a_log + if padded_num_v_heads > num_v_heads: + a_log_3d = jnp.pad( + a_log_3d, ((0, 0), (0, 0), (0, padded_num_v_heads - num_v_heads)) + ) + + if dt_bias.ndim == 1: + dt_bias_3d = jnp.broadcast_to( + dt_bias[None, None, :], (batch_size, 1, num_v_heads) + ) + elif dt_bias.ndim == 2: + dt_bias_3d = dt_bias[:, None, :] + else: + dt_bias_3d = dt_bias + if padded_num_v_heads > num_v_heads: + dt_bias_3d = jnp.pad( + dt_bias_3d, ((0, 0), (0, 0), (0, padded_num_v_heads - num_v_heads)) + ) + + t_inv_5d = t_inv.astype(jnp.float32) + + # Segment reset tensor for cross-document boundary gradient reset + if segment_ids is not None and num_chunks > 1: + end_idx = jnp.arange(1, num_chunks) * chunk_size - 1 + start_next_idx = jnp.arange(1, num_chunks) * chunk_size + boundaries = segment_ids[:, end_idx] != segment_ids[:, start_next_idx] + reset_mask = jnp.pad(boundaries, ((0, 0), (0, 1)), constant_values=False) + reset_hbm = jnp.pad( + reset_mask[:, :, None, None].astype(jnp.float32), + ((0, 0), (0, 0), (0, 0), (0, 127)), + ) + else: + reset_hbm = jnp.zeros((batch_size, num_chunks, 1, 128), dtype=jnp.float32) + + in_specs, out_specs, nin, nout = make_bwd_block_specs( + batch_size=batch_size, + num_chunks=num_chunks, + chunk_size=chunk_size, + dim_size=group_dim_size, + num_v_heads=num_v_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + padded_num_v_heads=padded_num_v_heads, + ) + + out_shapes = ( + jax.ShapeDtypeStruct( + (batch_size, num_chunks, chunk_size, group_dim_size), + qkv_conv.dtype, + ), + jax.ShapeDtypeStruct(b_4d.shape, b_4d.dtype), + jax.ShapeDtypeStruct(a_4d.shape, a_4d.dtype), + jax.ShapeDtypeStruct( + (batch_size, num_chunks, 1, padded_num_v_heads), + a_log_3d.dtype, + ), + jax.ShapeDtypeStruct( + (batch_size, num_chunks, 1, padded_num_v_heads), + dt_bias_3d.dtype, + ), + ) + + body = functools.partial( + _bwd_gdn_pipeline_body, + chunk_size=chunk_size, + dim_size=group_dim_size, + num_kq_heads=num_kq_heads, + num_v_heads=num_v_heads, + padded_num_v_heads=padded_num_v_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + + def outer(*refs): + pltpu.emit_pipeline( + body, + grid=(batch_size, num_chunks), + in_specs=in_specs, + out_specs=out_specs, + )(*refs[: nin + nout], scratches=tuple(refs[nin + nout :])) + + if vmem_limit_mb is not None and vmem_limit_mb <= 64: + vmem_limit_bytes = int(vmem_limit_mb) * 1024 * 1024 + else: + tpu_info = pltpu.get_tpu_info() + vmem_limit_bytes = int(0.85 * tpu_info.vmem_capacity_bytes) + + hbm = pltpu.MemorySpace.HBM + ( + dy_conv_chunks, + d_b_chunks, + d_a_chunks, + d_a_log_chunks, + d_dt_bias_chunks, + ) = pl.pallas_call( + outer, + grid=(), + out_shape=out_shapes, + in_specs=[pl.BlockSpec(memory_space=hbm)] * nin, + out_specs=[pl.BlockSpec(memory_space=hbm)] * nout, + scratch_shapes=[ + pltpu.VMEM((num_v_heads, kq_head_dim, v_head_dim), jnp.float32), + ], + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ), + interpret=interpret, + name=name, + )( + qkv_conv_4d, + b_4d, + a_4d, + do_5d, + chunk_states, + t_inv_5d, + a_log_3d, + dt_bias_3d, + reset_hbm, + ) + + d_a_log_reduced = ( + jnp.sum(d_a_log_chunks[..., 0, :num_v_heads], axis=(0, 1)) + .astype(a_log.dtype) + ) + d_dt_bias_reduced = ( + jnp.sum(d_dt_bias_chunks[..., 0, :num_v_heads], axis=(0, 1)) + .astype(dt_bias.dtype) + ) + d_b_flat = ( + d_b_chunks[..., :num_v_heads] + .reshape(batch_size, seq_len, num_v_heads) + .astype(b.dtype) + ) + d_a_flat = ( + d_a_chunks[..., :num_v_heads] + .reshape(batch_size, seq_len, num_v_heads) + .astype(a.dtype) + ) + dy_conv_flat = dy_conv_chunks.reshape( + batch_size, seq_len, group_dim_size + ).astype(qkv_conv.dtype) + + return ( + dy_conv_flat, + d_b_flat, + d_a_flat, + d_a_log_reduced, + d_dt_bias_reduced, + ) + + +def pallas_gdn_bwd_kernel( + qkv_conv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + chunk_states: 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, + vmem_limit_mb: Optional[int] = None, + head_tile: Optional[int] = None, + segment_ids: Optional[jax.Array] = None, + interpret: bool | pltpu.InterpretParams | None = None, +) -> Tuple[ + jax.Array, + jax.Array, + jax.Array, + jax.Array, + jax.Array, +]: + """Executes the Pallas reverse-chunk GDNv3 backward kernel using emit_pipeline with native contiguous streaming per group.""" + if interpret is None and jax.default_backend() == "cpu": + interpret = True + if interpret: + ensure_cpu_interpret_registered() + + 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 + + target_tile = 16 if head_tile is None else head_tile + max_possible = min(num_v_heads, target_tile) + tile_v_heads = None + for candidate in range(max_possible, 0, -1): + if num_v_heads % candidate == 0 and candidate % repeats == 0: + tile_v_heads = candidate + break + if tile_v_heads is None: + tile_v_heads = repeats if num_v_heads % repeats == 0 else num_v_heads + num_groups = num_v_heads // tile_v_heads + tile_kq_heads = tile_v_heads // repeats + + if num_groups == 1: + return _pallas_gdn_bwd_kernel_single_group( + qkv_conv=qkv_conv, + 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, + num_kq_heads=num_kq_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + vmem_limit_mb=vmem_limit_mb, + segment_ids=segment_ids, + interpret=interpret, + name="gdn_bwd_kernel", + ) + + q_size = num_kq_heads * kq_head_dim + k_size = num_kq_heads * kq_head_dim + tile_q_size = tile_kq_heads * kq_head_dim + tile_k_size = tile_kq_heads * kq_head_dim + tile_v_size = tile_v_heads * v_head_dim + + dq_list = [] + dk_list = [] + dv_list = [] + db_list = [] + da_list = [] + dal_list = [] + ddt_list = [] + + for g in range(num_groups): + vh_start = g * tile_v_heads + vh_end = (g + 1) * tile_v_heads + kqh_start = g * tile_kq_heads + kqh_end = (g + 1) * tile_kq_heads + + q_g = qkv_conv[:, :, kqh_start * kq_head_dim : kqh_end * kq_head_dim] + k_g = qkv_conv[ + :, :, q_size + kqh_start * kq_head_dim : q_size + kqh_end * kq_head_dim + ] + v_g = qkv_conv[ + :, + :, + q_size + k_size + vh_start * v_head_dim : q_size + + k_size + + vh_end * v_head_dim, + ] + qkv_g = jnp.concatenate([q_g, k_g, v_g], axis=-1) + + b_g = b[:, :, vh_start:vh_end] + a_g = a[:, :, vh_start:vh_end] + do_g = do[:, :, vh_start:vh_end, :] + chunk_states_g = chunk_states[:, :, vh_start:vh_end, :, :] + t_inv_g = t_inv[:, :, vh_start:vh_end, :, :] + + if a_log.ndim == 1: + a_log_g = a_log[vh_start:vh_end] + elif a_log.ndim == 2: + a_log_g = a_log[:, vh_start:vh_end] + else: + a_log_g = a_log[:, :, vh_start:vh_end] + + if dt_bias.ndim == 1: + dt_bias_g = dt_bias[vh_start:vh_end] + elif dt_bias.ndim == 2: + dt_bias_g = dt_bias[:, vh_start:vh_end] + else: + dt_bias_g = dt_bias[:, :, vh_start:vh_end] + + dy_g, db_g, da_g, dal_g, ddt_g = _pallas_gdn_bwd_kernel_single_group( + qkv_conv=qkv_g, + b=b_g, + a=a_g, + a_log=a_log_g, + dt_bias=dt_bias_g, + do=do_g, + chunk_states=chunk_states_g, + t_inv=t_inv_g, + num_v_heads=tile_v_heads, + num_kq_heads=tile_kq_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + vmem_limit_mb=vmem_limit_mb, + segment_ids=segment_ids, + interpret=interpret, + name=f"gdn_bwd_kernel_group_{g}", + ) + + dq_g = dy_g[:, :, :tile_q_size] + dk_g = dy_g[:, :, tile_q_size : tile_q_size + tile_k_size] + dv_g = dy_g[:, :, tile_q_size + tile_k_size :] + + dq_list.append(dq_g) + dk_list.append(dk_g) + dv_list.append(dv_g) + db_list.append(db_g) + da_list.append(da_g) + dal_list.append(dal_g) + ddt_list.append(ddt_g) + + dq_flat = jnp.concatenate(dq_list, axis=-1) + dk_flat = jnp.concatenate(dk_list, axis=-1) + dv_flat = jnp.concatenate(dv_list, axis=-1) + dy_conv_flat = jnp.concatenate([dq_flat, dk_flat, dv_flat], axis=-1).astype( + qkv_conv.dtype + ) + d_b_flat = jnp.concatenate(db_list, axis=-1) + d_a_flat = jnp.concatenate(da_list, axis=-1) + d_a_log_reduced = jnp.concatenate(dal_list, axis=-1) + d_dt_bias_reduced = jnp.concatenate(ddt_list, axis=-1) + + return ( + dy_conv_flat, + d_b_flat, + d_a_flat, + d_a_log_reduced, + d_dt_bias_reduced, + ) + + +pallas_gdn_bwd_computation = pallas_gdn_bwd_kernel + + +def pallas_fused_conv1d_gdn_bwd_kernel( + pre_conv_qkv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + chunk_states: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array] = None, + t_inv: Optional[jax.Array] = None, + qkv: Optional[jax.Array] = None, + seq_lens: Optional[jax.Array] = None, + *, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + kernel_size: int = 4, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, + vmem_limit_mb: Optional[int] = None, + head_tile: Optional[int] = None, + segment_ids: Optional[jax.Array] = None, + interpret: bool | pltpu.InterpretParams | None = None, +) -> Tuple[ + jax.Array, + jax.Array, + jax.Array, + jax.Array, + Optional[jax.Array], + jax.Array, + jax.Array, +]: + """Fused Conv1D + GDN backward combining decoupled GDN bwd and Conv1D bwd.""" + del seq_lens, qkv + _, qkv_conv = conv1d_silu_fwd( + qkv=pre_conv_qkv, + conv_weight=conv_weight, + conv_bias=conv_bias, + kernel_size=kernel_size, + ) + + dy_conv, d_b, d_a, d_a_log, d_dt_bias = ( + pallas_gdn_bwd_kernel( + qkv_conv=qkv_conv, + 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, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + vmem_limit_mb=vmem_limit_mb, + head_tile=head_tile, + segment_ids=segment_ids, + interpret=interpret, + ) + ) + + d_pre_conv_qkv, d_conv_weight, d_conv_bias = conv1d_silu_bwd( + qkv=pre_conv_qkv, + conv_weight=conv_weight, + conv_bias=conv_bias, + dy=dy_conv, + kernel_size=kernel_size, + ) + + return ( + d_pre_conv_qkv, + d_b, + d_a, + d_conv_weight, + d_conv_bias, + d_a_log, + d_dt_bias, + ) + + +pallas_fused_conv1d_gdn_bwd_computation = pallas_fused_conv1d_gdn_bwd_kernel + + +# ============================================================================== +# SECTION 4: Unified GDN Custom VJP Interface (gdn_fused_conv1d) +# ============================================================================== + + +def pure_jax_fused_conv1d_gdn( + 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 = jnp.float32, +) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: + """Pure-JAX composite of Conv1D + GDN used during backward pass autodiff.""" + del conv_state + batch, seq_len, _ = qkv.shape + key_dim = num_k_heads * head_k_dim + + # Conv1D in FP32 + conv_input = jnp.pad( + qkv.astype(jnp.float32), ((0, 0), (conv_kernel_size - 1, 0), (0, 0)) + ) + if conv_weight.ndim == 3: + conv_weight_3d = conv_weight.astype(jnp.float32) + else: + conv_weight_3d = conv_weight[:, None, :].astype(jnp.float32) + conv_out = sum( + conv_input[:, k : k + seq_len, :] * conv_weight_3d[k, 0, :] + for k in range(conv_kernel_size) + ) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(jnp.float32) + qkv_conv = jax.nn.silu(conv_out).astype(jnp.float32) + + q_conv, k_conv, v_conv = jnp.split(qkv_conv, [key_dim, 2 * key_dim], axis=-1) + + # Reshape for GDN + query = q_conv.reshape(batch, seq_len, num_k_heads, head_k_dim) + key = k_conv.reshape(batch, seq_len, num_k_heads, head_k_dim) + value = v_conv.reshape(batch, seq_len, num_v_heads, head_v_dim) + + a_log_cast = jnp.asarray(a_log, dtype=jnp.float32) + dt_bias_cast = jnp.asarray(dt_bias, dtype=jnp.float32) + beta = jax.nn.sigmoid(b.astype(jnp.float32)) + g = -jnp.exp(a_log_cast) * jax.nn.softplus( + a.astype(jnp.float32) + dt_bias_cast + ) + + if num_v_heads > num_k_heads and num_v_heads % num_k_heads == 0: + repeats = num_v_heads // num_k_heads + query = jnp.repeat(query, repeats, axis=2) + key = jnp.repeat(key, repeats, axis=2) + + try: + from maxtext.models import qwen3 + except ImportError: + try: + from maxtext.src.maxtext.models import qwen3 + except ImportError: + from . import qwen3 + + core_attn_out, next_recurrent_state = qwen3.jax_chunk_gated_delta_rule( + query=query, + key=key, + value=value, + g=g, + beta=beta, + chunk_size=chunk_size, + initial_state=( + recurrent_state.astype(jnp.float32) + if recurrent_state is not None + else None + ), + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=jnp.float32, + ) + + next_conv_state = ( + qkv[:, -(conv_kernel_size - 1) :, :] + if seq_len >= conv_kernel_size - 1 + else jnp.zeros( + (batch, conv_kernel_size - 1, qkv.shape[-1]), dtype=qkv.dtype + ) + ) + if next_recurrent_state is None: + next_recurrent_state = jnp.zeros( + (batch, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32 + ) + + return core_attn_out.astype(qkv.dtype), ( + next_conv_state.astype(qkv.dtype), + next_recurrent_state.astype(jnp.float32), + ) + + +def _compute_forward_conv_and_states( + 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, + 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 = False, + compute_dtype: jnp.dtype = jnp.float32, + cached_t_inv: Optional[jax.Array] = None, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Computes convolved QKV, inter-chunk states, and t_inv matrices in FP32.""" + del compute_dtype + batch_size, seq_len, dim_size = qkv.shape + num_chunks = seq_len // chunk_size + + # Conv1D in FP32 + conv_input = jnp.pad( + qkv.astype(jnp.float32), ((0, 0), (conv_kernel_size - 1, 0), (0, 0)) + ) + if conv_weight.ndim == 3: + conv_weight_3d = conv_weight.astype(jnp.float32) + else: + conv_weight_3d = conv_weight[:, None, :].astype(jnp.float32) + conv_out = sum( + conv_input[:, k : k + seq_len, :] * conv_weight_3d[k, 0, :] + for k in range(conv_kernel_size) + ) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(jnp.float32) + qkv_conv_f32 = jax.nn.silu(conv_out).astype(jnp.float32) + + # Chunk states progression in FP32 + num_kq_heads = num_k_heads + q_size = num_kq_heads * head_k_dim + k_size = num_kq_heads * head_k_dim + repeats = num_v_heads // num_kq_heads + + q = qkv_conv_f32[:, :, :q_size].reshape( + batch_size, num_chunks, chunk_size, num_kq_heads, head_k_dim + ) + k = qkv_conv_f32[:, :, q_size : q_size + k_size].reshape( + batch_size, num_chunks, chunk_size, num_kq_heads, head_k_dim + ) + v = qkv_conv_f32[:, :, q_size + k_size :].reshape( + batch_size, num_chunks, chunk_size, num_v_heads, head_v_dim + ) + + b_4d = b.astype(jnp.float32).reshape( + batch_size, num_chunks, chunk_size, num_v_heads + ) + a_4d = a.astype(jnp.float32).reshape( + batch_size, num_chunks, chunk_size, num_v_heads + ) + a_log_f32 = a_log.astype(jnp.float32) + dt_bias_f32 = dt_bias.astype(jnp.float32) + + if recurrent_state is None: + init_state = jnp.zeros( + (batch_size, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32 + ) + else: + init_state = recurrent_state.astype(jnp.float32) + + k_chunks = k.swapaxes(0, 1) + v_chunks = v.swapaxes(0, 1) + b_chunks = b_4d.swapaxes(0, 1) + a_chunks = a_4d.swapaxes(0, 1) + + if cached_t_inv is not None: + t_inv_chunks = cached_t_inv.astype(jnp.float32).swapaxes(0, 1) + + def chunk_cached_step( + k_single, v_single, b_single, a_single, s_prev, t_inv_single + ): + return chunk_state_forward_with_cached_tinv( + k=k_single, + v=v_single, + b_val=b_single, + a_val=a_single, + a_log_val=a_log_f32, + dt_bias_val=dt_bias_f32, + state_prev=s_prev, + t_inv=t_inv_single, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + + def scan_fn_cached(carry_state, chunk_inputs): + k_i, v_i, b_i, a_i, t_inv_i = chunk_inputs + next_state = jax.vmap(chunk_cached_step)( + k_i, v_i, b_i, a_i, carry_state, t_inv_i + ) + return next_state, carry_state + + _, chunk_states = jax.lax.scan( + scan_fn_cached, + init_state, + (k_chunks, v_chunks, b_chunks, a_chunks, t_inv_chunks), + ) + chunk_states = chunk_states.swapaxes(0, 1) + t_inv_all = cached_t_inv + else: + + def chunk_step(q_single, k_single, v_single, b_single, a_single, s_prev): + return chunk_forward_with_tinv( + q_single, + k_single, + v_single, + b_single, + a_single, + a_log_f32, + dt_bias_f32, + s_prev, + kq_head_dim=head_k_dim, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + + def scan_fn(carry_state, chunk_inputs): + q_i, k_i, v_i, b_i, a_i = chunk_inputs + _, next_state, t_inv_i = jax.vmap(chunk_step)( + q_i, k_i, v_i, b_i, a_i, carry_state + ) + return next_state, (carry_state, t_inv_i) + + q_chunks = q.swapaxes(0, 1) + _, (chunk_states, t_inv_all) = jax.lax.scan( + scan_fn, init_state, (q_chunks, k_chunks, v_chunks, b_chunks, a_chunks) + ) + chunk_states = chunk_states.swapaxes(0, 1) + t_inv_all = t_inv_all.swapaxes(0, 1) + + return qkv_conv_f32.astype(qkv.dtype), chunk_states, t_inv_all + + +def _run_local_gdn_fused_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, +) -> Tuple[ + Tuple[jax.Array, Tuple[jax.Array, jax.Array]], + Optional[jax.Array], + Optional[jax.Array], +]: + """Runs local GDN fused forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU.""" + if jax.extend.backend.get_backend().platform == "cpu": + out, states = pure_jax_fused_conv1d_gdn( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + conv_state=conv_state, + recurrent_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, + ) + _, chunk_states, t_inv = _compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_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, + ) + return (out, states), t_inv, chunk_states + + batch_size, seq_len, dim_size = qkv.shape + 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) + + query_start_loc = jnp.arange( + 0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32 + ) + state_indices = jnp.arange(num_seqs, dtype=jnp.int32) + seq_lens = jnp.full((num_seqs,), seq_len, 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, + 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, + compute_precision=jnp.dtype(jnp.float32), + mixed_tile_size=chunk_size, + is_prefill_only=True, + ) + + core_attn_out = core_attn_out_flat.reshape( + batch_size, seq_len, num_v_heads, head_v_dim + ) + t_inv = t_inv_raw.astype(jnp.float32).reshape( + batch_size, num_chunks, num_v_heads, chunk_size, chunk_size + ) + chunk_states = chunk_states_raw.astype(jnp.float32).reshape( + batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim + ) + return ( + ( + core_attn_out.astype(qkv.dtype), + ( + new_conv_state[1:].astype(qkv.dtype), + new_recurrent_state[1:].astype(jnp.float32), + ), + ), + t_inv, + chunk_states, + ) + + +@functools.partial( + jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16) +) +def gdn_fused_conv1d( + 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, +) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: + """Fused Conv1D + GDN with decoupled Pallas backward pass.""" + (out, states), _, _ = _run_local_gdn_fused_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, + ) + return out, states + + +def _gdn_fused_conv1d_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, +): + qkv = checkpoint_name(qkv, "gdn_fwd_conv") + (out, states), t_inv, chunk_states = _run_local_gdn_fused_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, + ) + out = checkpoint_name(out, "gdn_fwd_out") + residuals = ( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + t_inv, + chunk_states, + ) + return (out, states), residuals + + +def _gdn_fused_conv1d_bwd( + 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, + residuals: tuple, + cotangents: tuple, +): + if len(residuals) == 11: + ( + pre_conv_qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + t_inv_fwd, + chunk_states, + ) = residuals + else: + ( + pre_conv_qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + t_inv_fwd, + ) = residuals + chunk_states = None + + d_out, d_states = cotangents + d_conv_state, d_recurrent_state = d_states + del d_conv_state, d_recurrent_state + + # Recompute forward chunk states and t_inv if not cached in residuals + if chunk_states is None or t_inv_fwd is None: + qkv_conv, chunk_states_recomputed, t_inv_recomputed = ( + _compute_forward_conv_and_states( + qkv=pre_conv_qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_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, + cached_t_inv=t_inv_fwd, + ) + ) + if chunk_states is None: + chunk_states = chunk_states_recomputed + if t_inv_fwd is None: + t_inv_fwd = t_inv_recomputed + else: + _, qkv_conv = conv1d_silu_fwd( + qkv=pre_conv_qkv, + conv_weight=conv_weight, + conv_bias=conv_bias, + kernel_size=conv_kernel_size, + ) + t_inv = t_inv_fwd + + dy_conv, d_b, d_a, d_a_log, d_dt_bias = ( + pallas_gdn_bwd_kernel( + qkv_conv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=d_out, + chunk_states=chunk_states, + t_inv=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, + ) + ) + + d_pre_conv_qkv, d_conv_weight, d_conv_bias = conv1d_silu_bwd( + qkv=pre_conv_qkv, + conv_weight=conv_weight, + conv_bias=conv_bias, + dy=dy_conv, + kernel_size=conv_kernel_size, + ) + + d_conv_state_out = None if conv_state is None else jnp.zeros_like(conv_state) + d_recurrent_state_out = ( + None if recurrent_state is None else jnp.zeros_like(recurrent_state) + ) + return ( + d_pre_conv_qkv, + d_b, + d_a, + d_conv_weight, + d_conv_bias, + d_a_log, + d_dt_bias, + d_conv_state_out, + d_recurrent_state_out, + ) + + +gdn_fused_conv1d.defvjp( + _gdn_fused_conv1d_fwd, + _gdn_fused_conv1d_bwd, +) + +__all__ = [ + "chunk_forward", + "chunk_forward_with_tinv", + "pallas_gdn_bwd_kernel", + "pallas_gdn_bwd_computation", + "pallas_fused_conv1d_gdn_bwd_kernel", + "pallas_fused_conv1d_gdn_bwd_computation", + "conv1d_silu_fwd", + "conv1d_silu_bwd", + "pure_jax_fused_conv1d_gdn", + "gdn_fused_conv1d", + "ensure_cpu_interpret_registered", +] diff --git a/src/maxtext/models/kernels/gdn/memory_ref.py b/src/maxtext/models/kernels/gdn/memory_ref.py new file mode 100644 index 0000000000..eafea0b766 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/memory_ref.py @@ -0,0 +1,611 @@ +# 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. +# ============================================================================== + +"""Weight and state reference dataclasses for VMEM.""" + +import dataclasses +import functools +from typing import Any + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + except (ImportError, ModuleNotFoundError): + from . import config + + +def _flat_pos(shape: tuple[int, ...], indices: tuple[Any, ...]) -> Any: + """Row-major flat offset of `indices` into a logical array of `shape`.""" + strides = pl.strides_from_shape(shape) + assert len(strides) == len(indices) + + pos = 0 + for stride, idx in zip(strides, indices): + pos += stride * idx + return pos + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class ConvWeightsRef: + weight: Any + bias: Any | None = None + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class GDNWeightsRef: + a_log: Any + dt_bias: Any + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class WeightRefs: + conv: ConvWeightsRef + gdn: GDNWeightsRef + + +class FieldOffset: + """Descriptor returning the record field at ``data[pos + offset]``. + + Reads a single dynamically-indexed element rather than a slice, since JAX + can't slice a range with traced indices. Read-only: metadata is never + written. + """ + + def __init__(self, offset: int): + self.offset = offset + + def __get__(self, obj, objtype=None): + if obj is None: + return self + return obj.data[obj.pos + self.offset] + + +# Per-p_id metadata is an array of structs: each p_id's fields sit contiguously +# and FieldOffset(k) reads the k-th word of its struct. +# +# Packed struct: [r_base, packed_word]. +# Fields share packed_word to save SMEM: is_first_tile(0), is_last_tile(1), r_size(2..15), s_idx(16..31). +@dataclasses.dataclass(frozen=True) +class PackedPIdRecord: + """Packed struct [r_base, packed_word]; the four small fields bit-slice word. + + Each bit field masks after shifting, which also clears the sign bits that + ``>>`` extends on the signed int32 word. + """ + + STRUCT_SIZE = 2 + FIRST_TILE_SHIFT = 0 + LAST_TILE_SHIFT = 1 + R_SIZE_SHIFT = 2 + S_IDX_SHIFT = 16 + FLAG_MASK = 1 + R_SIZE_MASK = (1 << (S_IDX_SHIFT - R_SIZE_SHIFT)) - 1 + S_IDX_MASK = (1 << (32 - S_IDX_SHIFT)) - 1 + MAX_SEQS = S_IDX_MASK + 1 + + data: Any + pos: Any + r_base = FieldOffset(0) + word = FieldOffset(1) + + @property + def s_idx(self): + return (self.word >> self.S_IDX_SHIFT) & self.S_IDX_MASK + + @property + def r_size(self): + return (self.word >> self.R_SIZE_SHIFT) & self.R_SIZE_MASK + + @property + def is_first_tile(self): + return (self.word & self.FLAG_MASK) != 0 + + @property + def is_last_tile(self): + return ((self.word >> self.LAST_TILE_SHIFT) & self.FLAG_MASK) != 0 + + @classmethod + def pack( + cls, + s_idx: jax.Array, + r_size: jax.Array, + is_first_tile: jax.Array, + is_last_tile: jax.Array, + ) -> jax.Array: + """Packs s_idx, row size and two tile-state flags into one int32 word.""" + + s_idx = s_idx.reshape(-1).astype(jnp.int32) + r_size = r_size.reshape(-1).astype(jnp.int32) + is_first_tile = is_first_tile.reshape(-1).astype(jnp.int32) + is_last_tile = is_last_tile.reshape(-1).astype(jnp.int32) + word = s_idx << cls.S_IDX_SHIFT + word |= r_size << cls.R_SIZE_SHIFT + word |= is_last_tile << cls.LAST_TILE_SHIFT + word |= is_first_tile << cls.FIRST_TILE_SHIFT + return word + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class MetadataRef: + num_tiles: Any + # Array of structs holding every p_id's metadata + records: Any + s_idx_has_initial_state: Any + s_idx_to_state_indices: Any + shape: tuple[int, ...] = dataclasses.field(metadata=dict(static=True)) + + def get_record(self, p_id, idx) -> PackedPIdRecord: + """View of one p_id's metadata: .r_base / .s_idx / .r_size / .is_*_tile.""" + record_idx = _flat_pos(self.shape, (p_id, idx)) + return PackedPIdRecord( + self.records, record_idx * PackedPIdRecord.STRUCT_SIZE + ) + + @classmethod + def create( # pyrefly: ignore[bad-override] + cls, + cfgs: config.GDNConfig, + num_tiles: jax.Array, + p_id_to_s_idx: jax.Array, + p_id_to_r_base: jax.Array, + p_id_to_r_size: jax.Array, + p_id_is_first_tile: jax.Array, + p_id_is_last_tile: jax.Array, + s_idx_has_initial_state: jax.Array, + s_idx_to_state_indices: jax.Array, + ): + # NOTE: First dim does not matter when it comes to calculating stride. + shape = (1, cfgs.seq_tile_size) + assert s_idx_has_initial_state.shape[0] <= PackedPIdRecord.MAX_SEQS, ( + f"Number of sequences ({s_idx_has_initial_state.shape[0]}) exceeds" + f" PackedPIdRecord limit ({PackedPIdRecord.MAX_SEQS})." + ) + assert cfgs.tile_size <= PackedPIdRecord.R_SIZE_MASK, ( + f"Tile size ({cfgs.tile_size}) exceeds PackedPIdRecord limit" + f" ({PackedPIdRecord.R_SIZE_MASK})." + ) + + r_base = p_id_to_r_base.reshape(-1).astype(jnp.int32) + word = PackedPIdRecord.pack( + p_id_to_s_idx, p_id_to_r_size, p_id_is_first_tile, p_id_is_last_tile + ) + fields = [r_base, word] + # Interleave fields into one array of structs: [rec0_f0, rec0_f1, ...]. + records = jnp.stack(fields, axis=-1).reshape(-1) + + return cls( + num_tiles=num_tiles, + records=records, + s_idx_has_initial_state=s_idx_has_initial_state, + s_idx_to_state_indices=s_idx_to_state_indices, + shape=shape, + ) + + def __len__(self) -> int: + return len(jax.tree_util.tree_leaves(self)) + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class BaseBufferedRef(pltpu.BufferedRef): + + cfg: config.GDNConfig = dataclasses.field(metadata=dict(static=True)) + # NOTE: Despite being ref, metadata_ref should be set to static. This is + # because the memory will be allocated outside of kernel and metadata_ref + # merely points to the reference. + metadata_ref: MetadataRef = dataclasses.field(metadata=dict(static=True)) + + @classmethod + def create( # pyrefly: ignore[bad-override] + cls, + spec: pl.BlockSpec, + dtype_or_type: jax.Array, + buffer_type: pltpu.BufferType, + buffer_count: int, + use_lookahead: bool, + cfg: config.GDNConfig, + metadata_ref: MetadataRef, + ): + standard_ref = pltpu.BufferedRef.create( + spec=spec, + dtype_or_type=dtype_or_type, + buffer_type=buffer_type, + buffer_count=buffer_count, + grid_rank=1, + use_lookahead=use_lookahead, + ) + return cls( + cfg=cfg, + metadata_ref=metadata_ref, + **{ + f.name: getattr(standard_ref, f.name) + for f in dataclasses.fields(pltpu.BufferedRef) + }, + ) + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class InBufferedRef(BaseBufferedRef): + + def copy_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + 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): + record = self.metadata_ref.get_record(p_id, idx) + r_base = record.r_base + dma_size = record.r_size + pltpu.make_async_copy( + src_ref.at[pl.ds(r_base, dma_size)], + vmem_ref.at[idx, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).start() + + def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + 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] + p_id = grid_indices[0] + + dma_size = 0 + for idx in range(self.cfg.seq_tile_size): + dma_size += self.metadata_ref.get_record(p_id, idx).r_size + + pltpu.make_async_copy( + vmem_ref.at[0, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[0, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class OutBufferedRef(BaseBufferedRef): + + def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_copy_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + record = self.metadata_ref.get_record(p_id, idx) + r_base = record.r_base + dma_size = record.r_size + pltpu.make_async_copy( + vmem_ref.at[idx, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + dst_ref.at[pl.ds(r_base, dma_size)], + sem, + ).start() + + def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_wait_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + dma_size = 0 + for idx in range(self.cfg.seq_tile_size): + dma_size += self.metadata_ref.get_record(p_id, idx).r_size + + pltpu.make_async_copy( + vmem_ref.at[0, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[0, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class TInvBufferedRef(BaseBufferedRef): + """DMA buffer for triangular inverse matrix caching (t_inv).""" + + def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_copy_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + pltpu.make_async_copy( + vmem_ref.at[idx], + dst_ref.at[p_id + idx], + sem, + ).start() + + def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_wait_out_slot + sem = self.sem_sends.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 ChunkStatesBufferedRef(BaseBufferedRef): + """DMA buffer for caching intermediate recurrent chunk states.""" + + def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_copy_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + pltpu.make_async_copy( + vmem_ref.at[idx], + dst_ref.at[p_id + idx], + sem, + ).start() + + def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_wait_out_slot + sem = self.sem_sends.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 StateBufferedRef(BaseBufferedRef): + + def copy_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + 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): + record = self.metadata_ref.get_record(p_id, idx) + is_first_tile = record.is_first_tile + s_idx = record.s_idx + state_idx = self.metadata_ref.s_idx_to_state_indices[s_idx] + has_initial_state = self.metadata_ref.s_idx_has_initial_state[s_idx] + should_read = jnp.logical_and(is_first_tile, has_initial_state) + dma_size = jnp.where(should_read, 1, 0) + + pltpu.make_async_copy( + src_ref.at[pl.ds(state_idx, dma_size)], + vmem_ref.at[pl.ds(idx, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).start() + + def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + 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] + p_id = grid_indices[0] + + dma_size = 0 + for idx in range(self.cfg.seq_tile_size): + record = self.metadata_ref.get_record(p_id, idx) + is_first_tile = record.is_first_tile + s_idx = record.s_idx + has_initial_state = self.metadata_ref.s_idx_has_initial_state[s_idx] + should_read = jnp.logical_and(is_first_tile, has_initial_state) + dma_size += jnp.where(should_read, 1, 0) + + pltpu.make_async_copy( + vmem_ref.at[pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() + + def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_copy_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + record = self.metadata_ref.get_record(p_id, idx) + is_last_tile = record.is_last_tile + s_idx = record.s_idx + state_idx = self.metadata_ref.s_idx_to_state_indices[s_idx] + dma_size = jnp.where(is_last_tile, 1, 0) + + pltpu.make_async_copy( + vmem_ref.at[pl.ds(idx, dma_size)], # pyrefly: ignore[missing-attribute] + dst_ref.at[pl.ds(state_idx, dma_size)], + sem, + ).start() + + def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_wait_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + dma_size = 0 + for idx in range(self.cfg.seq_tile_size): + is_last_tile = self.metadata_ref.get_record(p_id, idx).is_last_tile + dma_size += jnp.where(is_last_tile, 1, 0) + + pltpu.make_async_copy( + vmem_ref.at[pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() + + +def create_allocs( + metadata_ref: MetadataRef, + qkv_ref: jax.Array, + b_ref: jax.Array, + a_ref: jax.Array, + out_ref: jax.Array, + 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, ...]: + 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 = ( + cfg.seq_tile_size, + cfg.chunk_size, + cfg.num_v_heads, + cfg.v_head_dim, + ) + conv_shape = (cfg.seq_tile_size, cfg.prev_kernel_size, 1, cfg.dim_size) + recurrent_shape = ( + cfg.seq_tile_size, + cfg.num_v_heads, + cfg.kq_head_dim, + cfg.v_head_dim, + ) + + pipeline_mode = pl.Buffered(buffer_count=cfg.num_buffers, use_lookahead=False) + + block_spec_partial = functools.partial( + pl.BlockSpec, + memory_space=pltpu.VMEM, + index_map=lambda i: (i,), + pipeline_mode=pipeline_mode, + ) + + qkv_spec = block_spec_partial(block_shape=qkv_shape) + ba_spec = block_spec_partial(block_shape=ba_shape) + in_buffered_partial = functools.partial( + InBufferedRef.input, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + 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) + state_buffered_partial = functools.partial( + StateBufferedRef.input_output, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + conv_alloc = state_buffered_partial( + spec=conv_spec, dtype_or_type=conv_state_ref + ) + recurrent_alloc = state_buffered_partial( + spec=recurrent_spec, dtype_or_type=recurrent_state_ref + ) + + allocs = [ + qkv_alloc, + b_alloc, + a_alloc, + conv_alloc, + recurrent_alloc, + out_alloc, + ] + + if t_inv_ref is not None: + t_inv_shape = ( + cfg.seq_tile_size, + cfg.num_v_heads, + cfg.chunk_size, + cfg.chunk_size, + ) + t_inv_alloc = TInvBufferedRef.output( + spec=block_spec_partial(block_shape=t_inv_shape), + dtype_or_type=t_inv_ref, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + allocs.append(t_inv_alloc) + + if chunk_states_ref is not None: + chunk_states_shape = ( + cfg.seq_tile_size, + cfg.num_v_heads, + cfg.kq_head_dim, + cfg.v_head_dim, + ) + chunk_states_alloc = ChunkStatesBufferedRef.output( + spec=block_spec_partial(block_shape=chunk_states_shape), + dtype_or_type=chunk_states_ref, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + allocs.append(chunk_states_alloc) + + return tuple(allocs) diff --git a/src/maxtext/models/kernels/gdn/metadata.py b/src/maxtext/models/kernels/gdn/metadata.py new file mode 100644 index 0000000000..d69dea7eb6 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/metadata.py @@ -0,0 +1,150 @@ +# 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. +# ============================================================================== + +"""Metadata references for sequence mapping and grid distribution.""" + +import jax +from jax.experimental import pallas as pl +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config + from maxtext.models.kernels.gdn import memory_ref +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + from maxtext.src.maxtext.models.kernels.gdn import memory_ref + except (ImportError, ModuleNotFoundError): + from . import config + from . import memory_ref + + +def compute_batched_seq_metadata( + cfg: config.GDNConfig, + seq_lens: jax.Array, + query_start_loc: jax.Array, + state_indices: jax.Array, + end_seq: jax.Array, +) -> memory_ref.MetadataRef: + """Metadata for computing multiple sequences per tile.""" + + max_seqs = seq_lens.size + all_seqs = jnp.arange(max_seqs) + + # NOTE: Only supports use case where query_lens[i] = 1 where i < end_seq. + # This must be guaranteed by the function caller. + # TODO(b/534541682): Add error handling when above condition is not met. + query_lens = query_start_loc[1:] - query_start_loc[:-1] + is_valid_seqs = jnp.where(all_seqs < end_seq, True, False) + has_initial_state = (seq_lens - query_lens) > 0 + all_valid_seqs = jnp.where(is_valid_seqs, all_seqs, 0) + + return memory_ref.MetadataRef.create( + cfgs=cfg, + num_tiles=pl.cdiv(end_seq, cfg.tile_size), + p_id_to_s_idx=all_valid_seqs, + p_id_to_r_base=all_valid_seqs, + p_id_to_r_size=jnp.where(is_valid_seqs, 1, 0), + p_id_is_first_tile=is_valid_seqs, + p_id_is_last_tile=is_valid_seqs, + s_idx_has_initial_state=has_initial_state, + s_idx_to_state_indices=state_indices, + ) + + +def compute_per_seq_metadata( + cfg: config.GDNConfig, + seq_lens: jax.Array, + query_start_loc: jax.Array, + state_indices: jax.Array, + start_seq: jax.Array, + end_seq: jax.Array, +) -> memory_ref.MetadataRef: + """Metadata for computing single sequence per tile.""" + + max_seqs = seq_lens.size + max_tokens = cfg.batch_size + all_seqs = jnp.arange(max_seqs) + all_tokens = jnp.arange(max_tokens) + + # Shift to ensure first element is for start_seq. + query_start_loc = jnp.roll(query_start_loc, shift=-start_seq) + seq_lens = jnp.roll(seq_lens, shift=-start_seq) + state_indices = jnp.roll(state_indices, shift=-start_seq) + + query_lens = query_start_loc[1:] - query_start_loc[:-1] + # NOTE: query_lens is used for calculating num_tiles. Defensive programming + # that masks out all the other values (seq_lens, state_indices) are not needed + # since they will not be visited as long as num_tiles is correct. + num_seqs = end_seq - start_seq + query_lens = jnp.where(all_seqs < num_seqs, query_lens, 0) + + # Calculate number of tiles needed for each sequence. + s_idx_to_num_tiles = pl.cdiv(query_lens, cfg.chunk_size) + # Calculate starting p_id of each sequence. + s_idx_to_start_p_id = jnp.cumulative_sum( + s_idx_to_num_tiles, include_initial=True + ) + # Map tile index to seq index. + # Consider following case: + # all_seqs = [0 1 2 3 4] + # s_idx_to_num_tiles = [1 2 3 0 1] + # jnp.repeat will return following results: + # p_id_to_s_idx = [0 1 1 2 2 2 4] + # This means p_id_to_s_idx[i] will point to its corresponding seq index. + + # NOTE: To make jnp.repeat jit compilable, we add total_repeat_length. This + # introduces padding to p_id_to_s_idx[i] where i >= num_tiles. Since the + # kernel only checks value up-to p_id_to_s_idx[num_tiles-1], padded value + # will not impact kernel execution. + p_id_to_s_idx = jnp.repeat( + all_seqs, s_idx_to_num_tiles, total_repeat_length=max_tokens + ) + # Map program id (p_id) to tile id of a sequence. + p_id_to_t_id = all_tokens - s_idx_to_start_p_id[p_id_to_s_idx] + # Map tile index to starting row of its activation. + p_id_to_r_base = ( + query_start_loc[p_id_to_s_idx] + p_id_to_t_id * cfg.chunk_size + ) + # Calculate number of rows to calculate / fetch for each tile. + p_id_to_r_size = jnp.minimum( + query_start_loc[p_id_to_s_idx + 1] - p_id_to_r_base, + cfg.tile_size, + ) + + # Calculate predicate used for state DMA. State is read if program id (p_id) + # is the first tile of a sequence and the sequence had been computed before + # (chunked prefill, decode, etc). State is written if the program id is the + # last tile of a sequence. + has_initial_state = (seq_lens - query_lens) > 0 + p_id_is_first_tile = p_id_to_t_id == 0 + p_id_is_last_tile = p_id_to_t_id == (s_idx_to_num_tiles[p_id_to_s_idx] - 1) + + # NOTE: Since query_lens[i] = 0 where i >= num_seqs, s_idx_to_num_tiles[i] + # where i >= num_seqs will also be 0. Therefore, s_idx_to_num_tiles.sum() + # will contain number of tiles for valid sequence. + num_tiles = s_idx_to_num_tiles.sum() + + return memory_ref.MetadataRef.create( + cfgs=cfg, + num_tiles=num_tiles, + p_id_to_s_idx=p_id_to_s_idx, + p_id_to_r_base=p_id_to_r_base, + p_id_to_r_size=p_id_to_r_size, + p_id_is_first_tile=p_id_is_first_tile, + p_id_is_last_tile=p_id_is_last_tile, + s_idx_has_initial_state=has_initial_state, + s_idx_to_state_indices=state_indices, + ) diff --git a/src/maxtext/models/kernels/gdn/pallas_mosaic_tpu.py b/src/maxtext/models/kernels/gdn/pallas_mosaic_tpu.py new file mode 100644 index 0000000000..5388c73137 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/pallas_mosaic_tpu.py @@ -0,0 +1,107 @@ +# 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. +# ============================================================================== +"""Pallas Mosaic TPU kernel implementation for Causal Conv1D Gated Delta Rule.""" + +import dataclasses +from typing import Any, Optional, override + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp +from tokamax._src.ops import op +from tokamax._src.ops.causal_conv1d_gated_delta_rule import base + +try: + from maxtext.models.kernels.gdn import config + from maxtext.models.kernels.gdn import wrapper +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + from maxtext.src.maxtext.models.kernels.gdn import wrapper + except (ImportError, ModuleNotFoundError): + from . import config + from . import wrapper + +GDNConfig = config.GDNConfig + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class PallasMosaicTpuCausalConv1dGatedDeltaRule( + base.CausalConv1dGatedDeltaRule[GDNConfig] +): + """Wrapper for the tokamax Op API for Pallas Mosaic TPU kernel.""" + + def _fwd( + self, + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_state: jax.Array, + recurrent_state: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + query_start_loc: jax.Array, + state_indices: jax.Array, + distribution: jax.Array, + seq_lens: jax.Array, + *, + 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 = 4, + mixed_tile_size: int = 64, + config: GDNConfig | None = None, + return_residuals: bool = False, + ) -> tuple[tuple[tuple[jax.Array, jax.Array], jax.Array], None]: + del return_residuals, config + out_act, states, *_ = wrapper.fused_conv1d_gdn( + qkv=qkv, + b=b, + a=a, + conv_state=conv_state, + recurrent_state=recurrent_state, + conv_weight=conv_weight, + 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=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, + ) + return (states, out_act), None + + @override + def supported_on(self, device: jax.Device) -> bool: + try: + return device.platform == "tpu" and pltpu.get_tpu_info().generation >= 6 + except Exception: + return False diff --git a/src/maxtext/models/kernels/gdn/tiling.py b/src/maxtext/models/kernels/gdn/tiling.py new file mode 100644 index 0000000000..34523a3973 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/tiling.py @@ -0,0 +1,357 @@ +# 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. +# ============================================================================== + +"""Dynamic tiling heuristics and VMEM memory estimation for Fused Conv1D-GDN.""" + +from jax.experimental import pallas as pl +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + except (ImportError, ModuleNotFoundError): + from . import config + + +def align_to(x: int, alignment: int) -> int: + """Aligns an integer upward to the nearest multiple of alignment.""" + return pl.cdiv(x, alignment) * alignment + + +def get_vmem_estimate_bytes( + tile_b: int, + chunk_sz: int, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + kernel_size: int, + act_in_bytes: int, + act_out_bytes: int, + conv_state_bytes: int, + rec_state_bytes: int, + num_lanes: int, + conv_state_dim_size: int, + is_decode: bool = False, +) -> int: + """Estimates total on-chip VMEM footprint in bytes for a GDN tile.""" + aligned_num_v_heads = align_to(n_v, num_lanes) + aligned_d_k = align_to(d_k, num_lanes) + aligned_d_v = align_to(d_v, num_lanes) + dim_size = align_to(2 * n_kq * d_k + n_v * d_v, num_lanes) + aligned_out_dim = align_to(n_v * d_v, num_lanes) + + # 1. Double-buffered input activation buffers (QKV, A, B). + qkv_bytes = 2 * (tile_b * chunk_sz * dim_size * act_in_bytes) + b_bytes = 2 * (tile_b * chunk_sz * aligned_num_v_heads * act_in_bytes) + a_bytes = 2 * (tile_b * chunk_sz * aligned_num_v_heads * act_in_bytes) + + # 2. Double-buffered state cache buffers (convolution and recurrent states). + conv_state_buffer_bytes = 2 * ( + tile_b * max(0, kernel_size - 1) * conv_state_dim_size * conv_state_bytes + ) + recurrent_state_buffer_bytes = 2 * ( + tile_b * n_v * d_v * d_k * rec_state_bytes + ) + + # 3. Double-buffered output activation buffer. + out_bytes = 2 * (tile_b * chunk_sz * aligned_out_dim * act_out_bytes) + + # 4. Temporary scratch buffers (allocated in non-batched mode). + if is_decode: + scratch_conv_bytes = 0 + scratch_recurrent_bytes = 0 + else: + scratch_conv_bytes = ( + tile_b * max(0, kernel_size - 1) * dim_size * conv_state_bytes + ) + scratch_recurrent_bytes = tile_b * n_v * d_v * d_k * rec_state_bytes + + # 5. Static weight cache references in on-chip memory. + weights_bytes = ( + ((kernel_size - 1) * dim_size * 4) + + (dim_size * 4) + + (aligned_num_v_heads * 8) + ) + + # 6. Working memory for intra-chunk recurrence and projections. + intermediate_bytes = ( + n_v + * (5 * chunk_sz * chunk_sz + 3 * chunk_sz * (aligned_d_v + aligned_d_k)) + * 4 + ) + + return ( + qkv_bytes + + b_bytes + + a_bytes + + conv_state_buffer_bytes + + recurrent_state_buffer_bytes + + out_bytes + + scratch_conv_bytes + + scratch_recurrent_bytes + + weights_bytes + + intermediate_bytes + ) + + +def calculate_decode_tile_size( + batch_size: int, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + conv_state_dim_size: int, + act_in_dtype: jnp.dtype, + act_out_dtype: jnp.dtype, + conv_state_dtype: jnp.dtype, + recurrent_state_dtype: jnp.dtype, + num_lanes: int, + vmem_capacity_limit_bytes: int, + kernel_size: int = 4, +) -> int: + """Derives optimal batch tile size for decode execution. + + Searches candidate batch tile sizes within maximum VMEM capacity limits. + + Args: + batch_size: Total batch size of the active decode sequence. + n_kq: Number of key/query heads. + n_v: Number of value heads. + d_k: Key head dimension. + d_v: Value head dimension. + conv_state_dim_size: Feature dimension size for conv state. + act_in_dtype: Data type for input activations. + act_out_dtype: Data type for output activations. + conv_state_dtype: Data type for conv state cache. + recurrent_state_dtype: Data type for recurrent state matrix. + num_lanes: Number of lanes for TPU vector layout alignment. + vmem_capacity_limit_bytes: Maximum allowed VMEM capacity in bytes. + kernel_size: 1D convolution kernel window size. + + Returns: + Derived batch tile size fitting within VMEM capacity limits. + """ + # Return a minimum valid tile size of 1 for empty or zero-length batches. + if batch_size <= 0: + return 1 + + act_in_bytes = jnp.dtype(act_in_dtype).itemsize + act_out_bytes = jnp.dtype(act_out_dtype).itemsize + conv_state_bytes = jnp.dtype(conv_state_dtype).itemsize + rec_state_bytes = jnp.dtype(recurrent_state_dtype).itemsize + + # Balance vector compute density against on-chip VMEM capacity: + # - Cap tile size across batch size tiers to maximize vector lane compute + # density. + # - Floor at tile_b = 4 for small batches to ensure compute density. + # - When value head count is large (n_v >= 64), recurrent state working + # memory scales up, so cap max_decode_b to 4 to prevent on-chip memory + # overflow. + if n_v >= 64: + max_decode_b = 2 + elif batch_size <= 64: + max_decode_b = 4 + elif batch_size <= 128: + max_decode_b = 8 + elif batch_size <= 256: + max_decode_b = 16 + else: + max_decode_b = 32 + + decode_candidates = [ + c for c in (32, 16, 8, 4, 2, 1) if c <= batch_size and c <= max_decode_b + ] + decode_tile_size = decode_candidates[-1] + + for cand in decode_candidates: + vmem_est = get_vmem_estimate_bytes( + tile_b=cand, + chunk_sz=1, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + act_in_bytes=act_in_bytes, + act_out_bytes=act_out_bytes, + conv_state_bytes=conv_state_bytes, + rec_state_bytes=rec_state_bytes, + num_lanes=num_lanes, + conv_state_dim_size=conv_state_dim_size, + is_decode=True, + ) + if vmem_est <= vmem_capacity_limit_bytes: + return cand + + return decode_tile_size + + +def calculate_mixed_tile_size( + seq_len: int, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + conv_state_dim_size: int, + act_in_dtype: jnp.dtype, + act_out_dtype: jnp.dtype, + conv_state_dtype: jnp.dtype, + recurrent_state_dtype: jnp.dtype, + num_lanes: int, + vmem_capacity_limit_bytes: int, + kernel_size: int = 4, +) -> int: + """Derives optimal chunk tile size for prefill and mixed execution. + + Searches candidate tile sizes within maximum VMEM capacity limits. + + Args: + seq_len: Sequence length of the active prefill or mixed sequence. + n_kq: Number of key/query heads. + n_v: Number of value heads. + d_k: Key head dimension. + d_v: Value head dimension. + conv_state_dim_size: Feature dimension size for conv state. + act_in_dtype: Data type for input activations. + act_out_dtype: Data type for output activations. + conv_state_dtype: Data type for conv state cache. + recurrent_state_dtype: Data type for recurrent state matrix. + num_lanes: Number of lanes for TPU vector layout alignment. + vmem_capacity_limit_bytes: Maximum allowed VMEM capacity in bytes. + kernel_size: 1D convolution kernel window size. + + Returns: + Derived chunk tile size fitting within VMEM capacity limits. + """ + # Return a minimum valid chunk size of 1 for empty or zero-length sequences. + if seq_len <= 0: + return 1 + + act_in_bytes = jnp.dtype(act_in_dtype).itemsize + act_out_bytes = jnp.dtype(act_out_dtype).itemsize + conv_state_bytes = jnp.dtype(conv_state_dtype).itemsize + rec_state_bytes = jnp.dtype(recurrent_state_dtype).itemsize + + # Limit chunk size to C <= 128: above 128, intra-chunk triangular + # solve operations and vector register pressure outweigh systolic compute + # density gains. + # When value head count is large (n_v >= 64), intra-chunk intermediate + # memory scales up, so cap chunk search space to C <= 64 to avoid on-chip + # memory overflow. + max_chunk_cap = 64 if n_v >= 64 else 128 + prefill_candidates = [ + c + for c in (128, 64, 32, 16, 8, 4, 2, 1) + if c <= seq_len and c <= max_chunk_cap + ] + mixed_tile_size = prefill_candidates[-1] + for candidate in prefill_candidates: + vmem_est = get_vmem_estimate_bytes( + tile_b=1, + chunk_sz=candidate, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + act_in_bytes=act_in_bytes, + act_out_bytes=act_out_bytes, + conv_state_bytes=conv_state_bytes, + rec_state_bytes=rec_state_bytes, + num_lanes=num_lanes, + conv_state_dim_size=conv_state_dim_size, + is_decode=False, + ) + if vmem_est <= vmem_capacity_limit_bytes: + return candidate + + return mixed_tile_size + + +def get_tile_sizes( + batch_size: int, + num_seqs: int, + padded_batch_size: int, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + kernel_size: int, + conv_state_dim_size: int, + act_in_dtype: jnp.dtype, + act_out_dtype: jnp.dtype, + conv_state_dtype: jnp.dtype, + recurrent_state_dtype: jnp.dtype, + num_lanes: int, + decode_tile_size: int | None = None, + mixed_tile_size: int | None = None, +) -> tuple[int, int]: + """Derives optimal decode and mixed tile sizes fitting within VMEM limits.""" + vmem_capacity_limit_bytes = config.get_vmem_limit_bytes() + + if decode_tile_size is None or decode_tile_size <= 0: + decode_tile_size = calculate_decode_tile_size( + batch_size=padded_batch_size, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + conv_state_dim_size=conv_state_dim_size, + act_in_dtype=act_in_dtype, + act_out_dtype=act_out_dtype, + conv_state_dtype=conv_state_dtype, + recurrent_state_dtype=recurrent_state_dtype, + num_lanes=num_lanes, + vmem_capacity_limit_bytes=vmem_capacity_limit_bytes, + kernel_size=kernel_size, + ) + + if mixed_tile_size is None or mixed_tile_size <= 0: + if batch_size <= num_seqs: + # When all sequences have length 1 (decode), size prefill chunks to 1. + effective_prefill_seq_len = 1 + else: + # Estimate maximum sequence length across uniform and mixed prefill + # batches. In an adversarial mixed batch of B tokens with N_seq sequences, + # the largest prefill sequence length is bounded by + # B - (N_seq - 1) = B - N_seq + 1. + effective_prefill_seq_len = max( + batch_size // max(1, num_seqs), + batch_size - num_seqs + 1, + ) + + mixed_tile_size = calculate_mixed_tile_size( + seq_len=effective_prefill_seq_len, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + conv_state_dim_size=conv_state_dim_size, + act_in_dtype=act_in_dtype, + act_out_dtype=act_out_dtype, + conv_state_dtype=conv_state_dtype, + recurrent_state_dtype=recurrent_state_dtype, + num_lanes=num_lanes, + vmem_capacity_limit_bytes=vmem_capacity_limit_bytes, + kernel_size=kernel_size, + ) + + # Guarantee strictly positive tile sizes (>= 1) for Pallas grid compilation. + decode_tile_size = max(1, min(decode_tile_size, batch_size)) + mixed_tile_size = max(1, min(mixed_tile_size, batch_size)) + return decode_tile_size, mixed_tile_size diff --git a/src/maxtext/models/kernels/gdn/vmem_ldst.py b/src/maxtext/models/kernels/gdn/vmem_ldst.py new file mode 100644 index 0000000000..73a2fe9c06 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/vmem_ldst.py @@ -0,0 +1,271 @@ +# 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. +# ============================================================================== + +"""VMEM load/store pre-processing logic.""" + +import jax +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config + from maxtext.models.kernels.gdn import memory_ref +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + from maxtext.src.maxtext.models.kernels.gdn import memory_ref + except (ImportError, ModuleNotFoundError): + from . import config + from . import memory_ref + + +def load_as_qkv_large( + qkv_vmem_ref: jax.Ref, cfgs: config.GDNConfig +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Split qkv and transpose by performing 1 load per chunk for large layout. + + Args: + qkv_vmem_ref: qkv reference in VMEM containing concatenated values of q, k, + and v of shape [seq_tile_size, chunk_size, 1, num_kq_heads * kq_head_dim * + 2 + num_v_heads * v_head_dim]. + cfgs: GDN configuration object. + + Returns: + q, k: [seq_tile_size, num_kq_heads, chunk_size, kq_head_dim] + v: [seq_tile_size, num_v_heads, chunk_size, v_head_dim] + """ + + num_lanes = pltpu.get_tpu_info().num_lanes + lanes_per_col = qkv_vmem_ref.shape[-1] // num_lanes + kq_lanes_per_head = cfgs.kq_head_dim // num_lanes + k_offset = cfgs.num_kq_heads * kq_lanes_per_head + + q_large_list = [] + k_large_list = [] + v_large_list = [] + + qkv_slot_flat_ref = qkv_vmem_ref.reshape(-1, num_lanes) # pyrefly: ignore[missing-attribute] + for kq_head in range(cfgs.num_kq_heads): + q_head_list = [] + k_head_list = [] + for lane in range(kq_lanes_per_head): + q_lane = kq_head * kq_lanes_per_head + lane + k_lane = k_offset + q_lane + + q_head_list.append(qkv_slot_flat_ref[q_lane::lanes_per_col]) + k_head_list.append(qkv_slot_flat_ref[k_lane::lanes_per_col]) + q_large_list.append(jnp.concat(q_head_list, axis=-1)) + k_large_list.append(jnp.concat(k_head_list, axis=-1)) + v_offset = kq_lanes_per_head * cfgs.num_kq_heads * 2 + v_lanes_per_head = cfgs.v_head_dim // num_lanes + for v_head in range(cfgs.num_v_heads): + v_head_list = [] + for lane in range(v_lanes_per_head): + v_lane = v_offset + v_head * v_lanes_per_head + lane + v_head_list.append(qkv_slot_flat_ref[v_lane::lanes_per_col]) + v_large_list.append(jnp.concat(v_head_list, axis=-1)) + + q_large = jnp.stack(q_large_list, axis=0) + k_large = jnp.stack(k_large_list, axis=0) + v_large = jnp.stack(v_large_list, axis=0) + + return q_large, k_large, v_large + + +def load_as_qkv_compact( + qkv_vmem_ref: jax.Ref, cfg: config.GDNConfig +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Split qkv and transpose by performing 1 load per head for compact layout. + + Args: + qkv_vmem_ref: qkv reference in VMEM containing concatenated values of q, k, + and v of shape [seq_tile_size, chunk_size, 1, num_kq_heads * kq_head_dim * + 2 + num_v_heads * v_head_dim]. + cfg: GDN configuration object. + + Returns: + q, k: [seq_tile_size, num_kq_heads, chunk_size, 1, kq_head_dim] + v: [seq_tile_size, num_v_heads, chunk_size, 1, v_head_dim] + """ + + k_offset = cfg.num_kq_heads * cfg.kq_head_dim + v_offset = cfg.num_kq_heads * 2 * cfg.kq_head_dim + + q_compact_list = [] + k_compact_list = [] + v_compact_list = [] + + for kq_head in range(cfg.num_kq_heads): + q_start = kq_head * cfg.kq_head_dim + q_end = q_start + cfg.kq_head_dim + k_start = k_offset + q_start + k_end = k_start + cfg.kq_head_dim + q_compact_list.append(qkv_vmem_ref[..., q_start:q_end]) + k_compact_list.append(qkv_vmem_ref[..., k_start:k_end]) + for v_head in range(cfg.num_v_heads): + v_start = v_offset + v_head * cfg.v_head_dim + v_end = v_start + cfg.v_head_dim + v_compact_list.append(qkv_vmem_ref[..., v_start:v_end]) + + q_compact = jnp.stack(q_compact_list, axis=1) + k_compact = jnp.stack(k_compact_list, axis=1) + v_compact = jnp.stack(v_compact_list, axis=1) + + return q_compact, k_compact, v_compact + + +def load_compact_to_large(vmem_ref: jax.Ref) -> jax.Array: + """Use strided load to convert compact to large layout without transpose.""" + + # NOTE: Only support 32-bits for now. + assert vmem_ref.dtype.itemsize == 4 + assert vmem_ref.shape[-2] == 1 + col_size = vmem_ref.shape[-1] + new_shape = vmem_ref.shape[:-2] + (col_size,) + tpu_info = pltpu.get_tpu_info() + num_lanes = tpu_info.num_lanes + + vreg_list = [] + vmem_ref = vmem_ref.reshape(-1, col_size) # pyrefly: ignore[missing-attribute] + for col_start in range(0, col_size, num_lanes): + col_end = min(col_start + num_lanes, col_size) + vreg = vmem_ref[..., col_start:col_end] + vreg_list.append(vreg) + return jnp.concat(vreg_list, axis=-1).reshape(new_shape) + + +def load_and_select_states( + metadata_ref: memory_ref.MetadataRef, + p_id: jax.Array, + conv_state_slot_ref: jax.Ref, + recurrent_slot_ref: jax.Ref, + carry_conv_scratch_ref: jax.Ref | None, + carry_recurrent_scratch_ref: jax.Ref | None, + cfg: config.GDNConfig, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Load correct states from HBM or prior tile, and masks invalid states. + + Reference metadata to select the appropriate prior states. If `is_first_tile` + is True, it selects states read from HBM. If it is False, it selects + carry states from previous tile. If `has_initial_state` is False, states are + zero initialized. + + Args: + metadata_ref: Metadata reference containing grid and sequence mappings. + p_id: Current Pallas program ID. + conv_state_slot_ref: Convolution state read from HBM of shape + [seq_tile_size, prev_kernel_size, 1, dim_size]. + recurrent_slot_ref: Recurrent state read from HBM of shape [seq_tile_size, + num_v_heads, kq_head_dim, v_head_dim]. + carry_conv_scratch_ref: Optional inter-tile convolution carry of shape + [seq_tile_size, prev_kernel_size, 1, dim_size]. + carry_recurrent_scratch_ref: Optional inter-tile recurrent state carry of + shape [seq_tile_size, num_v_heads, kq_head_dim, v_head_dim]. + cfg: GDN configuration object. + + Returns: + real_sizes: Valid token count per sequence tile of shape [seq_tile_size]. + prev_conv_state: Selected convolution state of shape [seq_tile_size, + prev_kernel_size, 1, dim_size] in float32. + prev_recurrent_state: Selected recurrent state of shape [seq_tile_size, + num_v_heads, kq_head_dim, v_head_dim]. + """ + + real_sizes_list = [] + prev_conv_state_list = [] + prev_recurrent_state_list = [] + + for idx in range(cfg.seq_tile_size): + record = metadata_ref.get_record(p_id, idx) + s_idx = record.s_idx + real_sizes = record.r_size + is_first_tile = record.is_first_tile + has_initial_state = metadata_ref.s_idx_has_initial_state[s_idx] + + # NOTE: Conv1D mandates fp32 due to its usage of compact layout. + hbm_conv_state = conv_state_slot_ref[idx].astype(jnp.float32) + prev_conv_state = jnp.where(has_initial_state, hbm_conv_state, 0) + + if carry_conv_scratch_ref is not None: + prev_tile_conv = carry_conv_scratch_ref[idx] + prev_conv_state = jnp.where( + is_first_tile, prev_conv_state, prev_tile_conv + ) + + hbm_recurrent_state = recurrent_slot_ref[idx] + prev_recurrent_state = jnp.where(has_initial_state, hbm_recurrent_state, 0) + + if carry_recurrent_scratch_ref is not None: + prev_tile_recurrent_scratch = carry_recurrent_scratch_ref[idx] + prev_recurrent_state = jnp.where( + is_first_tile, prev_recurrent_state, prev_tile_recurrent_scratch + ) + + real_sizes_list.append(real_sizes) + prev_conv_state_list.append(prev_conv_state) + prev_recurrent_state_list.append(prev_recurrent_state) + + real_sizes = jnp.stack(real_sizes_list, axis=0) + prev_conv_state = jnp.stack(prev_conv_state_list, axis=0) + prev_recurrent_state = jnp.stack(prev_recurrent_state_list, axis=0) + + return real_sizes, prev_conv_state, prev_recurrent_state + + +def load_activation_as_compact( + qkv_vreg: jax.Array, + qkv_vmem_ref: jax.Ref, + b_vmem_ref: jax.Ref, + a_vmem_ref: jax.Ref, + cfgs: config.GDNConfig, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Load activations from VMEM as a compact layout.""" + + qkv_vmem_ref[...] = qkv_vreg + q_compact, k_compact, v_compact = load_as_qkv_compact(qkv_vmem_ref, cfgs) + b_compact = jnp.expand_dims(b_vmem_ref[...], axis=1) + a_compact = jnp.expand_dims(a_vmem_ref[...], axis=1) + return q_compact, k_compact, v_compact, b_compact, a_compact + + +def load_activation_as_large( + qkv_vreg: jax.Array, + qkv_vmem_ref: jax.Ref, + b_vmem_ref: jax.Ref, + a_vmem_ref: jax.Ref, + cfgs: config.GDNConfig, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Load activations from VMEM as a large layout.""" + + qkv_vmem_ref[...] = qkv_vreg + + q_large_list = [] + k_large_list = [] + v_large_list = [] + for idx in range(cfgs.seq_tile_size): + q_large, k_large, v_large = load_as_qkv_large(qkv_vmem_ref.at[idx], cfgs) + q_large_list.append(q_large) + k_large_list.append(k_large) + v_large_list.append(v_large) + + q_large = jnp.stack(q_large_list, axis=0) + k_large = jnp.stack(k_large_list, axis=0) + v_large = jnp.stack(v_large_list, axis=0) + b_large = load_compact_to_large(b_vmem_ref) + a_large = load_compact_to_large(a_vmem_ref) + b_large = jnp.expand_dims(b_large, axis=1) + a_large = jnp.expand_dims(a_large, axis=1) + + return q_large, k_large, v_large, b_large, a_large diff --git a/src/maxtext/models/kernels/gdn/wrapper.py b/src/maxtext/models/kernels/gdn/wrapper.py new file mode 100644 index 0000000000..e44b474eb3 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/wrapper.py @@ -0,0 +1,550 @@ +# 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. +# ============================================================================== + +"""Top-level Pallas kernel wrapper for fused Conv1D-GDN with triangular inverse caching.""" + +import functools + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import compute_conv1d + from maxtext.models.kernels.gdn import compute_gdn + from maxtext.models.kernels.gdn import config + from maxtext.models.kernels.gdn import memory_ref + from maxtext.models.kernels.gdn import metadata + from maxtext.models.kernels.gdn import tiling + from maxtext.models.kernels.gdn import vmem_ldst +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import compute_conv1d + from maxtext.src.maxtext.models.kernels.gdn import compute_gdn + from maxtext.src.maxtext.models.kernels.gdn import config + from maxtext.src.maxtext.models.kernels.gdn import memory_ref + from maxtext.src.maxtext.models.kernels.gdn import metadata + from maxtext.src.maxtext.models.kernels.gdn import tiling + from maxtext.src.maxtext.models.kernels.gdn import vmem_ldst + except (ImportError, ModuleNotFoundError): + from . import compute_conv1d + from . import compute_gdn + from . import config + from . import memory_ref + from . import metadata + from . import tiling + from . import vmem_ldst + + +def inner_kernel( + # Inputs. + qkv_slot_ref: jax.Ref, # [seq, chunk, 1, dim_size] + 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] + *args, + cfg: config.GDNConfig, + **kwargs, +) -> None: + """Orchestrates computation of Conv1D and GDN for a single tile. + + This kernel acts as a facade adhering to strict separation of concerns. It + operates VMEM reference without knowledge on DMA logic. Furthermore, the + 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 + else: + 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 + + p_id = pl.program_id(0) + + # Prepare states. + real_sizes, prev_conv, prev_recurrent = vmem_ldst.load_and_select_states( + metadata_ref=metadata_ref, + p_id=p_id, + conv_state_slot_ref=conv_state_slot_ref, + recurrent_slot_ref=recurrent_slot_ref, + carry_conv_scratch_ref=carry_conv_scratch_ref, + carry_recurrent_scratch_ref=carry_recurrent_scratch_ref, + cfg=cfg, + ) + + # Step 1: Conv1D. + qkv_in_compact = qkv_slot_ref[...].astype(jnp.float32) + qkv_in_compact = jnp.concat([prev_conv, qkv_in_compact], axis=1) + + # Prepare conv1d weights. + conv_weight = weights_ref.conv.weight[...].astype(jnp.float32) + conv_bias = None + if weights_ref.conv.bias is not None: + conv_bias = weights_ref.conv.bias[...].astype(jnp.float32) + + qkv_out_compact, new_conv_state = compute_conv1d.causal_conv1d( + real_sizes=real_sizes, + lhs=qkv_in_compact, + conv_weight=conv_weight, + conv_bias=conv_bias, + cfg=cfg, + ) + + conv_state_slot_ref[...] = new_conv_state + if carry_conv_scratch_ref is not None: + carry_conv_scratch_ref[...] = new_conv_state + + # Apply activation function. + qkv_out_compact = jax.nn.silu(qkv_out_compact) + + # Step 2: GDN. + padding_size = cfg.aligned_num_v_heads - cfg.num_v_heads + a_log = jnp.pad(weights_ref.gdn.a_log[...], ((0, padding_size))) + dt_bias = jnp.pad(weights_ref.gdn.dt_bias[...], ((0, padding_size))) + + if cfg.chunk_size == 1: + q_compact, k_compact, v_compact, b_compact, a_compact = ( + vmem_ldst.load_activation_as_compact( + qkv_vreg=qkv_out_compact, + qkv_vmem_ref=qkv_slot_ref, + b_vmem_ref=b_slot_ref, + a_vmem_ref=a_slot_ref, + cfgs=cfg, + ) + ) + + out, new_recurrent_state, t_inv = compute_gdn.recurrent_gdn( + q_compact=q_compact, + k_compact=k_compact, + v_compact=v_compact, + b_compact=b_compact, + a_compact=a_compact, + state_prev=prev_recurrent, + a_log=a_log, + dt_bias=dt_bias, + cfg=cfg, + real_sizes=real_sizes, + ) + else: + q_large, k_large, v_large, b_large, a_large = ( + vmem_ldst.load_activation_as_large( + qkv_vreg=qkv_out_compact, + qkv_vmem_ref=qkv_slot_ref, + b_vmem_ref=b_slot_ref, + a_vmem_ref=a_slot_ref, + cfgs=cfg, + ) + ) + + out, new_recurrent_state, t_inv = compute_gdn.chunked_gdn( + q_large=q_large, + k_large=k_large, + v_large=v_large, + b_large=b_large, + a_large=a_large, + state_prev=prev_recurrent, + a_log=a_log, + dt_bias=dt_bias, + cfg=cfg, + real_sizes=real_sizes, + ) + + # Store output, recurrent, and t_inv to vmem. + 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: + chunk_states_slot_ref[...] = prev_recurrent.astype( + chunk_states_slot_ref.dtype + ) + + if carry_recurrent_scratch_ref is not None: + carry_recurrent_scratch_ref[...] = new_recurrent_state + + +def outer_kernel( + # Inputs. + metadata_ref: memory_ref.MetadataRef, + qkv_ref: jax.Array, + b_ref: jax.Array, + a_ref: jax.Array, + conv_state_ref: jax.Array, + 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, + *args, + carry_conv_scratch_ref: jax.Array | None = None, + carry_recurrent_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 + + chunk_states_ref = ( + args[0] + if (len(args) > 0 and cfg.mode == config.GDNMode.PER_SEQ) + else None + ) + + allocs = memory_ref.create_allocs( + metadata_ref=metadata_ref, + qkv_ref=qkv_ref, + b_ref=b_ref, + a_ref=a_ref, + out_ref=out_ref, + conv_state_ref=conv_state_ref, + recurrent_state_ref=recurrent_state_ref, + cfg=cfg, + t_inv_ref=t_inv_ref, + chunk_states_ref=chunk_states_ref, + ) + qkv_alloc = allocs[0] + b_alloc = allocs[1] + 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 + + 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, + cfg=cfg, + ), + grid=(num_tiles,), + in_specs=( + qkv_alloc.spec, + b_alloc.spec, + a_alloc.spec, + conv_alloc.spec, + recurrent_alloc.spec, + ), + out_specs=tuple(out_specs), + ) + + @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) + pipeline_func( + qkv_ref, + b_ref, + a_ref, + conv_state_ref, + recurrent_state_ref, + *out_args, + scratches=( + metadata_ref, + weights_ref, + carry_conv_scratch_ref, + carry_recurrent_scratch_ref, + ), + allocations=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", + ), +) +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, +) -> 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).""" + 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) + + # Step 1: Validate inputs. + num_seqs = state_indices.size + batch_size, dim = qkv.shape + assert conv_weight.shape == (dim, 1, kernel_size) + if conv_bias is not None: + assert conv_bias.shape == (dim,) + assert query_start_loc.shape == (num_seqs + 1,) + assert state_indices.shape == (num_seqs,) + assert distribution.shape == (3,) + + num_lanes = pltpu.get_tpu_info().num_lanes + packing = 4 // act_in_dtype.itemsize + padded_batch_size = pl.cdiv(batch_size, packing) * packing + conv_state_dim_size = conv_state.shape[-1] + + decode_tile_size, mixed_tile_size = tiling.get_tile_sizes( + batch_size=batch_size, + num_seqs=num_seqs, + padded_batch_size=padded_batch_size, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + conv_state_dim_size=conv_state_dim_size, + act_in_dtype=act_in_dtype, + act_out_dtype=act_out_dtype, + conv_state_dtype=conv_state.dtype, + recurrent_state_dtype=recurrent_state.dtype, + num_lanes=num_lanes, + decode_tile_size=decode_tile_size, + mixed_tile_size=mixed_tile_size, + ) + + batch_padding_size = padded_batch_size - batch_size + aligned_num_v_heads = tiling.align_to(n_v, 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))) + + qkv = qkv.reshape(padded_batch_size, 1, -1) + b = b.reshape(padded_batch_size, 1, -1) + a = a.reshape(padded_batch_size, 1, -1) + + # Step 3: States and weights pre-processing. + conv_state_shape = conv_state.shape + conv_state = conv_state.reshape(-1, kernel_size - 1, 1, dim) + conv_weight = conv_weight.swapaxes(0, 2).astype(jnp.float32) + conv_bias = conv_bias.astype(jnp.float32) if conv_bias is not None else None + + # Step 4: Wrap inputs for the kernel. + conv_weights = memory_ref.ConvWeightsRef(weight=conv_weight, bias=conv_bias) + gdn_weights = memory_ref.GDNWeightsRef(a_log=a_log, dt_bias=dt_bias) + weights = memory_ref.WeightRefs(conv=conv_weights, gdn=gdn_weights) + + # Step 5: Create specs. + smem_spec = pl.BlockSpec(memory_space=pltpu.SMEM) + vmem_spec = pl.BlockSpec(memory_space=pltpu.VMEM) + hbm_spec = pl.BlockSpec(memory_space=pltpu.HBM) + weights_spec = jax.tree.map(lambda _: vmem_spec, weights) + + def call_kernel( + in_conv_state: jax.Array, + in_recurrent_state: jax.Array, + in_act: jax.Array | None, + mode: config.GDNMode, + ) -> tuple[jax.Array, ...]: + if mode == config.GDNMode.BATCHED: + tile_size = decode_tile_size + else: + tile_size = mixed_tile_size + + cfg = config.GDNConfig( + mode=mode, + batch_size=padded_batch_size, + kernel_size=kernel_size, + tile_size=tile_size, + dim_size=dim, + num_kq_heads=n_kq, + num_v_heads=n_v, + kq_head_dim=d_k, + v_head_dim=d_v, + dtypes=config.Dtypes( + act_in=act_in_dtype, + act_out=act_out_dtype, + compute=compute_precision, + recurrent_state=in_recurrent_state.dtype, + conv_state=in_conv_state.dtype, + ), + ) + + if mode == config.GDNMode.BATCHED: + metadata_obj = metadata.compute_batched_seq_metadata( + cfg=cfg, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + state_indices=state_indices, + end_seq=distribution[0], + ) + else: + metadata_obj = metadata.compute_per_seq_metadata( + cfg=cfg, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + state_indices=state_indices, + start_seq=distribution[0], + end_seq=distribution[-1], + ) + + metadata_spec = jax.tree.map(lambda _: smem_spec, metadata_obj) + + in_out_spec = None + 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: + in_act = jnp.zeros_like(out_shape) + if in_act is not None: + out_shape = in_act + in_out_spec = hbm_spec + input_output_aliases[len(metadata_obj) + 5] = 0 + + num_chunks = cfg.batch_size // cfg.chunk_size + t_inv_shape = jax.ShapeDtypeStruct( + (num_chunks, cfg.num_v_heads, cfg.chunk_size, cfg.chunk_size), + cfg.dtypes.compute, + ) + + if 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, + ) + out_shape_tuple = ( + out_shape, + in_conv_state, + in_recurrent_state, + t_inv_shape, + chunk_states_shape, + ) + out_specs_tuple = (hbm_spec, hbm_spec, hbm_spec, hbm_spec, hbm_spec) + else: + out_shape_tuple = ( + out_shape, + in_conv_state, + in_recurrent_state, + t_inv_shape, + ) + out_specs_tuple = (hbm_spec, hbm_spec, hbm_spec, hbm_spec) + + return pl.pallas_call( + functools.partial(outer_kernel, cfg=cfg), + out_shape=out_shape_tuple, + in_specs=( + metadata_spec, + hbm_spec, + hbm_spec, + hbm_spec, + hbm_spec, + hbm_spec, + in_out_spec, + weights_spec, + ), + out_specs=out_specs_tuple, + scratch_shapes=cfg.get_scratch_shape_dict(), + input_output_aliases=input_output_aliases, + compiler_params=pltpu.CompilerParams( + disable_bounds_checks=True, + vmem_limit_bytes=config.get_vmem_limit_bytes(), + ), + name=cfg.get_kernel_name(), + metadata=cfg.get_metadata(), + )( + metadata_obj, + qkv, + b, + a, + in_conv_state, + in_recurrent_state, + in_act, + weights, + ) + + if not is_prefill_only: + try: + if int(distribution[0]) == 0: + is_prefill_only = True + except (TypeError, ValueError, jax.errors.TracerIntegerConversionError): + pass + + if not is_prefill_only: + out_act, out_conv_state, out_recurrent_state, _ = call_kernel( + conv_state, recurrent_state, None, config.GDNMode.BATCHED + ) + else: + out_act = None + 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 + ) + ) + + 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) + out_recurrent_state = out_recurrent_state.astype(recurrent_out_dtype) + + return out_act, (out_conv_state, out_recurrent_state), t_inv, chunk_states diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 1a04df4c7c..494460ed81 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -580,8 +580,8 @@ def a_log_init(key, shape, dtype=jnp.float32): a_vals = jax.random.uniform(key, shape=shape, dtype=dtype, minval=1e-9, maxval=16.0) return jnp.log(a_vals) - self.A_log = nnx.Param(a_log_init(rngs.params(), (self.num_v_heads,), dtype=cfg.weight_dtype)) - self.dt_bias = nnx.Param(nnx.initializers.ones(rngs.params(), (self.num_v_heads,), dtype=cfg.weight_dtype)) + self.A_log = nnx.Param(a_log_init(rngs.params(), (self.num_v_heads,), dtype=jnp.float32)) + self.dt_bias = nnx.Param(nnx.initializers.ones(rngs.params(), (self.num_v_heads,), dtype=jnp.float32)) self.norm = Qwen3NextRMSNormGated( num_features=self.head_v_dim, # Normalize over the head dimension (D_v) @@ -654,6 +654,12 @@ def __call__( # hidden_states: (B, S, E) cfg = self.config batch, seq_len, _ = hidden_states.shape + decay_dtype = getattr(cfg, "gdn_decay_dtype", jnp.float32) + if isinstance(decay_dtype, str): + decay_dtype = getattr(jnp, decay_dtype, jnp.float32) + state_dtype = getattr(cfg, "gdn_state_dtype", jnp.float32) + if isinstance(state_dtype, str): + state_dtype = getattr(jnp, state_dtype, jnp.float32) flat_sharding, head_sharding, state_sharding = ( self._explicit_activation_shardings(batch) ) @@ -847,8 +853,8 @@ def __call__( recurrent_state_paged, conv_weight, None, # conv_bias: MaxText conv1d uses use_bias=False. - jnp.asarray(self.A_log[...], dtype=cfg.dtype), - jnp.asarray(self.dt_bias[...], dtype=cfg.dtype), + jnp.asarray(self.A_log[...], dtype=decay_dtype), + jnp.asarray(self.dt_bias[...], dtype=decay_dtype), state_indices, query_start_loc, attention_metadata.request_distribution, # pyrefly: ignore[missing-attribute] @@ -925,181 +931,319 @@ def extract_state(c_in, v_len): else: conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) - # Perform the convolution. - conv_out = self.conv1d(conv_input, out_sharding=flat_sharding) - # Slice the output to match the original input sequence length. - conv_out = conv_out[:, -seq_len:, :] - qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(cfg.dtype) - # q_conv shape: (B, S, key_dim), k_conv shape: (B, S, key_dim), v_conv shape: (B, S, value_dim) - q_conv, k_conv, v_conv = jnp.split(qkv_conv, [self.key_dim, 2 * self.key_dim], axis=-1) - - # Reshape for multi-head processing - # query shape: (B, S, H_k, D_k) - query = jnp.reshape( - q_conv, - (batch, seq_len, self.num_k_heads, self.head_k_dim), - out_sharding=head_sharding, - ) - # key shape: (B, S, H_k, D_k) - key = jnp.reshape( - k_conv, - (batch, seq_len, self.num_k_heads, self.head_k_dim), - out_sharding=head_sharding, - ) - # value shape: (B, S, H_v, D_v) - value = jnp.reshape( - v_conv, - (batch, seq_len, self.num_v_heads, self.head_v_dim), - out_sharding=head_sharding, - ) - - # ========================================================================= - # STEP C: Gated Delta Rule Recurrence - # ========================================================================= - A_log = jnp.asarray(self.A_log[...], dtype=cfg.dtype) - dt_bias = jnp.asarray(self.dt_bias[...], dtype=cfg.dtype) - if cfg.shard_mode == ShardMode.EXPLICIT: - # Both are stored replicated but broadcast against (B, S, H_v) activations whose - # head axis is sharded, and explicit sharding requires broadcast operands to - # agree -- the same fix `_align_scale_with_normalized_axis` applies to the norm - # scales. - head_spec = jax.sharding.PartitionSpec(jax.typeof(a).sharding.spec[-1]) - A_log = jax.sharding.reshard(A_log, head_spec) - dt_bias = jax.sharding.reshard(dt_bias, head_spec) - # beta shape: (B, S, H_v) - beta = jax.nn.sigmoid(b) - # g shape: (B, S, H_v) - g = -jnp.exp(A_log) * jax.nn.softplus(a + dt_bias) - - if decoder_segment_ids is not None: - mask = decoder_segment_ids != 0 - # Apply mask by broadcasting to respective shapes - key = jnp.where(mask[..., None, None], key, 0.0) - value = jnp.where(mask[..., None, None], value, 0.0) - g = jnp.where(mask[..., None], g, 0.0) - - if self.num_v_heads > self.num_k_heads and self.num_v_heads % self.num_k_heads == 0: - repeats = self.num_v_heads // self.num_k_heads - # query shape after repeat: (B, S, H_v, D_k) - query = jnp.repeat(query, repeats, axis=2, out_sharding=head_sharding) - # key shape after repeat: (B, S, H_v, D_k) - key = jnp.repeat(key, repeats, axis=2, out_sharding=head_sharding) - - if seq_len == 1 and model_mode == MODEL_MODE_AUTOREGRESSIVE: - core_attn_out, next_recurrent_state = jax_ar_gated_delta_rule( - query, - key, - value, - g, - beta, - initial_state=recurrent_state, # pyrefly: ignore[bad-argument-type] - use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, - compute_dtype=cfg.dtype, + if getattr(cfg, "use_gdn_kernel", False): + try: + from maxtext.models.kernels.gdn.gdn_bwd_pallas import gdn_fused_conv1d + except ImportError: + try: + from maxtext.models.kernels.gdn.gdn_bwd_pallas import gdn_fused_conv1d + except ImportError: + try: + from .kernels.gdn.gdn_bwd_pallas import gdn_fused_conv1d + except ImportError: + from .kernels.gdn import gdn_fused_conv1d + + conv_state_arg = ( + conv_state + if conv_state is not None + else jnp.zeros( + (batch, self.config.gdn_conv_kernel_dim - 1, qkv.shape[-1]), + dtype=cfg.dtype, + ) ) - elif self.mesh is not None: - logical_rules = get_logical_axis_rules() recurrent_state_arg = ( - recurrent_state + recurrent_state.astype(state_dtype) if recurrent_state is not None else jnp.zeros( (batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), - dtype=cfg.dtype, - out_sharding=state_sharding, + dtype=state_dtype, ) ) - # LENGTH, not None. The sequence axis was hardcoded to replicated, so - # ici_context_parallelism could never shard the GDN sequence while still - # consuming the context axis from the mesh -- which is why raising ctx - # made memory worse instead of better. The scan handles a sharded - # sequence via the two-pass affine composition in kernels/attention/gdn_cp.py. - # Either context axis can carry the sequence. LENGTH maps to both in the - # logical rules, so this is correct whichever one is configured. - cp_axes = gdn_context_axes(cfg) - cp_len = LENGTH if cp_axes else None - qkv_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) - g_beta_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD), mesh=self.mesh, rules=logical_rules) - state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules) - # Keep every shard_map input/output batch spec consistent when replication is required. - qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec( - qkv_pspec, - query.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, + conv_bias_arg = ( + self.conv1d.bias.value + if getattr(self.conv1d, "bias", None) is not None + else None ) - g_beta_pspec = remove_incompatible_mesh_axes_from_partition_spec( - g_beta_pspec, - g.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, + + if self.mesh is not None: + logical_rules = self.config.logical_axis_rules + qkv_pspec = logical_to_mesh_axes( + (KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules + ) + b_a_pspec = logical_to_mesh_axes( + (KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules + ) + conv_state_pspec = logical_to_mesh_axes( + (KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules + ) + recurrent_state_pspec = logical_to_mesh_axes( + (KV_BATCH, None, None, None), mesh=self.mesh, rules=logical_rules + ) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=( + qkv_pspec, + b_a_pspec, + b_a_pspec, + P(), + P(), + P(), + P(), + conv_state_pspec, + recurrent_state_pspec, + ), + out_specs=( + qkv_pspec, + (conv_state_pspec, recurrent_state_pspec), + ), + check_vma=False, + ) + def shard_mapped_gdn( + qkv_val, + b_val, + a_val, + cw_val, + cb_val, + alog_val, + dt_val, + cs_val, + rs_val, + ): + return gdn_fused_conv1d( + qkv=qkv_val, + b=b_val, + a=a_val, + conv_weight=cw_val, + conv_bias=cb_val, + a_log=alog_val, + dt_bias=dt_val, + conv_state=cs_val, + recurrent_state=rs_val, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.config.gdn_conv_kernel_dim, + chunk_size=self.config.gdn_chunk_size, + use_qk_norm_in_gdn=self.config.use_qk_norm_in_gdn, + compute_dtype=state_dtype, + ) + + gdn_step_fn = shard_mapped_gdn + core_attn_out, (next_conv_state, next_recurrent_state) = ( + gdn_step_fn( + qkv, + b, + a, + self.conv1d.kernel.value, + conv_bias_arg, + self.A_log[...], + self.dt_bias[...], + conv_state_arg, + recurrent_state_arg, + ) + ) + else: + gdn_step_fn = gdn_fused_conv1d + core_attn_out, (next_conv_state, next_recurrent_state) = ( + gdn_step_fn( + qkv=qkv, + b=b, + a=a, + conv_weight=self.conv1d.kernel.value, + conv_bias=conv_bias_arg, + a_log=self.A_log[...], + dt_bias=self.dt_bias[...], + conv_state=conv_state_arg, + recurrent_state=recurrent_state_arg, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.config.gdn_conv_kernel_dim, + chunk_size=self.config.gdn_chunk_size, + use_qk_norm_in_gdn=self.config.use_qk_norm_in_gdn, + compute_dtype=state_dtype, + ) + ) + else: + # Perform the convolution. + conv_out = self.conv1d(conv_input, out_sharding=flat_sharding) + # Slice the output to match the original input sequence length. + conv_out = conv_out[:, -seq_len:, :] + qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(cfg.dtype) + # q_conv shape: (B, S, key_dim), k_conv shape: (B, S, key_dim), v_conv shape: (B, S, value_dim) + q_conv, k_conv, v_conv = jnp.split(qkv_conv, [self.key_dim, 2 * self.key_dim], axis=-1) + + # Reshape for multi-head processing + # query shape: (B, S, H_k, D_k) + query = jnp.reshape( + q_conv, + (batch, seq_len, self.num_k_heads, self.head_k_dim), + out_sharding=head_sharding, ) - state_pspec = remove_incompatible_mesh_axes_from_partition_spec( - state_pspec, - recurrent_state_arg.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, + # key shape: (B, S, H_k, D_k) + key = jnp.reshape( + k_conv, + (batch, seq_len, self.num_k_heads, self.head_k_dim), + out_sharding=head_sharding, + ) + # value shape: (B, S, H_v, D_v) + value = jnp.reshape( + v_conv, + (batch, seq_len, self.num_v_heads, self.head_v_dim), + out_sharding=head_sharding, ) + # ========================================================================= + # STEP C: Gated Delta Rule Recurrence + # ========================================================================= + A_log = jnp.asarray(self.A_log[...], dtype=decay_dtype) + dt_bias = jnp.asarray(self.dt_bias[...], dtype=decay_dtype) if cfg.shard_mode == ShardMode.EXPLICIT: - # shard_map manualises the mesh axes it is given and will not insert a reshard - # for an operand whose layout differs from `in_specs`, so hand it arrays that - # already match. - query = jax.sharding.reshard(query, qkv_pspec) - key = jax.sharding.reshard(key, qkv_pspec) - value = jax.sharding.reshard(value, qkv_pspec) - g = jax.sharding.reshard(g, g_beta_pspec) - beta = jax.sharding.reshard(beta, g_beta_pspec) - recurrent_state_arg = jax.sharding.reshard( - recurrent_state_arg, state_pspec + # Both are stored replicated but broadcast against (B, S, H_v) activations whose + # head axis is sharded, and explicit sharding requires broadcast operands to + # agree -- the same fix `_align_scale_with_normalized_axis` applies to the norm + # scales. + head_spec = jax.sharding.PartitionSpec(jax.typeof(a).sharding.spec[-1]) + A_log = jax.sharding.reshard(A_log, head_spec) + dt_bias = jax.sharding.reshard(dt_bias, head_spec) + # beta shape: (B, S, H_v) + beta = jax.nn.sigmoid(b) + # g shape: (B, S, H_v) + g = -jnp.exp(A_log) * jax.nn.softplus(a + dt_bias) + + if decoder_segment_ids is not None: + mask = decoder_segment_ids != 0 + # Apply mask by broadcasting to respective shapes + key = jnp.where(mask[..., None, None], key, 0.0) + value = jnp.where(mask[..., None, None], value, 0.0) + g = jnp.where(mask[..., None], g, 0.0) + + if self.num_v_heads > self.num_k_heads and self.num_v_heads % self.num_k_heads == 0: + repeats = self.num_v_heads // self.num_k_heads + # query shape after repeat: (B, S, H_v, D_k) + query = jnp.repeat(query, repeats, axis=2, out_sharding=head_sharding) + # key shape after repeat: (B, S, H_v, D_k) + key = jnp.repeat(key, repeats, axis=2, out_sharding=head_sharding) + + if seq_len == 1 and model_mode == MODEL_MODE_AUTOREGRESSIVE: + core_attn_out, next_recurrent_state = jax_ar_gated_delta_rule( + query, + key, + value, + g, + beta, + initial_state=recurrent_state, # pyrefly: ignore[bad-argument-type] + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + ) + elif self.mesh is not None: + logical_rules = get_logical_axis_rules() + recurrent_state_arg = ( + recurrent_state + if recurrent_state is not None + else jnp.zeros( + (batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), + dtype=state_dtype, + out_sharding=state_sharding, + ) + ) + # LENGTH, not None. The sequence axis was hardcoded to replicated, so + # ici_context_parallelism could never shard the GDN sequence while still + # consuming the context axis from the mesh -- which is why raising ctx + # made memory worse instead of better. The scan handles a sharded + # sequence via the two-pass affine composition in kernels/attention/gdn_cp.py. + # Either context axis can carry the sequence. LENGTH maps to both in the + # logical rules, so this is correct whichever one is configured. + cp_axes = gdn_context_axes(cfg) + cp_len = LENGTH if cp_axes else None + qkv_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) + g_beta_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD), mesh=self.mesh, rules=logical_rules) + state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules) + # Keep every shard_map input/output batch spec consistent when replication is required. + qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec( + qkv_pspec, + query.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, + ) + g_beta_pspec = remove_incompatible_mesh_axes_from_partition_spec( + g_beta_pspec, + g.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, + ) + state_pspec = remove_incompatible_mesh_axes_from_partition_spec( + state_pspec, + recurrent_state_arg.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, ) - @functools.partial( - jax.shard_map, - mesh=self.mesh, - in_specs=( - qkv_pspec, # query - qkv_pspec, # key - qkv_pspec, # value - g_beta_pspec, # g - g_beta_pspec, # beta - state_pspec, # initial_state - ), - out_specs=( - qkv_pspec, # core_attn_out - state_pspec, # final_state - ), - check_vma=False, - ) - def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): - return jax_chunk_gated_delta_rule( - query=q, - key=k, - value=v, - g=g_val, - beta=beta_val, + if cfg.shard_mode == ShardMode.EXPLICIT: + # shard_map manualises the mesh axes it is given and will not insert a reshard + # for an operand whose layout differs from `in_specs`, so hand it arrays that + # already match. + query = jax.sharding.reshard(query, qkv_pspec) + key = jax.sharding.reshard(key, qkv_pspec) + value = jax.sharding.reshard(value, qkv_pspec) + g = jax.sharding.reshard(g, g_beta_pspec) + beta = jax.sharding.reshard(beta, g_beta_pspec) + recurrent_state_arg = jax.sharding.reshard( + recurrent_state_arg, state_pspec + ) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=( + qkv_pspec, # query + qkv_pspec, # key + qkv_pspec, # value + g_beta_pspec, # g + g_beta_pspec, # beta + state_pspec, # initial_state + ), + out_specs=( + qkv_pspec, # core_attn_out + state_pspec, # final_state + ), + check_vma=False, + ) + def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): + return jax_chunk_gated_delta_rule( + query=q, + key=k, + value=v, + g=g_val, + beta=beta_val, + chunk_size=cfg.gdn_chunk_size, + initial_state=init_h, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + cp_axis=cp_axes or None, + ) + + core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg) + else: + core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( + query, + key, + value, + g, + beta, chunk_size=cfg.gdn_chunk_size, - initial_state=init_h, + initial_state=recurrent_state, use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, compute_dtype=cfg.dtype, - cp_axis=cp_axes or None, ) - core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg) - else: - core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( - query, - key, - value, - g, - beta, - chunk_size=cfg.gdn_chunk_size, - initial_state=recurrent_state, - use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, - compute_dtype=cfg.dtype, - ) - if model_mode != MODEL_MODE_TRAIN and active_cache is not None: assert next_conv_state is not None assert next_recurrent_state is not None diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index 8782f4c3e0..19aa650f56 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -85,6 +85,7 @@ "qwen3-next-80b-a3b": "Qwen/Qwen3-Next-80B-A3B-Instruct", "qwen3.5-397b-a17b": "Qwen/Qwen3.5-397B-A17B", "qwen3.5-35b-a3b": "Qwen/Qwen3.5-35B-A3B", + "qwen3.5-tiny": "Qwen/Qwen3.5-35B-A3B", "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistral-7b": "mistralai/Mistral-7B-v0.1", "mixtral-8x22b": "mistralai/Mixtral-8x22B-Instruct-v0.1", diff --git a/tests/unit/gdn_benchmark_test.py b/tests/unit/gdn_benchmark_test.py new file mode 100644 index 0000000000..d9aec24e76 --- /dev/null +++ b/tests/unit/gdn_benchmark_test.py @@ -0,0 +1,1825 @@ +# 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. + +"""Benchmarking and verification script for Canonical GDN kernel on Cloud TPU. + +Authoritative 2-way comparison: +1. Pure JAX GDN (Reference) +2. Canonical Decoupled GDN Kernel (use_gdn_kernel=True) +""" + +import argparse +import builtins +import functools +import glob +import os +import shutil +import sys +import time +import types +from typing import Any, Tuple + +print = functools.partial(builtins.print, flush=True) + +from absl import flags +from absl.testing import absltest + +os.environ["GLOG_minloglevel"] = "2" +os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" + +FLAGS = flags.FLAGS +flags.DEFINE_bool( + "enable_custom_call_tracing", + True, + "Whether to enable detailed LLO and custom call region tracing in XProf.", +) + + +def configure_custom_call_tracing(): + """Configures detailed LLO and custom call region tracing in XProf.""" + for flag_name in ( + "xla_xprof_register_llo_debug_info", + "xla_xprof_enable_custom_call_tracing", + ): + try: + FLAGS.set_default(flag_name, "True") + except Exception: + pass + try: + if hasattr(FLAGS, flag_name): + setattr(FLAGS, flag_name, True) + except Exception: + pass + + # Explicitly set xla_enable_custom_call_region_trace and xla_enable_transpose_trace + # to False to prevent b/523408873 (134-million bundle interpolation loop in + # convert_tpu_trace_to_xplane causing Forge OOM). + for disable_flag in ("xla_enable_custom_call_region_trace", "xla_enable_transpose_trace"): + try: + FLAGS.set_default(disable_flag, "False") + except Exception: + pass + try: + if hasattr(FLAGS, disable_flag): + setattr(FLAGS, disable_flag, False) + except Exception: + pass + + try: + FLAGS.set_default("tpu_chip_config_name", "megachip") + except Exception: + pass + try: + FLAGS.set_default("deepsea_chip_config_name", "megachip") + except Exception: + pass + + os.environ["TPU_CHIP_CONFIG_NAME"] = "megachip" + + extra_args = [ + "--xla_xprof_register_llo_debug_info=true", + "--xla_xprof_enable_custom_call_tracing=true", + "--xla_enable_custom_call_region_trace=false", + "--xla_enable_transpose_trace=false", + "--deepsea_chip_config_name=megachip", + ] + cur_libtpu = os.environ.get("LIBTPU_INIT_ARGS", "") + for arg in extra_args: + arg_key = arg.split("=")[0] + if arg_key not in cur_libtpu: + cur_libtpu = f"{cur_libtpu} {arg}".strip() + os.environ["LIBTPU_INIT_ARGS"] = cur_libtpu + + + +# Early top-level configuration before TPU / JAX initialization +_disable_custom_call = any( + arg in sys.argv + for arg in ( + "--noenable_custom_call_tracing", + "--enable_custom_call_tracing=false", + "--enable_custom_call_tracing=False", + "--enable_custom_call_tracing=0", + ) +) +if not _disable_custom_call: + configure_custom_call_tracing() + + +def setUpModule(): + """Module-level setup for absltest.""" + if getattr(FLAGS, "enable_custom_call_tracing", True): + configure_custom_call_tracing() + + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +# Force highest precision for TPU MXU 2/3-pass FP32 simulation +jax.config.update("jax_default_matmul_precision", "highest") + +try: + from maxtext.models.kernels.gdn import gdn_bwd_pallas + from maxtext.models import qwen3 +except ImportError: + from maxtext.src.maxtext.models.kernels.gdn import gdn_bwd_pallas + from maxtext.src.maxtext.models import qwen3 + + +def create_model_configs( + hidden_size: int = 4096, + num_key_heads: int = 16, + num_value_heads: int = 64, + head_dim: int = 128, + conv_kernel_dim: int = 4, + chunk_size: int = 64, + dtype: Any = jnp.float32, + weight_dtype: Any = None, + gdn_state_dtype: Any = jnp.float32, + gdn_decay_dtype: Any = jnp.float32, + use_qk_norm: bool = True, +) -> Tuple[types.SimpleNamespace, types.SimpleNamespace]: + """Creates configurations for Pure JAX (Reference) and Canonical GDN Kernel.""" + if dtype is None: + dtype = jnp.float32 + elif isinstance(dtype, str): + dtype = getattr(jnp, dtype) + if weight_dtype is None: + weight_dtype = dtype + elif isinstance(weight_dtype, str): + weight_dtype = getattr(jnp, weight_dtype) + if gdn_state_dtype is None: + gdn_state_dtype = jnp.float32 + elif isinstance(gdn_state_dtype, str): + gdn_state_dtype = getattr(jnp, gdn_state_dtype) + if gdn_decay_dtype is None: + gdn_decay_dtype = jnp.float32 + elif isinstance(gdn_decay_dtype, str): + gdn_decay_dtype = getattr(jnp, gdn_decay_dtype) + + base_dict = dict( + emb_dim=hidden_size, + gdn_num_value_heads=num_value_heads, + gdn_num_key_heads=num_key_heads, + gdn_key_head_dim=head_dim, + gdn_value_head_dim=head_dim, + gdn_conv_kernel_dim=conv_kernel_dim, + dtype=dtype, + weight_dtype=weight_dtype, + gdn_state_dtype=gdn_state_dtype, + gdn_decay_dtype=gdn_decay_dtype, + matmul_precision="highest", + normalization_layer_epsilon=1e-6, + gdn_chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm, + load_balance_loss_weight=0.0, + scan_layers=False, + using_pipeline_parallelism=False, + logical_axis_rules=(), + shard_mode="auto", + debug_sharding=False, + ) + + # 1. Pure JAX GDN config (Reference) + pure_jax_config = types.SimpleNamespace( + **base_dict, + use_gdn_kernel=False, + ) + + # 2. Canonical Decoupled GDN Kernel config (Decoupled v1.5) + gdn_kernel_config = types.SimpleNamespace( + **base_dict, + use_gdn_kernel=True, + ) + + return pure_jax_config, gdn_kernel_config + + +def create_jitted_train_step( + model: nnx.Module, + input_shape: Tuple[int, ...], + step_scope: str = "TrainStep", + fwd_scope: str = "Fwd", + bwd_scope: str = "Bwd", + remat: bool | str = False, +): + """Creates a pure functional, JIT-compiled training step with position-aware loss.""" + is_remat = ( + remat.lower() in ("full", "true", "yes", "1") + if isinstance(remat, str) + else bool(remat) + ) + graphdef, params = nnx.split(model) + + proj_key = jax.random.PRNGKey(99) + projection = jax.random.normal(proj_key, input_shape) + + @jax.jit + def pure_train_step(params, x): + with jax.named_scope(step_scope): + m = nnx.merge(graphdef, params) + + def loss_fn(m_inner): + with jax.named_scope(fwd_scope): + if is_remat: + out = jax.checkpoint(lambda m, inp: m(inp))(m_inner, x) + else: + out = m_inner(x) + y = out[0] if isinstance(out, tuple) else out + loss = jnp.mean(y * projection.astype(y.dtype)) + return loss, out + + with jax.named_scope(bwd_scope): + (loss, y), grads = nnx.value_and_grad(loss_fn, has_aux=True)(m) + return loss, y, grads + + return pure_train_step, params + + +def create_jitted_forward(model: nnx.Module, scope_name: str = "Fwd"): + """Creates a pure functional, JIT-compiled forward pass.""" + graphdef, params = nnx.split(model) + + @jax.jit + def pure_forward(params, x): + with jax.named_scope(scope_name): + m = nnx.merge(graphdef, params) + out = m(x) + return out + + return pure_forward, params + + +def print_numerical_correctness_table( + out_ref: Any, + out_test: Any, + loss_ref: Any, + loss_test: Any, + grads_ref: Any, + grads_test: Any, + tolerance: float = 1e-4, + abs_tolerance: float = 1e-5, + comparison_name: str = "Candidate vs Reference", + diff_records: dict[str, Any] | None = None, +) -> bool: + """Prints a numerical correctness comparison table between two implementations.""" + print( + "\n=========================================================================================" + ) + print(f">>> NUMERICAL CORRECTNESS: {comparison_name}") + print( + "=========================================================================================" + ) + header = ( + f" {'Tensor / Parameter':<40} | {'Max Abs Diff':<12} | {'Relative Diff':<13} |" + f" {'Tolerance':<10} | {'Status'}" + ) + sep = " " + "-" * (len(header) - 2) + print(sep) + print(header) + print(sep) + + rows = [] + + # 1. Forward Output + ref_t = np.asarray(out_ref[0] if isinstance(out_ref, tuple) else out_ref) + test_t = np.asarray(out_test[0] if isinstance(out_test, tuple) else out_test) + abs_d_fwd = float(np.max(np.abs(ref_t - test_t))) + rel_d_fwd = abs_d_fwd / (float(np.max(np.abs(ref_t))) + 1e-7) + match_fwd = (rel_d_fwd <= tolerance) or (abs_d_fwd <= abs_tolerance) + rows.append(("Forward Output", abs_d_fwd, rel_d_fwd, match_fwd)) + + # 2. Loss Scalar + lp = float(loss_ref) + la = float(loss_test) + abs_d_loss = abs(lp - la) + rel_d_loss = abs_d_loss / (abs(lp) + 1e-7) + match_loss = (rel_d_loss <= tolerance) or (abs_d_loss <= abs_tolerance) + rows.append(("Loss Scalar", abs_d_loss, rel_d_loss, match_loss)) + + # 3. Parameter Gradients + ref_leaves = jax.tree_util.tree_leaves_with_path(grads_ref) + test_leaves = jax.tree_util.tree_leaves_with_path(grads_test) + + for (path_ref, g_ref), (_, g_test) in zip(ref_leaves, test_leaves): + if not hasattr(g_ref, "shape") or not hasattr(g_test, "shape"): + continue + path_parts = [] + for k in path_ref: + if hasattr(k, "key"): + path_parts.append(str(k.key)) + elif hasattr(k, "name"): + path_parts.append(str(k.name)) + elif hasattr(k, "idx"): + path_parts.append(str(k.idx)) + else: + path_parts.append(str(k)) + name = ".".join(path_parts) + g_ref_np = np.asarray(g_ref) + g_test_np = np.asarray(g_test) + abs_d = float(np.max(np.abs(g_ref_np - g_test_np))) + rel_d = abs_d / (float(np.max(np.abs(g_ref_np))) + 1e-7) + is_m = (rel_d <= tolerance) or (abs_d <= abs_tolerance) + rows.append((name, abs_d, rel_d, is_m)) + + overall_diverged = False + for name, abs_d, rel_d, is_m in rows: + if not is_m: + overall_diverged = True + status = "❌ DIVERGED" + else: + status = "✅ MATCH" + if diff_records is not None: + diff_records[name] = { + "abs_diff": abs_d, + "rel_diff": rel_d, + "match": is_m, + "tolerance": tolerance, + "abs_tolerance": abs_tolerance, + } + print( + f" {name:<40} | {abs_d:<12.2e} | {rel_d:<13.2e} |" + f" {tolerance:<10.2e} | {status}" + ) + print(sep) + return overall_diverged + + +def get_device_memory_stats() -> dict[str, Any] | None: + """Returns memory stats dict from jax.devices()[0] if supported, else None.""" + try: + dev = jax.devices()[0] + if hasattr(dev, "memory_stats"): + stats = dev.memory_stats() + if stats and "bytes_in_use" in stats: + return stats + except Exception: + pass + return None + + +def get_compiled_memory_analysis(jit_fn: Any, params: Any, inputs: Any) -> Any | None: + """Extracts static HBM memory analysis from XLA compiler.""" + if hasattr(jit_fn, "memory_analysis"): + try: + return jit_fn.memory_analysis() + except Exception: + pass + if hasattr(jit_fn, "_cached_memory_analysis") and jit_fn._cached_memory_analysis is not None: + return jit_fn._cached_memory_analysis + try: + lowered = jit_fn.lower(params, inputs) + compiled = lowered.compile() + if hasattr(compiled, "memory_analysis"): + return compiled.memory_analysis() + except Exception: + pass + return None + + +def run_memory_profile_analysis( + kernel_names: list[str], + fwd_fns: list[Any], + train_fns: list[Any], + params_list: list[Any], + inputs: Any, + seq_len: int, + batch_size: int, + policy_label: str = "", +): + """Measures and displays comparative HBM memory usage across implementations.""" + title_suffix = ( + f" ({policy_label}S={seq_len}, B={batch_size}, Dtype=FP32)" + if policy_label + else f" (S={seq_len}, B={batch_size}, Dtype=FP32)" + ) + print( + "\n=========================================================================================" + ) + print( + f">>> HBM MEMORY PROFILING & COMPARATIVE ANALYSIS{title_suffix}" + ) + print( + "=========================================================================================" + ) + + fwd_act_mbs = [] + train_peak_mbs = [] + bwd_peak_mbs = [] + fwd_compiled_mbs = [] + train_compiled_mbs = [] + dev_peak_train_mbs = [] + breakdown_rows = [] + + for name, fwd_fn, train_fn, p in zip( + kernel_names, fwd_fns, train_fns, params_list + ): + mem_before_fwd = get_device_memory_stats() + out_fwd = fwd_fn(p, inputs) + jax.block_until_ready(out_fwd) + mem_after_fwd = get_device_memory_stats() + fwd_analysis = get_compiled_memory_analysis(fwd_fn, p, inputs) + + mem_before_train = get_device_memory_stats() + out_train = train_fn(p, inputs) + jax.block_until_ready(out_train) + mem_after_train = get_device_memory_stats() + train_analysis = get_compiled_memory_analysis(train_fn, p, inputs) + + dev_in_use_fwd = (mem_after_fwd["bytes_in_use"] / (1024**2)) if mem_after_fwd else 0.0 + dev_peak_fwd = (mem_after_fwd.get("peak_bytes_in_use", 0) / (1024**2)) if mem_after_fwd else 0.0 + dev_in_use_train = (mem_after_train["bytes_in_use"] / (1024**2)) if mem_after_train else 0.0 + dev_peak_train = (mem_after_train.get("peak_bytes_in_use", 0) / (1024**2)) if mem_after_train else 0.0 + + if fwd_analysis is not None: + fwd_act_mb = fwd_analysis.temp_size_in_bytes / (1024**2) + fwd_peak_compiled_mb = ( + fwd_analysis.argument_size_in_bytes + + fwd_analysis.temp_size_in_bytes + + fwd_analysis.output_size_in_bytes + ) / (1024**2) + else: + fwd_act_mb = dev_in_use_fwd + fwd_peak_compiled_mb = dev_peak_fwd + + if train_analysis is not None: + train_peak_compiled_mb = ( + train_analysis.argument_size_in_bytes + + train_analysis.temp_size_in_bytes + + train_analysis.output_size_in_bytes + ) / (1024**2) + train_peak_mb = train_peak_compiled_mb + bwd_peak_mb = max(train_peak_compiled_mb - fwd_peak_compiled_mb, 0.0) + else: + train_peak_mb = dev_peak_train if dev_peak_train > 0 else dev_in_use_train + bwd_peak_mb = max(train_peak_mb - fwd_act_mb, 0.0) + train_peak_compiled_mb = train_peak_mb + + fwd_act_mbs.append(fwd_act_mb) + train_peak_mbs.append(train_peak_mb) + bwd_peak_mbs.append(bwd_peak_mb) + fwd_compiled_mbs.append(fwd_peak_compiled_mb) + train_compiled_mbs.append(train_peak_compiled_mb) + dev_peak_train_mbs.append(dev_peak_train) + + if fwd_analysis is not None and train_analysis is not None: + breakdown_rows.append(( + name, + "Forward", + fwd_analysis.argument_size_in_bytes / (1024**2), + fwd_analysis.temp_size_in_bytes / (1024**2), + fwd_analysis.output_size_in_bytes / (1024**2), + fwd_peak_compiled_mb, + dev_in_use_fwd, + dev_peak_fwd, + )) + breakdown_rows.append(( + name, + "Backward (Est.)", + train_analysis.argument_size_in_bytes / (1024**2), + max(train_analysis.temp_size_in_bytes - fwd_analysis.temp_size_in_bytes, 0) / (1024**2), + train_analysis.output_size_in_bytes / (1024**2), + bwd_peak_mb, + dev_in_use_train, + dev_peak_train, + )) + breakdown_rows.append(( + name, + "Train Step", + train_analysis.argument_size_in_bytes / (1024**2), + train_analysis.temp_size_in_bytes / (1024**2), + train_analysis.output_size_in_bytes / (1024**2), + train_peak_compiled_mb, + dev_in_use_train, + dev_peak_train, + )) + else: + breakdown_rows.append((name, "Forward", 0.0, fwd_act_mb, 0.0, fwd_peak_compiled_mb, dev_in_use_fwd, dev_peak_fwd)) + breakdown_rows.append((name, "Backward (Est.)", 0.0, bwd_peak_mb, 0.0, bwd_peak_mb, dev_in_use_train, dev_peak_train)) + breakdown_rows.append((name, "Train Step", 0.0, train_peak_mb, 0.0, train_peak_mb, dev_in_use_train, dev_peak_train)) + + # 1. Comparative Summary Table + ref_fwd = fwd_act_mbs[0] if fwd_act_mbs[0] > 0 else 1.0 + ref_train = train_peak_mbs[0] if train_peak_mbs[0] > 0 else 1.0 + + summary_header = ( + f" {'Kernel Implementation':<36} | {'Fwd Activation Mem':<20} |" + f" {'Est. Backward Mem':<18} | {'Peak Train HBM':<18} |" + f" {'Fwd Ratio vs Pure':<20} | {'Train Ratio vs Pure'}" + ) + separator = " " + "-" * len(summary_header) + print("\nComparative HBM Memory Usage:") + print(separator) + print(summary_header) + print(separator) + + for i in range(len(kernel_names)): + f_mb = fwd_act_mbs[i] + b_mb = bwd_peak_mbs[i] + t_mb = train_peak_mbs[i] + f_ratio = f_mb / ref_fwd if ref_fwd > 0 else 1.0 + t_ratio = t_mb / ref_train if ref_train > 0 else 1.0 + f_pct = (f_ratio - 1.0) * 100.0 + t_pct = (t_ratio - 1.0) * 100.0 + + if i == 0: + f_str = "1.00x (ref)" + t_str = "1.00x (ref)" + else: + f_color = "🟢" if f_ratio <= 1.0 else "🔴" + t_color = "🟢" if t_ratio <= 1.0 else "🔴" + f_str = f"{f_ratio:.2f}x ({f_color} {f_pct:+.0f}%)" + t_str = f"{t_ratio:.2f}x ({t_color} {t_pct:+.0f}%)" + + print( + f" [{i + 1}] {kernel_names[i]:<32} | {f_mb:>16.2f} MB |" + f" {b_mb:>14.2f} MB | {t_mb:>14.2f} MB | {f_str:<20} | {t_str}" + ) + print(separator) + + # 2. Detailed Buffer Breakdown + if breakdown_rows: + print("\nDetailed Memory Breakdown (XLA Compiled Buffers & Allocator):") + b_header = ( + f" {'Implementation':<32} | {'Pass':<16} | {'Argument':<12} |" + f" {'Temp / Scratch':<14} | {'Output':<10} | {'Peak Total':<12} |" + f" {'Dev In-Use':<12} | {'Dev Peak'}" + ) + b_sep = " " + "-" * len(b_header) + print(b_sep) + print(b_header) + print(b_sep) + for impl, scope, arg, tmp, out, pk, dev_u, dev_pk in breakdown_rows: + print( + f" {impl:<32} | {scope:<16} | {arg:>9.2f} MB | {tmp:>11.2f} MB |" + f" {out:>7.2f} MB | {pk:>9.2f} MB | {dev_u:>9.2f} MB | {dev_pk:>9.2f} MB" + ) + print(b_sep) + + return fwd_act_mbs, train_peak_mbs, bwd_peak_mbs + + +def print_latency_comparison( + kernel_names: list[str], + fwd_lats: list[float], + bwd_lats: list[float], + train_lats: list[float], + policy_label: str = "", +) -> None: + """Prints a comprehensive 2-way latency and speedup comparison table.""" + print( + "\n=========================================================================================" + ) + print( + f">>> LATENCY & SPEEDUP: COMPARISON {policy_label}({kernel_names[0]} vs {kernel_names[1]})" + ) + print( + "=========================================================================================" + ) + header = ( + f" {'Pass / Step':<20} | {kernel_names[0]:<28} |" + f" {kernel_names[1]:<32} | {'Speedup':<12} | {'Winner'}" + ) + sep = " " + "-" * (len(header) - 2) + print(sep) + print(header) + print(sep) + + passes = [ + ("Forward Pass", [fwd_lats[0], fwd_lats[1]]), + ("Backward Pass", [bwd_lats[0], bwd_lats[1]]), + ("Full Training Step", [train_lats[0], train_lats[1]]), + ] + + for step_name, lats in passes: + p_val = lats[0] + k_val = lats[1] + + p_str = f"{p_val:>25.2f} ms" if not np.isnan(p_val) and p_val > 0 else "N/A" + k_str = f"{k_val:>29.2f} ms" if not np.isnan(k_val) and k_val > 0 else "FAILED" + + if (not np.isnan(p_val) and p_val > 0) and (not np.isnan(k_val) and k_val > 0): + speedup = p_val / k_val + speedup_str = f"{speedup:>9.2f}x" + if speedup > 1.05: + winner = f"🏆 {kernel_names[1]}" + elif speedup < 0.95: + winner = f"🏆 {kernel_names[0]}" + else: + winner = "≈ Parity" + else: + speedup_str = "N/A" + winner = "None" + + print( + f" {step_name:<20} | {p_str:>28} |" + f" {k_str:>32} | {speedup_str:<12} | {winner}" + ) + print(sep) + + +def print_tradeoff_table( + ref_name: str, + kernel_name: str, + fwd_ref: float, + fwd_k: float, + bwd_ref: float, + bwd_k: float, + train_ref: float, + train_k: float, + fwd_mem_ref: float, + fwd_mem_k: float, + train_mem_ref: float, + train_mem_k: float, + policy_label: str = "", +) -> None: + """Prints quantitative trade-off analysis of Canonical GDN Kernel vs Pure JAX Reference.""" + print( + "\n=========================================================================================" + ) + print( + f">>> QUANTITATIVE TRADE-OFF {policy_label}: {kernel_name} vs {ref_name}" + ) + print( + "=========================================================================================" + ) + header = ( + f" {'Metric':<30} | {ref_name:<28} | {kernel_name:<32} |" + f" {'Difference / Savings':<22} | {'Advantage'}" + ) + sep = " " + "-" * (len(header) - 2) + print(sep) + print(header) + print(sep) + + metrics = [ + ("Forward Pass Latency", fwd_ref, fwd_k, "ms", True), + ("Backward Pass Latency", bwd_ref, bwd_k, "ms", True), + ("Full Training Step Latency", train_ref, train_k, "ms", True), + ("Forward Activation Memory", fwd_mem_ref, fwd_mem_k, "MB", True), + ("Peak Train Step Memory", train_mem_ref, train_mem_k, "MB", True), + ] + + for label, val_ref, val_k, unit, lower_is_better in metrics: + val_ref_is_valid = not np.isnan(val_ref) and val_ref > 0 + val_k_is_valid = not np.isnan(val_k) and val_k > 0 + + val_ref_str = f"{val_ref:>25.2f} {unit}" if val_ref_is_valid else "N/A" + val_k_str = f"{val_k:>29.2f} {unit}" if val_k_is_valid else "FAILED" + + if val_ref_is_valid and val_k_is_valid: + diff = val_k - val_ref + pct = (diff / (val_ref + 1e-7)) * 100.0 + diff_str = f"{diff:+.2f} {unit} ({pct:+.1f}%)" + if abs(pct) < 1.0: + advantage = "≈ Parity" + elif (diff < 0 and lower_is_better) or (diff > 0 and not lower_is_better): + advantage = f"🏆 {kernel_name} ({abs(pct):.1f}% better)" + else: + advantage = f"🏆 {ref_name} ({abs(pct):.1f}% better)" + else: + diff_str = "N/A" + advantage = "N/A" + + print( + f" {label:<30} | {val_ref_str:>28} | {val_k_str:>32} |" + f" {diff_str:>22} | {advantage}" + ) + print(sep) + + +def print_cross_policy_summary(results: dict[str, dict[str, Any]]) -> None: + """Prints an executive summary table comparing results across remat policies.""" + print( + "\n=========================================================================================================================" + ) + print( + ">>> EXECUTIVE SUMMARY: CROSS-POLICY COMPARISON (remat=none vs remat=full)" + ) + print( + "=========================================================================================================================" + ) + # 1. Numerical Correctness + has_diffs = any("diff_records" in res and res["diff_records"] for res in results.values()) + if has_diffs: + num_header = ( + f" {'Tensor / Parameter':<40} | {'Rel Diff (none)':<16} |" + f" {'Rel Diff (full)':<16} | {'Tolerance':<10} | {'Status'}" + ) + num_sep = " " + "-" * (len(num_header) - 2) + print("\n1. Numerical Correctness (Relative Gradient Difference vs Pure JAX):") + print(num_sep) + print(num_header) + print(num_sep) + + all_names = [] + for policy in ("none", "full"): + if policy in results and "diff_records" in results[policy]: + for k in results[policy]["diff_records"]: + if k not in all_names: + all_names.append(k) + + for name in all_names: + none_rec = results.get("none", {}).get("diff_records", {}).get(name) + full_rec = results.get("full", {}).get("diff_records", {}).get(name) + none_str = f"{none_rec['rel_diff']:<16.2e}" if none_rec else f"{'N/A':<16}" + full_str = f"{full_rec['rel_diff']:<16.2e}" if full_rec else f"{'N/A':<16}" + tol_val = (full_rec or none_rec or {}).get("tolerance", 1e-4) + match_none = none_rec.get("match", True) if none_rec else True + match_full = full_rec.get("match", True) if full_rec else True + status = "✅ MATCH" if (match_none and match_full) else "❌ DIVERGED" + print( + f" {name:<40} | {none_str} | {full_str} |" + f" {tol_val:<10.2e} | {status}" + ) + print(num_sep) + + # 2. Latency & Speedup + header = ( + f" {'Configuration / Implementation':<42} | {'Remat':<6} |" + f" {'Fwd (ms)':<10} | {'Bwd (ms)':<10} | {'Train (ms)':<12} |" + f" {'Speedup vs Pure':<16} | {'Winner'}" + ) + sep = " " + "-" * (len(header) - 2) + print("\n2. Latency & Speedup Summary:") + print(sep) + print(header) + print(sep) + + for policy in ("none", "full"): + if policy not in results: + continue + res = results[policy] + fwd_p = res["t_fwd_pure"] + bwd_p = res["t_bwd_pure"] + trn_p = res["t_train_pure"] + fwd_k = res["t_fwd_kernel"] + bwd_k = res["t_bwd_kernel"] + trn_k = res["t_train_kernel"] + p_str = f"{trn_p:>10.2f} ms" if not np.isnan(trn_p) and trn_p > 0 else "FAILED" + k_str = f"{trn_k:>10.2f} ms" if not np.isnan(trn_k) and trn_k > 0 else "FAILED" + speedup = trn_p / trn_k if (not np.isnan(trn_p) and trn_p > 0 and not np.isnan(trn_k) and trn_k > 0) else 0.0 + speedup_str = f"{speedup:.2f}x" if speedup > 0 else "N/A" + winner = "🏆 Canonical GDN" if speedup > 1.05 else ("🏆 Pure JAX" if speedup < 0.95 and speedup > 0 else "≈ Parity") + + print( + f" {'Pure JAX GDN (Reference)':<42} | {policy:<6} |" + f" {fwd_p:>8.2f} ms | {bwd_p:>8.2f} ms | {p_str} |" + f" {'1.00x (ref)':<16} | ref" + ) + print( + f" {'Canonical GDN Kernel (use_gdn_kernel=True)':<42} | {policy:<6} |" + f" {fwd_k:>8.2f} ms | {bwd_k:>8.2f} ms | {k_str} |" + f" {speedup_str:<16} | {winner}" + ) + print(sep) + + # 3. HBM Memory Footprint Summary + mem_header = ( + f" {'Configuration / Implementation':<42} | {'Remat':<6} |" + f" {'Fwd Act (MB)':<12} | {'Est Bwd (MB)':<12} | {'Peak Train (MB)':<15} |" + f" {'Mem Ratio vs Pure':<18} | {'Savings vs Pure'}" + ) + mem_sep = " " + "-" * (len(mem_header) - 2) + print("\n3. HBM Memory Footprint Summary:") + print(mem_sep) + print(mem_header) + print(mem_sep) + + for policy in ("none", "full"): + if policy not in results: + continue + res = results[policy] + fwd_p_mem = res["fwd_act_mbs"][0] + bwd_p_mem = res["bwd_peak_mbs"][0] + trn_p_mem = res["train_peak_mbs"][0] + fwd_k_mem = res["fwd_act_mbs"][1] + bwd_k_mem = res["bwd_peak_mbs"][1] + trn_k_mem = res["train_peak_mbs"][1] + mem_ratio = trn_k_mem / trn_p_mem if trn_p_mem > 0 else 1.0 + mem_savings_pct = (1.0 - mem_ratio) * 100.0 + mem_savings_str = f"🟢 -{mem_savings_pct:.1f}%" if mem_savings_pct >= 0 else f"🔴 +{abs(mem_savings_pct):.1f}%" + + print( + f" {'Pure JAX GDN (Reference)':<42} | {policy:<6} |" + f" {fwd_p_mem:>10.2f} MB | {bwd_p_mem:>10.2f} MB | {trn_p_mem:>13.2f} MB |" + f" {'1.00x (ref)':<18} | ref" + ) + print( + f" {'Canonical GDN Kernel (use_gdn_kernel=True)':<42} | {policy:<6} |" + f" {fwd_k_mem:>10.2f} MB | {bwd_k_mem:>10.2f} MB | {trn_k_mem:>13.2f} MB |" + f" {f'{mem_ratio:.2f}x':<18} | {mem_savings_str}" + ) + print(mem_sep) + + +def print_3way_latency_summary( + pure_none_times: list[float], + pure_full_times: list[float], + kernel_none_times: list[float], + seq_len: int, + batch_size: int, + hidden_size: int, +) -> None: + """Prints a 3-way latency and speedup comparison table for the structured profile.""" + print( + "\n=========================================================================================================================" + ) + print( + f">>> STRUCTURED 3-WAY PROFILE LATENCY & SPEEDUP (Full Qwen3.5-397B Layer: S={seq_len}, B={batch_size}, H={hidden_size}, FP32)" + ) + print( + "=========================================================================================================================" + ) + header = ( + f" {'Section / Implementation':<42} | {'Remat':<6} | {'Steps':<5} |" + f" {'Avg Train (ms)':<15} | {'Min (ms)':<10} | {'Max (ms)':<10} |" + f" {'Speedup vs Pure(none)':<22} | {'Speedup vs Pure(full)'}" + ) + sep = " " + "-" * (len(header) - 2) + print(sep) + print(header) + print(sep) + + t_pn_avg = float(np.mean(pure_none_times)) if pure_none_times else float("nan") + t_pn_min = float(np.min(pure_none_times)) if pure_none_times else float("nan") + t_pn_max = float(np.max(pure_none_times)) if pure_none_times else float("nan") + + t_pf_avg = float(np.mean(pure_full_times)) if pure_full_times else float("nan") + t_pf_min = float(np.min(pure_full_times)) if pure_full_times else float("nan") + t_pf_max = float(np.max(pure_full_times)) if pure_full_times else float("nan") + + t_kn_avg = float(np.mean(kernel_none_times)) if kernel_none_times else float("nan") + t_kn_min = float(np.min(kernel_none_times)) if kernel_none_times else float("nan") + t_kn_max = float(np.max(kernel_none_times)) if kernel_none_times else float("nan") + + sp_kn_vs_pn = (t_pn_avg / t_kn_avg) if (t_kn_avg > 0 and t_pn_avg > 0) else 1.0 + sp_kn_vs_pf = (t_pf_avg / t_kn_avg) if (t_kn_avg > 0 and t_pf_avg > 0) else 1.0 + sp_pf_vs_pn = (t_pn_avg / t_pf_avg) if (t_pf_avg > 0 and t_pn_avg > 0) else 1.0 + + print( + f" {'Section 1: Pure JAX GDN (Reference)':<42} | {'none':<6} | {len(pure_none_times):<5} |" + f" {t_pn_avg:>11.2f} ms | {t_pn_min:>7.2f} ms | {t_pn_max:>7.2f} ms |" + f" {'1.00x (ref)':<22} | {sp_pf_vs_pn:>6.2f}x" + ) + print( + f" {'Section 2: Pure JAX GDN':<42} | {'full':<6} | {len(pure_full_times):<5} |" + f" {t_pf_avg:>11.2f} ms | {t_pf_min:>7.2f} ms | {t_pf_max:>7.2f} ms |" + f" {sp_pf_vs_pn:>6.2f}x | {'1.00x (ref)':<22}" + ) + kn_vs_pn_winner = f"{sp_kn_vs_pn:.2f}x (🏆 WINNER)" if sp_kn_vs_pn > 1.05 else f"{sp_kn_vs_pn:.2f}x" + kn_vs_pf_winner = f"{sp_kn_vs_pf:.2f}x (🏆 WINNER)" if sp_kn_vs_pf > 1.05 else f"{sp_kn_vs_pf:.2f}x" + print( + f" {'Section 3: Canonical GDN Kernel (v1.5)':<42} | {'none':<6} | {len(kernel_none_times):<5} |" + f" {t_kn_avg:>11.2f} ms | {t_kn_min:>7.2f} ms | {t_kn_max:>7.2f} ms |" + f" {kn_vs_pn_winner:<22} | {kn_vs_pf_winner}" + ) + print(sep) + + +def print_3mode_memory_table( + mem_pure_none: Any | None, + mem_pure_full: Any | None, + mem_kernel_none: Any | None, +) -> None: + """Prints comparative compiled HBM memory usage across the 3 benchmark modes.""" + print( + "\n=========================================================================================================================" + ) + print( + ">>> HBM MEMORY PROFILING & COMPILATION ANALYSIS (3-WAY COMPARISON, FP32)" + ) + print( + "=========================================================================================================================" + ) + header = ( + f" {'Section / Implementation':<42} | {'Remat':<6} | {'Argument (MB)':<14} |" + f" {'Temp/Scratch (MB)':<18} | {'Peak HBM (MB)':<14} | {'Ratio vs Pure(none)':<20} | {'Savings vs Pure(none)'}" + ) + sep = " " + "-" * (len(header) - 2) + print(sep) + print(header) + print(sep) + + def get_sizes(mem): + if mem is not None: + arg_mb = getattr(mem, "argument_size_in_bytes", 0) / (1024**2) + tmp_mb = getattr(mem, "temp_size_in_bytes", 0) / (1024**2) + out_mb = getattr(mem, "output_size_in_bytes", 0) / (1024**2) + pk_mb = arg_mb + tmp_mb + out_mb + return arg_mb, tmp_mb, out_mb, pk_mb + return 0.0, 0.0, 0.0, 0.0 + + arg_pn, tmp_pn, out_pn, pk_pn = get_sizes(mem_pure_none) + arg_pf, tmp_pf, out_pf, pk_pf = get_sizes(mem_pure_full) + arg_kn, tmp_kn, out_kn, pk_kn = get_sizes(mem_kernel_none) + + ref_pk = pk_pn if pk_pn > 0 else 1.0 + + # Section 1 + print( + f" {'Section 1: Pure JAX GDN (Reference)':<42} | {'none':<6} |" + f" {arg_pn:>11.2f} MB | {tmp_pn:>15.2f} MB | {pk_pn:>11.2f} MB |" + f" {'1.00x (ref)':<20} | {'ref'}" + ) + + # Section 2 + r_pf = pk_pf / ref_pk if ref_pk > 0 else 1.0 + sav_pf = (1.0 - r_pf) * 100.0 + sav_pf_str = f"🟢 -{sav_pf:.1f}%" if sav_pf >= 0 else f"🔴 +{abs(sav_pf):.1f}%" + print( + f" {'Section 2: Pure JAX GDN':<42} | {'full':<6} |" + f" {arg_pf:>11.2f} MB | {tmp_pf:>15.2f} MB | {pk_pf:>11.2f} MB |" + f" {f'{r_pf:.2f}x':<20} | {sav_pf_str}" + ) + + # Section 3 + r_kn = pk_kn / ref_pk if ref_pk > 0 else 1.0 + sav_kn = (1.0 - r_kn) * 100.0 + sav_kn_str = f"🟢 -{sav_kn:.1f}%" if sav_kn >= 0 else f"🔴 +{abs(sav_kn):.1f}%" + print( + f" {'Section 3: Canonical GDN Kernel (v1.5)':<42} | {'none':<6} |" + f" {arg_kn:>11.2f} MB | {tmp_kn:>15.2f} MB | {pk_kn:>11.2f} MB |" + f" {f'{r_kn:.2f}x':<20} | {sav_kn_str}" + ) + print(sep) + + +def run_gdn_comparison( + batch_size: int | None = None, + seq_len: int | None = None, + iters: int | None = None, + warmup: int | None = None, + dtype_str: str | None = None, + weight_dtype_str: str | None = None, + gdn_state_dtype_str: str | None = "float32", + gdn_decay_dtype_str: str | None = "float32", + hidden_size: int = 4096, + num_key_heads: int = 16, + num_value_heads: int = 64, + head_dim: int = 128, + conv_kernel_dim: int = 4, + chunk_size: int = 64, + remat_policy: str = "structured", + gap_seconds: float | None = None, +): + backend = jax.default_backend() + print(f"\nDevice: {jax.devices()[0]} ({backend})") + print( + "Precision: jax_default_matmul_precision = highest (TPU MXU multi-pass" + " FP32 simulation)" + ) + + if backend == "cpu": + gdn_bwd_pallas.ensure_cpu_interpret_registered() + + dtype = jnp.float32 if dtype_str is None else (getattr(jnp, dtype_str) if isinstance(dtype_str, str) else dtype_str) + weight_dtype = dtype if weight_dtype_str is None else (getattr(jnp, weight_dtype_str) if isinstance(weight_dtype_str, str) else weight_dtype_str) + gdn_state_dtype = jnp.float32 if gdn_state_dtype_str is None else (getattr(jnp, gdn_state_dtype_str) if isinstance(gdn_state_dtype_str, str) else gdn_state_dtype_str) + gdn_decay_dtype = jnp.float32 if gdn_decay_dtype_str is None else (getattr(jnp, gdn_decay_dtype_str) if isinstance(gdn_decay_dtype_str, str) else gdn_decay_dtype_str) + + # Hardware defaults: Dedicate strictly to 8k sequence length on TPU in FP32 + if backend == "tpu": + batch = 1 if batch_size is None else batch_size + slen = 8192 if seq_len is None else seq_len + num_iters = 5 if iters is None else iters + num_warmup = 2 if warmup is None else warmup + gap = 2.0 if gap_seconds is None else gap_seconds + else: + print("⚠️ Running on CPU: Using reduced dims and CPU interpret mode.") + batch = 1 if batch_size is None else batch_size + slen = 128 if seq_len is None else seq_len + num_iters = 2 if iters is None else iters + num_warmup = 1 if warmup is None else warmup + gap = 0.5 if gap_seconds is None else gap_seconds + + print( + f"Config: Batch={batch}, SeqLen={slen}, Dtype={dtype}," + f" WeightDtype={weight_dtype}, StateDtype={gdn_state_dtype}," + f" DecayDtype={gdn_decay_dtype}" + ) + print( + f"Model: H={hidden_size}, K_Heads={num_key_heads}," + f" V_Heads={num_value_heads}, HeadDim={head_dim}, ChunkSize={chunk_size}," + f" Structured 3-Mode Profiling ({num_iters} steps per mode, {gap:.1f}s gap)" + ) + + pure_jax_cfg, gdn_kernel_cfg = create_model_configs( + hidden_size=hidden_size, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + head_dim=head_dim, + conv_kernel_dim=conv_kernel_dim, + chunk_size=chunk_size, + dtype=dtype, + weight_dtype=weight_dtype, + gdn_state_dtype=gdn_state_dtype, + gdn_decay_dtype=gdn_decay_dtype, + use_qk_norm=True, + ) + + print("\nInitializing models...") + pure_jax_model = qwen3.Qwen3NextGatedDeltaNet( + config=pure_jax_cfg, rngs=nnx.Rngs(0) + ) + gdn_kernel_model = qwen3.Qwen3NextGatedDeltaNet( + config=gdn_kernel_cfg, rngs=nnx.Rngs(0) + ) + + _, params_state = nnx.split(gdn_kernel_model) + nnx.update(pure_jax_model, params_state) + print("✅ Both models synchronized with identical weights.") + + key = jax.random.PRNGKey(42) + inputs = jax.random.normal(key, (batch, slen, hidden_size), dtype=dtype) + + # Create the 3 JIT-compiled training step functions + print("\n--- Creating Functional JIT Training Steps ---") + # Mode 1: Pure JAX remat=none + jit_train_pure_none, params_pure_none = create_jitted_train_step( + pure_jax_model, + inputs.shape, + step_scope="PureJAX_RematNone_Step", + fwd_scope="PureJAX_RematNone_Fwd", + bwd_scope="PureJAX_RematNone_Bwd", + remat=False, + ) + # Mode 2: Pure JAX remat=full + jit_train_pure_full, params_pure_full = create_jitted_train_step( + pure_jax_model, + inputs.shape, + step_scope="PureJAX_RematFull_Step", + fwd_scope="PureJAX_RematFull_Fwd", + bwd_scope="PureJAX_RematFull_Bwd", + remat=True, + ) + # Mode 3: Canonical GDN Kernel remat=none + jit_train_kernel_none, params_kernel_none = create_jitted_train_step( + gdn_kernel_model, + inputs.shape, + step_scope="GdnKernel_RematNone_Step", + fwd_scope="GdnKernel_RematNone_Fwd", + bwd_scope="GdnKernel_RematNone_Bwd", + remat=False, + ) + + # ========================================================================= + # 1. WARMUP & JIT COMPILATION STRICTLY OUTSIDE THE TRACE + # ========================================================================= + print( + "\n=========================================================================================" + ) + print( + f">>> STEP 1: PRE-WARMUP & JIT COMPILATION (STRICTLY OUTSIDE TRACE, {num_warmup} warmups each)" + ) + print( + "=========================================================================================" + ) + + # Mode 1: Pure JAX remat=none + print(f"[{time.strftime('%X')}] Compiling Mode 1: Pure JAX (remat=none)...") + compiled_pure_none = jit_train_pure_none.lower(params_pure_none, inputs).compile() + mem_pure_none = compiled_pure_none.memory_analysis() if hasattr(compiled_pure_none, "memory_analysis") else None + jit_train_pure_none._cached_memory_analysis = mem_pure_none + print(f"[{time.strftime('%X')}] Warming up Mode 1: Pure JAX (remat=none) ({num_warmup} warmups)...") + for _ in range(num_warmup): + res_pure_none = compiled_pure_none(params_pure_none, inputs) + jax.block_until_ready(res_pure_none) + loss_pure_none, out_pure_none, grads_pure_none = res_pure_none + + # Mode 2: Pure JAX remat=full + print(f"[{time.strftime('%X')}] Compiling Mode 2: Pure JAX (remat=full)...") + compiled_pure_full = jit_train_pure_full.lower(params_pure_full, inputs).compile() + mem_pure_full = compiled_pure_full.memory_analysis() if hasattr(compiled_pure_full, "memory_analysis") else None + jit_train_pure_full._cached_memory_analysis = mem_pure_full + print(f"[{time.strftime('%X')}] Warming up Mode 2: Pure JAX (remat=full) ({num_warmup} warmups)...") + for _ in range(num_warmup): + res_pure_full = compiled_pure_full(params_pure_full, inputs) + jax.block_until_ready(res_pure_full) + loss_pure_full, out_pure_full, grads_pure_full = res_pure_full + + # Mode 3: Canonical GDN Kernel remat=none + print(f"[{time.strftime('%X')}] Compiling Mode 3: Canonical GDN Kernel (remat=none)...") + compiled_kernel_none = jit_train_kernel_none.lower(params_kernel_none, inputs).compile() + mem_kernel_none = compiled_kernel_none.memory_analysis() if hasattr(compiled_kernel_none, "memory_analysis") else None + jit_train_kernel_none._cached_memory_analysis = mem_kernel_none + print(f"[{time.strftime('%X')}] Warming up Mode 3: Canonical GDN Kernel (remat=none) ({num_warmup} warmups)...") + for _ in range(num_warmup): + res_kernel_none = compiled_kernel_none(params_kernel_none, inputs) + jax.block_until_ready(res_kernel_none) + loss_kernel_none, out_kernel_none, grads_kernel_none = res_kernel_none + print(f"[{time.strftime('%X')}] ✅ All 3 modes compiled and warmed up outside trace.") + + # Numerical correctness checks (outside trace) + if dtype == jnp.bfloat16: + tol = 1.5e-2 + abs_tol = 3e-2 + elif backend == "cpu": + tol = 1e-3 + abs_tol = 1e-5 + else: + tol = 1e-4 + abs_tol = 1e-5 + policy_diffs_kernel = {} + diverged_kernel = print_numerical_correctness_table( + out_ref=out_pure_none, + out_test=out_kernel_none, + loss_ref=loss_pure_none, + loss_test=loss_kernel_none, + grads_ref=grads_pure_none, + grads_test=grads_kernel_none, + tolerance=tol, + abs_tolerance=abs_tol, + comparison_name="Canonical GDN Kernel (remat=none) vs Pure JAX (remat=none)", + diff_records=policy_diffs_kernel, + ) + + policy_diffs_remat = {} + diverged_remat = print_numerical_correctness_table( + out_ref=out_pure_none, + out_test=out_pure_full, + loss_ref=loss_pure_none, + loss_test=loss_pure_full, + grads_ref=grads_pure_none, + grads_test=grads_pure_full, + tolerance=tol, + abs_tolerance=abs_tol, + comparison_name="Pure JAX (remat=full) vs Pure JAX (remat=none)", + diff_records=policy_diffs_remat, + ) + + overall_numerical_diverged = diverged_kernel or diverged_remat + + # Memory profiling table (outside trace) + print_3mode_memory_table( + mem_pure_none, + mem_pure_full, + mem_kernel_none, + ) + + # ========================================================================= + # 2. START TRACE (PRISTINE TRACE CONTAINING ONLY THE 3 STRUCTURED SECTIONS) + # ========================================================================= + log_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", "/tmp/xprof_traces") + os.makedirs(log_dir, exist_ok=True) + print( + "\n=========================================================================================" + ) + print(f">>> STEP 2: STARTING PRISTINE XPROF TRACE (log_dir={log_dir})") + print( + "=========================================================================================" + ) + + tracing_active = False + try: + jax.profiler.start_trace(log_dir) + tracing_active = True + print(f"[{time.strftime('%X')}] ✅ jax.profiler.start_trace active.") + except Exception as e: + print(f"⚠️ Failed to start JAX profiler trace: {e}") + + # --- Section 1: Pure JAX remat=none --- + print(f"\n[{time.strftime('%X')}] >>> Tracing Section 1: Pure JAX remat=none ({num_iters} steps)...") + pure_none_times = [] + for step_i in range(num_iters): + t_s = time.perf_counter() + with jax.named_scope("PureJAX_RematNone_Step"): + with jax.profiler.StepTraceAnnotation("PureJAX_RematNone_Step", step_num=step_i): + with jax.profiler.TraceAnnotation(f"PureJAX_RematNone_Step_{step_i}"): + res = compiled_pure_none(params_pure_none, inputs) + jax.block_until_ready(res) + dur = (time.perf_counter() - t_s) * 1000.0 + pure_none_times.append(dur) + print(f" [Section 1: Pure JAX remat=none] Step {step_i + 1}/{num_iters}: {dur:.2f} ms") + + # --- Gap 1: Intentional gap with device sync --- + print(f"\n[{time.strftime('%X')}] >>> Gap 1: Device sync & intentional {gap:.1f}s gap between Section 1 and Section 2...") + jax.block_until_ready(res) + with jax.profiler.TraceAnnotation(f"Gap1_Sleep_{gap:.1f}s"): + time.sleep(gap) + + # --- Section 2: Pure JAX remat=full --- + print(f"\n[{time.strftime('%X')}] >>> Tracing Section 2: Pure JAX remat=full ({num_iters} steps)...") + pure_full_times = [] + for step_i in range(num_iters): + t_s = time.perf_counter() + with jax.named_scope("PureJAX_RematFull_Step"): + with jax.profiler.StepTraceAnnotation("PureJAX_RematFull_Step", step_num=step_i): + with jax.profiler.TraceAnnotation(f"PureJAX_RematFull_Step_{step_i}"): + res = compiled_pure_full(params_pure_full, inputs) + jax.block_until_ready(res) + dur = (time.perf_counter() - t_s) * 1000.0 + pure_full_times.append(dur) + print(f" [Section 2: Pure JAX remat=full] Step {step_i + 1}/{num_iters}: {dur:.2f} ms") + + # --- Gap 2: Intentional gap with device sync --- + print(f"\n[{time.strftime('%X')}] >>> Gap 2: Device sync & intentional {gap:.1f}s gap between Section 2 and Section 3...") + jax.block_until_ready(res) + with jax.profiler.TraceAnnotation(f"Gap2_Sleep_{gap:.1f}s"): + time.sleep(gap) + + # --- Section 3: Canonical GDN Kernel remat=none --- + print(f"\n[{time.strftime('%X')}] >>> Tracing Section 3: Canonical GDN Kernel remat=none ({num_iters} steps)...") + kernel_none_times = [] + for step_i in range(num_iters): + t_s = time.perf_counter() + with jax.named_scope("GdnKernel_RematNone_Step"): + with jax.profiler.StepTraceAnnotation("GdnKernel_RematNone_Step", step_num=step_i): + with jax.profiler.TraceAnnotation(f"GdnKernel_RematNone_Step_{step_i}"): + res = compiled_kernel_none(params_kernel_none, inputs) + jax.block_until_ready(res) + dur = (time.perf_counter() - t_s) * 1000.0 + kernel_none_times.append(dur) + print(f" [Section 3: Canonical GDN Kernel remat=none] Step {step_i + 1}/{num_iters}: {dur:.2f} ms") + + # --- Stop Trace --- + if tracing_active: + try: + del res + except Exception: + pass + import gc + gc.collect() + print(f"\n[{time.strftime('%X')}] Stopping XProf trace...") + try: + jax.profiler.stop_trace() + print(f"[{time.strftime('%X')}] ✅ jax.profiler.stop_trace completed. Trace written to: {log_dir}") + except Exception as e: + print(f"⚠️ Failed to stop JAX profiler trace: {e}") + + # ========================================================================= + # 3. POST-TRACE ARTIFACTS & EXECUTIVE SUMMARY + # ========================================================================= + xplane_files = glob.glob(os.path.join(log_dir, "**/*.xplane.pb"), recursive=True) + print(f"\nDiscovered {len(xplane_files)} generated .xplane.pb file(s) in {log_dir}:") + for xf in xplane_files: + sz = os.path.getsize(xf) + print(f" 📁 {xf} ({sz:,} bytes)") + target_named_copy = os.path.join(log_dir, "ghostlite_structured_3modes.xplane.pb") + if os.path.abspath(xf) != os.path.abspath(target_named_copy): + try: + shutil.copy(xf, target_named_copy) + except Exception: + pass + target_traced_copy = os.path.join( + log_dir, "ghostlite_structured_3modes_custom_call_traced.xplane.pb" + ) + if os.path.abspath(xf) != os.path.abspath(target_traced_copy): + try: + shutil.copy(xf, target_traced_copy) + except Exception: + pass + try: + os.makedirs("/tmp/xprof_traces", exist_ok=True) + shutil.copy(xf, os.path.join("/tmp/xprof_traces", os.path.basename(xf))) + shutil.copy(xf, "/tmp/xprof_traces/ghostlite_structured_3modes.xplane.pb") + shutil.copy( + xf, + "/tmp/xprof_traces/ghostlite_structured_3modes_custom_call_traced.xplane.pb", + ) + except Exception: + pass + + print_3way_latency_summary( + pure_none_times=pure_none_times, + pure_full_times=pure_full_times, + kernel_none_times=kernel_none_times, + seq_len=slen, + batch_size=batch, + hidden_size=hidden_size, + ) + + return overall_numerical_diverged + + +# Backwards compatibility alias +run_analytical_comparison = run_gdn_comparison + + +def print_standalone_kernel_memory_table( + mem_fwd: Any | None, + mem_train: Any | None, + seq_len: int, +) -> None: + """Prints compiled HBM memory usage for standalone GDN kernel.""" + print( + "\n=========================================================================================================================" + ) + print( + f">>> HBM MEMORY PROFILING & COMPILATION ANALYSIS (STANDALONE CANONICAL GDN KERNEL, S={seq_len})" + ) + print( + "=========================================================================================================================" + ) + header = ( + f" {'Step / Graph':<30} | {'Argument (MB)':<14} |" + f" {'Temp/Scratch (MB)':<18} | {'Output (MB)':<14} | {'Peak HBM (MB)':<14}" + ) + sep = " " + "-" * (len(header) - 2) + print(sep) + print(header) + print(sep) + + def get_sizes(mem): + if mem is not None: + arg_mb = getattr(mem, "argument_size_in_bytes", 0) / (1024**2) + tmp_mb = getattr(mem, "temp_size_in_bytes", 0) / (1024**2) + out_mb = getattr(mem, "output_size_in_bytes", 0) / (1024**2) + pk_mb = getattr(mem, "peak_memory_in_bytes", 0) / (1024**2) + if pk_mb == 0: + pk_mb = arg_mb + tmp_mb + out_mb + return arg_mb, tmp_mb, out_mb, pk_mb + return 0.0, 0.0, 0.0, 0.0 + + arg_f, tmp_f, out_f, pk_f = get_sizes(mem_fwd) + arg_t, tmp_t, out_t, pk_t = get_sizes(mem_train) + print( + f" {'Forward Pass Only':<30} | {arg_f:>11.2f} MB | {tmp_f:>15.2f} MB |" + f" {out_f:>11.2f} MB | {pk_f:>11.2f} MB" + ) + print( + f" {'Training Step (Fwd + Bwd)':<30} | {arg_t:>11.2f} MB | {tmp_t:>15.2f} MB |" + f" {out_t:>11.2f} MB | {pk_t:>11.2f} MB" + ) + print(sep) + + +def run_standalone_kernel_profile( + batch_size: int | None = None, + seq_len: int | None = None, + iters: int | None = None, + warmup: int | None = None, + dtype_str: str | None = "bfloat16", + weight_dtype_str: str | None = "bfloat16", + gdn_state_dtype_str: str | None = "float32", + gdn_decay_dtype_str: str | None = "float32", + hidden_size: int = 4096, + num_key_heads: int = 16, + num_value_heads: int = 64, + head_dim: int = 128, + conv_kernel_dim: int = 4, + chunk_size: int = 64, +): + backend = jax.default_backend() + print(f"\nDevice: {jax.devices()[0]} ({backend})") + print("Precision: Standalone Canonical GDN Kernel Profile (Hybrid Precision)") + + if backend == "cpu": + gdn_bwd_pallas.ensure_cpu_interpret_registered() + + dtype = jnp.bfloat16 if dtype_str is None else (getattr(jnp, dtype_str) if isinstance(dtype_str, str) else dtype_str) + weight_dtype = dtype if weight_dtype_str is None else (getattr(jnp, weight_dtype_str) if isinstance(weight_dtype_str, str) else weight_dtype_str) + gdn_state_dtype = jnp.float32 if gdn_state_dtype_str is None else (getattr(jnp, gdn_state_dtype_str) if isinstance(gdn_state_dtype_str, str) else gdn_state_dtype_str) + gdn_decay_dtype = jnp.float32 if gdn_decay_dtype_str is None else (getattr(jnp, gdn_decay_dtype_str) if isinstance(gdn_decay_dtype_str, str) else gdn_decay_dtype_str) + + if backend == "tpu": + batch = 1 if batch_size is None else batch_size + slen = 65536 if seq_len is None else seq_len + num_iters = 3 if iters is None else iters + num_warmup = 1 if warmup is None else warmup + else: + print("⚠️ Running on CPU: Using reduced dims and CPU interpret mode.") + batch = 1 if batch_size is None else batch_size + slen = 128 if seq_len is None else seq_len + num_iters = 2 if iters is None else iters + num_warmup = 1 if warmup is None else warmup + + print( + f"Config: Batch={batch}, SeqLen={slen}, Dtype={dtype}," + f" WeightDtype={weight_dtype}, StateDtype={gdn_state_dtype}," + f" DecayDtype={gdn_decay_dtype}" + ) + print( + f"Model: H={hidden_size}, K_Heads={num_key_heads}," + f" V_Heads={num_value_heads}, HeadDim={head_dim}, ChunkSize={chunk_size}" + f" ({slen // chunk_size} chunks)" + ) + + _, gdn_kernel_cfg = create_model_configs( + hidden_size=hidden_size, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + head_dim=head_dim, + conv_kernel_dim=conv_kernel_dim, + chunk_size=chunk_size, + dtype=dtype, + weight_dtype=weight_dtype, + gdn_state_dtype=gdn_state_dtype, + gdn_decay_dtype=gdn_decay_dtype, + use_qk_norm=True, + ) + + print("\nInitializing Canonical GDN Kernel model...") + gdn_kernel_model = qwen3.Qwen3NextGatedDeltaNet( + config=gdn_kernel_cfg, rngs=nnx.Rngs(0) + ) + + key = jax.random.PRNGKey(42) + inputs = jax.random.normal(key, (batch, slen, hidden_size), dtype=dtype) + + print("\n--- Creating Functional JIT Training Steps (Forward & Train Step) ---") + jit_fwd_kernel, params_fwd = create_jitted_forward( + gdn_kernel_model, scope_name="GdnKernel_64k_Fwd" + ) + jit_train_kernel, params_train = create_jitted_train_step( + gdn_kernel_model, + inputs.shape, + step_scope="GdnKernel_64k_TrainStep", + fwd_scope="GdnKernel_64k_Fwd", + bwd_scope="GdnKernel_64k_Bwd", + remat=False, + ) + + print("\n--- Compiling and Warming Up ---") + print(f"[{time.strftime('%X')}] Compiling Forward Pass (S={slen})...") + compiled_fwd = jit_fwd_kernel.lower(params_fwd, inputs).compile() + mem_fwd = ( + compiled_fwd.memory_analysis() + if hasattr(compiled_fwd, "memory_analysis") + else None + ) + + print( + f"[{time.strftime('%X')}] Compiling Training Step (Forward + Backward," + f" S={slen})..." + ) + compiled_train = jit_train_kernel.lower(params_train, inputs).compile() + mem_train = ( + compiled_train.memory_analysis() + if hasattr(compiled_train, "memory_analysis") + else None + ) + + print(f"[{time.strftime('%X')}] Warming up ({num_warmup} steps)...") + for _ in range(num_warmup): + res_fwd = compiled_fwd(params_fwd, inputs) + jax.block_until_ready(res_fwd) + res_train = compiled_train(params_train, inputs) + jax.block_until_ready(res_train) + + print(f"[{time.strftime('%X')}] ✅ Warmed up successfully.") + + # Print Memory Profile + print_standalone_kernel_memory_table(mem_fwd, mem_train, slen) + + # Profiling with trace + log_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", "/tmp/xprof_traces") + os.makedirs(log_dir, exist_ok=True) + tracing_active = False + try: + jax.profiler.start_trace(log_dir) + tracing_active = True + print(f"[{time.strftime('%X')}] ✅ jax.profiler.start_trace active for 64k profile.") + except Exception as e: + print(f"⚠️ Failed to start JAX profiler trace: {e}") + + fwd_times = [] + train_times = [] + print( + f"\n[{time.strftime('%X')}] >>> Executing {num_iters} iterations of" + " 64k Canonical GDN Kernel..." + ) + for step_i in range(num_iters): + # Fwd + t_f = time.perf_counter() + with jax.named_scope("GdnKernel_64k_Fwd_Step"): + with jax.profiler.StepTraceAnnotation( + "GdnKernel_64k_Fwd_Step", step_num=step_i + ): + res_f = compiled_fwd(params_fwd, inputs) + jax.block_until_ready(res_f) + fwd_dur = (time.perf_counter() - t_f) * 1000.0 + fwd_times.append(fwd_dur) + + # Train (Fwd + Bwd) + t_t = time.perf_counter() + with jax.named_scope("GdnKernel_64k_Train_Step"): + with jax.profiler.StepTraceAnnotation( + "GdnKernel_64k_Train_Step", step_num=step_i + ): + res_t = compiled_train(params_train, inputs) + jax.block_until_ready(res_t) + train_dur = (time.perf_counter() - t_t) * 1000.0 + train_times.append(train_dur) + bwd_dur = max(0.0, train_dur - fwd_dur) + print( + f" Step {step_i + 1}/{num_iters}: Fwd = {fwd_dur:.2f} ms, Bwd =" + f" {bwd_dur:.2f} ms, Total Step = {train_dur:.2f} ms" + ) + + if tracing_active: + try: + del res_f, res_t + except Exception: + pass + import gc + gc.collect() + print(f"\n[{time.strftime('%X')}] Stopping XProf trace for 64k profile...") + try: + jax.profiler.stop_trace() + print(f"[{time.strftime('%X')}] ✅ Trace written to: {log_dir}") + except Exception as e: + print(f"⚠️ Failed to stop JAX profiler trace: {e}") + + # Check generated .xplane.pb + xplane_files = glob.glob(os.path.join(log_dir, "**/*.xplane.pb"), recursive=True) + print(f"\nDiscovered {len(xplane_files)} generated .xplane.pb file(s) in {log_dir}:") + for xf in xplane_files: + sz = os.path.getsize(xf) + print(f" 📁 {xf} ({sz:,} bytes)") + target_named_copy = os.path.join( + log_dir, "ghostfish_64k_canonical_gdn.xplane.pb" + ) + if os.path.abspath(xf) != os.path.abspath(target_named_copy): + try: + shutil.copy(xf, target_named_copy) + except Exception: + pass + try: + os.makedirs("/tmp/xprof_traces", exist_ok=True) + shutil.copy(xf, os.path.join("/tmp/xprof_traces", os.path.basename(xf))) + shutil.copy(xf, "/tmp/xprof_traces/ghostfish_64k_canonical_gdn.xplane.pb") + except Exception: + pass + + # Print Summary Table + avg_fwd = float(np.mean(fwd_times)) + avg_train = float(np.mean(train_times)) + avg_bwd = max(0.0, avg_train - avg_fwd) + print("\n" + "=" * 80) + print( + ">>> 64K CANONICAL GDN KERNEL PERFORMANCE SUMMARY (S=" + f"{slen}, Chunks={slen // chunk_size})" + ) + print("=" * 80) + print(f" Forward Pass Latency: {avg_fwd:.2f} ms") + print(f" Backward Pass Latency: {avg_bwd:.2f} ms") + print(f" Total Step Latency: {avg_train:.2f} ms") + print("=" * 80 + "\n") + + +class GdnBenchmarkTest(absltest.TestCase): + + def setUp(self): + super().setUp() + jax.config.update("jax_default_matmul_precision", "highest") + gdn_bwd_pallas.ensure_cpu_interpret_registered() + + def tearDown(self): + super().tearDown() + import gc + gc.collect() + try: + jax.clear_caches() + except Exception: + pass + + def test_structured_profile_3modes(self): + """Structured 3-mode profile benchmark on Cloud TPU.""" + self.test_benchmark_397b_8k_fp32() + + def test_benchmark_8k_fp32(self): + """Backwards compatibility alias.""" + self.test_benchmark_397b_8k_fp32() + + def test_benchmark_397b_8k_fp32(self): + """Primary benchmark testing Pure JAX vs Canonical GDN Kernel (Decoupled v1.5) in FP32 at full Qwen3.5-397B config (S=8192, H=4096, V=64, K=16).""" + backend = jax.default_backend() + if backend == "tpu": + print( + "\n=========================================================================================" + ) + print( + ">>> BENCHMARK: Dedicated 8k FP32 Comparison (Pure JAX vs Canonical GDN Kernel - Full Qwen3.5-397B Config)" + ) + print( + "=========================================================================================" + ) + bench_iters = 1 if getattr(FLAGS, "enable_custom_call_tracing", True) else 5 + gap_sec = 0.05 if getattr(FLAGS, "enable_custom_call_tracing", True) else 2.0 + diverged = run_gdn_comparison( + batch_size=1, + seq_len=8192, + iters=bench_iters, + warmup=2, + dtype_str="float32", + hidden_size=4096, + num_key_heads=16, + num_value_heads=64, + head_dim=128, + conv_kernel_dim=4, + chunk_size=64, + gap_seconds=gap_sec, + ) + else: + print( + "\n=========================================================================================" + ) + print(">>> CPU HERMETIC VERIFICATION: Structured 3-Mode FP32 Comparison (S=128, B=1)") + print( + "=========================================================================================" + ) + diverged = run_gdn_comparison( + batch_size=1, + seq_len=128, + iters=2, + warmup=1, + dtype_str="float32", + hidden_size=2048, + num_key_heads=8, + num_value_heads=16, + head_dim=128, + conv_kernel_dim=4, + chunk_size=64, + gap_seconds=0.5, + ) + self.assertFalse( + diverged, "GDN Kernel gradients diverged beyond tolerance in FP32!" + ) + + def test_benchmark_397b_8k_hybrid_precision(self): + """Tier 1: Dedicated 8k comparison of Pure JAX vs Canonical GDN Kernel in Hybrid Precision (BF16 activations/weights, FP32 recurrent states/decay).""" + backend = jax.default_backend() + if backend == "tpu": + print( + "\n=========================================================================================" + ) + print( + ">>> TIER 1 BENCHMARK: Dedicated 8k Hybrid Precision Comparison (Pure" + " JAX vs Canonical GDN Kernel - Full Qwen3.5-397B Config)" + ) + print( + "=========================================================================================" + ) + bench_iters = 1 if getattr(FLAGS, "enable_custom_call_tracing", True) else 5 + gap_sec = 0.05 if getattr(FLAGS, "enable_custom_call_tracing", True) else 2.0 + diverged = run_gdn_comparison( + batch_size=1, + seq_len=8192, + iters=bench_iters, + warmup=2, + dtype_str="bfloat16", + weight_dtype_str="bfloat16", + gdn_state_dtype_str="float32", + gdn_decay_dtype_str="float32", + hidden_size=4096, + num_key_heads=16, + num_value_heads=64, + head_dim=128, + conv_kernel_dim=4, + chunk_size=64, + gap_seconds=gap_sec, + ) + else: + print( + "\n=========================================================================================" + ) + print( + ">>> CPU HERMETIC VERIFICATION: Tier 1 Hybrid Precision Comparison" + " (S=128, B=1)" + ) + print( + "=========================================================================================" + ) + diverged = run_gdn_comparison( + batch_size=1, + seq_len=128, + iters=2, + warmup=1, + dtype_str="bfloat16", + weight_dtype_str="bfloat16", + gdn_state_dtype_str="float32", + gdn_decay_dtype_str="float32", + hidden_size=2048, + num_key_heads=8, + num_value_heads=16, + head_dim=128, + conv_kernel_dim=4, + chunk_size=64, + gap_seconds=0.5, + ) + self.assertFalse( + diverged, + "GDN Kernel gradients diverged beyond tolerance in Hybrid Precision!", + ) + + def test_benchmark_397b_64k_kernel_profile(self): + """Tier 2: Standalone Canonical GDN Kernel Profiling at S=65536 in Hybrid Precision on Cloud TPU v7x.""" + backend = jax.default_backend() + if backend == "tpu": + print( + "\n=========================================================================================" + ) + print( + ">>> TIER 2 BENCHMARK: 64k Standalone Kernel Profile (Canonical GDN" + " Kernel - 1024 Chunks)" + ) + print( + "=========================================================================================" + ) + bench_iters = 1 if getattr(FLAGS, "enable_custom_call_tracing", True) else 3 + run_standalone_kernel_profile( + batch_size=1, + seq_len=65536, + iters=bench_iters, + warmup=1, + dtype_str="bfloat16", + weight_dtype_str="bfloat16", + gdn_state_dtype_str="float32", + gdn_decay_dtype_str="float32", + hidden_size=4096, + num_key_heads=16, + num_value_heads=64, + head_dim=128, + conv_kernel_dim=4, + chunk_size=64, + ) + else: + print( + "\n=========================================================================================" + ) + print( + ">>> CPU HERMETIC VERIFICATION: Tier 2 Standalone Kernel Profile" + " (S=128, B=1)" + ) + print( + "=========================================================================================" + ) + run_standalone_kernel_profile( + batch_size=1, + seq_len=128, + iters=2, + warmup=1, + dtype_str="bfloat16", + weight_dtype_str="bfloat16", + gdn_state_dtype_str="float32", + gdn_decay_dtype_str="float32", + hidden_size=2048, + num_key_heads=8, + num_value_heads=16, + head_dim=128, + conv_kernel_dim=4, + chunk_size=64, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Benchmark GDN Kernel") + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--seq_len", type=int, default=None) + parser.add_argument("--iters", type=int, default=None) + parser.add_argument("--warmup", type=int, default=None) + parser.add_argument("--dtype", type=str, default=None) + parser.add_argument("--hidden_size", type=int, default=4096) + parser.add_argument("--num_key_heads", type=int, default=16) + parser.add_argument("--num_value_heads", type=int, default=64) + parser.add_argument("--head_dim", type=int, default=128) + parser.add_argument("--conv_kernel_dim", type=int, default=4) + parser.add_argument("--chunk_size", type=int, default=64) + parser.add_argument( + "--remat", + type=str, + default="structured", + choices=["full", "none", "both", "structured"], + help="Remat policy: 'full', 'none', 'both', or 'structured'", + ) + parser.add_argument( + "--gap_seconds", + type=float, + default=2.0, + help="Gap between sections in seconds (default: 2.0)", + ) + + if "--benchmark" in sys.argv: + sys.argv.remove("--benchmark") + args, _ = parser.parse_known_args() + run_gdn_comparison( + batch_size=args.batch_size, + seq_len=args.seq_len, + iters=args.iters, + warmup=args.warmup, + dtype_str=args.dtype, + hidden_size=args.hidden_size, + num_key_heads=args.num_key_heads, + num_value_heads=args.num_value_heads, + head_dim=args.head_dim, + conv_kernel_dim=args.conv_kernel_dim, + chunk_size=args.chunk_size, + remat_policy=args.remat, + gap_seconds=args.gap_seconds, + ) + else: + absltest.main() diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py new file mode 100644 index 0000000000..43a7403082 --- /dev/null +++ b/tests/unit/gdn_bwd_pallas_test.py @@ -0,0 +1,1162 @@ +# 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. + +"""Unit tests for gdn_bwd_pallas with canonical GDN backward pass.""" + +import functools +from absl.testing import absltest +import jax +import jax.numpy as jnp +import numpy as np + +try: + from maxtext.models.kernels.gdn import gdn_bwd_pallas + from maxtext.models import qwen3 +except ImportError: + from maxtext.src.maxtext.models.kernels.gdn import gdn_bwd_pallas + from maxtext.src.maxtext.models import qwen3 + + +class GdnBwdPallasTest(absltest.TestCase): + + def setUp(self): + super().setUp() + gdn_bwd_pallas.ensure_cpu_interpret_registered() + + def test_chunk_forward_matches_jax_reference(self): + key = jax.random.PRNGKey(42) + chunk_size = 64 + num_kq_heads = 2 + num_v_heads = 4 + kq_head_dim = 128 + v_head_dim = 128 + repeats = num_v_heads // num_kq_heads + + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + q = jax.random.normal( + k1, (chunk_size, num_kq_heads, kq_head_dim), dtype=jnp.float32 + ) + k = jax.random.normal( + k2, (chunk_size, num_kq_heads, kq_head_dim), dtype=jnp.float32 + ) + v = jax.random.normal( + k3, (chunk_size, num_v_heads, v_head_dim), dtype=jnp.float32 + ) + b_val = jax.random.normal(k4, (chunk_size, num_v_heads), dtype=jnp.float32) + a_val = jax.random.normal(k5, (chunk_size, num_v_heads), dtype=jnp.float32) + a_log_val = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias_val = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + state_prev = jax.random.normal( + k8, (num_v_heads, kq_head_dim, v_head_dim), dtype=jnp.float32 + ) + + out_emit, state_emit, t_inv = ( + gdn_bwd_pallas.chunk_forward_with_tinv( + q, + k, + v, + b_val, + a_val, + a_log_val, + dt_bias_val, + state_prev, + kq_head_dim=kq_head_dim, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + ) + + self.assertEqual(t_inv.shape, (num_v_heads, chunk_size, chunk_size)) + + # Reference computation using qwen3.jax_chunk_gated_delta_rule + q_4d = q[None, :, :, :] + k_4d = k[None, :, :, :] + v_4d = v[None, :, :, :] + q_rep_4d = jnp.repeat(q_4d, repeats, axis=2) + k_rep_4d = jnp.repeat(k_4d, repeats, axis=2) + beta_3d = jax.nn.sigmoid(b_val)[None, :, :] + log_g_3d = (-jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val))[ + None, :, : + ] + state_4d = state_prev[None, :, :, :] + + expected_out, expected_state = qwen3.jax_chunk_gated_delta_rule( + query=q_rep_4d, + key=k_rep_4d, + value=v_4d, + g=log_g_3d, + beta=beta_3d, + chunk_size=chunk_size, + initial_state=state_4d, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + + 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_compute_forward_conv_and_states(self): + batch_size = 1 + chunk_size = 16 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 128 + head_v_dim = 128 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(456) + k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, 7) + + qkv = jax.random.normal( + k1, (batch_size, seq_len, dim_size), dtype=jnp.float32 + ) + b = jax.random.normal( + k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + a = jax.random.normal( + k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + conv_weight = jax.random.normal( + k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32 + ) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + + qkv_conv, chunk_states, t_inv = ( + gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + 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, + ) + ) + + self.assertEqual( + chunk_states.shape, + (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), + ) + self.assertEqual( + t_inv.shape, + (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), + ) + + conv_input = jnp.pad( + qkv.astype(jnp.float32), ((0, 0), (conv_kernel_size - 1, 0), (0, 0)) + ) + expected_conv_out = jax.lax.conv_general_dilated( + lhs=conv_input, + rhs=conv_weight.astype(jnp.float32), + window_strides=(1,), + padding="VALID", + dimension_numbers=("NWC", "WIO", "NWC"), + feature_group_count=dim_size, + ) + expected_conv_out = expected_conv_out + conv_bias + expected_qkv_conv = jax.nn.silu(expected_conv_out) + np.testing.assert_allclose( + qkv_conv, expected_qkv_conv, rtol=1e-5, atol=1e-5 + ) + + def test_fused_conv1d_gdn_kernel_gradient_against_autodiff(self): + """Compares gdn_fused_conv1d custom VJP against JAX autodiff on pure JAX.""" + batch_size = 1 + chunk_size = 64 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 2 + num_v_heads = 4 + head_k_dim = 128 + head_v_dim = 128 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(789) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal( + k1, (batch_size, seq_len, dim_size), dtype=jnp.float32 + ) + b = jax.random.normal( + k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + a = jax.random.normal( + k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + conv_weight = jax.random.normal( + k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32 + ) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal( + k8, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32 + ) + + # 1. Golden Reference Gradient via Autodiff on pure JAX implementation + def loss_pure(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.pure_jax_fused_conv1d_gdn( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + 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, + ) + return jnp.sum(out * do) + + exp_dqkv, exp_db, exp_da, exp_dcw, exp_dcb, exp_dal, exp_ddt = jax.grad( + loss_pure, argnums=(0, 1, 2, 3, 4, 5, 6) + )(qkv, b, a, conv_weight, conv_bias, a_log, dt_bias) + + # 2. Kernel Gradients via GDN Kernel custom VJP + def loss_gdn_kernel(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = ( + gdn_bwd_pallas.gdn_fused_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + 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, + ) + ) + return jnp.sum(out * do) + + act_dqkv, act_db, act_da, act_dcw, act_dcb, act_dal, act_ddt = jax.grad( + loss_gdn_kernel, argnums=(0, 1, 2, 3, 4, 5, 6) + )(qkv, b, a, conv_weight, conv_bias, a_log, dt_bias) + + print( + "\n--- GDN Kernel Custom VJP vs Pure JAX Autodiff Breakdown ---" + ) + comparisons = [ + ("beta (d_b)", exp_db, act_db), + ("alpha (d_a)", exp_da, act_da), + ("a_log (d_a_log)", exp_dal, act_dal), + ("dt_bias (d_dt_bias)", exp_ddt, act_ddt), + ("qkv (d_qkv)", exp_dqkv, act_dqkv), + ("conv_weight (d_conv_weight)", exp_dcw, act_dcw), + ("conv_bias (d_conv_bias)", exp_dcb, act_dcb), + ] + + for name, exp_g, act_g in comparisons: + self.assertIsNotNone(act_g, f"{name} actual gradient is None") + abs_diff = float(jnp.max(jnp.abs(exp_g - act_g))) + rel_diff = abs_diff / (float(jnp.max(jnp.abs(exp_g))) + 1e-7) + status = "✅ MATCH" if rel_diff < 1e-3 else "❌ DIVERGED" + print( + f" {name:<28}: MaxAbsDiff = {abs_diff:.2e} | RelDiff =" + f" {rel_diff:.2e} | {status}" + ) + self.assertLess( + rel_diff, + 1e-3, + f"{name} relative difference {rel_diff:.2e} exceeds tolerance 1e-3", + ) + + print( + "✅ All 7 parameter gradients match Pure JAX autodiff within 0.1% on" + " CPU!" + ) + + def test_fused_conv1d_gdn_kernel_conv_bias_none(self): + """Verifies GDN kernel backward executes correctly when conv_bias is None.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(101) + k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, 7) + + qkv = jax.random.normal( + k1, (batch_size, seq_len, dim_size), dtype=jnp.float32 + ) + b = jax.random.normal( + k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + a = jax.random.normal( + k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + conv_weight = jax.random.normal( + k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32 + ) + a_log = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal( + k7, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32 + ) + + def loss_fn(qkv_in, b_in, a_in, cw_in, al_in, dt_in): + out, _ = ( + gdn_bwd_pallas.gdn_fused_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=None, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + 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, + ) + ) + return jnp.sum(out * do) + + grads = jax.grad(loss_fn, argnums=(0, 1, 2, 3, 4, 5))( + qkv, b, a, conv_weight, a_log, dt_bias + ) + for g in grads: + self.assertIsNotNone(g) + self.assertFalse(np.any(np.isnan(np.array(g)))) + + def test_fused_conv1d_gdn_kernel_multi_batch(self): + """Verifies GDN kernel backward handles batch_size > 1.""" + batch_size = 2 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(202) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal( + k1, (batch_size, seq_len, dim_size), dtype=jnp.float32 + ) + b = jax.random.normal( + k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + a = jax.random.normal( + k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + conv_weight = jax.random.normal( + k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32 + ) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal( + k8, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32 + ) + + def loss_fn(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = ( + gdn_bwd_pallas.gdn_fused_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + 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, + ) + ) + return jnp.sum(out * do) + + grads = jax.grad(loss_fn, argnums=(0, 1, 2, 3, 4, 5, 6))( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias + ) + for g in grads: + self.assertIsNotNone(g) + self.assertFalse(np.any(np.isnan(np.array(g)))) + + def test_chunk_state_forward_with_cached_tinv_matches_chunk_forward(self): + """Verifies chunk_state_forward_with_cached_tinv matches chunk_forward_with_tinv state.""" + key = jax.random.PRNGKey(999) + chunk_size = 64 + num_kq_heads = 2 + num_v_heads = 4 + kq_head_dim = 128 + v_head_dim = 128 + repeats = num_v_heads // num_kq_heads + + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + q = jax.random.normal( + k1, (chunk_size, num_kq_heads, kq_head_dim), dtype=jnp.float32 + ) + k = jax.random.normal( + k2, (chunk_size, num_kq_heads, kq_head_dim), dtype=jnp.float32 + ) + v = jax.random.normal( + k3, (chunk_size, num_v_heads, v_head_dim), dtype=jnp.float32 + ) + b_val = jax.random.normal(k4, (chunk_size, num_v_heads), dtype=jnp.float32) + a_val = jax.random.normal(k5, (chunk_size, num_v_heads), dtype=jnp.float32) + a_log_val = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias_val = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + state_prev = jax.random.normal( + k8, (num_v_heads, kq_head_dim, v_head_dim), dtype=jnp.float32 + ) + + _, expected_state, t_inv = ( + gdn_bwd_pallas.chunk_forward_with_tinv( + q, + k, + v, + b_val, + a_val, + a_log_val, + dt_bias_val, + state_prev, + kq_head_dim=kq_head_dim, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + ) + + actual_state = ( + gdn_bwd_pallas.chunk_state_forward_with_cached_tinv( + k=k, + v=v, + b_val=b_val, + a_val=a_val, + a_log_val=a_log_val, + dt_bias_val=dt_bias_val, + state_prev=state_prev, + t_inv=t_inv, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + ) + + np.testing.assert_allclose( + actual_state, expected_state, rtol=1e-6, atol=1e-6 + ) + + def test_compute_forward_conv_and_states_with_cached_tinv(self): + """Verifies _compute_forward_conv_and_states with cached_t_inv matches uncached version.""" + batch_size = 1 + chunk_size = 16 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(888) + k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, 7) + + qkv = jax.random.normal( + k1, (batch_size, seq_len, dim_size), dtype=jnp.float32 + ) + b = jax.random.normal( + k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + a = jax.random.normal( + k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + conv_weight = jax.random.normal( + k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32 + ) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + + qkv_conv_ref, chunk_states_ref, t_inv_ref = ( + gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + 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, + ) + ) + + qkv_conv_cached, chunk_states_cached, t_inv_cached = ( + gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + 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, + cached_t_inv=t_inv_ref, + ) + ) + + np.testing.assert_allclose( + qkv_conv_cached, qkv_conv_ref, rtol=1e-6, atol=1e-6 + ) + np.testing.assert_allclose( + chunk_states_cached, chunk_states_ref, rtol=1e-6, atol=1e-6 + ) + np.testing.assert_allclose(t_inv_cached, t_inv_ref, rtol=1e-6, atol=1e-6) + + def test_fused_conv1d_gdn_kernel_bwd_with_cached_tinv_in_residuals(self): + """Verifies _gdn_fused_conv1d_bwd gives identical grads with cached t_inv.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(777) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal( + k1, (batch_size, seq_len, dim_size), dtype=jnp.float32 + ) + b = jax.random.normal( + k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + a = jax.random.normal( + k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + conv_weight = jax.random.normal( + k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32 + ) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal( + k8, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32 + ) + + _, chunk_states, t_inv = ( + gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + 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, + ) + ) + + res_none = ( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + None, + None, + None, + ) + res_cached = ( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + None, + None, + t_inv, + ) + res_cached_all = ( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + None, + None, + t_inv, + chunk_states, + ) + cotangents = (do, (None, None)) + + grads_none = ( + gdn_bwd_pallas._gdn_fused_conv1d_bwd( + 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, + residuals=res_none, + cotangents=cotangents, + ) + ) + + grads_cached = ( + gdn_bwd_pallas._gdn_fused_conv1d_bwd( + 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, + residuals=res_cached, + cotangents=cotangents, + ) + ) + + grads_cached_all = ( + gdn_bwd_pallas._gdn_fused_conv1d_bwd( + 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, + residuals=res_cached_all, + cotangents=cotangents, + ) + ) + + for g_none, g_cached, g_cached_all in zip( + grads_none, grads_cached, grads_cached_all + ): + if g_none is not None and g_cached is not None: + np.testing.assert_allclose(g_none, g_cached, rtol=1e-5, atol=1e-5) + if g_none is not None and g_cached_all is not None: + np.testing.assert_allclose(g_none, g_cached_all, rtol=1e-5, atol=1e-5) + + def test_run_local_gdn_fused_fwd_returns_cached_chunk_states(self): + """Verifies _run_local_gdn_fused_fwd returns properly shaped chunk_states.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(101) + k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, 7) + + qkv = jax.random.normal( + k1, (batch_size, seq_len, dim_size), dtype=jnp.float32 + ) + b = jax.random.normal( + k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + a = jax.random.normal( + k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + conv_weight = jax.random.normal( + k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32 + ) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + + (out, states), t_inv, chunk_states = ( + gdn_bwd_pallas._run_local_gdn_fused_fwd( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + conv_state=None, + recurrent_state=None, + 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, + ) + ) + + self.assertIsNotNone(chunk_states) + self.assertIsNotNone(t_inv) + self.assertEqual( + chunk_states.shape, + (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), + ) + self.assertEqual( + t_inv.shape, + (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), + ) + + # Verify against golden pure JAX _compute_forward_conv_and_states + _, exp_chunk_states, exp_t_inv = ( + gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + 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, + ) + ) + np.testing.assert_allclose( + chunk_states, exp_chunk_states, rtol=1e-5, atol=1e-5 + ) + np.testing.assert_allclose(t_inv, exp_t_inv, rtol=1e-5, atol=1e-5) + + def test_fused_conv1d_gdn_kernel_gradient_with_initial_states(self): + """Verifies custom VJP gradients when initial conv_state and recurrent_state are provided.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(999) + k1, k2, k3, k4, k5, k6, k7, k8, k9, k10 = jax.random.split(key, 10) + + qkv = jax.random.normal( + k1, (batch_size, seq_len, dim_size), dtype=jnp.float32 + ) + b = jax.random.normal( + k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + a = jax.random.normal( + k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32 + ) + conv_weight = jax.random.normal( + k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32 + ) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal( + k8, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32 + ) + conv_state = jax.random.normal( + k9, (batch_size, conv_kernel_size - 1, dim_size), dtype=jnp.float32 + ) + recurrent_state = jax.random.normal( + k10, + (batch_size, num_v_heads, head_k_dim, head_v_dim), + dtype=jnp.float32, + ) + + def loss_pure(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.pure_jax_fused_conv1d_gdn( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=conv_state, + recurrent_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=True, + compute_dtype=jnp.float32, + ) + return jnp.sum(out * do) + + def loss_gdn_kernel(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = ( + gdn_bwd_pallas.gdn_fused_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=conv_state, + recurrent_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=True, + compute_dtype=jnp.float32, + ) + ) + return jnp.sum(out * do) + + exp_grads = jax.grad(loss_pure, argnums=(0, 1, 2, 3, 4, 5, 6))( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias + ) + act_grads = jax.grad(loss_gdn_kernel, argnums=(0, 1, 2, 3, 4, 5, 6))( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias + ) + + for exp_g, act_g in zip(exp_grads, act_grads): + self.assertIsNotNone(act_g) + np.testing.assert_allclose(exp_g, act_g, rtol=1e-3, atol=1e-3) + + def test_gdn_kernel_bwd_multi_group_head_parallel(self): + """Verifies multi-group head-parallel grid dispatch matches reference.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 4 + num_v_heads = 8 + head_k_dim = 64 + head_v_dim = 64 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(1234) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + 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_v_dim), dtype=jnp.float32) + chunk_states = jax.random.normal( + k7, (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32 + ) + t_inv = jax.random.normal( + k8, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), dtype=jnp.float32 + ) + + # 1. Dispatch with head_tile = 4 -> 2 head groups + dy1, db1, da1, dal1, ddt1 = ( + gdn_bwd_pallas.pallas_gdn_bwd_kernel( + 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, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + head_tile=4, + ) + ) + + # 2. Dispatch with head_tile = 8 -> 1 head group + dy2, db2, da2, dal2, ddt2 = ( + gdn_bwd_pallas.pallas_gdn_bwd_kernel( + 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, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + head_tile=8, + ) + ) + + np.testing.assert_allclose(dy1, dy2, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(db1, db2, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(da1, da2, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(dal1, dal2, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(ddt1, ddt2, rtol=1e-3, atol=1e-3) + + def test_gdn_kernel_bwd_variable_length_segment_ids_reset(self): + """Verifies segment_ids document boundaries reset carried state gradient to prevent leakage.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(5678) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a_log = jax.random.normal(k4, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32) + chunk_states = jax.random.normal( + k6, (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32 + ) + t_inv = jax.random.normal( + k7, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), dtype=jnp.float32 + ) + + # Only chunk 1 has non-zero incoming gradients; chunk 0 do is all zeros + do_chunk1 = jax.random.normal( + k8, (batch_size, chunk_size, num_v_heads, head_v_dim), dtype=jnp.float32 + ) + do_chunk0 = jnp.zeros((batch_size, chunk_size, num_v_heads, head_v_dim), dtype=jnp.float32) + do = jnp.concatenate([do_chunk0, do_chunk1], axis=1) + + # 1. No document reset: gradient flows backwards from chunk 1 into chunk 0 + dy_no_reset, _, _, _, _ = ( + gdn_bwd_pallas.pallas_gdn_bwd_kernel( + 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, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + segment_ids=None, + ) + ) + # Chunk 0 gradient is non-zero due to recurrent state carrying gradients from chunk 1 + self.assertGreater(float(jnp.max(jnp.abs(dy_no_reset[:, :chunk_size, :]))), 1e-4) + + # 2. With segment_ids boundary between chunk 0 (doc 0) and chunk 1 (doc 1) + seg_doc0 = jnp.zeros((batch_size, chunk_size), dtype=jnp.int32) + seg_doc1 = jnp.ones((batch_size, chunk_size), dtype=jnp.int32) + segment_ids = jnp.concatenate([seg_doc0, seg_doc1], axis=1) + + dy_reset, _, _, _, _ = ( + gdn_bwd_pallas.pallas_gdn_bwd_kernel( + 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, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + segment_ids=segment_ids, + ) + ) + # Chunk 0 gradient is strictly zero because boundary reset eliminated cross-document leakage! + np.testing.assert_allclose( + dy_reset[:, :chunk_size, :], + jnp.zeros_like(dy_reset[:, :chunk_size, :]), + atol=1e-6, + ) + + def test_fused_conv1d_gdn_kernel_bwd_with_head_tile(self): + """Verifies pallas_fused_conv1d_gdn_bwd_kernel forwards head_tile correctly.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 4 + num_v_heads = 8 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(999) + k1, k2, k3, k4, k5, k6, k7, k8, k9, k10 = jax.random.split(key, 10) + + pre_conv_qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + 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_v_dim), dtype=jnp.float32) + chunk_states = jax.random.normal( + k7, (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32 + ) + conv_weight = jax.random.normal(k8, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k9, (dim_size,), dtype=jnp.float32) + t_inv = jax.random.normal( + k10, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), dtype=jnp.float32 + ) + + res1 = ( + gdn_bwd_pallas.pallas_fused_conv1d_gdn_bwd_kernel( + pre_conv_qkv=pre_conv_qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + conv_weight=conv_weight, + conv_bias=conv_bias, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + head_tile=4, + ) + ) + + res2 = ( + gdn_bwd_pallas.pallas_fused_conv1d_gdn_bwd_kernel( + pre_conv_qkv=pre_conv_qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + conv_weight=conv_weight, + conv_bias=conv_bias, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + head_tile=8, + ) + ) + + for g1, g2 in zip(res1, res2): + if g1 is not None and g2 is not None: + np.testing.assert_allclose(g1, g2, rtol=2e-2, atol=2.5e-1) + + +if __name__ == "__main__": + absltest.main() + + diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index 7e272fc581..e5b5e90fa2 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -45,6 +45,7 @@ def test_gmm_v2_heuristic_tiling_requires_gmm_v2(self): with self.assertRaisesRegex(ValueError, "`use_gmm_v2_heuristic_tiling=True` requires `use_gmm_v2=True`."): pyconfig.initialize( [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + skip_jax_distributed_system=True, use_gmm_v2_heuristic_tiling=True, use_gmm_v2=False, ) diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py index 70653e4ec7..88adf575a6 100644 --- a/tests/unit/train_compile_test.py +++ b/tests/unit/train_compile_test.py @@ -1274,6 +1274,61 @@ def test_qwen3_5_explicit_sharding_zero1(self): ), )) + def test_qwen3_5_gdn_kernel(self): + """AOT test for qwen3-5 with analytical GDN kernel and GMM v2""" + compiled_trainstep_file = "/tmp/test_qwen3_5_gdn_kernel" + train_compile_main( + ( + "", + get_test_config_path(), + f"compiled_trainstep_file={compiled_trainstep_file}", + "compile_topology=tpu7x-512", + "compile_topology_num_slices=1", + "model_name=qwen3.5-397b-a17b", + "per_device_batch_size=0.5", + "max_target_length=4096", + "ici_tensor_parallelism=2", + "sparse_matmul=True", + "megablox=True", + "use_tokamax_gmm=True", + "use_gmm_v2=True", + "use_tokamax_splash=True", + "base_num_decoder_layers=2", + "override_model_config=True", + "enable_nnx=False", + "pure_nnx_decoder=False", + "pure_nnx=False", + "use_gdn_kernel=True", + ) + ) + + def test_qwen3_5_gdn_kernel_v5p(self): + """AOT test for qwen3-5 with analytical GDN kernel and GMM v2""" + compiled_trainstep_file = "/tmp/test_qwen3_5_gdn_kernel" + train_compile_main( + ( + "", + get_test_config_path(), + f"compiled_trainstep_file={compiled_trainstep_file}", + "compile_topology=tpu7x-64", + "compile_topology_num_slices=1", + "model_name=qwen3.5-35b-a3b", + "per_device_batch_size=8", + "max_target_length=4096", + "ici_tensor_parallelism=2", + "sparse_matmul=True", + "megablox=True", + "use_tokamax_gmm=True", + "use_gmm_v2=True", + "use_tokamax_splash=True", + "override_model_config=True", + "enable_nnx=False", + "pure_nnx_decoder=False", + "pure_nnx=False", + "use_gdn_kernel=True", + ) + ) + def test_serialization_and_deserialization_formats(self): """Tests that our custom binary save/load functions work securely and legacy fallback triggers warning."""