feat(kda): integrate KDA attention with tokamax backend and CP support by chiaotung97 · Pull Request #5128 · AI-Hypercomputer/maxtext

chiaotung97 · 2026-09-03T18:22:41Z

Description

This PR integrates KDA (Kimi Delta Attention, a recurrent linear attention mechanism) into MaxText end to end: the attention layer, the tokamax-backed Pallas TPU kernel, context parallelism (CP), and decoder wiring via attention_type='kda'. 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 layer 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, always-on QK L2 normalization, per-head RMSNorm + output gate, gate parameters `A_log` / `dt_bias` (matching the Megatron reference); also hosts `halo_exchange_for_conv` for CP
`src/maxtext/kernels/kda/`	New `chunk_kda` entry point + tokamax adapter: `[B,T,H,D]` ↔ `[H,B,T,D]` layout translation, lazy import of the KDA API, explicit `mosaic` implementation selection inside `shard_map`
`src/maxtext/common/common_types.py`	`AttentionType.KDA`
`src/maxtext/configs/types.py`	New `KdaAttention` config (`linear_conv_kernel_dim`, `use_kda_safe_gate`, `kda_lower_bound`, reserved `use_kda_lora`) + validators; `attention_type` Literal gains `"kda"`; `attention_type='kda'` with `scan_layers=true` is rejected at config time
`src/maxtext/configs/base.yml`	KDA flags registered; `attention_type` supported list
`src/maxtext/layers/nnx_decoders.py`	`NNXDecoderLayer` gains a per-layer `attention_type` override and a KDA branch: builds `KimiDeltaAttention` (no KV cache — recurrent state is carried by the kernel); non-KDA types are untouched
`tests/unit/kda_attention_test.py`	Layer/kernel/CP suite (see below)
`tests/unit/kda_decoder_integration_test.py`	Decoder-integration suite: config guards + a real-decoder training run
`docs/reference/kda_cp_support.md`	Design doc, linked into the Reference toctree

Using KDA

Select it like any other attention variant — this mirrors how attention_type='mla' is consumed:

attention_type: kda
scan_layers: false        # required (KDA layers are not validated in scanned stacks yet)
use_kda_safe_gate: true   # recommended: bounded-decay gate, keeps the recurrence stable
kda_lower_bound: -5.0     # common choice, required in [-5, 0) when the safe gate is on

Notes:

KDA always L2-normalizes Q/K — this is part of the KDA architecture (the Delta-Rule recurrence diverges with unbounded q/k), independent of the shared use_qk_norm flag which belongs to dot-product attention.
context_parallel_load_balance is rejected with KDA+CP (token order is load-bearing for the recurrent state).
KDA has no KV cache; autoregressive decoding is not implemented yet (raises a clear NotImplementedError). Train/prefill paths are supported.

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
Under CP without user segments, an all-ones segment tensor is synthesized outside shard_map so the kernel can always derive per-rank cu_seqlens / chain fields
ShortConvolution pulls kernel_size-1 left-context tokens from the previous CP rank via ppermute inside a dedicated shard_map; a clear ValueError guards the kernel_size-1 > T_local case (multi-rank receptive fields are not implemented)

Tests and Validation

Unit tests — 35 TPU tests + 23 CPU tests, all passing (4×TPU v6e, 2026-09-04):

Layer/kernel: forward/backward (activations + weights), determinism, padding, packed-segment isolation (cross-row and within-row), precision vs a pure-XLA recurrent reference (FP32/BF16, ULP-based fallback)
CP: kernel-level equivalence parametrized CP=2/CP=4, CP backward, full-layer CP without segments (dummy synthesis path), full-layer CP with real packed segments (segment spanning the rank boundary + boundary exactly at the split; fwd + input/weight grads vs non-CP), oversized-halo rejection, load-balance rejection, CP metadata-missing refusal
Composition: full-layer Mosaic vs tokamax XLA reference parity (same weights, fwd + grads)
Kernel guards: initial_state / output_final_state rejected with clear errors
Markers: tpu_only is applied per-test, so the pure config/pure-op tests also run in regular CPU CI

Decoder integration (tests/unit/kda_decoder_integration_test.py):

