Add fused dsa

[kunlunl]

What does this PR do ?

Separate PR for adding absorbed-mla #3193 (merged)

main PR: #3747

1. TL;DR

• What: Adds AbsorbedMLASelfAttention, a variant of MLA that absorbs K's up projection into Q and applies V's up projection after core attention, enabling MQA-style computation.
• Why: MQA-style attention (single KV head) can be more efficient for certain sparse attention variants like DSA (Dense Sparse Attention).
• Impact: Enables DSA + MLA combination; adds fused Tilelang kernels for DSA-MQA; refactors base Attention class to move module init to subclasses.

---

2. Big Picture

2.1 Before vs After Architecture

graph TB
    subgraph "Standard MLA (Before)"
        H1[hidden_states] --> Q1[Q projection]
        H1 --> KV1[KV projection]
        KV1 --> KVup1["KV up proj<br/>[K,V] = W_up @ kv_compressed"]
        Q1 --> Qup1[Q up proj]
        Qup1 --> CA1[Core Attention<br/>MHA: n heads for Q,K,V]
        KVup1 --> CA1
        CA1 --> OUT1[Output proj]
    end
    
    subgraph "Absorbed MLA (After)"
        H2[hidden_states] --> Q2[Q projection]
        H2 --> KV2[KV projection]
        Q2 --> Qup2[Q up proj]
        Qup2 --> ABS["Absorb K_up into Q<br/>Q' = Q @ K_up^T"]
        ABS --> CA2[Core Attention<br/>MQA: n heads Q, 1 head KV]
        KV2 --> CA2
        CA2 --> Vup["Apply V_up after attn<br/>out = attn_out @ V_up"]
        Vup --> OUT2[Output proj]
    end

Loading

Key insight: By absorbing K's up-projection into Q, the attention operates in MQA form (Q has n heads, K/V have 1 head). This is mathematically equivalent but enables more efficient sparse attention patterns.

2.2 Change Scope Summary

Category	Files	Description
New	experimental_attention_variant/absorbed_mla.py	Core AbsorbedMLA implementation (~980 lines)
New	experimental_attention_variant/dsa_fused_kernels.py	Tilelang fused kernels for DSA-MQA (~550 lines)
New	tests/.../test_absorbed_mla.py	Unit tests (~390 lines)
Modified	attention.py	Move core_attention/linear_proj init from base to subclasses
Modified	dsa.py	Add MQA mode support for DSA
Modified	multi_latent_attention.py	Remove unused position_ids parameter
Modified	experimental_attention_variant_module_specs.py	Use AbsorbedMLA for DSA spec

---

3. Design Rationale

3.1 Problem Background

Standard MLA applies KV up-projection before attention:

K = W_k_up @ kv_compressed  # [seq, batch, n_heads, qk_head_dim]
V = W_v_up @ kv_compressed  # [seq, batch, n_heads, v_head_dim]

This produces multi-head K/V tensors, which don't work well with DSA's sparse attention pattern that benefits from MQA-style computation.

3.2 Solution: Matrix Absorption

The mathematical trick:

# Standard: score = softmax(Q @ K^T)
# Where K = W_k_up @ kv_compressed

# Absorbed: Q' = Q @ W_k_up^T, then score = softmax(Q' @ kv_compressed^T)
# Result: K/V remain single-head (MQA form)

V up-projection is moved after attention:

# Standard: out = score @ V, where V = W_v_up @ kv_compressed  
# Absorbed: out' = score @ kv_compressed, then out = out' @ W_v_up

3.3 Key Design Points

1. Separate K/V up projections: linear_kv_up_proj split into linear_k_up_proj + linear_v_up_proj
2. Core attention takes MQA input: Q shape [s,b,n,d], K/V shape [s,b,1,d]
3. Checkpoint compatibility: _load_from_state_dict handles splitting combined KV weights

---

4. Execution Path Deep Dive

4.1 Call Chain

