iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Triton] Migrating MiniMax-M3 kernels from make_block_ptr to tensor descriptors (#59285)

Browse Source 
Signed-off-by: jpvillam <juan.villamizar@amd.com>
Signed-off-by: Daniel <danichan@amd.com>
Co-authored-by: jpvillam <juan.villamizar@amd.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Rohan Potdar <rohan.potdar@amd.com>
Co-authored-by: Daniel <danichan@amd.com>
This commit is contained in:
Daniel Chan
2026-10-07 11:49:59 -05:00
committed by GitHub
co-authored by jpvillam Claude Opus 5.5 Rohan Potdar Daniel
parent 240ee8de9e
commit d64a18832b
3 changed files with 80 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/minimax_m3/amd/ops/index_topk.py
+8 -7
View File
@@ -229,15 +229,16 @@ def _index_block_score_kernel(
if BLOCK_SIZE_Q * pid_q >= q_len:
return
q_ptrs = tl.make_block_ptr(
q_desc = tl.make_tensor_descriptor(
base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h,
shape=(q_len, head_dim),
strides=(stride_q_n, stride_q_d),
offsets=(pid_q * BLOCK_SIZE_Q, 0),
block_shape=(BLOCK_SIZE_Q, head_dim),
order=(1, 0),
shape=[q_len, head_dim],
strides=[stride_q_n, stride_q_d],
block_shape=[BLOCK_SIZE_Q, head_dim],
padding_option="zero",
)
q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero")
# 2. Perform the load by providing the multi-dimensional offset directly
q = q_desc.load([pid_q * BLOCK_SIZE_Q, 0])
q_start = prefix_len + pid_q * BLOCK_SIZE_Q
off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len
vllm/models/minimax_m3/amd/ops/sparse_attn.py
+16 -14
View File
@@ -139,15 +139,15 @@ def _gqa_sparse_fwd_kernel(
off_t = tl.arange(0, BLOCK_SIZE_T)
topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1)
real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
q_ptrs = tl.make_block_ptr(
q_desc = tl.make_tensor_descriptor(
base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
shape=(q_len, gqa_group_size, head_dim),
strides=(stride_qn, stride_qh, stride_qd),
offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(2, 1, 0),
shape=[q_len, gqa_group_size, head_dim],
strides=[stride_qn, stride_qh, stride_qd],
block_shape=[BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D],
padding_option="zero",
)
q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero")
q = q_desc.load([pid_q_j * BLOCK_SIZE_Q, 0, 0])
m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32)
@@ -228,15 +228,17 @@ def _gqa_sparse_fwd_kernel(
lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None]
acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D)
o_ptrs = tl.make_block_ptr(
o_desc = tl.make_tensor_descriptor(
base=o_ptr + q_start * stride_on + pid_h * stride_oh,
shape=(q_len, gqa_group_size, head_dim),
strides=(stride_on, stride_oh, stride_od),
offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(2, 1, 0),
shape=[q_len, gqa_group_size, head_dim],
strides=[stride_on, stride_oh, stride_od],
block_shape=[BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D],
)
o_desc.store(
offsets=[pid_q_j * BLOCK_SIZE_Q, 0, 0],
value=acc_o.to(o_ptr.dtype.element_ty),
)
tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2))
@torch.no_grad()
vllm/models/minimax_m3/common/ops/sparse_attn.py
+56 -49
View File
@@ -110,7 +110,9 @@ def _gqa_sparse_fwd_kernel(
bt_row = block_table_ptr + pid_b * stride_bt_b
off_n = tl.arange(0, BLOCK_SIZE_K)
off_d = tl.arange(0, BLOCK_SIZE_D)
off_h = tl.arange(0, BLOCK_SIZE_H)
d_mask = off_d < head_dim
h_mask = off_h < gqa_group_size
for j in range(real_q_loop):
pid_q_j = pid_q * num_q_loop + j
t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th
@@ -118,15 +120,20 @@ def _gqa_sparse_fwd_kernel(
q_abs = prefix_len + pid_q_j * BLOCK_SIZE_Q
valid_blocks = (q_abs + BLOCK_SIZE_K) // BLOCK_SIZE_K
real_topk = tl.minimum(max_topk, valid_blocks)
q_ptrs = tl.make_block_ptr(
base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
shape=(q_len, gqa_group_size, head_dim),
strides=(stride_qn, stride_qh, stride_qd),
offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(2, 1, 0),
off_qi = pid_q_j * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)
qo_mask = (
(off_qi < q_len)[:, None, None]
& h_mask[None, :, None]
& d_mask[None, None, :]
)
q = tl.load(
q_ptr
+ (q_start + off_qi)[:, None, None] * stride_qn
+ (pid_h + off_h)[None, :, None] * stride_qh
+ off_d[None, None, :] * stride_qd,
mask=qo_mask,
other=0.0,
)
q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero")
off_q = (
tl.arange(0, BLOCK_SIZE_Q)[:, None]
+ pid_q_j * BLOCK_SIZE_Q
@@ -203,15 +210,14 @@ def _gqa_sparse_fwd_kernel(
lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None]
acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D)
o_ptrs = tl.make_block_ptr(
base=o_ptr + q_start * stride_on + pid_h * stride_oh,
shape=(q_len, gqa_group_size, head_dim),
strides=(stride_on, stride_oh, stride_od),
offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(2, 1, 0),
tl.store(
o_ptr
+ (q_start + off_qi)[:, None, None] * stride_on
+ (pid_h + off_h)[None, :, None] * stride_oh
+ off_d[None, None, :] * stride_od,
acc_o.to(o_ptr.dtype.element_ty),
mask=qo_mask,
)
tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2))
# ---------------------------------------------------------------------------
@@ -306,21 +312,23 @@ def _gqa_sparse_decode_kernel(
off_n = tl.arange(0, BLOCK_SIZE_K)
off_d = tl.arange(0, BLOCK_SIZE_D)
off_h = tl.arange(0, BLOCK_SIZE_H)
d_mask = off_d < head_dim
h_mask = off_h < gqa_group_size
hd_mask = h_mask[:, None] & d_mask[None, :]
bt_row = block_table_ptr + req_id * stride_bt_b
m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32)
q_ptrs = tl.make_block_ptr(
base=q_ptr + pid_b * stride_qn + pid_h * stride_qh,
shape=(gqa_group_size, head_dim),
strides=(stride_qh, stride_qd),
offsets=(0, 0),
block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(1, 0),
q = tl.load(
q_ptr
+ pid_b * stride_qn
+ (pid_h + off_h)[:, None] * stride_qh
+ off_d[None, :] * stride_qd,
mask=hd_mask,
other=0.0,
)
q = tl.load(q_ptrs, boundary_check=(0, 1), padding_option="zero")
cur_idx_ptr = idx_base + chunk_start_topk * stride_tk
for _ in tl.range(chunk_start_topk, chunk_end_topk):
@@ -392,24 +400,23 @@ def _gqa_sparse_decode_kernel(
# can hit 0 * NaN. All-empty padded rows may still produce NaNs in merge.
scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), tl.zeros_like(lse_i))
acc_o = acc_o * scale[:, None]
o_ptrs = tl.make_block_ptr(
base=o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h,
shape=(gqa_group_size, head_dim),
strides=(stride_o_h, stride_o_d),
offsets=(0, 0),
block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
order=(1, 0),
tl.store(
o_ptr
+ pid_c * stride_o_c
+ pid_b * stride_o_b
+ (pid_h + off_h)[:, None] * stride_o_h
+ off_d[None, :] * stride_o_d,
acc_o.to(o_ptr.dtype.element_ty),
mask=hd_mask,
)
tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1))
lse_ptrs = tl.make_block_ptr(
base=lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h,
shape=(gqa_group_size,),
strides=(stride_l_h,),
offsets=(0,),
block_shape=(BLOCK_SIZE_H,),
order=(0,),
tl.store(
lse_ptr
+ pid_c * stride_l_c
+ pid_b * stride_l_b
+ (pid_h + off_h) * stride_l_h,
lse_i.to(lse_ptr.dtype.element_ty),
mask=h_mask,
)
tl.store(lse_ptrs, lse_i.to(lse_ptr.dtype.element_ty), boundary_check=(0,))
@triton.heuristics(
@@ -444,16 +451,16 @@ def _merge_topk_attn_out_kernel(
off_c = tl.arange(0, NUM_TOPK_CHUNKS)
off_d = tl.arange(0, BLOCK_SIZE_D)
o_ptrs = tl.make_block_ptr(
base=o_ptr + pid_b * stride_o_b + pid_h * stride_o_h,
shape=(NUM_TOPK_CHUNKS, head_dim),
strides=(stride_o_c, stride_o_d),
offsets=(0, 0),
block_shape=(NUM_TOPK_CHUNKS, BLOCK_SIZE_D),
order=(1, 0),
d_mask = off_d < head_dim
o_ptrs = (
o_ptr
+ pid_b * stride_o_b
+ pid_h * stride_o_h
+ off_c[:, None] * stride_o_c
+ off_d[None, :] * stride_o_d
)
lse_ptrs = lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c
o = tl.load(o_ptrs, boundary_check=(0, 1), padding_option="zero")
o = tl.load(o_ptrs, mask=d_mask[None, :], other=0.0)
lse = tl.load(lse_ptrs) # empty chunks contribute -inf -> weight 0
lse_max = tl.max(lse, axis=0)
weights = tl.exp2(lse - lse_max)
@@ -462,7 +469,7 @@ def _merge_topk_attn_out_kernel(
out_ptrs = (
out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d
)
tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim)
tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=d_mask)
# ---------------------------------------------------------------------------