[Bug] DSV4 unfused CSA indexer loss uses a compressed-only teacher instead of the full CSA attention distribution

@NVIDIA/mcore-oncall

## Describe the bug

The DeepSeek-V4 unfused / small-operator CSA path computes the indexer-loss teacher by recomputing attention over compressed KV only. This teacher is not equivalent to the real CSA attention distribution described by the DeepSeek-V4 paper, because the real CSA core attention also includes the original-KV sliding window and the attention sink in a shared softmax denominator.

This is a semantic mismatch in the teacher distribution, not a floating-point precision difference.

Static analysis scope:

* Megatron-LM dev commit: fd1121b8ff7e3a4f83a28d35aed172d7bc0260e1
* Primary path: unfused CSA indexer loss, both SBHD and THD
* No runtime or convergence claim is required to establish the mismatch below

Related DSV4 tracking issue: NVIDIA/Megatron-LM#4468

## Paper reference

DeepSeek-V3.2 defines the Lightning Indexer teacher from the real main-attention distribution: per-head attention probabilities are aggregated across heads and then L1-normalized (Section 2.1, equations (3) and (4)):

[https://arxiv.org/html/2512.02556v1#S2.SS1.SSS1]()

DeepSeek-V4 keeps the Lightning Indexer objective while changing the CSA core attention. The CSA attention contains all of the following components:

1. Compressed KV selected by the Lightning Indexer:
[https://arxiv.org/html/2606.19348v1#S2.SS3.SSS1.Px2]()
2. Original KV from the sliding window:
[https://arxiv.org/html/2606.19348v1#S2.SS3.SSS3.Px3]()
3. An attention sink in the same normalization denominator:
[https://arxiv.org/html/2606.19348v1#S2.SS3.SSS3.Px4]()

For a query and attention head h, let c(h,j) be a compressed-key logit, w(h,r) a window-key logit, and a(h) the sink logit. The compressed part of the paper-equivalent teacher must retain the real per-head CSA denominator:

```math
Z_h = \exp(a_h) + \sum_{j \in C}\exp(c_{h,j}) + \sum_{r \in W}\exp(w_{h,r})
```

```math
t_j =
\frac{
\sum_h \frac{\exp(c_{h,j})}{Z_h}
}{
\sum_{k \in C}\sum_h \frac{\exp(c_{h,k})}{Z_h}
}
```

Consequently, changing only the window logits or sink logits can change the compressed teacher distribution because it changes the relative contribution of each attention head.

## Megatron-LM unfused / small-operator implementation

### 1. The loss receives compressed KV only

In the SBHD unfused path, key_for_loss is constructed only from compressed_kv and passed to FusedDSAIndexerLoss:

[https://github.com/NVIDIA/Megatron-LM/blob/fd1121b8ff7e3a4f83a28d35aed172d7bc0260e1/megatron/core/transformer/experimental_attention_variant/csa.py#L1606-L1640]()

Relevant data flow:

```
key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1)

topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply(
...,
query.detach(),
key_for_loss.detach(),
...
)
```

Neither window KV/logits nor attn_sink is an input to this loss call.

The THD unfused path does the same with key_for_loss_thd:

[https://github.com/NVIDIA/Megatron-LM/blob/fd1121b8ff7e3a4f83a28d35aed172d7bc0260e1/megatron/core/transformer/experimental_attention_variant/csa.py#L1927-L2000]()

### 2. The small operators recompute a compressed-only softmax

compute_dsa_indexer_loss calculates QK scores from the supplied query and key, applies the compressed causal/Top-k masks, softmaxes over that key axis, sums attention heads, and L1-normalizes:

[https://github.com/NVIDIA/Megatron-LM/blob/fd1121b8ff7e3a4f83a28d35aed172d7bc0260e1/megatron/core/transformer/experimental_attention_variant/dsa.py#L275-L353]()

Therefore, ignoring masks for clarity, the implemented teacher is:

```math
\hat{t}_j =
\frac{
\sum_h
\frac{\exp(c_{h,j})}
{\sum_{l \in C}\exp(c_{h,l})}
}{
\sum_{k \in C}\sum_h
\frac{\exp(c_{h,k})}
{\sum_{l \in C}\exp(c_{h,l})}
}
```

With sparse loss enabled, the Top-k mask is applied before this per-head softmax, so C above becomes the selected Top-k subset.

This differs from the paper teacher because the window and sink terms have been removed from every per-head denominator. It forces every head to contribute unit compressed probability mass before head aggregation, even when a head's real CSA probability mass is mostly assigned to its sliding window or sink.

### 3. The real unfused CSA attention uses a different denominator

Only after the indexer loss has already been computed, the implementation concatenates the window indices with the compressed Top-k and calls the actual sparse attention with attn_sink:

[https://github.com/NVIDIA/Megatron-LM/blob/fd1121b8ff7e3a4f83a28d35aed172d7bc0260e1/megatron/core/transformer/experimental_attention_variant/csa.py#L1652-L1673]()

```
topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1)

output = unfused_compressed_sparse_attn(
query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale
)
```

Thus the output attention and the indexer-loss teacher do not use the same attention distribution:

| Component | Real unfused CSA output | Unfused indexer-loss teacher |
| -- | -- | -- |
| Compressed KV | Yes | Yes |
| Original-KV sliding window | Yes | No |
| Attention sink | Yes | No |
| Shared CSA softmax denominator | Yes | No |

## Static reproduction / invariant

Hold query and compressed_kv fixed, and perturb only either:

* the original-KV window values/logits; or
* attn_sink for one attention head.

Expected from the paper equations:

* the per-head CSA denominator changes;
* the relative head weights in the compressed teacher can change;
* therefore the final L1-normalized teacher can change.

Current unfused implementation:

* FusedDSAIndexerLoss receives exactly the same query and key_for_loss;
* the recomputed compressed-only teacher is exactly unchanged.

This invariant violation follows directly from the function inputs and does not depend on a particular GPU kernel or numerical tolerance.

## Additional Top-k concern

When dsa_indexer_topk is smaller than the number of visible compressed keys, sparse loss applies the Top-k mask before the per-head attention softmax. This is generally different from:

1. computing the teacher with the full real CSA denominator;
2. gathering the selected compressed entries; and
3. applying one final L1 normalization.

The current short/full-selection cases can hide this additional difference.

## Expected behavior

The indexer-loss teacher should be derived from the real CSA attention probabilities, or from an equivalent recomputation that retains the same per-head denominator containing compressed attention, window attention, and sink.

For sparse loss, selected compressed entries should be derived from that teacher distribution and then normalized according to the Lightning Indexer objective, rather than rebuilding independent per-head softmax denominators over the selected compressed subset.

## Additional context: fused path

The fused path appears to have the same semantic issue. It places compressed Top-k entries before window entries, obtains both lse and lse_indexer from FlashMLA, but passes lse_indexer to the target recomputation:

[https://github.com/NVIDIA/Megatron-LM/blob/fd1121b8ff7e3a4f83a28d35aed172d7bc0260e1/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py#L1135-L1257]()

FlashMLA defines lse_indexer as the LSE over the first indexer_topk entries only, while attn_sink does not affect lse/lse_indexer:

[https://github.com/deepseek-ai/FlashMLA/blob/b7643bd54521f563b839b98289b5cd048c062ba2/flash_mla/flash_mla_interface.py#L176-L215]()

Therefore the fused target also appears to omit the sliding-window and sink competition from the teacher denominator. The primary bug reported here is the statically explicit unfused/small-operator path, but the fused path likely needs the same semantic correction.