Skip to content

vllm.models.minimax_m3.amd.ops.index_topk_decode.prefill_kernel

Index top-k for MiniMax-M3 (FlyDSL, gfx950).

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.

_below(mask, base)

The set bits of mask below this lane, plus base (v_mbcnt).

Source code in vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py 
53
54
55
56
57
58
def _below(mask, base):
    """The set bits of ``mask`` below this lane, plus ``base`` (v_mbcnt)."""
    lo = mask.to(fx.Int32)
    hi = (mask >> 32).to(fx.Int32)
    r = fx.rocdl.mbcnt_lo(T.i32, lo.ir_value(), fx.Int32(base).ir_value())
    return fx.Int32(fx.rocdl.mbcnt_hi(T.i32, hi.ir_value(), r))

_comp(k, i)

(key, index) as one i64 whose signed order is AITER's.

Source code in vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py 
69
70
71
def _comp(k, i):
    """(key, index) as one i64 whose signed order is AITER's."""
    return (k.to(fx.Int64) << 32) | (i.to(fx.Int64) & 0xFFFFFFFF)

_key(bits)

Signed key in AITER's order; NaNs land above +inf or below -inf.

Source code in vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py 
33
34
35
def _key(bits):
    """Signed key in AITER's order; NaNs land above +inf or below -inf."""
    return bits ^ ((bits >> 31) & 0x7FFFFFFF)

_popc(m)

Set bits of a wave mask; two 32-bit counts keep the compare scalar.

Source code in vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py 
46
47
48
49
50
def _popc(m):
    """Set bits of a wave mask; two 32-bit counts keep the compare scalar."""
    lo = fx.math.ctpop(m.to(fx.Int32)).to(fx.Int32)
    hi = fx.math.ctpop((m >> 32).to(fx.Int32)).to(fx.Int32)
    return fx.Int32(lo) + fx.Int32(hi)