[WIP] DeepSeek-V3.2 model support

[sdmyzlp]

Updated in 7.14:

Moved code into deepseek_v3, and with first attempt for indexer loss.

---

Adds the DeepSeek-V3.2 (DeepSeek-V3.2-Exp) transformer model with Lightning Indexer + DeepSeek Sparse Attention (DSA). Based on DeepSeek-V3 modeling in torchtitan, and this PR depends on #3787.

Indexer — Lightning Indexer for sparse token selection

• Projects query (via wq_b from the MLA latent query), and key (via wk) into a low-rank index space.
• Applies a LayerNorm (k_norm), non-interleaved RoPE (rotate-half, using CosSinRoPE with yarn-scaled frequencies), and a Hadamard rotation to both index query and key.
• Produces per-token index weights (weights_proj).
• Static select() computes \sum_h \mathrm{relu}(q_h\cdot k) \cdot w_h, masks with the causal+document attention mask, and returns a per-token boolean selection via top-k.

DSAFlexAttention — fused sparse inner attention

• Extends FlexAttention. Calls Indexer.select to get the per-token selection, composes it with the document mask into a combined BlockMask, then delegates to standard flex_attention with that sparse mask.
• Static build_dsa_block_mask builds the combined BlockMask.

Attention — V3 MLA + Indexer + DSA

• Extends DeepSeek-V3's Attention. In addition to the MLA q/k/v projections, it runs the Indexer on (x, qr) and passes the index tensors into DSAFlexAttention.

Sharding

• Indexer weights: wq_b / weights_proj colwise S(0)@TP; wk / k_norm / rope cache Replicate@TP.
• Inner attention: all positional inputs are head-allgathered to Replicate@TP (q/k/v are gathered for the indexer KL loss; idx_* are gathered for select()'s reduce_sum over index heads). The out_src is Replicate@TP (truthful, since the inputs are replicated), and out_dst reshards to S(2)@TP to feed the rowwise wo.

TODO

• Indexer loss support — add the KL-loss between indexer scores and full attention scores (requires per-head attention score output and a head-dim allgather, which the current in_dst = Replicate@TP on q/k/v already prepares for).

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/8gpu

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

[tianyu-l]

Thanks for the proposal. I wonder what's your thought on supporting DeepSeek V4 #3634 vs. V3.2

[sdmyzlp]

@tianyu-l Thanks for the quick response. For context: we've been maintaining V3.2 and V4 on our internal version of torchtitan, and would like to polish and contribute them upstream.

We're also actively improving the V4 PR (#3634); these are parallel efforts. Currently DeepSeek-V3.2 is put in a separate folder. Would you prefer merging it into the existing DeepSeek-V3 directory instead, since it's more of a variant of the original V3, to keep the model directory cleaner? FMPV, V3.2 is simpler, while also being able to exercise new features like indexer loss and sparse attentions.

We'd appreciate your thoughts on any potential concerns, and are willing to iterate based on the project's direction.

[tianyu-l]

we are migrating from full_dtensor to spmd_types -- would you be interested in help migrate this model?

[sdmyzlp]

Sure, will investigate this.

[tianyu-l]

this is debugmodel only, can we add real configs?

[tianyu-l]

would benefit from a README.md listing out the distributed training techniques supported and verified, including parallelisms, per-layer compilation, activation checkpointing, etc.

[tianyu-l]

Please report numerics compared with HF native impl. See https://github.com/pytorch/torchtitan/tree/main/scripts/checkpoint_conversion for more details. We don't have to land the custom scripts, but good to convince ourselves.

[sdmyzlp]

OK, thanks for the pointer, I briefly tested with fp32 against DeepSeek-V3.2 HF, confirmed bitwise-exact logits given CosSinRoPE is modified per #3787. Will share a more comprehensive result later.