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 7 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. |
bc9356c to
eb2b63f
Compare
- 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
eb2b63f to
ecc2171
Compare
…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.
fdz-1999
left a comment
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.
| Returns: | ||
| (o, None) where o is [B, T, H, V]. | ||
| """ | ||
| from tokamax._src.ops.experimental.kda.api import kimi_delta_attention |
There was a problem hiding this comment.
[P1] This imports the KDA API introduced by openxla/tokamax#1103, but the declared dependency still allows tokamax>=0.0.12, which does not guarantee that this module exists. Could we bump the minimum version to the first public tokamax release containing KDA and regenerate the derived requirement files? Otherwise a clean MaxText installation can satisfy the requirements but fail at runtime when KDA is used.
| dtype=cfg.dtype, | ||
| weight_dtype=cfg.weight_dtype, | ||
| rngs=rngs, | ||
| ) |
There was a problem hiding this comment.
[P1] This task is solvable from the current token alone: because t[i+1] = perm[t[i]], the residual embedding and MLP can learn the mapping even if KDA always returns zero. The loss decrease therefore verifies that the training loop runs, but not that recurrent attention works. Could we use a history-dependent task such as delayed copy/associative recall, or add an ablation showing that zeroing KDA fails? Otherwise this should be described as an optimizer/compilation smoke test rather than an end-to-end KDA correctness test.
|
|
||
| @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): |
There was a problem hiding this comment.
[P1] The tests cover segment-aware ShortConvolution and CP KDA separately, but this full-layer CP test only exercises synthesized all-ones segments. Could we add a full KimiDeltaAttention CP case with multiple real packed segments—including a segment spanning the rank boundary and a boundary exactly at the split—and compare forward and input/parameter gradients with non-CP? This validates convolution masking, segment forwarding, padding, and recurrent-state reset when composed.
| else: | ||
| # CP without varlen: pass None; the shard_map function | ||
| # synthesises a dummy seg internally. | ||
| kda_args = kda_args + (None,) |
There was a problem hiding this comment.
[P2] Kernel parity against the recurrent reference does not validate the complete layer composition. Could we add one full-layer numerical parity test using the same parameters with the tokamax XLA implementation/reference, covering QKV projection, ShortConvolution, gate/beta transforms, output RMSNorm/gating, and output projection? This would catch wiring errors that kernel-only tests cannot detect.
| # The chunk_kda kernel tests exercise tokamax Pallas TPU kernels and multi-chip | ||
| # CP; mark the module tpu_only so CPU-only testbeds skip them (consistent with | ||
| # kernels_test.py) while they run on TPU hosts. | ||
| pytestmark = pytest.mark.tpu_only |
There was a problem hiding this comment.
[P2] The module-level tpu_only marker also skips the pure configuration, normalization, and non-CP ShortConvolution tests on CPU CI. Could we apply tpu_only only to tests that invoke the Mosaic kernel or multi-device CP? Keeping the pure tests in regular CI would provide faster regression coverage.
| spec = list(pspec) | ||
| if spec[t_axis] is None: | ||
| spec[t_axis] = "context" | ||
| return jax.sharding.PartitionSpec(*spec) |
There was a problem hiding this comment.
[P2] This snippet is stale relative to the implementation. The code now overwrites the T-axis partition unconditionally and uses cfg.context_sharding, while the document only replaces None and hard-codes "context". Could we update the example and surrounding text so the documented contract also matches expert-as-context?
| cp_size = jax.lax.psum(1, axis_name=axis_name) | ||
| if cp_size == 1: | ||
| return zero_padded | ||
|
|
There was a problem hiding this comment.
[P2] linear_conv_kernel_dim has no upper bound, but this exchange only reads from the immediately preceding rank and assumes halo_size <= T_local. With a larger kernel or short local sequence, the required history can span multiple ranks. Could we either implement that case or fail clearly when kernel_size - 1 > T_local, with a boundary test?
| 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]) |
There was a problem hiding this comment.
[P3/nit] Could we avoid unconditional print calls in this assertion helper? They make normal test output noisy. Please use logging, or include the diagnostics only in the assertion failure message.
| seqs = np.empty((num_seqs, seq_len + 1), dtype=np.int32) | ||
| seqs[:, 0] = start | ||
| for i in range(seq_len): | ||
| seqs[:, i + 1] = perm[seqs[:, i]] |
There was a problem hiding this comment.
[P3/nit] Please use an explicit ValueError or parser.error rather than assert for user-input validation, since assertions can be disabled with optimized Python.
2 participants
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
tokamax>=pin insrc/dependencies/requirements/must be bumped to the first release containing the KDA API.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.