CPU-runnable: config acceptance, scan_layers guard, dispatch of NNXDecoderLayer to KimiDeltaAttention (and unchanged default for other types)
TPU: a real-decoder training run (decoder_block=default, 2 layers, delayed-copy task — i.i.d. tokens with t[i] = t[i-delay], delay=5 beyond the conv receptive field, so only the recurrent state can solve it; masked loss): loss collapses from 5.28 to < 1.0 in 400 steps through the actual decoder

Reproducing

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

# 1. TPU-capable JAX (validation 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 kernels — the KDA API has landed on openxla/tokamax
#    main but no release contains it yet; install from source:
git clone https://github.com/openxla/tokamax.git && pip install -e tokamax
#    (or the equivalent branch validated here: -b antgroup/kda-pallas-kernel, tip 939da5c)

# 4. Tests (kernel/CP tests are tpu_only and skip cleanly elsewhere)
pytest tests/unit/kda_attention_test.py tests/unit/kda_decoder_integration_test.py -v

Dependencies

The KDA kernels come from openxla/tokamax#1103 (Kimi Delta Attention Pallas kernels, by @Fred33146). Status: the code has landed on openxla/tokamax main (signature-compatible with this PR's adapter, verified), but no public tokamax release contains it yet. MaxText's pin therefore cannot express the dependency; the adapter imports the KDA API lazily, so installs without it are unaffected until KDA is actually used (clear ImportError with instructions at that point). Once the first release ships, the tokamax>= pin and derived requirement files will be bumped.
Validated with JAX 0.11.0 + libtpu 0.0.44.1

Hardware / Shape Constraints (mosaic kernel)

The adapter selects the "mosaic" Pallas implementation explicitly (no silent fallback). Constraints — surfaced as clear NotImplementedErrors from tokamax at 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

Autoregressive/inference path not implemented (train/prefill only; AR raises a clear error)
initial_state / output_final_state not yet supported (explicit NotImplementedError)
scan_layers must be false with attention_type='kda' (scanned KDA stacks not validated yet)
bf16 training-recipe tuning for KDA models (the integration test validates in fp32 with the safe gate)
Hybrid per-layer KDA + softmax models (e.g. Ling-style interleaving) and a production KDA model config
tokamax version pin bump once the first public release with the KDA API ships

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 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.

- 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

…smoke

…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.

gemini-code-assist

Code Review

This pull request integrates Kimi Delta Attention (KDA) into MaxText with tokamax backend and context parallelism (CP) support, introducing the KimiDeltaAttention layer, ShortConvolution, and CP-aware causal convolution boundary handling. The review feedback identifies several critical issues: a subscripting error on nnx.Param in ShortConvolution, a JAX compilation error when passing None to shard_map under CP without user-provided segment IDs, and an incorrect parameter count calculation in the smoke test script due to Flax NNX Variable wrapping. Addressing these issues by accessing .value on NNX parameters/variables and synthesizing dummy segment IDs outside of shard_map will ensure correctness and successful compilation.

gemini-code-assist · 2026-09-03T18:24:18Z

+    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


Under context parallelism (cp_size > 1) without user-provided decoder_segment_ids, passing None to shard_map while specifying seg_pspec in in_specs will cause a JAX compilation error because shard_map expects all sharded inputs to be JAX arrays.

To prevent this, synthesize the dummy decoder_segment_ids (all-ones) outside of shard_map and set n_max = 1 here. This ensures a valid JAX array is sharded and passed to shard_map normally.

Suggested change

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

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 decoder_segment_ids is None and cp_size > 1:

decoder_segment_ids = jnp.ones((B, T), dtype=jnp.int32)

n_max = 1

gemini-code-assist · 2026-09-03T18:24:18Z

+      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
+        max_num_segments = n_max
+
+        # CP: provide a dummy seg (all-ones) so the kernel has
+        # segment_ids to derive cu_seqlens from, even when the user
+        # hasn't supplied real segmentation info.
+        if seg is None and cp_size > 1:
+          seg = jnp.ones(q.shape[:2], dtype=jnp.int32)
+          max_num_segments = 1


Since decoder_segment_ids is now synthesized outside of shard_map, we can simplify the internal helper function by removing the None check and the local dummy array creation.

Suggested change

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

max_num_segments = n_max

# CP: provide a dummy seg (all-ones) so the kernel has

# segment_ids to derive cu_seqlens from, even when the user

# hasn't supplied real segmentation info.

if seg is None and cp_size > 1:

seg = jnp.ones(q.shape[:2], dtype=jnp.int32)

max_num_segments = 1

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

gemini-code-assist · 2026-09-03T18:24:18Z

+      if has_seg:
+        if decoder_segment_ids is not None:
+          kda_args = kda_args + (decoder_segment_ids,)
+        else:
+          # CP without varlen: pass None; the shard_map function
+          # synthesises a dummy seg internally.
+          kda_args = kda_args + (None,)
+      o = _shard_map_chunk_kda(*kda_args)


With decoder_segment_ids guaranteed to be a valid JAX array when has_seg is True, we can simplify the argument construction and avoid appending None to kda_args.

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,)

has_seg is now simply decoder_segment_ids is not None (CP without user segments synthesizes the tensor earlier), so kda_args never appends None.

gemini-code-assist · 2026-09-03T18:24:18Z

+  rngs = nnx.Rngs(0)
+  with mesh:
+    model = TinyKdaLM(cfg, mesh, args.num_layers, rngs=rngs)
+  n_params = sum(v.size for v in jax.tree.leaves(nnx.state(model)) if isinstance(v, (jax.Array, np.ndarray)))


In Flax NNX, nnx.state(model) returns a State dict containing Variable objects (like Param). Checking isinstance(v, (jax.Array, np.ndarray)) on these variables will evaluate to False, resulting in n_params being calculated as 0. Access the underlying array using v.value for any nnx.Variable instead.

Suggested change

n_params = sum(v.size for v in jax.tree.leaves(nnx.state(model)) if isinstance(v, (jax.Array, np.ndarray)))

n_params = sum(v.value.size for v in jax.tree.leaves(nnx.state(model)) if isinstance(v, nnx.Variable))

codecov · 2026-09-03T18:29:01Z

Codecov Report

❌ Patch coverage is 79.69925% with 54 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_kda.py	82.96%	30 Missing and 9 partials ⚠️
src/maxtext/kernels/kda/tokamax.py	42.10%	10 Missing and 1 partial ⚠️
src/maxtext/kernels/kda/__init__.py	77.77%	1 Missing and 1 partial ⚠️
src/maxtext/layers/nnx_decoders.py	75.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

NuojCheng · 2026-09-03T18:54:20Z

+  print(f"devices: {devices}")
+  mesh = jax.sharding.Mesh(np.array(devices), ("x",))
+  rngs = nnx.Rngs(0)
+  with mesh:


with jax.set_mesh(mesh):

with mesh is deprecated

Removed entirely rather than migrating to jax.set_mesh: the mesh is passed explicitly to the layer/shard_map.
It forces Manual axis types that clash with the layer's explicit Auto-typed shard maps. Consistent with upstream CP tests, which also use no ambient mesh context.

NuojCheng · 2026-09-03T18:55:44Z

@@ -0,0 +1,235 @@
+# Copyright 2026 Ant Group. All Rights Reserved.


could you either update this as a test, or move it to src/maxtext/examples/.

Converted to a test: TestKdaE2eSmoke::test_delayed_copy_loss_collapses (tpu_only, ~1 min on v6e). It keeps the history-dependent delayed-copy task (delay chosen to exceed the conv receptive field so only the KDA recurrent state can solve it) and asserts the loss collapses. The scripts/dev/ script is removed.

NuojCheng · 2026-09-03T18:56:25Z

  )


+class KdaAttention(BaseModel):


please also add new flags in configs/base.yml

Done, linear_conv_kernel_dim, use_kda_safe_gate, kda_lower_bound, use_kda_lora are registered in configs/base.yml next to the other attention options, matching the KdaAttention field names/defaults in types.py.

NuojCheng · 2026-09-03T18:57:44Z

@@ -0,0 +1,95 @@
+# Copyright 2026 Ant Group. All Rights Reserved.


if this util is not for kda, I suggest merging it into attention_kda.py

Yes, it's KDA-specific, so it is good to merge into attention_kda.py and removed utils/cp_utils.py, Thanks.

NuojCheng

Thank you for contribution! Some general architecture comments before diving deep

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

chiaotung97 · 2026-09-04T06:03:57Z

Thank you for contribution! Some general architecture comments before diving deep

