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 7, 02:45 UTC
[ROCm][CI] Fix tool use test stability - disable skinny GEMM, prefix caching, eliminate batch variance (#35553)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 6, 2026 at 01:15 UTC
807d6803376ff8610efbf9da23f772a5dbd7b5ea
5 files changed
+33
-17
docs/design/attention_backends.md
+1
-1
@@ -171,7 +171,7 @@ Priority is **1 = highest** (tried first).
171
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
172
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
173
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
174
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
174
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
175
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
176
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 544 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
177
| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
requirements/rocm-test.txt
+2
@@ -108,3 +108,5 @@ bitsandbytes==0.49.2
108
tensorizer==2.10.1
109
# Multi-modal models test (`allendou/FireRedASR2-LLM-vllm`)
110
kaldi-native-fbank==1.22.3
111
+# Pinning numpy version
112
+numpy==2.2.6
tests/entrypoints/openai/test_completion_with_function_calling.py
+8
-16
@@ -9,14 +9,13 @@ import openai # use the official client for correctness check
9
import pytest
10
import pytest_asyncio
11
12
-from vllm.platforms import current_platform
13
-
12
# downloading lora to test lora requests
15
-from ...utils import RemoteOpenAIServer
13
+from ...utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
14
15
# any model with a chat template should work here
16
MODEL_NAME = "Qwen/Qwen3-0.6B"
17
18
+
19
tools = [
20
{
21
"type": "function",
@@ -142,19 +141,11 @@ def server():
141
"--gpu-memory-utilization",
142
"0.4",
143
"--enforce-eager",
145
- ]
144
+ ] + ROCM_EXTRA_ARGS
145
147
- rocm_args = {
148
- "--max-num-seqs": "1",
149
- "--no-enable-prefix-caching": None,
150
- }
151
- if current_platform.is_rocm():
152
- for k, v in rocm_args.items():
153
- args.append(k)
154
- if v is not None:
155
- args.append(v)
156
-
157
- with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
146
+ with RemoteOpenAIServer(
147
+ MODEL_NAME, args, env_dict=ROCM_ENV_OVERRIDES
148
+ ) as remote_server:
149
yield remote_server
150
151
@@ -239,12 +230,13 @@ def k2_server():
230
"qwen3",
231
"--gpu-memory-utilization",
232
"0.4",
242
- ]
233
+ ] + ROCM_EXTRA_ARGS
234
# hack to test kimi_k2 tool use tool_id format.
235
# avoid error in is_deepseek_mla check by setting kv_lora_rank=null
236
with RemoteOpenAIServer(
237
MODEL_NAME,
238
args,
239
+ env_dict=ROCM_ENV_OVERRIDES,
240
override_hf_configs={"model_type": "kimi_k2", "kv_lora_rank": None},
241
) as remote_server:
242
yield remote_server
tests/utils.py
+14
@@ -109,6 +109,20 @@ else:
109
VLLM_PATH = Path(__file__).parent.parent
110
"""Path to root of the vLLM repository."""
111
112
+# ROCm: disable skinny GEMM to avoid non-deterministic results from
113
+# atomic reductions in wvSplitKrc kernel.
114
+# See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
115
+ROCM_ENV_OVERRIDES = (
116
+ {"VLLM_ROCM_USE_SKINNY_GEMM": "0"} if current_platform.is_rocm() else {}
117
+)
118
+# ROCm: disable prefix caching and eliminate batch variance to reduce
119
+# test flakiness.
120
+ROCM_EXTRA_ARGS = (
121
+ ["--no-enable-prefix-caching", "--max-num-seqs", "1"]
122
+ if current_platform.is_rocm()
123
+ else []
124
+)
125
+
126
127
class RemoteVLLMServer:
128
"""Base class for launching vLLM server subprocesses for testing.
vllm/v1/attention/backends/rocm_aiter_fa.py
+8
@@ -741,6 +741,14 @@ class AiterFlashAttentionBackend(AttentionBackend):
741
"fp8_e5m2",
742
]
743
744
+ @classmethod
745
+ def supports_attn_type(cls, attn_type: str) -> bool:
746
+ """ROCM AITER FA supports decoder and encoder-decoder (cross) attention."""
747
+ return attn_type in (
748
+ AttentionType.DECODER,
749
+ AttentionType.ENCODER_DECODER,
750
+ )
751
+
752
@staticmethod
753
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
754
return [16, 32]