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

[JavaZeroo]

@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: fd1121b
• 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: #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:

$$Z_h = \exp(a_h) + \sum_{j \in C}\exp(c_{h,j}) + \sum_{r \in W}\exp(w_{h,r})$$ $$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:

Megatron-LM/megatron/core/transformer/experimental_attention_variant/csa.py

Lines 1606 to 1640 in fd1121b

	causal_mask = (
	torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1)
	)
	positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1)
	causal_mask = (
	torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0)
	.unsqueeze(0)
	.expand(b, -1, -1)
	) # [b, sq, n_compressed]
	if self.training and torch.is_grad_enabled():
	q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk(
	x_det, qr_det, packed_seq_params
	)
	indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0)
	key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1)
	# ``FusedDSAIndexerLoss`` does not accept a separate
	# indexer_softmax_scale; apply it here via the
	# weights-scaling trick so the effective weights match
	# the pre-scale-split behaviour.
	weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale
	topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply(
	q_indexer,
	weights_for_unfused,
	k_indexer,
	query.detach(),
	key_for_loss.detach(),
	self.softmax_scale,
	min(self.indexer.index_topk, n_compressed),
	indexer_loss_coeff,
	causal_mask,
	getattr(self.config, "dsa_indexer_use_sparse_loss", True),
	self.indexer.pg_collection,
	self.config.calculate_per_token_loss,
	)

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:

Megatron-LM/megatron/core/transformer/experimental_attention_variant/csa.py

