[dev] DeepSeek V3.2 support

[kunlunl]

DeepSeek V3.2 Sparse Attention Support

Main PR: #2440

1. TL;DR

What: This PR adds support for DeepSeek V3.2-style sparse attention (DSA) to Megatron-LM, enabling models to use learned attention sparsity patterns via a lightweight indexer module.

Why: Dense attention has O(n²) complexity which becomes computationally prohibitive for long sequences. DSA reduces this by learning to predict which key-value pairs are most relevant for each query, allowing the model to attend to only top-k tokens instead of all tokens.

Impact: Users can now train models with DeepSeek V3.2's sparse attention mechanism in Megatron-LM, which combines Multi-Latent Attention with a trainable sparse indexer. The PR also refactors the linear attention infrastructure to be more extensible for future attention variants.

---

2. Big Picture

2.1 Before vs After Architecture

graph TB
    subgraph "Before: Standard Dense Attention"
        A1[Hidden States] --> B1[QKV Projection]
        B1 --> C1[Dense Attention]
        C1 --> D1[Output Projection]
    end
    
    subgraph "After: DeepSeek Sparse Attention"
        A2[Hidden States] --> B2[QKV Projection]
        A2 --> E2[DSA Indexer]
        B2 --> E2
        E2 --> F2[Top-K Selection]
        B2 --> G2[Sparse Attention]
        F2 --> G2
        G2 --> D2[Output Projection]
        E2 --> H2[Indexer Loss]
        H2 --> D2
    end
    
    style E2 fill:#90EE90
    style F2 fill:#90EE90
    style H2 fill:#FFD700

Loading

Key Changes:

• NEW: DSA Indexer module that learns to predict important tokens
• NEW: Sparse attention module that only computes attention for top-k tokens
• NEW: KL divergence auxiliary loss to train the indexer
• MODIFIED: Multi-Latent Attention to support DSA variant
• REFACTORED: Linear attention specs → Experimental attention variant specs (more general)

2.2 Change Scope Summary

Category	Files	Description
New Core Module	megatron/core/transformer/experimental_attention_variant/dsa.py	DSA indexer, DSA sparse attention module, loss computation
New Spec File	megatron/core/models/gpt/experimental_attention_variant_module_specs.py	Module specs for attention variants
New Test	tests/unit_tests/transformer/test_attention_variant_dsa.py	Comprehensive DSA unit tests
Modified Core	megatron/core/transformer/multi_latent_attention.py	MLA integration with DSA
Modified Config	megatron/core/transformer/transformer_config.py	Added DSA config parameters
Modified Args	megatron/training/arguments.py	CLI arguments for DSA
Modified Training	megatron/training/training.py	Loss logging for indexer
Modified Specs	megatron/core/models/gpt/gpt_layer_specs.py	Renamed linear_attention → experimental_attention_variant
Modified Builder	gpt_builders.py	Updated to use new attention variant system
Deleted	megatron/core/models/gpt/linear_attention_module_specs.py	Replaced by more general experimental_attention_variant_module_specs.py

---

3. Key Design Points

Core Abstractions Introduced:

1. DSAIndexer: Computes index scores to identify top-k most relevant tokens

   • Input: Hidden states x [seqlen, batch, hidden_size] + compressed query qr [seqlen, batch, q_lora_rank]
   • Output: Top-k indices [batch, seqlen, index_topk]
   • Uses its own small transformer-like architecture with Q/K projections + RoPE + Hadamard rotation

2. DSAttention: Sparse attention mechanism using indexer outputs

   • Wraps DSAIndexer and applies sparse attention kernel
   • Attaches KL divergence loss to train indexer

3. DSAIndexerLossAutoScaler: Custom autograd function

   • Allows indexer loss to backpropagate independently of main loss
   • Scales indexer loss gradient separately

Interface Contracts:

# DSAIndexer.forward
def forward(x, qr, mask=None, packed_seq_params=None) -> topk_indices
    """
    x: [seqlen, batch, hidden_size] - Main hidden states (DETACHED)
    qr: [seqlen, batch, q_lora_rank] - Compressed query (DETACHED)
    mask: [batch, seqlen, seqlen] - Attention mask (FP32 with -inf for masked positions)
    
    Returns: [batch, seqlen, index_topk] - Indices of top-k tokens to attend to
    """

# DSAttention.forward
def forward(query, key, value, x, qr, attention_mask, ...) -> output
    """
    query: [sq, b, np, hn] - Full query tensor from MLA
    key: [sk, b, np, hn] - Full key tensor from MLA
    value: [sk, b, np, hnv] - Full value tensor from MLA
    x: [sq, b, hidden_size] - Original hidden states for indexer
    qr: [sq, b, q_lora_rank] - Compressed query for indexer
    
    Returns: [sq, b, hidden_size] - Attention output with indexer loss attached
    """

