[Fix] Honor FlashMLA natural-log LSE in DCP reduction (#33065)

python/sglang/kernels/ops/attention/dcp_kernels.py

@@ -202,6 +202,7 @@ def _correct_attn_cp_out_kernel(
     lse_idx,
     HEAD_DIM: tl.constexpr,
     N_ROUNDED: tl.constexpr,
+    IS_LSE_BASE_ON_E: tl.constexpr,
 ):
     """
     Apply the all-gathered lses to correct each local rank's attention
@@ -242,9 +243,9 @@ def _correct_attn_cp_out_kernel(
     lse_max = tl.max(lse, axis=0)
     lse_max = tl.where(lse_max == neg_inf, 0.0, lse_max)
     lse = lse - lse_max
-    lse_exp = tl.exp2(lse)
+    lse_exp = tl.exp(lse) if IS_LSE_BASE_ON_E else tl.exp2(lse)
     lse_acc = tl.sum(lse_exp, axis=0)
-    final_lse = tl.log2(lse_acc) + lse_max
+    final_lse = (tl.log(lse_acc) if IS_LSE_BASE_ON_E else tl.log2(lse_acc)) + lse_max
 
     # Compute correction factor
     lse_offset = lse_idx * lses_stride_N + b_i32 * lses_stride_B + h_i32 * lses_stride_H
@@ -255,7 +256,7 @@ def _correct_attn_cp_out_kernel(
         neg_inf,
         lse_diff,
     )
-    factor = tl.exp2(lse_diff)
+    factor = tl.exp(lse_diff) if IS_LSE_BASE_ON_E else tl.exp2(lse_diff)
 
     # Store final LSE
     tl.store(vlse_ptr + b_i32 * lses_stride_B + h_i32 * lses_stride_H, final_lse)
@@ -298,6 +299,7 @@ def correct_attn_out(
     cp_rank: int,
     ctx: Optional[CPTritonContext],
     new_output: torch.Tensor = None,
+    is_lse_base_on_e: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Correct the attention output using the all-gathered lses.
 
@@ -361,7 +363,11 @@ def correct_attn_out(
         no_sD,
         cp_rank,
     )
-    const_args = {"HEAD_DIM": D, "N_ROUNDED": N}
+    const_args = {
+        "HEAD_DIM": D,
+        "N_ROUNDED": N,
+        "IS_LSE_BASE_ON_E": is_lse_base_on_e,
+    }
 
     ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args)
     return new_output, lse

python/sglang/srt/layers/dcp/comm.py

@@ -113,11 +113,16 @@ def cp_lse_ag_out_rs_mla(
     cp_attn_lse: torch.Tensor,
     cp_group: GroupCoordinator,
     ctx: Optional[CPTritonContext] = None,
+    is_lse_base_on_e: bool = False,
 ):
-    """Merge DCP partial attention outputs via Triton correction (PR #14194).
+    """Merge DCP partial attention outputs with the LSE's actual log base.
 
     cp_attn_out: [ B, H, D ]
     cp_attn_lse: [ B, H ]
+
+    FlashInfer MLA returns base-2 LSE, while FlashMLA returns natural-log LSE.
+    The correction kernel must use the matching exp/log pair or it computes
+    incorrect cross-rank softmax weights.
     """
     if cp_group.world_size == 1:
         return cp_attn_out
@@ -133,7 +138,12 @@ def cp_lse_ag_out_rs_mla(
         cp_attn_lse = cp_attn_lse.to(torch.float32)
     lses = _ag_lse(cp_attn_lse, cp_group)
     out, _ = correct_attn_out(
-        cp_attn_out, lses, cp_group.rank_in_group, ctx, new_output
+        cp_attn_out,
+        lses,
+        cp_group.rank_in_group,
+        ctx,
+        new_output,
+        is_lse_base_on_e=is_lse_base_on_e,
     )
     out = cp_group.reduce_scatter_along_dim(out, dim=0)
     return out.to(cp_attn_out.dtype)

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -111,6 +111,12 @@ def _is_dcp_mla_decode_phase(forward_batch: ForwardBatch) -> bool:
     )
 
 
+def _is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool:
+    # FlashMLA exposes natural-log softmax LSE. FlashInfer MLA and the other
+    # currently supported MLA DCP decode backends expose base-2 LSE.
+    return attention_backend == "flashmla"
+
+
 if _is_cuda:
     from sglang.kernels.ops.gemm import bmm_fp8
 
@@ -1030,19 +1036,22 @@ class DeepseekMLAForwardMixin:
                 self.kv_lora_rank,
             )
             dcp_comm_backend = get_parallel().dcp_comm_backend
+            is_lse_base_on_e = _is_mla_dcp_lse_base_on_e(self.current_attention_backend)
             if dcp_comm_backend in ("a2a", "fi_a2a"):
                 # A2A exchange of head partials + LSE, then local Triton combine.
-                # MLA decode LSE is base-2 (FlashInfer-MLA/FlashMLA) -> base_on_e=False.
                 attn_output = dcp_a2a_lse_reduce(
                     attn_output.contiguous(),
                     lse.contiguous(),
                     get_parallel().dcp_group,
-                    is_lse_base_on_e=False,
+                    is_lse_base_on_e=is_lse_base_on_e,
                     comm_backend=dcp_comm_backend,
                 )
             else:
                 attn_output = cp_lse_ag_out_rs_mla(
-                    attn_output, lse, get_parallel().dcp_group
+                    attn_output,
+                    lse,
+                    get_parallel().dcp_group,
+                    is_lse_base_on_e=is_lse_base_on_e,
                 )
                 attn_output = attn_output.transpose(0, 1)
         attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)

test/registered/kernels/test_dcp_lse_combine.py

@@ -93,6 +93,27 @@ class TestLSECombineTritonVsCPU(CustomTestCase):
     def test_n4_large_head_dim(self):
         self._run_combine_test(N=4, B=8, H_local=8, D=512, is_base_e=True)
 
+    def test_flashmla_natural_log_lse_correction(self):
+        """Natural-log LSEs log(2), log(8) give an 8/10 local weight."""
+        from sglang.kernels.ops.attention.dcp_kernels import correct_attn_out
+
+        local_output = torch.tensor(
+            [[[10.0]]], device=self.device, dtype=torch.bfloat16
+        )
+        lses = torch.log(torch.tensor([[[2.0]], [[8.0]]], device=self.device))
+        corrected, _ = correct_attn_out(
+            local_output,
+            lses,
+            cp_rank=1,
+            ctx=None,
+            new_output=torch.empty((1, 1, 1), device=self.device),
+            is_lse_base_on_e=True,
+        )
+
+        torch.testing.assert_close(
+            corrected.cpu(), torch.tensor([[[8.0]]]), atol=1e-5, rtol=1e-5
+        )
+
 
 class TestLSECombineSingleShard(CustomTestCase):
     """N=1 should return input unchanged."""
@@ -244,6 +265,14 @@ class TestCPUReference(CustomTestCase):
 
         self.assertFalse(torch.allclose(result_e, result_2, atol=1e-3))
 
+    def test_flashmla_selects_natural_log_lse(self):
+        from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
+            _is_mla_dcp_lse_base_on_e,
+        )
+
+        self.assertTrue(_is_mla_dcp_lse_base_on_e("flashmla"))
+        self.assertFalse(_is_mla_dcp_lse_base_on_e("flashinfer_mla"))
+
     def test_nan_lse_handled(self):
         from sglang.kernels.ops.attention.dcp_kernels import _lse_weighted_combine_cpu