diff --git a/docs/reference.md b/docs/reference.md index fe8d74faa6..f94068caca 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -49,6 +49,13 @@ High-level overview of MaxText design, JAX/XLA choices, and how components inter Key concepts including checkpointing strategies, quantization, tiling, and Mixture of Experts (MoE) configuration. ``` + +```{grid-item-card} 🔁 KDA Context Parallelism +:link: reference/kda_cp_support +:link-type: doc + +Design of Kimi Delta Attention (KDA) CP support: halo exchange, per-rank recurrent state, segment handling, and constraints. +``` ```` ## 📚 API Reference @@ -64,5 +71,6 @@ reference/performance_metrics reference/models reference/architecture reference/core_concepts +reference/kda_cp_support reference/api.rst ``` diff --git a/docs/reference/kda_cp_support.md b/docs/reference/kda_cp_support.md new file mode 100644 index 0000000000..35139e622b --- /dev/null +++ b/docs/reference/kda_cp_support.md @@ -0,0 +1,188 @@ +# 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 `ContextParallelMetadata` 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 + → ContextParallelMetadata(mesh, cfg.context_sharding) ← constructed outside shard_map + → _inject_cp_axis_on_T + _wsc ← partition spec fixup + → SHARD_MAP(chunk_kda) ← context_parallel_metadata 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, `ContextParallelMetadata` lets the kernel coordinate recurrent state across ranks during forward/backward. + +### Plan 1: `halo_exchange_for_conv` (in `layers/attention_kda.py`, KDA-specific) + +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 CP axis (the `cfg.context_sharding` mesh axis, default `"context"`). See Plan 2. + +**Constraint**: the exchange only reads from the immediately preceding rank, so it requires `halo_size <= T_local` (i.e. `linear_conv_kernel_dim - 1` must not exceed the per-rank sequence length). Larger receptive fields would span multiple ranks and are not implemented; `halo_exchange_for_conv` raises a `ValueError` in that case. + +### Plan 2: ShortConvolution CP Wrapper (`layers/attention_kda.py`) + +`ShortConvolution.__call__` internally calls `halo_exchange_for_conv`, which requires the CP axis scope (`cfg.context_sharding`). 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__`. + +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 ContextParallelMetadata + Partition Spec (`attention_kda.py`) + +#### 3a. ContextParallelMetadata Construction (outside shard_map) + +```python +try: + from tokamax._src.ops.experimental.kda.cp_utils import ( + ContextParallelMetadata as TokamaxContextParallelMetadata, + ) +except ImportError: + TokamaxContextParallelMetadata = None + +cp_axis_name = cfg.context_sharding # default "context"; "expert" for expert-as-context +if cp_size > 1: + if TokamaxContextParallelMetadata is None: + raise ImportError(...) # refuse to run: CP would silently break state + cp_ctx = TokamaxContextParallelMetadata(mesh=self.mesh, axis_name=cp_axis_name) +``` + +`ContextParallelMetadata` is a frozen dataclass. `mesh` and `axis_name` are set at construction time; chain metadata fields are populated internally by `chunk_kda`. The `axis_name` comes from `cfg.context_sharding`, so expert-as-context meshes bind the metadata to the `"expert"` axis. + +#### 3b. Partition Spec Injection + +`nnx.logical_to_mesh_axes` may map the T axis to `None` (or to a mesh axis that does not carry the sequence shard) due to Flax rule priority + size-1 axis stripping, but shard_map requires the T axis to carry the CP axis: + +```python +def _inject_cp_axis_on_T(pspec, t_axis=1): + spec = list(pspec) + spec[t_axis] = cp_axis_name # overwritten unconditionally + 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. + +`cp_axis_name` is `cfg.context_sharding` (default `"context"`; may be `"expert"` for expert-as-context). The T axis is overwritten **unconditionally** rather than only when it maps to `None`: the `activation_norm_length` logical-axis rules do not cover every CP strategy (notably expert-as-context), so an unconditional overwrite guarantees the shard_map always sees the per-rank sequence shards on the axis the collectives (halo exchange, cross-rank state merge) actually use. + +#### 3c. chunk_kda shard_map + +Under CP, pass through `context_parallel_metadata=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 +cp_size = self.mesh.shape.get(cfg.context_sharding, 1) +if cp_size > 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) + │ + ├── T-padding: pad sequence to a multiple of the chunk alignment (64) + │ + ├── ShortConvolution: halo_exchange_for_conv(segment_ids) + │ cross-segment boundary masking inside conv + │ + ├── _inject_cp_axis_on_T + _wsc: inject cfg.context_sharding axis onto the T axis + │ + └── 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 | ~743 | +| `kernels/kda/__init__.py` | **New**: `chunk_kda()` entry point | ~99 | +| `kernels/kda/tokamax.py` | **New**: tokamax backend adapter (lazy import) | ~143 | +| `configs/types.py` | **Modified**: `KdaAttention` config class + validators | +~90 | +| `tests/unit/kda_attention_test.py` | **New**: layer + conv halo + CP fwd/bwd + packed-seg CP + parity + e2e smoke test | ~1675 | +| `docs/reference/kda_cp_support.md` | **New**: design doc | — | +| `**Total**` | | **~3035** | + +## Key Constraints + +1. **ContextParallelMetadata availability**: Raise `ImportError` with a clear message when ContextParallelMetadata 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 `ContextParallelMetadata`. 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 `cp_size > 1` (derived from the mesh's `context_sharding` axis) +- ContextParallelMetadata import: `try/except`, raise `ImportError` with clear message if unavailable +- segment_ids dummy: auto-construct `jnp.ones` when no varlen + CP + +## Test Plan + +Only tests that invoke the Mosaic Pallas kernel or multi-device CP carry the `tpu_only` marker; pure config / pure-op / non-CP tests run in regular CPU CI as well. + +| 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; parametrized over segment layouts: uniform, boundary on the rank split, and a segment spanning both ranks (halo + segment-mask interaction) | +| `test_short_conv_cp_rejects_oversized_halo` | `halo_size > T_local` under CP raises a clear ValueError (multi-rank receptive field not implemented) | +| `test_kda_cp_equivalence` | kernel-level CP multi-rank forward equals single-rank, parametrized CP=2 and CP=4 | +| `test_kda_cp_backward` | CP gradients (dq/dk/dv/dg/dbeta) equal the non-CP reference | +| `test_kda_cp_full_layer_dummy_segments` | full layer under CP with no user segment_ids: covers the internal dummy-segment synthesis path, forward equivalence and backward finiteness | +| `test_kda_cp_full_layer_packed_segments` | full layer under CP with multiple real packed segments — one spanning the rank boundary, one boundary exactly at the split; forward + input/weight gradients equal the non-CP reference | +| `test_full_layer_mosaic_vs_xla_parity` | full layer with identical weights, Mosaic kernel vs tokamax XLA reference implementation: forward + gradients match | +| `test_kda_cp_rejects_load_balance` | CP+load_balance raises ValueError | +| `test_packed_segment_no_leak_within_row` | packed segments inside one row are structurally isolated in both directions | +| `test_l2_normalize_produces_unit_norm` | `_l2_normalize` yields unit L2 norm and preserves direction | +| `TestKdaConfigGuards` | config-time guards: safe-gate/lower_bound range, `use_kda_lora=True` rejection, packing without `max_segments_per_seq` | diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index ccb439a9ac..4b893d3c2f 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -137,6 +137,7 @@ class AttentionType(enum.Enum): LOCAL_SLIDING = "local_sliding" CHUNK = "chunk" MLA = "mla" + KDA = "kda" # Kimi Delta Attention: recurrent Delta-Rule attention, tokamax Pallas kernel COMPRESSED = "compressed" FULL = "full" BLOCK_DIFFUSION = "block_diffusion" diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 5669e012ae..2f2c1352bf 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -408,10 +408,16 @@ param_scan_axis: 1 # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores # The attention_type parameter determines the variants of attention, e.g. global or local_sliding attention: 'autoselected' # Supported attention: autoselected, dot_product, flash, cudnn_flash_te -attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla, full, compressed, block_diffusion +attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla, kda, full, compressed, block_diffusion share_kv_projections: false # Note: Not compatible with attention_type='mla' attention_bias: false # If true, adds a learnable bias to the query, key, and value projections attention_sink: false +# KDA (Kimi Delta Attention) options, consumed by the KimiDeltaAttention layer +# selected via attention_type: 'kda'. +linear_conv_kernel_dim: 4 # Kernel size of the causal depthwise short convolution applied to Q/K/V in KDA; 0 disables it +use_kda_safe_gate: false # If true, uses the sigmoid lower-bound ("safe") gate path of the KDA kernel +kda_lower_bound: 0.0 # Lower bound of the KDA sigmoid gate; only used when use_kda_safe_gate=true (must be in [-5, 0); -5.0 is a common choice) +use_kda_lora: false # Reserved: the KDA LoRA path is not implemented and must stay false sliding_window_size: 0 chunk_attn_window_size: 0 # Token block size B for block-causal attention in Block Diffusion (arXiv:2503.09573). diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 353656c665..5b1f3db75e 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -632,8 +632,8 @@ class Attention(BaseModel): "autoselected", description="The attention algorithm to use (dot_product, flash, cudnn_flash_te, vllm_rpa, vllm_batched_rpa, etc).", ) - attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed", "block_diffusion"] = Field( - "global", description="The variant of attention to use." + attention_type: Literal["global", "local_sliding", "chunk", "mla", "kda", "full", "compressed", "block_diffusion"] = ( + Field("global", description="The variant of attention to use.") ) share_kv_projections: bool = Field( False, @@ -746,6 +746,81 @@ 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 depthwise causal 1D convolution applied to Q, K and V " + "for local dependency modeling. Default 4 matches the reference Megatron implementation." + ), + ) + use_kda_lora: bool = Field( + False, + description=( + "Reserved for a future LoRA (Low-Rank Adaptation) KDA variant. " + "The current KimiDeltaAttention layer only implements the full-rank " + "(no-LoRA) path and does not read this flag." + ), + ) + use_kda_safe_gate: bool = Field( + False, + description=( + "Whether to use the numerically safe (sigmoid lower-bound) gate path in KDA " + "layers instead of the standard softplus activation. When True, " + "``kda_lower_bound`` is passed to the tokamax kernel as ``lower_bound``; " + "when False the kernel uses its standard gate activation." + ), + ) + kda_lower_bound: float = Field( + 0.0, + description=( + "Lower bound for the sigmoid gate path in KDA layers, used only when " + "``use_kda_safe_gate=True``. Passed to the tokamax kernel as " + "``lower_bound`` (which requires a value in ``[-5, 0)``). " + "-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 + + @field_validator("use_kda_lora") + @classmethod + def _check_use_kda_lora_not_set(cls, v: bool) -> bool: + if v: + raise ValueError( + "use_kda_lora=True is not implemented: KimiDeltaAttention only " + "implements the full-rank (no-LoRA) path. Leave it False." + ) + return v + + @model_validator(mode="after") + def _check_safe_gate_lower_bound(self): + """Cross-field guard: the sigmoid gate path requires lower_bound in [-5, 0). + + Rejects invalid combinations at config time instead of failing deep in + tokamax kernel binding (tokamax enforces the same range on `lower_bound`). + """ + if self.use_kda_safe_gate and (self.kda_lower_bound < -5.0 or self.kda_lower_bound >= 0.0): + raise ValueError( + "use_kda_safe_gate=True requires kda_lower_bound in [-5, 0) " + f"(the tokamax sigmoid gate path constraint), got " + f"kda_lower_bound={self.kda_lower_bound}. A common choice is -5.0." + ) + return self + + class AttentionIndexer(BaseModel): """Configuration for DeepSeek Sparse Attention (DSA): DeepSeek3.2-style MLA with indexer.""" @@ -3174,6 +3249,7 @@ class MaxTextConfig( # Attention Mechanisms Attention, MlaAttention, + KdaAttention, CompressedAttention, MoBa, AttentionIndexer, @@ -4658,6 +4734,12 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"QK-Clip is only supported when attention_type='mla', but found attention_type='{self.attention_type}'." ) + if self.attention_type == "kda" and self.scan_layers: + raise ValueError( + "attention_type='kda' requires scan_layers=false: KDA layers have not been validated " + "inside a scanned layer stack. Set scan_layers: false." + ) + if self.use_qk_clip and self.attn_logits_soft_cap is not None: raise ValueError( "QK-Clip monitors raw dot products, but attn_logits_soft_cap is enabled. " diff --git a/src/maxtext/kernels/kda/__init__.py b/src/maxtext/kernels/kda/__init__.py new file mode 100644 index 0000000000..a08ce23847 --- /dev/null +++ b/src/maxtext/kernels/kda/__init__.py @@ -0,0 +1,99 @@ +# Copyright 2026 Ant Group. All Rights Reserved. +# +# 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. + +"""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 ``context_parallel_metadata``. +""" + +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, + a_log: jnp.ndarray | None = None, + delta_time_bias: jnp.ndarray | None = None, + use_gate_in_kernel: bool = False, + use_qk_l2norm: bool = False, + lower_bound: float | None = None, + segment_ids: jnp.ndarray | None = None, + max_num_segments: int | None = None, + context_parallel_metadata: 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 ``K ** -0.5``). + initial_state: must be None (not yet supported). + output_final_state: must be False (not yet supported). + a_log: [H] per-head log decay-rate parameter. Required when + ``use_gate_in_kernel=True``. + delta_time_bias: [H*K] optional per-head, per-key-channel gate bias. + use_gate_in_kernel: whether ``g`` is activated (``a_log`` / + ``delta_time_bias``) inside the kernel. + use_qk_l2norm: whether to L2-normalize q/k in-kernel. + lower_bound: optional sigmoid-gate lower bound in ``[-5, 0)``. When + None, the standard ``softplus`` gate path is used. + segment_ids: [B, T] 1-based segment IDs for varlen mode (0=padding). + max_num_segments: static upper bound on varlen segments. Required when + ``segment_ids`` is provided without ``initial_state``. + context_parallel_metadata: optional ``ContextParallelMetadata`` 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, + a_log=a_log, + delta_time_bias=delta_time_bias, + use_gate_in_kernel=use_gate_in_kernel, + use_qk_l2norm=use_qk_l2norm, + lower_bound=lower_bound, + segment_ids=segment_ids, + max_num_segments=max_num_segments, + context_parallel_metadata=context_parallel_metadata, + ) diff --git a/src/maxtext/kernels/kda/tokamax.py b/src/maxtext/kernels/kda/tokamax.py new file mode 100644 index 0000000000..806b560841 --- /dev/null +++ b/src/maxtext/kernels/kda/tokamax.py @@ -0,0 +1,152 @@ +# Copyright 2026 Ant Group. All Rights Reserved. +# +# 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. + +"""Tokamax KDA backend for maxtext. + +Wraps ``tokamax._src.ops.experimental.kda.api.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.numpy as jnp + +# TODO(kda): `kimi_delta_attention` lives on tokamax's experimental path +# `tokamax._src.ops.experimental.kda`. The KDA change has landed on +# openxla/tokamax main (the original PR #1103 was left open, but the API is +# on main, signature-compatible with this adapter). No public tokamax release +# contains it yet, so no pip version specifier can express this dependency +# and the requirement pins are left unchanged. This import is deliberately +# lazy: a clean MaxText installation without the KDA API still works for +# everything else, and only KDA use fails — with this ImportError as the +# symptom. Once tokamax cuts a release containing the KDA API: (1) switch +# this adapter to the stable public entry point if one is added, (2) bump the +# tokamax pins under src/dependencies/requirements/ to the first release +# containing it and regenerate the derived requirement files. Keep the lazy +# import and this note in sync until then. + + +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, + a_log: jnp.ndarray | None = None, + delta_time_bias: jnp.ndarray | None = None, + use_gate_in_kernel: bool = False, + use_qk_l2norm: bool = False, + lower_bound: float | None = None, + segment_ids: jnp.ndarray | None = None, + max_num_segments: int | None = None, + context_parallel_metadata: 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] queries. + k: [B, T, H, K] keys. + v: [B, T, H, V] values. + g: [B, T, H, K] raw gate values (activated in-kernel when + ``use_gate_in_kernel=True``; otherwise already in log space). + beta: [B, T, H] delta-rule mixing coefficient in [0, 1]. + scale: attention scale (default ``K ** -0.5``). + initial_state: must be None (not yet supported in maxtext). + output_final_state: must be False (not yet supported in maxtext). + a_log: [H] per-head log decay-rate parameter. Required when + ``use_gate_in_kernel=True``. + delta_time_bias: [H*K] optional per-head, per-key-channel gate bias. + Used only when ``use_gate_in_kernel=True``. + use_gate_in_kernel: whether ``g`` is raw delta-time input that should + be activated with ``a_log`` / ``delta_time_bias`` inside the kernel. + use_qk_l2norm: whether to L2-normalize q/k on the last dim in-kernel. + lower_bound: optional sigmoid-gate lower bound in ``[-5, 0)``. When + None, the standard ``softplus`` gate path is used. + segment_ids: [B, T] 1-based, 0=padding. Passed directly to tokamax. + max_num_segments: static upper bound on varlen segments. Required when + ``segment_ids`` is provided without ``initial_state``. + context_parallel_metadata: optional ``tokamax ... ContextParallelMetadata`` + for context parallelism. + + Returns: + (o, None) where o is [B, T, H, V]. + """ + # Input validation before the lazy import so the guards fire even on + # installs without tokamax. + 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") + + # Deliberately lazy: importing the KDA API must not fail at module import + # time on installs without tokamax; only an actual KDA call requires it. + try: + from tokamax._src.ops.experimental.kda.api import ( # pylint: disable=import-outside-toplevel + kimi_delta_attention, + ) + except ImportError as exc: + raise ImportError( + "KDA requires the tokamax KDA API (tokamax._src.ops.experimental.kda.api), " + "which is not available in the installed tokamax build. The KDA kernels " + "have landed on openxla/tokamax main but no public release contains them " + "yet — install tokamax from source (openxla/tokamax main, or the " + "antgroup/kda-pallas-kernel branch) until the first release ships." + ) from exc + + q_h, k_h, v_h, g_h, beta_h = _to_tokamax(q, k, v, g, beta) + + o_h, _ = kimi_delta_attention( + query=q_h, + key=k_h, + value=v_h, + gate=g_h, + beta=beta_h, + a_log=a_log, + delta_time_bias=delta_time_bias, + scale=scale, + segment_ids=segment_ids, + use_gate_in_kernel=use_gate_in_kernel, + use_qk_l2norm=use_qk_l2norm, + lower_bound=lower_bound, + max_num_segments=max_num_segments, + implementation="mosaic", # tokamax's Pallas Mosaic TPU kernel ("mosaic" resolves to the mosaic_tpu impl) + context_parallel_metadata=context_parallel_metadata, + ) + + 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..41f3aa17a3 --- /dev/null +++ b/src/maxtext/layers/attention_kda.py @@ -0,0 +1,747 @@ +# Copyright 2026 Ant Group. All Rights Reserved. +# +# 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. + +KDA is a linear attention mechanism with Delta Rule correction, featuring: + - Depthwise causal 1D convolution for local dependency modeling + - Optional numerically safe gate (sigmoid lower-bound) mechanism + - Optional Q/K L2 normalization + +The layer computes Q/K/V projections, short convolutions, gate/beta +projections, and delegates the chunk-parallel Delta Rule recurrence to +``tokamax._src.ops.experimental.kda.api.kimi_delta_attention`` via +``maxtext.kernels.kda.chunk_kda``. +""" + + +import functools +import math +import warnings + +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 ( + ContextParallelMetadata as TokamaxContextParallelMetadata, + ) +except ImportError: + TokamaxContextParallelMetadata = 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 + + +# Sequence is padded to a multiple of this size before the KDA kernel, so +# TPU-friendly fixed shapes are used (matching the Megatron chunk 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) + + +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. + + KDA's ``ShortConvolution`` is the only user today; the helper lives in this + module accordingly. The caller 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. + + Constraint: the exchange only reads from the immediately preceding rank, + so ``halo_size`` must not exceed the local sequence length. A larger + receptive field (kernel_size - 1 > T_local) would need tokens from + multiple previous ranks, which is not implemented; a ``ValueError`` is + raised instead of silently reading the wrong context. + + 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 + + t_local = x.shape[seq_axis] + if halo_size > t_local: + raise ValueError( + f"halo_exchange_for_conv: halo_size ({halo_size}) exceeds the local " + f"sequence length ({t_local}) on the '{axis_name}' axis. The causal " + "convolution receptive field would span multiple CP ranks, which is " + "not implemented. Use a smaller linear_conv_kernel_dim, a longer " + "sequence, or a smaller CP size." + ) + + # 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) + + +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, + cp_axis_name: str = "context", + ) -> 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=cp_axis_name) + + if segment_ids is not None: + seg_padded = halo_exchange_for_conv(segment_ids, self.kernel_size - 1, axis_name=cp_axis_name) + # 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). Params keep the + # Megatron names (A_log / dt_bias) but are passed to tokamax as `a_log` + # and `delta_time_bias` inside the shard_map below. + # 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 + + # Context-parallel size derived from the mesh (the CP axis name is + # cfg.context_sharding, default "context"; it may also be "expert" for + # expert-as-context). This mirrors attention_op.py. + cp_axis_name = cfg.context_sharding + cp_size = self.mesh.shape.get(cp_axis_name, 1) + + # 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 up front. + if cp_size > 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." + ) + + if model_mode == MODEL_MODE_AUTOREGRESSIVE: + raise NotImplementedError("KDA autoregressive mode not yet implemented.") + + # Packed/varlen execution: tokamax requires a static positive + # max_num_segments whenever segment_ids are supplied without + # initial_state (see tokamax .../kda/api.py). maxtext surfaces this as + # `max_segments_per_seq`, which defaults to -1 (unset) — fail fast with + # a config-level message instead of erroring deep in kernel binding. + if decoder_segment_ids is not None and cfg.max_segments_per_seq <= 0: + raise ValueError( + "KDA with packed sequences (decoder_segment_ids) requires " + f"`max_segments_per_seq` to be a positive integer, got " + f"{cfg.max_segments_per_seq}. Set `max_segments_per_seq` in your " + "config to a static upper bound on the number of packed segments " + "per sequence." + ) + + def _inject_cp_axis_on_T(pspec, t_axis=1): + """Overwrite the T axis of *pspec* with the CP axis name. + + logical_to_mesh_axes may map the LENGTH logical axis to a different + mesh axis, or to None, because the activation_norm_length rules do not + cover every CP strategy (notably expert-as-context). Overwrite + unconditionally so shard_map always sees the correct per-rank sequence + shards on the axis the collectives (halo exchange, CP state merge) use. + """ + spec = list(pspec) + spec[t_axis] = cp_axis_name + return jax.sharding.PartitionSpec(*spec) + + 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 CP mesh axis (cfg.context_sharding). 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 cp_size > 1: + conv_flat_pspec = _inject_cp_axis_on_T( + self._logical_to_mesh_axes(("activation_batch", "activation_norm_length", None)) + ) + conv_seg_pspec = ( + _inject_cp_axis_on_T(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, cp_axis_name=cp_axis_name) + kf = k_conv_mod(kf, segment_ids=seg, cp_axis_name=cp_axis_name) + vf = v_conv_mod(vf, segment_ids=seg, cp_axis_name=cp_axis_name) + 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, cp_axis_name=cp_axis_name) + k_flat = self.k_conv(k_flat, segment_ids=decoder_segment_ids, cp_axis_name=cp_axis_name) + v_flat = self.v_conv(v_flat, segment_ids=decoder_segment_ids, cp_axis_name=cp_axis_name) + + 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). Always on for KDA: the Delta-Rule recurrence is + # numerically unstable with unbounded q/k (training diverges to NaN in + # bf16), and QK L2-norm is part of the KDA architecture in the reference. + # This is independent of the shared `use_qk_norm` flag, which belongs to + # dot-product attention. + 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: + 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.max_segments_per_seq if cfg.max_segments_per_seq > 0 else None + if cp_size > 1 and decoder_segment_ids is None: + # Under CP the tokamax kernel derives per-rank cu_seqlens / chain + # metadata from segment_ids, so a seg tensor must always be present. + # Without user segmentation, synthesize a single all-ones segment — + # done outside shard_map so a real array is sharded through. + decoder_segment_ids = jnp.ones((B, T), dtype=jnp.int32) + n_max = 1 + + # 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",)) + delta_time_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 the dt_bias param from [H*K] to [H, K] for head-dim sharding. + # (Tokamax's argument name is delta_time_bias; the nnx param is named + # dt_bias to match the Megatron reference.) + delta_time_bias_2d = self.dt_bias.value.reshape(self.num_key_heads, self.key_head_dim) + + # Force the CP axis onto the T axis of every pspec so the shard_map + # sees the correct per-rank shard layout and the collectives run on + # the axis the sequence is actually sharded over. + if cp_size > 1: + qkv_pspec = _inject_cp_axis_on_T(qkv_pspec) + beta_pspec = _inject_cp_axis_on_T(beta_pspec) + seg_pspec = _inject_cp_axis_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) + + # Under CP a seg tensor is always present (synthesized above when + # the user supplies none), so the kernel can derive cu_seqlens / chain + # fields via one small all_gather. + has_seg = decoder_segment_ids is not None + base_in_specs = ( + qkv_pspec, + qkv_pspec, + qkv_pspec, + qkv_pspec, + beta_pspec, + a_log_pspec, + delta_time_bias_2d_pspec, + ) + in_specs = base_in_specs + ((seg_pspec,) if has_seg else ()) + + # ContextParallelMetadata 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 metadata to the kernel. + cp_ctx = None + if cp_size > 1: + if TokamaxContextParallelMetadata is None: + raise ImportError( + "KDA context parallelism requires " + "tokamax._src.ops.experimental.kda.cp_utils.ContextParallelMetadata, " + "but it failed to import. Refusing to run: CP would silently " + "break recurrent state across ranks." + ) + cp_ctx = TokamaxContextParallelMetadata(mesh=self.mesh, axis_name=cp_axis_name) + + @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, delta_time_bias_2d, *rest = args + seg = rest[0] if rest else None + + delta_time_bias_flat = delta_time_bias_2d.reshape(-1) + o, _ = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + a_log=a_log, + delta_time_bias=delta_time_bias_flat, + segment_ids=seg, + scale=scale, + initial_state=None, + output_final_state=False, + use_qk_l2norm=False, + use_gate_in_kernel=True, + lower_bound=lower_bound, + max_num_segments=n_max, + context_parallel_metadata=cp_ctx, + ) + return o + + kda_args = (q, k, v, g, beta, self.A_log.value, delta_time_bias_2d) + if has_seg: + kda_args = kda_args + (decoder_segment_ids,) + 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/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1e06020e72..63e0f653b0 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -29,6 +29,7 @@ import jax.numpy as jnp from jax.sharding import Mesh from maxtext.common.common_types import ( + AttentionType, Config, DecoderBlockType, MODEL_MODE_AUTOREGRESSIVE, @@ -38,7 +39,7 @@ ShardMode, ) from maxtext.configs.types import check_forced_routing_support -from maxtext.layers import initializers, linears, mhc, moe, normalizations, quantizations +from maxtext.layers import attention_kda, initializers, linears, mhc, moe, normalizations, quantizations from maxtext.layers import nnx_scan, nnx_wrappers from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding @@ -90,6 +91,7 @@ def __init__( model_mode: str, quant: None | Quant = None, name: str = "decoder_layer", + attention_type: AttentionType | str | None = None, *, rngs: nnx.Rngs, ): @@ -100,6 +102,12 @@ def __init__( cfg = self.config + # Per-layer attention type override (hybrid models); defaults to the + # global config value. + self.attention_type = ( + AttentionType(attention_type) if attention_type is not None else AttentionType(cfg.attention_type) + ) + self.pre_self_attention_norm = RMSNorm( num_features=cfg.emb_dim, dtype=cfg.dtype, @@ -109,34 +117,45 @@ def __init__( rngs=rngs, ) - self.self_attention = Attention( - config=self.config, - num_query_heads=cfg.num_query_heads, - num_kv_heads=cfg.num_kv_heads, - head_dim=cfg.head_dim, - max_target_length=cfg.max_target_length, - max_prefill_predict_length=cfg.max_prefill_predict_length, - attention_kernel=cfg.attention, - inputs_q_shape=(1, 1, cfg.emb_dim), - inputs_kv_shape=(1, 1, cfg.emb_dim), - mesh=mesh, - dtype=cfg.dtype, - weight_dtype=cfg.weight_dtype, - dropout_rate=cfg.dropout_rate, - float32_qk_product=cfg.float32_qk_product, - float32_logits=cfg.float32_logits, - quant=self.quant, - kv_quant=quantizations.configure_kv_quant(cfg), - prefill_cache_axis_order=tuple(map(int, cfg.prefill_cache_axis_order.split(","))), - ar_cache_axis_order=tuple(map(int, cfg.ar_cache_axis_order.split(","))), - compute_axis_order=tuple(map(int, cfg.compute_axis_order.split(","))), - reshape_q=cfg.reshape_q, - use_mrope=cfg.use_mrope, - mrope_section=cfg.mrope_section, - share_kv_projections=cfg.share_kv_projections, - model_mode=model_mode, - rngs=rngs, - ) + if self.attention_type == AttentionType.KDA: + # KDA is a self-contained recurrent attention layer (its own QKV/conv/ + # gate/beta stack + tokamax Delta-Rule kernel); it carries recurrent + # state instead of a KV cache. + self.self_attention = attention_kda.KimiDeltaAttention( + config=self.config, + layer_idx=0, + mesh=mesh, + rngs=rngs, + ) + else: + self.self_attention = Attention( + config=self.config, + num_query_heads=cfg.num_query_heads, + num_kv_heads=cfg.num_kv_heads, + head_dim=cfg.head_dim, + max_target_length=cfg.max_target_length, + max_prefill_predict_length=cfg.max_prefill_predict_length, + attention_kernel=cfg.attention, + inputs_q_shape=(1, 1, cfg.emb_dim), + inputs_kv_shape=(1, 1, cfg.emb_dim), + mesh=mesh, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + dropout_rate=cfg.dropout_rate, + float32_qk_product=cfg.float32_qk_product, + float32_logits=cfg.float32_logits, + quant=self.quant, + kv_quant=quantizations.configure_kv_quant(cfg), + prefill_cache_axis_order=tuple(map(int, cfg.prefill_cache_axis_order.split(","))), + ar_cache_axis_order=tuple(map(int, cfg.ar_cache_axis_order.split(","))), + compute_axis_order=tuple(map(int, cfg.compute_axis_order.split(","))), + reshape_q=cfg.reshape_q, + use_mrope=cfg.use_mrope, + mrope_section=cfg.mrope_section, + share_kv_projections=cfg.share_kv_projections, + model_mode=model_mode, + rngs=rngs, + ) self.mlp = linears.MlpBlock( in_features=cfg.emb_dim, @@ -194,16 +213,26 @@ def __call__( lnx = self.pre_self_attention_norm(inputs) lnx = _maybe_shard_with_logical(lnx, logical_axis_names) - attention_lnx, kv_cache = self.self_attention( - lnx, - lnx, - decoder_positions, - decoder_segment_ids=decoder_segment_ids, - deterministic=deterministic, - model_mode=model_mode, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - ) + if self.attention_type == AttentionType.KDA: + # KDA has no KV cache; its recurrent state is carried by the kernel. + attention_lnx, _ = self.self_attention( + lnx, + decoder_positions, + deterministic=deterministic, + model_mode=model_mode, + decoder_segment_ids=decoder_segment_ids, + ) + else: + attention_lnx, kv_cache = self.self_attention( + lnx, + lnx, + decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + model_mode=model_mode, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) attention_lnx = _maybe_shard_with_logical(attention_lnx, logical_axis_names) mlp_lnx = self.mlp(lnx, deterministic=deterministic) diff --git a/tests/unit/kda_attention_test.py b/tests/unit/kda_attention_test.py new file mode 100644 index 0000000000..aa1aec6f17 --- /dev/null +++ b/tests/unit/kda_attention_test.py @@ -0,0 +1,1665 @@ +# Copyright 2026 Ant Group. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for 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 + - Context parallelism: kernel-level CP equivalence and load_balance rejection + +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 +from types import SimpleNamespace + +import pytest +import jax +import jax.numpy as jnp +import numpy as np +import optax +import ml_dtypes +from flax import nnx + +try: + from tokamax._src.ops.experimental.kda import api as tokamax_kda_api + from tokamax._src.ops.experimental.kda.cp_utils import ContextParallelMetadata + + TOKAMAX_AVAILABLE = True +except ImportError: + tokamax_kda_api = None + ContextParallelMetadata = None + TOKAMAX_AVAILABLE = False + +from maxtext.configs.types import KdaAttention +from maxtext.kernels.kda import chunk_kda +from maxtext.kernels.kda.tokamax import tokamax_chunk_kda +from maxtext.layers import attention_kda +from maxtext.layers.attention_kda import ShortConvolution, _l2_normalize, halo_exchange_for_conv +from maxtext.layers.normalizations import RMSNorm + +# Marker policy: `tpu_only` is applied per test/class — only where a test +# invokes the tokamax Mosaic Pallas kernel or multi-device CP. Pure +# config / pure-op / non-CP tests (init checks, the naive recurrence, the L2 +# norm helper, standalone ShortConvolution, config guards) also run in +# regular CPU CI, keeping fast regression coverage outside TPU testbeds. + +# --------------------------------------------------------------------------- +# 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. + + Diagnostics are collected silently and surfaced only in the assertion + failure message, so passing tests keep normal pytest output quiet. + """ + 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()) + + close_mask = diff <= atol + rtol * np.abs(expected_f32) + if close_mask.all(): + 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 + + assert over_rate <= max_ulp_fail_rate, ( + f"{label}: max_abs={max_abs:.6e} mean_abs={mean_abs:.6e}; " + f"{n_fail}/{n_total} elements fail allclose (atol={atol}, rtol={rtol}), " + f"{n_mis} have bf16 ULP diff, worst_ulp={worst_ulp}, " + f"{n_over}/{n_fail} elements ({over_rate:.2e}) exceed {max_ulp} ULP " + f"(threshold {max_ulp_fail_rate:.2e})" + ) + + +def _assert_rel_l2_close(actual, expected, label, tol=2e-2): + """Assert the relative L2 distance between two arrays is at most ``tol``. + + Suited for gradient comparisons whose per-element absolute tails are + dominated by accumulation (e.g. weight gradients summed over tokens): the + norm ratio is scale- and tail-insensitive, while genuine wiring errors + (relative diff O(1)) still fail hard. + """ + a = np.asarray(actual, dtype=np.float32).ravel() + e = np.asarray(expected, dtype=np.float32).ravel() + rel = float(np.linalg.norm(a - e) / max(float(np.linalg.norm(e)), 1e-8)) + assert rel <= tol, ( + f"{label}: relative L2 diff {rel:.3e} exceeds {tol:.3e} " f"(max_abs={float(np.abs(a - e).max()):.6e})" + ) + + +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.max_segments_per_seq = 25 + self.context_sharding = "context" + + 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) + 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 + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + def test_forward_no_conv(self, mesh): + """Forward with linear_conv_kernel_dim=0 (the conv path is skipped entirely).""" + attn = self._make_attn(mesh, linear_conv_kernel_dim=0) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + output, _ = attn(x) + assert output.shape == (1, 64, 128) + assert jnp.isfinite(output).all() + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + def test_lower_bound_warns_without_safe_gate(self, mesh): + """use_kda_safe_gate=False with a non-zero kda_lower_bound must warn.""" + attn = self._make_attn(mesh, use_kda_safe_gate=False, kda_lower_bound=-1.0) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + with pytest.warns(UserWarning, match="use_kda_safe_gate=False"): + output, _ = attn(x) + assert jnp.isfinite(output).all() + + 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") + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + 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)) + output, aux = attn(x) + assert output.shape == (B, T, D) + assert aux is None + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + def test_forward_no_nan_inf(self, mesh): + attn = self._make_attn(mesh) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + output, _ = attn(x) + assert not jnp.any(jnp.isnan(output)) + assert not jnp.any(jnp.isinf(output)) + assert jnp.any(output != 0) + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + 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 the KDA chunk alignment (64) + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, D)) + output, _ = attn(x) + assert output.shape == (B, T, D) + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + def test_deterministic(self, mesh): + attn = self._make_attn(mesh) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + o1, _ = attn(x) + o2, _ = attn(x) + assert jnp.allclose(o1, o2, atol=1e-5) + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + 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, + ) + 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() + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + 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 the KDA chunk alignment (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) + 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() + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + 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)) + o1, _ = attn(x, decoder_segment_ids=None) + o2, _ = attn(x) # Default None + assert jnp.allclose(o1, o2, atol=1e-5) + + @pytest.mark.tpu_only + @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) + + 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" + ) + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + def test_packed_segment_no_leak_within_row(self, mesh): + """Within a single row, changing tokens in segment 2 must not affect segment 1's + output (and vice versa). Complements test_row_independence (cross-row) by + proving packed segments inside ONE row are structurally isolated.""" + attn = self._make_attn(mesh) + T, hidden_dim = 64, 128 + + # One row with two packed segments: positions [0,32)=seg1, [32,64)=seg2. + seg = jnp.array([[1] * 32 + [2] * 32], dtype=jnp.int32) + x_base = jax.random.normal(jax.random.PRNGKey(2), (1, T, hidden_dim)) + + key = jax.random.PRNGKey(3) + x_mod2 = x_base.at[0, 32:, :].set(jax.random.normal(key, (32, hidden_dim))) # change seg2 tokens + x_mod1 = x_base.at[0, :32, :].set(jax.random.normal(key, (32, hidden_dim))) # change seg1 tokens + + o_base, _ = attn(x_base, decoder_segment_ids=seg) + o_mod2, _ = attn(x_mod2, decoder_segment_ids=seg) + o_mod1, _ = attn(x_mod1, decoder_segment_ids=seg) + + # Changing segment 2 must leave segment 1's positions bit-exact unchanged. + assert jnp.allclose( + o_base[0, :32], o_mod2[0, :32], atol=0.0 + ), "Segment 1 output changed when only segment 2's tokens changed (same row)." + # Symmetric: changing segment 1 must leave segment 2 unchanged. + assert jnp.allclose( + o_base[0, 32:], o_mod1[0, 32:], atol=0.0 + ), "Segment 2 output changed when only segment 1's tokens changed (same row)." + + 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 +# --------------------------------------------------------------------------- + + +@pytest.mark.tpu_only +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): + + 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) + 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.tpu_only + @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.""" + + 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) + 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.tpu_only + @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).""" + + 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) + 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",)) + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + 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) + 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)) + output, _ = attn(x) + assert output.shape == (B, T, D) + assert not jnp.any(jnp.isnan(output)) + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + def test_qk_l2norm_always_on_regardless_of_flag(self, mesh): + """KDA always L2-normalizes Q/K: the shared use_qk_norm flag does not turn it off. + + The Delta-Rule recurrence diverges without bounded q/k, so the norm is part + of the KDA architecture regardless of the dot-product-attention flag. + """ + 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) + 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) + attn_off = attention_kda.KimiDeltaAttention(config=cfg_off, layer_idx=0, mesh=mesh, rngs=rngs_off) + out_off, _ = attn_off(x) + + # Identical weights + input, only the flag differs -> outputs must match. + assert jnp.allclose(out_on, out_off, atol=1e-5), "KDA must L2-normalize Q/K regardless of use_qk_norm" + + def test_l2_normalize_produces_unit_norm(self): + """Direct check that _l2_normalize yields unit L2 norm along the last axis. + + The layer applies this to Q/K before the kernel; verifying the helper's + norm (not just output shape/NaN) is the actual correctness property. + """ + + x = jax.random.normal(jax.random.PRNGKey(11), (2, 16, 4, 128)) + normed = _l2_normalize(x) + + norms = jnp.linalg.norm(normed.astype(jnp.float32), axis=-1) + _assert_close(norms, jnp.ones_like(norms), "l2_unit_norm", atol=1e-4, rtol=1e-4) + + # Direction preserved: normalized vector stays parallel to the input. + in_norm = jnp.linalg.norm(x.astype(jnp.float32), axis=-1, keepdims=True) + expected = x.astype(jnp.float32) / in_norm + _assert_close(normed.astype(jnp.float32), expected, "l2_direction", atol=1e-4, rtol=1e-4) + + +# --------------------------------------------------------------------------- +# Backward (VJP) tests +# --------------------------------------------------------------------------- + + +@pytest.mark.tpu_only +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) + 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) + 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)) + _, 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)) + _, 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)) + 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) + 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" + + +# --------------------------------------------------------------------------- +# Full-layer kernel parity (Mosaic vs tokamax XLA reference) +# --------------------------------------------------------------------------- + + +@pytest.mark.tpu_only +class TestKdaLayerParity: + """Full-layer parity between the Mosaic Pallas kernel and the tokamax XLA + reference implementation. + + Kernel-vs-recurrent-reference parity (TestNaiveKda) validates the kernel + in isolation, and the CP equivalence tests validate sharding. Neither + exercises the *composed layer* (QKV projection, ShortConvolution, gate / + beta transforms, output RMSNorm + gating, output projection) against an + independent implementation. Running the same layer with identical weights + on both tokamax implementations catches composition-level argument and + constraint bugs that kernel-only tests cannot see. + """ + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + def test_full_layer_mosaic_vs_xla_parity(self, monkeypatch): + mesh = jax.sharding.Mesh(jax.devices(), ("x",)) + cfg = _MockKdaConfig() + rngs = nnx.Rngs(0) + attn = attention_kda.KimiDeltaAttention(config=cfg, layer_idx=0, mesh=mesh, rngs=rngs) + + B, T, D = 2, 128, 128 + x = jax.random.normal(jax.random.PRNGKey(13), (B, T, D)) + # Two packed segments so varlen handling is part of the parity check. + seg_ids = jnp.array([[1] * 64 + [2] * 64, [1] * 128], dtype=jnp.int32) + + def _loss(model, x_in): + o, _ = model(x_in, decoder_segment_ids=seg_ids) + return o.astype(jnp.float32).sum() + + # Input gradients: differentiate the output sum w.r.t. the input. + def _loss_x(x_in): + return _loss(attn, x_in) + + # Weight gradients via the split-params pattern (see TestKdaBackward). + graphdef, params, other = nnx.split(attn, nnx.Param, ...) + + def _loss_params(params): + return _loss(nnx.merge(graphdef, params, other), x) + + def _run(): + o, _ = attn(x, decoder_segment_ids=seg_ids) + return ( + jax.device_get(o), + jax.device_get(jax.grad(_loss_x)(x)), + [jax.device_get(g) for g in jax.tree.leaves(jax.grad(_loss_params)(params))], + ) + + # --- Mosaic (production path, the adapter's default) --- + o_mosaic, grad_x_mosaic, leaves_mosaic = _run() + + # --- XLA reference implementation, same module and weights --- + orig = tokamax_kda_api.kimi_delta_attention + + def _xla_impl(*args, **kwargs): + kwargs["implementation"] = "xla" + return orig(*args, **kwargs) + + # The adapter imports kimi_delta_attention lazily inside the function, + # so patching the module attribute switches the implementation. + monkeypatch.setattr(tokamax_kda_api, "kimi_delta_attention", _xla_impl) + o_xla, grad_x_xla, leaves_xla = _run() + + # The two tokamax implementations are not bitwise identical — tokamax's + # own CI validates the mosaic kernel against the XLA reference at ~5e-3 + # RMS, and the layer's gated norm/projections propagate that rounding. + # Use a generous absolute tolerance; gross wiring errors deviate by O(0.1) + # and still fail hard. + _assert_close(o_mosaic, o_xla, "full_layer_mosaic_vs_xla_fwd", atol=2e-2, rtol=1e-2) + # Gradients: relative L2 norm comparison (accumulated weight-gradient + # tails make absolute tolerances unreliable here; see helper docstring). + _assert_rel_l2_close(grad_x_mosaic, grad_x_xla, "full_layer_mosaic_vs_xla_dx", tol=2e-2) + + assert len(leaves_mosaic) == len(leaves_xla) + for i, (gm, gx) in enumerate(zip(leaves_mosaic, leaves_xla)): + _assert_rel_l2_close(gm, gx, f"full_layer_mosaic_vs_xla_param_{i}", tol=2e-2) + + +# --------------------------------------------------------------------------- +# 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.""" + + 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" + + def test_short_conv_rejects_wrong_features(self): + """Input with the wrong feature count must fail loudly, not silently.""" + rngs = nnx.Rngs(0) + conv = ShortConvolution(kernel_size=4, features=8, dtype=jnp.float32, weight_dtype=jnp.float32, rngs=rngs) + with pytest.raises(ValueError, match="Input features"): + conv(jnp.zeros((1, 16, 7))) + + def test_short_conv_kernel_size_one(self): + """kernel_size=1 -> halo_size=0: the exchange returns the input untouched.""" + rngs = nnx.Rngs(0) + conv = ShortConvolution(kernel_size=1, features=8, dtype=jnp.float32, weight_dtype=jnp.float32, rngs=rngs) + B, T = 2, 16 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, 8)) + out = conv(x) + assert out.shape == (B, T, 8) + assert jnp.isfinite(out).all() + + def test_halo_exchange_single_rank_shard_map(self): + """Inside a shard_map whose CP axis has size 1 the exchange degrades to zero-pad.""" + mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("context",)) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 8, 4)) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=jax.sharding.PartitionSpec(None, "context", None), + out_specs=jax.sharding.PartitionSpec(None, "context", None), + check_vma=False, + ) + def _exchange(x_local): + return halo_exchange_for_conv(x_local, halo_size=2, axis_name="context", seq_axis=1) + + out = _exchange(x) + assert out.shape == (1, 10, 4) + assert jnp.allclose(out[:, :2, :], 0.0, atol=0.0), "halo rows must be zeros on the only rank" + assert jnp.allclose(out[:, 2:, :], x, atol=0.0) + + @pytest.mark.tpu_only + @pytest.mark.skipif(len(jax.devices()) < 2, reason="need >=2 devices for CP test") + @pytest.mark.parametrize("layout", ["uniform", "rank_boundary", "spanning_ranks"]) + def test_short_conv_cp_halo(self, layout): + """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. + """ + + 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)) + + # Segment layouts exercising different halo/masking interactions. + # T_local = T // cp_size = 16 per rank. + if layout == "uniform": + # No cross-segment masking; halo tokens are the true left context. + seg_ids = jnp.ones((B, T), dtype=jnp.int32) + elif layout == "rank_boundary": + # Segment boundary exactly at the rank boundary: rank 1's first + # tokens must NOT attend into rank 0 (halo must be masked out). + seg_ids = jnp.broadcast_to(jnp.array([1] * 16 + [2] * 16, dtype=jnp.int32), (B, T)) + elif layout == "spanning_ranks": + # Segment 1 spans both ranks: rank 1's leading tokens of segment 1 + # MUST read rank 0's tail through the halo. + seg_ids = jnp.broadcast_to(jnp.array([1] * 24 + [2] * 8, dtype=jnp.int32), (B, T)) + else: + raise ValueError(f"unknown layout {layout}") + seg_ids = seg_ids.astype(jnp.int32) + + # Reference: conv on the full (non-sharded) sequence with the same + # segment_ids. + ref_out = jax.device_get(conv(x, segment_ids=seg_ids)) + + # 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)), + ) + + 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())) + ) + + # For every layout, CP conv with halo (and halo'd segment masking) must + # match the non-sharded reference exactly — halos carry the true left + # context, and segment masks must suppress cross-segment halo reads. + assert jnp.allclose(cp_out_full, ref_out, atol=1e-5), ( + f"ShortConvolution CP halo output differs from reference " + f"(layout={layout}). " + f"max_diff={float(jnp.abs(cp_out_full - ref_out).max()):.2e}" + ) + + @pytest.mark.tpu_only + @pytest.mark.skipif(len(jax.devices()) < 2, reason="need >=2 devices for CP test") + def test_short_conv_cp_rejects_oversized_halo(self): + """halo_size > T_local under CP must fail clearly, not silently convolve + with wrong context. + + The halo exchange only reads from the immediately preceding rank; a + receptive field (kernel_size - 1) larger than the per-rank sequence + length would span multiple ranks, which is not implemented. + """ + + 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 = 8, 8 # halo_size = 7 > T_local = 4 below + rngs = nnx.Rngs(0) + conv = ShortConvolution( + kernel_size=kernel_size, + features=features, + dtype=jnp.float32, + weight_dtype=jnp.float32, + rngs=rngs, + ) + + B, T = 1, 8 # T_local = T // cp_size = 4 < kernel_size - 1 = 7 + x = jax.random.normal(jax.random.PRNGKey(0), (B, T, features)) + xs = jax.lax.with_sharding_constraint( + x, + jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "context", None)), + ) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=jax.sharding.PartitionSpec(None, "context", None), + out_specs=jax.sharding.PartitionSpec(None, "context", None), + check_vma=False, + ) + def _conv_cp(x_local): + return conv(x_local) + + with pytest.raises(ValueError, match="halo_size"): + _conv_cp(xs) + + +# --------------------------------------------------------------------------- +# CP (Context Parallelism) tests +# --------------------------------------------------------------------------- + + +@pytest.mark.tpu_only +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")) + + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="tokamax not available") + @pytest.mark.skipif(len(jax.devices()) < 2, reason="need >=2 devices for CP test") + @pytest.mark.parametrize("cp_size", [2, 4]) + def test_kda_cp_equivalence(self, cp_size): + """KDA with CP should produce equivalent output to non-CP KDA.""" + if len(jax.devices()) < cp_size: + pytest.skip(f"need >={cp_size} devices for CP={cp_size}") + + mesh_cp = self._cp_mesh(cp_size=cp_size) + + 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, segment_ids=seg_ids, max_num_segments=1) + + # --- CP run: shard along T, call chunk_kda with context_parallel_metadata --- + + cp_ctx = ContextParallelMetadata(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, + segment_ids=seg_loc, + max_num_segments=1, + context_parallel_metadata=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") + @pytest.mark.skipif(len(jax.devices()) < 2, reason="need >=2 devices for CP test") + def test_kda_cp_backward(self): + """CP backward: gradients match the non-CP reference (fwd+bwd through the CP kernels).""" + + cp_size = 2 + mesh_cp = self._cp_mesh(cp_size=cp_size) + + B, T, H, K, V = 2, 128, 4, 128, 128 + key = jax.random.PRNGKey(43) + 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) + args = (q, k, v, g, beta) + + # --- Reference: non-CP gradients --- + def _loss_ref(q, k, v, g, beta): + o, _ = chunk_kda(q, k, v, g, beta, scale=scale, segment_ids=seg_ids, max_num_segments=1) + return o.astype(jnp.float32).sum() + + ref_grads = jax.grad(_loss_ref, argnums=(0, 1, 2, 3, 4))(*args) + + # --- CP gradients: shard inputs along T, run the kernel under CP --- + cp_ctx = ContextParallelMetadata(mesh=mesh_cp, axis_name="context") + 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)) + + def _loss_cp(q, k, v, g, beta): + @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, + segment_ids=seg_loc, + max_num_segments=1, + context_parallel_metadata=cp_ctx, + ) + return o_loc + + o = _kda_cp( + _shard(q, pspec_4d), + _shard(k, pspec_4d), + _shard(v, pspec_4d), + _shard(g, pspec_4d), + _shard(beta, pspec_3d), + _shard(seg_ids, pspec_2d), + ) + return o.astype(jnp.float32).sum() + + cp_grads = jax.grad(_loss_cp, argnums=(0, 1, 2, 3, 4))(*args) + + # dq/dk/dv tighter than dg/dbeta (matches tokamax CI tolerances). + names = ("dq", "dk", "dv", "dg", "dbeta") + tols = (8e-3, 8e-3, 8e-3, 2e-2, 2e-2) + for name, tol, cg, rg in zip(names, tols, cp_grads, ref_grads): + _assert_close(jax.device_get(cg), jax.device_get(rg), f"kda_cp_bwd_{name}", atol=tol, rtol=1e-3) + + @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_full_layer_dummy_segments(self): + """Full KimiDeltaAttention under CP without user segment_ids. + + Covers the layer's internal dummy-segment synthesis path (cp_size > 1 + with decoder_segment_ids=None): conv halo exchange, T-axis pspec + injection and kernel-side CP metadata derivation must combine to + reproduce the non-CP output, and the CP backward must be finite. + """ + cp_size = 2 + mesh_cp = self._cp_mesh(cp_size=cp_size) + B, T, D = 2, 128, 128 + x = jax.random.normal(jax.random.PRNGKey(7), (B, T, D)) + + def _build(mesh): + # Same rng seed on both meshes -> identical weights. + # head_dim must be a multiple of 128 under CP (mosaic kernel constraint). + cfg = _MockKdaConfig(head_dim=128) + rngs = nnx.Rngs(0) + return attention_kda.KimiDeltaAttention(config=cfg, layer_idx=0, mesh=mesh, rngs=rngs) + + attn_cp = _build(mesh_cp) + # Non-CP reference: same weights on a mesh without a CP axis. + mesh_ref = jax.sharding.Mesh(np.array(jax.devices()), ("x",)) + attn_ref = _build(mesh_ref) + + o_cp, _ = attn_cp(x) + o_full = jax.device_get( + jax.lax.with_sharding_constraint(o_cp, jax.sharding.NamedSharding(mesh_cp, jax.sharding.PartitionSpec())) + ) + o_ref, _ = attn_ref(x) + + _assert_close(o_full, jax.device_get(o_ref), "kda_cp_full_layer_dummy_seg", atol=5e-3, rtol=1e-3) + assert not np.any(np.isnan(o_full)), "NaN in full-layer CP output" + + # CP backward through the whole layer (conv shard_map + kernel shard_map). + def _sum_cp(x): + o, _ = attn_cp(x) + return o.astype(jnp.float32).sum() + + grad_x = jax.device_get(jax.grad(_sum_cp)(x)) + assert np.all(np.isfinite(grad_x)), "non-finite gradient through full-layer CP backward" + + @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_full_layer_packed_segments(self): + """Full KimiDeltaAttention under CP with multiple real packed segments. + + Unlike the dummy-segment test, the layouts here stress the composition + of CP mechanics inside one layer: + - row 0: segment 2 spans the rank boundary, so recurrent state and + conv halo must cross ranks within one segment; + - row 1: a segment boundary exactly at the rank split, so the halo + tokens pulled across ranks belong to a different segment and must + be masked out by the conv segment logic, and the kernel must reset + recurrent state at the rank edge. + Forward output and input/weight gradients must match the non-CP run. + """ + cp_size = 2 + mesh_cp = self._cp_mesh(cp_size=cp_size) + mesh_ref = jax.sharding.Mesh(np.array(jax.devices()), ("x",)) + B, T, D = 2, 128, 128 + x = jax.random.normal(jax.random.PRNGKey(11), (B, T, D)) + + def _build(mesh): + # Same rng seed on both meshes -> identical weights. + # head_dim must be a multiple of 128 under CP (mosaic kernel constraint). + cfg = _MockKdaConfig(head_dim=128) + rngs = nnx.Rngs(0) + return attention_kda.KimiDeltaAttention(config=cfg, layer_idx=0, mesh=mesh, rngs=rngs) + + attn_cp = _build(mesh_cp) + attn_ref = _build(mesh_ref) + + # T_local = 64 per rank. + seg_ids = jnp.array( + [ + # seg 2 [30, 90) spans the rank boundary at 64. + [1] * 30 + [2] * 60 + [3] * 38, + # segment boundary exactly at the rank split: [1]*64 | [2]*64. + [1] * 64 + [2] * 64, + ], + dtype=jnp.int32, + ) + + # --- Forward: CP vs non-CP --- + o_cp, _ = attn_cp(x, decoder_segment_ids=seg_ids) + o_cp_full = jax.device_get( + jax.lax.with_sharding_constraint(o_cp, jax.sharding.NamedSharding(mesh_cp, jax.sharding.PartitionSpec())) + ) + o_ref, _ = attn_ref(x, decoder_segment_ids=seg_ids) + # Full-layer tolerance: 2x the kernel-level CP tolerance (5e-3, tokamax CI + # baseline) since the gated norm / output projection propagate the + # cross-rank chunk-boundary rounding through the rest of the layer. + _assert_close(o_cp_full, jax.device_get(o_ref), "kda_cp_full_layer_packed_seg_fwd", atol=1e-2, rtol=1e-3) + assert not np.any(np.isnan(o_cp_full)), "NaN in full-layer CP packed-segment output" + + # --- Backward: input and weight gradients, CP vs non-CP --- + def _grads(attn, mesh): + graphdef, params, other = nnx.split(attn, nnx.Param, ...) + + def loss_fn(params, x): + model = nnx.merge(graphdef, params, other) + o, _ = model(x, decoder_segment_ids=seg_ids) + return o.astype(jnp.float32).sum() + + _, (grad_params, grad_x) = jax.value_and_grad(loss_fn, argnums=(0, 1))(params, x) + return jax.tree.leaves(grad_params), jax.device_get(grad_x) + + leaves_cp, grad_x_cp = _grads(attn_cp, mesh_cp) + leaves_ref, grad_x_ref = _grads(attn_ref, mesh_ref) + + _assert_close(grad_x_cp, grad_x_ref, "kda_cp_full_layer_packed_seg_dx", atol=2e-2, rtol=1e-2) + + # Weight gradients are summed over the whole sequence, so their absolute + # tails are larger than per-token kernel tolerances; compare by + # relative L2 norm instead. + assert len(leaves_cp) == len(leaves_ref), "CP and non-CP param gradient trees differ" + for i, (gc, gr) in enumerate(zip(leaves_cp, leaves_ref)): + _assert_rel_l2_close( + jax.device_get(gc), + jax.device_get(gr), + f"kda_cp_full_layer_packed_seg_param_{i}", + tol=2e-2, + ) + + @pytest.mark.skipif(len(jax.devices()) < 2, reason="need >=2 devices for CP test") + def test_kda_cp_requires_tokamax_metadata(self, monkeypatch): + """If ContextParallelMetadata failed to import, CP must refuse to run. + + CP without cross-rank metadata would silently split the recurrent + state across ranks, so the layer raises instead of degrading. + """ + mesh = self._cp_mesh(cp_size=2) + cfg = _MockKdaConfig(head_dim=128) + rngs = nnx.Rngs(0) + attn = attention_kda.KimiDeltaAttention(config=cfg, layer_idx=0, mesh=mesh, rngs=rngs) + + monkeypatch.setattr(attention_kda, "TokamaxContextParallelMetadata", None) + x = jax.random.normal(jax.random.PRNGKey(0), (2, 128, 128)) + with pytest.raises(ImportError, match="ContextParallelMetadata"): + attn(x) + + @pytest.mark.skipif(len(jax.devices()) < 2, reason="need >=2 devices for CP test") + def test_kda_cp_rejects_load_balance(self): + """KDA CP should raise ValueError when load_balance is enabled.""" + mesh = self._cp_mesh(cp_size=2) + cfg = _MockKdaConfig(context_parallel_load_balance=True) + rngs = nnx.Rngs(0) + 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() + rngs = nnx.Rngs(0) + attn = attention_kda.KimiDeltaAttention( + config=cfg, + layer_idx=0, + mesh=mesh, + rngs=rngs, + ) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + output, _ = attn(x) + assert output.shape == (1, 64, 128) + assert jnp.isfinite(output).all() + + +# --------------------------------------------------------------------------- +# Config guard tests +# --------------------------------------------------------------------------- + + +class TestKdaConfigGuards: + """Config-time guards for invalid KDA combinations.""" + + def test_safe_gate_requires_valid_lower_bound(self): + """use_kda_safe_gate=True with kda_lower_bound outside [-5,0) must be rejected.""" + + with pytest.raises(ValueError, match="kda_lower_bound"): + KdaAttention(use_kda_safe_gate=True, kda_lower_bound=0.0) + with pytest.raises(ValueError, match="kda_lower_bound"): + KdaAttention(use_kda_safe_gate=True, kda_lower_bound=-6.0) + # Valid combinations pass. + KdaAttention(use_kda_safe_gate=True, kda_lower_bound=-5.0) + KdaAttention(use_kda_safe_gate=True, kda_lower_bound=-1.0) + KdaAttention(use_kda_safe_gate=False) + + def test_use_kda_lora_true_rejected(self): + """use_kda_lora=True is an unimplemented no-op and must be rejected.""" + + with pytest.raises(ValueError, match="use_kda_lora"): + KdaAttention(use_kda_lora=True) + KdaAttention(use_kda_lora=False) + + @pytest.fixture + def mesh(self): + return jax.sharding.Mesh(jax.devices(), ("x",)) + + def test_packing_requires_max_segments_per_seq(self, mesh): + """Layer must fail fast when packed sequences are used without max_segments_per_seq.""" + cfg = _MockKdaConfig(max_segments_per_seq=-1) + rngs = nnx.Rngs(0) + attn = attention_kda.KimiDeltaAttention(config=cfg, layer_idx=0, mesh=mesh, rngs=rngs) + x = jax.random.normal(jax.random.PRNGKey(0), (1, 64, 128)) + seg = jnp.ones((1, 64), dtype=jnp.int32) + with pytest.raises(ValueError, match="max_segments_per_seq"): + attn(x, decoder_segment_ids=seg) + + +# --------------------------------------------------------------------------- +# Kernel input-guard tests (raise before any kernel call; run in CPU CI) +# --------------------------------------------------------------------------- + + +class TestKdaKernelGuards: + """``initial_state`` / ``output_final_state`` are rejected before dispatch. + + These guards fire before any tokamax kernel is touched, so they are safe + to run on CPU (no tpu_only marker). + """ + + @staticmethod + def _dummy_inputs(): + """Small random q/k/v/g/beta tensors shaped for the kernel interface.""" + B, T, H, K, V = 1, 64, 2, 16, 16 + key = jax.random.PRNGKey(0) + keys = jax.random.split(key, 5) + q = jax.random.normal(keys[0], (B, T, H, K)) + k = jax.random.normal(keys[1], (B, T, H, K)) + v = jax.random.normal(keys[2], (B, T, H, V)) + g = jax.random.normal(keys[3], (B, T, H, K)) + beta = jax.random.normal(keys[4], (B, T, H)) + return q, k, v, g, beta + + def test_chunk_kda_rejects_initial_state(self): + q, k, v, g, beta = self._dummy_inputs() + with pytest.raises(NotImplementedError, match="initial_state"): + chunk_kda(q, k, v, g, beta, scale=0.25, initial_state=jnp.zeros((2, 16, 16))) + + def test_chunk_kda_rejects_output_final_state(self): + q, k, v, g, beta = self._dummy_inputs() + with pytest.raises(NotImplementedError, match="output_final_state"): + chunk_kda(q, k, v, g, beta, scale=0.25, output_final_state=True) + + def test_tokamax_adapter_rejects_initial_state(self): + q, k, v, g, beta = self._dummy_inputs() + with pytest.raises(NotImplementedError, match="initial_state"): + tokamax_chunk_kda(q, k, v, g, beta, scale=0.25, initial_state=jnp.zeros((2, 16, 16))) + + def test_tokamax_adapter_rejects_output_final_state(self): + q, k, v, g, beta = self._dummy_inputs() + with pytest.raises(NotImplementedError, match="output_final_state"): + tokamax_chunk_kda(q, k, v, g, beta, scale=0.25, output_final_state=True) + + +# --------------------------------------------------------------------------- +# End-to-end training smoke (delayed-copy task; replaces the former +# scripts/dev/kda_e2e_smoke.py) +# --------------------------------------------------------------------------- + +_KDA_SMOKE_VOCAB = 128 + + +class _KdaBlock(nnx.Module): + """Pre-norm transformer block: RMSNorm -> KimiDeltaAttention -> MLP.""" + + def __init__(self, cfg, mesh, layer_idx, *, rngs): + self.attn_norm = RMSNorm( + num_features=cfg.base_emb_dim, + epsilon=cfg.normalization_layer_epsilon, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + rngs=rngs, + ) + self.attn = attention_kda.KimiDeltaAttention(cfg, layer_idx=layer_idx, mesh=mesh, rngs=rngs) + self.mlp_norm = RMSNorm( + num_features=cfg.base_emb_dim, + epsilon=cfg.normalization_layer_epsilon, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + rngs=rngs, + ) + hidden = 4 * cfg.base_emb_dim + self.wi = nnx.Linear(cfg.base_emb_dim, hidden, dtype=cfg.dtype, param_dtype=cfg.weight_dtype, rngs=rngs) + self.wo = nnx.Linear(hidden, cfg.base_emb_dim, dtype=cfg.dtype, param_dtype=cfg.weight_dtype, rngs=rngs) + + def __call__(self, x): + attn_out, _ = self.attn(self.attn_norm(x).astype(self.attn.config.dtype)) + x = x + attn_out.astype(x.dtype) + h = nnx.gelu(self.wi(self.mlp_norm(x))) + x = x + self.wo(h).astype(x.dtype) + return x + + +class _TinyKdaLM(nnx.Module): + """Embed -> N x _KdaBlock -> RMSNorm -> lm_head.""" + + def __init__(self, cfg, mesh, num_layers, *, rngs): + self.embed = nnx.Embed(_KDA_SMOKE_VOCAB, cfg.base_emb_dim, dtype=cfg.dtype, param_dtype=cfg.weight_dtype, rngs=rngs) + self.blocks = nnx.List([_KdaBlock(cfg, mesh, i, rngs=rngs) for i in range(num_layers)]) + self.final_norm = RMSNorm( + num_features=cfg.base_emb_dim, + epsilon=cfg.normalization_layer_epsilon, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + rngs=rngs, + ) + self.lm_head = nnx.Linear( + cfg.base_emb_dim, _KDA_SMOKE_VOCAB, dtype=cfg.dtype, param_dtype=cfg.weight_dtype, rngs=rngs + ) + + def __call__(self, tokens): + x = self.embed(tokens) + for block in self.blocks: + x = block(x) + return self.lm_head(self.final_norm(x)) + + +def _delayed_copy_dataset(seed, num_seqs, seq_len, delay): + """Delayed-copy sequences over random tokens: t[i] = t[i-delay].""" + rng = np.random.default_rng(seed) + total = seq_len + 1 + seqs = rng.integers(0, _KDA_SMOKE_VOCAB, size=(num_seqs, total), dtype=np.int32) + for i in range(delay, total): + seqs[:, i] = seqs[:, i - delay] + return seqs + + +class TestKdaE2eSmoke: + """End-to-end training smoke for the KimiDeltaAttention layer. + + The task is delayed copy: i.i.d. tokens with ``t[i] = t[i-delay]`` where + ``delay`` exceeds the short convolution's receptive field. Neither a + memoryless model nor the convolution alone can predict the next token, so + the loss collapses to near zero only if the KDA recurrent state carries + history — validating the full forward/backward/optimizer chain through the + real Pallas kernels. + """ + + @pytest.mark.tpu_only + @pytest.mark.skipif(not TOKAMAX_AVAILABLE, reason="KDA API not available in the installed tokamax") + def test_delayed_copy_loss_collapses(self): + seq_len, delay, steps, batch, num_layers = 64, 5, 300, 32, 2 + cfg = SimpleNamespace( + base_emb_dim=256, + base_num_query_heads=8, + head_dim=64, + dtype=jnp.float32, + weight_dtype=jnp.float32, + attention_bias=False, + shard_mode="auto", + matmul_precision="default", + normalization_layer_epsilon=1e-6, + logical_axis_rules=[], + linear_conv_kernel_dim=4, + use_qk_norm=True, + use_kda_safe_gate=True, + kda_lower_bound=-5.0, + max_segments_per_seq=25, + context_sharding="context", + ) + # Delay beyond the conv receptive field, or the task is solvable without + # any KDA state (see review of the original permutation task). + assert delay > cfg.linear_conv_kernel_dim + + mesh = jax.sharding.Mesh(np.array(jax.devices()), ("x",)) + rngs = nnx.Rngs(0) + model = _TinyKdaLM(cfg, mesh, num_layers, rngs=rngs) + + data = _delayed_copy_dataset(seed=42, num_seqs=4096, seq_len=seq_len, delay=delay) + # Label position j predicts token j+1 = t[j+1-delay]; positions with + # j+1 < delay have random targets (irreducible), so mask them out. + loss_mask = jnp.asarray(np.arange(seq_len) >= delay - 1, dtype=jnp.float32)[None, :] + optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param) + + def masked_ce(logits, labels): + ce = optax.softmax_cross_entropy_with_integer_labels(logits=logits.astype(jnp.float32), labels=labels) + return (ce * loss_mask).sum() / loss_mask.sum() / labels.shape[0] + + @nnx.jit + def train_step(model, optimizer, tokens): + def loss_fn(model): + logits = model(tokens[:, :-1]) + return masked_ce(logits, tokens[:, 1:]), logits + + (loss, _), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model) + optimizer.update(model, grads) + return loss + + perm_rng = np.random.default_rng(1) + losses = [] + for step in range(steps): + idx = perm_rng.integers(0, data.shape[0], size=batch) + loss_val = float(train_step(model, optimizer, jnp.asarray(data[idx]))) + assert np.isfinite(loss_val), f"non-finite loss {loss_val} at step {step}" + losses.append(loss_val) + + init_loss, final_loss = losses[0], float(np.mean(losses[-20:])) + assert final_loss < 0.5 * init_loss and final_loss < 1.0, ( + f"delayed-copy loss did not collapse: init={init_loss:.4f} final={final_loss:.4f} " + "(the KDA recurrent state is not carrying history through training)" + ) diff --git a/tests/unit/kda_decoder_integration_test.py b/tests/unit/kda_decoder_integration_test.py new file mode 100644 index 0000000000..d4c23dc534 --- /dev/null +++ b/tests/unit/kda_decoder_integration_test.py @@ -0,0 +1,187 @@ +# Copyright 2026 Ant Group. All Rights Reserved. +# +# 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. + +"""Decoder integration tests for KDA (attention_type='kda'). + +These exercise the *real* MaxText decoder path — `NNXDecoderLayer` selecting +`KimiDeltaAttention` via the `attention_type` dispatch — rather than the +hand-rolled block used in the standalone layer smoke test. + +CPU-runnable tests cover config acceptance, the scan_layers guard, and the +layer-branch dispatch (no kernel needed at construction). TPU-only tests run +an actual forward/backward training loop through the full model. + +Run with: python -m pytest tests/unit/kda_decoder_integration_test.py -v +""" + +import sys + +import pytest +import jax +import jax.numpy as jnp +import numpy as np +import optax +from flax import nnx +from jax.sharding import Mesh + +from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.configs import pyconfig +from maxtext.layers.attentions import Attention +from maxtext.layers.attention_kda import KimiDeltaAttention +from maxtext.layers import nnx_decoders +from maxtext.models import models +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +try: + from tokamax._src.ops.experimental.kda import api # noqa: F401 # pylint: disable=unused-import + + KDA_API_AVAILABLE = True +except ImportError: + KDA_API_AVAILABLE = False + +# Small shared hyper-parameters for the integration model. +_KDA_TEST_VOCAB = 128 +_KDA_TEST_SEQ = 64 # multiple of the KDA chunk size + + +def _kda_pyconfig(**kwargs): + """Build a tiny train config that selects the KDA attention variant.""" + defaults = { + "per_device_batch_size": 4.0, + "run_name": "test", + "enable_checkpointing": False, + "decoder_block": "default", # generic NNXDecoderLayer, where the KDA branch lives + "base_num_decoder_layers": 2, + "attention": "dot_product", + "attention_type": "kda", + "scan_layers": False, + "max_target_length": _KDA_TEST_SEQ, + "base_emb_dim": 128, + "base_num_query_heads": 4, + "base_num_kv_heads": 4, + "head_dim": 128, + "vocab_size": _KDA_TEST_VOCAB, + "max_prefill_predict_length": 4, + # Bounded-decay (safe) gate keeps the Delta-Rule recurrence stable, + # matching the standalone layer smoke. fp32 for the same reason; bf16 + # hyperparameter tuning for KDA models belongs to the model-landing + # follow-up, not the integration validation. + "use_kda_safe_gate": True, + "kda_lower_bound": -5.0, + "dtype": "float32", + "weight_dtype": "float32", + } + defaults.update(kwargs) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **defaults) + + +class TestKdaDecoderConfig: + """Config-level acceptance and guards (CPU-runnable, no kernel).""" + + def test_kda_attention_type_accepted(self): + """attention_type='kda' with scan_layers=false builds a valid config.""" + cfg = _kda_pyconfig() + assert cfg.attention_type == "kda" + assert cfg.scan_layers is False + + def test_kda_requires_scan_layers_false(self): + """KDA layers are not validated in a scanned stack; the config must reject it.""" + with pytest.raises(ValueError, match="scan_layers"): + _kda_pyconfig(scan_layers=True) + + def test_kda_decoder_layer_dispatches_to_kimi_delta(self): + """NNXDecoderLayer builds KimiDeltaAttention when attention_type='kda'.""" + cfg = _kda_pyconfig(per_device_batch_size=1.0) + mesh = Mesh(np.array(jax.devices()), (cfg.mesh_axes[0],)) + layer = nnx_decoders.NNXDecoderLayer( + config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, attention_type="kda", rngs=nnx.Rngs(0) + ) + assert isinstance(layer.self_attention, KimiDeltaAttention) + + def test_default_decoder_layer_keeps_attention(self): + """Non-KDA attention types still build the regular Attention module.""" + cfg = _kda_pyconfig(per_device_batch_size=1.0, attention_type="global") + mesh = Mesh(np.array(jax.devices()), (cfg.mesh_axes[0],)) + layer = nnx_decoders.NNXDecoderLayer(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(0)) + assert isinstance(layer.self_attention, Attention) + + +@pytest.mark.tpu_only +@pytest.mark.skipif(not KDA_API_AVAILABLE, reason="KDA API not available in the installed tokamax") +class TestKdaDecoderTraining: + """End-to-end training through the real decoder with the KDA variant. + + Uses the same history-dependent delayed-copy task as the standalone smoke: + the loss collapses only if the KDA recurrent state carries history through + the full forward/backward/optimizer chain inside the actual decoder. + """ + + @staticmethod + def _delayed_copy_dataset(seed, num_seqs, seq_len, delay): + rng = np.random.default_rng(seed) + total = seq_len + 1 + seqs = rng.integers(0, _KDA_TEST_VOCAB, size=(num_seqs, total), dtype=np.int32) + for i in range(delay, total): + seqs[:, i] = seqs[:, i - delay] + return seqs + + def test_kda_decoder_train_loss_decreases(self): + cfg = _kda_pyconfig() + delay = 5 + assert delay > cfg.linear_conv_kernel_dim + + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + model = models.Transformer(config=cfg, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(0)) + + # Every decoder layer instance must carry a KDA attention module. + for i in range(cfg.num_decoder_layers): + layer = getattr(model.decoder, f"layers_{i}") + assert isinstance(layer.self_attention, KimiDeltaAttention) + + data = self._delayed_copy_dataset(seed=42, num_seqs=512, seq_len=_KDA_TEST_SEQ, delay=delay) + # Position j predicts token j+1 = t[j+1-delay]; mask the irreducible prefix. + loss_mask = jnp.asarray(np.arange(_KDA_TEST_SEQ) >= delay - 1, dtype=jnp.float32)[None, :] + optimizer = nnx.Optimizer(model, optax.adamw(2e-3), wrt=nnx.Param) + + def masked_ce(logits, labels): + ce = optax.softmax_cross_entropy_with_integer_labels(logits=logits.astype(jnp.float32), labels=labels) + return (ce * loss_mask).sum() / loss_mask.sum() / labels.shape[0] + + @nnx.jit + def train_step(model, optimizer, tokens): + def loss_fn(model): + positions = jnp.broadcast_to(jnp.arange(_KDA_TEST_SEQ, dtype=jnp.int32)[None, :], tokens[:, :-1].shape) + logits = model(tokens[:, :-1], positions, enable_dropout=False, model_mode=MODEL_MODE_TRAIN) + return masked_ce(logits, tokens[:, 1:]), logits + + (loss, _), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model) + optimizer.update(model, grads) + return loss + + perm_rng = np.random.default_rng(1) + steps, batch = 400, int(cfg.global_batch_size_to_train_on) + losses = [] + for step in range(steps): + idx = perm_rng.integers(0, data.shape[0], size=batch) + loss_val = float(train_step(model, optimizer, jnp.asarray(data[idx]))) + assert np.isfinite(loss_val), f"non-finite loss {loss_val} at step {step}" + losses.append(loss_val) + + init_loss, final_loss = losses[0], float(np.mean(losses[-20:])) + assert final_loss < 0.5 * init_loss and final_loss < 1.0, ( + f"kda decoder loss did not collapse: init={init_loss:.4f} final={final_loss:.4f} " + "(the KDA recurrent state is not carrying history through the real decoder)" + )