[Bugfix][Attention] Normalize FlashInfer prefill LSE before merging (#52796)
Signed-off-by: yimdev <5779256+yimdev@users.noreply.github.com> Co-authored-by: usberkeley <150880684+usberkeley@users.noreply.github.com>
yimdev committed
Aug 21, 2026 at 23:31 UTC
574e6a00efd3d594c4d7f72e766c3e605316a118
5 files changed
+23
-10
tests/v1/attention/test_mla_backends.py
+6
-4
@@ -1662,6 +1662,12 @@ def _run_backend_correctness(
1662
kv_cache_per_block_size[block_size] = kv_cache
1663
1664
# 4. Run vLLM backends and compare
1665
+ rtol = 1e-2
1666
+ atol = {
1667
+ "auto": 1e-2,
1668
+ "fp8": 1.5e-1,
1669
+ "fp8_e4m3": 1.5e-1,
1670
+ }[kv_cache_dtype]
1671
failures = []
1672
for backend_idx, backend_name in enumerate(backends_to_test):
1673
# Skip backends that don't support spec decode for spec decode tests
@@ -1726,10 +1732,6 @@ def _run_backend_correctness(
1732
f"[{backend_name}] produced non-finite values"
1733
)
1734
1729
- # Check numerical similarity
1730
- rtol = 1e-2
1731
- atol = 5e-1
1732
-
1735
max_diff = torch.max(torch.abs(backend_output - expected_output)).item()
1736
max_rel_diff = torch.max(
1737
torch.abs(backend_output - expected_output) / torch.abs(expected_output)
vllm/v1/attention/backends/flashinfer.py
+3
-2
@@ -73,6 +73,7 @@ from vllm.v1.attention.backends.utils import (
73
get_num_attention_heads_from_layers,
74
get_per_layer_parameters,
75
infer_global_hyperparameters,
76
+ log2_lse_to_ln,
77
split_decodes_and_prefills,
78
)
79
from vllm.v1.attention.ops.dcp import (
@@ -370,7 +371,7 @@ class BatchDCPPrefillWrapper:
371
get_dcp_group(),
372
return_lse=True,
373
)
373
- lse_context = lse_context.transpose(0, 1).contiguous()
374
+ lse_context = log2_lse_to_ln(lse_context.transpose(0, 1).contiguous())
375
376
output_query, lse_query = self._new_tokens.run(
377
prefill_query,
@@ -378,7 +379,7 @@ class BatchDCPPrefillWrapper:
379
value,
380
return_lse=True,
381
)
381
- lse_query = lse_query.transpose(0, 1).contiguous()
382
+ lse_query = log2_lse_to_ln(lse_query.transpose(0, 1).contiguous())
383
384
merge_attn_states(
385
out,
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+3
-2
@@ -15,6 +15,7 @@ from vllm.v1.attention.backends.utils import (
15
PerLayerParameters,
16
get_per_layer_parameters,
17
infer_global_hyperparameters,
18
+ log2_lse_to_ln,
19
)
20
from vllm.v1.worker.workspace import current_workspace_manager
21
@@ -219,7 +220,7 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
220
221
if isinstance(ret, tuple):
222
# Convert from (q_len, num_heads) to (num_heads, q_len)
222
- return ret[0], ret[1].transpose(0, 1)
223
+ return ret[0], log2_lse_to_ln(ret[1].transpose(0, 1))
224
return ret
225
226
def run_prefill_context_chunk(
@@ -239,4 +240,4 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
240
)
241
242
# Convert from (q_len, num_heads) to (num_heads, q_len)
242
- return attn_out, lse.transpose(0, 1)
243
+ return attn_out, log2_lse_to_ln(lse.transpose(0, 1))
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+3
-2
@@ -11,6 +11,7 @@ from vllm.v1.attention.backends.mla.prefill.base import (
11
MLADimensions,
12
MLAPrefillBackend,
13
)
14
+from vllm.v1.attention.backends.utils import log2_lse_to_ln
15
from vllm.v1.worker.workspace import current_workspace_manager
16
17
if TYPE_CHECKING:
@@ -138,7 +139,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
139
140
if isinstance(ret, tuple):
141
# Convert from (q_len, num_heads) to (num_heads, q_len)
141
- return ret[0], ret[1].transpose(0, 1)
142
+ return ret[0], log2_lse_to_ln(ret[1].transpose(0, 1))
143
return ret
144
145
def run_prefill_context_chunk(
@@ -182,4 +183,4 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
183
)
184
185
# Convert from (q_len, num_heads) to (num_heads, q_len)
185
- return attn_out, lse.transpose(0, 1)
186
+ return attn_out, log2_lse_to_ln(lse.transpose(0, 1))
vllm/v1/attention/backends/utils.py
+8
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import functools
4
+import math
5
from collections.abc import Callable
6
from dataclasses import dataclass, field, fields, make_dataclass
7
from typing import (
@@ -45,6 +46,13 @@ _KV_CACHE_LAYOUT_OVERRIDE: KVCacheLayoutType | None = None
46
PAD_SLOT_ID = -1
47
NULL_BLOCK_ID = 0
48
49
+_LN_2 = math.log(2.0)
50
+
51
+
52
+def log2_lse_to_ln(lse: torch.Tensor) -> torch.Tensor:
53
+ """Convert a base-2 log-sum-exp tensor to natural-log units."""
54
+ return lse * _LN_2
55
+
56
57
def compute_mm_prefix_range_tensor(
58
mm_prefix_range: dict[int, list[tuple[int, int]]] | None,