diff --git a/docs/design/kda_cp_support.md b/docs/design/kda_cp_support.md new file mode 100644 index 0000000000..582ab2ecfe --- /dev/null +++ b/docs/design/kda_cp_support.md @@ -0,0 +1,167 @@ +# Design Doc: KDA CP (Context Parallelism) Support + +## Summary + +This PR integrates KDA (Kimi Delta Attention) into MaxText with tokamax backend and CP (context parallelism) support. It adds the `KimiDeltaAttention` layer, `ShortConvolution`, QKV/beta/gate projections, and CP-aware causal convolution boundary handling. The `CPContext` mechanism passes context information to the `chunk_kda` kernel for coordinated recurrent state across CP ranks. + +## Design + +### CP Data Flow Overview + +``` +No CP: + [B, T, E] → QKV proj → ShortConv → SiLU → L2Norm → chunk_kda → output + +CP (cp_size > 1): + [B, T/cp, E] → QKV proj → SHARD_MAP(ShortConv w/ halo) ← independent conv shard_map + → SiLU + L2Norm + → CPContext(mesh, "context") ← constructed outside shard_map + → _inject_context_on_T + _wsc ← partition spec fixup + → SHARD_MAP(chunk_kda) ← cp_context passed in + → [B, T/cp, E] +``` + +Key difference from MLA CP: MLA relies on splash attention kernel internally doing implicit all_gather K/V → local attention; KDA does not rely on all_gather. Instead, `CPContext` lets the kernel coordinate recurrent state across ranks during forward/backward. + +### Plan 1: `halo_exchange_for_conv` (`utils/cp_utils.py`, new file) + +ShortConvolution is a causal 1D depthwise convolution. Under CP sharding, each rank lacks the preceding `kernel_size-1` historical tokens at its left boundary. + +``` +rank 0: [t0 t1 t2 t3] pad: [0 0 t0 t1 t2 t3] ← zeros (sequence start) +rank 1: [t4 t5 t6 t7] pad: [t2 t3 t4 t5 t6 t7] ← pull t2, t3 from rank 0 +``` + +**Algorithm**: +1. `jnp.pad(x, (halo_size, 0))` — left zero-pad +2. Outside CP scope or cp_size==1 → return padded directly (degenerate causal padding) +3. Inside CP scope: `ppermute` forward ring — rank i sends its last `halo_size` tokens to rank i+1, rank 0's halo is set to zero +4. `return jnp.concatenate([halo, x], axis=seq_axis)` + +`ppermute` is a collective op and must be called inside a scope that exposes the `"context"` axis. See Plan 2. + +### Plan 2: ShortConvolution CP Wrapper (`layers/attention_kda.py`) + +`ShortConvolution.__call__` internally calls `halo_exchange_for_conv`, which requires the `"context"` axis scope. Inside `KimiDeltaAttention.__call__`, when CP is enabled, wrap the q/k/v conv calls in an independent `jax.shard_map`. + +Change location: the conv call segment after QKV projection in `KimiDeltaAttention.__call__` (see dev branch implementation `attention_kda.py:407-429`). + +Key design decisions: +- **conv shard_map and chunk_kda shard_map are independent**: two separate `jax.shard_map` invocations, freeing conv's ppermute buffer in between +- `check_vma=False`: FlashAttention custom rules may falsely report VMA errors +- Zero-overhead fallback when no CP: follows the original path exactly + +### Plan 3: chunk_kda CPContext + Partition Spec (`attention_kda.py`) + +#### 3a. CPContext Construction (outside shard_map) + +```python +try: + from cp_utils import CPContext +except ImportError: + CPContext = None + +cp_ctx = CPContext(mesh=self.mesh, axis_name="context") +``` + +`CPContext` is a frozen dataclass. `mesh` and `axis_name` are set at construction time; chain metadata fields are populated internally by `chunk_kda`. + +#### 3b. Partition Spec Injection + +`nnx.logical_to_mesh_axes` may map the T axis to `None` due to Flax rule priority + size-1 axis stripping, but shard_map requires the T axis to have `"context"` sharding: + +```python +def _inject_context_on_T(pspec, t_axis=1): + spec = list(pspec) + if spec[t_axis] is None: + spec[t_axis] = "context" + return jax.sharding.PartitionSpec(*spec) +``` + +Applied to `qkv_pspec`, `beta_pspec`, `seg_pspec` when CP is enabled, followed by `with_sharding_constraint` to ensure tensor physical layout matches. + +#### 3c. chunk_kda shard_map + +Under CP, pass through `cp_context=cp_ctx` and `segment_ids` to the `chunk_kda` kernel. + +segment_ids handling: +- **varlen**: pass through as-is +- **non-varlen + CP**: construct dummy `jnp.ones(q.shape[:2], dtype=jnp.int32)` (used internally by the kernel to derive per-rank cu_seqlens) + +### Plan 4: CP and load_balance Mutual Exclusion + +The Delta Rule's recurrent state `S_t = f(S_{t-1}, k_t, v_t, beta_t)` depends on strict token ordering. load_balance's DUAL_CHUNK_SWAP reorder scrambles token order, breaking the sequential dependency. + +Runtime check (added at the `__call__` entry of `attention_kda.py`): + +```python +if (getattr(cfg, "context_parallel_size", 1) > 1 + and getattr(cfg, "context_parallel_load_balance", False)): + raise ValueError( + "KDA CP does not support context_parallel_load_balance. " + "Recurrent state S depends on exact token order; DUAL_CHUNK_SWAP " + "reorder breaks the sequential dependency. Set " + "context_parallel_load_balance=false when using KDA with CP." + ) +``` + +## segment_ids Data Flow + +``` +batch["inputs_segmentation"] ← [B, T], seg=0 = padding + │ + ▼ +KimiDeltaAttention.__call__(decoder_segment_ids) + │ + ├── chunk_size padding: pad to chunk size (64) multiple + │ + ├── ShortConvolution: halo_exchange_for_conv(segment_ids) + │ cross-segment boundary masking inside conv + │ + ├── _inject_context_on_T + _wsc: inject "context" sharding + │ + └── shard_map(chunk_kda): + - real seg → pass chunk_kda(segment_ids=seg) + - no seg + CP → pass dummy jnp.ones +``` + +## Files Changed + +| File | Change | Lines | +|------|--------|:----:| +| `layers/attention_kda.py` | **New**: `KimiDeltaAttention`, `ShortConvolution`, CP support | ~586 | +| `kernels/kda/__init__.py` | **New**: `chunk_kda()` entry point | ~84 | +| `kernels/kda/tokamax.py` | **New**: tokamax backend adapter | ~99 | +| `utils/cp_utils.py` | **New**: `halo_exchange_for_conv` | ~66 | +| `configs/types.py` | **Modified**: `KdaAttention` config class | +48 | +| `tests/unit/kda_attention_test.py` | **New**: KDA layer + conv halo + CP equivalence tests | ~914 | +| `docs/design/kda_cp_support.md` | **New**: design doc | — | +| **Total** | | **~1979** | + +## Key Constraints + +1. **CPContext availability**: Raise `ImportError` with a clear message when CPContext is unavailable; do not silently fall back. + +2. **ShortConvolution halo shard_map is required**: Under CP, conv needs to read historical tokens across ranks. Without shard_map → each rank independently left-zero-pads → causal sequence is split into independent segments → **correctness bug**. Without CP, falls back to `jnp.pad`, zero overhead. + +3. **conv and chunk_kda are two independent shard_maps**: Non-nested. conv only needs `ppermute`; chunk_kda needs `CPContext`. Separate shard_maps give independent XLA boundaries with resource release in between. + +4. **KDA does not use the `apply_attention` dispatcher**: KDA has its own QKV projection + SiLU + L2Norm + beta/gate projections and does not share the interface with `AttentionOp`. + +5. **CP + load_balance are mutually exclusive**: Recurrent state sequential dependency is irreversible. Runtime `ValueError`. + +## Backward Compatibility + +- `halo_exchange_for_conv`: degrades to `jnp.pad` when no CP, zero overhead +- ShortConv shard_map: only activated when `context_parallel_size > 1` +- CPContext import: `try/except`, raise `ImportError` with clear message if unavailable +- segment_ids dummy: auto-construct `jnp.ones` when no varlen + CP + +## Test Plan + +| Test | Coverage | +|------|----------| +| `test_short_conv_no_cp` | halo degrades to causal pad without CP | +| `test_short_conv_cp_halo` | conv under CP>1 equals single-rank reference | +| `test_kda_cp_equivalence` | CP multi-rank forward equals single-rank | +| `test_kda_cp_rejects_load_balance` | CP+load_balance raises ValueError | \ No newline at end of file diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index f8f1c7e021..6f8f3b71df 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -690,6 +690,55 @@ class CompressedAttention(BaseModel): ) +class KdaAttention(BaseModel): + """KDA (Kimi Delta Attention) configuration. + + These fields are placed in a separate class from MlaAttention for clear responsibility separation. + """ + + linear_conv_kernel_dim: int = Field( + 4, + ge=0, + description=( + "Convolution kernel dimension for linear attention layers (KDA). " + "This specifies the size of the 1D convolution applied to keys for local dependency modeling. " + "Default 4 matches the reference Megatron implementation." + ), + ) + use_kda_lora: bool = Field( + False, + description=( + "Whether to use LoRA (Low-Rank Adaptation) style decomposition in KDA layers. " + "When True, uses low-rank factorization for KDA computation. " + "When False, uses full-rank projections. " + "Default matches the reference Megatron implementation." + ), + ) + use_kda_safe_gate: bool = Field( + False, + description=( + "Whether to use numerically safe gate computation in KDA layers. " + "When True, applies value clamping and safe operations to prevent gate value explosion " + "during training." + ), + ) + kda_lower_bound: float = Field( + 0.0, + description=( + "Lower bound for gate values in KDA layers. Prevents gate values from " + "becoming too small (highly negative) during training, which can cause numerical instability. " + "Default 0.0 means no lower bound; -5.0 is a common choice." + ), + ) + + @field_validator("kda_lower_bound") + @classmethod + def _check_kda_lower_bound_finite(cls, v: float) -> float: + if not math.isfinite(v): + raise ValueError(f"kda_lower_bound must be finite, got {v}") + return v + + class AttentionIndexer(BaseModel): """Configuration for DeepSeek Sparse Attention (DSA): DeepSeek3.2-style MLA with indexer.""" @@ -2561,6 +2610,7 @@ class MaxTextConfig( # Attention Mechanisms Attention, MlaAttention, + KdaAttention, CompressedAttention, MoBa, AttentionIndexer, diff --git a/src/maxtext/kernels/kda/__init__.py b/src/maxtext/kernels/kda/__init__.py new file mode 100644 index 0000000000..9ff9dcb82f --- /dev/null +++ b/src/maxtext/kernels/kda/__init__.py @@ -0,0 +1,89 @@ +"""KDA (Kimi Delta Attention) kernels. + +Entry point that delegates to tokamax ``kimi_delta_attention`` with +native ``[B, T]`` segment_ids (head-first layout internally). + +Supports CP (context parallelism) via ``cp_context``. +""" + +from __future__ import annotations + +import jax.numpy as jnp +from maxtext.kernels.kda.tokamax import tokamax_chunk_kda + + +def chunk_kda( + q: jnp.ndarray, + k: jnp.ndarray, + v: jnp.ndarray, + g: jnp.ndarray, + beta: jnp.ndarray, + scale: float | None = None, + initial_state: jnp.ndarray | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + A_log: jnp.ndarray | None = None, + dt_bias: jnp.ndarray | None = None, + use_gate_in_kernel: bool = False, + use_qk_l2norm_in_kernel: bool = False, + safe_gate: bool = False, + lower_bound: float | None = None, + segment_ids: jnp.ndarray | None = None, + disable_recompute: bool = False, + N_max: int | None = None, + cp_context: object | None = None, +) -> tuple[jnp.ndarray, jnp.ndarray | None]: + """KDA entry point via tokamax backend. + + Tokamax natively accepts ``[B, T]`` segment_ids so no B*T flatten + or per-batch offset computation is needed. + + Args: + q: [B, T, H, K] queries. + k: [B, T, H, K] keys. + v: [B, T, H, V] values. + g: [B, T, H, K] gate values. + beta: [B, T, H] delta rule mixing coefficient. + scale: attention scale (default 1/sqrt(K)). + initial_state: must be None. + output_final_state: must be False. + chunk_size: chunk size (64). + A_log: [H] learnable decay in log space. + dt_bias: [H*K] dt bias. + use_gate_in_kernel: apply gate inside kernel. + use_qk_l2norm_in_kernel: apply L2 norm to q/k in kernel. + safe_gate: numerically safe gate mode. + lower_bound: gate value lower bound. + segment_ids: [B, T] segment IDs for varlen mode (2D, 1-based, 0=padding). + N_max: max segments per sample. + cp_context: Optional ``CPContext`` for CP. When set, the + kernel derives cross-rank metadata from ``segment_ids`` + and coordinates recurrent state across CP ranks. + + Returns: + (o, final_state) where o is [B, T, H, V] and final_state is None. + """ + if initial_state is not None: + raise NotImplementedError("initial_state is not supported") + if output_final_state: + raise NotImplementedError("output_final_state is not supported") + + return tokamax_chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + chunk_size=chunk_size, + A_log=A_log, + dt_bias=dt_bias, + use_gate_in_kernel=use_gate_in_kernel, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + safe_gate=safe_gate, + lower_bound=lower_bound, + segment_ids=segment_ids, + disable_recompute=disable_recompute, + N_max=N_max, + cp_context=cp_context, + ) diff --git a/src/maxtext/kernels/kda/tokamax.py b/src/maxtext/kernels/kda/tokamax.py new file mode 100644 index 0000000000..9c7c43d4c3 --- /dev/null +++ b/src/maxtext/kernels/kda/tokamax.py @@ -0,0 +1,101 @@ +"""Tokamax KDA backend for maxtext. + +Wraps ``tokamax.kimi_delta_attention`` with a maxtext-compatible interface +(batch-first ``[B, T, H, K]`` layout). Tokamax natively supports +``[B, T]`` segment_ids, so no +B*T flatten / per-batch offset is needed. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp + + +def _to_tokamax(q, k, v, g, beta): + """[B, T, H, K] -> [H, B, T, K]; [B, T, H] -> [H, B, T].""" + return ( + jnp.transpose(q, (2, 0, 1, 3)), + jnp.transpose(k, (2, 0, 1, 3)), + jnp.transpose(v, (2, 0, 1, 3)), + jnp.transpose(g, (2, 0, 1, 3)), + jnp.transpose(beta, (2, 0, 1)), + ) + + +def _to_maxtext(o_h): + """[H, B, T, V] -> [B, T, H, V].""" + return jnp.transpose(o_h, (1, 2, 0, 3)) + + +def tokamax_chunk_kda( + q: jnp.ndarray, + k: jnp.ndarray, + v: jnp.ndarray, + g: jnp.ndarray, + beta: jnp.ndarray, + scale: float | None = None, + initial_state: jnp.ndarray | None = None, + output_final_state: bool = False, + chunk_size: int = 64, + A_log: jnp.ndarray | None = None, + dt_bias: jnp.ndarray | None = None, + use_gate_in_kernel: bool = False, + use_qk_l2norm_in_kernel: bool = False, + safe_gate: bool = False, + lower_bound: float | None = None, + segment_ids: jnp.ndarray | None = None, + disable_recompute: bool = False, + N_max: int | None = None, + cp_context: object | None = None, +) -> tuple[jnp.ndarray, jnp.ndarray | None]: + """KDA via tokamax, batch-first interface matching ``kernels.kda.chunk_kda``. + + Tokamax accepts ``[B, T]`` segment_ids natively so no B*T flatten + or per-batch offset computation is needed. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + g: [B, T, H, K] + beta: [B, T, H] + segment_ids: [B, T] 1-based, 0=padding. Passed directly to tokamax. + cp_context: Optional ``tokamax...CPContext`` for context parallelism. + (other args match the backend interface) + + Returns: + (o, None) where o is [B, T, H, V]. + """ + from tokamax._src.ops.experimental.kda.api import kimi_delta_attention + + if initial_state is not None: + raise NotImplementedError("initial_state is not supported with tokamax backend") + if output_final_state: + raise NotImplementedError("output_final_state is not supported with tokamax backend") + + q_h, k_h, v_h, g_h, beta_h = _to_tokamax(q, k, v, g, beta) + + o_h, _final_state = kimi_delta_attention( + q=q_h, + k=k_h, + v=v_h, + g=g_h, + beta=beta_h, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + segment_ids=segment_ids, + chunk_size=chunk_size, + use_gate_in_kernel=use_gate_in_kernel, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + safe_gate=safe_gate, + lower_bound=lower_bound, + disable_recompute=disable_recompute, + N_max=N_max, + implementation="pallas_tpu", # "pallas_tpu" is tokamax's internal TPU kernel implementation name (see tokamax._src.ops.experimental.kda.api.Implementation) + cp_context=cp_context, + ) + + o = _to_maxtext(o_h) + return o, None diff --git a/src/maxtext/layers/attention_kda.py b/src/maxtext/layers/attention_kda.py new file mode 100644 index 0000000000..1ee36a8ae9 --- /dev/null +++ b/src/maxtext/layers/attention_kda.py @@ -0,0 +1,595 @@ +# Copyright 2023–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. + +"""Kimi Delta Attention (KDA) Layer Implementation. + +This module implements the KDA (Kimi Delta Attention) layer. KDA is a linear attention mechanism with Delta Rule correction, +featuring: + - Depthwise causal 1D convolution for local dependency modeling + - Numerically safe gate mechanism + - Q/K L2 normalization + +The implementation delegates to ``tokamax.kimi_delta_attention`` for the +chunk-parallel Delta Rule computation. +""" + + +import functools +import math + +from flax import nnx +import jax +import jax.numpy as jnp +from jax.ad_checkpoint import checkpoint_name +from jax.sharding import Mesh +from maxtext.kernels.kda import chunk_kda + +# KDA depends on tokamax at runtime, but import should succeed because +# tokamax is a mandatory dependency for KDA models. +try: + from tokamax._src.ops.experimental.kda.cp_utils import CPContext as TokamaxCPContext +except ImportError: + TokamaxCPContext = None + +from maxtext.common.common_types import Config, MODEL_MODE_AUTOREGRESSIVE +from maxtext.layers import linears +from maxtext.layers.normalizations import RMSNorm +from maxtext.utils.sharding import logical_to_mesh_axes +from maxtext.utils.cp_utils import halo_exchange_for_conv + + +# Chunk size for KDA kernel (matching Megatron convention). +_KDA_CHUNK_SIZE = 64 + + +def _l2_normalize(x, axis=-1, eps=1e-6): + x_f = x.astype(jnp.float32) + rstd = jax.lax.rsqrt(jnp.sum(x_f * x_f, axis=axis, keepdims=True) + eps) + return (x_f * rstd).astype(x.dtype) + + +class ShortConvolution(nnx.Module): + """Depthwise causal 1D convolution for local dependency modeling in KDA. + + Each channel is convolved independently (no cross-channel mixing), + matching Megatron's Conv1d with groups=in_channels. Position i can + only attend to positions <= i (causal). When segment_ids is provided, + cross-segment contributions are masked to prevent leakage across + document boundaries (matches Megatron causal_conv1d_fn seq_idx). + """ + + def __init__( + self, + kernel_size: int, + features: int, + *, + dtype: jnp.dtype = jnp.bfloat16, + weight_dtype: jnp.dtype = jnp.bfloat16, + rngs: nnx.Rngs, + ): + self.kernel_size = kernel_size + self.features = features + self.dtype = dtype + + self.kernel = nnx.Param( + nnx.initializers.lecun_normal()( + rngs.params(), + (kernel_size, features), + weight_dtype, + ) + ) + + def __call__(self, x: jnp.ndarray, segment_ids: jnp.ndarray | None = None) -> jnp.ndarray: + B, T, F = x.shape + if F != self.features: + raise ValueError(f"Input features {F} != {self.features}") + + x_padded = halo_exchange_for_conv(x, self.kernel_size - 1, axis_name="context") + + if segment_ids is not None: + seg_padded = halo_exchange_for_conv(segment_ids, self.kernel_size - 1, axis_name="context") + # Stack per-tap masks once so the loop body has no tap-dependent + # broadcasts beyond the slice itself. + masks = [ + (seg_padded[:, k : k + T] == segment_ids).astype(x.dtype)[:, :, None] + for k in range(self.kernel_size - 1, -1, -1) + ] + + output = jnp.zeros((B, T, F), dtype=x.dtype) + for k in range(self.kernel_size): + offset = self.kernel_size - 1 - k + x_slice = x_padded[:, offset : offset + T, :] + if segment_ids is not None: + x_slice = x_slice * masks[k] + output = output + x_slice * self.kernel[k] + + return output.astype(self.dtype) + + +class KimiDeltaAttention(nnx.Module): + """Kimi Delta Attention (KDA) layer implementation. + + KDA is a linear attention mechanism that uses the Delta Rule for state + correction: + S' = S * exp(g_t) + residual = v_t - k_t^T @ S' + S = S' + beta_t * k_t (x) residual + o_t = scale * q_t^T @ S + + Attributes: + config: Model configuration containing KDA parameters. + layer_idx: Index of this layer in the decoder stack. + mesh: JAX device mesh for sharding. + """ + + def __init__( + self, + config: Config, + layer_idx: int, + mesh: Mesh, + *, + rngs: nnx.Rngs, + ): + self.config = config + self.layer_idx = layer_idx + self.mesh = mesh + + cfg = self.config + + # KDA head dimensions derived from global config (matching Megatron convention): + # key_head_dim = value_head_dim = config.head_dim (kv_channels) + # num_key_heads = num_value_heads = config.base_num_query_heads (num_attention_heads) + self.key_head_dim = cfg.head_dim + self.value_head_dim = cfg.head_dim + self.num_key_heads = cfg.base_num_query_heads + self.num_value_heads = cfg.base_num_query_heads + self.num_query_heads = self.num_key_heads + + # Short convolution for local dependency modeling + if cfg.linear_conv_kernel_dim > 0: + self.q_conv = ShortConvolution( + kernel_size=cfg.linear_conv_kernel_dim, + features=self.num_query_heads * self.key_head_dim, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + rngs=rngs, + ) + self.k_conv = ShortConvolution( + kernel_size=cfg.linear_conv_kernel_dim, + features=self.num_key_heads * self.key_head_dim, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + rngs=rngs, + ) + self.v_conv = ShortConvolution( + kernel_size=cfg.linear_conv_kernel_dim, + features=self.num_value_heads * self.value_head_dim, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + rngs=rngs, + ) + else: + self.q_conv = None + self.k_conv = None + self.v_conv = None + + # QKV projections + # Separate projections for Q, K, V (not fused) to allow independent conv + self.q_proj = linears.DenseGeneral( + in_features_shape=cfg.base_emb_dim, + out_features_shape=(self.num_query_heads, self.key_head_dim), + axis=-1, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=("embed", "heads", "kv"), + use_bias=cfg.attention_bias, + shard_mode=cfg.shard_mode, + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + + self.k_proj = linears.DenseGeneral( + in_features_shape=cfg.base_emb_dim, + out_features_shape=(self.num_key_heads, self.key_head_dim), + axis=-1, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=("embed", "heads", "kv"), + use_bias=cfg.attention_bias, + shard_mode=cfg.shard_mode, + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + + self.v_proj = linears.DenseGeneral( + in_features_shape=cfg.base_emb_dim, + out_features_shape=(self.num_value_heads, self.value_head_dim), + axis=-1, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=("embed", "heads", "kv"), + use_bias=cfg.attention_bias, + shard_mode=cfg.shard_mode, + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + + # Output projection + self.o_proj = linears.DenseGeneral( + in_features_shape=(self.num_value_heads, self.value_head_dim), + out_features_shape=cfg.base_emb_dim, + axis=(-2, -1), + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=("heads", "kv", "embed"), + use_bias=cfg.attention_bias, + shard_mode=cfg.shard_mode, + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + + # Gate projection for generating g (log-space gate) + # g has shape [B, T, H, K] - per-head, per-dim gate + self.g_proj = linears.DenseGeneral( + in_features_shape=cfg.base_emb_dim, + out_features_shape=(self.num_key_heads, self.key_head_dim), + axis=-1, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=("embed", "heads", "kv"), + use_bias=False, + shard_mode=cfg.shard_mode, + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + + # Beta projection for generating beta (Delta rule mixing coefficient) + # beta has shape [B, T, H] - per-head scalar + self.b_proj = linears.DenseGeneral( + in_features_shape=cfg.base_emb_dim, + out_features_shape=(self.num_key_heads,), + axis=-1, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=("embed", "heads"), + use_bias=False, + shard_mode=cfg.shard_mode, + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + + # Q/K L2 normalization is applied outside chunk_kda (matching Megatron) + + # Output gate projection: gate shape [B, T, H, V] (matching Megatron no_kda_lora path) + self.gate_proj = linears.DenseGeneral( + in_features_shape=cfg.base_emb_dim, + out_features_shape=(self.num_value_heads, self.value_head_dim), + axis=-1, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + kernel_axes=("embed", "heads", "kv"), + use_bias=cfg.attention_bias, + shard_mode=cfg.shard_mode, + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + + # Output norm (per-head RMSNorm, applied before gating) + self.out_norm = RMSNorm( + num_features=self.value_head_dim, + epsilon=cfg.normalization_layer_epsilon, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + rngs=rngs, + ) + + # Gate parameters (matching Megatron kda.py:299-330) + # A_log: [num_key_heads] — log of diagonal decay matrix + A_init_range = (1.0, 16.0) + A = jax.random.uniform( + rngs.params(), + shape=(self.num_key_heads,), + minval=A_init_range[0], + maxval=A_init_range[1], + ) + self.A_log = nnx.Param(jnp.log(A)) + + # dt_bias: [num_key_heads * key_head_dim] — gate bias + # Initialize via inverse softplus of uniform(dt_min, dt_max) + dt_min, dt_max, dt_init_floor = 0.001, 0.1, 1e-4 + dt = jnp.exp( + jax.random.uniform( + rngs.params(), + shape=(self.num_key_heads * self.key_head_dim,), + ) + * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min) + ) + dt = jnp.clip(dt, min=dt_init_floor) + # Inverse softplus: x = dt + log(-expm1(-dt)) + inv_dt = dt + jnp.log(-jnp.expm1(-dt)) + self.dt_bias = nnx.Param(inv_dt) + + # Axis names for shard_map (tokamax kernels cannot be auto-partitioned). + self.qkv_axis_names = ("activation_batch", "activation_norm_length", "activation_heads", "activation_kv") + self.beta_axis_names = ("activation_batch", "activation_norm_length", "activation_heads") + + def _logical_to_mesh_axes(self, logical_name): + return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=self.config.logical_axis_rules) + + def __call__( + self, + hidden_states: jnp.ndarray, + decoder_positions: jnp.ndarray | None = None, + deterministic: bool = True, + model_mode: str = "train", + *, + layer_idx: int | None = None, + decoder_segment_ids: jnp.ndarray | None = None, + ) -> tuple[jnp.ndarray, None]: + """Forward pass for KDA attention. + + Args: + hidden_states: Input tensor of shape [B, T, emb_dim]. + decoder_positions: Position indices for RoPE (not used in KDA). + deterministic: Whether to use deterministic mode. + model_mode: Model mode (train/prefill/autoregressive). + layer_idx: Optional layer index override. + decoder_segment_ids: Optional segment IDs for packed sequences. + + Returns: + Tuple of (output, None) where output has shape [B, T, emb_dim]. + """ + del decoder_positions # KDA doesn't use RoPE + del deterministic # No dropout in KDA currently + del layer_idx # Not used + + cfg = self.config + + if model_mode == MODEL_MODE_AUTOREGRESSIVE: + raise NotImplementedError("KDA autoregressive mode not yet implemented.") + + B, T_orig, _ = hidden_states.shape + T = T_orig + + if T % _KDA_CHUNK_SIZE != 0: + pad_len = _KDA_CHUNK_SIZE - (T % _KDA_CHUNK_SIZE) + hidden_states = jnp.pad(hidden_states, ((0, 0), (0, pad_len), (0, 0))) + if decoder_segment_ids is not None: + decoder_segment_ids = jnp.pad(decoder_segment_ids, ((0, 0), (0, pad_len)), constant_values=0) + T = hidden_states.shape[1] + + # QKV projections + with jax.named_scope("qkv_proj"): + q = self.q_proj(hidden_states) # [B, T, H, K] + k = self.k_proj(hidden_states) # [B, T, H, K] + v = self.v_proj(hidden_states) # [B, T, H, V] + + # Names must match decoders.minimal_policy so remat policies save these. + q = checkpoint_name(q, "query_proj") + k = checkpoint_name(k, "key_proj") + v = checkpoint_name(v, "value_proj") + + # Apply short convolution if enabled (before activation, matching Megatron) + if self.q_conv is not None: + with jax.named_scope("short_conv"): + # Reshape for conv: [B, T, H*D] -> conv -> [B, T, H*D] -> reshape back + q_flat = q.reshape(B, T, -1) + k_flat = k.reshape(B, T, -1) + v_flat = v.reshape(B, T, -1) + + # Under CP, ShortConvolution needs to pull kernel_size-1 tokens + # of left context from the previous CP rank via ppermute — which + # is a collective and so must run inside a shard_map that exposes + # the "context" mesh axis. Without this wrap, halo_exchange_for_conv + # would silently degrade to zero-pad (causal-conv at every CP shard + # boundary would be wrong). This applies to all CP strategies. + if getattr(cfg, "context_parallel_size", 1) > 1: + conv_flat_pspec = self._logical_to_mesh_axes(("activation_batch", "activation_norm_length", None)) + conv_seg_pspec = ( + self._logical_to_mesh_axes(("activation_batch", "activation_norm_length")) + if decoder_segment_ids is not None + else None + ) + q_conv_mod, k_conv_mod, v_conv_mod = self.q_conv, self.k_conv, self.v_conv + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=(conv_flat_pspec, conv_flat_pspec, conv_flat_pspec, conv_seg_pspec), + out_specs=(conv_flat_pspec, conv_flat_pspec, conv_flat_pspec), + check_vma=False, + ) + def _conv_with_halo(qf, kf, vf, seg): + qf = q_conv_mod(qf, segment_ids=seg) + kf = k_conv_mod(kf, segment_ids=seg) + vf = v_conv_mod(vf, segment_ids=seg) + return qf, kf, vf + + q_flat, k_flat, v_flat = _conv_with_halo(q_flat, k_flat, v_flat, decoder_segment_ids) + else: + q_flat = self.q_conv(q_flat, segment_ids=decoder_segment_ids) + k_flat = self.k_conv(k_flat, segment_ids=decoder_segment_ids) + v_flat = self.v_conv(v_flat, segment_ids=decoder_segment_ids) + + q = q_flat.reshape(B, T, self.num_query_heads, self.key_head_dim) + k = k_flat.reshape(B, T, self.num_key_heads, self.key_head_dim) + v = v_flat.reshape(B, T, self.num_value_heads, self.value_head_dim) + + # Apply SiLU activation after conv on q, k, v (matching Megatron) + q = jax.nn.silu(q) + k = jax.nn.silu(k) + v = jax.nn.silu(v) + + # Apply L2 normalization to Q/K outside the kernel (matching Megatron kda.py:824-828) + if cfg.use_qk_norm: + q = _l2_normalize(q) + k = _l2_normalize(k) + + # Generate gate g (raw projection, gate transform done inside kernel) + with jax.named_scope("gate_proj"): + g = self.g_proj(hidden_states) # [B, T, H, K] + + # Generate output gate (for gated norm after KDA kernel) + with jax.named_scope("output_gate_proj"): + output_gate = self.gate_proj(hidden_states) # [B, T, H, V] + + # Generate beta (Delta rule mixing coefficient) + with jax.named_scope("beta_proj"): + beta = self.b_proj(hidden_states) # [B, T, H] + beta = beta.astype(jnp.float32) + beta = jax.nn.sigmoid(beta) # Ensure (0, 1) range, in fp32 + + scale = self.key_head_dim**-0.5 + safe_gate = cfg.use_kda_safe_gate + lower_bound = cfg.kda_lower_bound if safe_gate else None + if not safe_gate and cfg.kda_lower_bound != 0.0: + import warnings + + warnings.warn( + f"kda_lower_bound={cfg.kda_lower_bound} is ignored because use_kda_safe_gate=False. " + "Set use_kda_safe_gate=True to enable lower_bound clamping.", + stacklevel=2, + ) + n_max = cfg.packing_max_segments_per_sample if cfg.packing_max_segments_per_sample > 0 else None + + # KDA Delta Rule relies on sequential recurrent state S_t = f(S_{t-1}, ...). + # load_balance's DUAL_CHUNK_SWAP reorder breaks token order, invalidating + # the sequential dependency. Reject this combination at runtime. + if getattr(cfg, "context_parallel_size", 1) > 1 and getattr(cfg, "context_parallel_load_balance", False): + raise ValueError( + "KDA CP does not support context_parallel_load_balance. " + "Recurrent state S depends on exact token order; DUAL_CHUNK_SWAP " + "reorder breaks the sequential dependency. Set " + "context_parallel_load_balance=false when using KDA with CP." + ) + + # Call KDA kernel via shard_map (tokamax kernels cannot be auto-partitioned). + with jax.named_scope("kda_kernel"): + qkv_pspec = self._logical_to_mesh_axes(self.qkv_axis_names) + beta_pspec = self._logical_to_mesh_axes(self.beta_axis_names) + a_log_pspec = self._logical_to_mesh_axes(("activation_heads",)) + dt_bias_2d_pspec = self._logical_to_mesh_axes(("activation_heads", "activation_kv")) + seg_pspec = self._logical_to_mesh_axes(("activation_batch", "activation_norm_length")) + + # Reshape dt_bias from [H*K] to [H, K] for proper head-dim sharding. + dt_bias_2d = self.dt_bias.value.reshape(self.num_key_heads, self.key_head_dim) + + # Inject "context" on the T axis so the shard_map sees the correct + # per-rank shard layout. logical_to_mesh_axes may map the LENGTH + # axis to None (size-1 stripping / Flax rule precedence) — we + # explicitly set it to "context" when CP is active. + def _inject_context_on_T(pspec, t_axis=1): + spec = list(pspec) + if spec[t_axis] is None: + spec[t_axis] = "context" + return jax.sharding.PartitionSpec(*spec) + + if getattr(cfg, "context_parallel_size", 1) > 1: + qkv_pspec = _inject_context_on_T(qkv_pspec) + beta_pspec = _inject_context_on_T(beta_pspec) + seg_pspec = _inject_context_on_T(seg_pspec) + + def _wsc(x, pspec): + return jax.lax.with_sharding_constraint(x, jax.sharding.NamedSharding(self.mesh, pspec)) + + q, k, v, g = _wsc(q, qkv_pspec), _wsc(k, qkv_pspec), _wsc(v, qkv_pspec), _wsc(g, qkv_pspec) + beta = _wsc(beta, beta_pspec) + if decoder_segment_ids is not None: + decoder_segment_ids = _wsc(decoder_segment_ids, seg_pspec) + + # CP: tokamax kernel derives CP metadata from segment_ids. + # Always pass a seg arg when CP is active so the kernel can + # compute cu_seqlens / chain fields via one small all_gather. + has_seg = decoder_segment_ids is not None or getattr(cfg, "context_parallel_size", 1) > 1 + base_in_specs = (qkv_pspec, qkv_pspec, qkv_pspec, qkv_pspec, beta_pspec, a_log_pspec, dt_bias_2d_pspec) + in_specs = base_in_specs + ((seg_pspec,) if has_seg else ()) + + # CPContext lives outside shard_map — it is a frozen dataclass that + # holds the mesh identity. tokamax's chunk_kda derives the per-rank + # chain fields (cu_seqlens, is_first_rank, …) internally from + # segment_ids, then passes the completed context to the kernel. + cp_ctx = None + if getattr(cfg, "context_parallel_size", 1) > 1: + if TokamaxCPContext is None: + raise ImportError( + "KDA context parallelism requires " + "tokamax._src.ops.experimental.kda.cp_utils.CPContext, " + "but it failed to import. Refusing to run: CP would silently " + "break recurrent state across ranks." + ) + cp_ctx = TokamaxCPContext(mesh=self.mesh, axis_name="context") + + @functools.partial(jax.shard_map, mesh=self.mesh, in_specs=in_specs, out_specs=qkv_pspec, check_vma=False) + def _shard_map_chunk_kda(*args): + q, k, v, g, beta, A_log, dt_bias_2d, *rest = args + seg = rest[0] if rest else None + + # CP: provide a dummy seg (all-ones) so the kernel has + # segment_ids to derive cu_seqlens from, even when the user + # hasn't supplied real segmentation info. + if seg is None and getattr(cfg, "context_parallel_size", 1) > 1: + seg = jnp.ones(q.shape[:2], dtype=jnp.int32) + + dt_bias_flat = dt_bias_2d.reshape(-1) + o, _ = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias_flat, + segment_ids=seg, + scale=scale, + chunk_size=_KDA_CHUNK_SIZE, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=True, + safe_gate=safe_gate, + lower_bound=lower_bound, + disable_recompute=True, + N_max=n_max, + cp_context=cp_ctx, + ) + return o + + kda_args = (q, k, v, g, beta, self.A_log.value, dt_bias_2d) + if has_seg: + if decoder_segment_ids is not None: + kda_args = kda_args + (decoder_segment_ids,) + else: + # CP without varlen: pass None; the shard_map function + # synthesises a dummy seg internally. + kda_args = kda_args + (None,) + o = _shard_map_chunk_kda(*kda_args) + + # Analogous to MLA's `context` (see attention_op.py); a remat boundary + # right after the KDA kernel so its result survives `minimal_with_context`. + o = checkpoint_name(o, "context") + + # Output gated norm (matching Megatron _apply_gated_norm): + # per-head RMSNorm over the value dim, then sigmoid gate. + with jax.named_scope("output_gated_norm"): + o_dtype = o.dtype + o_normed = self.out_norm(o) + o = (o_normed * jax.nn.sigmoid(output_gate.astype(jnp.float32))).astype(o_dtype) + + # Output projection + with jax.named_scope("o_proj"): + output = self.o_proj(o) + output = checkpoint_name(output, "out_proj") + + return output[:, :T_orig, :], None diff --git a/src/maxtext/utils/cp_utils.py b/src/maxtext/utils/cp_utils.py new file mode 100644 index 0000000000..33a764ac08 --- /dev/null +++ b/src/maxtext/utils/cp_utils.py @@ -0,0 +1,65 @@ +"""Context Parallelism utilities — halo exchange for ShortConvolution.""" + +import jax +import jax.numpy as jnp + + +def _has_named_axis(axis_name: str) -> bool: + """Check whether *axis_name* is bound in the current shard_map / mesh scope.""" + try: + jax.lax.axis_index(axis_name) + return True + except NameError: + return False + + +def halo_exchange_for_conv( + x: jax.Array, + halo_size: int, + axis_name: str = "context", + seq_axis: int = 1, +) -> jax.Array: + """Prepend ``halo_size`` tokens from the previous CP rank for causal conv. + + The caller (ShortConvolution) receives ``[halo_size + T_local, …]`` so the + per-tap loop naturally reads the correct context window. Halos are fetched + via a forward-ring ``ppermute``: rank *i* sends its last ``halo_size`` + tokens to rank *i+1*; rank 0 receives zeros (sequence start). + + When no CP axis is in scope or ``cp_size == 1`` the function degrades to + left zero-padding, which is the correct causal-convolution boundary for a + single-device / no-CP run. + + Args: + x: Tensor shaped ``[B, T, …]`` (seq_axis = 1). + halo_size: Number of tokens to pull from the previous rank. + axis_name: Mesh axis along which the sequence is sharded. + seq_axis: The sequence dimension index (default 1). + + Returns: + ``x`` with ``halo_size`` context tokens prepended along *seq_axis*. + """ + if halo_size <= 0: + return x + + # Left zero-pad — works correctly for both no-CP and CP. + pad_width = [(0, 0)] * x.ndim + pad_width[seq_axis] = (halo_size, 0) + zero_padded = jnp.pad(x, pad_width) + + if not _has_named_axis(axis_name): + return zero_padded + + cp_size = jax.lax.psum(1, axis_name=axis_name) + if cp_size == 1: + return zero_padded + + # Forward ring: each rank sends its tail to the next rank. + tail = jax.lax.dynamic_slice_in_dim(x, x.shape[seq_axis] - halo_size, halo_size, axis=seq_axis) + perm = [(i, (i + 1) % cp_size) for i in range(cp_size)] + halo = jax.lax.ppermute(tail, axis_name=axis_name, perm=perm) + + cp_rank = jax.lax.axis_index(axis_name) + halo = jnp.where(cp_rank == 0, jnp.zeros_like(halo), halo) + + return jnp.concatenate([halo, x], axis=seq_axis) diff --git a/tests/unit/kda_attention_test.py b/tests/unit/kda_attention_test.py new file mode 100644 index 0000000000..5fe9781c70 --- /dev/null +++ b/tests/unit/kda_attention_test.py @@ -0,0 +1,926 @@ +"""Unit tests for KDA (Kimi Delta Attention) module. + +Tests cover: + - KimiDeltaAttention: initialization, forward pass, padding, determinism + - chunk_kda kernel: basic operation, chunk vs recurrent comparison + - Naive KDA: recurrent Delta Rule reference impl vs kernel precision + - Backward (VJP): activation gradients, weight gradients, determinism, bf16 + - QK L2 norm: applied outside kernel, matching Megatron + - Real config: loading a KDA model config and running KimiDeltaAttention + +Precision comparison uses _assert_close (atol+rtol+ULP fallback), +adapted from gla_compare_test.py. + +Run with: python -m pytest tests/unit/kda_attention_test.py -v +""" + +import functools + +import pytest +import jax +import jax.numpy as jnp +import numpy as np +import ml_dtypes +from flax import nnx + +try: + import tokamax + + TOKAMAX_AVAILABLE = True +except ImportError: + TOKAMAX_AVAILABLE = False + +from maxtext.layers import attention_kda + + +# --------------------------------------------------------------------------- +# Precision comparison utilities (adapted from gla_compare_test.py) +# --------------------------------------------------------------------------- + + +def _bf16_bits_to_ordered(u16): + magnitude = (u16 & 0x7FFF).astype(np.int64) + return np.where(u16 & 0x8000, -magnitude, magnitude) + + +def bf16_ulp_diff(actual_f32, expected_f32): + """Compute per-element ULP distance at bf16 precision.""" + a_u16 = np.ascontiguousarray(actual_f32.astype(ml_dtypes.bfloat16)).view(np.uint16) + b_u16 = np.ascontiguousarray(expected_f32.astype(ml_dtypes.bfloat16)).view(np.uint16) + mismatch_mask = a_u16 != b_u16 + n_mismatch = int(mismatch_mask.sum()) + n_total = a_u16.size + if n_mismatch == 0: + return n_mismatch, n_total, 0, np.array([], dtype=np.int64) + a_ordered = _bf16_bits_to_ordered(a_u16[mismatch_mask]) + b_ordered = _bf16_bits_to_ordered(b_u16[mismatch_mask]) + abs_ulp = np.abs(a_ordered - b_ordered) + return n_mismatch, n_total, int(abs_ulp.max()), abs_ulp + + +def _assert_close(actual, expected, label, atol=1e-2, rtol=1e-5, max_ulp=2, max_ulp_fail_rate=1e-3): + """Assert two arrays match via allclose with bf16 ULP diff fallback.""" + actual_f32 = np.asarray(actual, dtype=np.float32) + expected_f32 = np.asarray(expected, dtype=np.float32) + + diff = np.abs(actual_f32 - expected_f32) + max_abs = float(diff.max()) + mean_abs = float(diff.mean()) + print(f" {label}: max_abs={max_abs:.6e} mean_abs={mean_abs:.6e}") + + close_mask = diff <= atol + rtol * np.abs(expected_f32) + if close_mask.all(): + print(f" {label}: all close ({atol=}, {rtol=})") + return + + n_fail = int((~close_mask).sum()) + n_total = actual_f32.size + fail_actual = actual_f32[~close_mask] + fail_expected = expected_f32[~close_mask] + n_mis, _, worst_ulp, abs_ulps = bf16_ulp_diff(fail_actual, fail_expected) + + n_over = int((abs_ulps > max_ulp).sum()) if n_mis > 0 else 0 + over_rate = n_over / n_fail if n_fail > 0 else 0.0 + + if n_mis > 0: + print( + f" {label} ULP: {n_fail}/{n_total} fail allclose, " + f"{n_mis} have ULP diff, max_ulp={worst_ulp}, " + f"over {max_ulp} ULP: {n_over}/{n_fail} ({over_rate:.2e})" + ) + + assert over_rate <= max_ulp_fail_rate, ( + f"{label}: {n_over}/{n_fail} elements ({over_rate:.2e}) exceed " + f"{max_ulp} ULP (threshold {max_ulp_fail_rate:.2e})" + ) + + +class _MockKdaConfig: + """Minimal mock config for KDA testing. + + KDA derives head dims from global config (matching Megatron): + key_head_dim = value_head_dim = head_dim + num_key_heads = num_value_heads = base_num_query_heads + """ + + def __init__(self, **overrides): + self.base_emb_dim = 128 + self.base_num_query_heads = 4 + self.head_dim = 32 + self.dtype = jnp.float32 + self.weight_dtype = jnp.float32 + self.attention_bias = False + self.shard_mode = "auto" + self.matmul_precision = "default" + self.normalization_layer_epsilon = 1e-6 + self.logical_axis_rules = [] + + # KDA-specific + self.linear_conv_kernel_dim = 4 + self.use_qk_norm = True + self.use_kda_safe_gate = True + self.kda_lower_bound = -5.0 + self.packing_max_segments_per_sample = 25 + + for k, v in overrides.items(): + setattr(self, k, v) + + +# --------------------------------------------------------------------------- +# KimiDeltaAttention tests +# --------------------------------------------------------------------------- + + +class TestKimiDeltaAttention: + """Tests for KimiDeltaAttention module.""" + + @pytest.fixture + def mesh(self): + return jax.sharding.Mesh(jax.devices(), ("x",)) + + def _make_attn(self, mesh, **config_overrides): + cfg = _MockKdaConfig(**config_overrides) + rngs = nnx.Rngs(0) + with mesh: + return attention_kda.KimiDeltaAttention( + config=cfg, + layer_idx=0, + mesh=mesh, + rngs=rngs, + ) + + def test_init_head_dims(self, mesh): + """Head dims derived from global config: head_dim=32, base_num_query_heads=4.""" + attn = self._make_attn(mesh) + assert attn.num_query_heads == 4 + assert attn.num_key_heads == 4 + assert attn.num_value_heads == 4 + assert attn.key_head_dim == 32 + assert attn.value_head_dim == 32 + + def test_init_no_conv(self, mesh): + attn = self._make_attn(mesh, linear_conv_kernel_dim=0) + assert attn.q_conv is None + + def test_init_has_gate_and_norm(self, mesh): + """Output gate projection and out_norm should always be present.""" + attn = self._make_attn(mesh) + assert hasattr(attn, "gate_proj") + assert hasattr(attn, "out_norm") + assert hasattr(attn, "A_log") + assert hasattr(attn, "dt_bias") + + def test_forward_shape(self, mesh): + attn = self._make_attn(mesh) + B, T, D = 2, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + with mesh: + output, aux = attn(x) + assert output.shape == (B, T, D) + assert aux is None + + def test_forward_no_nan_inf(self, mesh): + attn = self._make_attn(mesh) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + with mesh: + output, _ = attn(x) + assert not jnp.any(jnp.isnan(output)) + assert not jnp.any(jnp.isinf(output)) + assert jnp.any(output != 0) + + def test_sequence_padding(self, mesh): + """Non-divisible sequence lengths should be handled via padding.""" + attn = self._make_attn(mesh) + B, T, D = 1, 100, 128 # 100 not divisible by chunk_size=64 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + with mesh: + output, _ = attn(x) + assert output.shape == (B, T, D) + + def test_deterministic(self, mesh): + attn = self._make_attn(mesh) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + with mesh: + o1, _ = attn(x) + o2, _ = attn(x) + assert jnp.allclose(o1, o2, atol=1e-5) + + def test_packed_sequences_supported(self, mesh): + """Test that KDA supports packed sequences with segment_ids.""" + attn = self._make_attn(mesh) + B, T, hidden_dim = 2, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, hidden_dim)) + # 1-based segment_ids, 0 = padding + seg_ids = jnp.array( + [[1, 1, 1, 2, 2, 2, 3, 3] + [0] * (T - 8), [1, 1, 2, 2, 2, 2, 3, 3] + [0] * (T - 8)], dtype=jnp.int32 + ) + with mesh: + o, _ = attn(x, decoder_segment_ids=seg_ids) + # Output shape should match input + assert o.shape == (B, T, hidden_dim) + # No NaN or Inf + assert jnp.isfinite(o).all() + + def test_segment_ids_padding_alignment(self, mesh): + """When T % 64 != 0, segment_ids should be padded along with hidden_states.""" + attn = self._make_attn(mesh) + B, T, hidden_dim = 1, 100, 128 # 100 not divisible by chunk_size=64 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, hidden_dim)) + # segment_ids shorter than padded length (100 -> 128 after pad) + seg_ids = jnp.array([[1, 1, 1, 2, 2, 2, 3, 3] + [0] * (T - 8)], dtype=jnp.int32) + with mesh: + o, _ = attn(x, decoder_segment_ids=seg_ids) + # Output shape should match input (unpadded back from 128 to 100) + assert o.shape == (B, T, hidden_dim) + # First 8 positions should have segment info, rest may be affected by padding + # but output should still be finite + assert jnp.isfinite(o).all() + + def test_segment_ids_none_fallback(self, mesh): + """Test that segment_ids=None falls back to legacy behavior.""" + attn = self._make_attn(mesh) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + with mesh: + o1, _ = attn(x, decoder_segment_ids=None) + o2, _ = attn(x) # Default None + assert jnp.allclose(o1, o2, atol=1e-5) + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_row_independence(self, mesh): + """Hard verification: row0 and row1 use different inputs; only change row1's seg, + assert row0 output unchanged. + + Construct batch=[row_a, row_b], only modify row_b's segment_ids, + assert row_a's output is bit-exact unchanged. Proves segment_ids-based structural + isolation is effective. + """ + attn = self._make_attn(mesh) + T, hidden_dim = 64, 128 + + # Critical: two rows use different inputs (prevents XLA caching optimization) + x0 = jax.random.normal(jax.random.PRNGKey(0), (1, T, hidden_dim)) + x1 = jax.random.normal(jax.random.PRNGKey(1), (1, T, hidden_dim)) + x = jnp.concatenate([x0, x1], axis=0) # [2, T, hidden_dim] + + # row0: fixed segment; row1: varying segment (keeping padding zeros identical) + seg_base = jnp.array([[1] * T, [1, 1, 2, 2, 2, 3, 3, 3] + [0] * (T - 8)], dtype=jnp.int32) + seg_modified = jnp.array([[1] * T, [1, 1, 2, 2, 2, 4, 4, 4] + [0] * (T - 8)], dtype=jnp.int32) + + with mesh: + o1, _ = attn(x, decoder_segment_ids=seg_base) + o2, _ = attn(x, decoder_segment_ids=seg_modified) + + # Hard verification: row0 output is bit-exact unchanged (atol=0 means strict equality) + assert jnp.allclose(o1[0], o2[0], atol=0.0), ( + "Row 0 changed when only row 1's segment changed; " "this indicates segment-based isolation violation" + ) + + def test_autoregressive_not_supported(self, mesh): + attn = self._make_attn(mesh) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + with pytest.raises(NotImplementedError, match="autoregressive"): + attn(x, model_mode="autoregressive") + + +# --------------------------------------------------------------------------- +# Kernel-level tests +# --------------------------------------------------------------------------- + + +class TestChunkKda: + """Direct tests for the chunk_kda kernel via tokamax backend.""" + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_basic(self): + from maxtext.kernels.kda import chunk_kda + from maxtext.layers.attention_kda import _l2_normalize + + B, T, H, K, V = 1, 2048, 4, 128, 128 + key = jax.random.PRNGKey(42) + keys = jax.random.split(key, 5) + q = jax.nn.silu(jax.random.normal(keys[0], (B, T, H, K), dtype=jnp.float32)) + k = jax.nn.silu(jax.random.normal(keys[1], (B, T, H, K), dtype=jnp.float32)) + q = _l2_normalize(q) + k = _l2_normalize(k) + v = jax.random.normal(keys[2], (B, T, H, V), dtype=jnp.float32) + g = jax.nn.log_sigmoid(jax.random.normal(keys[3], (B, T, H, K))) * 0.3 + beta = jax.nn.sigmoid(jax.random.normal(keys[4], (B, T, H))) + + o, _ = chunk_kda(q, k, v, g, beta, scale=K**-0.5, chunk_size=64) + assert o.shape == (B, T, H, V) + assert not jnp.any(jnp.isnan(o)) + + +# --------------------------------------------------------------------------- +# Naive KDA reference implementation and precision tests +# --------------------------------------------------------------------------- + + +def _naive_kda_recurrent(q, k, v, g, beta, scale): + """Naive Python implementation of KDA Delta Rule (recurrent form). + + Implements the exact recurrence from the KDA docstring: + S' = S * exp(g_t) (gated decay) + residual = v_t - S'^T @ k_t (delta residual) + S = S' + beta_t * k_t outer residual (state update) + o_t = scale * S @ q_t (output) + + Args: + q: [B, T, H, K] query + k: [B, T, H, K] key + v: [B, T, H, V] value + g: [B, T, H, K] gate (log-space, negative) + beta: [B, T, H] delta rule mixing coefficient + scale: float output scaling factor + + Returns: + o: [B, T, H, V] output + """ + B, T, H, K = q.shape + V = v.shape[-1] + o = jnp.zeros((B, T, H, V), dtype=jnp.float32) + + # S: [B, H, K, V] recurrent state + S = jnp.zeros((B, H, K, V), dtype=jnp.float32) + + for t in range(T): + # Extract per-step tensors + q_t = q[:, t, :, :] # [B, H, K] + k_t = k[:, t, :, :] # [B, H, K] + v_t = v[:, t, :, :] # [B, H, V] + g_t = g[:, t, :, :] # [B, H, K] + beta_t = beta[:, t, :] # [B, H] + + # Gated decay: S' = S * exp(g_t) + # g_t is [B, H, K], S is [B, H, K, V] -> broadcast over V + S = S * jnp.exp(g_t)[..., None] # [B, H, K, V] + + # Delta residual: residual = v_t - S^T @ k_t + # S^T @ k_t: [B, H, V, K] @ [B, H, K] -> [B, H, V] + # Equivalently: einsum('bhkv,bhk->bhv', S, k_t) + Sk = jnp.einsum("bhkv,bhk->bhv", S, k_t) # [B, H, V] + residual = v_t - Sk # [B, H, V] + + # State update: S = S + beta_t * k_t outer residual + # k_t: [B, H, K], residual: [B, H, V] -> outer: [B, H, K, V] + outer = k_t[..., None] * residual[..., None, :] # [B, H, K, V] + S = S + beta_t[..., None, None] * outer # [B, H, K, V] + + # Output: o_t = scale * S @ q_t + # einsum('bhkv,bhk->bhv', S, q_t) + o_t = scale * jnp.einsum("bhkv,bhk->bhv", S, q_t) # [B, H, V] + o = o.at[:, t, :, :].set(o_t) + + return o + + +class TestNaiveKda: + """Compare tokamax chunk_kda kernel against naive recurrent KDA implementation.""" + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_chunk_kda_vs_naive(self): + """Verify chunk_kda matches the naive Delta Rule recurrence.""" + from maxtext.kernels.kda import chunk_kda + from maxtext.layers.attention_kda import _l2_normalize + + B, T, H, K, V = 1, 64, 2, 16, 16 + key = jax.random.PRNGKey(0) + keys = jax.random.split(key, 5) + + q = jax.nn.silu(jax.random.normal(keys[0], (B, T, H, K), dtype=jnp.float32)) + k = jax.nn.silu(jax.random.normal(keys[1], (B, T, H, K), dtype=jnp.float32)) + q = _l2_normalize(q) + k = _l2_normalize(k) + v = jax.random.normal(keys[2], (B, T, H, V), dtype=jnp.float32) + g = jax.nn.log_sigmoid(jax.random.normal(keys[3], (B, T, H, K), dtype=jnp.float32)) * 0.3 + beta = jax.nn.sigmoid(jax.random.normal(keys[4], (B, T, H), dtype=jnp.float32)) + + scale = K**-0.5 + + o_kernel, _ = chunk_kda(q, k, v, g, beta, scale=scale, chunk_size=64) + o_naive = _naive_kda_recurrent(q, k, v, g, beta, scale) + + assert not jnp.any(jnp.isnan(o_naive)), "Naive output contains NaN" + assert not jnp.any(jnp.isnan(o_kernel)), "Kernel output contains NaN" + _assert_close(o_kernel, o_naive, "chunk_kda_vs_naive", atol=5e-3, rtol=1e-3) + + def test_naive_kda_basic_properties(self): + """Verify naive KDA implementation has correct basic properties.""" + B, T, H, K, V = 1, 8, 2, 4, 4 + key = jax.random.PRNGKey(42) + keys = jax.random.split(key, 5) + + q = jax.random.normal(keys[0], (B, T, H, K), dtype=jnp.float32) * 0.1 + k = jax.random.normal(keys[1], (B, T, H, K), dtype=jnp.float32) * 0.1 + v = jax.random.normal(keys[2], (B, T, H, V), dtype=jnp.float32) * 0.1 + g = -jnp.abs(jax.random.normal(keys[3], (B, T, H, K), dtype=jnp.float32)) * 0.1 + beta = jax.nn.sigmoid(jax.random.normal(keys[4], (B, T, H), dtype=jnp.float32)) + + scale = K**-0.5 + o = _naive_kda_recurrent(q, k, v, g, beta, scale) + + assert o.shape == (B, T, H, V) + assert not jnp.any(jnp.isnan(o)), "Output contains NaN" + assert not jnp.any(jnp.isinf(o)), "Output contains Inf" + # First position should be non-zero (state starts empty but gets updated) + assert jnp.any(o[:, 0, :, :] != 0), "First position output should be non-zero" + + def test_naive_kda_zero_gate_accumulates(self): + """With g=0 (no decay), state should accumulate without forgetting.""" + B, H, K, V = 1, 1, 2, 2 + T = 4 + + q = jnp.ones((B, T, H, K), dtype=jnp.float32) + k = jnp.ones((B, T, H, K), dtype=jnp.float32) * 0.1 + v = jnp.ones((B, T, H, V), dtype=jnp.float32) * 0.1 + g = jnp.zeros((B, T, H, K), dtype=jnp.float32) # no decay + beta = jnp.ones((B, T, H), dtype=jnp.float32) # full update + + scale = 1.0 + o = _naive_kda_recurrent(q, k, v, g, beta, scale) + + # Output magnitude should grow over time as state accumulates + norms = jnp.linalg.norm(o[0, :, 0, :], axis=-1) # [T] + # Later positions should have larger or equal output norm + assert norms[-1] >= norms[0], f"With zero gate, output norm should grow: first={norms[0]:.4f}, last={norms[-1]:.4f}" + + def test_naive_kda_large_negative_gate_decays(self): + """With very negative g, state should decay rapidly.""" + B, H, K, V = 1, 1, 2, 2 + T = 4 + + q = jnp.ones((B, T, H, K), dtype=jnp.float32) + k = jnp.zeros((B, T, H, K), dtype=jnp.float32) # no new info + v = jnp.zeros((B, T, H, V), dtype=jnp.float32) + g = jnp.full((B, T, H, K), -10.0, dtype=jnp.float32) # aggressive decay + beta = jnp.ones((B, T, H), dtype=jnp.float32) + + # Manually set initial state by making first step contribute + k = k.at[:, 0, :, :].set(1.0) + v = v.at[:, 0, :, :].set(1.0) + + scale = 1.0 + o = _naive_kda_recurrent(q, k, v, g, beta, scale) + + # After step 0, large negative gate should make state decay to ~0 + norm_0 = jnp.linalg.norm(o[0, 0, 0, :]) + norm_last = jnp.linalg.norm(o[0, -1, 0, :]) + assert norm_last < norm_0 * 0.01, ( + f"Large negative gate should decay state: t=0 norm={norm_0:.6f}, " f"t={T-1} norm={norm_last:.6f}" + ) + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_chunk_kda_vs_naive_bf16(self): + """Verify chunk_kda matches naive in bfloat16 (training dtype).""" + from maxtext.kernels.kda import chunk_kda + from maxtext.layers.attention_kda import _l2_normalize + + B, T, H, K, V = 1, 64, 2, 16, 16 + key = jax.random.PRNGKey(0) + keys = jax.random.split(key, 5) + + q = jax.nn.silu(jax.random.normal(keys[0], (B, T, H, K), dtype=jnp.float32)) + k = jax.nn.silu(jax.random.normal(keys[1], (B, T, H, K), dtype=jnp.float32)) + q = _l2_normalize(q) + k = _l2_normalize(k) + q = q.astype(jnp.bfloat16) + k = k.astype(jnp.bfloat16) + v = jax.random.normal(keys[2], (B, T, H, V), dtype=jnp.bfloat16) + g = jax.nn.log_sigmoid(jax.random.normal(keys[3], (B, T, H, K), dtype=jnp.float32)) * 0.3 + beta = jax.nn.sigmoid(jax.random.normal(keys[4], (B, T, H), dtype=jnp.float32)) + + scale = K**-0.5 + + o_kernel, _ = chunk_kda(q, k, v, g, beta, scale=scale, chunk_size=64) + o_naive = _naive_kda_recurrent( + q.astype(jnp.float32), + k.astype(jnp.float32), + v.astype(jnp.float32), + g, + beta, + scale, + ) + + assert not jnp.any(jnp.isnan(o_kernel)), "Kernel bf16 output contains NaN" + _assert_close(o_kernel, o_naive, "chunk_kda_bf16_vs_naive", atol=1e-2, rtol=1e-2) + + +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# QK L2 norm tests +# --------------------------------------------------------------------------- + + +class TestQkL2Norm: + """Verify QK L2 normalization is applied outside the kernel.""" + + @pytest.fixture + def mesh(self): + return jax.sharding.Mesh(jax.devices(), ("x",)) + + def test_qk_l2norm_applied_outside_kernel(self, mesh): + """With use_qk_norm=True, Q and K should be L2-normalized before kernel call.""" + cfg = _MockKdaConfig(use_qk_norm=True) + rngs = nnx.Rngs(0) + with mesh: + attn = attention_kda.KimiDeltaAttention(config=cfg, layer_idx=0, mesh=mesh, rngs=rngs) + B, T, D = 1, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + with mesh: + output, _ = attn(x) + assert output.shape == (B, T, D) + assert not jnp.any(jnp.isnan(output)) + + def test_qk_l2norm_skipped_when_disabled(self, mesh): + """With use_qk_norm=False, forward pass should still work without L2 norm.""" + cfg = _MockKdaConfig(use_qk_norm=False) + rngs = nnx.Rngs(0) + with mesh: + attn = attention_kda.KimiDeltaAttention(config=cfg, layer_idx=0, mesh=mesh, rngs=rngs) + B, T, D = 1, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + with mesh: + output, _ = attn(x) + assert output.shape == (B, T, D) + assert not jnp.any(jnp.isnan(output)) + + def test_l2norm_changes_output(self, mesh): + """Enabling vs disabling L2 norm should produce different outputs.""" + B, T, D = 1, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + + rngs_on = nnx.Rngs(0) + cfg_on = _MockKdaConfig(use_qk_norm=True) + with mesh: + attn_on = attention_kda.KimiDeltaAttention(config=cfg_on, layer_idx=0, mesh=mesh, rngs=rngs_on) + out_on, _ = attn_on(x) + + rngs_off = nnx.Rngs(0) + cfg_off = _MockKdaConfig(use_qk_norm=False) + with mesh: + attn_off = attention_kda.KimiDeltaAttention(config=cfg_off, layer_idx=0, mesh=mesh, rngs=rngs_off) + out_off, _ = attn_off(x) + + assert not jnp.allclose(out_on, out_off, atol=1e-4), "L2 norm on/off should produce different outputs" + + +# --------------------------------------------------------------------------- +# Backward (VJP) tests +# --------------------------------------------------------------------------- + + +class TestKdaBackward: + """Backward pass tests for KimiDeltaAttention (learning from GLA test patterns).""" + + @pytest.fixture + def mesh(self): + return jax.sharding.Mesh(jax.devices(), ("x",)) + + def _make_attn(self, mesh, **config_overrides): + cfg = _MockKdaConfig(**config_overrides) + rngs = nnx.Rngs(0) + with mesh: + return attention_kda.KimiDeltaAttention( + config=cfg, + layer_idx=0, + mesh=mesh, + rngs=rngs, + ) + + def _run_vjp(self, module, inp, mesh): + """Run gradient using value_and_grad instead of vjp.""" + graphdef, params, other = nnx.split(module, nnx.Param, ...) + + def forward_fn(params, x): + model = nnx.merge(graphdef, params, other) + with mesh: + out, _ = model(x) + # Return scalar loss for gradient computation + return jnp.sum(out) + + # Use value_and_grad instead of vjp (matching training code) + grad_fn = jax.value_and_grad(forward_fn, argnums=(0, 1), has_aux=False) + # Returns (loss, (grad_params, grad_input)) + _, grads = grad_fn(params, inp) + grad_params, grad_input = grads + return grad_params, grad_input + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_backward_no_nan(self, mesh): + """Activation gradient should be free of NaN/Inf and non-zero.""" + attn = self._make_attn(mesh) + B, T, D = 1, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + with mesh: + _, grad_input = self._run_vjp(attn, x, mesh) + assert not jnp.any(jnp.isnan(grad_input)), "grad_input contains NaN" + assert not jnp.any(jnp.isinf(grad_input)), "grad_input contains Inf" + assert jnp.any(grad_input != 0), "grad_input is all zeros" + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_backward_deterministic(self, mesh): + """Two VJP runs should produce identical gradients.""" + attn = self._make_attn(mesh) + B, T, D = 1, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + with mesh: + _, grad1 = self._run_vjp(attn, x, mesh) + _, grad2 = self._run_vjp(attn, x, mesh) + assert jnp.allclose(grad1, grad2, atol=1e-5), "Backward is not deterministic" + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_weight_grads_no_nan(self, mesh): + """Every parameter gradient should be free of NaN/Inf and non-zero.""" + attn = self._make_attn(mesh) + B, T, D = 1, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + with mesh: + grad_params, _ = self._run_vjp(attn, x, mesh) + + flat_grads = jax.tree.leaves(grad_params) + for i, g in enumerate(flat_grads): + assert not jnp.any(jnp.isnan(g)), f"weight grad {i} contains NaN" + assert not jnp.any(jnp.isinf(g)), f"weight grad {i} contains Inf" + assert jnp.any(g != 0), f"weight grad {i} is all zeros" + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_backward_bf16(self, mesh): + """bf16 backward should produce valid gradients.""" + attn = self._make_attn(mesh, dtype=jnp.bfloat16, weight_dtype=jnp.bfloat16) + B, T, D = 1, 64, 128 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D), dtype=jnp.bfloat16) + with mesh: + grad_params, grad_input = self._run_vjp(attn, x, mesh) + assert not jnp.any(jnp.isnan(grad_input)), "bf16 grad_input contains NaN" + assert not jnp.any(jnp.isinf(grad_input)), "bf16 grad_input contains Inf" + assert jnp.any(grad_input != 0), "bf16 grad_input is all zeros" + + flat_grads = jax.tree.leaves(grad_params) + for i, g in enumerate(flat_grads): + assert not jnp.any(jnp.isnan(g)), f"bf16 weight grad {i} contains NaN" + + +# --------------------------------------------------------------------------- +# ShortConvolution tests (standalone) +# --------------------------------------------------------------------------- + + +class TestShortConvolution: + """Tests for ShortConvolution module, including CP halo exchange.""" + + def test_short_conv_no_cp(self): + """ShortConvolution without CP should produce correct output and respect segment masks.""" + from maxtext.layers.attention_kda import ShortConvolution + + rngs = nnx.Rngs(0) + kernel_size, features = 4, 8 + conv = ShortConvolution( + kernel_size=kernel_size, + features=features, + dtype=jnp.float32, + weight_dtype=jnp.float32, + rngs=rngs, + ) + + B, T = 2, 16 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, features)) + + # Without segment_ids: causal depthwise conv on full sequence. + out = conv(x) + assert out.shape == (B, T, features) + assert jnp.isfinite(out).all() + # Output should differ from input (conv applied). + assert not jnp.allclose(out, x, atol=1e-6) + + # With segment_ids: cross-segment contributions should be masked out. + seg_ids = jnp.array( + [[1, 1, 1, 1, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 0, 0, 0, 0, 0, 0]], + dtype=jnp.int32, + ) + out_seg = conv(x, segment_ids=seg_ids) + assert out_seg.shape == (B, T, features) + assert jnp.isfinite(out_seg).all() + # Segment masking should change output. + assert not jnp.allclose(out_seg, out, atol=1e-6) + + # Row independence: changing row 1's segment_ids should not affect row 0. + seg_alt = jnp.array( + [[1, 1, 1, 1, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 0, 0, 0, 0, 0, 0]], + dtype=jnp.int32, + ) + out_alt = conv(x, segment_ids=seg_alt) + assert jnp.allclose(out_seg[0], out_alt[0], atol=0.0), "Row 0 output changed when only row 1 segments changed" + + @pytest.mark.skipif(len(jax.devices()) < 2, reason="need >=2 devices for CP test") + def test_short_conv_cp_halo(self): + """ShortConvolution under CP: shard_map with halo exchange matches reference. + + Verifies that when ShortConvolution runs inside a shard_map with the + "context" axis, ``halo_exchange_for_conv`` pulls left-context tokens + from the previous CP rank so the causal-conv output is identical to + running on the full (non-sharded) sequence. + """ + from maxtext.layers.attention_kda import ShortConvolution + + devices = jax.devices() + cp_size = 2 + n_devices = (len(devices) // cp_size) * cp_size + mesh = jax.sharding.Mesh(np.array(devices[:n_devices]).reshape(cp_size, -1), ("context", "x")) + + kernel_size, features = 4, 8 + rngs = nnx.Rngs(0) + conv = ShortConvolution( + kernel_size=kernel_size, + features=features, + dtype=jnp.float32, + weight_dtype=jnp.float32, + rngs=rngs, + ) + + B, T = 2, 32 # divisible by cp_size + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, features)) + + # Reference: conv on full sequence without CP sharding. + ref_out = jax.device_get(conv(x)) + + # CP: shard input along T, run conv inside shard_map with "context" axis. + xs = jax.lax.with_sharding_constraint( + x, jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "context", None)) + ) + + # Uniform segment_ids — no cross-segment masking, so halo tokens are + # the true left-context and output should match reference. + seg_ids = jnp.ones((B, T), dtype=jnp.int32) + segs = jax.lax.with_sharding_constraint( + seg_ids, jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "context")) + ) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=( + jax.sharding.PartitionSpec(None, "context", None), + jax.sharding.PartitionSpec(None, "context"), + ), + out_specs=jax.sharding.PartitionSpec(None, "context", None), + check_vma=False, + ) + def _conv_cp(x_local, seg_local): + return conv(x_local, segment_ids=seg_local) + + cp_out = _conv_cp(xs, segs) + # All-gather: replicate across context axis so we can compare. + cp_out_full = jax.device_get( + jax.lax.with_sharding_constraint(cp_out, jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())) + ) + + # With uniform segment_ids (no segmentation), CP conv with halo should + # match reference exactly — the halo provides the true left context. + assert jnp.allclose(cp_out_full, ref_out, atol=1e-5), ( + f"ShortConvolution CP halo output differs from reference. " + f"max_diff={float(jnp.abs(cp_out_full - ref_out).max()):.2e}" + ) + + +# --------------------------------------------------------------------------- +# CP (Context Parallelism) tests +# --------------------------------------------------------------------------- + + +class TestKdaCp: + """Tests for KDA context parallelism.""" + + def _cp_mesh(self, cp_size=2): + devices = jax.devices() + n_devices = (len(devices) // cp_size) * cp_size + return jax.sharding.Mesh(np.array(devices[:n_devices]).reshape(cp_size, -1), ("context", "x")) + + def _cp_config(self, **overrides): + return _MockKdaConfig( + context_parallel_size=2, + context_parallel_strategy="all_gather", + context_parallel_load_balance=False, + **overrides, + ) + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + @pytest.mark.skipif(len(jax.devices()) < 2, reason="need >=2 devices for CP test") + def test_kda_cp_equivalence(self): + """KDA with CP should produce equivalent output to non-CP KDA.""" + from maxtext.layers.attention_kda import _l2_normalize + from maxtext.kernels.kda import chunk_kda + + cp_size = 2 + mesh_cp = self._cp_mesh(cp_size=cp_size) + mesh_ref = jax.sharding.Mesh(jax.devices(), ("x",)) + + B, T, H, K, V = 2, 128, 4, 128, 128 + key = jax.random.PRNGKey(42) + keys = jax.random.split(key, 5) + q = jax.nn.silu(jax.random.normal(keys[0], (B, T, H, K), dtype=jnp.float32)) + k = jax.nn.silu(jax.random.normal(keys[1], (B, T, H, K), dtype=jnp.float32)) + q = _l2_normalize(q) + k = _l2_normalize(k) + v = jax.random.normal(keys[2], (B, T, H, V), dtype=jnp.float32) + g = jax.nn.log_sigmoid(jax.random.normal(keys[3], (B, T, H, K))) * 0.3 + beta = jax.nn.sigmoid(jax.random.normal(keys[4], (B, T, H))) + seg_ids = jnp.ones((B, T), dtype=jnp.int32) + scale = float(K**-0.5) + + # --- Reference: non-CP run --- + ref_o, _ = chunk_kda(q, k, v, g, beta, scale=scale, chunk_size=64, segment_ids=seg_ids, N_max=1) + + # --- CP run: shard along T, call chunk_kda with cp_context --- + from tokamax._src.ops.experimental.kda.cp_utils import CPContext + + cp_ctx = CPContext(mesh=mesh_cp, axis_name="context") + + # Shard all inputs along T (axis 1). + pspec_4d = jax.sharding.PartitionSpec(None, "context", None, None) + pspec_3d = jax.sharding.PartitionSpec(None, "context", None) + pspec_2d = jax.sharding.PartitionSpec(None, "context") + + def _shard(arr, pspec): + return jax.lax.with_sharding_constraint(arr, jax.sharding.NamedSharding(mesh_cp, pspec)) + + qs = _shard(q, pspec_4d) + ks = _shard(k, pspec_4d) + vs = _shard(v, pspec_4d) + gs = _shard(g, pspec_4d) + betas = _shard(beta, pspec_3d) + segs = _shard(seg_ids, pspec_2d) + + @functools.partial( + jax.shard_map, + mesh=mesh_cp, + in_specs=(pspec_4d, pspec_4d, pspec_4d, pspec_4d, pspec_3d, pspec_2d), + out_specs=pspec_4d, + check_vma=False, + ) + def _kda_cp(q_loc, k_loc, v_loc, g_loc, beta_loc, seg_loc): + o_loc, _ = chunk_kda( + q_loc, + k_loc, + v_loc, + g_loc, + beta_loc, + scale=scale, + chunk_size=64, + segment_ids=seg_loc, + disable_recompute=True, + N_max=1, + cp_context=cp_ctx, + ) + return o_loc + + cp_o = _kda_cp(qs, ks, vs, gs, betas, segs) + cp_o_full = jax.device_get( + jax.lax.with_sharding_constraint(cp_o, jax.sharding.NamedSharding(mesh_cp, jax.sharding.PartitionSpec())) + ) + + # CP output should match reference within tolerance. + _assert_close(cp_o_full, ref_o, "kda_cp_equivalence", atol=5e-3, rtol=1e-3) + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_kda_cp_rejects_load_balance(self): + """KDA CP should raise ValueError when load_balance is enabled.""" + mesh = jax.sharding.Mesh(jax.devices(), ("x",)) + cfg = _MockKdaConfig( + context_parallel_size=2, + context_parallel_strategy="all_gather", + context_parallel_load_balance=True, + ) + rngs = nnx.Rngs(0) + with mesh: + attn = attention_kda.KimiDeltaAttention( + config=cfg, + layer_idx=0, + mesh=mesh, + rngs=rngs, + ) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + with pytest.raises(ValueError, match="load_balance"): + attn(x) + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_kda_no_cp_without_load_balance_ok(self): + """KDA without CP (cp_size=1) should succeed.""" + mesh = jax.sharding.Mesh(jax.devices(), ("x",)) + cfg = _MockKdaConfig( + context_parallel_size=1, # CP=1, no actual CP but validates config check path + ) + rngs = nnx.Rngs(0) + with mesh: + attn = attention_kda.KimiDeltaAttention( + config=cfg, + layer_idx=0, + mesh=mesh, + rngs=rngs, + ) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + with mesh: + output, _ = attn(x) + assert output.shape == (1, 64, 128) + assert jnp.isfinite(output).all()