Skip to content
Merged
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions vllm/v1/attention/backends/rocm_aiter_fa.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,15 +31,14 @@

if current_platform.is_rocm():
import aiter
from aiter.ops.triton.utils.device_info import get_num_sms

from vllm.triton_utils import tl, triton

def block_size(x, head_dim):
return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))

def num_programs(head_dim):
return min(head_dim, get_num_sms())
def num_programs(total_tokens):
return min(total_tokens, current_platform.get_cu_count())

@triton.jit
def cp_mha_gather_cache_kernel(
Expand Down