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 13, 08:15 UTC
[ROCm][BugFix]Fix `get_cu_count` in rocm_aiter_fa.py (#28618)
Signed-off-by: ganyi <ygan@amd.com>
Pleaplusone committed
Nov 13, 2025 at 22:18 UTC
8da2f28f53c14e2c21c50821d89e3909d9c84af6
1 file changed
+2
-1
vllm/v1/attention/backends/rocm_aiter_fa.py
+2
-1
@@ -18,6 +18,7 @@ from vllm.config import VllmConfig
18
from vllm.logger import init_logger
19
from vllm.platforms import current_platform
20
from vllm.utils.math_utils import cdiv
21
+from vllm.utils.platform_utils import get_cu_count
22
from vllm.v1.attention.backends.utils import (
23
AttentionCGSupport,
24
AttentionMetadataBuilder,
@@ -38,7 +39,7 @@ if current_platform.is_rocm():
39
return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))
40
41
def num_programs(total_tokens):
41
- return min(total_tokens, current_platform.get_cu_count())
42
+ return min(total_tokens, get_cu_count())
43
44
@triton.jit
45
def cp_mha_gather_cache_kernel(