[Hardware][AMD][CI] Use Triton-based AITER MHA for LM Eval Qwen-3.5 Models Tests (#46691)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Jun 25, 2026 at 17:08 UTC
27da2a2ac4776faa4265cba38ba86dd3a7119c4f
3 files changed
+9
-6
.buildkite/test-amd.yaml
+3
-4
@@ -2790,13 +2790,12 @@ steps:
2790
- uv pip install --system 'gpt-oss[eval]==0.0.5'
2791
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt
2792
2793
-- label: LM Eval Qwen3-5 Models (B200-MI355) %N # TBD
2794
- timeout_in_minutes: 180
2793
+- label: LM Eval Qwen3-5 Models (B200-MI355) # TBD
2794
+ timeout_in_minutes: 120
2795
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
2796
agent_pool: mi355_2
2797
num_gpus: 2
2798
optional: true
2799
- parallelism: 4
2799
working_dir: "/vllm-workspace/tests"
2800
source_file_dependencies:
2801
- vllm/model_executor/models/qwen3_5.py
@@ -2811,7 +2810,7 @@ steps:
2810
- vllm/_aiter_ops.py
2811
- vllm/platforms/rocm.py
2812
commands:
2814
- - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
2813
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-mi355.txt
2814
2815
- label: LM Eval Small Models (2xB200-2xMI355) # TBD
2816
timeout_in_minutes: 180
tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml
+1
-1
@@ -3,7 +3,6 @@ accuracy_threshold: 0.89
3
tolerance: 0.03
4
num_questions: 1319
5
num_fewshot: 5
6
-startup_max_wait_seconds: 3600
6
server_args: >-
7
--max-model-len 4096
8
--tensor-parallel-size 2
@@ -11,3 +10,4 @@ server_args: >-
10
--moe-backend aiter
11
env:
12
VLLM_ROCM_USE_AITER: "1"
13
+ ENABLE_CK: "0" # Avoid AITER CK-based MHA JIT compilation to save time
vllm/v1/attention/ops/vit_attn_wrappers.py
+5
-1
@@ -12,6 +12,8 @@ latencies by ~7% (see qwen2_5_vl for example usage)
12
To use these ops, you must have a recent version of PyTorch installed (>= 2.4.0)
13
"""
14
15
+from typing import Any
16
+
17
import einops
18
import torch
19
import torch.nn.functional as F
@@ -31,9 +33,11 @@ def flash_attn_maxseqlen_wrapper(
33
cu_seqlens: torch.Tensor | None = None,
34
max_seqlen: torch.Tensor | None = None,
35
) -> torch.Tensor:
34
- kwargs = {}
36
+ kwargs: dict[str, Any] = {}
37
if is_rocm_aiter:
38
from aiter import flash_attn_varlen_func
39
+
40
+ kwargs["window_size"] = (-1, -1)
41
else:
42
from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
43