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] Remove the usage of `device_info` from aiter (#28383)
Signed-off-by: ganyi <ygan@amd.com>
Pleaplusone committed
Nov 13, 2025 at 13:43 UTC
ca00b1bfc69e71d860485340f0a197bf584ec004
1 file changed
+5
-6
vllm/v1/attention/backends/rocm_aiter_fa.py
+5
-6
@@ -31,15 +31,14 @@ _CP_TOKENS_PER_ITER_ROCM = 32 * 1024
31
32
if current_platform.is_rocm():
33
import aiter
34
- from aiter.ops.triton.utils.device_info import get_num_sms
34
35
from vllm.triton_utils import tl, triton
36
37
def block_size(x, head_dim):
38
return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))
39
41
- def num_programs(head_dim):
42
- return min(head_dim, get_num_sms())
40
+ def num_programs(total_tokens):
41
+ return min(total_tokens, current_platform.get_cu_count())
42
43
@triton.jit
44
def cp_mha_gather_cache_kernel(
@@ -58,11 +57,11 @@ if current_platform.is_rocm():
57
x,
58
max_block_num,
59
num_tokens,
60
+ num_programs,
61
DEQUANT: tl.constexpr,
62
PAGE_SIZE: tl.constexpr,
63
CACHE_FORMAT: tl.constexpr,
64
BLOCK_SIZE: tl.constexpr,
65
- NUM_PRGMS: tl.constexpr,
65
):
66
bid = tl.program_id(0)
67
col_offsets = tl.arange(0, BLOCK_SIZE)
@@ -70,7 +69,7 @@ if current_platform.is_rocm():
69
k_scale = tl.load(k_scale_ptr)
70
v_scale = tl.load(v_scale_ptr)
71
73
- for token_id in tl.range(bid, num_tokens, NUM_PRGMS):
72
+ for token_id in tl.range(bid, num_tokens, num_programs):
73
key_ptr_offset = key_ptr + token_id * head_size * num_heads
74
value_ptr_offset = value_ptr + token_id * head_size * num_heads
75
batch_idx = tl.load(token_to_batch_ptr + token_id)
@@ -162,11 +161,11 @@ if current_platform.is_rocm():
161
x,
162
block_tables.size(1),
163
total_tokens,
164
+ NUM_PRGMS,
165
DEQUANT=dequant,
166
PAGE_SIZE=page_size,
167
CACHE_FORMAT=kv_cache_layout,
168
BLOCK_SIZE=BLOCK_SIZE,
169
- NUM_PRGMS=NUM_PRGMS,
169
)
170
171