From 2262927b508943a43effb59973901678b9952213 Mon Sep 17 00:00:00 2001
From: xielaixin <xielx@shanghaitech.edu.cn>
Date: Tue, 20 Jan 2026 20:32:13 +0800
Subject: [PATCH 1/5] refactor: split indexer topk and loss for further fused
 kernel

---
 .../experimental_attention_variant/dsa.py     | 401 +++++++++++++-----
 1 file changed, 303 insertions(+), 98 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py
index 88b4713dc60..67965d09c43 100644
--- a/megatron/core/transformer/experimental_attention_variant/dsa.py
+++ b/megatron/core/transformer/experimental_attention_variant/dsa.py
@@ -252,6 +252,238 @@ def compute_dsa_indexer_loss(
     return indexer_loss
 
 
+def _compute_index_scores(
+    q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor
+) -> torch.Tensor:
+    """
+    Perform index score using BF16 precision.
+
+    Reference:
+        https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/kernel.py#L254-L274
+    This is a BF16 implementation of the `fp8_index` logic:
+        1. Compute attention scores: q @ k^T;
+        2. Apply ReLU activation;
+        3. Weight by attention weights;
+        4. Sum across attention heads.
+
+    Args:
+        q: BF16 [seqlen_q, batch, index_n_heads, index_head_dim], the query tensor.
+        weights: BF16 [seqlen_q, batch, index_n_heads], the attention weights.
+        k: BF16 [seqlen_k, batch, index_head_dim], the key tensor.
+
+    Returns:
+        index_scores: FP32 [batch, seqlen_q, seqlen_k], the index scores.
+    """
+    # Compute attention scores: q @ k^T
+    # [seqlen_q, batch, index_n_heads, index_head_dim] @ [seqlen_k, batch, index_head_dim]^T
+    #   -> [seqlen_q, batch, index_n_heads, seqlen_k]
+    index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float())
+
+    # Apply ReLU activation.
+    index_scores = torch.relu(index_scores)
+
+    # Weight each head by attention weights.
+    # [seqlen_q, batch, index_n_heads, seqlen_k] * [seqlen_q, batch, index_n_heads, 1]
+    #   -> [seqlen_q, batch, index_n_heads, seqlen_k]
+    index_scores = index_scores * weights.unsqueeze(-1)
+
+    # Sum across attention heads.
+    # [seqlen_q, batch, index_n_heads, seqlen_k] -> [seqlen_q, batch, seqlen_k]
+    index_scores = index_scores.sum(dim=2)
+
+    # Transpose to [batch, seqlen_q, seqlen_k].
+    index_scores = index_scores.transpose(0, 1)
+
+    return index_scores
+
+
+def fused_qk_topk_native(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    weights: torch.Tensor,
+    index_topk: int,
+    mask: Optional[torch.Tensor] = None,
+):
+    seqlen = q.size(0)
+    # =========================================
+    # Compute index scores
+    # =========================================
+    # [batch, seqlen, seqlen]
+    index_scores = _compute_index_scores(q, weights, k)
+    if mask is not None:
+        assert mask.dtype == index_scores.dtype, "Mask dtype must match index scores dtype"
+        index_scores = index_scores + mask
+
+    # =========================================
+    # Select top-k indices
+    # =========================================
+    topk_k = min(index_topk, seqlen)
+    # [batch, seqlen, index_topk]
+    topk_indices = index_scores.topk(topk_k, dim=-1)[1]
+
+    return index_scores, topk_indices
+
+
+def fwd_fused_indexer_loss_native(q, weights, k, query, key, topk, softmax_scale, loss_coeff, mask, sparse_loss, pg_collection):
+    index_scores, topk_indices = fused_qk_topk_native(q, k, weights, topk, mask)
+
+    indexer_loss = compute_dsa_indexer_loss(
+        index_scores, topk_indices, query, key, softmax_scale, loss_coeff, sparse_loss, pg_collection
+    )
+
+    return topk_indices, indexer_loss
+
+
+def bwd_fused_indexer_loss_native(q, weights, k, query, key, topk_indices, softmax_scale, loss_coeff, sparse_loss, grad_loss, pg_collection):
+    index_scores = _compute_index_scores(q, weights, k)  # [B, Sq, Sk]
+
+    sq, b, np, hn = query.size()
+    sk = key.size(0)
+
+    # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn]
+    query_reshaped = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn)
+    # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk]
+    key_reshaped = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk)
+    # Compute attention scores [b * np, sq, sk]
+    attention_scores = torch.bmm(query_reshaped.float(), key_reshaped.float()) * softmax_scale
+    # Reshape to [b, np, sq, sk]
+    attention_scores = attention_scores.reshape(b, np, sq, sk)
+
+    # causal_mask [sq, sk]
+    causal_mask = torch.triu(
+        torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device),
+        diagonal=1,
+    )
+    # index_mask [b, sq, sk]
+    index_mask = torch.full(
+        (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device
+    ).scatter_(-1, topk_indices, 0)
+
+    # Apply causal mask to both attention and index scores
+    # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv]
+    attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk)
+    # [b, sq, sk] + [1, sq, sk] -> [b, sq, sk]  
+    index_scores = index_scores + causal_mask.unsqueeze(0)
+    
+    if sparse_loss:
+        # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk]
+        attention_scores = attention_scores + index_mask.view(b, 1, sq, sk)
+        # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk]
+        index_scores = index_scores + index_mask
+    
+    # Compute softmax for both
+    attention_scores_softmax = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32)
+    index_scores_softmax = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32)
+    
+    # Sum attention scores across heads: [b, np, sq, sk] -> [b, sq, sk]
+    attention_scores_sum = attention_scores_softmax.sum(dim=1)
+
+    if pg_collection.tp.size() > 1:
+        # attention scores are scattered to TP ranks in head dimension.
+        torch.distributed.all_reduce(attention_scores_sum.contiguous(), group=pg_collection.tp)
+
+    # L1 normalize
+    attention_scores_normalized = attention_scores_sum / attention_scores_sum.sum(dim=-1, keepdim=True)
+    
+    # Backward through loss = kl_div * loss_coeff
+    # where kl_div = kl_per_element.sum(dim=-1).mean()
+    grad_kl_div = grad_loss * loss_coeff  # scalar
+    
+    # Backward through mean: distribute gradient equally
+    grad_kl_per_row = grad_kl_div / (b * sq)  # scalar value for each row
+    
+    # Backward through sum(dim=-1): broadcast back to [b, sq, sk]
+    # Each element in a row contributes to the sum, so gradient is same for all
+    grad_kl_per_element = torch.full((b, sq, sk), grad_kl_per_row.item(), 
+                                      device=index_scores.device, dtype=torch.float32)
+    
+    # Backward through kl_per_element = target * (log(target) - log(index))
+    # ∂kl/∂index_softmax = -target / index_softmax
+    grad_index_scores_softmax = -attention_scores_normalized / (index_scores_softmax + 1e-10) * grad_kl_per_element
+    
+    # Backward through softmax: ∂L/∂x = softmax * (∂L/∂softmax - sum(∂L/∂softmax * softmax))
+    sum_grad = (grad_index_scores_softmax * index_scores_softmax).sum(dim=-1, keepdim=True)
+    grad_index_scores_logits = index_scores_softmax * (grad_index_scores_softmax - sum_grad)
+    
+    # Zero out gradients for masked positions
+    # Create a mask for valid (non-masked) positions
+    # Causal mask: position (i, j) is valid if j <= i
+    causal_valid_mask = torch.tril(torch.ones((sq, sk), device=index_scores.device, dtype=torch.bool))  # [sq, sk]
+    if sparse_loss:
+        # Also apply index mask - only topk positions are valid
+        index_valid_mask = (index_mask == 0)  # [b, sq, sk]
+        valid_mask = causal_valid_mask.unsqueeze(0) & index_valid_mask  # [b, sq, sk]
+    else:
+        valid_mask = causal_valid_mask.unsqueeze(0).expand(b, sq, sk)  # [b, sq, sk]
+    
+    grad_index_scores_logits = grad_index_scores_logits * valid_mask.float()
+    
+    # Transpose from [b, sq, sk] to [sq, b, sk]
+    grad_index_scores = grad_index_scores_logits.transpose(0, 1)  # [sq, b, sk]
+    
+    # Backward through sum over heads: expand gradient
+    grad_weighted_scores = grad_index_scores.unsqueeze(2)  # [sq, b, 1, sk]
+    
+    # Compute forward values needed for backward
+    scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float())  # [sq, b, h, sk]
+    scores_after_relu = torch.relu(scores)
+    
+    # Backward through multiplication by weights: index_scores_per_head * weights
+    # ∂L/∂weights = grad * relu_scores (sum over sk)
+    grad_weights = (grad_weighted_scores * scores_after_relu).sum(dim=-1)  # [sq, b, h]
+    
+    # ∂L/∂relu_scores = grad * weights
+    grad_scores_after_relu = grad_weighted_scores * weights.unsqueeze(-1)  # [sq, b, h, sk]
+    
+    # Backward through ReLU
+    relu_mask = (scores > 0).float()
+    grad_scores = grad_scores_after_relu * relu_mask  # [sq, b, h, sk]
+    
+    # Backward through einsum 'sbhd,tbd->sbht'
+    # ∂L/∂q = einsum('sbht,tbd->sbhd', grad_scores, k)
+    grad_q = torch.einsum('sbht,tbd->sbhd', grad_scores, k.float())  # [sq, b, h, d]
+    # ∂L/∂k = einsum('sbht,sbhd->tbd', grad_scores, q)
+    grad_k = torch.einsum('sbht,sbhd->tbd', grad_scores, q.float())  # [sk, b, d]
+
+    return grad_q.to(q.dtype), grad_weights.to(weights.dtype), grad_k.to(k.dtype)
+
+
+class FusedDSAIndexerLoss(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, weights, k, query, key, softmax_scale, topk, loss_coeff, mask, sparse_loss, pg_collection):
+        """
+        Fused forward: index_scores never materialized in full.
+        """
+
+        topk_indices, loss = fwd_fused_indexer_loss_native(
+            q, weights, k, query, key, topk, softmax_scale, loss_coeff, mask, sparse_loss, pg_collection,
+        )
+
+        # Save for backward (recomputation strategy)
+        ctx.save_for_backward(q, weights, k, query, key, topk_indices)
+        ctx.softmax_scale = softmax_scale
+        ctx.loss_coeff = loss_coeff
+        ctx.sparse_loss = sparse_loss
+        ctx.pg_collection = pg_collection
+
+        return topk_indices, loss
+
+    @staticmethod
+    def backward(ctx, grad_topk_indices, grad_loss):
+        """
+        Backward: Recompute what we need.
+        """
+        q, weights, k, query, key, topk_indices = ctx.saved_tensors
+
+        grad_q , grad_weights, grad_k = bwd_fused_indexer_loss_native(
+            q, weights, k, query, key, topk_indices, 
+            ctx.softmax_scale, ctx.loss_coeff, ctx.sparse_loss, grad_loss, ctx.pg_collection,
+        )
+
+        # query and key are detached in forward, so return None for their gradients
+        return grad_q, grad_weights, grad_k, None, None, None, None, None, None, None, None
+
+
 class DSAIndexerLossAutoScaler(torch.autograd.Function):
     """An AutoScaler that triggers the backward pass and scales the grad for indexer loss.
 
@@ -471,74 +703,12 @@ def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, mscale: flo
         x = torch.cat([x_nope, x_pe], dim=-1)
         return x
 
-    def _compute_index_scores(
-        self, q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor
-    ) -> torch.Tensor:
-        """
-        Perform index score using BF16 precision.
-
-        Reference:
-            https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/kernel.py#L254-L274
-        This is a BF16 implementation of the `fp8_index` logic:
-            1. Compute attention scores: q @ k^T;
-            2. Apply ReLU activation;
-            3. Weight by attention weights;
-            4. Sum across attention heads.
-
-        Args:
-            q: BF16 [seqlen_q, batch, index_n_heads, index_head_dim], the query tensor.
-            weights: BF16 [seqlen_q, batch, index_n_heads], the attention weights.
-            k: BF16 [seqlen_k, batch, index_head_dim], the key tensor.
-
-        Returns:
-            index_scores: FP32 [batch, seqlen_q, seqlen_k], the index scores.
-        """
-        # Compute attention scores: q @ k^T
-        # [seqlen_q, batch, index_n_heads, index_head_dim] @ [seqlen_k, batch, index_head_dim]^T
-        #   -> [seqlen_q, batch, index_n_heads, seqlen_k]
-        index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float())
-
-        # Apply ReLU activation.
-        index_scores = torch.relu(index_scores)
-
-        # Weight each head by attention weights.
-        # [seqlen_q, batch, index_n_heads, seqlen_k] * [seqlen_q, batch, index_n_heads, 1]
-        #   -> [seqlen_q, batch, index_n_heads, seqlen_k]
-        index_scores = index_scores * weights.unsqueeze(-1)
-
-        # Sum across attention heads.
-        # [seqlen_q, batch, index_n_heads, seqlen_k] -> [seqlen_q, batch, seqlen_k]
-        index_scores = index_scores.sum(dim=2)
-
-        # Transpose to [batch, seqlen_q, seqlen_k].
-        index_scores = index_scores.transpose(0, 1)
-
-        return index_scores
-
-    def forward_with_scores(
+    def forward_before_topk(
         self,
         x: torch.Tensor,
         qr: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for DSA Indexer that returns both index scores and top-k indices.
-
-        This is used when KL loss is enabled to compare indexer scores with true attention scores.
-
-        Args:
-            x: hidden states [seqlen, batch, hidden_size].
-            qr: Low-rank query tensor [seqlen, batch, q_lora_rank].
-            mask: Attention mask [batch, seqlen, seqlen].
-            packed_seq_params: Packed sequence parameters for variable length sequences.
-
-        Returns:
-            index_scores: Index scores [batch, seqlen, seqlen].
-            topk_indices: Top-k indices [batch, seqlen, index_topk].
-        """
-        assert packed_seq_params is None, "Packed sequence is not supported for DSAttention"
-
         # =========================================
         # Prepare RoPE params
         # =========================================
