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 15, 08:30 UTC
[BugFix] FA2 MLA Accuracy Issue (#18807)
Signed-off-by: LucasWilkinson <lwilkinson@neuralmagic.com>
Lucas Wilkinson committed
May 28, 2025 at 04:59 UTC
ce75efeecb57acb5421aeb545a95e922f3dc8b3e
3 files changed
+16
-8
csrc/attention/merge_attn_states.cu
+8
@@ -143,6 +143,14 @@ void merge_attn_states_launcher(torch::Tensor& output,
143
const uint pack_size = 16 / sizeof(scalar_t);
144
TORCH_CHECK(head_size % pack_size == 0,
145
"headsize must be multiple of pack_size:", pack_size);
146
+ TORCH_CHECK(output.stride(-2) == head_size && output.stride(-1) == 1,
147
+ "output heads must be contiguous in memory");
148
+ TORCH_CHECK(
149
+ prefix_output.stride(-2) == head_size && prefix_output.stride(-1) == 1,
150
+ "prefix_output heads must be contiguous in memory");
151
+ TORCH_CHECK(
152
+ suffix_output.stride(-2) == head_size && suffix_output.stride(-1) == 1,
153
+ "suffix_output heads must be contiguous in memory");
154
float* output_lse_ptr = nullptr;
155
if (output_lse.has_value()) {
156
output_lse_ptr = output_lse.value().data_ptr<float>();
vllm/attention/backends/mla/common.py
+4
-4
@@ -1093,10 +1093,6 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1093
if isinstance(attn_out, tuple):
1094
attn_out, *rest = attn_out
1095
1096
- # unpad if necessary
1097
- if self._pad_v:
1098
- attn_out = attn_out[..., :v.shape[-1]]
1099
-
1096
# Remain consistent with old `flash_attn_varlen_func` where there
1097
# is only one output tensor if `return_softmax_lse` is False.
1098
if return_softmax_lse:
@@ -1294,6 +1290,10 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1290
suffix_lse=suffix_lse,
1291
)
1292
1293
+ # unpad if necessary
1294
+ if self._pad_v:
1295
+ output = output[..., :v.shape[-1]]
1296
+
1297
return output.flatten(start_dim=-2)
1298
1299
@abstractmethod
vllm/v1/attention/backends/mla/common.py
+4
-4
@@ -653,10 +653,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
653
if isinstance(attn_out, tuple):
654
attn_out, lse = attn_out[0], attn_out[1]
655
656
- # unpad if necessary
657
- if self._pad_v:
658
- attn_out = attn_out[..., :v.shape[-1]]
659
-
656
# Remain consistent with old `flash_attn_varlen_func` where there
657
# is only one output tensor if `return_softmax_lse` is False.
658
if return_softmax_lse:
@@ -839,6 +835,10 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
835
suffix_lse=suffix_lse,
836
)
837
838
+ # unpad if necessary
839
+ if self._pad_v:
840
+ output = output[..., :v.shape[-1]]
841
+
842
return output.flatten(start_dim=-2)
843
844
@abstractmethod