sequenceDiagram
    participant F as forward()
    participant GQKV as get_query_key_value_tensors()
    participant UP as qkv_up_proj_and_rope_apply()
    participant CA as core_attention()
    participant VP as V up proj
    
    F->>GQKV: hidden_states
    GQKV->>GQKV: Q down proj → q_compressed
    GQKV->>GQKV: KV down proj → kv_compressed, k_pos_emb
    GQKV->>GQKV: Apply layernorms
    GQKV->>UP: q_compressed, kv_compressed
    UP->>UP: Q up proj → q [s,b,n,qk+rope_dim]
    UP->>UP: Absorb K_up into Q: q' = einsum("...nd,ndk->...nk", q_nope, K_up_weight)
    UP->>UP: Apply RoPE to q_rope and k_pos_emb
    UP->>UP: Concat: q_absorbed = [q', q_rope], kv = [kv_compressed, k_rope]
    UP-->>GQKV: q_absorbed, kv_compressed
    GQKV-->>F: q_absorbed, kv_compressed
    F->>CA: q_absorbed [s,b,n,kv_rank+rope], kv [s,b,1,kv_rank+rope]
    CA-->>F: attn_out [s,b,n,kv_rank]
    F->>VP: einsum("...nc,ndc->...nd", attn_out, V_up_weight)
    VP-->>F: out [s,b,n,v_dim]
    F->>F: linear_proj → output

Loading

4.2 Core Code: K Absorption

# absorbed_mla.py:764-789
def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb):
    # Q up projection: [num_tokens, q_lora_rank] -> [num_tokens, n, qk_head_dim+rope_dim]
    q, _ = self.linear_q_up_proj(q_compressed)
    q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim)
    
    # Prepare K up weight: [n * qk_head_dim, kv_lora_rank] -> [n, qk_head_dim, kv_lora_rank]
    k_up_weight = self.linear_k_up_proj.weight.view(
        self.num_attention_heads_per_partition,
        self.config.qk_head_dim,
        self.config.kv_lora_rank,
    )
    
    # Split Q into no-position-embedding and position-embedding parts
    q_no_pe, q_pos_emb = torch.split(q, [qk_head_dim, rope_dim], dim=-1)
    
    # KEY: Absorb K_up into Q
    # q_no_pe: [num_tokens, n, qk_head_dim]
    # k_up_weight: [n, qk_head_dim, kv_lora_rank]
    # q_absorbed: [num_tokens, n, kv_lora_rank]
    q_absorbed = torch.einsum("...nd,ndk->...nk", q_no_pe, k_up_weight)
    
    # Apply RoPE and concatenate
    q_absorbed = torch.cat([q_absorbed, q_pos_emb_with_rope], dim=-1)  # [tokens, n, kv_rank+rope]
    kv_compressed = torch.cat([kv_compressed, k_pos_emb_with_rope], dim=-1)  # [tokens, 1, kv_rank+rope]

4.3 Core Code: V Up-Projection After Attention

# absorbed_mla.py:1033-1047
# After core_attention returns: attn_out shape [tokens, n * kv_lora_rank]

# Reshape V up weight: [n * v_head_dim, kv_lora_rank] -> [n, v_head_dim, kv_lora_rank]
v_up_weight = self.linear_v_up_proj.weight.view(
    self.num_attention_heads_per_partition, self.config.v_head_dim, self.config.kv_lora_rank
)

# Reshape attn_out: [tokens, n, kv_lora_rank]
core_attn_out = core_attn_out.view(*core_attn_out.shape[:-1], n, kv_lora_rank)

# Apply V up projection: [tokens, n, kv_lora_rank] @ [n, v_head_dim, kv_lora_rank]^T -> [tokens, n, v_head_dim]
core_attn_out = torch.einsum("...nc,ndc->...nd", core_attn_out, v_up_weight)

---

5. Module Relationships

