[ROCm][AITER] fix wrong argument passed to AITER flash_attn_varlen_func

[vllmellm]

Purpose

This #31465 PR changed arguments name into wrong one that doesn't exist in aiter package.
This PR fixes the issue.

Test Plan

commands:

ROCM_AITER_MLA backend

export VLLM_USE_V1=1 export SAFETENSORS_FAST_GPU=1 export VLLM_ROCM_USE_AITER=1 export VLLM_ATTENTION_BACKEND=ROCM_AITER_MLA vllm serve deepseek-ai/DeepSeek-V3 \ --block-size 128 \ -tp 8 >

ROCM_AITER_TRITON_MLA
export VLLM_USE_V1=1 export SAFETENSORS_FAST_GPU=1 export VLLM_ROCM_USE_AITER=1 export VLLM_ATTENTION_BACKEND=ROCM_AITER_TRITON_MLA vllm serve deepseek-ai/DeepSeek-V3 \ --block-size 128 \ -tp 8

lm_eval --model local-completions \ --tasks gsm8k \ --model_args model=deepseek-ai/DeepSeek-V3,base_url=http://localhost:8000/v1/completions \ --trust_remote_code \ --num_fewshot 5 \ --limit 300 \ --batch_size 128

Test Result

ROCM_AITER_MLA backend

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9433	±	0.0134
		strict-match	5	exact_match	↑	0.9400	±	0.0137

ROCM_AITER_TRITON_MLA

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9533	±	0.0122
		strict-match	5	exact_match	↑	0.9533	±	0.0122

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request corrects a critical bug in the AITER attention backends for ROCm. The argument return_softmax_lse was being incorrectly passed to flash_attn_varlen_func from the aiter package. This has been fixed by changing it to the correct argument name, return_lse. This change prevents a runtime TypeError and is essential for the correct functioning of the AITER backends. The fix is applied in both aiter_triton_mla.py and rocm_aiter_mla.py. The changes are correct and necessary.

[tjtanaa]

LGTM