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, )