classDiagram
    class Attention {
        +config
        +layer_number
        +attn_mask_type
        -checkpoint_core_attention
        -offload_core_attention
    }
    
    class MultiLatentAttention {
        +rotary_pos_emb
        +core_attention
        +linear_proj
        +softmax_scale
    }
    
    class MLASelfAttention {
        +linear_q_down_proj
        +linear_q_up_proj
        +linear_kv_down_proj
        +linear_kv_up_proj
        +get_query_key_value_tensors()
    }
    
    class AbsorbedMLASelfAttention {
        +linear_k_up_proj [NEW]
        +linear_v_up_proj [NEW]
        +get_query_key_value_tensors()
        +_split_kv_weights()
        +_combine_kv_weights()
    }
    
    class DSAttention {
        +k_channels
        +v_channels
        +forward() supports MQA
    }
    
    Attention <|-- MultiLatentAttention
    MultiLatentAttention <|-- MLASelfAttention
    Attention <|-- AbsorbedMLASelfAttention : extends directly
    AbsorbedMLASelfAttention ..> DSAttention : uses for MQA mode

Loading

Note: AbsorbedMLASelfAttention extends Attention directly (not MultiLatentAttention) because it builds its own core_attention and linear_proj with MQA-specific parameters.

---

6. Risks & Edge Cases

Risk Category	Specific Concern	Status
Inference	cache_mla_latents not supported	Asserted, raises error
FP8/FP4	recompute_up_proj incompatible	Asserted in code
Checkpoint	Loading from standard MLA checkpoint	Handled via _load_from_state_dict
Fused kernels	Tilelang dependency for DSA-MQA	Optional import, falls back to unfused
Correctness	Unit tests verify cosine similarity > 0.9999	Tested for TP/CP combinations

---

7. Quick Reference

7.1 File Change Summary

experimental_attention_variant/absorbed_mla.py     [+979]  AbsorbedMLA implementation
experimental_attention_variant/dsa_fused_kernels.py [+554] Tilelang DSA-MQA kernels
experimental_attention_variant/dsa.py              [+80]   MQA mode for DSA
attention.py                                       [+100, -60] Move init to subclasses
multi_latent_attention.py                          [-4]    Remove position_ids
experimental_attention_variant_module_specs.py     [+8, -6] Use AbsorbedMLA for DSA
tests/.../test_absorbed_mla.py                     [+389]  Unit tests

7.2 Key Functions

Function	Location	Purpose
AbsorbedMLASelfAttention.__init__	absorbed_mla.py:350	Build separate K/V up projections
get_query_key_value_tensors	absorbed_mla.py:596	Q/KV down proj + absorption
qkv_up_proj_and_rope_apply	absorbed_mla.py:740	K absorption + RoPE
forward	absorbed_mla.py:953	Core attn + V up proj
_split_kv_weights	absorbed_mla.py:1157	Load combined KV weights from checkpoint
unfused_dsa_fn_mqa	dsa.py:699	MQA sparse attention
FusedDSAMQA	dsa_fused_kernels.py:1915	Tilelang fused DSA-MQA

7.3 Related Code

• megatron/core/transformer/multi_latent_attention.py - Standard MLA for comparison
• megatron/core/transformer/attention.py - Base attention class
• DeepSeek-V3 paper - MLA architecture reference

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

[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.

[jaredcasper]

What is the rational for this change? The base Attention's forward class depends on there being a self.core_attention. This also leads to code duplication. Perhaps something that doesn't have a core attention block followed by linear_proj should just be a different thing all together?

[kunlunl]

I just noticed that the MLA and base attention are both creating core attention and out_proj, so I made a quick fix. It has nothing to do with this PR, I will revert this change.

[yanring]

@kunlunl Did you push the commits?

[kunlunl]

Not in this PR yet. I revert this change in the standalone absorbed_mla PR. We will merge that PR first right? Once that PR merged, I'll rebase this one.

[kunlunl]

#3193

[kunlunl]

Rebased. To avoid any potential concerns.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[kunlunl]

/ok to test ac68a51

[kunlunl]

/ok to test 14e9792

[kunlunl]

/ok to test fdaa74c

[kunlunl]

/ok to test 73ea27c

[kunlunl]

/ok to test bf76770

[kunlunl]

/ok to test 1d37420

[kunlunl]

/ok to test 46c2e3c

[kunlunl]

/ok to test 3ead04d