DSV4: Add indexer_save_selection and indexer_sharded_topk for the CSA indexer - #4915
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces configuration flags and implementation details for sharding DeepSeek-V4 CSA indexer activations and top-k query rows, as well as saving indexer selections as rematerialization residuals. It also adds comprehensive unit tests to verify these features. The review feedback highlights a critical bug where a shape mismatch crash can occur in autoregressive decoding mode when slicing the attention mask, and points out several unused imports, constants, and functions in the newly added test files that should be cleaned up.
| entry_indices_mask = jnp.arange(compressed_len) | ||
| future_mask = entry_indices_mask[None, None, :] >= jnp.expand_dims(causal_threshold, axis=-1) | ||
|
|
||
| segment_mask = attention_mask[:, :, :compressed_len] if attention_mask is not None else None |
There was a problem hiding this comment.
In autoregressive decoding mode (MODEL_MODE_AUTOREGRESSIVE), attention_mask (which is compressed_segment_mask) has a sequence length dimension of 1 and is downsampled to 0 elements along the compressed dimension (since 1 // compress_rate is 0). However, compressed_len represents the total length of the compressed KV cache (which is greater than 0). Slicing attention_mask[:, :, :compressed_len] will return a tensor of shape (batch_size, 1, 0). Attempting to add this to index_scores of shape (batch_size, 1, compressed_len) will raise a ValueError due to shape mismatch.\n\nTo prevent this runtime crash, we should ensure segment_mask is only sliced and applied when not in autoregressive mode.
segment_mask = (\n attention_mask[:, :, :compressed_len]\n if attention_mask is not None and model_mode != MODEL_MODE_AUTOREGRESSIVE\n else None\n )| from pathlib import Path | ||
| import unittest | ||
|
|
||
| import jax | ||
| import jax.numpy as jnp | ||
| import numpy as np | ||
| import pytest | ||
| import yaml | ||
| from flax import nnx | ||
| from flax.linen import partitioning as nn_partitioning | ||
| from flax.linen.partitioning import logical_to_mesh_axes |
There was a problem hiding this comment.
The imports Path, yaml, and logical_to_mesh_axes are imported but never used in this test file. We should clean them up to keep the codebase tidy and maintainable.
import unittest\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport pytest\nfrom flax import nnx\nfrom flax.linen import partitioning as nn_partitioning| SCORE_AXES = ("activation_batch", "activation_heads", "activation_length", None) | ||
| CUSTOM_RULE_DIR = Path(pyconfig.__file__).parent / "custom_mesh_and_rule" |
There was a problem hiding this comment.
The constants SCORE_AXES and CUSTOM_RULE_DIR are defined but never used anywhere in this test file. They should be removed.
| def custom_rule_sets(): | ||
| """Loads logical-axis rules from each custom mesh preset.""" | ||
| out = {} | ||
| for path in sorted(CUSTOM_RULE_DIR.glob("*.yml")): | ||
| with path.open("r", encoding="utf-8") as file: | ||
| rules = yaml.safe_load(file)["logical_axis_rules"] | ||
| out[path.stem] = [(name, tuple(axes) if isinstance(axes, list) else axes) for name, axes in rules] | ||
| return out |
There was a problem hiding this comment.
The function custom_rule_sets is defined but never called or used anywhere in this test file. It should be removed to keep the test suite clean.
…ss heads Constrains the indexer's quadratic per-head score tensor over the activation_heads axis. Selections and masks are unchanged. Without it, seq 16384 exceeds the v6e per-chip memory limit by 4.42 GiB. Co-authored-by: Sudarsanan <[email protected]> Co-authored-by: Armin <[email protected]> Co-authored-by: utlz <[email protected]>
…V4 CSA indexer Save the indexer's selection as a remat residual instead of recomputing the chain in the backward pass, and shard the selection tail's rows over the tensor axis. Selections are unchanged. Measured together: 8.75 to 7.86 s/step (1.11x) on v6e-128 DeepSeek-V4 LoRA at seq 16384. Co-authored-by: Sudarsanan <[email protected]> Co-authored-by: Armin <[email protected]> Co-authored-by: utlz <[email protected]>
0ce0eaa to
c5d5205
Compare
2 participants
Description
Stacked on #4909 (the indexer sharding PR); the diff includes its commit until it merges. Review the top commit.
Adds
indexer_save_selectionto preserve the DeepSeek-V4 CSA indexer's final top-k selection across decoder rematerialization, avoiding recomputation of the indexer chain during the backward pass. The option composes with existing save and offload policies.Adds
indexer_sharded_topkto shard query rows for masking and top-k selection over the tensor mesh. Selected indices return to the standard sequence layout before sparse attention consumes them, while autoregressive selection remains unsharded because each query contains one row.Performance
Step time improved from 8.75 to 7.86 s/step (1.11x) with both
indexer_save_selectionandindexer_sharded_topkenabled.Measured on v6e-128 with DeepSeek-V4-Flash 284B LoRA fine-tuning, the dynamic-splash attention path enabled,
ici_expert_parallelism=16,ici_tensor_parallelism=8,ici_fsdp_parallelism=1,max_target_length=16384,per_device_batch_size=0.125(global batch 16), andLIBTPU_INIT_ARGS=--xla_tpu_scoped_vmem_limit_kib=98304, comparing both flags disabled and enabled.Tests
PYTHONPATH=$PWD/src JAX_PLATFORMS=cpu /home/jeremy/git/maxtext-16kmem/.venv/bin/python3 -m pytest -q tests/unit/indexer_selection_flags_test.py— 5 passed.top_krecomputation and match packed and unpacked selections on an 8-device tensor mesh./home/jeremy/git/maxtext-16kmem/.venv/bin/python3 -m pyink --check --pyink-indentation=2 --line-length=122 src/maxtext/configs/types.py src/maxtext/layers/attention_compressed.py src/maxtext/layers/nnx_decoders.py tests/unit/indexer_selection_flags_test.py— passed.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.