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 14, 05:30 UTC
[BugFix] Fix FI accuracy issue when used for MLA prefill (#26063)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Lucas Wilkinson committed
Oct 2, 2025 at 13:18 UTC
decf7f794bfff2dd591f79e591c2e5292fe9fe0f
1 file changed
+9
-2
vllm/v1/attention/backends/mla/common.py
+9
-2
@@ -1211,13 +1211,18 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1211
k, v, return_softmax_lse):
1212
assert isinstance(prefill, FlashInferPrefillMetadata)
1213
assert prefill.prefill_main is not None
1214
- return prefill.prefill_main.run(
1214
+ ret = prefill.prefill_main.run(
1215
q=q,
1216
k=k,
1217
v=v,
1218
return_lse=return_softmax_lse,
1219
)
1220
1221
+ if isinstance(ret, tuple):
1222
+ # Convert from (q_len, num_heads) to (num_heads, q_len)
1223
+ return ret[0], ret[1].transpose(0, 1).contiguous()
1224
+ return ret
1225
+
1226
def _run_prefill_new_tokens_cudnn(self, prefill: MLACommonPrefillMetadata,
1227
q, k, v, return_softmax_lse):
1228
assert isinstance(prefill, CudnnPrefillMetadata)
@@ -1260,12 +1265,14 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1265
def _run_prefill_context_chunk_fi(self, prefill: MLACommonPrefillMetadata,
1266
chunk_idx: int, q, k, v):
1267
assert isinstance(prefill, FlashInferPrefillMetadata)
1263
- return prefill.prefill_chunks[chunk_idx].run(
1268
+ attn_out, lse = prefill.prefill_chunks[chunk_idx].run(
1269
q=q,
1270
k=k,
1271
v=v,
1272
return_lse=True,
1273
)
1274
+ # Convert from (q_len, num_heads) to (num_heads, q_len)
1275
+ return attn_out, lse.transpose(0, 1).contiguous()
1276
1277
def _run_prefill_context_chunk_cudnn(self,
1278
prefill: MLACommonPrefillMetadata,