feat(kda): integrate KDA attention with tokamax backend and CP support - #1
feat(kda): integrate KDA attention with tokamax backend and CP support#1chiaotung97 wants to merge 14 commits into
Conversation
Unit Test Results30/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 Testssource ../maxtext_venv/bin/activate
cd ~/maxtext
python -m pytest tests/unit/kda_attention_test.py -vResultsEnvironment
Key Coverage
|
|
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. |
eb2b63f to
ecc2171
Compare
fdz-1999
left a comment
There was a problem hiding this comment.
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.
- 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.
47766c3 to
402bc86
Compare
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).
f91b160 to
176ec2f
Compare
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:
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
src/maxtext/layers/attention_kda.pyKimiDeltaAttentionlayer andShortConvolution: QKV/beta/gate/output-gate projections, depthwise causal 1D convolution, SiLU activation, optional QK L2 normalization, per-head RMSNorm + output gate, and gate parametersA_log/dt_bias(matching the Megatron reference)src/maxtext/kernels/kda/chunk_kdakernel entry point and tokamax adapter:[B,T,H,D]↔[H,B,T,D]layout translation; the non-auto-partitionable tokamax kernel is invoked insideshard_mapwith explicit partition specssrc/maxtext/configs/types.pyKdaAttentionconfig (linear_conv_kernel_dim,use_kda_safe_gate,kda_lower_bound, reserveduse_kda_lora) plus validators: safe gate requireskda_lower_bound ∈ [-5, 0);use_kda_lora=Trueis rejected as unimplementedsrc/maxtext/utils/cp_utils.pyhalo_exchange_for_conv: under CP, pullskernel_size-1tokens of left context from the previous CP rank viappermuteso causal convolution stays correct at CP shard boundaries (without it the exchange would silently degrade to zero-pad)tests/unit/kda_attention_test.pytpu_only)docs/design/kda_cp_support.mdscripts/dev/kda_e2e_smoke.pyContext Parallelism Support
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, andContextParallelMetadataContextParallelMetadatapasses mesh information to thechunk_kdakernel; tokamax derives per-rankcu_seqlens/ chain fields internally fromsegment_idsand coordinates recurrent state across CP rankscontext_parallel_load_balanceis rejected up front (DUAL_CHUNK_SWAP reordering would break the sequential dependency)segment_ids, the CP path synthesizes all-ones segment ids so the kernel can still derive its metadataShortConvolutionhalo exchange is wrapped in ashard_mapexposing the CP axis, sinceppermuteis a collectiveTests and Validation
Unit tests —
30 passed + 10 added in review round = 40/40 passing:82a9fc1e(test log:pytest tests/unit/kda_attention_test.py -v)_l2_normalizeuse_kda_lorarejection, packing withoutmax_segments_per_seqload_balancerejectionEnd-to-end smoke (4×TPU v6e): a 5.4M-param tiny model built from real
KimiDeltaAttentionlayers, trained for 400 steps on a fully learnable synthetic next-token task: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):
Dependencies
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 clearNotImplementedErrors from tokamax at kernel bind time:Known Limitations / Follow-ups
initial_state/output_final_statenot yet supported (explicitNotImplementedError)Checklist
gemini-reviewlabel.