Lines 1927 to 2000 in fd1121b

	if self.training and torch.is_grad_enabled():
	q_indexer, k_indexer, weights_indexer, cu_seqlens_compressed_idx = (
	self.indexer.forward_before_topk(x_det, qr_det, packed_seq_params)
	)
	if k_indexer is None:
	raise RuntimeError(
	"CompressedSparseAttention THD unfused Path B requires "
	"at least one segment with compressed indexer K."
	)
	q_thd = q_indexer.squeeze(1)
	w_thd = weights_indexer.squeeze(1)
	k_thd = k_indexer.squeeze(1)
	key_for_loss_thd = compressed_kv.unsqueeze(1).expand(-1, np_, -1)
	weights_for_unfused = w_thd * self.indexer.softmax_scale
	indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0)
	# ``_forward_thd`` (caller) absorbs trailing padded
	# tokens into the last ``cu_seqlens_q[-1]`` bucket
	# so ``batch_of_row`` doesn't OOB; the Compressor
	# does the same to ``cu_seqlens_compressed_idx[-1]``.
	# Both are correct for the sparse-attention path,
	# but the per-segment indexer-loss loop in
	# ``fwd/bwd_fused_indexer_loss_naive_thd`` would
	# then iterate a "fake" absorbed-padding segment
	# with ``seqlen_k_b < topk`` — triggering a write
	# shape mismatch and a downstream ``scatter_`` OOB
	# on its ``-1`` entries.
	#
	# Restore the original (pre-absorption) cu_seqlens
	# for the loss path so the segment loop's
	# ``if seqlen_k_b == 0: continue`` guard skips the
	# padding-only iteration. When no padding exists,
	# ``packed_seq_params`` already equals the absorbed
	# version and this is a no-op.
	# Rebuild compressed cu_seqlens from *unpadded* Q lengths.
	# This may disagree with the Indexer's cu_seqlens_compressed_idx
	# (which uses padded lengths) for the last segment — that's
	# intentional: the extra compressed tokens from padding sit at
	# the tail of k_thd and are simply never visited by the loss
	# loop, which is correct since they don't represent real data.
	# Non-last segments are unaffected (padding absorption only
	# extends the final segment).
	cu_seqlens_q_for_loss = packed_seq_params.cu_seqlens_q
	seg_lens_q = cu_seqlens_q_for_loss[1:] - cu_seqlens_q_for_loss[:-1]
	cu_seqlens_compressed_idx_for_loss = torch.cat(
	[
	torch.zeros(
	1,
	dtype=cu_seqlens_q_for_loss.dtype,
	device=cu_seqlens_q_for_loss.device,
	),
	(seg_lens_q // self.compress_ratio)
	.cumsum(0)
	.to(cu_seqlens_q_for_loss.dtype),
	]
	)
	topk_indices_cmp, indexer_loss = FusedDSAIndexerLoss.apply(
	q_thd,
	weights_for_unfused,
	k_thd,
	query.detach(),
	key_for_loss_thd.detach(),
	self.softmax_scale,
	min(self.indexer.index_topk, max_seqlen_compressed_idx),
	indexer_loss_coeff,
	None,
	getattr(self.config, "dsa_indexer_use_sparse_loss", True),
	self.indexer.pg_collection,
	self.config.calculate_per_token_loss,
	cu_seqlens_q_for_loss,
	cu_seqlens_compressed_idx_for_loss,
	self.compress_ratio,
	)

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:

Megatron-LM/megatron/core/transformer/experimental_attention_variant/dsa.py

Lines 275 to 353 in fd1121b

	# [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk]
	key = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk)
	# Compute attention scores [b * np, sq, sk]
	attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale
	# Reshape to [b, np, sq, sk]
	attention_scores = attention_scores.reshape(b, np, sq, sk)
	# causal_mask: use caller-provided mask when available (handles compressed KV),
	# otherwise fall back to standard upper-triangular causal mask.
	if causal_mask_override is not None:
	causal_mask = causal_mask_override.to(dtype=torch.float32) # [b, sq, sk]
	else:
	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 attention_scores
	# causal_mask: [b, sq, sk] (from causal_mask_override) or [sq, sk] (from triu)
	if causal_mask.dim() == 3:
	attention_scores = attention_scores + causal_mask.unsqueeze(1) # [b,1,sq,sk]
	else:
	attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk)
	if sparse_loss:
	# [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk]
	attention_scores += index_mask.view(b, 1, sq, sk)
	# [b, sq, sk] + [b, sq, sk] -> [b, sq, sk]
	index_scores += index_mask
	# Identify rows where all KV positions are masked (e.g., early query positions with
	# compress_ratio=4 have zero valid compressed KV entries). These rows would produce NaN
	# from softmax(all -inf). We zero out their logits before softmax and mask out their
	# contributions after, so NaN is never produced.
	# row_valid: [b, sq] or [sq] — True if the row has at least one unmasked position.
	row_valid = (causal_mask > float('-inf')).any(dim=-1)
	if row_valid.dim() == 1:
	# [sq] -> broadcast for attention_scores [b, np, sq, sk] and index_scores [b, sq, sk]
	attn_row_mask = row_valid.view(1, 1, sq, 1) # [1, 1, sq, 1]
	idx_row_mask = row_valid.view(1, sq, 1) # [1, sq, 1]
	else:
	# [b, sq]
	attn_row_mask = row_valid.view(b, 1, sq, 1) # [b, 1, sq, 1]
	idx_row_mask = row_valid.view(b, sq, 1) # [b, sq, 1]
	# Zero out fully-masked rows before softmax so it produces valid uniform distribution
	attention_scores = attention_scores.masked_fill(~attn_row_mask, 0.0)
	index_scores = index_scores.masked_fill(~idx_row_mask, 0.0)
	# [b, np, sq, sk] -> [b, np, sq, sk]
	attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32)
	# [b, sq, sk] -> [b, sq, sk]
	index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32)
	# Zero out invalid rows so they contribute nothing to loss/gradients
	attention_scores = attention_scores * attn_row_mask.float()
	index_scores = index_scores * idx_row_mask.float()
	# Sum attention scores across heads.
	# [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k]
	attention_scores = attention_scores.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.contiguous(), group=pg_collection.tp)
	# L1 normalize target on the last dimension. Doesn't use abs() because attention_scores are
	# obtained from softmax so they are already non-negative.
	attention_scores = attention_scores / (
	attention_scores.sum(dim=-1, keepdim=True).clamp(min=1e-10)
	)
	# Compute KL divergence: KL(target || index) = target(x) * log(target(x) / index(x))
	# kl_per_element [b, sq, sk]
	kl_per_element = attention_scores * (
	torch.log(attention_scores + 1e-10) - torch.log(index_scores + 1e-10)

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

$$\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:

Megatron-LM/megatron/core/transformer/experimental_attention_variant/csa.py

Lines 1652 to 1673 in fd1121b

	n_valid_per_pos = positions // self.compress_ratio # [sq, 1]
	valid = topk_indices_compressed < n_valid_per_pos
	compress_topk_idxs = torch.where(
	valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device)
	)
	else:
	compress_topk_idxs = get_compress_topk_idxs(
	self.compress_ratio, b, sq, offset, query.device
	)
	topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1)
	nvtx_range_pop("compressed_indices")
	else:
	topk_idxs = window_idxs
	topk_idxs = topk_idxs.int()
	nvtx_range_push("sparse_attn_kernel")
	output = unfused_compressed_sparse_attn(
	query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale
	)
	nvtx_range_pop("sparse_attn_kernel")

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:

Megatron-LM/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py

Lines 1135 to 1257 in fd1121b

	)
	combined_local = torch.cat([compress_topk_idxs, window_idxs], dim=-1)
	global_idxs = local_to_global_flat(combined_local, b)
	# ---- 4. FlashMLA forward (flat layout for both SBHD and THD). --------
	if is_thd:
	q_flat = query
	kv_flat = kv_full
	else:
	q_flat = query.reshape(sq * b, np_, d)
	kv_flat = kv_full.reshape(skv * b, d)
	out_flat, lse, lse_indexer = _dsa_fwd_flash_mla(
	q_flat,
	kv_flat,
	global_idxs,
	softmax_scale,
	attn_sink=attn_sink,
	topk_length=None,
	indexer_topk=indexer_topk,
	)
	# ---- 4b. Derive padding-row mask for loss exclusion. -----------------
	# When CUDA-graph padding makes cu_seqlens_q cover all total_q rows
	# (including padding), cu_seqlens_q_unpadded supplies the true
	# boundaries. Padding rows must not contribute to the indexer KL
	# loss or backward gradients — only the sparse-attention output
	# needs them for static-shape compatibility.
	# The caller only passes cu_seqlens_q_unpadded when it differs from
	# cu_seqlens_q (checked via data_ptr), so no GPU→CPU sync is needed.
	padding_row_mask: Optional[Tensor] = None # True = padding (excluded from loss)
	if is_thd and cu_seqlens_q_unpadded is not None:
	real_seg_lens = cu_seqlens_q_unpadded[1:] - cu_seqlens_q_unpadded[:-1]
	row_idx = torch.arange(total_q, device=query.device, dtype=torch.int32)
	row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q)
	pos_in_seg = row_idx - cu_seqlens_q[row_batch_ids].to(torch.int32)
	# Rows whose intra-segment position >= real segment length
	# are padding (including rows in a dummy trailing segment
	# whose real length is 0).
	real_len_per_row = real_seg_lens[row_batch_ids].to(torch.int32)
	padding_row_mask = pos_in_seg >= real_len_per_row
	# ---- 5. Derive predict from indexer_scores, compute target. ----------
	# Layout-specific attn tensors (detached — loss is not differentiable
	# through them).
	if is_thd:
	assert compressed_kv is not None, "compressed_kv is required for THD"
	q_attn_det = query.detach()
	k_attn_compressed_det = compressed_kv.detach()
	lse_indexer_det = lse_indexer.detach()
	else:
	q_attn_det = query.detach().permute(1, 0, 2, 3).contiguous()
	k_attn_compressed_det = kv_full[kv_offset:].detach().permute(1, 0, 2).contiguous()
	lse_indexer_det = lse_indexer.reshape(sq, b, np_).permute(1, 0, 2)
	# Invalidate padding rows for the loss/backward path. The sparse
	# attention (steps 3-4) has already built global_idxs from the
	# original topk_indices_cmp, so this mutation only affects steps 5-7.
	if padding_row_mask is not None:
	topk_indices_cmp = topk_indices_cmp.clone()
	topk_indices_cmp[padding_row_mask] = -1
	indexer_scores = indexer_scores.clone()
	indexer_scores[padding_row_mask] = float('-inf')
	if sparse_loss:
	# Derive predict: gather topk scores from indexer_scores → softmax.
	safe_indices = topk_indices_cmp.clamp(min=0).long()
	gathered_scores = torch.gather(indexer_scores, dim=-1, index=safe_indices)
	gathered_scores = torch.where(
	topk_indices_cmp >= 0, gathered_scores, torch.finfo(torch.float32).min
	)
	predict = torch.softmax(gathered_scores, dim=-1)
	# THD: _compute_attn_target's kernel addresses K by flat ids over
	# the packed (total_k, D) buffer, so promote per-segment-local
	# indices to flat-global against cu_seqlens_compressed_idx.
	if is_thd:
	topk_for_target = local_to_global_flat(
	topk_indices_cmp,
	batch_size=-1,
	cu_seqlens_q=cu_seqlens_q,
	cu_seqlens_kv=cu_seqlens_compressed_idx,
	)
	else:
	topk_for_target = topk_indices_cmp
	target = _compute_attn_target(
	q_attn_det,
	k_attn_compressed_det,
	lse_indexer_det,
	topk_for_target,
	softmax_scale,
	qhead_per_kv_head=np_,
	topk_indices_global=is_thd,
	)
	if loss_coeff > 0:
	indexer_loss = _kl_loss_from_target_predict(
	target, predict, topk_indices_cmp, loss_coeff, calculate_per_token_loss
	)
	else:
	indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32)
	else:
	index_score = indexer_scores
	index_lse = torch.logsumexp(indexer_scores, dim=-1)
	k_unsqueeze_dim = 1 if is_thd else 2
	dense_attn_kwargs = {}
	if is_thd:
	dense_attn_kwargs = dict(
	cu_seqlens_q=cu_seqlens_q,
	cu_seqlens_kv=cu_seqlens_compressed_idx,
	max_seqlen_q=int(max_seqlen_q),
	max_seqlen_kv=int(max_seqlen_compressed_idx),
	)
	attn_score, attn_l1norm = _compute_dense_attn_score(
	q_attn_det,
	k_attn_compressed_det.unsqueeze(k_unsqueeze_dim),
	lse_indexer_det,
	qhead_per_kv_head=np_,
	softmax_scale=softmax_scale,
	ratio=ratio,
	**dense_attn_kwargs,
	)

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.