Important Invariants:

• Indexer inputs (x, qr) are always detached - gradients don't flow back to main model
• Indexer loss is attached via DSAIndexerLossAutoScaler.apply() - backpropagates separately
• Top-k selection uses masked index scores (causal mask applied before topk)
• DSA currently requires multi_latent_attention=True and context_parallel_size=1

---

4. Execution Path Deep Dive

4.1 Entry Point

DSA is triggered when creating a GPT model with --experimental-attention-variant dsa flag:

# Entry: gpt_builders.py::gpt_builder()
def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None):
    # ...
    linear_attention_variants = ["gated_delta_net"]
    if args.num_experts or args.experimental_attention_variant in linear_attention_variants:
        transformer_layer_spec = get_gpt_decoder_block_spec(...)  # Uses MoE path
    elif:
        # ...
    else:
        transformer_layer_spec = _get_transformer_layer_spec(
            # ...
            experimental_attention_variant=args.experimental_attention_variant,  # 'dsa'
            # ...
        )

4.2 Call Chain Visualization

sequenceDiagram
    participant User
    participant GPTBuilder
    participant LayerSpec
    participant AttentionSpec
    participant MLAAttention
    participant DSAttention
    participant DSAIndexer
    
    User->>GPTBuilder: gpt_builder(args with --experimental-attention-variant dsa)
    GPTBuilder->>LayerSpec: get_gpt_layer_with_transformer_engine_spec(...)
    LayerSpec->>AttentionSpec: get_attention_module_spec_for_backend(experimental_attention_variant='dsa')
    AttentionSpec->>AttentionSpec: get_experimental_attention_variant_module_spec_for_backend()
    AttentionSpec->>AttentionSpec: get_dsa_module_spec_for_backend()
    
    Note over AttentionSpec: Creates ModuleSpec for MLASelfAttention<br/>with DSAttention as core_attention
    
    AttentionSpec-->>LayerSpec: ModuleSpec(MLASelfAttention + DSAttention)
    LayerSpec-->>GPTBuilder: transformer_layer_spec
    
    Note over GPTBuilder: At runtime (forward pass)...
    
    User->>MLAAttention: forward(hidden_states, ...)
    MLAAttention->>MLAAttention: get_query_key_value_tensors(return_compressed_tensors=True)
    Note over MLAAttention: Returns query, key, value, q_compressed, _
    
    MLAAttention->>DSAttention: forward(query, key, value, x=hidden_states, qr=q_compressed, ...)
    DSAttention->>DSAIndexer: forward_with_scores(x, qr, mask)
    DSAIndexer->>DSAIndexer: Compute index scores via indexer network
    DSAIndexer-->>DSAttention: (index_scores, topk_indices)
    
    DSAttention->>DSAttention: unfused_dsa_fn(query, key, value, topk_indices)
    Note over DSAttention: Sparse attention only on top-k tokens
    
    DSAttention->>DSAttention: compute_dsa_indexer_loss(index_scores, topk_indices, query, key, ...)
    Note over DSAttention: KL divergence between indexer & true attention
    
    DSAttention->>DSAttention: DSAIndexerLossAutoScaler.apply(output, indexer_loss)
    DSAttention-->>MLAAttention: output (with loss attached)
    MLAAttention-->>User: attention_output

Loading

4.3 Data Flow

graph TD
    A["Input: hidden_states<br/>[sq, b, hidden]"] --> B["MLA Q Compression<br/>linear_q_proj→linear_q_down_proj"]
    A --> C["MLA KV Compression<br/>linear_kv_down_proj"]
    
    B --> D["q_compressed<br/>[sq, b, q_lora_rank]"]
    C --> E["kv_compressed<br/>[sq, b, kv_lora_rank]"]
    
    D --> F["MLA Q Upsampling<br/>linear_q_up_proj + RoPE"]
    E --> G["MLA KV Upsampling<br/>linear_kv_up_proj + RoPE"]
    
    F --> H["query<br/>[sq, b, np, hn]"]
    G --> I["key<br/>[sk, b, np, hn]"]
    G --> J["value<br/>[sk, b, np, hnv]"]
    
    A --> K["x.detach()"]
    D --> L["q_compressed.detach()"]
    
    K --> M["DSAIndexer"]
    L --> M
    M --> N["Indexer Q Proj<br/>linear_wq_b<br/>[sq, b, index_n_heads, index_head_dim]"]
    K --> O["Indexer K Proj<br/>linear_wk + k_norm<br/>[sk, b, index_head_dim]"]
    K --> P["Indexer Weights<br/>linear_weights_proj<br/>[sq, b, index_n_heads]"]
    
    N --> Q["Apply RoPE"]
    O --> R["Apply RoPE"]
    Q --> S["rotate_activation<br/>&#40;Hadamard transform&#41;"]
    R --> T["rotate_activation<br/>&#40;Hadamard transform&#41;"]
    
    S --> U["Index Scores<br/>q @ k^T → ReLU → weighted sum<br/>[b, sq, sk]"]
    T --> U
    P --> U
    
    U --> V["TopK Selection<br/>[b, sq, index_topk]"]
    
    H --> W["Sparse Attention"]
    I --> W
    J --> W
    V --> W
    
    W --> X["attention_output<br/>[sq, b, hidden]"]
    
    U --> Y["KL Divergence Loss<br/>KL&#40;true_attn || index_scores&#41;"]
    V --> Y
    H --> Y
    I --> Y
    
    Y --> Z["indexer_loss<br/>scalar"]
    
    X --> AA["DSAIndexerLossAutoScaler.apply"]
    Z --> AA
    AA --> AB["Final Output<br/>&#40;with loss attached&#41;"]
    
    style K fill:#FFE4B5
    style L fill:#FFE4B5
    style M fill:#90EE90
    style U fill:#87CEEB
    style V fill:#87CEEB
    style W fill:#FFD700
    style Y fill:#FF6347
    style Z fill:#FF6347
    style AA fill:#DDA0DD

