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 10, 14:00 UTC
[ROCm][AITER] fix wrong argument passed to AITER `flash_attn_varlen_func` (#31880)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
Jan 7, 2026 at 12:25 UTC
41cfa50632c26c6064cabbbc43c9fc29c7792a2d
2 files changed
+2
-2
vllm/v1/attention/backends/mla/aiter_triton_mla.py
+1
-1
@@ -54,7 +54,7 @@ class AiterTritonMLAImpl(AiterMLAImpl):
54
k,
55
v,
56
softmax_scale=softmax_scale,
57
- return_softmax_lse=return_softmax_lse,
57
+ return_lse=return_softmax_lse,
58
**kwargs,
59
)
60
# Transpose the LSE if Triton MHA is used:
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1
-1
@@ -236,7 +236,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
236
k=k,
237
v=v,
238
softmax_scale=softmax_scale,
239
- return_softmax_lse=return_softmax_lse,
239
+ return_lse=return_softmax_lse,
240
**kwargs,
241
)
242