[Feature-request] Context Parallel support for DSAttention

[slikhite-1]

Megatron-Core currently does not support Context Parallelism (CP) with DSAttention. This blocks long-context training and post-training for models that use the experimental DSA path, such as GLM-style sparse attention models. This also affects Megatron Bridge/ Nemo-RL usage, since the Megatron Bridge GLM recipe points to experimental_attention_variant == "dsa".

Current Behavior

When experimental_attention_variant == "dsa", transformer_config.py asserts that context_parallel_size == 1:

assert self.context_parallel_size == 1, "Currently context parallelism is not supported by DSAttention!"

Possible Solution / Reference

One possible reference is the Slime RL PR that added GLM 5.1 support:
THUDM/slime#1599

From a preliminary read, Slime does not use mcore DSAttention directly for that recipe. Instead, it builds mcore GPT/MLA layers but replaces the attention path with a custom GLM5 DSA implementation. They also add a Slime-specific --allgather-cp flag to support Context Parallelism

This is only a possible design reference, not necessarily a request to implement the Slime approach directly.

[guihong-nv]

I checked THUDM/slime#1599. Slime’s GLM5 support is useful as a reference, but it is not a drop-in kernel for MCore DSAttention. For Megatron-LM, we need a native DSA CP implementation: gather/globalize KV and indexer keys for CP, make DSA work with THD packed seq params and global causal offsets, then optionally port/adapt the TileLang sparse MLA/indexer kernels with generic dimensions and tests. I would suggest you contacting the kernel development team for this matter.

[Prachi-kushwaha]

Hi, I'd like to take a stab at implementing this. Here issmy proposed approach before I star happy to adjust based on feedback.

Strategy: allgather CP mode

Each CP rank holds local_seq = global_seq / cp_size tokens. Before the indexer runs, KV/x/qr are allgathered to the full global sequence. The indexer then produces topk_indices in global-sequence coordinates. The sparse attention kernel runs with local Q against global KV, with the causal mask diagonal shifted by cp_rank * local_seq.

Steps I plan to follow:

1. transformer_config.py :-- relax the context_parallel_size == 1 assert to only block unsupported comm types; allow allgather
2. dsa.py / unfused_dsa_fn:-- add a causal_offset` param to shift the causal mask diagonal for the correct CP rank
3. dsa.py / compute_dsa_indexer_loss :-- same causal_offset fix + add ReduceOp.AVG reduce over the CP group after the existing TP reduce, so the KL loss doesn't scale with cp_size
4. dsa.py / DSAttention.forward :-- allgather KV, x, qr before the indexer; derive causal_offset = cp_rank * local_seq; pass it through to unfused_dsa_fn and the loss
5. dsa.py :--- wrap in a torch.autograd.Function so the backward does a reduce_scatter on dK/dV back to local-sequence shards
6. New file dsa_sparse_attn_kernel.py :--- Triton fused sparse forward kernel that avoids materializing the full [b, np, sq, skv] score matrix by iterating only over topk indices; follows the same pattern as megatron/core/fusions/fused_mla_yarn_rope_apply.py
7. Tests :-- extend test_attention_variant_dsa.py with CP correctness cases; single-GPU test compares causal_offset=sq_local output against the corresponding slice of a full single-rank run

On TileLang vs Triton ------ The feature request references Slime/DeepSeek which use TileLang for the DSA kernel. Megatron-Core currently has zero TileLang usage and all existing custom kernels are Triton. I plan to use Triton to stay consistent with the repo conventions, but happy to discuss

Let me know if this approach looks reasonable and I'll open a draft PR.
@slikhite-1 @guihong-nv

[guihong-nv]

@slikhite-1 @Prachi-kushwaha We are actively working on this topic: #3674 and #5099. Thanks for your suggestions!