From 3be332fe4194e07f615aa49341e6cb15e71b89d2 Mon Sep 17 00:00:00 2001 From: chao-xu-spec Date: Sat, 31 Oct 2026 12:00:05 +0000 Subject: [PATCH] [ROCm][Perf][MiniMax-M3.1] Four-wave decode top-k for the NVFP4 indexer The decode top-k walked each row with one wave. Run four waves per row, each over every fourth window with the next window already loading, and merge their sorted top-16 lists by an exact rank over (key, index), which keeps AITER's order including ties, NaN and infinities. Prefill is unchanged. Extend the decode tests with 100 requests and the window edges where the waves meet. Co-authored-by: Claude Opus 5.5 Co-authored-by: Claude Sonnet 5.5 Signed-off-by: chao-xu-spec --- .../test_minimax_m3_index_topk_decode.py | 3 + .../amd/ops/index_topk_decode/host.py | 2 +- .../ops/index_topk_decode/prefill_kernel.py | 162 ++++++++++++++---- 3 files changed, 130 insertions(+), 37 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3_index_topk_decode.py b/tests/kernels/attention/test_minimax_m3_index_topk_decode.py index afbd90c009f3..9bfdf580cfcb 100644 --- a/tests/kernels/attention/test_minimax_m3_index_topk_decode.py +++ b/tests/kernels/attention/test_minimax_m3_index_topk_decode.py @@ -111,6 +111,9 @@ def _aiter_topk(score, bt, seq, msl, nkv, pad, **kw): (8, 1, 200000, (200000, 262000), 4, 1), (4, 2, 0, (0, 0), 1, 1), # a zero max_seq_len (dummy batch) (1, 1, 1, (0, 1), 1, 1), + (100, 4, 1 << 20, (1, 1 << 20), 1, 1), # 100 requests, every window count + (16, 1, 1 << 20, (1020 * BLOCK, 1030 * BLOCK), 1, 1), # waves 0 and 1 meet + (16, 1, 1 << 20, (4090 * BLOCK, 4100 * BLOCK), 1, 1), # second trip of the waves ] KINDS = [ "real", diff --git a/vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py b/vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py index f3669b733c96..03472dc079f3 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py +++ b/vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py @@ -27,7 +27,7 @@ def _get_prefill_kernel(strips: int): @functools.lru_cache(maxsize=1) def _get_decode_kernel(): - return _kernel.compile_topk_prefill(26, decode=True) + return _kernel.compile_topk_prefill(16, decode=True) def supports_index_topk(topk_blocks: int, block_size: int) -> bool: diff --git a/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py b/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py index b18db79e055a..d4dd66723cba 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py +++ b/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py @@ -4,6 +4,10 @@ One 64-lane workgroup per (row, head) returns the 16 highest-scoring block ids in AITER's order, walking the row in windows of STRIPS x 64 blocks. + +Decode gives each row four waves: wave w walks windows w, w + 4, ... with its +own candidates and running top 16, then the four lists merge by an exact rank +over the (key, index) order. """ import flydsl.compiler as flyc @@ -72,8 +76,10 @@ def _comp(k, i): def compile_topk_prefill(strips: int, *, decode: bool = False): - assert strips == 26 if decode else strips in PREFILL_STRIPS + assert strips == 16 if decode else strips in PREFILL_STRIPS STRIPS = strips + WAVES = 4 if decode else 1 + NTHR = 64 * WAVES WIN = STRIPS * 64 RANK_BLOCK = 8 # Candidates, rank-loop padding, then one discard slot per lane. @@ -83,10 +89,10 @@ def compile_topk_prefill(strips: int, *, decode: bool = False): @fx.struct class Shared: - cand: fx.Array[fx.Int32, 2 * CAND, 16] - sel: fx.Array[fx.Int32, 4 * TOPK, 16] + cand: fx.Array[fx.Int32, WAVES * 2 * CAND, 16] + sel: fx.Array[fx.Int32, WAVES * 4 * TOPK, 16] - @flyc.kernel(name=name, known_block_size=[64, 1, 1]) + @flyc.kernel(name=name, known_block_size=[NTHR, 1, 1]) def kernel( arg_score: fx.Int64, i64_s_head: fx.Int64, @@ -100,7 +106,19 @@ def kernel( ): smem = fx.SharedAllocator().allocate(Shared).peek() n, h = fx.block_idx.x, fx.block_idx.y - lane = fx.thread_idx.x + if const_expr(decode): + tid = fx.Int32(fx.thread_idx.x) + lane = tid & 63 + # Wave index, scalar. + wv = fx.Int32(fx.rocdl.readfirstlane(T.i32, (tid >> 6).ir_value())) + c_off = wv * CAND # this wave's candidate pairs + s_off = wv * TOPK # this wave's running top 16 + w_off = wv * WIN # this wave's first window + else: + lane = fx.thread_idx.x + c_off = 0 + s_off = 0 + w_off = 0 n64 = fx.Int32(n).to(fx.Int64) h64 = fx.Int32(h).to(fx.Int64) @@ -131,7 +149,7 @@ def load_window(srs, wbase): ) # The first window is loaded before the row's length is known. - bits = load_window(srs, 0) + bits = load_window(srs, w_off) def scalar_at(base, idx64, nbytes): return fx.Int32( @@ -161,9 +179,9 @@ def view(ptr, n_dw, width): fx.make_layout(width, 1), ) - pairs = view(smem.cand.ptr, 2 * CAND, 2) - quads = view(smem.cand.ptr, 2 * CAND, 4) - sels = view(smem.sel.ptr, 4 * TOPK, 4) + pairs = view(smem.cand.ptr, WAVES * 2 * CAND, 2) + quads = view(smem.cand.ptr, WAVES * 2 * CAND, 4) + sels = view(smem.sel.ptr, WAVES * 4 * TOPK, 4) atom64 = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32) atom128 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) @@ -179,7 +197,7 @@ def lds_get(atom, v, w): return [fx.Int32(x[e]) for e in range_constexpr(w)] def put_pair(pos, i, k): - lds_put(atom64, fx.slice(pairs, (None, pos)), 2, [i, k]) + lds_put(atom64, fx.slice(pairs, (None, c_off + pos)), 2, [i, k]) def wave_sum(c): """Sum over the wave of a per-lane count in [0, 63].""" @@ -242,9 +260,11 @@ def exact_cut(keys, wbase): b = b - 1 return v, ie - def lane_max_bound(m): - p = fx.Int32(0) - for d in range_constexpr(12): + def lane_max_bound(m, d1=12, d0=0, key=INT_MIN): + """The 16th largest lane maximum m, to its top d1 bits (from d0 on, + above the bound ``key`` of the top d0).""" + p = fx.Int32(key ^ INT_MIN) + for d in range_constexpr(d0, d1): t = p | _s32(1 << (31 - d)) count = _popc(_ballot(m >= (t ^ INT_MIN))) p = (count >= TOPK).select(t, p) @@ -260,15 +280,25 @@ def window(bits, wbase, nrun, tprev): for j in range_constexpr(1, STRIPS): m = _imax(m, keys[j]) # Bound: the lane maxima in the first window, the running 16th later. - tw = fx.Int32(INT_MIN) - if nrun < TOPK: - tw = lane_max_bound(m) + if const_expr(decode): + tw = lane_max_bound(m, 16) # every window; 16 bits, not 12 + else: + tw = fx.Int32(INT_MIN) + if nrun < TOPK: + tw = lane_max_bound(m) ta = _imax(tw, tprev) total = compact(keys, wbase, nrun, lambda k, i: k >= ta) + if const_expr(decode): # noqa: SIM102 + if total > CAP: # too loose: the bound to all 32 bits + ta = _imax(lane_max_bound(m, 32, 16, tw), tprev) + total = compact(keys, wbase, nrun, lambda k, i: k >= ta) # Exact only if 16 keys above -inf bound it and no NaN is above +inf. nan_hi = _ballot(m > KEY_POS_INF) != fx.Int64(0) if (total > CAP) | (ta <= KEY_NEG_INF) | nan_hi: wl = nlive - wbase + if const_expr(decode): + # A wave's windows past the row's end are empty. + wl = (wl < 0).select(fx.Int32(0), wl) if wl <= CAP - nrun: # Every live block of the window is a candidate. put_pair( @@ -299,14 +329,18 @@ def window(bits, wbase, nrun, tprev): if lane < RANK_BLOCK: put_pair(total + lane, fx.Int32(0), fx.Int32(INT_MIN)) fx.gpu.barrier() - mi, mk = lds_get(atom64, fx.slice(pairs, (None, lane)), 2) + mi, mk = lds_get(atom64, fx.slice(pairs, (None, c_off + lane)), 2) mw = _comp(mk, mi) rank = fx.Int32(0) for g in range_constexpr(CAP // RANK_BLOCK): if total > RANK_BLOCK * g: es = [ lds_get( - atom128, fx.slice(quads, (None, RANK_BLOCK // 2 * g + u)), 4 + atom128, + fx.slice( + quads, (None, c_off // 2 + RANK_BLOCK // 2 * g + u) + ), + 4, ) for u in range_constexpr(RANK_BLOCK // 2) ] @@ -316,18 +350,29 @@ def window(bits, wbase, nrun, tprev): rank = rank + _one(_comp(e[1], e[0]) > mw) rank = rank + _one(_comp(e[3], e[2]) > mw) if (lane < total) & (rank < TOPK): - lds_put(atom128, fx.slice(sels, (None, rank)), 4, [mi, mk, 0, mi]) + lds_put( + atom128, fx.slice(sels, (None, s_off + rank)), 4, [mi, mk, 0, mi] + ) return (total < TOPK).select(total, fx.Int32(TOPK)) - wbase = fx.Int32(0) + # Every wave runs the same trip count (the row's), so the barriers match. + gbase = fx.Int32(0) nrun = fx.Int32(0) tprev = fx.Int32(INT_MIN + 1) # dead keys are INT_MIN - while (wbase == 0) | (wbase < nlive): - if wbase > 0: - bits = load_window(srs_live, wbase) + while (gbase == 0) | (gbase < nlive): + wbase = gbase + w_off + if const_expr(decode): + # The next trip's window is in flight while this one ranks; + # past the row's end it reads nothing. + nxt = load_window(srs_live, wbase + WAVES * WIN) + if gbase > 0: + if const_expr(not decode): + bits = load_window(srs_live, wbase) # The running top 16 back into cand; its 16th key bounds the window. fx.gpu.barrier() - r = lds_get(atom128, fx.slice(sels, (None, lane & (TOPK - 1))), 4) + r = lds_get( + atom128, fx.slice(sels, (None, s_off + (lane & (TOPK - 1)))), 4 + ) if lane < nrun: put_pair(lane, r[0], r[1]) k15 = fx.Int32(fx.rocdl.readlane(T.i32, r[1], fx.Int32(TOPK - 1))) @@ -339,17 +384,62 @@ def window(bits, wbase, nrun, tprev): nrun, tprev, ) - wbase = wbase + WIN - - fx.gpu.barrier() - r = lds_get(atom128, fx.slice(sels, (None, lane & (TOPK - 1))), 4) - blk = (lane < nrun).select(r[0], fx.Int32(-1)) - buffer_ops.buffer_store( - blk, - rsrc(arg_topk + (h64 * i64_t_head + n64 * i64_t_row) * 4, TOPK * 4), - lane, - mask=lane < TOPK, - ) + if const_expr(decode): + bits = nxt + gbase = gbase + WAVES * WIN + + out_rs = rsrc(arg_topk + (h64 * i64_t_head + n64 * i64_t_row) * 4, TOPK * 4) + if const_expr(not decode): + fx.gpu.barrier() + r = lds_get(atom128, fx.slice(sels, (None, lane & (TOPK - 1))), 4) + blk = (lane < nrun).select(r[0], fx.Int32(-1)) + buffer_ops.buffer_store(blk, out_rs, lane, mask=lane < TOPK) + else: + # Pad each wave's list with dead entries, then merge the four sorted + # lists: an entry's slot is its rank among all the entries. + if (lane < TOPK) & (lane >= nrun): + lds_put( + atom128, + fx.slice(sels, (None, s_off + lane)), + 4, + [0, INT_MIN, 0, 0], + ) + fx.gpu.barrier() + if nlive <= WIN: + # Only wave 0 saw live blocks: its list is the answer. + if tid < TOPK: + w0 = lds_get(atom128, fx.slice(sels, (None, tid)), 4) + live0 = w0[1] > INT_MIN + buffer_ops.buffer_store( + live0.select(w0[0], fx.Int32(-1)), out_rs, tid, mask=tid < TOPK + ) + elif tid < TOPK * WAVES: + e = lds_get(atom128, fx.slice(sels, (None, tid)), 4) + cv = _comp(e[1], e[0]) + # Each list is sorted: binary-lift its count above cv. + cs = [fx.Int32(0) for _ in range_constexpr(WAVES)] + for st in (16, 8, 4, 2, 1): + prb = [cs[w] + (st - 1) for w in range_constexpr(WAVES)] + es = [ + lds_get( + atom128, + fx.slice(sels, (None, w * TOPK + (prb[w] & (TOPK - 1)))), + 4, + ) + for w in range_constexpr(WAVES) + ] + fx.rocdl.sched_barrier(0) + for w in range_constexpr(WAVES): + hit = (prb[w] < TOPK) & (_comp(es[w][1], es[w][0]) > cv) + cs[w] = cs[w] + hit.select(fx.Int32(st), fx.Int32(0)) + rank = cs[0] + for w in range_constexpr(1, WAVES): + rank = rank + cs[w] + live = e[1] > INT_MIN + buffer_ops.buffer_store(e[0], out_rs, rank, mask=live & (rank < TOPK)) + tot = _popc(_ballot(live)) # tid < 64 is wave 0 + if (tid < TOPK) & (tid >= tot): + buffer_ops.buffer_store(fx.Int32(-1), out_rs, tid, mask=tid < TOPK) @flyc.jit def launch( @@ -378,7 +468,7 @@ def launch( i64_t_row, ).launch( grid=(fx.Int64(i32_rows), fx.Int64(i32_heads), 1), - block=(64, 1, 1), + block=(NTHR, 1, 1), stream=stream, )