diff --git a/python/sglang/kernels/ops/attention/dcp_kernels.py b/python/sglang/kernels/ops/attention/dcp_kernels.py index 44cce302c..f65c1777c 100644 --- a/python/sglang/kernels/ops/attention/dcp_kernels.py +++ b/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 diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index 518801709..952ff94a1 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/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) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index da8ee0793..07d1502f1 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/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) diff --git a/test/registered/kernels/test_dcp_lse_combine.py b/test/registered/kernels/test_dcp_lse_combine.py index d50f46184..3c6eb5086 100644 --- a/test/registered/kernels/test_dcp_lse_combine.py +++ b/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