Thank you for detailed comments. All four addressed in the latest push:
(1) dropped the deprecated with mesh: the mesh is passed explicitly and no ambient context is needed.
(2) the smoke script is now a proper test TestKdaE2eSmoke::test_delayed_copy_loss_collapses.
(3) the KDA flags are registered in configs/base.yml.
(4) halo_exchange_for_conv is KDA-specific, so it's merged into attention_kda.py. Details in the inline replies.

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.

chiaotung97 · 2026-09-04T09:05:03Z

Codecov Report

❌ Patch coverage is 77.77778% with 58 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/layers/attention_kda.py 77.88% 37 Missing and 9 partials ⚠️
src/maxtext/kernels/kda/tokamax.py 50.00% 7 Missing and 1 partial ⚠️
src/maxtext/kernels/kda/init.py 77.77% 1 Missing and 1 partial ⚠️
src/maxtext/utils/cp_utils.py 92.85% 1 Missing and 1 partial ⚠️
📢 Thoughts on this report? Let us know!

The patch number reflects what CI can measure, not what the code covers: every KDA kernel-invoking path is exercised by tpu_only tests, additionally gated by skipif(KDA API available). CI installs the public tokamax release, which does not contain the KDA module yet (the kernels landed on openxla/tokamax main but no release has been cut — see the Dependencies section), so those lines are correctly skipped in CI rather than executed.

Running the full suite on an actual TPU host (4×TPU v6e, both TPU mode and CPU mode) measures 99% line coverage on the new KDA source files:

File	Coverage
`src/maxtext/kernels/kda/__init__.py`	100%
`src/maxtext/kernels/kda/tokamax.py`	100%
`src/maxtext/layers/attention_kda.py` (layer + CP halo exchange)	99%

The only uncovered lines are the except ImportError fallback for tokamax-less installs, which is unreachable by construction whenever the KDA API is present. The decoder-wiring path (attention_type='kda') is covered by tests/unit/kda_decoder_integration_test.py (CPU guards on every runner + a real-decoder training run on TPU).

Once tokamax ships the first release containing the KDA API and CI can install it, the kernel paths will be exercised in CI as well and the patch number will reflect them.

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 added 11 commits September 4, 2026 02:17

fix: pyink formatting and doc consistency

77456fc

- Apply pyink auto-formatting (line-length=122, indent=2) - Fix design doc: Assert -> raise ImportError for CPContext check

fix(kda): use mosaic implementation name, mark UTs tpu_only, add e2e …

233be2b

…smoke

chiaotung97 requested review from A9isha, NuojCheng, RissyRan, SurbhiJainUSC, aireenmei, bvandermoon, darisoy, gagika, gobbleturk, huytransformer, igorts-git, jacoguzo, jiangjy1982, parambole, richjames0, shralex, shuningjin, vipannalla and xibinliu as code owners September 3, 2026 18:22

chiaotung97 requested review from abhinavclemson, dipannita08, hengtaoguo and khatwanimohit as code owners September 3, 2026 18:22

gemini-code-assist Bot reviewed Sep 3, 2026

View reviewed changes

NuojCheng reviewed Sep 3, 2026

View reviewed changes

chiaotung97 changed the title ~~feat(kda): integrate KDA attention with tokamax backend and CP support- #1~~ feat(kda): integrate KDA attention with tokamax backend and CP support Sep 4, 2026

style(kda): use dict literal in test config helper (pylint R1735)

176ec2f

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).

	n_params = sum(v.size for v in jax.tree.leaves(nnx.state(model)) if isinstance(v, (jax.Array, np.ndarray)))
	n_params = sum(v.value.size for v in jax.tree.leaves(nnx.state(model)) if isinstance(v, nnx.Variable))

		@@ -0,0 +1,235 @@
		# Copyright 2026 Ant Group. All Rights Reserved.

		@@ -0,0 +1,95 @@
		# Copyright 2026 Ant Group. All Rights Reserved.

Uh oh!

Conversation

chiaotung97 commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Key Changes

Using KDA

Context Parallelism Support

Tests and Validation

Reproducing

Dependencies

Hardware / Shape Constraints (mosaic kernel)

Known Limitations / Follow-ups

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist Bot Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng left a comment

Choose a reason for hiding this comment

Uh oh!

chiaotung97 commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chiaotung97 commented Sep 4, 2026

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

chiaotung97 commented Sep 3, 2026 •

edited

Loading

codecov Bot commented Sep 3, 2026 •

edited

Loading

chiaotung97 commented Sep 4, 2026 •

edited

Loading