Loading

---

5. Module Relationships

classDiagram
    class TransformerConfig {
        +int num_layers
        +int hidden_size
        +str experimental_attention_variant
    }
    
    class MLATransformerConfig {
        +int q_lora_rank
        +int kv_lora_rank
        +int dsa_indexer_n_heads
        +int dsa_indexer_head_dim
        +int dsa_indexer_topk
        +float dsa_indexer_loss_coeff
    }
    
    class Attention {
        <<abstract>>
        +forward()*
    }
    
    class MultiLatentAttention {
        +get_query_key_value_tensors()
        +forward()
    }
    
    class MLASelfAttention {
        +linear_q_proj
        +linear_kv_down_proj
        +core_attention
        +get_query_key_value_tensors(return_compressed_tensors)
    }
    
    class DSAttention {
        +indexer: DSAIndexer
        +softmax_scale: float
        +forward(query, key, value, x, qr, ...)
    }
    
    class DSAIndexer {
        +linear_wq_b
        +linear_wk
        +k_norm
        +linear_weights_proj
        +rotary_pos_emb
        +forward(x, qr, mask)
        +forward_with_scores(x, qr, mask)
        -_apply_rope()
        -_compute_index_scores()
    }
    
    class DSAIndexerSubmodules {
        +linear_wq_b: ModuleSpec
        +linear_wk: ModuleSpec
        +k_norm: ModuleSpec
        +linear_weights_proj: ModuleSpec
    }
    
    class DSAttentionSubmodules {
        +indexer: ModuleSpec
    }
    
    class MLASelfAttentionSubmodules {
        +core_attention: ModuleSpec
        +linear_q_proj
        +linear_kv_down_proj
        +q_layernorm
        +kv_layernorm
    }
    
    class RotaryEmbedding {
        +forward(seq_len)
    }
    
    class DSAIndexerLossAutoScaler {
        <<autograd.Function>>
        +forward(output, loss)$
        +backward(grad_output)$
        +set_loss_scale(scale)$
        +main_loss_backward_scale$
    }
    
    class DSAIndexerLossLoggingHelper {
        +save_loss_to_tracker()$
        +reduce_loss_in_tracker()$
        +track_indexer_metrics()$
        +tracker: dict$
    }
    
    TransformerConfig <|-- MLATransformerConfig : extends
    Attention <|-- MultiLatentAttention : extends
    MultiLatentAttention <|-- MLASelfAttention : extends
    Attention <|-- DSAttention : implements (core_attention)
    
    MLASelfAttention --> DSAttention : uses as core_attention
    MLASelfAttention --> MLASelfAttentionSubmodules : configured by
    DSAttention --> DSAIndexer : contains
    DSAttention --> DSAttentionSubmodules : configured by
    DSAIndexer --> DSAIndexerSubmodules : configured by
    DSAIndexer --> RotaryEmbedding : uses
    DSAttention ..> DSAIndexerLossAutoScaler : uses
    DSAttention ..> DSAIndexerLossLoggingHelper : logs to
    
    MLASelfAttention ..> MLATransformerConfig : reads config
    DSAttention ..> MLATransformerConfig : reads config
    DSAIndexer ..> MLATransformerConfig : reads config

Loading

Key Relationships:

1. Composition:

   • MLASelfAttention contains DSAttention as its core_attention module
   • DSAttention contains DSAIndexer for computing sparse indices

