Implement QK attention head chunking in MLA to reduce HBM footprint#4564
Implement QK attention head chunking in MLA to reduce HBM footprint#4564zcjhao wants to merge 1 commit into
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
🤖 Hi @zcjhao, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
This pull request introduces head chunking over the QK attention dot product dimension in Multi-Head Latent Attention (MLA). By using jax.lax.scan to compute scores in blocks of attention heads, it prevents the full materialization of the query-key score tensor, substantially reducing the high-bandwidth memory (HBM) footprint for large context windows.
🔍 General Feedback
- Excellent Memory Efficiency: The introduction of chunking via sequential scan is a highly effective, mathematically correct design choice to mitigate OOMs at scale.
- Accidental Config Disabling: Please revert the change to
enable_checkpointinginbase.yml. Leaving checkpointing disabled by default is a major hazard for all production workloads. - Add Parity and Unit Tests: Although manual verification is mentioned, it is highly recommended to include automated regression coverage in
tests/unit/deepseek32_vs_reference_test.py. For example, a test namedtest_indexer_chunked_matchcould setqk_head_chunk_sizeto a divisor of the indexer heads (e.g.,16) and assert that the outputs are identical to the unchunked implementation. - Silent Fallback: As noted in the inline review, consider raising an explicit
ValueErrorwhenqk_head_chunk_sizeis misconfigured (e.g., if it exceeds the number of heads or does not divide them evenly), instead of falling back silently to unchunked materialization.
c601858 to
94eeccd
Compare
|
🤖 Hi @zcjhao, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This pull request implements QK attention head chunking in MLA (Multi-Head Latent Attention) using jax.lax.scan to significantly reduce the HBM footprint for long contexts. While the overall design and mathematical approach are excellent and provide great memory efficiency, a critical refactoring bug was introduced where the dynamic head importance weights became undefined, which completely breaks the forward pass.
🔍 General Feedback
- Crucial Definition Restored: Ensure the
weightstensor is properly calculated using theweights_projprojection layer before applying chunking to avoid a runtimeNameError. - Performance Optimizations: The proposed code changes optimize the scan loop bodies by formulating the
einsumcontractions directly on the chunked layouts, removing 4 transpose operations inside sequential loops. - Testing & Verification: Please add automated unit tests (e.g., in
tests/unit/deepseek32_vs_reference_test.py) that explicitly run with non-zeroqk_head_chunk_sizevalues and verify numerical parity with the unchunked version.
|
🤖 I'm sorry @zcjhao, but I was unable to process your request. Please see the logs for more details. |
94eeccd to
bab9fb9
Compare
|
🤖 Hi @zcjhao, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @zcjhao, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @zcjhao, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
bab9fb9 to
af8a014
Compare
There was a problem hiding this comment.
This pull request implements a highly effective chunked attention dot product for DeepSeek-V3.2 MLA. By chunking along the attention heads dimension during QK product evaluation, the implementation reduces the maximum HBM footprint of the intermediate attention scores tensor, allowing the model to handle much larger context lengths with native JAX primitives.
🔍 General Feedback
- High Quality and Parity: The implementation contains robust unit tests comparing the outputs of both the indexer and MLA module against a reference PyTorch implementation.
- Optimal Memory Trade-offs: The addition of
qk_head_chunk_sizeallows users to scale and control memory-compute trade-offs cleanly. - Clean Fallback Paths: Keeping the native unchunked branch as the default ensures backwards compatibility and zero performance regression when chunking is not enabled.
af8a014 to
e7a7b54
Compare
|
🤖 Hi @zcjhao, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This PR introduces QK attention head chunking to MLA to reduce the HBM footprint, particularly beneficial for long context lengths. By iteratively evaluating head groups via jax.lax.scan and calculating softmax partials sequentially, the design achieves identical mathematical parity to full tensor materialization with significantly reduced memory footprint.
🔍 General Feedback
- High-Quality Design: The mathematical implementation using
jax.lax.scanis precise and exactly matches the un-chunked fallback branch. - Testing Coverage Gap: Although the PR description mentions manually verifying parity, the existing unit test suite in
tests/unit/deepseek32_vs_reference_test.pydoes not parameterize overqk_head_chunk_sizenor test any non-zero chunk sizes in CI. To ensure this features does not regress, we suggest addingqk_head_chunk_sizeas a test parameter indeepseek32_vs_reference_test.pyas detailed below.
Suggested Test Improvements:
-
In
tests/unit/deepseek32_vs_reference_test.py, updateget_cfg_and_meshto supportqk_head_chunk_size:def get_cfg_and_mesh(config, run_name, dtype, batch_size, seq_len, attention, indexer_topk, qk_head_chunk_size=0): """Returns MaxText configuration and mesh.""" cfg = pyconfig.initialize( [None, get_test_config_path()], run_name=run_name, enable_checkpointing=False, model_name="default", dtype=dtype, weight_dtype="float32", matmul_precision="highest", float32_qk_product=True, float32_logits=True, per_device_batch_size=batch_size, max_target_length=seq_len, max_prefill_predict_length=seq_len, attention=attention, indexer_topk=indexer_topk, qk_head_chunk_size=qk_head_chunk_size, **asdict(config), ) # ... rest of method ...
-
Parameterize
test_indexer_matchto verify chunking configurations (e.g.,0,2,4):class DeepseekV32IndexerTest(DeepseekTestBase): """Tests for the Sparse Indexer (Top-K Selection).""" @parameterized.parameters(0, 2, 4) def test_indexer_match(self, qk_head_chunk_size, seq_len=8): """Verifies Indexer output matches PyTorch output.""" # ... cfg, mesh = get_cfg_and_mesh( config=self.config, run_name="deepseek_indexer_test", dtype=self.dtype, batch_size=self.batch_size, seq_len=self.seq_len, attention="dot_product", indexer_topk=4, qk_head_chunk_size=qk_head_chunk_size, ) # ...
-
Update
test_mla_parityparameter list and test cases to verify different chunk sizes:def test_mla_parity(self, attention, seq_len, indexer_topk, check_norm=False, qk_head_chunk_size=0): # ... cfg, mesh = get_cfg_and_mesh( config=self.config, run_name="deepseek_mla_test", dtype=self.dtype, batch_size=self.batch_size, seq_len=self.seq_len, attention=attention, indexer_topk=indexer_topk, qk_head_chunk_size=qk_head_chunk_size, ) # ...
53dffbd to
fc14e2e
Compare
|
🤖 Hi @zcjhao, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @zcjhao, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @zcjhao, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This pull request implements a highly effective memory optimization for Multi-Layer Attention (MLA) in DeepSeek models by introducing QK head chunking. By utilizing jax.lax.scan to sequentially compute the query-key dot product and its softmax over head chunks, the attention scores tensor footprint is drastically reduced from
🔍 General Feedback
- High Code Quality: The JAX implementation of chunking over the heads dimension is exceptionally clean, well-documented, and mathematically identical to the full tensor materialization.
-
Rematerialization Strategy: The use of
jax.checkpointon the scan body of the sparse indexer is extremely clever and essential to prevent the compiler from caching the intermediate logits, thereby successfully securing the memory reductions during training backpropagation. -
Comprehensive Verification: The addition of parameterized reference tests in
tests/unit/deepseek32_vs_reference_test.pyvalidates equivalence across multiple chunk sizes (0, 2, and 4) against PyTorch reference results, ensuring high reliability.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
fc14e2e to
9f2ab58
Compare
|
🤖 Hi @zcjhao, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This pull request implements head-wise attention chunking using jax.lax.scan for both the DeepSeek-V3.2 MLA sparse indexer scores computation and its KL divergence loss function. The implementation is highly efficient, mathematically equivalent to the native (un-chunked) implementation, and successfully reduces HBM footprint to avoid out-of-memory (OOM) failures under long context lengths.
🔍 General Feedback
- Excellent JAX Idiomatic Design: Leverages
jax.lax.scanto sequentially process slices of head dimension, resulting in a compiled loop (rather than Python unrolling) which keeps compile times and the XLA HLO graph extremely small. - Optimal Gradient Isolation and Checkpointing: Strategically applies
jax.checkpointonscan_body_indexerto avoid saving large intermediate logits on the backward pass, while correctly omitting checkpointing onscan_body_headswhere gradients are already stopped. - Robust Validation Guardrails: Comprehensive boundary checks (divisibility and size) are integrated on both
SparseIndexerandcalculate_indexer_lossblocks to prevent incorrect reshaping or compilation-level failures. - Thorough Test Coverage: New test cases successfully verify parity under multiple sweep sizes (
qk_head_chunk_sizeof 0, 2, 4) against the PyTorch reference implementation.
|
🤖 I'm sorry @zcjhao, but I was unable to process your request. Please see the logs for more details. |
96c937f to
75548aa
Compare
Currently, evaluating the QK dot product natively in MLA materializes the full [batch, count_of_heads, q_len, kv_len] attention scores tensor. This causes severe memory constraints (HBM OOM) for large context lengths. This change introduces the Config variable `qk_head_chunk_size` alongside a `jax.lax.scan` algorithm. Since the `heads` dimension is locally dense and unsharded (unlike `q_len` which is Context Parallel), we scan over groups of heads iteratively computing query-key projections and their softmax partials, avoiding vast intermediate allocations.
75548aa to
192491b
Compare
3 participants
Description
Currently, the QK dot product natively in MLA materializes the full
[batch, count_of_heads, q_len, kv_len]attention scores tensor. This causes memory constraints for large context lengths.This change introduces the Config variable
qk_head_chunk_sizealongside ajax.lax.scanalgorithm. Since theheadsdimension is locally unsharded (unlikeq_lenwhich is Context Parallelized), we scan over groups of heads iteratively over the QK dimension to compute query-key projections and their softmax partials.If
qk_head_chunk_sizeis left at the default0, the attention layer automatically falls back to the native, un-chunked tensor multiplication.
Tests
tests/unit/deepseek32_vs_reference_test.py::DeepseekV32IndexerTest::test_indexer_matchto automatically sweep multiple chunk sizes (qk_head_chunk_size=0,2, and4).qk_head_chunk_size=4) toDeepseekV32MLATest::test_mla_parityto verify mathematical equivalence with the PyTorch reference.Test output:
The two flash failures are unrelated to this PR, we do not use flash attention for QK in MLA currently. Running
python -m pytest tests/unit/deepseek32_vs_reference_test.py -k "test_mla_parity_flash" -vwould also fail in a clean workspace synced with the HEAD.qk_head_chunk_size=16):After
Before
Checklist
Before submitting this PR, please make sure (put X in square brackets):
[✓] 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 https://maxtext.readthedocs.
io/en/latest/development.html#adding-new-documentation-files.