@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 41cfa50632
Fork 0

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