[JavaZeroo]

@hxbai @kunlunl Could you please help confirm and triage this DSV4 correctness issue?
cc @wujingyue @sbhavani

[cspades]

@mehraakash do you have any thoughts about this? Can we implement this?

[JavaZeroo]

@cspades @mehraakash is there any update?

[chengcuiping]

Hi @mehraakash, I’m interested in helping with this issue. I can first add a focused unfused SBHD/THD regression oracle showing that changing only the sliding-window or sink logits must change the compressed-key teacher distribution, then implement the smallest denominator-correct fix without touching the fused path. Are you already working on an implementation, or would this scoped contribution be useful?

[mehraakash]

Hi @chengcuiping,
I do not have a fix planned yet. Could you add more details about the change here.

[chengcuiping]

Thanks, @mehraakash. Here is the concrete scope I have in mind.

First, a correction to my earlier wording: a window/sink perturbation does not necessarily change the final L1-normalized target. The denominator cancels in the single-head case. A valid regression oracle therefore needs at least two attention heads and two causally valid compressed target entries, with different compressed-key conditional distributions across the heads. Under those conditions, changing the window or sink mass of one head has a provable effect on the final teacher.

For the first scoped change, I propose to cover the two non-CP unfused CSA callers of FusedDSAIndexerLoss: _forward_unfused_csa and _forward_unfused_csa_thd.

For each query and local attention head, the caller would compute a detached FP32 log-mass for the non-compressed part of the real CSA denominator:

$$ u_h = \log\left( \exp(a_h) + \sum_{r \in W_{\mathrm{valid}}}\exp(w_{h,r}) \right). $$

This computation must reuse the same attention query/KV values, softmax_scale, window indices, invalid-index masking, and sink semantics as unfused_compressed_sparse_attn. In particular, the packed THD path must use the same per-segment/globalized window indexing as the actual attention path.

The narrowest interface I plan to validate is an optional CSA-only non_compressed_lse input to the naive target computation behind FusedDSAIndexerLoss. After compressed Top-k selection, the target would use

$$ \ell_h = \log\left( \exp(u_h) + \sum_{j \in C}\exp(c_{h,j}) \right), \qquad p_{h,j} = \exp(c_{h,j} - \ell_h). $$

For sparse loss, $C$ is the selected Top-k intersected with the causally valid compressed entries used by the CSA core attention. For dense-loss mode, $C$ is the set of all causally visible compressed entries, preserving the current dense support and masking policy.

The existing sum across local attention heads, TP all-reduce, and single final L1 normalization over the compressed support would remain unchanged. The indexer-score softmax and KL definition would also remain unchanged. Non-CSA callers would provide no additional denominator term and retain their current behavior.

All teacher-side inputs, including the non-compressed LSE, would remain detached. The custom backward would reconstruct the same corrected target and continue to produce gradients only for the indexer Q, weights, and K. It would return no gradients for the main-attention query/KV, window contribution, or sink.

The focused tests would cover:

