A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 09:45 UTC
bugfix: correct attn output with base 2 or e (#28840)
Signed-off-by: augusto.yjh <augusto.yjh@antgroup.com>
Augusto Yao committed
Nov 29, 2025 at 07:52 UTC
9726e64530e391a4d617cf7ecb862b433698ff6d
3 files changed
+47
-13
vllm/attention/ops/common.py
+32
-10
@@ -21,6 +21,7 @@ def _correct_attn_cp_out_kernel(
21
lse_idx,
22
HEAD_DIM: tl.constexpr,
23
N_ROUNDED: tl.constexpr,
24
+ IS_BASE_E: tl.constexpr,
25
):
26
"""
27
Apply the all-gathered lses to correct each local rank's attention
@@ -55,9 +56,14 @@ def _correct_attn_cp_out_kernel(
56
lse_max = tl.max(lse, axis=0)
57
lse_max = tl.where(lse_max == -float("inf"), 0, lse_max)
58
lse -= lse_max
58
- lse_exp = tl.exp(lse)
59
- lse_acc = tl.sum(lse_exp, axis=0)
60
- lse = tl.log(lse_acc)
59
+ if IS_BASE_E:
60
+ lse_exp = tl.exp(lse)
61
+ lse_acc = tl.sum(lse_exp, axis=0)
62
+ lse = tl.log(lse_acc)
63
+ else:
64
+ lse_exp = tl.exp2(lse)
65
+ lse_acc = tl.sum(lse_exp, axis=0)
66
+ lse = tl.log2(lse_acc)
67
lse += lse_max
68
69
lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H
@@ -81,7 +87,7 @@ def _correct_attn_cp_out_kernel(
87
-float("inf"),
88
lse_finally,
89
)
84
- factor = tl.exp(lse_finally)
90
+ factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally)
91
output = tl.load(outputs_ptr + output_offsets)
92
output = output * factor
93
@@ -102,7 +108,11 @@ class CPTritonContext:
108
109
110
def correct_attn_out(
105
- out: torch.Tensor, lses: torch.Tensor, cp_rank: int, ctx: CPTritonContext
111
+ out: torch.Tensor,
112
+ lses: torch.Tensor,
113
+ cp_rank: int,
114
+ ctx: CPTritonContext,
115
+ is_lse_base_on_e: bool = True,
116
) -> tuple[torch.Tensor, torch.Tensor]:
117
"""Correct the attention output using the all-gathered lses.
118
@@ -163,8 +173,7 @@ def correct_attn_out(
173
l_sH,
174
cp_rank,
175
)
166
- const_args = {"HEAD_DIM": D, "N_ROUNDED": N}
167
-
176
+ const_args = {"HEAD_DIM": D, "N_ROUNDED": N, "IS_BASE_E": is_lse_base_on_e}
177
ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args)
178
return out, lse
179
@@ -174,6 +183,7 @@ def _cp_lse_common(
183
cp_attn_lse: torch.Tensor,
184
cp_group: GroupCoordinator,
185
ctx: CPTritonContext | None = None,
186
+ is_lse_base_on_e=True,
187
):
188
"""
189
cp_attn_out: [ B, H, D ]
@@ -193,7 +203,13 @@ def _cp_lse_common(
203
204
cp_attn_lse = cp_attn_lse.contiguous()
205
lses = cp_group.all_gather(cp_attn_lse, dim=0).view_as(lses)
196
- out, lse = correct_attn_out(cp_attn_out, lses, cp_group.rank_in_group, ctx)
206
+ out, lse = correct_attn_out(
207
+ cp_attn_out,
208
+ lses,
209
+ cp_group.rank_in_group,
210
+ ctx,
211
+ is_lse_base_on_e=is_lse_base_on_e,
212
+ )
213
return out, lse
214
215
@@ -203,12 +219,15 @@ def cp_lse_ag_out_rs(
219
cp_group: GroupCoordinator,
220
ctx: CPTritonContext | None = None,
221
return_lse: bool = False,
222
+ is_lse_base_on_e=True,
223
):
224
"""
225
cp_attn_out: [ B, H, D ]
226
cp_attn_lse: [ B, H ]
227
"""
211
- out, lse = _cp_lse_common(cp_attn_out, cp_attn_lse, cp_group, ctx=ctx)
228
+ out, lse = _cp_lse_common(
229
+ cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e
230
+ )
231
out = cp_group.reduce_scatter(out, dim=1)
232
233
if return_lse:
@@ -225,12 +244,15 @@ def cp_lse_ag_out_ar(
244
cp_group: GroupCoordinator,
245
ctx: CPTritonContext | None = None,
246
return_lse: bool = False,
247
+ is_lse_base_on_e=True,
248
):
249
"""
250
cp_attn_out: [ B, H, D ]
251
cp_attn_lse: [ B, H ]
252
"""
233
- out, lse = _cp_lse_common(cp_attn_out, cp_attn_lse, cp_group, ctx=ctx)
253
+ out, lse = _cp_lse_common(
254
+ cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e
255
+ )
256
out = cp_group.all_reduce(out)
257
258
if return_lse:
vllm/v1/attention/backends/flashinfer.py
+9
-2
@@ -249,7 +249,11 @@ class BatchDCPPrefillWrapper:
249
return_lse=True,
250
)
251
output_context, lse_context = cp_lse_ag_out_rs(
252
- output_context_tmp, lse_context_tmp, get_dcp_group(), return_lse=True
252
+ output_context_tmp,
253
+ lse_context_tmp,
254
+ get_dcp_group(),
255
+ return_lse=True,
256
+ is_lse_base_on_e=False,
257
)
258
lse_context = lse_context.transpose(0, 1).contiguous()
259
@@ -1335,7 +1339,10 @@ class FlashInferImpl(AttentionImpl):
1339
return_lse=True,
1340
)
1341
output[:num_decode_tokens] = cp_lse_ag_out_rs(
1338
- output_tmp, lse, get_dcp_group()
1342
+ output_tmp,
1343
+ lse,
1344
+ get_dcp_group(),
1345
+ is_lse_base_on_e=False,
1346
)
1347
else:
1348
decode_wrapper.run(
vllm/v1/attention/backends/mla/common.py
+6
-1
@@ -2057,7 +2057,12 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
2057
2058
# correct dcp attn_out with lse.
2059
if self.dcp_world_size > 1:
2060
- attn_out = cp_lse_ag_out_rs(attn_out, lse, get_dcp_group())
2060
+ attn_out = cp_lse_ag_out_rs(
2061
+ attn_out,
2062
+ lse,
2063
+ get_dcp_group(),
2064
+ is_lse_base_on_e=not self._use_fi_prefill,
2065
+ )
2066
2067
# v_up projection
2068
self._v_up_proj(attn_out, out=output[:num_decode_tokens])