diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py index ce27ea1220..d2b4602c56 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk.py +++ b/vllm/models/minimax_m3/amd/ops/index_topk.py @@ -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 diff --git a/vllm/models/minimax_m3/amd/ops/sparse_attn.py b/vllm/models/minimax_m3/amd/ops/sparse_attn.py index 0719374496..b4355d18f6 100644 --- a/vllm/models/minimax_m3/amd/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/amd/ops/sparse_attn.py @@ -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() diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index af703c0862..8dfd6e0c19 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -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) # ---------------------------------------------------------------------------