From c63cd4906c2a67f18b3714786cc036c1ad97a64f Mon Sep 17 00:00:00 2001 From: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Date: Thu, 25 Jun 2026 04:56:00 -0400 Subject: [PATCH] [ROCm][ [Perf] sparse attention optimization on minimax-m3 (#46546) Signed-off-by: Hongxia Yang Signed-off-by: tjtanaa Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: yueliu14 Co-authored-by: tjtanaa --- vllm/models/minimax_m3/amd/ops/index_topk.py | 939 ++++++++++++++++++ vllm/models/minimax_m3/amd/ops/sparse_attn.py | 271 +++++ vllm/models/minimax_m3/common/indexer.py | 19 +- .../minimax_m3/common/ops/sparse_attn.py | 26 - .../minimax_m3/common/sparse_attention.py | 19 +- 5 files changed, 1238 insertions(+), 36 deletions(-) create mode 100644 vllm/models/minimax_m3/amd/ops/index_topk.py create mode 100644 vllm/models/minimax_m3/amd/ops/sparse_attn.py diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py new file mode 100644 index 00000000000..2b076a38b89 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/index_topk.py @@ -0,0 +1,939 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton kernels for MiniMax M3 lightning-indexer block scoring + top-k. + +Index queries score each 128-token block of index keys (max over the block), +then the top-k blocks (plus forced init/local blocks) are selected per query +token. Adapted to vLLM's paged KV cache: the KV page size is forced to equal the +sparse block size (128), so one sparse block maps to exactly one page. + +Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head). + +Only the paths MiniMax M3 uses are implemented: score_type="max", index value +disabled (score-only indexer), single shared index head. The selected block ids +feed the block-sparse attention kernels in ``sparse_attn``. +""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import round_up + +# One sparse block == one KV page. +SPARSE_BLOCK_SIZE = 128 + + +# --------------------------------------------------------------------------- +# Bitonic top-k helpers (layout-agnostic). +# --------------------------------------------------------------------------- +@triton.jit +def _compare_and_swap(x, ids, flip, i: tl.constexpr, n_dims: tl.constexpr): + n_outer: tl.constexpr = x.numel >> n_dims + shape: tl.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)] + y = tl.reshape(x, shape) + mask = tl.arange(0, 2)[None, :, None] + left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype) + right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype) + left = tl.reshape(left, x.shape) + right = tl.reshape(right, x.shape) + y_idx = tl.reshape(ids, shape) + left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape) + right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape) + left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype) + right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype) + idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) + ileft = left.to(idtype, bitcast=True) + iright = right.to(idtype, bitcast=True) + ix = x.to(idtype, bitcast=True) + cond = (left > right) != flip + ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix)) + new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids)) + return ret.to(x.dtype, bitcast=True), new_ids + + +@triton.jit +def _bitonic_merge( + x, ids, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr +): + n_outer: tl.constexpr = x.numel >> n_dims + tl.static_assert(stage <= n_dims) + if order == 2: + shape: tl.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage] + flip = tl.reshape( + tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape + ) + else: + flip = order + for i in tl.static_range(stage): + x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims) + return x, ids + + +# --------------------------------------------------------------------------- +# Index block-score kernel (paged). score[h, token, block] = max over the +# 128-token block of (idx_q . index_k), causal-masked. BLOCK_SIZE_K == 128 so +# each K-tile is exactly one page (BLOCKS_PER_K_BLOCK == 1). +# --------------------------------------------------------------------------- +# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens +# might lose pointer alignment, which trigger Triton recompiles. we don't actually +# need pointer alignment for those tensors anyway because we do scalar load. +@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"]) +def _index_block_score_kernel( + q_ptr, # idx_q: [total_q, num_idx_heads, head_dim] + ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim] + score_ptr, # [num_idx_heads, total_q, max_block] + block_table_ptr, # [num_reqs, max_blocks] + cu_seqlens, # [batch+1] query start offsets + seq_lens, # [batch] total K length + prefix_lens, # [batch] context length before this chunk's queries + num_idx_heads, + head_dim: tl.constexpr, + stride_q_n, + stride_q_h, + stride_q_d, + stride_ik_blk, + stride_ik_pos, + stride_ik_d, + stride_s_h, + stride_s_n, + stride_s_k, + stride_bt_b, + BLOCK_SIZE_Q: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) +): + pid_q = tl.program_id(0) + pid_bh = tl.program_id(1) + pid_b = pid_bh // num_idx_heads + pid_h = pid_bh % num_idx_heads + + seq_start = tl.load(cu_seqlens + pid_b) + q_len = tl.load(cu_seqlens + pid_b + 1) - seq_start + seq_len = tl.load(seq_lens + pid_b) + prefix_len = tl.load(prefix_lens + pid_b) + if BLOCK_SIZE_Q * pid_q >= q_len: + return + + q_ptrs = tl.make_block_ptr( + 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), + ) + q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") + q_start = prefix_len + pid_q * BLOCK_SIZE_Q + + off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len + off_k = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, head_dim) + # Block table row for this request. + bt_row = block_table_ptr + pid_b * stride_bt_b + # Causal window: only blocks up to the last query token's position. + hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q) + for i in tl.range(0, hi, BLOCK_SIZE_K): + blk = i // BLOCK_SIZE_K + page = tl.load(bt_row + blk).to(tl.int64) + pos = i + off_k + # index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed) + # we don't need masked load for K, because KV cache ensures + # allocation is multiple of BLOCK_SIZE_K. + # for tokens beyond seqlen, they will be masked in qk later. + k = tl.load( + ik_cache_ptr + + page * stride_ik_blk + + off_k[None, :] * stride_ik_pos + + off_d[:, None] * stride_ik_d, + ) + qk = tl.dot(q, k) + # apply causal mask as needed + if q_start < i + BLOCK_SIZE_K: + qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf")) + # one sparse block per K-tile -> max over the 128 positions + score = tl.max(qk, axis=1) # [BLOCK_SIZE_Q] + s_ptrs = ( + score_ptr + + pid_h * stride_s_h + + (seq_start + pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) + * stride_s_n + + blk * stride_s_k + ) + q_store_mask = (pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) < q_len + tl.store(s_ptrs, score, mask=q_store_mask) + + +# --------------------------------------------------------------------------- +# Top-k selection over per-token block scores (layout-agnostic). block_size_q +# is 1 for M3, so top-k is computed per query token. +# --------------------------------------------------------------------------- +# since prefill metadata is sliced from mixed batch metadata, prefix_lens +# might lose pointer alignment, which trigger Triton recompiles. we don't actually +# need pointer alignment for those tensors anyway because we do scalar load. +@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])}) +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE_K": 2048}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 1024}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 512}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2), + ], + key=["BLOCK_SIZE_T"], +) +@triton.jit(do_not_specialize_on_alignment=["prefix_lens"]) +def _topk_index_kernel( + s_ptr, # [num_heads, total_q, max_block] + ti_ptr, # [num_heads, total_q, topk] + sample_interval: tl.constexpr, # block_size_q (1 for M3) + block_size: tl.constexpr, # sparse block size (128) + cu_seqlens, + cu_seqblocks_q, + prefix_lens, + topk, + init_blocks: tl.constexpr, + local_blocks: tl.constexpr, + stride_s_h, + stride_s_n, + stride_s_k, + stride_ti_h, + stride_ti_n, + stride_ti_t, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + MASK_INIT: tl.constexpr, + MASK_LOCAL: tl.constexpr, +): + tl.static_assert(BLOCK_SIZE_K > BLOCK_SIZE_T) + pid_q = tl.program_id(0) + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + seq_start = tl.load(cu_seqlens + pid_b) + block_start = tl.load(cu_seqblocks_q + pid_b) + block_num = tl.load(cu_seqblocks_q + pid_b + 1) - block_start + prefix_len = tl.load(prefix_lens + pid_b) + if pid_q >= block_num: + return + off_k = tl.arange(0, BLOCK_SIZE_K) + off_t = tl.arange(0, BLOCK_SIZE_T) + s_ptrs = ( + s_ptr + + (seq_start + pid_q * sample_interval) * stride_s_n + + pid_h * stride_s_h + + off_k * stride_s_k + ) + topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32) + topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32) + left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2 + valid_blocks = (prefix_len + pid_q * sample_interval + block_size) // block_size + for i in tl.range(0, valid_blocks, BLOCK_SIZE_K): + causal_mask = i + off_k < valid_blocks + local_mask = i + off_k >= max(0, valid_blocks - local_blocks) + init_mask = i + off_k < init_blocks + score = tl.load(s_ptrs, mask=causal_mask, other=-1e30).to(tl.float32) + score = tl.where(score != score, -1e30, score) + s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K + if MASK_INIT: + score = tl.where(causal_mask & init_mask, score - 1e29, score) + else: + score = tl.where(causal_mask & init_mask, 1e30, score) + if MASK_LOCAL: + score = tl.where(causal_mask & local_mask, score - 1e28, score) + else: + score = tl.where(causal_mask & local_mask, 1e29, score) + topk_score, last_topk_score = score, topk_score + topk_idx, last_topk_idx = (tl.where(causal_mask, i + off_k + 1, 0), topk_idx) + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) + for j in tl.static_range(1, n_dims): + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), j, 2, n_dims + ) + if i != 0: + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims + ) + topk_score_new = last_topk_score * left_half_mask + topk_score * ( + 1 - left_half_mask + ) + topk_idx_new = last_topk_idx * left_half_mask + topk_idx * ( + 1 - left_half_mask + ) + topk_score, topk_idx = _bitonic_merge( + topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims + ) + else: + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims + ) + topk_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 + topk_idx = tl.sum( + topk_mask[:, None] + * tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), + axis=0, + ) + ti_ptrs = ( + ti_ptr + + (block_start + pid_q) * stride_ti_n + + pid_h * stride_ti_h + + off_t * stride_ti_t + ) + store_mask = off_t < topk + valid_mask = off_t < valid_blocks + topk_idx = tl.where(store_mask & valid_mask, topk_idx, -1) + tl.store(ti_ptrs, topk_idx.to(ti_ptrs.dtype.element_ty), mask=store_mask) + + +# --------------------------------------------------------------------------- +# Decode index-score kernel (split-K over seq blocks). Decode batches are +# flattened request-major, with a runtime query length used to map each query +# token back to its request metadata. Chunk counts depend only on shape +# constants so the grid is fixed within a cuda graph. The score scale is omitted +# because decode only consumes block ordering. +# --------------------------------------------------------------------------- +@triton.jit(do_not_specialize=["num_kv_chunks", "decode_query_len"]) +def _decode_index_score_kernel( + q_ptr, # idx_q: [total_q, num_idx_heads, head_dim] + ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim] + score_ptr, # [num_idx_heads, total_q, max_block] + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + num_idx_heads: tl.constexpr, + head_dim: tl.constexpr, + init_blocks, + local_blocks, + decode_query_len, + stride_q_n, + stride_q_h, + stride_q_d, + stride_ik_blk, + stride_ik_pos, + stride_ik_d, + stride_s_h, + stride_s_n, + stride_s_k, + stride_bt_b, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + BLOCK_SIZE_Q: tl.constexpr, + num_kv_chunks, + USE_PDL: tl.constexpr, +): + BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) + hq_offsets = tl.arange(0, BLOCK_SIZE_HQ) + h_offsets = hq_offsets // BLOCK_SIZE_Q + q_offsets = hq_offsets % BLOCK_SIZE_Q + q_mask = q_offsets < decode_query_len + q_ids = pid_r * decode_query_len + q_offsets + + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + pid_r) + query_pos = seq_len - decode_query_len + q_offsets + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + num_blocks_q = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + kv_len_max = tl.max(tl.where(q_mask, kv_len, 0), axis=0) + num_blocks = (kv_len_max + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + + # block-aligned fixed-count split: grid independent of seq_len (cuda graph). + chunk_size_blocks = (num_blocks + num_kv_chunks - 1) // num_kv_chunks + chunk_start_block = pid_c * chunk_size_blocks + chunk_end_block = tl.minimum(chunk_start_block + chunk_size_blocks, num_blocks) + if chunk_start_block >= chunk_end_block: + return + off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block + off_d = tl.arange(0, head_dim) + bt_row = block_table_ptr + pid_r * stride_bt_b + # Force-select init (1e30) and local (1e29, higher priority) blocks. + local_start = tl.maximum(0, num_blocks_q - local_blocks) + # Query vectors for all index heads in a small spec-decode block. + q = tl.load( + q_ptr + + q_ids[None, :] * stride_q_n + + h_offsets[None, :] * stride_q_h + + off_d[:, None] * stride_q_d, + mask=q_mask[None, :], + other=0.0, + ) # [D,HQ] + for blk in tl.range(chunk_start_block, chunk_end_block): + page = tl.load(bt_row + blk).to(tl.int64) + pos = blk * BLOCK_SIZE_K + off_k + pos_mask = pos[:, None] < kv_len[None, :] + # we don't need masked load for K, because KV cache ensures + # allocation is multiple of BLOCK_SIZE_K. + # for tokens beyond seqlen, they will be masked in qk later. + k = tl.load( + ik_cache_ptr + + page * stride_ik_blk + + off_k[:, None] * stride_ik_pos + + off_d * stride_ik_d, + ) # [N,D] + if BLOCK_SIZE_HQ == 1: + # Degenerate GEMV (q is [D,1]): vectorized fp32 multiply + reduce + # instead of an MFMA tile. Numerically equivalent to tl.dot. + q_vec = tl.sum(q, axis=1).to(tl.float32) # [D] + kq = tl.sum(k.to(tl.float32) * q_vec[None, :], axis=1)[:, None] # [N,1] + else: + # fp32 accumulation is required for the fp8 (e4m3) index cache: q/k + # are loaded in their stored dtype (bf16 or e4m3) and the MMA + # accumulates in fp32 so the per-block max score is exact for the + # fp8 indexer too. + kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ] + kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf")) + score = tl.max(kq, axis=0) # [HQ] + is_visible_block = blk < num_blocks_q + is_init = (blk < init_blocks) & is_visible_block + is_local = (blk >= local_start) & is_visible_block + score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score)) + tl.store( + score_ptr + h_offsets * stride_s_h + q_ids * stride_s_n + blk * stride_s_k, + score, + mask=q_mask, + ) + + +# --------------------------------------------------------------------------- +# Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local +# blocks are already encoded in the scores. +# --------------------------------------------------------------------------- +@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])}) +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2), + ], + key=["topk"], +) +@triton.jit(do_not_specialize=["chunk_blocks", "decode_query_len"]) +def _topk_index_partial_kernel( + s_ptr, # score: [num_idx_heads, total_q, max_block] + ts_partial_ptr, # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T] + ti_partial_ptr, # partial idx out (1-indexed global, 0=invalid): same shape + seq_lens, # [num_reqs] + block_size: tl.constexpr, # sparse block size (128) + topk: tl.constexpr, + chunk_blocks, # how many score-blocks each chunk owns + decode_query_len, + stride_s_h, + stride_s_b, + stride_s_k, + stride_ts_c, + stride_ts_h, + stride_ts_b, + stride_ts_t, + stride_ti_c, + stride_ti_h, + stride_ti_b, + stride_ti_t, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + USE_PDL: tl.constexpr, +): + tl.static_assert(topk < BLOCK_SIZE_K) + pid_b = tl.program_id(0) # flattened query-token id + pid_h = tl.program_id(1) + pid_chunk = tl.program_id(2) + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + num_blocks = (kv_len + block_size - 1) // block_size + + # Slice this chunk owns within [0, num_blocks). + chunk_start = pid_chunk * chunk_blocks + chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks) + chunk_actual = tl.maximum(chunk_end - chunk_start, 0) + + off_k = tl.arange(0, BLOCK_SIZE_K) + off_t = tl.arange(0, BLOCK_SIZE_T) + + s_ptrs = ( + s_ptr + + pid_b * stride_s_b + + pid_h * stride_s_h + + (chunk_start + off_k) * stride_s_k + ) + + topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32) + topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32) + left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2 + + # Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty + # chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0. + for i in tl.range(0, chunk_actual, BLOCK_SIZE_K): + mask = off_k < chunk_actual - i + score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32) + score = tl.where(score != score, -1e30, score) + s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K + topk_score, last_topk_score = score, topk_score + topk_idx, last_topk_idx = ( + tl.where(mask, chunk_start + i + off_k + 1, 0), # 1-indexed global + topk_idx, + ) + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) + for j in tl.static_range(1, n_dims): + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), j, 2, n_dims + ) + if i != 0: + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims + ) + topk_score_new = last_topk_score * left_half_mask + topk_score * ( + 1 - left_half_mask + ) + topk_idx_new = last_topk_idx * left_half_mask + topk_idx * ( + 1 - left_half_mask + ) + topk_score, topk_idx = _bitonic_merge( + topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims + ) + else: + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims + ) + + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + # Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort). + topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 + final_score = tl.sum( + topk_mask_extract[:, None] + * tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), + axis=0, + ) + final_idx = tl.sum( + topk_mask_extract[:, None] + * tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), + axis=0, + ) + + # Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0 + # sentinels and lose to real scores in the merge stage. + ts_ptrs = ( + ts_partial_ptr + + pid_chunk * stride_ts_c + + pid_b * stride_ts_b + + pid_h * stride_ts_h + + off_t * stride_ts_t + ) + ti_ptrs = ( + ti_partial_ptr + + pid_chunk * stride_ti_c + + pid_b * stride_ti_b + + pid_h * stride_ti_h + + off_t * stride_ti_t + ) + tl.store(ts_ptrs, final_score) + tl.store(ti_ptrs, final_idx) + + +@triton.heuristics( + { + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]), + "BLOCK_SIZE_K": lambda args: triton.next_power_of_2( + args["num_topk_chunks"] * triton.next_power_of_2(args["topk"]) + ), + } +) +@triton.jit(do_not_specialize=["num_topk_chunks", "decode_query_len"]) +def _topk_index_merge_kernel( + ts_partial_ptr, # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T] + ti_partial_ptr, # partial idx (1-indexed global, 0=invalid): same shape + ti_final_ptr, # final idx (0-indexed, -1=invalid): [num_idx_heads, total_q, topk] + seq_lens, # [num_reqs] + block_size: tl.constexpr, # sparse block size (128) + topk: tl.constexpr, + decode_query_len, + stride_ts_c, + stride_ts_h, + stride_ts_b, + stride_ts_t, + stride_ti_c, + stride_ti_h, + stride_ti_b, + stride_ti_t, + stride_tif_h, + stride_tif_b, + stride_tif_t, + num_topk_chunks, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + USE_PDL: tl.constexpr, +): + pid_b = tl.program_id(0) # flattened query-token id + pid_h = tl.program_id(1) + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + num_blocks = (kv_len + block_size - 1) // block_size + + # Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K. + # Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T, + # in_chunk = p % BLOCK_SIZE_T. + off = tl.arange(0, BLOCK_SIZE_K) + chunk_idx = off // BLOCK_SIZE_T + in_chunk_idx = off % BLOCK_SIZE_T + valid = chunk_idx < num_topk_chunks + + score_offset = ( + chunk_idx * stride_ts_c + + pid_h * stride_ts_h + + pid_b * stride_ts_b + + in_chunk_idx * stride_ts_t + ) + idx_offset = ( + chunk_idx * stride_ti_c + + pid_h * stride_ti_h + + pid_b * stride_ti_b + + in_chunk_idx * stride_ti_t + ) + + score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to( + tl.float32 + ) + score = tl.where(score != score, -1e30, score) + idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32) + + # Full bitonic descending sort of BLOCK_SIZE_K items. + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) + for j in tl.static_range(1, n_dims): + score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims) + score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims) + + # Extract first BLOCK_SIZE_T positions — these are the global top-K. + extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 + topk_idx_final = tl.sum( + extract_mask[:, None] + * tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), + axis=0, + ) + + off_t = tl.arange(0, BLOCK_SIZE_T) + tif_ptrs = ( + ti_final_ptr + + pid_h * stride_tif_h + + pid_b * stride_tif_b + + off_t * stride_tif_t + ) + store_mask = off_t < topk + topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1) + tl.store( + tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask + ) + + +# --------------------------------------------------------------------------- +# Python wrappers +# --------------------------------------------------------------------------- +@torch.no_grad() +def minimax_m3_index_score( + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] + block_table: torch.Tensor, # [batch, max_blocks] + cu_seqlens_q: torch.Tensor, # [batch+1] int32 + seq_lens: torch.Tensor, # [batch] int32 + prefix_lens: torch.Tensor, # [batch] int32 + max_query_len: int, + max_seq_len: int, + num_kv_heads: int, +) -> torch.Tensor: + """Compute per-token index scores for each visible sparse block. + + Returns score [num_kv_heads, total_q, max_block], where each score is the + max over a 128-token index-K block. M3 has num_idx_heads == num_kv_heads. + """ + total_q, num_idx_heads, head_dim = idx_q.shape + assert num_idx_heads == num_kv_heads, ( + "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)" + ) + batch = cu_seqlens_q.shape[0] - 1 + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) + + # Keep score strides 16-divisible to avoid Triton recompiles. + score_block_stride = round_up(max_block, 16) + score = torch.empty( + (num_idx_heads, total_q, score_block_stride), + dtype=torch.float32, + device=idx_q.device, + ) + BLOCK_SIZE_Q = 64 + grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads) + _index_block_score_kernel[grid_score]( + idx_q, + index_kv_cache, + score, + block_table, + cu_seqlens_q, + seq_lens, + prefix_lens, + num_idx_heads, + head_dim, + idx_q.stride(0), + idx_q.stride(1), + idx_q.stride(2), + index_kv_cache.stride(0), + index_kv_cache.stride(1), + index_kv_cache.stride(2), + score.stride(0), + score.stride(1), + score.stride(2), + block_table.stride(0), + BLOCK_SIZE_Q=BLOCK_SIZE_Q, + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + ) + return score + + +@torch.no_grad() +def minimax_m3_index_topk( + score: torch.Tensor, # [num_idx_heads, total_q, max_block] + cu_seqlens_q: torch.Tensor, # [batch+1] int32 + prefix_lens: torch.Tensor, # [batch] int32 + max_query_len: int, + topk: int, + init_blocks: int, + local_blocks: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Select index top-k from a precomputed score tensor. + + When ``out`` is provided (a ``[num_idx_heads, >=total_q, topk]`` buffer), the + result is written into ``out[:, :total_q, :]`` instead of a fresh tensor -- + used to keep the top-k output at a stable address for cudagraph capture. + """ + num_idx_heads = score.shape[0] + batch = cu_seqlens_q.shape[0] - 1 + total_q = score.shape[1] + if out is not None: + topk_idx = out[:, :total_q, :] + else: + topk_idx = torch.empty( + (num_idx_heads, total_q, topk), + dtype=torch.int32, + device=score.device, + ) + # block_size_q == 1 -> query blocks coincide with query tokens. + grid_topk = (max_query_len, batch, num_idx_heads) + _topk_index_kernel[grid_topk]( + score, + topk_idx, + 1, # sample_interval (block_size_q) + SPARSE_BLOCK_SIZE, + cu_seqlens_q, + cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1 + prefix_lens, + topk, + init_blocks, + local_blocks, + score.stride(0), + score.stride(1), + score.stride(2), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + MASK_INIT=False, + MASK_LOCAL=False, + ) + return topk_idx + + +@torch.no_grad() +def minimax_m3_index_decode( + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + max_seq_len: int, + topk: int, + init_blocks: int, + local_blocks: int, + num_kv_heads: int, + decode_query_len: int, + max_decode_query_len: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Decode index block-score + top-k, both split-K (cudagraph-safe). + + Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad). + When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into + ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating. + """ + total_q, num_idx_heads, head_dim = idx_q.shape + assert num_idx_heads == num_kv_heads, ( + "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)" + ) + assert decode_query_len <= max_decode_query_len + assert total_q == seq_lens.shape[0] * decode_query_len + batch = total_q + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) + use_pdl = current_platform.is_arch_support_pdl() + # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA + # SM9+); this ROCm Triton rejects it even when False ("Keyword argument + # launch_pdl was specified but unrecognised"). Only pass it when PDL is + # actually supported -- on ROCm use_pdl is always False, so it's omitted. + pdl_kwargs: dict[str, bool | int] = {} + if use_pdl: + pdl_kwargs.update({"launch_pdl": True}) + # TP=1 spec decode scores a wide 4-head x 4-position query tile per K block; + # reduce stages to ease memory/register pressure. Keep no-spec and TP=4 + # single-head codegen unchanged. + score_kwargs = pdl_kwargs.copy() + if num_idx_heads > 1 and max_decode_query_len > 1: + score_kwargs.update({"num_warps": 4, "num_stages": 2}) + + # Keep score strides 16-divisible to avoid Triton recompiles. + score_block_stride = round_up(max_block, 16) + score = torch.empty( + (num_idx_heads, total_q, score_block_stride), + dtype=torch.float32, + device=idx_q.device, + ) + # split-K over seq blocks; chunk count depends only on shape constants so + # the grid is fixed within a cuda graph. + TARGET_GRID = 512 + MAX_NUM_KV_CHUNKS = 256 + # Use the configured max decode length to avoid Triton recompiles when + # switching between qlen=1 and spec-decode verification batches. + BLOCK_SIZE_Q = triton.next_power_of_2(max_decode_query_len) + score_ctas_per_chunk = seq_lens.shape[0] + target = max( + 1, + min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, score_ctas_per_chunk)), + ) + num_kv_chunks = 1 << (target.bit_length() - 1) + grid_score = (seq_lens.shape[0], num_kv_chunks) + _decode_index_score_kernel[grid_score]( + idx_q, + index_kv_cache, + score, + block_table, + seq_lens, + num_idx_heads, + head_dim, + init_blocks, + local_blocks, + decode_query_len, + idx_q.stride(0), + idx_q.stride(1), + idx_q.stride(2), + index_kv_cache.stride(0), + index_kv_cache.stride(1), + index_kv_cache.stride(2), + score.stride(0), + score.stride(1), + score.stride(2), + block_table.stride(0), + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + BLOCK_SIZE_Q=BLOCK_SIZE_Q, + num_kv_chunks=num_kv_chunks, + USE_PDL=use_pdl, + **score_kwargs, + ) + + if out is not None: + topk_idx = out[:, :total_q, :] + else: + topk_idx = torch.empty( + (num_idx_heads, total_q, topk), + dtype=torch.int32, + device=idx_q.device, + ) + # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts + # pow2(num_topk_chunks * pow2(topk)) candidates. + TOPK_TARGET_GRID = 64 + MAX_NUM_TOPK_CHUNKS = 16 + topk_target = max( + 1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads)) + ) + num_topk_chunks = 1 << (topk_target.bit_length() - 1) + block_size_t = triton.next_power_of_2(topk) + chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks + topk_score_partial = torch.empty( + num_topk_chunks, + num_idx_heads, + batch, + block_size_t, + dtype=torch.float32, + device=idx_q.device, + ) + topk_idx_partial = torch.empty( + num_topk_chunks, + num_idx_heads, + batch, + block_size_t, + dtype=torch.int32, + device=idx_q.device, + ) + _topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)]( + score, + topk_score_partial, + topk_idx_partial, + seq_lens, + SPARSE_BLOCK_SIZE, + topk, + chunk_blocks, + decode_query_len, + score.stride(0), + score.stride(1), + score.stride(2), + topk_score_partial.stride(0), + topk_score_partial.stride(1), + topk_score_partial.stride(2), + topk_score_partial.stride(3), + topk_idx_partial.stride(0), + topk_idx_partial.stride(1), + topk_idx_partial.stride(2), + topk_idx_partial.stride(3), + USE_PDL=use_pdl, + **pdl_kwargs, + ) + _topk_index_merge_kernel[(batch, num_idx_heads)]( + topk_score_partial, + topk_idx_partial, + topk_idx, + seq_lens, + SPARSE_BLOCK_SIZE, + topk, + decode_query_len, + topk_score_partial.stride(0), + topk_score_partial.stride(1), + topk_score_partial.stride(2), + topk_score_partial.stride(3), + topk_idx_partial.stride(0), + topk_idx_partial.stride(1), + topk_idx_partial.stride(2), + topk_idx_partial.stride(3), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + num_topk_chunks=num_topk_chunks, + USE_PDL=use_pdl, + **pdl_kwargs, + ) + return topk_idx diff --git a/vllm/models/minimax_m3/amd/ops/sparse_attn.py b/vllm/models/minimax_m3/amd/ops/sparse_attn.py new file mode 100644 index 00000000000..015a4c516bc --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/sparse_attn.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm gfx942/gfx950 block-sparse GQA prefill kernel for MiniMax-M3. + +Only the prefill path is specialized on CDNA: each 128-token KV block is split +into SUB_K-token sub-tiles to right-size the per-block QK/PV MFMAs. Everything +else -- the decode split-K kernels, the FP8 dtype set, the sparse block size -- +is reused unchanged from ``common.ops.sparse_attn``. +""" + +import torch + +from vllm.models.minimax_m3.common.ops.sparse_attn import ( + _FP8_DTYPES, + SPARSE_BLOCK_SIZE, + minimax_m3_sparse_attn_decode, +) +from vllm.platforms.rocm import on_gfx950, on_mi3xx +from vllm.triton_utils import tl, triton + +__all__ = ["minimax_m3_sparse_attn", "minimax_m3_sparse_attn_decode"] + + +# Sub-tile width for the prefill kernel's per-block QK/PV GEMMs. gfx950 -> 64, +# gfx942 -> 32 (re-tune with tune_sparse_attn.py). Must divide SPARSE_BLOCK_SIZE. +_SPARSE_ATTN_SUB_K = SPARSE_BLOCK_SIZE // 2 if on_gfx950() else SPARSE_BLOCK_SIZE // 4 + +_SPARSE_ATTN_PREFILL_KWARG: dict | None = None + + +def _sparse_attn_prefill_kwargs() -> dict: + """MFMA + pipeline launch params for the sub-tiled prefill kernel. + + gfx942 and gfx950 share the same params: ``num_warps=1`` keeps one wave + resident on the small per-sub-tile GEMM, ``matrix_instr_nonkdim=16`` / + ``kpack=2`` select the MFMA_16x16 path, and ``num_stages=1`` fits LDS and is + fastest in the sweep. Only the sub-tile width (``_SPARSE_ATTN_SUB_K``) + differs by arch. Empty on other AMD archs. Cached: arch is fixed per process. + """ + global _SPARSE_ATTN_PREFILL_KWARG + if _SPARSE_ATTN_PREFILL_KWARG is None: + kwarg: dict = {} + if on_mi3xx(): + kwarg = { + "num_warps": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2, + "num_stages": 1, + } + _SPARSE_ATTN_PREFILL_KWARG = kwarg + return _SPARSE_ATTN_PREFILL_KWARG + + +# --------------------------------------------------------------------------- +# GQA block-sparse attention (paged). Main heads attend only to the selected +# blocks. BLOCK_SIZE_K == 128 so each selected block is one page. +# --------------------------------------------------------------------------- +# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens +# might lose pointer alignment, which trigger Triton recompiles. we don't actually +# need pointer alignment for those tensors anyway because we do scalar load. +@triton.heuristics( + { + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]), + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]), + "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"] + * triton.next_power_of_2(args["gqa_group_size"]), + } +) +@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"]) +def _gqa_sparse_fwd_kernel( + q_ptr, # [total_q, num_heads, head_dim] + kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim] + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # [total_q, num_heads, head_dim] + block_table_ptr, # [num_reqs, max_blocks] + cu_seqlens_q, + cu_seqblocks_q, + seq_lens, + prefix_lens, + num_kv_heads, + gqa_group_size, + head_dim, + max_topk, + num_q_loop, + sm_scale, + stride_qn, + stride_qh, + stride_qd, + stride_kv_blk, + stride_kv_kv, + stride_kv_pos, + stride_kv_h, + stride_kv_d, + stride_th, + stride_tn, + stride_tk, + stride_on, + stride_oh, + stride_od, + stride_bt_b, + BLOCK_SIZE_Q: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + BLOCK_SIZE_D: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_QH: tl.constexpr, + USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load + SUB_K: tl.constexpr, # CDNA only: KV sub-tile width (see _IS_MI3XX) +): + sm_scale_log2e = sm_scale * 1.4426950409 + pid_q = tl.program_id(0) + pid_kh = tl.program_id(1) + pid_b = tl.program_id(2) + pid_h = pid_kh * gqa_group_size + q_start = tl.load(cu_seqlens_q + pid_b) + q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start + q_block_start = tl.load(cu_seqblocks_q + pid_b) + q_block_len = tl.load(cu_seqblocks_q + pid_b + 1) - q_block_start + seq_len = tl.load(seq_lens + pid_b) + prefix_len = tl.load(prefix_lens + pid_b) + if pid_q * num_q_loop >= q_block_len: + return + real_q_loop = min(num_q_loop, q_block_len - pid_q * num_q_loop) + bt_row = block_table_ptr + pid_b * stride_bt_b + off_d = tl.arange(0, BLOCK_SIZE_D) + d_mask = off_d < head_dim + 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 + 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( + 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), + ) + q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero") + 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) + q = tl.reshape(q, BLOCK_SIZE_QH, BLOCK_SIZE_D) + + # CDNA: process each 128-token KV block in SUB_K-token sub-tiles so + # each QK/PV MFMA is right-sized. Numerically equivalent to the dense + # path below (flash-softmax reassociation). + NUM_SUB: tl.constexpr = BLOCK_SIZE_K // SUB_K + for _ in tl.range(real_topk): + blk = tl.load(t_ptr_j).to(tl.int32) + t_ptr_j = t_ptr_j + stride_tk + c = blk * BLOCK_SIZE_K + page = tl.load(bt_row + blk).to(tl.int64) + kv_base = kv_cache_ptr + page * stride_kv_blk + pid_kh * stride_kv_h + for sub_i in range(NUM_SUB): + off_sub = tl.arange(0, SUB_K) + sub_i * SUB_K + pos_sub = c + off_sub + pos_mask_sub = pos_sub < seq_len + k_sub = tl.load( + kv_base + + 0 * stride_kv_kv + + off_sub[None, :] * stride_kv_pos + + off_d[:, None] * stride_kv_d, + mask=d_mask[:, None] & pos_mask_sub[None, :], + other=0.0, + ) + if USE_FP8: + k_sub = k_sub.to(q.dtype) + off_q_sub = ( + tl.arange(0, BLOCK_SIZE_Q)[:, None] + + pid_q_j * BLOCK_SIZE_Q + + prefix_len + - off_sub[None, :] + ) + qk_sub = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, SUB_K), dtype=tl.float32) + # causal: q_abs_pos - k_off >= block_start (c) + qk_sub += tl.where(off_q_sub[:, None, :] >= c, 0, float("-inf")) + qk_sub = tl.reshape(qk_sub, BLOCK_SIZE_QH, SUB_K) + qk_sub += tl.dot(q, k_sub) * sm_scale_log2e + qk_sub += tl.where(pos_mask_sub[None, :], 0, float("-inf")) + m_ij = tl.maximum(m_i, tl.max(qk_sub, axis=1)) + p_sub = tl.exp2(qk_sub - m_ij[:, None]) + l_ij = tl.sum(p_sub, axis=1) + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + v_sub = tl.load( + kv_base + + 1 * stride_kv_kv + + off_sub[:, None] * stride_kv_pos + + off_d[None, :] * stride_kv_d, + mask=pos_mask_sub[:, None] & d_mask[None, :], + other=0.0, + ) + if USE_FP8: + v_sub = v_sub.to(q.dtype) + acc_o += tl.dot(p_sub.to(v_sub.dtype), v_sub) + m_i = m_ij + 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_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2)) + + +@torch.no_grad() +def minimax_m3_sparse_attn( + q: torch.Tensor, # [total_q, num_heads, head_dim] + kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim] + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [batch, max_blocks] + cu_seqlens_q: torch.Tensor, # [batch+1] int32 + seq_lens: torch.Tensor, # [batch] int32 + prefix_lens: torch.Tensor, # [batch] int32 + max_query_len: int, + num_kv_heads: int, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] +) -> None: + """GQA block-sparse attention over the selected blocks. block_size_q == 1.""" + total_q, num_heads, head_dim = q.shape + batch = cu_seqlens_q.shape[0] - 1 + topk = topk_idx.shape[-1] + gqa_group_size = num_heads // num_kv_heads + use_fp8 = kv_cache.dtype in _FP8_DTYPES + grid = (max_query_len, num_kv_heads, batch) + _gqa_sparse_fwd_kernel[grid]( + q, + kv_cache, + topk_idx, + output, + block_table, + cu_seqlens_q, + cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1 + seq_lens, + prefix_lens, + num_kv_heads, + gqa_group_size, + head_dim, + topk, + 1, # num_q_loop + sm_scale, + q.stride(0), + q.stride(1), + q.stride(2), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + kv_cache.stride(3), + kv_cache.stride(4), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + output.stride(0), + output.stride(1), + output.stride(2), + block_table.stride(0), + BLOCK_SIZE_Q=1, + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + USE_FP8=use_fp8, + SUB_K=_SPARSE_ATTN_SUB_K, + **_sparse_attn_prefill_kwargs(), + ) diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index 4a72b6bc2c9..bb1ed619320 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -27,12 +27,21 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.models.minimax_m3.common.ops.index_topk import ( - minimax_m3_index_decode, - minimax_m3_index_score, - minimax_m3_index_topk, -) from vllm.platforms import current_platform + +if current_platform.is_rocm(): + from vllm.models.minimax_m3.amd.ops.index_topk import ( + minimax_m3_index_decode, + minimax_m3_index_score, + minimax_m3_index_topk, + ) +else: + from vllm.models.minimax_m3.common.ops.index_topk import ( + minimax_m3_index_decode, + minimax_m3_index_score, + minimax_m3_index_topk, + ) + from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index 7b6fb73cba9..f04652c89a3 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -31,30 +31,6 @@ _FP8_DTYPES = ( torch.float8_e5m2fnuz, ) -_SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None - - -def _sparse_attn_num_stages_kwarg() -> dict: - """Triton ``num_stages`` override for the sparse-attn GEMM kernels. - - Forced only where required: CDNA3 (gfx942) caps LDS at - 64 KB, and the default 2-stage pipeline double-buffers the 128x128 K/V tiles - to ~66 KB ("out of resource: shared memory"), so pin gfx942 to a single - stage (~32 KB, which fits). Everywhere else (NVIDIA, CDNA4 gfx950) return an - empty kwarg and let Triton keep its own default -- don't second-guess it. - Cached: the arch is fixed per process. - """ - global _SPARSE_ATTN_NUM_STAGES_KWARG - if _SPARSE_ATTN_NUM_STAGES_KWARG is None: - kwarg: dict = {} - if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx942 - - if on_gfx942(): - kwarg = {"num_stages": 1} - _SPARSE_ATTN_NUM_STAGES_KWARG = kwarg - return _SPARSE_ATTN_NUM_STAGES_KWARG - # --------------------------------------------------------------------------- # GQA block-sparse attention (paged). Main heads attend only to the selected @@ -498,7 +474,6 @@ def minimax_m3_sparse_attn( BLOCK_SIZE_Q=1, BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, USE_FP8=use_fp8, - **_sparse_attn_num_stages_kwarg(), ) @@ -574,7 +549,6 @@ def minimax_m3_sparse_attn_decode( NUM_TOPK_CHUNKS=num_topk_chunks, USE_FP8=use_fp8, USE_PDL=use_pdl, - **_sparse_attn_num_stages_kwarg(), **pdl_launch, ) merge_grid = (total_q, num_heads) diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 55542230885..88109772eb1 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -24,12 +24,21 @@ from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.minimax_m3.common.ops.sparse_attn import ( - SPARSE_BLOCK_SIZE, - minimax_m3_sparse_attn, - minimax_m3_sparse_attn_decode, -) +from vllm.models.minimax_m3.common.ops.sparse_attn import SPARSE_BLOCK_SIZE from vllm.platforms import current_platform + +# AMD/ROCm uses the gfx942/gfx950-optimized block-sparse kernels in amd.ops; +# every other platform uses the generic common.ops implementation. +if current_platform.is_rocm(): + from vllm.models.minimax_m3.amd.ops.sparse_attn import ( + minimax_m3_sparse_attn, + minimax_m3_sparse_attn_decode, + ) +else: + from vllm.models.minimax_m3.common.ops.sparse_attn import ( + minimax_m3_sparse_attn, + minimax_m3_sparse_attn_decode, + ) from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport,