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 11, 06:15 UTC
[ROCm][CI] Reduce Flakiness For test_async_scheduling Using ROCM_ATTN With FP32 (#30811)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Dec 17, 2025 at 20:27 UTC
fd8afdf38dad8bf7ccc4e7fcc3d4aaa4d6d9e0d8
2 files changed
+7
-11
tests/v1/e2e/test_async_scheduling.py
+2
-10
@@ -148,7 +148,7 @@ def run_tests(
148
# Use TRITON_ATTN for spec decoding test for consistency
149
attention_config = {"backend": "TRITON_ATTN"}
150
else:
151
- attention_config = {"backend": "ROCM_AITER_FA"}
151
+ attention_config = {"backend": "ROCM_ATTN"}
152
else:
153
attention_config = {"backend": "FLEX_ATTENTION"}
154
@@ -284,14 +284,6 @@ def run_test(
284
print(f"---- TESTING {test_str}: {test_config}")
285
print("-" * 80)
286
287
- # On ROCm: use float16 for first test (ROCM_AITER_FA), but float32 for
288
- # spec decoding test (TRITON_ATTN) for better precision.
289
- # On others: always use float32.
290
- if current_platform.is_rocm() and not is_testing_with_spec_decoding:
291
- dtype = "float16"
292
- else:
293
- dtype = "float32"
294
-
287
with VllmRunner(
288
model,
289
max_model_len=512,
@@ -301,7 +293,7 @@ def run_test(
293
# enforce_eager=True,
294
async_scheduling=async_scheduling,
295
distributed_executor_backend=executor,
304
- dtype=dtype,
296
+ dtype="float32",
297
speculative_config=spec_config,
298
disable_log_stats=False,
299
attention_config=attention_config,
vllm/v1/attention/backends/rocm_attn.py
+5
-1
@@ -152,7 +152,11 @@ class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadat
152
153
class RocmAttentionBackend(AttentionBackend):
154
accept_output_buffer: bool = True
155
- supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
155
+ supported_dtypes: ClassVar[list[torch.dtype]] = [
156
+ torch.float16,
157
+ torch.bfloat16,
158
+ torch.float32,
159
+ ]
160
161
@classmethod
162
def get_supported_head_sizes(cls) -> list[int]: