[Bugfix][MLA] Fix LSE log-base mismatch in DCP + FlashInfer MLA decode (#47079)
Signed-off-by: girasoley <girasoleyang@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Summer Yang committed
Jun 29, 2026 at 19:15 UTC
5b4cb6952310ff20e054818eb34b8e70d3c06a1e
3 files changed
+20
-2
vllm/model_executor/layers/attention/mla_attention.py
+2
-2
@@ -818,14 +818,14 @@ class MLAAttention(nn.Module, AttentionLayerBase):
818
attn_out,
819
lse,
820
get_dcp_group(),
821
- is_lse_base_on_e=True,
821
+ is_lse_base_on_e=self.impl.lse_base_on_e,
822
)
823
else:
824
attn_out = cp_lse_ag_out_rs(
825
attn_out,
826
lse,
827
get_dcp_group(),
828
- is_lse_base_on_e=True,
828
+ is_lse_base_on_e=self.impl.lse_base_on_e,
829
)
830
831
# v_up projection
vllm/v1/attention/backend.py
+11
@@ -756,6 +756,17 @@ class AttentionImplBase(ABC, Generic[T]):
756
# Some features like decode context parallelism require the softmax lse.
757
can_return_lse_for_decode: bool = False
758
759
+ # Base of the logarithm used by this backend when returning softmax lse.
760
+ # True => natural log (lse = ln(sum(exp(qk))))
761
+ # -- e.g. Triton MLA, FlashAttention, FlashMLA, Cutlass MLA
762
+ # False => base 2 (lse = log2(sum(exp(qk))))
763
+ # -- e.g. FlashInfer trtllm-gen MLA
764
+ # The DCP combine kernel (cp_lse_ag_out_rs / dcp_a2a_lse_reduce in
765
+ # vllm/v1/attention/ops/common.py) branches on this via its IS_BASE_E
766
+ # constexpr; getting it wrong silently corrupts the cross-shard
767
+ # softmax denominator.
768
+ lse_base_on_e: bool = True
769
+
770
# Whether the attention impl supports Prefill Context Parallelism.
771
supports_pcp: bool = False
772
# Whether the attention impl(or ops) supports MTP
vllm/v1/attention/backends/mla/flashinfer_mla.py
+7
@@ -114,6 +114,13 @@ g_fi_workspace = torch.zeros(
114
115
class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
116
can_return_lse_for_decode: bool = True
117
+ # trtllm-gen MLA decode emits LSE in log2 (per flashinfer's own
118
+ # reference at flashinfer/trace/templates/attention.py:81:
119
+ # `logsumexp / log(2.0)`). Override the AttentionImplBase default
120
+ # so MLAAttention's DCP combine branches on the correct base
121
+ # (IS_BASE_E=False uses tl.exp2/tl.log2 natively, avoiding an FP
122
+ # multiply per decode step).
123
+ lse_base_on_e: bool = False
124
125
def __init__(
126
self,