• an explicit FP32 two-head/two-entry oracle with different per-head compressed-key preferences, where perturbing one head’s sink changes the teacher;
• the equivalent window-mass perturbation;
• exact construction of the window-plus-sink LSE from the same SBHD and packed-THD indexing used by the reference attention;
• sparse Top-k and causal masking followed by exactly one final compressed-axis L1 normalization, plus preservation of the current dense support/masking policy;
• equivalent single-segment SBHD and packed-THD targets;
• custom-backward versus autograd-reference parity for the indexer Q/weights/K gradients, with no teacher-side gradients.

For the same inputs and parameter state, this change would not alter the current Top-k selection or CSA forward output. It would change only the detached auxiliary teacher, the indexer loss, and the resulting indexer gradients.

This first change would cover only the two non-CP callers of FusedDSAIndexerLoss. The separate THD+CP fallback _unfused_indexer_sparse_attn_from_topk, which already includes the sink but still omits the window contribution from its teacher denominator, would remain an explicitly documented follow-up. The FlashMLA/fused path and its ABI would also remain unchanged. I therefore would not claim that this initial scoped PR fully closes #5776.

If this matches the intended CSA training semantics, I can start with the failing oracle and then validate the smallest denominator-correct interface.

[JavaZeroo]

Thanks, @mehraakash. Here is the concrete scope I have in mind.

First, a correction to my earlier wording: a window/sink perturbation does not necessarily change the final L1-normalized target. The denominator cancels in the single-head case. A valid regression oracle therefore needs at least two attention heads and two causally valid compressed target entries, with different compressed-key conditional distributions across the heads. Under those conditions, changing the window or sink mass of one head has a provable effect on the final teacher.

For the first scoped change, I propose to cover the two non-CP unfused CSA callers of FusedDSAIndexerLoss: _forward_unfused_csa and _forward_unfused_csa_thd.

For each query and local attention head, the caller would compute a detached FP32 log-mass for the non-compressed part of the real CSA denominator:

u h = log ⁡ ( exp ⁡ ( a h ) + ∑ r ∈ W valid exp ⁡ ( w h , r ) ) .

This computation must reuse the same attention query/KV values, softmax_scale, window indices, invalid-index masking, and sink semantics as unfused_compressed_sparse_attn. In particular, the packed THD path must use the same per-segment/globalized window indexing as the actual attention path.

The narrowest interface I plan to validate is an optional CSA-only non_compressed_lse input to the naive target computation behind FusedDSAIndexerLoss. After compressed Top-k selection, the target would use

ℓ h = log ⁡ ( exp ⁡ ( u h ) + ∑ j ∈ C exp ⁡ ( c h , j ) ) , p h , j = exp ⁡ ( c h , j − ℓ h ) .

For sparse loss, C is the selected Top-k intersected with the causally valid compressed entries used by the CSA core attention. For dense-loss mode, C is the set of all causally visible compressed entries, preserving the current dense support and masking policy.

The existing sum across local attention heads, TP all-reduce, and single final L1 normalization over the compressed support would remain unchanged. The indexer-score softmax and KL definition would also remain unchanged. Non-CSA callers would provide no additional denominator term and retain their current behavior.

All teacher-side inputs, including the non-compressed LSE, would remain detached. The custom backward would reconstruct the same corrected target and continue to produce gradients only for the indexer Q, weights, and K. It would return no gradients for the main-attention query/KV, window contribution, or sink.

The focused tests would cover:

• an explicit FP32 two-head/two-entry oracle with different per-head compressed-key preferences, where perturbing one head’s sink changes the teacher;
• the equivalent window-mass perturbation;
• exact construction of the window-plus-sink LSE from the same SBHD and packed-THD indexing used by the reference attention;
• sparse Top-k and causal masking followed by exactly one final compressed-axis L1 normalization, plus preservation of the current dense support/masking policy;
• equivalent single-segment SBHD and packed-THD targets;
• custom-backward versus autograd-reference parity for the indexer Q/weights/K gradients, with no teacher-side gradients.

For the same inputs and parameter state, this change would not alter the current Top-k selection or CSA forward output. It would change only the detached auxiliary teacher, the indexer loss, and the resulting indexer gradients.

This first change would cover only the two non-CP callers of FusedDSAIndexerLoss. The separate THD+CP fallback _unfused_indexer_sparse_attn_from_topk, which already includes the sink but still omits the window contribution from its teacher denominator, would remain an explicitly documented follow-up. The FlashMLA/fused path and its ABI would also remain unchanged. I therefore would not claim that this initial scoped PR fully closes #5776.

If this matches the intended CSA training semantics, I can start with the failing oracle and then validate the smallest denominator-correct interface.

I've implemented a non-fused version #5960 , if you'd like, you can implement the fused version.

