[Fix] Pass natural-log LSE to the cuDNN sparse MLA backward (#2147)

b03bf01/fe0037f0 switched _cudnn_dsa_sparse_mla_backward_op from the
log2-space LSE contract to natural-log LSE, but this backend kept
forwarding the log2 value, so every GLM-5.3 training silently fed the
cuDNN backward a 1.4427x-inflated LSE. Forward is unaffected; backward
gradients come out systematically shrunk (grad_norm ~20-30% low) and
loss drifts (+0.043 over 300 steps vs the pre-change baseline).
Save and pass the natural-log softmax_lse instead.

Co-authored-by: liutong <liutong@pjlab.org.cn>

Author: Denny991

xtuner/v1/ops/sparse_mla/flash_mla_cudnn.py

@@ -69,8 +69,9 @@ def _flash_mla_cudnn_forward(
         scaling,
         d_v=value_dim,
     )
-    # FlashMLA returns natural-log LSE; cuDNN's backward (below) expects log2-space LSE,
-    # matching the same contract flash_mla.py uses for its (different) TileLang backward.
+    # FlashMLA returns natural-log LSE. cuDNN's backward (below) also expects natural-log
+    # LSE since b03bf018 switched _cudnn_dsa_sparse_mla_backward_op off the log2 contract;
+    # the third slot keeps the legacy log2 value only for custom-op schema stability.
     return raw_output, softmax_lse, softmax_lse * 1.4426950408889634
 
 
@@ -84,20 +85,20 @@ def _(q: Tensor, kv: Tensor, indices: Tensor, scaling: float, value_dim: int) ->
 
 def _setup_flash_mla_cudnn_context(ctx, inputs, output) -> None:
     q, kv, indices, scaling, _ = inputs
-    raw_output, _, lse_log2 = output
+    raw_output, softmax_lse, _ = output
     ctx.scaling = scaling
-    ctx.save_for_backward(q, kv, indices, raw_output, lse_log2)
+    ctx.save_for_backward(q, kv, indices, raw_output, softmax_lse)
 
 
 def _flash_mla_cudnn_backward(ctx, grad_output: Tensor, grad_lse: Tensor, grad_lse_log2: Tensor):
-    q, kv, indices, raw_output, lse_log2 = ctx.saved_tensors
+    q, kv, indices, raw_output, softmax_lse = ctx.saved_tensors
     dq, dkv = _cudnn_dsa_sparse_mla_backward_op(
         q,
         kv,
         raw_output,
         grad_output.contiguous(),
         indices,
-        lse_log2,
+        softmax_lse,
         ctx.scaling,
     )
     return dq, dkv, None, None, None