2. Utility Classes:

   • DSAIndexerLossAutoScaler: Custom autograd for loss attachment
   • DSAIndexerLossLoggingHelper: Singleton for collecting losses across layers

New Dependencies Introduced:

• fast_hadamard_transform (optional): For Hadamard rotation activation
• Fallback: Mock implementation in tests
• Production: Uses optimized CUDA kernel

---

6. Examples

6.1 Configuration Parameters

CLI Arguments Example (added in arguments.py):

--experimental-attention-variant dsa          # Enable DSA (DeepSeek Sparse Attention)
--dsa-indexer-n-heads 8                       # Number of indexer heads (default: num-attention-heads)
--dsa-indexer-head-dim 64                     # Dimension per indexer head (default: kv-channels)
--dsa-indexer-topk 32                         # Top-k tokens to select per query
--dsa-indexer-loss-coeff 1.0                # Coefficient for KL divergence loss (0 = disabled)
--dsa-indexer-use-sparse-loss                 # Use sparse KL loss (only on top-k positions)

TransformerConfig Example:

config = MLATransformerConfig(
    # ... standard MLA params ...
    experimental_attention_variant='dsa',      # 'dsa' | 'gated_delta_net' | None
    dsa_indexer_n_heads=8,                    # Must divide by TP size
    dsa_indexer_head_dim=64,                  # Typically same as kv_channels
    dsa_indexer_topk=32,                      # k in O(n·k) complexity
    dsa_indexer_loss_coeff=1.0,             # Typical range: 0.0001 - 0.01
    dsa_indexer_use_sparse_loss=False,        # True = sparse, False = dense KL loss
)

6.2 Example Usage

Training a GPT model with DSA:

python pretrain_gpt.py \
    --num-layers 32 \
    --hidden-size 4096 \
    --num-attention-heads 32 \
    --seq-length 8192 \
    \
    # Enable Multi-Latent Attention (required for DSA)
    --multi-latent-attention \
    --q-lora-rank 512 \
    --kv-lora-rank 512 \
    --qk-head-dim 128 \
    --qk-pos-emb-head-dim 64 \
    --v-head-dim 128 \
    \
    # Enable DeepSeek Sparse Attention
    --experimental-attention-variant dsa \
    --dsa-indexer-n-heads 16 \
    --dsa-indexer-head-dim 128 \
    --dsa-indexer-topk 256 \
    --dsa-indexer-loss-coeff 0.001 \
    \
    # Standard training args
    --micro-batch-size 1 \
    --global-batch-size 512 \
    --lr 1.0e-4 \
    --train-iters 100000 \
    --lr-decay-iters 100000 \
    --lr-decay-style cosine \
    --min-lr 1.0e-5 \
    --weight-decay 0.1 \
    --clip-grad 1.0 \
    --bf16

Expected Behavior:

• Each layer will use sparse attention with top-256 tokens (instead of full 8192)
• Indexer loss will be logged to TensorBoard as indexer loss

---

Further Reading

• DeepSeek V3.2 Technical Report: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/DeepSeek_V3_2.pdf

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[iansheng]

Great work—this has been really helpful to me. Quick question: what are the typical values used for indexer-loss-coeff?

[kunlunl]

Great work—this has been really helpful to me. Quick question: what are the typical values used for indexer-loss-coeff?

@iansheng Thanks for your interest. Currently, I don't have relevant experience about indexer loss coefficient, still working on refining this MR. and haven't actually used it to train a model yet.

[yuzhongw-nvidia]

Thanks for your quick and great work. I've preliminarily view the code and it overall LGTM. Please add UTs for the new features, especially for the functionality of DSA and its parallel correctness. I'll do a second-round review after everything is ready.

[yuzhongw-nvidia]

Please also do not forget to create a mirror PR to the main branch once ready for review. Thanks!

[kunlunl]

/ok to test 6fe2f35

[kunlunl]

/ok to test 33684d2

[kunlunl]

/ok to test a0b6fd9

[kunlunl]

/ok to test 1d2f01c

[kunlunl]

/ok to test 23ba310

[kunlunl]

/ok to test cbfa053

[ninangezaici]

Please adapt the MTP function of DSA as soon as possible

[ninangezaici]

Would it be possible to also make this code compatible with the first-stage training of DSA？

[kunlunl]

Would it be possible to also make this code compatible with the first-stage training of DSA？

Under the current mcore architecture, it’s difficult to have an independent computational graph (the indexer loss) perform a backward pass in isolation. It would require a major architectural overhaul, especially since this graph isn't attached to the end of the model.

As a workaround, you could try multiplying the main model loss by zero. Although the main model will still undergo the backward pass, its gradients will all be zero, effectively training only the indexer. Given that the training data for the first stage is relatively small, the redundant computation is a manageable trade-off.