@@ -592,23 +762,43 @@ def forward_with_scores(
         k = rotate_activation(k)
 
         # =========================================
-        # Compute index scores
+        # Prepare weights for index scores
         # =========================================
         # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads]
         weights, _ = self.linear_weights_proj(x)
         weights = weights * (self.index_n_heads**-0.5) * self.softmax_scale
-        # [batch, seqlen, seqlen]
-        index_scores = self._compute_index_scores(q, weights, k)
-        if mask is not None:
-            assert mask.dtype == index_scores.dtype, "Mask dtype must match index scores dtype"
-            index_scores = index_scores + mask
 
-        # =========================================
-        # Select top-k indices
-        # =========================================
-        topk_k = min(self.index_topk, seqlen)
-        # [batch, seqlen, index_topk]
-        topk_indices = index_scores.topk(topk_k, dim=-1)[1]
+        return q, k, weights
+
+    def forward_with_scores(
+        self,
+        x: torch.Tensor,
+        qr: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass for DSA Indexer that returns both index scores and top-k indices.
+
+        This is used when KL loss is enabled to compare indexer scores with true attention scores.
+
+        Args:
+            x: hidden states [seqlen, batch, hidden_size].
+            qr: Low-rank query tensor [seqlen, batch, q_lora_rank].
+            mask: Attention mask [batch, seqlen, seqlen].
+            packed_seq_params: Packed sequence parameters for variable length sequences.
+
+        Returns:
+            index_scores: Index scores [batch, seqlen, seqlen].
+            topk_indices: Top-k indices [batch, seqlen, index_topk].
+        """
+        assert packed_seq_params is None, "Packed sequence is not supported for DSAttention"
+
+        # [seqlen, batch, index_n_heads * index_head_dim], [seqlen, batch, index_head_dim], [seqlen, batch, index_n_heads]
+        q, k, weights = self.forward_before_topk(x, qr, packed_seq_params)
+
+        # [batch, seqlen, seqlen], [batch, seqlen, index_topk]
+        index_scores, topk_indices = fused_qk_topk_native(q, k, weights, self.index_topk, mask)
 
         return index_scores, topk_indices
 
@@ -781,31 +971,27 @@ def forward(
                 mask, float('-inf')
             )
 
-        # ===================================
-        # Get index scores and top-k indices
-        # ===================================
-        index_scores, topk_indices = self.indexer.forward_with_scores(
-            x, qr, mask=float_mask, packed_seq_params=packed_seq_params
-        )
-
-        # ===================================
-        # Run sparse attention kernel
-        # ===================================
-        output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
-
-        # ===================================
-        # Attach indexer loss
-        # ===================================
         if self.training and torch.is_grad_enabled():
-            # Compute KL divergence loss between indexer scores and true attention scores
+            # ===================================
+            # Prepare inputs for indexer loss
+            # ===================================
+            q, k, weights = self.indexer.forward_before_topk(x, qr, packed_seq_params)
             indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0)
-            indexer_loss = compute_dsa_indexer_loss(
-                index_scores,
-                topk_indices,
-                query.detach(),
-                key.detach(),
-                self.softmax_scale,
-                indexer_loss_coeff,
+
+            # ===================================
+            # Attach indexer topk and loss
+            # ===================================
+            # Compute KL divergence loss between indexer scores and true attention scores
+            topk_indices, indexer_loss = FusedDSAIndexerLoss.apply(
+                q, 
+                weights, 
+                k, 
+                query.detach(), 
+                key.detach(), 
+                self.softmax_scale, 
+                self.indexer.index_topk, 
+                indexer_loss_coeff, 
+                float_mask, 
                 getattr(self.config, "dsa_indexer_use_sparse_loss", False),
                 self.indexer.pg_collection,
             )
@@ -816,7 +1002,26 @@ def forward(
                     layer_number=self.layer_number,
                     num_layers=self.config.num_layers,
                 )
+
+            # ===================================
+            # Run sparse attention kernel
+            # ===================================
+            output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
+            
             # Attach loss to output
             output = DSAIndexerLossAutoScaler.apply(output, indexer_loss)
 
+        else:
+            # ===================================
+            # Get index scores and top-k indices
+            # ===================================
+            _, topk_indices = self.indexer.forward_with_scores(
+                x, qr, mask=float_mask, packed_seq_params=packed_seq_params
+            )
+
+            # ===================================
+            # Run sparse attention kernel
+            # ===================================
+            output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
+
         return output

From 4c1cb2ca0671ede792a8b76a9d75ebce0060e227 Mon Sep 17 00:00:00 2001
From: kunlunl <kunlunl@nvidia.com>
Date: Wed, 4 Feb 2026 02:32:36 +0800
Subject: [PATCH 2/5] Fix lint error

---
 .../experimental_attention_variant/dsa.py     | 176 ++++++++++++------
 1 file changed, 120 insertions(+), 56 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py
index 67965d09c43..274059713c3 100644
--- a/megatron/core/transformer/experimental_attention_variant/dsa.py
+++ b/megatron/core/transformer/experimental_attention_variant/dsa.py
@@ -252,9 +252,7 @@ def compute_dsa_indexer_loss(
     return indexer_loss
 
 
-def _compute_index_scores(
-    q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor
-) -> torch.Tensor:
+def _compute_index_scores(q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
     """
     Perform index score using BF16 precision.
 
@@ -297,13 +295,14 @@ def _compute_index_scores(
     return index_scores
 
 
-def fused_qk_topk_native(
+def fused_qk_topk_naive(
     q: torch.Tensor,
     k: torch.Tensor,
     weights: torch.Tensor,
     index_topk: int,
     mask: Optional[torch.Tensor] = None,
 ):
+    """Naive implementation of QK Topk."""
     seqlen = q.size(0)
     # =========================================
     # Compute index scores
@@ -324,17 +323,40 @@ def fused_qk_topk_native(
     return index_scores, topk_indices
 
 
-def fwd_fused_indexer_loss_native(q, weights, k, query, key, topk, softmax_scale, loss_coeff, mask, sparse_loss, pg_collection):
-    index_scores, topk_indices = fused_qk_topk_native(q, k, weights, topk, mask)
+def fwd_fused_indexer_loss_naive(
+    q, weights, k, query, key, topk, softmax_scale, loss_coeff, mask, sparse_loss, pg_collection
+):
+    """Naive implementation of forward pass for indexer loss."""
+    index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, topk, mask)
 
     indexer_loss = compute_dsa_indexer_loss(
-        index_scores, topk_indices, query, key, softmax_scale, loss_coeff, sparse_loss, pg_collection
+        index_scores,
+        topk_indices,
+        query,
+        key,
+        softmax_scale,
+        loss_coeff,
+        sparse_loss,
+        pg_collection,
     )
 
     return topk_indices, indexer_loss
 
 
-def bwd_fused_indexer_loss_native(q, weights, k, query, key, topk_indices, softmax_scale, loss_coeff, sparse_loss, grad_loss, pg_collection):
+def bwd_fused_indexer_loss_naive(
+    q,
+    weights,
+    k,
+    query,
+    key,
+    topk_indices,
+    softmax_scale,
+    loss_coeff,
+    sparse_loss,
+    grad_loss,
+    pg_collection,
+):
+    """Naive implementation of backward pass for indexer loss."""
     index_scores = _compute_index_scores(q, weights, k)  # [B, Sq, Sk]
 
     sq, b, np, hn = query.size()
@@ -362,19 +384,21 @@ def bwd_fused_indexer_loss_native(q, weights, k, query, key, topk_indices, softm
     # Apply causal mask to both attention and index scores
     # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv]
     attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk)
-    # [b, sq, sk] + [1, sq, sk] -> [b, sq, sk]  
+    # [b, sq, sk] + [1, sq, sk] -> [b, sq, sk]
     index_scores = index_scores + causal_mask.unsqueeze(0)
-    
+
     if sparse_loss:
         # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk]
         attention_scores = attention_scores + index_mask.view(b, 1, sq, sk)
         # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk]
         index_scores = index_scores + index_mask
-    
+
     # Compute softmax for both
-    attention_scores_softmax = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32)
+    attention_scores_softmax = torch.nn.functional.softmax(
+        attention_scores, dim=-1, dtype=torch.float32
+    )
     index_scores_softmax = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32)
-    
+
     # Sum attention scores across heads: [b, np, sq, sk] -> [b, sq, sk]
     attention_scores_sum = attention_scores_softmax.sum(dim=1)
 
@@ -383,62 +407,69 @@ def bwd_fused_indexer_loss_native(q, weights, k, query, key, topk_indices, softm
         torch.distributed.all_reduce(attention_scores_sum.contiguous(), group=pg_collection.tp)
 
     # L1 normalize
-    attention_scores_normalized = attention_scores_sum / attention_scores_sum.sum(dim=-1, keepdim=True)
-    
+    attention_scores_normalized = attention_scores_sum / attention_scores_sum.sum(
+        dim=-1, keepdim=True
+    )
+
     # Backward through loss = kl_div * loss_coeff
     # where kl_div = kl_per_element.sum(dim=-1).mean()
     grad_kl_div = grad_loss * loss_coeff  # scalar
-    
+
     # Backward through mean: distribute gradient equally
     grad_kl_per_row = grad_kl_div / (b * sq)  # scalar value for each row
-    
+
     # Backward through sum(dim=-1): broadcast back to [b, sq, sk]
     # Each element in a row contributes to the sum, so gradient is same for all
-    grad_kl_per_element = torch.full((b, sq, sk), grad_kl_per_row.item(), 
-                                      device=index_scores.device, dtype=torch.float32)
-    
+    grad_kl_per_element = torch.full(
+        (b, sq, sk), grad_kl_per_row.item(), device=index_scores.device, dtype=torch.float32
+    )
+
     # Backward through kl_per_element = target * (log(target) - log(index))
     # ∂kl/∂index_softmax = -target / index_softmax
-    grad_index_scores_softmax = -attention_scores_normalized / (index_scores_softmax + 1e-10) * grad_kl_per_element
-    
+    grad_index_scores_softmax = (
+        -attention_scores_normalized / (index_scores_softmax + 1e-10) * grad_kl_per_element
+    )
+
     # Backward through softmax: ∂L/∂x = softmax * (∂L/∂softmax - sum(∂L/∂softmax * softmax))
     sum_grad = (grad_index_scores_softmax * index_scores_softmax).sum(dim=-1, keepdim=True)
     grad_index_scores_logits = index_scores_softmax * (grad_index_scores_softmax - sum_grad)
-    
+
     # Zero out gradients for masked positions
     # Create a mask for valid (non-masked) positions
     # Causal mask: position (i, j) is valid if j <= i
-    causal_valid_mask = torch.tril(torch.ones((sq, sk), device=index_scores.device, dtype=torch.bool))  # [sq, sk]
+    causal_valid_mask = torch.tril(
+        torch.ones((sq, sk), device=index_scores.device, dtype=torch.bool)
+    )  # [sq, sk]
     if sparse_loss:
         # Also apply index mask - only topk positions are valid
-        index_valid_mask = (index_mask == 0)  # [b, sq, sk]
+        index_valid_mask = index_mask == 0  # [b, sq, sk]
         valid_mask = causal_valid_mask.unsqueeze(0) & index_valid_mask  # [b, sq, sk]
     else:
         valid_mask = causal_valid_mask.unsqueeze(0).expand(b, sq, sk)  # [b, sq, sk]
-    
+
     grad_index_scores_logits = grad_index_scores_logits * valid_mask.float()
-    
+
     # Transpose from [b, sq, sk] to [sq, b, sk]
     grad_index_scores = grad_index_scores_logits.transpose(0, 1)  # [sq, b, sk]
-    
+
     # Backward through sum over heads: expand gradient
     grad_weighted_scores = grad_index_scores.unsqueeze(2)  # [sq, b, 1, sk]
-    
+
     # Compute forward values needed for backward
     scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float())  # [sq, b, h, sk]
     scores_after_relu = torch.relu(scores)
-    
+
     # Backward through multiplication by weights: index_scores_per_head * weights
     # ∂L/∂weights = grad * relu_scores (sum over sk)
     grad_weights = (grad_weighted_scores * scores_after_relu).sum(dim=-1)  # [sq, b, h]
-    
+
     # ∂L/∂relu_scores = grad * weights
     grad_scores_after_relu = grad_weighted_scores * weights.unsqueeze(-1)  # [sq, b, h, sk]
-    
+
     # Backward through ReLU
     relu_mask = (scores > 0).float()
     grad_scores = grad_scores_after_relu * relu_mask  # [sq, b, h, sk]
-    
+
     # Backward through einsum 'sbhd,tbd->sbht'
     # ∂L/∂q = einsum('sbht,tbd->sbhd', grad_scores, k)
     grad_q = torch.einsum('sbht,tbd->sbhd', grad_scores, k.float())  # [sq, b, h, d]
@@ -449,14 +480,38 @@ def bwd_fused_indexer_loss_native(q, weights, k, query, key, topk_indices, softm
 
 
 class FusedDSAIndexerLoss(torch.autograd.Function):
+    """Fused implementation of DSA Indexer Loss."""
+
     @staticmethod
-    def forward(ctx, q, weights, k, query, key, softmax_scale, topk, loss_coeff, mask, sparse_loss, pg_collection):
+    def forward(
+        ctx,
+        q,
+        weights,
+        k,
+        query,
+        key,
+        softmax_scale,
+        topk,
+        loss_coeff,
+        mask,
+        sparse_loss,
+        pg_collection,
+    ):
         """
         Fused forward: index_scores never materialized in full.
         """
-
-        topk_indices, loss = fwd_fused_indexer_loss_native(
-            q, weights, k, query, key, topk, softmax_scale, loss_coeff, mask, sparse_loss, pg_collection,
+        topk_indices, loss = fwd_fused_indexer_loss_naive(
+            q,
+            weights,
+            k,
+            query,
+            key,
+            topk,
+            softmax_scale,
+            loss_coeff,
+            mask,
+            sparse_loss,
+            pg_collection,
         )
 
         # Save for backward (recomputation strategy)
@@ -475,9 +530,18 @@ def backward(ctx, grad_topk_indices, grad_loss):
         """
         q, weights, k, query, key, topk_indices = ctx.saved_tensors
 
-        grad_q , grad_weights, grad_k = bwd_fused_indexer_loss_native(
-            q, weights, k, query, key, topk_indices, 
-            ctx.softmax_scale, ctx.loss_coeff, ctx.sparse_loss, grad_loss, ctx.pg_collection,
+        grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive(
+            q,
+            weights,
+            k,
+            query,
+            key,
+            topk_indices,
+            ctx.softmax_scale,
+            ctx.loss_coeff,
+            ctx.sparse_loss,
+            grad_loss,
+            ctx.pg_collection,
         )
 
         # query and key are detached in forward, so return None for their gradients
@@ -704,11 +768,9 @@ def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, mscale: flo
         return x
 
     def forward_before_topk(
-        self,
-        x: torch.Tensor,
-        qr: torch.Tensor,
-        packed_seq_params: Optional[PackedSeqParams] = None,
+        self, x: torch.Tensor, qr: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None
     ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """All computations before topk."""
         # =========================================
         # Prepare RoPE params
         # =========================================
@@ -794,11 +856,13 @@ def forward_with_scores(
         """
         assert packed_seq_params is None, "Packed sequence is not supported for DSAttention"
 
-        # [seqlen, batch, index_n_heads * index_head_dim], [seqlen, batch, index_head_dim], [seqlen, batch, index_n_heads]
+        # [seqlen, batch, index_n_heads * index_head_dim]
+        # [seqlen, batch, index_head_dim]
+        # [seqlen, batch, index_n_heads]
         q, k, weights = self.forward_before_topk(x, qr, packed_seq_params)
 
         # [batch, seqlen, seqlen], [batch, seqlen, index_topk]
-        index_scores, topk_indices = fused_qk_topk_native(q, k, weights, self.index_topk, mask)
+        index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, self.index_topk, mask)
 
         return index_scores, topk_indices
 
@@ -983,15 +1047,15 @@ def forward(
             # ===================================
             # Compute KL divergence loss between indexer scores and true attention scores
             topk_indices, indexer_loss = FusedDSAIndexerLoss.apply(
-                q, 
-                weights, 
-                k, 
-                query.detach(), 
-                key.detach(), 
-                self.softmax_scale, 
-                self.indexer.index_topk, 
-                indexer_loss_coeff, 
-                float_mask, 
+                q,
+                weights,
+                k,
+                query.detach(),
+                key.detach(),
+                self.softmax_scale,
+                self.indexer.index_topk,
+                indexer_loss_coeff,
+                float_mask,
                 getattr(self.config, "dsa_indexer_use_sparse_loss", False),
                 self.indexer.pg_collection,
             )
@@ -1007,7 +1071,7 @@ def forward(
             # Run sparse attention kernel
             # ===================================
             output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
-            
+
             # Attach loss to output
             output = DSAIndexerLossAutoScaler.apply(output, indexer_loss)
 

From ee2753dc009d5b3efe637fec4347e00d2b841665 Mon Sep 17 00:00:00 2001
From: kunlunl <kunlunl@nvidia.com>
Date: Wed, 4 Feb 2026 03:22:13 +0800
Subject: [PATCH 3/5] Minor fix

---
 .../core/transformer/experimental_attention_variant/dsa.py    | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py
index 274059713c3..e1d18014542 100644
--- a/megatron/core/transformer/experimental_attention_variant/dsa.py
+++ b/megatron/core/transformer/experimental_attention_variant/dsa.py
@@ -420,9 +420,7 @@ def bwd_fused_indexer_loss_naive(
 
     # Backward through sum(dim=-1): broadcast back to [b, sq, sk]
     # Each element in a row contributes to the sum, so gradient is same for all
-    grad_kl_per_element = torch.full(
-        (b, sq, sk), grad_kl_per_row.item(), device=index_scores.device, dtype=torch.float32
-    )
+    grad_kl_per_element = grad_kl_per_row.view(1, 1, 1).expand(b, sq, sk)
 
     # Backward through kl_per_element = target * (log(target) - log(index))
     # ∂kl/∂index_softmax = -target / index_softmax

From 0c0b99511a8505e77164fa7042613430f41618e8 Mon Sep 17 00:00:00 2001
From: kunlunl <kunlunl@nvidia.com>
Date: Wed, 4 Feb 2026 03:52:48 +0800
Subject: [PATCH 4/5] Add UT for manual bwd of indexer loss

---
 .../transformer/test_attention_variant_dsa.py | 317 ++++++++++++++++++
 1 file changed, 317 insertions(+)

diff --git a/tests/unit_tests/transformer/test_attention_variant_dsa.py b/tests/unit_tests/transformer/test_attention_variant_dsa.py
index bd106aa6f0e..96253a4ca10 100644
--- a/tests/unit_tests/transformer/test_attention_variant_dsa.py
+++ b/tests/unit_tests/transformer/test_attention_variant_dsa.py
@@ -17,7 +17,10 @@
     DSAIndexerSubmodules,
     DSAttention,
     DSAttentionSubmodules,
+    FusedDSAIndexerLoss,
+    _compute_index_scores,
     compute_dsa_indexer_loss,
+    fused_qk_topk_naive,
     rotate_activation,
 )
 from megatron.core.transformer.transformer_config import MLATransformerConfig
@@ -265,6 +268,320 @@ def test_backward_pass(self):
         ), f"Gradient should be scaled by loss scale, expected {expected_grad_per_element}, got {dummy_input.grad[0].item()}"
 
 
+@pytest.mark.parametrize("seqlen_and_topk", [[16, 8], [32, 16], [64, 32]])
+@pytest.mark.parametrize("sparse_loss", [False, True])
+class TestFusedDSAIndexerLossGradient:
+    """Test that FusedDSAIndexerLoss manual backward matches autograd backward."""
+
+    @pytest.fixture(scope='function', autouse=True)
+    def setup_method(self):
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, pipeline_model_parallel_size=1
+        )
+        self.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp'])
+        yield
+        Utils.destroy_model_parallel()
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_fused_indexer_loss_gradient_matches_autograd(self, seqlen_and_topk, sparse_loss):
+        """
+        Test that the manually written backward in FusedDSAIndexerLoss produces
+        the same gradients as PyTorch autograd on the unfused implementation.
+        """
+        seqlen = seqlen_and_topk[0]
+        index_topk = seqlen_and_topk[1]
+        batch_size = 2
+        num_heads = 4
+        head_dim = 64
+        index_n_heads = 8
+        index_head_dim = 64
+        softmax_scale = head_dim**-0.5
+        loss_coeff = 1.0
+
+        torch.manual_seed(42)
+
+        # Create inputs for indexer
+        # q: [seqlen, batch, index_n_heads, index_head_dim]
+        q_ref = (
+            torch.randn(seqlen, batch_size, index_n_heads, index_head_dim, dtype=torch.float32)
+            .cuda()
+            .requires_grad_(True)
+        )
+        # weights: [seqlen, batch, index_n_heads]
+        weights_ref = (
+            torch.randn(seqlen, batch_size, index_n_heads, dtype=torch.float32)
+            .cuda()
+            .requires_grad_(True)
+        )
+        # k: [seqlen, batch, index_head_dim]
+        k_ref = (
+            torch.randn(seqlen, batch_size, index_head_dim, dtype=torch.float32)
+            .cuda()
+            .requires_grad_(True)
+        )
+        # query: [seqlen, batch, num_heads, head_dim] - detached, not trained
+        query = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda()
+        # key: [seqlen, batch, num_heads, head_dim] - detached, not trained
+        key = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda()
+
+        # Create causal mask
+        mask = torch.triu(
+            torch.full((seqlen, seqlen), float('-inf'), dtype=torch.float32).cuda(), diagonal=1
+        )
+
+        # =============================================
+        # Method 1: Autograd (reference)
+        # =============================================
+        # Compute index scores and apply mask (matches fused_qk_topk_naive behavior)
+        index_scores_ref = _compute_index_scores(q_ref, weights_ref, k_ref)
+        # Apply mask
+        index_scores_masked = index_scores_ref + mask.unsqueeze(0)
+        # Get topk indices from masked scores
+        topk_k = min(index_topk, seqlen)
+        topk_indices = index_scores_masked.topk(topk_k, dim=-1)[1]
+
+        # Compute loss using autograd
+        loss_ref = compute_dsa_indexer_loss(
+            index_scores=index_scores_masked,
+            topk_indices=topk_indices,
+            query=query,
+            key=key,
+            softmax_scale=softmax_scale,
+            loss_coeff=loss_coeff,
+            sparse_loss=sparse_loss,
+            pg_collection=self.pg_collection,
+        )
+
+        # Backward with autograd
+        loss_ref.backward()
+
+        # Save reference gradients
+        grad_q_ref = q_ref.grad.clone()
+        grad_weights_ref = weights_ref.grad.clone()
+        grad_k_ref = k_ref.grad.clone()
+
+        # =============================================
+        # Method 2: FusedDSAIndexerLoss (manual backward)
+        # =============================================
+        # Clone tensors from ref (detach and require grad again)
+        q_fused = q_ref.detach().clone().requires_grad_(True)
+        weights_fused = weights_ref.detach().clone().requires_grad_(True)
+        k_fused = k_ref.detach().clone().requires_grad_(True)
+
+        # Use FusedDSAIndexerLoss
+        topk_indices_fused, loss_fused = FusedDSAIndexerLoss.apply(
+            q_fused,
+            weights_fused,
+            k_fused,
+            query.detach(),
+            key.detach(),
+            softmax_scale,
+            index_topk,
+            loss_coeff,
+            mask,
+            sparse_loss,
+            self.pg_collection,
+        )
+
+        # Backward with manual implementation
+        loss_fused.backward()
+
+        # Get fused gradients
+        grad_q_fused = q_fused.grad
+        grad_weights_fused = weights_fused.grad
+        grad_k_fused = k_fused.grad
+
+        # =============================================
+        # Compare gradients
+        # =============================================
+        # Check loss values match
+        assert torch.allclose(
+            loss_fused, loss_ref, rtol=1e-5, atol=1e-5
+        ), f"Loss mismatch: fused={loss_fused.item()}, ref={loss_ref.item()}"
+
+        # Check topk indices match
+        assert torch.equal(
+            topk_indices_fused, topk_indices
+        ), "Top-k indices mismatch between fused and reference"
+
+        # Check gradients match
+        assert torch.allclose(
+            grad_q_fused, grad_q_ref, rtol=1e-5, atol=1e-5
+        ), f"grad_q mismatch: max diff = {(grad_q_fused - grad_q_ref).abs().max().item()}"
+
+        assert torch.allclose(
+            grad_weights_fused, grad_weights_ref, rtol=1e-5, atol=1e-5
+        ), f"grad_weights mismatch: max diff = {(grad_weights_fused - grad_weights_ref).abs().max().item()}"
+
+        assert torch.allclose(
+            grad_k_fused, grad_k_ref, rtol=1e-5, atol=1e-5
+        ), f"grad_k mismatch: max diff = {(grad_k_fused - grad_k_ref).abs().max().item()}"
+
+
+@pytest.mark.parametrize("tensor_model_parallel_size", [2, 4])
+@pytest.mark.parametrize("sparse_loss", [False, True])
+class TestFusedDSAIndexerLossGradientTP:
+    """Test FusedDSAIndexerLoss gradient consistency across different TP sizes."""
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_fused_indexer_loss_gradient_tp_consistency(
+        self, tensor_model_parallel_size, sparse_loss
+    ):
+        """
+        Test that FusedDSAIndexerLoss produces consistent gradients across TP ranks
+        and matches TP=1 baseline.
+        """
+        seqlen = 64
+        index_topk = 32
+        batch_size = 2
+        num_heads = 8
+        head_dim = 64
+        index_n_heads = 8
+        index_head_dim = 64
+        softmax_scale = head_dim**-0.5
+        loss_coeff = 1.0
+
+        # =============================================
+        # First run with TP=1 to get baseline
+        # =============================================
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, pipeline_model_parallel_size=1
+        )
+        torch.manual_seed(42)
+        model_parallel_cuda_manual_seed(42)
+
+        pg_collection_tp1 = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp'])
+
+        # Create inputs
+        q_input = torch.randn(
+            seqlen, batch_size, index_n_heads, index_head_dim, dtype=torch.float32
+        ).cuda()
+        weights_input = torch.randn(seqlen, batch_size, index_n_heads, dtype=torch.float32).cuda()
+        k_input = torch.randn(seqlen, batch_size, index_head_dim, dtype=torch.float32).cuda()
+        query_input = torch.randn(
+            seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16
+        ).cuda()
+        key_input = torch.randn(
+            seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16
+        ).cuda()
+        mask = torch.triu(
+            torch.full((seqlen, seqlen), float('-inf'), dtype=torch.float32).cuda(), diagonal=1
+        )
+
+        # Clone for TP=1
+        q_tp1 = q_input.clone().requires_grad_(True)
+        weights_tp1 = weights_input.clone().requires_grad_(True)
+        k_tp1 = k_input.clone().requires_grad_(True)
+
+        # Forward and backward with TP=1
+        topk_indices_tp1, loss_tp1 = FusedDSAIndexerLoss.apply(
+            q_tp1,
+            weights_tp1,
+            k_tp1,
+            query_input.detach(),
+            key_input.detach(),
+            softmax_scale,
+            index_topk,
+            loss_coeff,
+            mask,
+            sparse_loss,
+            pg_collection_tp1,
+        )
+        loss_tp1.backward()
+
+        # Save TP=1 results
+        grad_q_tp1 = q_tp1.grad.clone()
+        grad_weights_tp1 = weights_tp1.grad.clone()
+        grad_k_tp1 = k_tp1.grad.clone()
+        loss_tp1_value = loss_tp1.detach().clone()
+
+        Utils.destroy_model_parallel()
+
+        # =============================================
+        # Run with target TP size
+        # =============================================
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tensor_model_parallel_size, pipeline_model_parallel_size=1
+        )
+        torch.manual_seed(42)
+        model_parallel_cuda_manual_seed(42)
+
+        pg_collection_tpn = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp'])
+        tp_rank = parallel_state.get_tensor_model_parallel_rank()
+
+        # Clone inputs for TP=N (same values as TP=1)
+        q_tpn = q_input.clone().requires_grad_(True)
+        weights_tpn = weights_input.clone().requires_grad_(True)
+        k_tpn = k_input.clone().requires_grad_(True)
+
+        # query and key need to be split along heads for TP
+        head_per_rank = num_heads // tensor_model_parallel_size
+        start_head = tp_rank * head_per_rank
+        end_head = (tp_rank + 1) * head_per_rank
+        query_tpn = query_input[:, :, start_head:end_head, :].clone()
+        key_tpn = key_input[:, :, start_head:end_head, :].clone()
+
+        # Forward and backward with TP=N
+        topk_indices_tpn, loss_tpn = FusedDSAIndexerLoss.apply(
+            q_tpn,
+            weights_tpn,
+            k_tpn,
+            query_tpn.detach(),
+            key_tpn.detach(),
+            softmax_scale,
+            index_topk,
+            loss_coeff,
+            mask,
+            sparse_loss,
+            pg_collection_tpn,
+        )
+        loss_tpn.backward()
+
+        # =============================================
+        # Compare results
+        # =============================================
+        # Loss should be the same
+        assert torch.allclose(
+            loss_tpn, loss_tp1_value, rtol=1e-5, atol=1e-5
+        ), f"Loss mismatch: TP={tensor_model_parallel_size} got {loss_tpn.item()}, TP=1 got {loss_tp1_value.item()}"
+
+        # Top-k indices should be the same
+        assert torch.equal(
+            topk_indices_tpn, topk_indices_tp1
+        ), "Top-k indices mismatch between TP=1 and TP=N"
+
+        # Gradients should match exactly (indexer params are duplicated across TP)
+        assert torch.allclose(
+            q_tpn.grad, grad_q_tp1, rtol=1e-5, atol=1e-5
+        ), f"grad_q mismatch: max diff = {(q_tpn.grad - grad_q_tp1).abs().max().item()}"
+
+        assert torch.allclose(
+            weights_tpn.grad, grad_weights_tp1, rtol=1e-5, atol=1e-5
+        ), f"grad_weights mismatch: max diff = {(weights_tpn.grad - grad_weights_tp1).abs().max().item()}"
+
+        assert torch.allclose(
+            k_tpn.grad, grad_k_tp1, rtol=1e-5, atol=1e-5
+        ), f"grad_k mismatch: max diff = {(k_tpn.grad - grad_k_tp1).abs().max().item()}"
+
+        # Check gradients are identical across all TP ranks
+        tp_size = parallel_state.get_tensor_model_parallel_world_size()
+        if tp_size > 1:
+            for grad_tensor, name in [
+                (q_tpn.grad, "grad_q"),
+                (weights_tpn.grad, "grad_weights"),
+                (k_tpn.grad, "grad_k"),
+            ]:
+                grad_list = [torch.zeros_like(grad_tensor) for _ in range(tp_size)]
+                torch.distributed.all_gather(grad_list, grad_tensor, group=pg_collection_tpn.tp)
+
+                for i in range(1, tp_size):
+                    assert torch.allclose(
+                        grad_list[0], grad_list[i], rtol=0, atol=0
+                    ), f"{name} differs between TP rank 0 and rank {i}"
+
+        Utils.destroy_model_parallel()
+
+
 @pytest.mark.parametrize("seqlen", [16, 64])
 class TestDSAIndexer:
     """Test DSA Indexer module basic functionality with TP=1."""

From 7673ae4d561f0c5ab1e2ab451f2eeaa34b781439 Mon Sep 17 00:00:00 2001
From: kunlunl <kunlunl@nvidia.com>
Date: Wed, 4 Feb 2026 05:08:10 +0800
Subject: [PATCH 5/5] Free unused tensors immediately to save memory

---
 .../experimental_attention_variant/dsa.py     | 36 +++++++++++++++++--
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py
index e1d18014542..3734db7043f 100644
--- a/megatron/core/transformer/experimental_attention_variant/dsa.py
+++ b/megatron/core/transformer/experimental_attention_variant/dsa.py
@@ -368,6 +368,9 @@ def bwd_fused_indexer_loss_naive(
     key_reshaped = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk)
     # Compute attention scores [b * np, sq, sk]
     attention_scores = torch.bmm(query_reshaped.float(), key_reshaped.float()) * softmax_scale
+    # Free reshaped tensors - no longer needed after bmm
+    del query_reshaped, key_reshaped
+
     # Reshape to [b, np, sq, sk]
     attention_scores = attention_scores.reshape(b, np, sq, sk)
 
@@ -386,6 +389,8 @@ def bwd_fused_indexer_loss_naive(
     attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk)
     # [b, sq, sk] + [1, sq, sk] -> [b, sq, sk]
     index_scores = index_scores + causal_mask.unsqueeze(0)
+    # Free causal_mask - no longer needed
+    del causal_mask
 
     if sparse_loss:
         # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk]
@@ -397,10 +402,17 @@ def bwd_fused_indexer_loss_naive(
     attention_scores_softmax = torch.nn.functional.softmax(
         attention_scores, dim=-1, dtype=torch.float32
     )
+    # Free attention_scores immediately
+    del attention_scores
+
     index_scores_softmax = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32)
+    # Free index_scores - no longer needed after softmax
+    del index_scores
 
     # Sum attention scores across heads: [b, np, sq, sk] -> [b, sq, sk]
     attention_scores_sum = attention_scores_softmax.sum(dim=1)
+    # Free attention_scores_softmax
+    del attention_scores_softmax
 
     if pg_collection.tp.size() > 1:
         # attention scores are scattered to TP ranks in head dimension.
@@ -410,6 +422,8 @@ def bwd_fused_indexer_loss_naive(
     attention_scores_normalized = attention_scores_sum / attention_scores_sum.sum(
         dim=-1, keepdim=True
     )
+    # Free attention_scores_sum - no longer needed after normalization
+    del attention_scores_sum
 
     # Backward through loss = kl_div * loss_coeff
     # where kl_div = kl_per_element.sum(dim=-1).mean()
@@ -427,35 +441,49 @@ def bwd_fused_indexer_loss_naive(
     grad_index_scores_softmax = (
         -attention_scores_normalized / (index_scores_softmax + 1e-10) * grad_kl_per_element
     )
+    # Free attention_scores_normalized - no longer needed
+    del attention_scores_normalized
 
     # Backward through softmax: ∂L/∂x = softmax * (∂L/∂softmax - sum(∂L/∂softmax * softmax))
     sum_grad = (grad_index_scores_softmax * index_scores_softmax).sum(dim=-1, keepdim=True)
     grad_index_scores_logits = index_scores_softmax * (grad_index_scores_softmax - sum_grad)
+    # Free intermediate tensors
+    del index_scores_softmax, grad_index_scores_softmax, sum_grad
 
     # Zero out gradients for masked positions
     # Create a mask for valid (non-masked) positions
     # Causal mask: position (i, j) is valid if j <= i
     causal_valid_mask = torch.tril(
-        torch.ones((sq, sk), device=index_scores.device, dtype=torch.bool)
+        torch.ones((sq, sk), device=q.device, dtype=torch.bool)
     )  # [sq, sk]
     if sparse_loss:
         # Also apply index mask - only topk positions are valid
         index_valid_mask = index_mask == 0  # [b, sq, sk]
+        del index_mask  # Free index_mask immediately after use
         valid_mask = causal_valid_mask.unsqueeze(0) & index_valid_mask  # [b, sq, sk]
+        del index_valid_mask
     else:
+        del index_mask  # Free index_mask even if not used for sparse_loss
         valid_mask = causal_valid_mask.unsqueeze(0).expand(b, sq, sk)  # [b, sq, sk]
+    del causal_valid_mask
 
     grad_index_scores_logits = grad_index_scores_logits * valid_mask.float()
+    del valid_mask
 
     # Transpose from [b, sq, sk] to [sq, b, sk]
     grad_index_scores = grad_index_scores_logits.transpose(0, 1)  # [sq, b, sk]
+    del grad_index_scores_logits
 
     # Backward through sum over heads: expand gradient
     grad_weighted_scores = grad_index_scores.unsqueeze(2)  # [sq, b, 1, sk]
+    del grad_index_scores
 
     # Compute forward values needed for backward
     scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float())  # [sq, b, h, sk]
+    # Compute relu_mask before relu (saves memory vs keeping both scores and relu output)
+    relu_mask = scores > 0
     scores_after_relu = torch.relu(scores)
+    del scores
 
     # Backward through multiplication by weights: index_scores_per_head * weights
     # ∂L/∂weights = grad * relu_scores (sum over sk)
@@ -463,16 +491,18 @@ def bwd_fused_indexer_loss_naive(
 
     # ∂L/∂relu_scores = grad * weights
     grad_scores_after_relu = grad_weighted_scores * weights.unsqueeze(-1)  # [sq, b, h, sk]
+    del grad_weighted_scores, scores_after_relu
 
     # Backward through ReLU
-    relu_mask = (scores > 0).float()
-    grad_scores = grad_scores_after_relu * relu_mask  # [sq, b, h, sk]
+    grad_scores = grad_scores_after_relu * relu_mask.float()  # [sq, b, h, sk]
+    del grad_scores_after_relu, relu_mask
 
     # Backward through einsum 'sbhd,tbd->sbht'
     # ∂L/∂q = einsum('sbht,tbd->sbhd', grad_scores, k)
     grad_q = torch.einsum('sbht,tbd->sbhd', grad_scores, k.float())  # [sq, b, h, d]
     # ∂L/∂k = einsum('sbht,sbhd->tbd', grad_scores, q)
     grad_k = torch.einsum('sbht,sbhd->tbd', grad_scores, q.float())  # [sk, b, d]
+    del grad_scores
 
     return grad_q.to(q.dtype), grad_weights.to(weights.dtype), grad_k.to(k.dtype)