fix[ROCm]: Remove unconditional aiter import (#32902)
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
@@ -32,10 +32,11 @@ from vllm.v1.kv_cache_interface import AttentionSpec
|
||||
_PARTITION_SIZE_ROCM = 256
|
||||
_CP_TOKENS_PER_ITER_ROCM = 32 * 1024
|
||||
if current_platform.is_rocm():
|
||||
import aiter
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
if rocm_aiter_ops.is_enabled():
|
||||
import aiter
|
||||
|
||||
def block_size(x, head_dim):
|
||||
return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user