karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[ROCm][ [Perf] sparse attention optimization on minimax-m3 (#46546)

Browse Source 
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: yueliu14 <yue.liu4@amd.com>
Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
This commit is contained in:
Hongxia Yang
2026-06-25 16:56:00 +08:00
committed by GitHub
co-authored by Claude Opus 4.8 yueliu14 tjtanaa
parent 638b1a99cc
commit c63cd4906c
5 changed files with 1238 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/minimax_m3/amd/ops/index_topk.py
+939
View File
@@ -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
vllm/models/minimax_m3/amd/ops/sparse_attn.py
+271
View File
@@ -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(),
)
vllm/models/minimax_m3/common/indexer.py
+14 -5
View File
@@ -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,
vllm/models/minimax_m3/common/ops/sparse_attn.py
-26
View File
@@ -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)
vllm/models/minimax_m3/common/sparse_attention.py
+14 -5
View File
@@ -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,