From 8da2f28f53c14e2c21c50821d89e3909d9c84af6 Mon Sep 17 00:00:00 2001 From: Pleaplusone Date: Thu, 13 Nov 2025 22:18:20 +0800 Subject: [PATCH] [ROCm][BugFix]Fix `get_cu_count` in rocm_aiter_fa.py (#28618) Signed-off-by: ganyi --- vllm/v1/attention/backends/rocm_aiter_fa.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index c7f925817a6a..ad454daa582e 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -18,6 +18,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv +from vllm.utils.platform_utils import get_cu_count from vllm.v1.attention.backends.utils import ( AttentionCGSupport, AttentionMetadataBuilder, @@ -38,7 +39,7 @@ if current_platform.is_rocm(): return min(65536 // x.element_size(), triton.next_power_of_2(head_dim)) def num_programs(total_tokens): - return min(total_tokens, current_platform.get_cu_count()) + return min(total_tokens, get_cu_count()) @triton.jit def cp_mha_gather_cache_kernel(