Skip to content

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

Open
chiaotung97 wants to merge 14 commits into
mainfrom
feature_kda_integration
Open

feat(kda): integrate KDA attention with tokamax backend and CP support#1
chiaotung97 wants to merge 14 commits into
mainfrom
feature_kda_integration

Conversation

@chiaotung97

@chiaotung97 chiaotung97 commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Description

This PR integrates KDA (Kimi Delta Attention, a linear attention mechanism) into MaxText with a tokamax-backed Pallas TPU kernel and context parallelism (CP) support. KDA updates its recurrent state with the Delta Rule:

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

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.

Key Changes

File Description
src/maxtext/layers/attention_kda.py New KimiDeltaAttention layer and ShortConvolution: QKV/beta/gate/output-gate projections, depthwise causal 1D convolution, SiLU activation, optional QK L2 normalization, per-head RMSNorm + output gate, and gate parameters A_log / dt_bias (matching the Megatron reference)
src/maxtext/kernels/kda/ New chunk_kda kernel entry point and tokamax adapter: [B,T,H,D][H,B,T,D] layout translation; the non-auto-partitionable tokamax kernel is invoked inside shard_map with explicit partition specs
src/maxtext/configs/types.py New KdaAttention config (linear_conv_kernel_dim, use_kda_safe_gate, kda_lower_bound, reserved use_kda_lora) plus validators: safe gate requires kda_lower_bound ∈ [-5, 0); use_kda_lora=True is rejected as unimplemented
src/maxtext/utils/cp_utils.py New halo_exchange_for_conv: under CP, pulls kernel_size-1 tokens of left context from the previous CP rank via ppermute so causal convolution stays correct at CP shard boundaries (without it the exchange would silently degrade to zero-pad)
tests/unit/kda_attention_test.py 40 unit tests (marked tpu_only)
docs/design/kda_cp_support.md Design doc
scripts/dev/kda_e2e_smoke.py Standalone end-to-end smoke training script (dev use)

Context Parallelism Support

  • The CP mesh axis is taken from cfg.context_sharding (default "context"; "expert" works for expert-as-context) and is threaded consistently through the conv halo exchange, the T-axis partition-spec injection, and ContextParallelMetadata
  • ContextParallelMetadata passes mesh information to the chunk_kda kernel; tokamax derives per-rank cu_seqlens / chain fields internally from segment_ids and coordinates recurrent state across CP ranks
  • The KDA recurrent state depends on exact token order, so context_parallel_load_balance is rejected up front (DUAL_CHUNK_SWAP reordering would break the sequential dependency)
  • When the user supplies no segment_ids, the CP path synthesizes all-ones segment ids so the kernel can still derive its metadata
  • ShortConvolution halo exchange is wrapped in a shard_map exposing the CP axis, since ppermute is a collective

Tests and Validation

Unit tests30 passed + 10 added in review round = 40/40 passing:

40 passed in 200.21s
  • Run on 4×TPU v6e, 2026-08-31, at head 82a9fc1e (test log: pytest tests/unit/kda_attention_test.py -v)
  • Coverage:
    • Precision vs. a pure-XLA recurrent reference (token-by-token Delta Rule, no chunking), FP32 / BF16, with ULP-based fallback checks
    • Forward / backward (activation and weight gradients), determinism
    • QK L2 normalization — including a direct unit-norm + direction-preservation assertion on _l2_normalize
    • Within-row packed-segment isolation (both directions), complementing cross-row independence
    • ShortConvolution CP halo-exchange equivalence, parametrized over segment layouts: uniform, a boundary exactly on the rank split, and a segment spanning both ranks
    • Kernel-level CP forward equivalence, parametrized CP=2 and CP=4
    • CP backward: dq/dk/dv/dg/dbeta equal the non-CP reference
    • Full-layer CP without user segment_ids: exercises the internal dummy-segment synthesis path (forward equivalence + backward finiteness)
    • Config guards: safe-gate/lower-bound range, use_kda_lora rejection, packing without max_segments_per_seq
    • load_balance rejection

End-to-end smoke (4×TPU v6e): a 5.4M-param tiny model built from real KimiDeltaAttention layers, trained for 400 steps on a fully learnable synthetic next-token task:

step 0:   loss 5.38  (random baseline ln(128) ≈ 4.85)
step 399: loss 0.0007  → PASS

This validates the full forward / backward / optimizer chain through the real Pallas kernels.

Reproducing these results

On a TPU host (Python ≥ 3.12; verified on 4×TPU v6e):

# 1. TPU-capable JAX (the repo's tpu requirements pin jax>=0.11.1;
#    the validation run here used JAX 0.11.0 + libtpu 0.0.44.1)
pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# 2. MaxText with TPU dependencies
pip install -e <maxtext checkout>
pip install -r <maxtext checkout>/src/dependencies/requirements/generated_requirements/tpu-requirements.txt

# 3. tokamax with the KDA Pallas kernels (until openxla/tokamax#1103 merges;
#    branch tip validated here: 939da5c)
git clone -b antgroup/kda-pallas-kernel https://github.com/antgroup/tokamax.git
pip install -e tokamax

# 4. Unit tests — marked tpu_only since they exercise the Pallas TPU
#    kernels; they are skipped on CPU/GPU-only hosts
pytest tests/unit/kda_attention_test.py -v

# 5. (optional) end-to-end smoke training
python scripts/dev/kda_e2e_smoke.py

Dependencies

  • Depends on the KDA kernels that have landed on openxla/tokamax main (originally Implement KDA Pallas Kernel openxla/tokamax#1103, left open; the module tokamax._src.ops.experimental.kda is on main and this PR's adapter targets its API). No public tokamax release contains the KDA API yet, so reproduction installs tokamax from source (main), see steps below. MaxText's declared pin cannot express an unreleased API; the adapter imports KDA lazily, so installs without it are unaffected until KDA is actually used. Once the first release containing KDA ships, the tokamax>= pin will be bumped and the derived requirement files regenerated (TODO tracked in src/maxtext/kernels/kda/tokamax.py).
  • Validated with JAX 0.11.0 + libtpu 0.0.44.1

Hardware / Shape Constraints (mosaic kernel)

The adapter explicitly selects the "mosaic" Pallas implementation (no silent fallback to the slow XLA reference during training). Its constraints — all surfaced as clear NotImplementedErrors from tokamax at kernel bind time:

  • TPU generation ≥ 6 (validated on v6e)
  • Key dimension ≤ 256
  • Under CP: key and value head dims must be multiples of 128
  • Sequence length padded internally to a multiple of chunk size 64

Known Limitations / Follow-ups

  • initial_state / output_final_state not yet supported (explicit NotImplementedError)
  • Autoregressive mode not yet implemented
  • The KDA layer is not yet wired into the decoder (this PR delivers the layer and kernels; model integration is a follow-up)

Checklist

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@antgroup antgroup deleted a comment from qiaotonggg Jul 28, 2026
@chiaotung97

chiaotung97 commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator Author

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

@github-actions

Copy link
Copy Markdown

This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions.

@github-actions github-actions Bot added the stale label Aug 28, 2026
@chiaotung97
chiaotung97 force-pushed the feature_kda_integration branch 2 times, most recently from eb2b63f to ecc2171 Compare August 31, 2026 07:44
@github-actions github-actions Bot removed the stale label Aug 31, 2026
@fdz-1999

fdz-1999 commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator

@fdz-1999 fdz-1999 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review summary (prioritized):

P1: The required Code Quality Check is currently failing because mdformat changes docs/design/kda_cp_support.md. Since the workflow exits there, the remaining changed-file linters and downstream test jobs have not run. Please apply mdformat and confirm the complete required workflow passes.

P2: KimiDeltaAttention is not reachable from the standard MaxText decoder/model configuration; the smoke script builds a separate temporary model. Please either wire it into a production decoder following the hybrid Ling3 pattern, or narrow the PR title/description to state that this is a standalone layer/kernel integration and link a concrete decoder-integration follow-up.

Please also update the dependency section: the KDA change has landed internally, but the public GitHub status/release containing the API is not yet available. Distinguishing internal submission from the public version that MaxText can declare would make the reproducibility status clearer.

Comment thread src/maxtext/kernels/kda/tokamax.py Outdated
Comment thread scripts/dev/kda_e2e_smoke.py Outdated
Comment thread tests/unit/kda_attention_test.py
Comment thread src/maxtext/layers/attention_kda.py Outdated
Comment thread tests/unit/kda_attention_test.py Outdated
Comment thread docs/reference/kda_cp_support.md
Comment thread src/maxtext/utils/cp_utils.py Outdated
Comment thread tests/unit/kda_attention_test.py
Comment thread scripts/dev/kda_e2e_smoke.py Outdated
Comment thread scripts/dev/kda_e2e_smoke.py Outdated
- Add KimiDeltaAttention layer (attention_kda.py) with QKV projections,
  ShortConvolution, gate/beta/output-gate projections
- Add KDA kernel dispatch (kernels/kda/__init__.py) delegating to tokamax
- Add tokamax backend adapter (kernels/kda/tokamax.py) with layout translation
- Add CP utilities (cp_utils.py) for halo exchange and AG-CP support
- Add KdaAttention config class (types.py) with kda_backend field
- Add base.yml config entry for kda_backend
- Add comprehensive unit tests (kda_attention_test.py)
- Add KDA+CP support design doc (docs/design/kda_cp_support.md)
P0 fixes:
- Replace all AG-CP/All-Gather CP references with CP (23 occurrences)
- Remove tops/pallas-kernel references from base.yml and types.py
- Add comment explaining tokamax's pallas_tpu implementation name

P1 fixes:
- Remove unused kda_backend parameter from chunk_kda and config
- Update design doc scope to reflect one-time KDA+CP integration

P2 fixes:
- Replace assert statements with raise (NotImplementedError, ValueError, ImportError)
- Fix misleading test name (test_kda_cp_no_load_balance_ok -> test_kda_no_cp_without_load_balance_ok)
- Fix test method name: test_kda_ag_cp_equivalence -> test_kda_cp_equivalence
- Add warning when kda_lower_bound is set but safe_gate=False
- Add ge=0 constraint on linear_conv_kernel_dim in types.py
- Add field_validator for kda_lower_bound to reject NaN/Inf
- Apply pyink auto-formatting (line-length=122, indent=2)
- Fix design doc: Assert -> raise ImportError for CPContext check
…tention

Renames stale parameters to the finalized tokamax API (a_log,
delta_time_bias, use_qk_l2norm, max_num_segments,
context_parallel_metadata), updates config docs to the sigmoid
lower-bound gate semantics, and adds license headers.
- Thread cfg.context_sharding through the conv halo exchange, T-axis
  pspec injection (now an unconditional overwrite) and
  ContextParallelMetadata, fixing latent breakage under expert-as-context
  sharding.
- Fail fast with a config-level message when packed sequences are used
  without a positive max_segments_per_seq.
- Config validators: use_kda_safe_gate=True requires kda_lower_bound in
  [-5, 0); reject use_kda_lora=True (unimplemented no-op).
- Fix linear_conv_kernel_dim docs (convolution applies to Q/K/V, not
  only keys); refresh design doc file/test tables.
- New tests: CP=2/4 parametrized forward equivalence, CP gradient
  equivalence, full-layer CP with the internal dummy-segment path,
  parametrized ShortConv cross-rank segment boundaries, l2-norm
  unit-norm assertion, within-row packed-segment isolation, and config
  guard tests.
- pyink the e2e smoke script.
- e2e smoke: replace the permutation task with a history-dependent
  delayed-copy task (i.i.d. tokens, t[i] = t[i-delay], unpredictable
  positions masked out of the loss) so the smoke can no longer be
  solved without the recurrent KDA state; validate args via
  parser.error instead of assert
- tests: add full-layer CP with real packed segments (a segment
  spanning the rank boundary + a boundary exactly at the split;
  forward/input/weight grads vs non-CP), full-layer Mosaic vs tokamax
  XLA reference parity, and oversized CP-halo rejection; suite grows
  from 30 to 43 items, all passing on 4xTPU v6e
- tests: refine the tpu_only marker from module-level to per-test so
  pure config/pure-op/non-CP tests run in regular CPU CI; drop the
  unconditional prints in _assert_close (diagnostics now only in the
  assertion failure message); add _assert_rel_l2_close for accumulated
  weight-gradient comparisons
- cp_utils: raise a clear ValueError when halo_size > T_local under
  CP (multi-rank receptive field is not implemented)
- docs/design/kda_cp_support.md: sync snippets with the implementation
  (unconditional T-axis overwrite, cfg.context_sharding / expert-as-
  context), add new test entries, mdformat-clean
- tokamax adapter: document the deliberate lazy import; the pin bump
  must wait for the first public tokamax release containing KDA
  (openxla/tokamax#1103 is still open)
- types.py: drop superfluous parens after not (C0325, the finding that
  failed the Code Quality Check pylint step)
- tokamax adapter: validate initial_state/output_final_state before the
  lazy tokamax import so the guards fire on installs without tokamax
- tests: kernel input-guard tests (initial_state / output_final_state on
  both chunk_kda and the tokamax adapter), ShortConvolution
  feature-mismatch and kernel_size=1 cases, and a single-rank shard_map
  halo-exchange case — all CPU-runnable; plus tpu_only coverage for the
  no-conv forward path, the safe-gate warning, and the CP
  metadata-missing refusal
- The delayed-copy task with delay=4 could be solved by the 4-tap causal
  ShortConvolution alone (its window [j-3, j] contains the target
  t[j-3]), so the loss drop was not evidence of recurrent-state carry.
  Default the delay to 8, outside the receptive field, and enforce
  delay > linear_conv_kernel_dim at startup; bump the default step
  count to 600 for the harder task. Rerun on 4xTPU v6e: 5.43 -> 0.42,
  PASS.
- tokamax adapter TODO: the KDA API has landed on openxla/tokamax main
  (PR AI-Hypercomputer#1103 was left open but its content is on main, signature and
  ContextParallelMetadata both compatible); no public release contains
  it yet, so the pin bump still waits on the first release.
The reviewer's checklist item asks for new documentation pages to be added
to the relevant toctree. Move docs/design/kda_cp_support.md to
docs/reference/kda_cp_support.md (next to the peer MTP CP doc) and register
it in the Reference section: a navigation card plus a hidden-toctree entry,
same as the existing reference pages. Refresh the doc's Files Changed stats
and its self-reference to the new path.
@chiaotung97
chiaotung97 force-pushed the feature_kda_integration branch from 47766c3 to 402bc86 Compare September 3, 2026 18:20
Review-driven changes:
- scripts/dev/kda_e2e_smoke.py is converted into a proper test
  (TestKdaE2eSmoke::test_delayed_copy_loss_collapses), tuned to converge in
  ~1 min on v6e; the standalone script is removed (NuojCheng)
- KDA flags registered in configs/base.yml: linear_conv_kernel_dim,
  use_kda_safe_gate, kda_lower_bound, use_kda_lora (NuojCheng)
- halo_exchange_for_conv merged into attention_kda.py — it is KDA-specific;
  utils/cp_utils.py removed (NuojCheng)
- deprecated `with mesh:` context is gone from the tests: the mesh is
  passed explicitly everywhere and no ambient context is needed (verified
  empirically; jax.set_mesh cannot be used here — it forces Manual axis
  types that clash with the layer's explicit Auto-typed shard_maps)
- design doc moved to docs/reference/kda_cp_support.md and linked in the
  Reference toctree/card grid (checklist item)

Upstream CI fixes (the tpu-unit / pathways failures):
- every kernel-invoking test now carries
  skipif(not TOKAMAX_AVAILABLE) where TOKAMAX_AVAILABLE means the KDA API
  is importable — upstream CI installs a released tokamax without the KDA
  module, so those tests now skip cleanly there (they still run on hosts
  with KDA-enabled tokamax)
- the adapter's lazy import now raises an explanatory ImportError naming
  the missing tokamax KDA API and how to get it

Gemini review:
- CP without user segment_ids now synthesizes the all-ones segment tensor
  outside shard_map (n_max=1), simplifying the shard_map body and the
  kda_args construction (G2-G4)
- the G1 kernel.value suggestion is not applied: on current flax both
  `.value` and `[...]` access warn, and `[...]` is the form flax itself
  recommends for array variables; the original code was correct
Full decoder integration so attention_type='kda' trains through the real
MaxText stack (mirrors the MLA selection pattern):

- common_types: AttentionType.KDA
- types.py: attention_type Literal gains 'kda'; cross-validator rejects
  attention_type='kda' with scan_layers=true (KDA layers are not validated
  inside a scanned stack yet)
- base.yml: attention_type supported list + KDA flag block restored on top
  of pristine content (an earlier lint run had corrupted the YAML)
- nnx_decoders.NNXDecoderLayer: per-layer attention_type override + KDA
  branch that builds KimiDeltaAttention (no KV cache; recurrent state is
  carried by the kernel); non-KDA types unaffected
- KimiDeltaAttention now always L2-normalizes Q/K: the Delta-Rule
  recurrence diverges to NaN in bf16 with unbounded q/k (verified on v6e),
  and QK L2-norm is part of the KDA reference architecture; this is
  independent of the shared use_qk_norm flag, which belongs to dot-product
  attention

New tests (tests/unit/kda_decoder_integration_test.py):
- CPU: config acceptance, scan_layers guard, layer dispatch to
  KimiDeltaAttention vs regular Attention
- TPU: real-decoder training run (delayed-copy task) — loss collapses from
  5.28 to well under 1.0 in 400 steps through the actual decoder

Verified on 4xTPU v6e: 35 TPU tests + 23 CPU tests pass across both KDA
test files.
Upstream Code Quality flagged use-dict-literal in
kda_decoder_integration_test.py; convert the defaults dict(...) call to a
literal (pylint passes with the exact pre-commit hook invocation).
@chiaotung97
chiaotung97 force-pushed the feature_kda_integration branch from f91b160 to 176ec2f Compare September 7, 2026 03:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants