[ROCm][AITER] fix wrong argument passed to AITER `flash_attn_varlen_func` (#31880)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

vllm/v1/attention/backends/mla/aiter_triton_mla.py

@@ -54,7 +54,7 @@ class AiterTritonMLAImpl(AiterMLAImpl):
             k,
             v,
             softmax_scale=softmax_scale,
-            return_softmax_lse=return_softmax_lse,
+            return_lse=return_softmax_lse,
             **kwargs,
         )
         # Transpose the LSE if Triton MHA is used:

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -236,7 +236,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
             k=k,
             v=v,
             softmax_scale=softmax_scale,
-            return_softmax_lse=return_softmax_lse,
+            return_lse=return_softmax_lse,
             **kwargs,
         )