[chengcuiping]

@JavaZeroo Thanks for implementing this. I reviewed the current commit 9e81463. The core denominator correction matches the semantics I had in mind: the compressed teacher retains the per-head competition from the sliding window and attention sink, followed by head aggregation and one final normalization over the compressed support.

I do, however, see two implementation-contract issues that look important to resolve before treating this as a drop-in implementation of the scoped contract.

1. Autograd and peak-memory behavior

The two non-CP callers now bypass FusedDSAIndexerLoss and run fused_qk_topk_naive under ordinary autograd. This records the full per-indexer-head score computation in the graph. The relevant intermediate is approximately [Sq, B, H_index, Sk_compressed]. With B=1, 64 indexer heads, and a compression ratio of 4, one such FP32 tensor is already about 1 GiB at sequence length 4096 and 4 GiB at sequence length 8192, before additional saved tensors and multiple layers are considered.

The existing custom autograd function avoids retaining this internal score graph and recomputes the required values during backward. The new implementation therefore appears to introduce a potentially serious training-memory regression even though the teacher formula itself is correct.

The optional detached non_compressed_lse interface in my earlier proposal was intended to preserve that recomputation contract. The caller would provide only the FP32 window-plus-sink sufficient statistic, while FusedDSAIndexerLoss would reconstruct the corrected target during both forward and backward, save only the tensors required for recomputation, and continue to return gradients only for indexer Q, weights, and K.

I think extending the custom function in this way would be safer than replacing it with an ordinary autograd graph. A representative GPU peak-memory regression test would also be valuable.

2. KL numerical contract

The new helper changes the predictor-side KL stabilization from the existing log(p + 1e-10) behavior to log(p.clamp_min(torch.finfo(torch.float32).tiny)). The asserted loss for a predictor-logit gap of 40 explicitly exercises and locks in this changed behavior.

Using tiny for the target-mass normalization may be appropriate to avoid losing an extremely small but nonzero compressed mass. However, changing the predictor-side stabilization is independent of the denominator correction and materially changes tail losses and gradients.

For a denominator-only correctness fix, I think the existing indexer softmax and KL behavior should be preserved. If the predictor-side stabilization also needs changing, it would be clearer to justify and test that as a separate numerical change with both forward and gradient comparisons.

Scope and reduction clarifications

A few additional points seem worth clarifying:

• In the padded-THD fixture, the real segment losses for lengths [9, 11] are weighted by those real lengths, but the final result is divided by the physical total_q=24. Thus, in mean mode, padding rows contribute zero to the numerator but remain part of the physical-row divisor. Is that the intended reduction contract? If so, the PR description should clarify what “excluded from loss computation” means, and the tests should encode the expected padded/unpadded scaling under both reduction modes. Changing the divisor to the valid-row count would be a separate reduction-contract change.

• The THD path now converts w_thd to FP32 before applying indexer.softmax_scale, whereas the previous path scaled the original dtype. Near a Top-k boundary, this can change selection. The padded-THD correction can also intentionally change previously invalid selections, so the statement that Top-k and the forward output are unchanged should be qualified or protected by parity tests.

• The PR additionally modifies the THD+CP fallback while CP+TP still lacks the target all-reduce across TP-sharded attention heads. Since that path has a separate distributed-reduction contract, it may be safer to keep the first correction limited to the two non-CP callers and handle CP in a separately validated follow-up, or explicitly limit the CP correctness claim to TP=1.

For a memory-safe custom implementation, I would plan to add checked-in tests covering:

• custom-backward versus autograd-reference parity for indexer Q/weights/K gradients;
• no auxiliary-loss gradients to the main-attention query/KV or sink, and no gradient returned for non_compressed_lse;
• an independent SBHD/THD teacher or loss oracle rather than having the THD wrapper oracle call the same target helper;
• padded and unpadded behavior under both loss-reduction modes, according to the documented divisor contract;
• representative GPU peak memory.

@mehraakash, if this direction matches the intended implementation contract, I would be happy to prepare the memory-safe non_compressed_lse extension, backward-parity tests, and an A100 peak-memory regression test as signed-off commits on my fork for cherry-picking into this PR. If the maintainers prefer a separate narrow PR instead, I can follow that workflow. I will avoid opening an overlapping PR without explicit direction.

I would also prefer to settle the unfused interface and backward contract before starting the fused FlashMLA work, since the fused path has a separate LSE/ABI question and cannot be runtime-validated on my current SM80 environment.