feat(kda): integrate KDA attention with tokamax backend and CP support

[chiaotung97]

Summary

This PR integrates KDA (Kimi Delta Attention) into MaxText with tokamax backend and context parallelism (CP) support.

• KimiDeltaAttention layer with QKV/beta/gate projections and depthwise causal 1D convolution
• tokamax backend adapter with [B,T,H,D] ↔ [H,B,T,D] layout translation
• CP support via CPContext, shard_map wrapping, and ppermute-based halo exchange
• configuration model (KdaAttention) with kernel dim, safe-gate, lower-bound, and LoRA controls
• design documentation and 914-line unit test suite

Dependency: Requires openxla/tokamax#1103 to be merged into tokamax main branch first — until then, tokamax must be installed from the PR branch.

Motivation

KDA is a recurrent linear-attention operator that augments the standard Delta Rule with a learned gating mechanism and depthwise convolution for local dependency modeling. Integrating KDA into MaxText enables:

• alternative attention backend for models that benefit from linear-complexity sequence processing
• context-parallel training with proper recurrent state propagation across ranks
• packed variable-length sequence support via segment_ids

The integration follows the Megatron KDA reference and delegates kernel execution to tokamax's Pallas TPU implementation, keeping MaxText free of low-level kernel code.

Public API

KimiDeltaAttention (nnx.Module)

class KimiDeltaAttention(nnx.Module):
  def __init__(self, config: Config, layer_idx: int, mesh: Mesh, *, rngs: nnx.Rngs)

  def __call__(
      self,
      hidden_states: jnp.ndarray,       # [B, T, D]
      decoder_segment_ids: jnp.ndarray | None = None,  # [B, T], 1-based, 0=padding
  ) -> tuple[jnp.ndarray, None]:        # (output [B, T, D], None)

chunk_kda (kernel entry point)

def chunk_kda(
    q, k, v, g, beta,
    scale=None, initial_state=None, output_final_state=False,
    chunk_size=64, A_log=None, dt_bias=None,
    use_gate_in_kernel=False, use_qk_l2norm_in_kernel=False,
    safe_gate=False, lower_bound=None, segment_ids=None,
    disable_recompute=False, N_max=None, cp_context=None,
) -> tuple[jnp.ndarray, jnp.ndarray | None]

halo_exchange_for_conv (CP utility)

def halo_exchange_for_conv(x: jnp.ndarray, halo_size: int, axis_name: str = "context") -> jnp.ndarray

Configuration (KdaAttention)

Field	Type	Default	Description
linear_conv_kernel_dim	int (≥0)	4	Causal 1D conv kernel size; 0 disables
use_kda_lora	bool	False	Low-rank factorization for KDA computation
use_kda_safe_gate	bool	False	Numerically safe gate clamping
kda_lower_bound	float	0.0	Gate lower bound (only active when safe_gate=True)

Implementation Details

Layer architecture (layers/attention_kda.py)

The KimiDeltaAttention forward pass follows a six-stage pipeline:

1. Short convolution — Q, K, V each pass through a causal depthwise 1D conv (ShortConvolution) with SiLU activation, modeling local (kernel_dim) token dependencies
2. Gate generation — a separate projection produces raw gate values g; gate activation (sigmoid from A_log, dt_bias) is handled inside tokamax
3. Output gate — a sigmoid-gated projection produces per-head output scaling applied after the KDA kernel
4. Beta projection — sigmoid-activated scalar per head controls the Delta Rule mixing coefficient
5. Q/K L2 normalization — optional normalization applied outside the kernel (matching Megatron convention)
6. KDA kernel — delegates to chunk_kda → tokamax kimi_delta_attention(implementation="pallas_tpu")

Delta Rule formulation:

S' = S * exp(g_t)
residual = v_t - k_t^T @ S'
S = S' + beta_t * k_t ⊗ residual
o_t = scale * q_t^T @ S

tokamax backend adapter (kernels/kda/tokamax.py)

• Translates batch-first maxtext layout [B, T, H, D] to head-first tokamax layout [H, B, T, D]
• Passes segment_ids directly as [B, T] (tokamax accepts this natively — no per-batch offset needed)
• Forwards CPContext for context-parallel execution

Context parallelism (utils/cp_utils.py)

Two independent jax.shard_map invocations provide per-op sharding:

• ShortConvolution halo exchange — halo_exchange_for_conv uses ppermute-based forward ring to fetch kernel_dim-1 tokens from the left neighbor, enabling correct causal convolution at rank boundaries
• KDA kernel shard_map — wraps chunk_kda with explicit partition specs so tokamax kernels (not auto-partitionable) receive correctly sharded tensors

CP metadata flow:

1. _inject_context_on_T ensures the T axis carries "context" sharding (overrides size-1 stripping)
2. with_sharding_constraint applies explicit pspecs to Q/K/V/G and beta before entering shard_map
3. CPContext(mesh, axis_name="context") is constructed outside shard_map and passed to tokamax, which derives per-rank chain fields internally

Runtime guards

• CP + load_balance raises ValueError — recurrent state depends on exact token order; DUAL_CHUNK_SWAP reorder breaks sequential dependency
• kda_lower_bound with safe_gate=False emits a warning (value ignored without clamping)

Test Coverage

tests/unit/kda_attention_test.py (914 lines, 30 tests) covers:

Category	Tests	Description
Initialization	3	Head dims derivation, no-conv mode, gate+norm layers
Forward	7	Shape, no NaN/Inf, padding, deterministic, packed seqs, segment_ids, row independence
Kernel precision	5	chunk_kda vs naive recurrent reference (FP32 + BF16)
QK L2 norm	3	Applied outside kernel, skipped when disabled, output changes
Backward	4	No NaN, deterministic, weight grads non-zero, BF16
ShortConvolution	2	No-CP and CP halo exchange
CP integration	3	cp_size=1 vs cp_size=2 equivalence, load_balance rejection, no-CP without load_balance ok
Guards	1	Autoregressive mode not supported
Properties	2	Naive KDA: zero-gate accumulates, large-negative-gate decays

