From fcc112e2c8f1f4c0ef72e6356239cfbb17f26855 Mon Sep 17 00:00:00 2001 From: chiaotung97 Date: Tue, 15 Sep 2026 11:19:03 +0800 Subject: [PATCH 1/4] feat(kda): integrate KDA attention with tokamax backend and CP support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kimi Delta Attention (KDA) end to end in MaxText: the attention layer, a tokamax-backed Pallas TPU kernel adapter, context parallelism (CP), and decoder wiring via attention_type='kda'. KDA is a recurrent linear attention that updates its state with the Delta Rule, so it carries recurrent state instead of a KV cache. Layer (src/maxtext/layers/attention_kda.py): QKV / forget-gate / beta / output-gate projections, a depthwise causal short convolution with segment masking for packed sequences, SiLU, unconditional Q/K L2 normalization, the chunked recurrence via the kernel, per-head RMSNorm scaled by a sigmoid output gate, and the output projection. The stage order and gate math follow the public KDA reference (Kimi Linear, arXiv:2510.26692; fla/layers/kda.py); the class docstring lists where this implementation departs from it. Kernel (src/maxtext/kernels/kda/): a thin adapter over tokamax api.kimi_delta_attention (implementation="mosaic", Pallas/TPU). It translates between MaxText's batch-first [B, T, H, D] and tokamax's head-first [H, B, K, D] layouts and imports the API lazily: the first tokamax release shipping it is 0.0.14, the decoupled-mode CI environment pins an older tokamax on purpose, and layers/nnx_decoders.py imports this module unconditionally — so the module stays importable everywhere and using KDA without the API raises an error naming the release required, the version installed and the fix. Context parallelism: the sequence stays sharded. tokamax coordinates the recurrent state across ranks via ContextParallelMetadata, but conv-side CP is left to the caller, so halo_exchange_for_conv fetches the kernel_size-1 tokens of left context from the previous rank with one ppermute around the CP ring (rank 0 gets zeros, the true sequence start) and degrades to left zero-padding without CP. Conv and kernel run in two independent shard_maps because ppermute needs the CP axis in scope. CP rejects context_parallel_load_balance up front: the DUAL_CHUNK_SWAP reorder breaks the recurrence's token order. Config: linear_conv_kernel_dim, use_kda_safe_gate, kda_lower_bound and use_kda_lora, with cross-field validation — the sigmoid gate path needs lower_bound in [-5, 0), packed sequences need a positive max_segments_per_seq, and attention_type='kda' with scan_layers is rejected because KDA layers are not validated inside a scanned stack. use_kda_lora selects the low-rank bottleneck that the public references build their gate projections with; only the full-rank path is implemented, so it is rejected at config time instead of being silently ignored. Decoder wiring (src/maxtext/layers/nnx_decoders.py): attention_type='kda' builds KimiDeltaAttention for that layer, with no KV cache. Other attention types are unaffected. Tests (tests/unit/kda_attention_test.py, kda_decoder_integration_test.py): kernel and layer numerics against a token-by-token Delta Rule reference in fp32/bf16; forward and weight gradients under CP versus non-CP, including a segment spanning the rank boundary; full-layer Mosaic versus tokamax XLA parity; conv halo behaviour and its oversized-halo guard; config guards; the version-aware API error; and a real two-layer decoder training run on a delayed-copy task whose loss collapses, so only a history-dependent solution scores. Kernel-invoking tests are marked tpu_only and skip cleanly elsewhere. Validated on 4xTPU v6e with JAX 0.11.0 + libtpu 0.0.44.1. Design doc: docs/reference/kda_cp_support.md. Later review round, applied on top of the above: - types.py: reject two KDA misconfigurations at parse time instead of at the first forward step. `context_parallel_load_balance` defaults to true, so a KDA CP run would otherwise pass config validation and fail only after device allocation and weight init; the guard mirrors the GatedDeltaNet one. Packing without a positive `max_segments_per_seq` is likewise rejected, mirroring the existing cudnn_flash_te packing guard. The layer keeps its runtime checks. - attention_kda.py: initialize A_log per gate path, matching the reference layer — log(U(1, 16)) for the softplus gate, where exp(A_log) scales the decay magnitude; zeros (A = 1) for the safe gate, where lower_bound already bounds the decay and A would only sharpen the sigmoid. Revalidated: the end-to-end delayed-copy training run still collapses the loss on 4xv6e. - attention_kda.py: build halo_exchange_for_conv's zero left-pad lazily. Under CP with more than one rank the fetched halo replaces it, so padding up front only added a discarded array to the graph. - attention_kda.py: document that the segment_ids synthesized under CP is load-bearing twice over — tokamax aligns segments to its 64-token chunk only on the varlen path, while its non-varlen path requires T_local % 64 == 0, which global padding cannot guarantee once the sequence is sharded. - kernels/kda: chunk_kda and tokamax_chunk_kda take `implementation` (default "mosaic") so tests and debugging can select tokamax's "xla" reference. It stays off the config surface deliberately: letting tokamax choose would silently fall back to the token-by-token reference recurrence. - nnx_decoders.py: NNXDecoderLayer accepts and forwards the real layer index to KimiDeltaAttention instead of tagging every KDA layer 0; the KDA call path also sets kv_cache = None explicitly. - Tests: config-time guards for the two new validations, and one pinning the A_log initialization rule. --- docs/reference.md | 8 + docs/reference/kda_cp_support.md | 205 +++ src/maxtext/common/common_types.py | 1 + src/maxtext/configs/base.yml | 8 +- src/maxtext/configs/types.py | 109 +- src/maxtext/kernels/kda/__init__.py | 107 ++ src/maxtext/kernels/kda/tokamax.py | 229 +++ src/maxtext/layers/attention_kda.py | 854 ++++++++++ src/maxtext/layers/nnx_decoders.py | 115 +- tests/unit/kda_attention_test.py | 1749 ++++++++++++++++++++ tests/unit/kda_decoder_integration_test.py | 191 +++ 11 files changed, 3534 insertions(+), 42 deletions(-) create mode 100644 docs/reference/kda_cp_support.md create mode 100644 src/maxtext/kernels/kda/__init__.py create mode 100644 src/maxtext/kernels/kda/tokamax.py create mode 100644 src/maxtext/layers/attention_kda.py create mode 100644 tests/unit/kda_attention_test.py create mode 100644 tests/unit/kda_decoder_integration_test.py diff --git a/docs/reference.md b/docs/reference.md index e3a152a869..a7f8a97d6c 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 @@ -65,5 +72,6 @@ reference/models reference/architecture reference/core_concepts reference/mtp_cp_packing +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..14c8c12794 --- /dev/null +++ b/docs/reference/kda_cp_support.md @@ -0,0 +1,205 @@ +# 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 | +| -------------------------------------------- | ------------------------------------------------------------------------------------------- | +| `layers/attention_kda.py` | **New**: `KimiDeltaAttention`, `ShortConvolution`, and the conv halo exchange for CP | +| `kernels/kda/__init__.py` | **New**: `chunk_kda()` entry point | +| `kernels/kda/tokamax.py` | **New**: tokamax backend adapter (lazy import; version-aware error when the API is missing) | +| `layers/nnx_decoders.py` | **Modified**: `attention_type='kda'` dispatches to `KimiDeltaAttention` | +| `common/common_types.py` | **Modified**: `AttentionType.KDA` | +| `configs/types.py` | **Modified**: `KdaAttention` config class + validators | +| `configs/base.yml` | **Modified**: `attention_type` supported list + the four KDA flags | +| `tests/unit/kda_attention_test.py` | **New**: layer + conv halo + CP fwd/bwd + packed-seg CP + parity + e2e smoke test | +| `tests/unit/kda_decoder_integration_test.py` | **New**: decoder dispatch, config guards, and a training run through the real decoder | +| `docs/reference/kda_cp_support.md` | **New**: this design doc | +| `docs/reference.md` | **Modified**: toctree and card entry for this doc | + +Line counts are deliberately not tracked here: they rot on every rebase, and the PR diff is the +authoritative list. + +## Key Constraints + +1. **ContextParallelMetadata availability**: raise `kda_api_unavailable` (an `ImportError` naming the first tokamax release that ships the KDA API, the version installed, and the fix) when it is unavailable; never fall back silently. See "tokamax KDA API availability" below. + +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` keeps the module importable; using KDA then raises `kda_api_unavailable` with the required version +- segment_ids dummy: auto-construct `jnp.ones` when no varlen + CP + +## tokamax KDA API availability + +`tokamax._src.ops.experimental.kda` first shipped in **tokamax 0.0.14** (0.0.12 and 0.0.13 do not contain it, and the 0.1.0 upload was yanked). MaxText deliberately does not express that as a `tokamax>=0.0.14` requirement: + +- the floor would live in `base_requirements/`, which feeds `generate_requirements.sh`; raising it re-pins the whole lock set, JAX included, +- the decoupled-mode environment pins tokamax to an older release on purpose (`generate_decoupled_requirements.py`: newer tokamax imports `xprof` at module scope), and `layers/nnx_decoders.py` imports `attention_kda` unconditionally — so a hard failure at import time would break every NNX model there, not just KDA. + +The module-level import therefore stays soft and the failure is deferred to use time: `kda_api_unavailable` names the release required, the version installed, and the fix command. The installed version is diagnostic only — a source install of a KDA branch can report an older number while providing the API — so "predates it" is asserted only when the version string parses and compares below 0.0.14. Until the generated locks are raised, CI resolves the lowest allowed tokamax and skips the kernel-invoking tests. + +## 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` | +| `TestKdaKernelGuards` | guards that fire before any kernel dispatch: `initial_state` / `output_final_state` rejected by both `chunk_kda` and the adapter; the KDA-API-unavailable error names the required tokamax release and an actionable fix, keeps the caller's detail, and chains `__cause__` | 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 b473999dfb..3811c35e83 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -414,10 +414,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 # Low-rank gate projections (used by the public KDA references) are not implemented; true is rejected at config time 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 c022f688be..835694556c 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -649,8 +649,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, @@ -771,6 +771,87 @@ 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 implementation." + ), + ) + use_kda_lora: bool = Field( + False, + description=( + "Selects the low-rank variant of KDA's forget-gate and output-gate " + "projections (hidden -> head_dim -> num_heads*head_dim), which is what " + "the public KDA reference implementations use. Not implemented: " + "KimiDeltaAttention provides the full-rank path only, so True is " + "rejected at config time instead of being silently ignored." + ), + ) + 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: + """Guard: the low-rank KDA gate path is not implemented, so reject True at config time.""" + if v: + raise ValueError( + "use_kda_lora=True is not implemented: KimiDeltaAttention only " + "implements the full-rank gate projections, not the low-rank " + "bottleneck variant (hidden -> head_dim -> num_heads*head_dim). " + "Rejected here so the request fails at config time instead of " + "silently training a different architecture. 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): MLA or Compressed Attention with indexer.""" @@ -3282,6 +3363,7 @@ class MaxTextConfig( # Attention Mechanisms Attention, MlaAttention, + KdaAttention, CompressedAttention, MoBa, AttentionIndexer, @@ -4852,6 +4934,29 @@ 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." + ) + + kda_context_parallel_size = self.ici_context_parallelism * self.dcn_context_parallelism + if self.attention_type == "kda" and kda_context_parallel_size > 1 and self.context_parallel_load_balance: + raise ValueError( + "attention_type='kda' with context parallelism requires context_parallel_load_balance=false. " + "The KDA recurrence composes state in token order, so device i must hold the sequence chunk " + "that follows device i-1; DUAL_CHUNK_SWAP hands device 0 the first and last chunks, which " + "composes the recurrent state out of order. `context_parallel_load_balance` defaults to true, " + "so set it explicitly. The run would still train and the loss would still fall, which is why " + "this is rejected here rather than left to the layer's runtime check." + ) + if self.attention_type == "kda" and self.packing and self.max_segments_per_seq <= 0: + raise ValueError( + "attention_type='kda' with packing=true requires a positive max_segments_per_seq: the KDA " + "kernel derives per-rank segment metadata from segment_ids and needs a static upper bound " + "on the number of packed segments per sequence. Set max_segments_per_seq to that bound." + ) + 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..223c213a3b --- /dev/null +++ b/src/maxtext/kernels/kda/__init__.py @@ -0,0 +1,107 @@ +# 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, + implementation: str = "mosaic", +) -> 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. + implementation: which tokamax backend to run, ``"mosaic"`` (the Pallas + TPU kernel) or ``"xla"`` (tokamax's pure-JAX reference recurrence). + Defaults to ``"mosaic"`` and is not exposed as a config flag: asking + tokamax to choose would let it fall back to the reference + implementation silently, which is orders of magnitude slower. Pass + ``"xla"`` explicitly from tests or debugging only. + + 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, + implementation=implementation, + ) diff --git a/src/maxtext/kernels/kda/tokamax.py b/src/maxtext/kernels/kda/tokamax.py new file mode 100644 index 0000000000..ab795184ae --- /dev/null +++ b/src/maxtext/kernels/kda/tokamax.py @@ -0,0 +1,229 @@ +# 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 importlib.metadata + +import jax.numpy as jnp +from packaging.version import InvalidVersion, Version + +# `kimi_delta_attention` lives on tokamax's experimental path +# `tokamax._src.ops.experimental.kda`, which first shipped in release 0.0.14 +# (0.0.12 / 0.0.13 do not contain it, and the empty 0.1.0 upload was yanked). +# The import in `tokamax_chunk_kda` is deliberately lazy so an environment +# with an older tokamax can still import MaxText: `layers/nnx_decoders.py` +# imports this adapter's caller unconditionally, so a failure at module import +# time would break every NNX model instead of just KDA. Only an actual KDA +# call needs the API, and that call raises `kda_api_unavailable`, which names +# the version required. +# +# This is intentionally NOT expressed as a `tokamax>=0.0.14` requirement pin: +# the decoupled-mode environment pins tokamax to an older release on purpose +# (see scripts/generate_decoupled_requirements.py — newer tokamax imports +# xprof at module scope) and still has to import MaxText, and a version floor +# would also reject source installs of tokamax main, which work fine. The +# capability test is therefore the import itself; the version is diagnostic. +_MIN_TOKAMAX_VERSION = "0.0.14" + + +def _installed_tokamax_version() -> str | None: + """Return the installed tokamax version, or None when it is not installed.""" + try: + return importlib.metadata.version("tokamax") + except importlib.metadata.PackageNotFoundError: + return None + + +def _predates_kda(version: str) -> bool: + """Whether *version* is a distribution known to predate the KDA module. + + An unparsable version is deliberately not treated as predating it: a source + build (``pip install -e`` of a KDA branch) can report an older number while + still providing the API, and claiming otherwise would send the user to + ``pip install`` over a working checkout. + """ + try: + return Version(version) < Version(_MIN_TOKAMAX_VERSION) + except InvalidVersion: + return False + + +def kda_api_unavailable(cause: BaseException | None = None, *, detail: str = "") -> ImportError: + """Build the ImportError raised when the tokamax KDA API cannot be imported. + + Names the first tokamax release that ships the API and the version actually + installed, so the fix is obvious. The installed version is reported for + diagnosis only — source installs may report a placeholder version while + still providing the API, which is why this is raised on a failed import + rather than on a version comparison. + + Args: + cause: The original import failure, attached as ``__cause__``. + detail: Optional sentence prefixed to the message (e.g. what the caller + needed the API for and why degrading silently is not acceptable). + + Returns: + An ``ImportError`` with an actionable message. + """ + found = _installed_tokamax_version() + if found is None: + installed = "tokamax is not installed" + elif _predates_kda(found): + # Verified: the 0.0.12 and 0.0.13 distributions do not ship the KDA module. + installed = f"the installed tokamax {found} predates it" + else: + installed = ( + f"tokamax {found} is installed but its KDA API could not be imported — a source build " + "may be incomplete or a dependency of it may be missing" + ) + message = ( + "KDA requires the tokamax KDA API (tokamax._src.ops.experimental.kda), first shipped in " + f"tokamax {_MIN_TOKAMAX_VERSION}; {installed}. " + f"Fix with: pip install -U 'tokamax>={_MIN_TOKAMAX_VERSION}'; for a source checkout, " + "reinstall openxla/tokamax main." + ) + if detail: + message = f"{detail} {message}" + # `raise X from Y` is statement-only syntax, so chain the cause by hand: + # this mirrors what the interpreter does for that statement. + error = ImportError(message) + if cause is not None: + error.__cause__ = cause + error.__suppress_context__ = True + return error + + +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, + implementation: str = "mosaic", +) -> 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. + implementation: tokamax backend to run — ``"mosaic"`` (Pallas TPU) or + ``"xla"`` (tokamax's pure-JAX reference recurrence). The default is + fixed on purpose: passing None would let tokamax pick its + ``("mosaic", "xla")`` default, which silently falls back to the + reference recurrence — a token-by-token triple ``fori_loop`` — when + the Pallas kernel declines the shapes or hardware. Training would + still be correct and unusably slow, so a fallback has to be asked + for explicitly. + + 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 kda_api_unavailable(exc) 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=implementation, # "mosaic" resolves to the mosaic_tpu Pallas kernel + 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..2ef815b1c3 --- /dev/null +++ b/src/maxtext/layers/attention_kda.py @@ -0,0 +1,854 @@ +# 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 + - Q/K L2 normalization (always applied; see ``KimiDeltaAttention``) + +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 +from maxtext.kernels.kda.tokamax import kda_api_unavailable + +# KDA needs tokamax's KDA API (first shipped in tokamax 0.0.14), but this +# module must stay importable without it: nnx_decoders imports it +# unconditionally, so a hard failure here would break every NNX model on any +# environment with an older tokamax (the decoupled-mode CI environment is +# pinned to one deliberately). The missing API is reported at use time by +# kda_api_unavailable, which names the version required. +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 + + +# The kernel evaluates the recurrence in fixed chunks of this size, so the +# sequence is padded up to a multiple of it before the call. Keeping T a +# compile-time multiple also gives TPU-friendly static shapes. +_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: + """Give a causal convolution the tokens that come before its shard. + + A causal convolution at position ``i`` reads positions ``i - halo_size`` + through ``i``. When the sequence axis is sharded over context parallel (CP) + ranks, the first ``halo_size`` positions of a shard need context that lives + on the *previous* rank, so the shard on its own is not enough. This helper + fetches those tokens and prepends them, so the caller sees + ``[halo_size + T_local, …]`` and its per-tap loop reads the right window at + every position without having to know anything about sharding. + + Each rank contributes its own last ``halo_size`` tokens to the next rank + (one ``jax.lax.ppermute`` around the ring) and uses whatever it receives as + its left context. Rank 0 has no predecessor — it holds the true start of the + sequence — so it receives zeros, which is exactly the boundary condition a + causal convolution expects there. + + With no CP axis in scope, or a single CP rank, the same call reduces to left + zero-padding: there is no previous rank, and the shard already starts at + position 0 of the sequence. That is the correct answer for a single-device + run and for anything called outside a CP ``shard_map``. + + Restriction: a rank only borrows from the rank immediately before it, so the + window has to fit inside one shard (``halo_size <= T_local``). A convolution + wider than a shard would need tokens from two or more previous ranks; that + case is not implemented and raises ``ValueError`` rather than silently + convolving over the wrong context. + + ``ShortConvolution`` below is the only caller today, which is why this lives + here instead of in a shared utils module. + + Args: + x: The local shard, shaped ``[B, T_local, …]`` with the sequence on + ``seq_axis``. + halo_size: How many preceding tokens the convolution needs, i.e. + ``kernel_size - 1``. + axis_name: Name of the mesh axis the sequence is sharded over. + seq_axis: Index of the sequence dimension in ``x`` (default 1). + + Returns: + ``x`` with ``halo_size`` tokens prepended along ``seq_axis``, shaped + ``[B, halo_size + T_local, …]``. + """ + if halo_size <= 0: + return x + + def _zero_pad(): + """Left zero-pad: the causal-convolution boundary when there is no predecessor. + + Built lazily — under CP with more than one rank the halo replaces it, so + padding here would only add a discarded array to the graph. + """ + pad_width = [(0, 0)] * x.ndim + pad_width[seq_axis] = (halo_size, 0) + return jnp.pad(x, pad_width) + + if not _has_named_axis(axis_name): + return _zero_pad() + + cp_size = jax.lax.psum(1, axis_name=axis_name) + if cp_size == 1: + return _zero_pad() + + 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), i.e. + a grouped convolution 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 — + the same guarantee a varlen causal-conv kernel gives per packed document. + """ + + 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. + + 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 + + This layer runs the reference layer's stages — Q/K/V projections, depthwise + causal short convolution, SiLU, Q/K L2 normalization, forget-gate and beta + projections, the chunked Delta-Rule recurrence, per-head RMSNorm scaled by a + sigmoid output gate, and the output projection — with the differences listed + below. + + References: + Moonshot AI, `Kimi Linear: An Expressive, Efficient Attention Architecture + `_, 2025 + Layer implementation: + https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/kda.py + + The recurrent kernel itself is not implemented here: it is delegated to + tokamax (``api.kimi_delta_attention`` with ``implementation="mosaic"``, a + Pallas/TPU kernel) through ``maxtext.kernels.kda.chunk_kda``. This module + therefore owns only the surrounding projections, convolution, normalization + and sharding, which are MaxText code rather than a port. Where they depart + from the reference layer: + + - Kernel backend: the reference calls fla's Triton ``chunk_kda`` on GPU, + this calls tokamax's Pallas kernel on TPU (generation >= 6 — the adapter + names a single implementation, which disables tokamax's XLA fallback). + Equivalent stage graph, not numerically equivalent. + - Q/K L2 normalization is unconditional here. The reference exposes it as + ``use_qk_l2norm_in_kernel``; this layer normalizes in JAX and passes + ``use_qk_l2norm=False``, because the Delta-Rule recurrence diverges to + NaN in bf16 with unbounded q/k. + - Beta's sigmoid and the output gate's sigmoid are applied here in fp32 + instead of inside the kernel. + - Context parallelism is this layer's concern, not the reference layer's: + the reference is single-device and shards nothing. Here the sequence + stays sharded over the CP axis, so the convolution's left context is + fetched from the previous rank by ``halo_exchange_for_conv``, while + tokamax coordinates the recurrent state across CP ranks (it explicitly + leaves conv-side CP to the caller). + - The sequence is padded up to a multiple of ``_KDA_CHUNK_SIZE`` and sliced + back afterwards, since the kernel needs a static chunk multiple. + - Not implemented: the low-rank bottleneck variant of the two gate + projections that references such as fla use unconditionally. This layer + keeps the full-rank map, matching the published Ling-3.0-flash checkpoint + (``no_kda_lora: true``): + https://huggingface.co/inclusionAI/Ling-3.0-flash + ``use_kda_lora`` selects the low-rank variant and is rejected at config + time rather than ignored, so a run cannot silently train the other + architecture. Also missing: FP8 projections, fused gated-norm kernels, and + autoregressive decode (``__call__`` raises ``NotImplementedError``). + + 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. KDA uses one head count + # and one head dim for q, k and v alike (no GQA-style grouping): + # 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 the log-space gate g, shape [B, T, H, K] (per-head, + # per-dim). Public KDA checkpoints name this `f_proj` and the output gate + # below `g_proj`, so a conversion utility has to swap the two names. + 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 in this layer, before chunk_kda (see + # __call__), so the kernel is passed use_qk_l2norm=False. + + # Output gate projection: full-rank gate of shape [B, T, H, V] (`g_proj` in + # public KDA checkpoints). The low-rank bottleneck variant is not + # implemented — see the class docstring; use_kda_lora=True is rejected at + # config time. + 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, named and initialized as in the KDA reference + # implementation. Params keep the reference 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 the diagonal decay. The two gate paths + # need different initializations, as in the reference layer + # (fla/layers/kda.py): + # softplus gate: log_decay = -exp(A_log) * softplus(g + dt_bias), so + # exp(A_log) scales the decay magnitude -> draw log(U(1, 16)), the + # Mamba-style convention. + # safe gate: log_decay = lower_bound * sigmoid(exp(A_log) * (g + dt_bias)), + # so lower_bound already fixes the magnitude and A only sharpens the + # sigmoid -> start at A = 1 (A_log = 0) and leave the gate in its linear + # regime instead of saturating it at step 0. + # Kept in fp32 either way, like the reference. + if cfg.use_kda_safe_gate: + self.A_log = nnx.Param(jnp.zeros((self.num_key_heads,), dtype=jnp.float32)) + else: + 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] — per-channel gate bias. + # Initialized via inverse softplus of exp(U(log dt_min, log dt_max)) + # clamped at dt_init_floor — again the reference implementation's defaults. + 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 (SiLU follows it, see below) + 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. The reference layer fuses + # this activation into its conv kernel; here it is a separate step. + q = jax.nn.silu(q) + k = jax.nn.silu(k) + v = jax.nn.silu(v) + + # Apply L2 normalization to Q/K outside the kernel (the reference layer + # does it in-kernel via use_qk_l2norm_in_kernel; this layer normalizes in + # JAX and passes use_qk_l2norm=False). 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. + # 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. + # + # This synthesis also carries a second, load-bearing consequence: with + # segment_ids present, tokamax takes its varlen path and aligns each + # segment up to the 64-token chunk internally, so a per-rank slice + # shorter than one chunk is legal. Its non-varlen path instead requires + # T_local % 64 == 0, which the global padding above does NOT guarantee + # once the sequence is sharded (T_global=128, cp_size=4 -> T_local=32). + # test_kda_cp_equivalence[cp_size=4] covers exactly that shape. If this + # synthesis is ever removed, CP needs an explicit per-rank alignment. + 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 keeps the + # reference implementation's name, dt_bias.) + 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 derives the per-rank + # chain fields (cu_seqlens, is_first_rank, …) device-side from + # segment_ids, then passes the completed metadata to the kernel. + cp_ctx = None + if cp_size > 1: + if TokamaxContextParallelMetadata is None: + raise kda_api_unavailable( + detail=( + "KDA context parallelism requires " + "tokamax._src.ops.experimental.kda.cp_utils.ContextParallelMetadata. " + "Refusing to run: without it CP would silently split the " + "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: per-head RMSNorm over the value dim, then a sigmoid + # output gate. The reference layer uses a fused gated RMSNorm for this. + 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 cbf731bfa6..c8baf6d4b2 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 linears, mhc, moe, normalizations, quantizations +from maxtext.layers import attention_kda, 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,8 @@ def __init__( model_mode: str, quant: None | Quant = None, name: str = "decoder_layer", + attention_type: AttentionType | str | None = None, + layer_idx: int = 0, *, rngs: nnx.Rngs, ): @@ -97,9 +100,16 @@ def __init__( self.mesh = mesh self.model_mode = model_mode self.quant = quant + self.layer_idx = layer_idx 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 +119,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=layer_idx, + 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 +215,27 @@ 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, + ) + kv_cache = None + 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) @@ -879,6 +911,11 @@ def _init_sequential_generic(self, decoder_block_classes, rngs): elif config.decoder_block == DecoderBlockType.OLMO3: layer_kwargs = {"attention_type": olmo3.get_attention_type(layer_id=lyr)} + if AttentionType(config.attention_type) == AttentionType.KDA: + # KimiDeltaAttention records which layer of the stack it is; forward the + # real index rather than leaving every KDA layer tagged 0. + layer_kwargs["layer_idx"] = lyr + self._create_and_register_layer(layer_cls, rngs, "layers", lyr, **layer_kwargs) def _init_gemma4_small_layers(self, rngs): diff --git a/tests/unit/kda_attention_test.py b/tests/unit/kda_attention_test.py new file mode 100644 index 0000000000..097581fbac --- /dev/null +++ b/tests/unit/kda_attention_test.py @@ -0,0 +1,1749 @@ +# 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 the kernel, as the layer documents + - 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 +import sys +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 import pyconfig +from maxtext.configs.types import KdaAttention +from maxtext.kernels.kda import chunk_kda +from maxtext.kernels.kda.tokamax import _MIN_TOKAMAX_VERSION, kda_api_unavailable, 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 +from tests.utils.test_helpers import get_test_config_path + +# 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 (one head count and head dim + for q, k and v alike): + 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_a_log_init_follows_gate_path(self, mesh): + """A_log init depends on the gate path, as in the reference layer. + + The softplus gate multiplies the decay by exp(A_log), so A_log is drawn as + log(U(1, 16)); the safe gate already bounds the decay by `lower_bound` and + A only sharpens the sigmoid, so it starts at A = 1. + """ + safe = self._make_attn(mesh, use_kda_safe_gate=True) + assert jnp.all(safe.A_log.value == 0.0) + assert safe.A_log.value.shape == (4,) + + softplus = self._make_attn(mesh, use_kda_safe_gate=False) + values = softplus.A_log.value + assert values.shape == (4,) + assert bool(jnp.all(values >= 0.0)), "log(U(1, 16)) must be non-negative" + assert bool(jnp.all(values <= jnp.log(16.0) + 1e-6)) + assert bool(jnp.any(values > 0.0)), "A_log should be drawn, not all zeros" + + 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 +# --------------------------------------------------------------------------- + + +def _kda_pyconfig(**kwargs): + """Build a config selecting KDA, for exercising the config-time guards. + + `base.yml` defaults `packing=true` and `max_segments_per_seq=-1`, a + combination KDA now rejects at config time, so the helper pins a valid bound + by default; individual tests override it to exercise the guard. + """ + defaults = { + "attention_type": "kda", + "scan_layers": False, + "max_segments_per_seq": 4, + } + defaults.update(kwargs) + return pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + skip_jax_distributed_system=True, + **defaults, + ) + + +class TestKdaConfigGuards: + """Config-time guards for invalid KDA combinations.""" + + def test_kda_cp_rejects_load_balance_at_config_time(self): + """KDA under CP must reject load balancing when the config is parsed. + + `context_parallel_load_balance` defaults to true, so without the config + guard a multi-node KDA run would only fail at its first forward step — + after device allocation and weight initialization. The layer keeps its own + runtime check as defense in depth. + """ + with pytest.raises(ValueError, match="context_parallel_load_balance"): + _kda_pyconfig(ici_context_parallelism=2, context_parallel_load_balance=True) + cfg = _kda_pyconfig(ici_context_parallelism=2, context_parallel_load_balance=False) + assert cfg.ici_context_parallelism == 2 + + def test_kda_packing_requires_max_segments_at_config_time(self): + """A packed KDA run needs a static segment bound before the data pipeline is built. + + The kernel derives per-rank segment metadata from segment_ids and needs + `max_num_segments`; failing here beats failing after the pipeline and the + weights are up. + """ + with pytest.raises(ValueError, match="max_segments_per_seq"): + _kda_pyconfig(packing=True, max_segments_per_seq=-1) + cfg = _kda_pyconfig(packing=True, max_segments_per_seq=8) + assert cfg.max_segments_per_seq == 8 + + 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 selects an unimplemented path 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) + + def test_missing_kda_api_error_names_required_version(self): + """The unavailability error must say which tokamax release is needed.""" + with pytest.raises(ImportError) as excinfo: + raise kda_api_unavailable() + message = str(excinfo.value) + assert _MIN_TOKAMAX_VERSION in message, "error must name the first release shipping the KDA API" + assert "tokamax._src.ops.experimental.kda" in message + assert "pip install" in message, "error must give an actionable upgrade command" + + def test_missing_kda_api_error_keeps_detail_and_cause(self): + """Callers can explain what needed the API, and the cause stays chained.""" + cause = ImportError("no module named tokamax._src.ops.experimental.kda") + with pytest.raises(ImportError) as excinfo: + raise kda_api_unavailable(cause, detail="ContextParallelMetadata is unavailable.") + assert str(excinfo.value).startswith("ContextParallelMetadata is unavailable.") + assert excinfo.value.__cause__ is cause + + +# --------------------------------------------------------------------------- +# 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..56da7fe76f --- /dev/null +++ b/tests/unit/kda_decoder_integration_test.py @@ -0,0 +1,191 @@ +# 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, + # base.yml defaults packing=true with max_segments_per_seq=-1; KDA rejects + # that combination at config time because the kernel needs a static upper + # bound on packed segments. + "max_segments_per_seq": 4, + "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)" + ) From 226b517eb5e98a986036634b659451648b12813a Mon Sep 17 00:00:00 2001 From: chiaotung97 Date: Wed, 23 Sep 2026 11:12:44 +0800 Subject: [PATCH 2/4] fix(kda): size the layer off the scaled config dims and the CP-aware axis Behavior-preserving at the default `global_parameter_scale=1`; all three are correctness fixes for the non-default configurations. - attention_kda.py: build the projections from `config.emb_dim` and `config.num_query_heads` instead of `base_emb_dim` / `base_num_query_heads`. The derived fields are what fold in `global_parameter_scale` (see `MaxTextConfig.set_derived_and_validate_values`), so sizing off the base ones built every KDA projection for the unscaled model and would mismatch the hidden states the decoder hands it at any other scale. The gate, beta and output-gate projections carried the same bug as the q/k/v/o ones and are fixed with them. - types.py: the KDA + `context_parallel_load_balance` guard now reuses the `context_parallel_size` already derived earlier in the same validator, which resolves the CP axis from `context_sharding`. It hardcoded `ici_context_parallelism * dcn_context_parallelism`, so an expert-as-context mesh (`context_sharding: "expert"`) skipped the guard entirely and would have composed the recurrent state out of token order. - attention_kda.py: `halo_exchange_for_conv` reads its CP size with `jax.lax.axis_size` instead of `jax.lax.psum(1, ...)`. - tests: the mock config and the smoke-test namespace now expose `emb_dim` / `num_query_heads`, and the smoke-test scaffold layers read the same width the KDA layer does, so the two cannot drift apart. --- src/maxtext/configs/types.py | 6 ++++-- src/maxtext/layers/attention_kda.py | 27 ++++++++++++++---------- tests/unit/kda_attention_test.py | 32 ++++++++++++++++------------- 3 files changed, 38 insertions(+), 27 deletions(-) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 835694556c..8afb625a76 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -4940,8 +4940,10 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "inside a scanned layer stack. Set scan_layers: false." ) - kda_context_parallel_size = self.ici_context_parallelism * self.dcn_context_parallelism - if self.attention_type == "kda" and kda_context_parallel_size > 1 and self.context_parallel_load_balance: + # Reuses the `context_parallel_size` derived above, which already resolves the + # CP axis from `context_sharding` ("context" by default, "expert" for + # expert-as-context) instead of assuming the named axis. + if self.attention_type == "kda" and context_parallel_size > 1 and self.context_parallel_load_balance: raise ValueError( "attention_type='kda' with context parallelism requires context_parallel_load_balance=false. " "The KDA recurrence composes state in token order, so device i must hold the sequence chunk " diff --git a/src/maxtext/layers/attention_kda.py b/src/maxtext/layers/attention_kda.py index 2ef815b1c3..af92cb40c1 100644 --- a/src/maxtext/layers/attention_kda.py +++ b/src/maxtext/layers/attention_kda.py @@ -142,7 +142,7 @@ def _zero_pad(): if not _has_named_axis(axis_name): return _zero_pad() - cp_size = jax.lax.psum(1, axis_name=axis_name) + cp_size = jax.lax.axis_size(axis_name=axis_name) if cp_size == 1: return _zero_pad() @@ -310,11 +310,16 @@ def __init__( # KDA head dimensions derived from global config. KDA uses one head count # and one head dim for q, k and v alike (no GQA-style grouping): # 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) + # num_key_heads = num_value_heads = config.num_query_heads (num_attention_heads) + # The derived `num_query_heads` / `emb_dim` below rather than their `base_*` + # counterparts: `base_*` are the unscaled model dims, and the derived ones + # fold in `global_parameter_scale`, so sizing the projections off `base_*` + # would build a layer mismatched to the decoder's hidden states whenever the + # scale is anything but 1. 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_key_heads = cfg.num_query_heads + self.num_value_heads = cfg.num_query_heads self.num_query_heads = self.num_key_heads # Short convolution for local dependency modeling @@ -348,7 +353,7 @@ def __init__( # 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, + in_features_shape=cfg.emb_dim, out_features_shape=(self.num_query_heads, self.key_head_dim), axis=-1, dtype=cfg.dtype, @@ -361,7 +366,7 @@ def __init__( ) self.k_proj = linears.DenseGeneral( - in_features_shape=cfg.base_emb_dim, + in_features_shape=cfg.emb_dim, out_features_shape=(self.num_key_heads, self.key_head_dim), axis=-1, dtype=cfg.dtype, @@ -374,7 +379,7 @@ def __init__( ) self.v_proj = linears.DenseGeneral( - in_features_shape=cfg.base_emb_dim, + in_features_shape=cfg.emb_dim, out_features_shape=(self.num_value_heads, self.value_head_dim), axis=-1, dtype=cfg.dtype, @@ -389,7 +394,7 @@ def __init__( # 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, + out_features_shape=cfg.emb_dim, axis=(-2, -1), dtype=cfg.dtype, weight_dtype=cfg.weight_dtype, @@ -404,7 +409,7 @@ def __init__( # per-dim). Public KDA checkpoints name this `f_proj` and the output gate # below `g_proj`, so a conversion utility has to swap the two names. self.g_proj = linears.DenseGeneral( - in_features_shape=cfg.base_emb_dim, + in_features_shape=cfg.emb_dim, out_features_shape=(self.num_key_heads, self.key_head_dim), axis=-1, dtype=cfg.dtype, @@ -419,7 +424,7 @@ def __init__( # 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, + in_features_shape=cfg.emb_dim, out_features_shape=(self.num_key_heads,), axis=-1, dtype=cfg.dtype, @@ -439,7 +444,7 @@ def __init__( # implemented — see the class docstring; use_kda_lora=True is rejected at # config time. self.gate_proj = linears.DenseGeneral( - in_features_shape=cfg.base_emb_dim, + in_features_shape=cfg.emb_dim, out_features_shape=(self.num_value_heads, self.value_head_dim), axis=-1, dtype=cfg.dtype, diff --git a/tests/unit/kda_attention_test.py b/tests/unit/kda_attention_test.py index 097581fbac..55d42c1d2f 100644 --- a/tests/unit/kda_attention_test.py +++ b/tests/unit/kda_attention_test.py @@ -147,12 +147,16 @@ class _MockKdaConfig: KDA derives head dims from global config (one head count and head dim for q, k and v alike): key_head_dim = value_head_dim = head_dim - num_key_heads = num_value_heads = base_num_query_heads + num_key_heads = num_value_heads = num_query_heads + + `emb_dim` / `num_query_heads` are what MaxText's config exposes after + folding in `global_parameter_scale` — the values `KimiDeltaAttention` sizes + its projections from. This mock stands at scale 1. """ def __init__(self, **overrides): - self.base_emb_dim = 128 - self.base_num_query_heads = 4 + self.emb_dim = 128 + self.num_query_heads = 4 self.head_dim = 32 self.dtype = jnp.float32 self.weight_dtype = jnp.float32 @@ -197,7 +201,7 @@ def _make_attn(self, mesh, **config_overrides): ) def test_init_head_dims(self, mesh): - """Head dims derived from global config: head_dim=32, base_num_query_heads=4.""" + """Head dims derived from global config: head_dim=32, num_query_heads=4.""" attn = self._make_attn(mesh) assert attn.num_query_heads == 4 assert attn.num_key_heads == 4 @@ -1613,7 +1617,7 @@ class _KdaBlock(nnx.Module): def __init__(self, cfg, mesh, layer_idx, *, rngs): self.attn_norm = RMSNorm( - num_features=cfg.base_emb_dim, + num_features=cfg.emb_dim, epsilon=cfg.normalization_layer_epsilon, dtype=cfg.dtype, weight_dtype=cfg.weight_dtype, @@ -1621,15 +1625,15 @@ def __init__(self, cfg, mesh, layer_idx, *, 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, + num_features=cfg.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) + hidden = 4 * cfg.emb_dim + self.wi = nnx.Linear(cfg.emb_dim, hidden, dtype=cfg.dtype, param_dtype=cfg.weight_dtype, rngs=rngs) + self.wo = nnx.Linear(hidden, cfg.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)) @@ -1643,17 +1647,17 @@ 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.embed = nnx.Embed(_KDA_SMOKE_VOCAB, cfg.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, + num_features=cfg.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 + cfg.emb_dim, _KDA_SMOKE_VOCAB, dtype=cfg.dtype, param_dtype=cfg.weight_dtype, rngs=rngs ) def __call__(self, tokens): @@ -1689,8 +1693,8 @@ class TestKdaE2eSmoke: 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, + emb_dim=256, + num_query_heads=8, head_dim=64, dtype=jnp.float32, weight_dtype=jnp.float32, From 4950aaacf6330caeb7b80683bce05ceee1966643 Mon Sep 17 00:00:00 2001 From: chiaotung97 Date: Wed, 23 Sep 2026 17:36:13 +0800 Subject: [PATCH 3/4] style(kda): apply pyink formatting to the smoke test's lm_head Renaming the projection width to cfg.emb_dim shortened this call enough to fit the line budget, so pyink wants it on one line. --- tests/unit/kda_attention_test.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit/kda_attention_test.py b/tests/unit/kda_attention_test.py index 55d42c1d2f..4227c43cef 100644 --- a/tests/unit/kda_attention_test.py +++ b/tests/unit/kda_attention_test.py @@ -1656,9 +1656,7 @@ def __init__(self, cfg, mesh, num_layers, *, rngs): weight_dtype=cfg.weight_dtype, rngs=rngs, ) - self.lm_head = nnx.Linear( - cfg.emb_dim, _KDA_SMOKE_VOCAB, dtype=cfg.dtype, param_dtype=cfg.weight_dtype, rngs=rngs - ) + self.lm_head = nnx.Linear(cfg.emb_dim, _KDA_SMOKE_VOCAB, dtype=cfg.dtype, param_dtype=cfg.weight_dtype, rngs=rngs) def __call__(self, tokens): x = self.embed(tokens) From ffad407aff628d7e4e1eb578a755a68a0d822fa0 Mon Sep 17 00:00:00 2001 From: chiaotung97 Date: Wed, 23 Sep 2026 18:00:50 +0800 Subject: [PATCH 4/4] docs(kda): correct the CP collective, check_vma and pspec rationale Three review corrections in the design doc, plus the comments in the layer that pointed at the same reasoning. Prose, docstrings and comments only, no behavior change. - Remove the claim that KDA does not all-gather. It gathers a fixed-size state summary rather than K/V, so the section now lists each gather with its shape and the tokamax source line, and states the cost as O(cp_size * B * H * K * (V + K)) per step, independent of sequence length. - Replace the claim that FlashAttention custom rules falsely report VMA errors with the measurement on 4xTPU v6e: 9 passed and 3 failed with check_vma=True on the kernel shard_map, identical when the conv side is enabled too. Both root causes are in tokamax, the fori_loop in _derive_cp_metadata_from_segment_ids and the Pallas launcher's ShapeDtypeStructs missing manual_axis_type. Also records that every other attention shard_map in MaxText hardcodes False and that config.check_vma is read only by moe.py. - Replace the nnx.logical_to_mesh_axes rationale, which described a problem that does not exist here. MaxText's own logical_to_mesh_axes already puts the sequence on the "context" axis, so the injection is a no-op there and only matters for expert-as-context, where activation_norm_length has no expert rule. The three measured pspecs are tabulated. - Rewrite the _inject_cp_axis_on_T docstring to match, and comment both check_vma=False call sites with a pointer to the doc. --- docs/reference/kda_cp_support.md | 39 +++++++++++++++++++++++------ src/maxtext/layers/attention_kda.py | 21 ++++++++++++---- 2 files changed, 48 insertions(+), 12 deletions(-) diff --git a/docs/reference/kda_cp_support.md b/docs/reference/kda_cp_support.md index 14c8c12794..067e7da09a 100644 --- a/docs/reference/kda_cp_support.md +++ b/docs/reference/kda_cp_support.md @@ -21,7 +21,15 @@ CP (cp_size > 1): → [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. +The difference from MLA CP is what the collective carries, not whether one happens at all. MLA all-gathers K/V for the sharded sequence, so its payload grows with sequence length. KDA all-gathers a fixed-size summary of the recurrent state instead. The shapes below are each rank's local shape, and `jax.lax.all_gather` adds a leading `cp_size` axis to the result: + +| Path | What is gathered | Shapes | Source | +| -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------- | ------------------------------------------- | +| Forward | `S_ext`, the affine contribution of each rank's last segment, and `M`, its transition matrix. `_merge_initial_state` composes them into each rank's incoming state. | `[H, B, K, V]` and `[H, B, K, K]` | `pallas_mosaic_tpu_fwd_kernel.py:411-414` | +| Backward | `dS_ext` and `dM`, packed along the last axis into one tensor so a single gather covers both | `[B, H, K, V + K]` | `pallas_mosaic_tpu_bwd_kernel.py:1752-1754` | +| Chain metadata | each rank's first and last segment id, used to derive `cu_seqlens`, `pre_num_ranks` and the related per-rank fields | `2 x cp_size` int32 | `cp_utils.py:260-261` | + +None of the three carries T, since `chunk_gated_delta_rule_fwd_h_pre_process` returns `[H, B, K, V]` and `[H, B, K, K]` whatever `T_local` is. The CP collective cost is therefore `O(cp_size * B * H * K * (V + K))` per step rather than something proportional to sequence length. That is what `ContextParallelMetadata` coordinates: the merge of recurrent state across ranks in forward and backward, with the sequence itself staying sharded. ### Plan 1: `halo_exchange_for_conv` (in `layers/attention_kda.py`, KDA-specific) @@ -52,9 +60,20 @@ Change location: the conv call segment after QKV projection in `KimiDeltaAttenti 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 +- **both pass `check_vma=False`**, see the `check_vma` section below - Zero-overhead fallback when no CP: follows the original path exactly +#### Why `check_vma=False` + +`check_vma=True` cannot be used while KDA runs through tokamax. Measured on 4xTPU v6e over the CP selection, `pytest tests/unit/kda_attention_test.py -k "Cp or cp or short_conv or halo"`, which collects 36 tests and skips 24 of them as `cpu_only`. With `check_vma=True` on the kernel-side shard_map, 9 passed and 3 failed. Enabling it on the conv side as well changes nothing: the same 9 pass and the same 3 fail. So the conv-side shard_map tolerates it and the failure is entirely kernel side, for two reasons that are both outside this repo: + +1. `tokamax/_src/ops/experimental/kda/cp_utils.py:297`. The `fori_loop` inside `_derive_cp_metadata_from_segment_ids` enters with a replicated carry (`bool[]`, `int32[]`) and returns one that varies on the CP axis (`bool[]{V:context}`, `int32[]{V:context}`). JAX's VMA scan check rejects the type mismatch, and its own error message suggests `jax.lax.pcast(..., ('context',), to='varying')` on the initial carry. +2. `jax/_src/pallas/core.py:1888`. With `check_vma=True` on a shard_map, every `jax.ShapeDtypeStruct` must set `manual_axis_type`, which tokamax's KDA Pallas launcher does not. + +The three failures are `test_kda_cp_full_layer_dummy_segments`, `test_kda_cp_full_layer_packed_segments` and `test_kda_no_cp_without_load_balance_ok`. The last one runs at `cp_size=1`, which shows that reason 2 applies to any KDA forward pass and not only under CP. + +Until both are fixed upstream, `False` is what every other attention shard_map in MaxText uses. `layers/attention_op.py` lines 1778, 1843 and 2298 and `kernels/tokamax_splash_attention/splash_attention_kernel.py:2155` all hardcode it, and `config.check_vma` (default `False`, `configs/base.yml:719`, documented as covering "EP / FSDP ICI parallelisms") is consumed only by `layers/moe.py:2625`. It is an MoE knob today rather than an attention one, so adopting it across the attention path is a separate cleanup that needs the two tokamax fixes first. + ### Plan 3: chunk_kda ContextParallelMetadata + Partition Spec (`attention_kda.py`) #### 3a. ContextParallelMetadata Construction (outside shard_map) @@ -78,18 +97,24 @@ if cp_size > 1: #### 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: +MaxText's own `logical_to_mesh_axes` (imported from `maxtext.utils.sharding` at `attention_kda.py:57`) already resolves the T axis correctly when the CP axis is named `"context"`. Measured on 4xTPU v6e using the config's own `logical_axis_rules`: + +| Config | Resolved pspec for `(activation_batch, activation_norm_length, None)` | Effect of `_inject_cp_axis_on_T` | +| ------------------------------------------------------- | --------------------------------------------------------------------- | -------------------------------- | +| `ici_context_parallelism=2` | `P('fsdp', 'context', None)` | none, already correct | +| `ici_context_parallelism=4` | `P(None, 'context', None)` | none, already correct | +| `context_sharding='expert'`, `ici_expert_parallelism=2` | `P(('fsdp', 'expert'), None, None)` | overwrites T with `'expert'` | + +The third row is why the injection exists. `activation_norm_length` maps to `["tensor_sequence", "context", "context_usp_ulysses"]` (`configs/types.py:1383`), and that list has no `expert` entry, so expert-as-context resolves T to `None`. Without the overwrite the shard_map would run replicated over the very axis its collectives use. ```python def _inject_cp_axis_on_T(pspec, t_axis=1): spec = list(pspec) - spec[t_axis] = cp_axis_name # overwritten unconditionally + spec[t_axis] = cp_axis_name 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. +Applied to `qkv_pspec`, `beta_pspec` and `seg_pspec` when CP is enabled, followed by `with_sharding_constraint` to ensure tensor physical layout matches. `cp_axis_name` is `cfg.context_sharding`, which defaults to `"context"` and is `"expert"` for expert-as-context. Overwriting on every strategy rather than only when T resolves to `None` keeps both cases on one code path, and on `"context"` it is a no-op by construction, as the table above shows. #### 3c. chunk_kda shard_map diff --git a/src/maxtext/layers/attention_kda.py b/src/maxtext/layers/attention_kda.py index af92cb40c1..373492a3cf 100644 --- a/src/maxtext/layers/attention_kda.py +++ b/src/maxtext/layers/attention_kda.py @@ -591,11 +591,13 @@ def __call__( 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. + A no-op on the default "context" axis, since activation_norm_length + already lists it (configs/types.py:1383) and logical_to_mesh_axes + resolves T to "context" on its own. Load-bearing for expert-as-context, + where that rule list has no "expert" entry, T resolves to None, and the + shard_map would otherwise run replicated over the axis its collectives + use. Overwriting on both strategies keeps one code path. The measured + pspecs are tabulated in docs/reference/kda_cp_support.md. """ spec = list(pspec) spec[t_axis] = cp_axis_name @@ -662,6 +664,8 @@ def _inject_cp_axis_on_T(pspec, t_axis=1): conv_seg_pspec, ), out_specs=(conv_flat_pspec, conv_flat_pspec, conv_flat_pspec), + # Same as the kernel-side shard_map below, see the check_vma + # section in docs/reference/kda_cp_support.md. check_vma=False, ) def _conv_with_halo(qf, kf, vf, seg): @@ -803,6 +807,13 @@ def _wsc(x, pspec): ) cp_ctx = TokamaxContextParallelMetadata(mesh=self.mesh, axis_name=cp_axis_name) + # check_vma=False is required while KDA calls into tokamax. Enabling it + # fails on the fori_loop in cp_utils._derive_cp_metadata_from_segment_ids, + # whose carry turns from replicated into varying on the CP axis, and on + # the Pallas launcher's ShapeDtypeStructs, which carry no manual_axis_type. + # Both are outside this repo, and every other attention shard_map in + # MaxText hardcodes False as well. The measured failures and what they + # cover are written up in docs/reference/kda_cp_support.md. @functools.partial( jax.shard_map, mesh=self.mesh,