Implement QK attention head chunking for CSA [Deepseek v4] - #5116
Open
octatrifan wants to merge 3 commits into
Open
Implement QK attention head chunking for CSA [Deepseek v4]#5116octatrifan wants to merge 3 commits into
octatrifan wants to merge 3 commits into
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
4 tasks
octatrifan
force-pushed
the
octatrifan-dsv4-qk-chunking
branch
2 times, most recently
from
September 3, 2026 22:25
f2ab1ae to
1317a0f
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
octatrifan
force-pushed
the
octatrifan-dsv4-qk-chunking
branch
from
September 3, 2026 23:26
08002f1 to
9979754
Compare
octatrifan
marked this pull request as ready for review
September 4, 2026 03:45
octatrifan
requested review from
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
khatwanimohit,
parambole,
richjames0,
shralex,
shuningjin,
vipannalla and
xibinliu
as code owners
September 4, 2026 03:45
octatrifan
requested review from
A9isha,
darisoy and
dipannita08
as code owners
September 4, 2026 03:45
octatrifan
force-pushed
the
octatrifan-dsv4-qk-chunking
branch
from
September 4, 2026 20:53
9979754 to
ad65b6c
Compare
dandragona
reviewed
Sep 4, 2026
| raw_probs = jax.nn.softmax(safe_scores.astype(jnp.float32), axis=-1) | ||
| raw_probs = jnp.where(valid_tokens_mask[:, None, :, None], raw_probs, 0.0) | ||
| target_probs = jnp.sum(raw_probs, axis=1) | ||
| target_probs = jax.lax.optimization_barrier(target_probs) |
Collaborator
There was a problem hiding this comment.
Do we need this in the chunking branch too? Tbh not sure.
dandragona
reviewed
Sep 4, 2026
| float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product | ||
| float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax | ||
| mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in the Indexer across the unsharded local heads dimension natively. | ||
| mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in MLA and CSA Indexers across the unsharded local heads dimension natively. |
Collaborator
There was a problem hiding this comment.
Might already exist, but raise an error when we validate flags to ensure this divides head size.
dandragona
reviewed
Sep 4, 2026
| float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product | ||
| float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax | ||
| mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in the Indexer across the unsharded local heads dimension natively. | ||
| mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in MLA and CSA Indexers across the unsharded local heads dimension natively. |
Collaborator
There was a problem hiding this comment.
Also I think we have overloaded this flag as it is for MLA originally and now its name is confusing.
I think you should rename it and update:
google3/cloud/cluster/supercomputer/ubench/core/workloads/maxtext_training/configs/tpu7x-4x8x8/deepseekv3.2_671b_128K.yaml
Or just split to its own flag.
dandragona
reviewed
Sep 4, 2026
| np.testing.assert_array_equal(np.array(topk_chunked), np.array(topk_native)) | ||
| np.testing.assert_allclose(np.array(scores_chunked), np.array(scores_native), rtol=1e-5, atol=1e-5) | ||
|
|
||
| def test_csa_indexer_chunked_gradients_flow(self): |
Collaborator
There was a problem hiding this comment.
I don't think we are testing whether the chunked gradients are equal to the non-chunked gradients.
dandragona
reviewed
Sep 4, 2026
dandragona
left a comment
dandragona
left a comment
Collaborator
There was a problem hiding this comment.
Mostly minor comments.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Stacked on PR 5033 .$[B, H, S, S]$ QK interaction tensors across all heads.
Dense attention scoring materializes full
This causes OOM issues at longer context lengths.
PR #4564 addressed this for MLA by chunking over the heads dimension with
jax.lax.scan. This PR extends that optimization to DeepSeek-V4 Compressed Sparse Attention (CSA).This requires some changes as CSA uses compressed blocks.
BUGS: b/556431288
Changes
Chunking is applied to two sites using
mla_qk_head_chunk_size:DeepseekV4Indexer: Scans overindex_n_heads // head_chunk_sizeusingjax.checkpointon the scan body so backward activation memory stays bounded to 1 chunk.calculate_csa_indexer_loss: Scans overheads // head_chunk_size. Since teacher inputs usestop_gradient,jax.checkpointis omitted.Defaults to
0(native unchunked path).Tests
Added 3 unit tests in
tests/unit/deepseek_v4_indexer_loss_test.py:test_csa_indexer_scoring_head_chunking_parity: Verifies mathematical parity (test_csa_indexer_loss_head_chunking_parity: Verifies mathematical parity (test_csa_indexer_chunked_gradients_flow: Verifies gradient flow through the chunked scan body into query, key, and indexer projections.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.