Naive KDA Reference

A pure-XLA recurrent reference implementation (_naive_kda_recurrent) is embedded in the test suite. It implements the Delta Rule one token at a time with no chunking, providing an independent correctness baseline for the chunked tokamax kernel. Tests compare output against this reference at both FP32 and BF16.

TPU Results

Environment:

JAX      0.11.0
libtpu   0.0.44.1
TPU      4 × v6e (2×2×1)
tokamax antgroup/kda-pallas-kernel

30 passed in 113.18s

Documentation

docs/design/kda_cp_support.md covers:

• CP architecture overview (two independent shard_map invocations, pspec injection)
• halo exchange design for causal convolution
• CPContext lifecycle and tokamax integration
• runtime guard rationale (load_balance incompatibility)
• current limitations and future extensions (CP load balancing, GQA support)

Scope and Compatibility

• All new files are additive — no existing attention implementations are modified
• KDA is a standalone attention module, not wired into decoder pipelines yet (subsequent PR)
• Requires tokamax at runtime; ImportError is raised with a clear message if tokamax is missing
• initial_state and output_final_state are not yet supported, which only for inference (raise NotImplementedError)

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

Co-authors

Co-authored-by: chiaotung97 qt533360@antgroup.com

[chiaotung97]

Unit Test Results

30/30 passed on a 4-chip TPU v6e VM (113s).

Test VM Setup

# 1. Install Python 3.12
sudo apt-get update -qq && sudo apt-get install -y -qq python3.12 python3.12-venv

# 2. Install uv
pip install uv

# 3. Clone maxtext KDA branch
git clone --depth=1 --branch=feature_kda_integration \
  https://github.com/antgroup/maxtext.git maxtext

# 4. Create venv + install maxtext TPU deps
cd maxtext
uv venv --python 3.12 --seed ../maxtext_venv
source ../maxtext_venv/bin/activate
uv pip install -e ".[tpu]"

# 5. Install tokamax from PR branch (required until openxla/tokamax#1103 is merged)
uv pip install git+https://github.com/antgroup/tokamax.git@antgroup/kda-pallas-kernel

Run Tests

source ../maxtext_venv/bin/activate
cd ~/maxtext
python -m pytest tests/unit/kda_attention_test.py -v

Results

tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_init_head_dims PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_init_no_conv PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_init_has_gate_and_norm PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_forward_shape PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_forward_no_nan_inf PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_sequence_padding PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_deterministic PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_packed_sequences_supported PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_segment_ids_padding_alignment PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_segment_ids_none_fallback PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_row_independence PASSED
tests/unit/kda_attention_test.py::TestKimiDeltaAttention::test_autoregressive_not_supported PASSED
tests/unit/kda_attention_test.py::TestChunkKda::test_basic PASSED
tests/unit/kda_attention_test.py::TestNaiveKda::test_chunk_kda_vs_naive PASSED
tests/unit/kda_attention_test.py::TestNaiveKda::test_naive_kda_basic_properties PASSED
tests/unit/kda_attention_test.py::TestNaiveKda::test_naive_kda_zero_gate_accumulates PASSED
tests/unit/kda_attention_test.py::TestNaiveKda::test_naive_kda_large_negative_gate_decays PASSED
tests/unit/kda_attention_test.py::TestNaiveKda::test_chunk_kda_vs_naive_bf16 PASSED
tests/unit/kda_attention_test.py::TestQkL2Norm::test_qk_l2norm_applied_outside_kernel PASSED
tests/unit/kda_attention_test.py::TestQkL2Norm::test_qk_l2norm_skipped_when_disabled PASSED
tests/unit/kda_attention_test.py::TestQkL2Norm::test_l2norm_changes_output PASSED
tests/unit/kda_attention_test.py::TestKdaBackward::test_backward_no_nan PASSED
tests/unit/kda_attention_test.py::TestKdaBackward::test_backward_deterministic PASSED
tests/unit/kda_attention_test.py::TestKdaBackward::test_weight_grads_no_nan PASSED
tests/unit/kda_attention_test.py::TestKdaBackward::test_backward_bf16 PASSED
tests/unit/kda_attention_test.py::TestShortConvolution::test_short_conv_no_cp PASSED
tests/unit/kda_attention_test.py::TestShortConvolution::test_short_conv_cp_halo PASSED
tests/unit/kda_attention_test.py::TestKdaCp::test_kda_cp_equivalence PASSED
tests/unit/kda_attention_test.py::TestKdaCp::test_kda_cp_rejects_load_balance PASSED
tests/unit/kda_attention_test.py::TestKdaCp::test_kda_no_cp_without_load_balance_ok PASSED

======================== 30 passed in 113.18s ========================

Environment

Component	Version
Python	3.12.13
JAX	0.11.0
libtpu	0.0.44.1
tokamax	antgroup/kda-pallas-kernel (openxla/tokamax#1103)
TPU	4 × v6e (2×2×1)

Key Coverage

• Forward: shape, no NaN/Inf, sequence padding, deterministic, packed sequences, segment_ids alignment, row independence
• Backward: no NaN, deterministic, weight grads non-zero, BF16
• Kernel precision: chunk_kda vs naive recurrent reference (FP32 + BF16)
• QK L2 norm: applied outside kernel, skipped when disabled, changes output
• ShortConvolution: no-CP and CP halo exchange
• CP equivalence: cp_size=1 vs cp_size=2 forward output matches
• CP guarding: load_balance rejected with CP, no-CP allowed without load_balance