vllm.models.minimax_m3.amd.specialized_decode.mono.ffn

MiniMax-M3 FFN launch: one persistent kernel from the unreduced o_proj partial to the reduced MoE output (MiniMaxM3DecoderLayer.forward after its attention).

BLOCKS x THREADS, S rows, one launch per sparse layer and rank:

push    (192 x 32 columns): the partial's bf16 pairs to every rank
reduce  (192): the ranks' pairs in rank order, f32 -> bf16; + residual
        in f32 -> residual out (bf16), the f32 sum, its sum of squares
norm    (S): Gemma RMSNorm of the f32 sum -> X (bf16)
router  (8 x 16 experts x 16 K parts): FP32 weight x f32(X), f32 MFMA
route   (S): the parts summed in order, sigmoid, + bias for selection,
        top-4, renormalized x 2.0
table   (every CTA): the union of picked experts in LDS, slot 0 the shared
        expert (expert 128) over all S rows
up      (48 x slots): 16 gate + 16 up rows of a slot's expert, A16W4 MFMA,
        SwiGLU-OAI -> bf16 MID
down    (384 x 16 rows): every slot's A16W4 GEMV of MID, routed weights
        and shared expert combined -> pushed to every rank
final   (192): the ranks' partials in rank order -> ``out``

Classes:

• Experts –

  Which expert terms the combine sums: all of them in production; one side

Experts

Bases: Enum

Which expert terms the combine sums: all of them in production; one side alone for stage parity.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

class Experts(Enum):
    """Which expert terms the combine sums: all of them in production; one side
    alone for stage parity."""

    ALL = "all"
    SHARED = "shared"
    ROUTED = "routed"

await_x(c, tokens_n)

XRDY of the first tokens_n tokens (thread t polls token t).

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def await_x(c, tokens_n):
    """XRDY of the first ``tokens_n`` tokens (thread t polls token t)."""
    tid = c["tid"]
    if tid < tokens_n:
        c["mb"].poll([(c["xrdy"], tid, 1)], site=SITE_ROUTER)
        c["mb"].acquire(Space.SCRATCH)
    cta_sync()

bf16_bits_f32(bits)

A bf16's 16 bits (the low half of an int) as f32.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

def bf16_bits_f32(bits):
    """A bf16's 16 bits (the low half of an int) as f32."""
    return ((bits & 0xFFFF) << 16).bitcast(fx.Float32)

block_sum(c, v, slot)

v summed over the CTA: each wave's butterfly, then waves in order (LDS red words slot * WAVES ..). Every thread gets the sum.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def block_sum(c, v, slot):
    """``v`` summed over the CTA: each wave's butterfly, then waves in order
    (LDS ``red`` words ``slot * WAVES ..``). Every thread gets the sum."""
    lane, wave, red = c["lane"], c["wave"], c["red"]
    v = butterfly(v, (32, 16, 8, 4, 2, 1))
    if lane == 0:
        fx.ptr_store(v, red + (slot * WAVES + wave))
    gpu.barrier()
    tot = fx.Float32(0.0)
    for w in range_constexpr(WAVES):
        tot = tot + fx.ptr_load(red + (slot * WAVES + w))
    return tot

epoch_end(c, epoch, tag)

CTA 0, once every CTA marked this launch: the epoch moves on. A mark missing past the wait bound sets the error word instead; the epoch stays.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def epoch_end(c, epoch, tag):
    """CTA 0, once every CTA marked this launch: the epoch moves on. A mark
    missing past the wait bound sets the error word instead; the epoch stays."""
    tid, red = c["tid"], c["red"]
    if fx.block_idx.x == 0:
        if tid == 0:
            fx.ptr_store(fx.Int32(0).bitcast(fx.Float32), red)
        gpu.barrier()
        if tid < BLOCKS:
            t0 = memrealtime()
            mark = load32(epoch, EPOCH_MARKS + tid, "agent")
            late = fx.Int64(0)
            while (mark != tag) & (late < c["mb"].ticks):
                rocdl.s_sleep(1)
                mark = load32(epoch, EPOCH_MARKS + tid, "agent")
                late = memrealtime() - t0
            if mark != tag:
                fx.ptr_store(fx.Int32(1).bitcast(fx.Float32), red)
            sync_acquire("agent")
        cta_sync()
        if tid == 0:
            if fx.ptr_load(red).bitcast(fx.Int32) == 0:
                store32(epoch, 0, tag, "agent")
            else:
                c["mb"].fail(SITE_EPOCH)

finish(c, key, epoch, tag, handshake, inject, peers)

Every CTA marks the launch done after its work; CTA 0 waits for the marks, then moves the epoch on. The step's last launch (handshake) then trades every rank's failure flag over the peers: a fault anywhere in the step is in every rank's error word before this launch ends.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def finish(c, key, epoch, tag, handshake, inject, peers):
    """Every CTA marks the launch done after its work; CTA 0 waits for the
    marks, then moves the epoch on. The step's last launch (``handshake``)
    then trades every rank's failure flag over the peers: a fault anywhere in
    the step is in every rank's error word before this launch ends."""
    tid, bid = c["tid"], fx.block_idx.x
    if (handshake != 0) & (bid == 0) & (tid == 0):
        injected_fault(c, inject)
    cta_sync(drain=True)
    if (tid == 0) & (bid != key.drop_mark):
        sync_release("agent")
        store32(epoch, EPOCH_MARKS + bid, tag, "agent")
    epoch_end(c, epoch, tag)
    if (handshake != 0) & (bid == 0) & (tid == 0):
        trade_faults(c, tag, peers)

flag_after_stores(c, addr, i, who)

Pair i := 1 (thread who), released after the CTA's earlier stores.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def flag_after_stores(c, addr, i, who):
    """Pair i := 1 (thread ``who``), released after the CTA's earlier stores."""
    cta_sync(drain=True)
    if who:
        c["mb"].publish(addr, i, fx.Int32(1))

injected_fault(c, inject)

Fault injection (qualification): rank inject[0] - 1 fails after inject[1] ticks at the end of the step's last launch.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def injected_fault(c, inject):
    """Fault injection (qualification): rank ``inject[0] - 1`` fails after
    ``inject[1]`` ticks at the end of the step's last launch."""
    target = load32(inject, 0, "agent")
    if target == c["rank"] + 1:
        ticks = fx.Int64(load32(inject, 1, "agent"))
        t0 = memrealtime()
        while memrealtime() - t0 < ticks:
            rocdl.s_sleep(8)
        c["mb"].fail(SITE_INJECT)

load_table(c)

Every valid token's picks -> the slots in LDS: slot 0 the shared expert over the valid rows, slots 1 .. n_u the picked experts ascending; a slot's columns its rows ascending, each column's pick (t * 4 + k; the shared expert's 4 S + t). A padded row's picks are -1: in no slot.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def load_table(c):
    """Every valid token's picks -> the slots in LDS: slot 0 the shared expert
    over the valid rows, slots 1 .. n_u the picked experts ascending; a slot's
    columns its rows ascending, each column's pick (t * 4 + k; the shared
    expert's 4 S + t). A padded row's picks are -1: in no slot."""
    s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"]
    n = s * TOPK
    if tid < n:
        t = tid // TOPK
        got = c["mb"].poll(
            [
                (c["route"], t * 2 * TOPK + tid % TOPK, 1),
                (c["route"], t * 2 * TOPK + TOPK + tid % TOPK, 1),
            ],
            site=SITE_TABLE,
        )
        c["mb"].acquire(Space.SCRATCH)
        tab_st(c, got[0][0], T_PICK_E + tid)
        tab_st(c, got[1][0], T_PICK_W + tid)
    if tid < EXPERTS:
        tab_st(c, 0, T_MASK + tid)
    if tid < WAVES:
        tab_st(c, 0, T_WAVE + tid)
    gpu.barrier()
    if tid < n:
        e = tab_ld(c, T_PICK_E + tid)
        if e >= 0:
            fx.llvm.atomic_or(
                c["tab"] + (T_MASK + e),
                fx.Int32(1) << (tid // TOPK),
                syncscope="workgroup",
            )
    gpu.barrier()
    e = tid
    m = fx.Int32(0)
    if tid < EXPERTS:
        m = tab_ld(c, T_MASK + e)
    ball = ballot(m != 0)
    if (wave == 0) & (lane == 0):
        tab_st(c, popcount(ball), T_MISC + 1)
    gpu.barrier()
    below = (wave == 1).select(tab_ld(c, T_MISC + 1), fx.Int32(0))
    slot = 1 + below + lanes_below(ball, lane)
    if tid < EXPERTS:
        tab_st(c, slot, T_ESLOT + e)
        if m != 0:
            tab_st(c, e, T_SLOT_E + slot)
            tab_st(c, m, T_SLOT_MASK + slot)
    if (wave == 1) & (lane == 0):
        tab_st(c, below + popcount(ball), T_MISC)
    shared_slot(c)
    if tid < n:
        t = tid // TOPK
        e = tab_ld(c, T_PICK_E + tid)
        if e >= 0:
            sl = tab_ld(c, T_ESLOT + e)
            col = popcount(tab_ld(c, T_SLOT_MASK + sl) & ((fx.Int32(1) << t) - 1))
            tab_st(c, t, T_SLOT_TOK + sl * 16 + col)
            tab_st(c, tid, T_SLOT_PICK + sl * 16 + col)
    gpu.barrier()

peer_half_bytes(tp)

One parity's peer regions: the attention partials (triple pairs), then the MoE's (bf16-pair pairs).

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

def peer_half_bytes(tp: int) -> int:
    """One parity's peer regions: the attention partials (triple pairs), then
    the MoE's (bf16-pair pairs)."""
    return tp * MAX_TOKENS * (TRIPLES + HIDDEN // 2) * 8

put_word(c, addr, i, word)

Pair i of addr := (word, tag) at its space's scope.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

def put_word(c, addr, i, word):
    """Pair i of ``addr`` := (word, tag) at its space's scope."""
    c["mb"].put_words(addr, i, [word])

shared_slot(c)

Slot 0: the shared expert over the valid rows (columns ascending; an unused column reads row 0).

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def shared_slot(c):
    """Slot 0: the shared expert over the valid rows (columns ascending; an
    unused column reads row 0)."""
    s, tid = c["S"], c["tid"]
    if tid == 0:
        tab_st(c, SHARED, T_SLOT_E)
        tab_st(c, c["vmask"], T_SLOT_MASK)
    if tid < 16:
        tab_st(c, 0, T_SLOT_TOK + tid)
    gpu.barrier()
    if (tid < s) & row_valid(c, fx.min(tid, s - 1)):
        col = popcount(c["vmask"] & ((fx.Int32(1) << tid) - 1))
        tab_st(c, tid, T_SLOT_TOK + col)
        tab_st(c, TOPK * s + tid, T_SLOT_PICK + col)
    gpu.barrier()

slot_col(c, cnt)

The lane's column of a slot of cnt columns (past them: the last).

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

def slot_col(c, cnt):
    """The lane's column of a slot of ``cnt`` columns (past them: the last)."""
    return fx.max(fx.min(c["lane"] % 16, cnt - 1), 0)

stage_down(c, tasks)

Rows 16 task .. of every task in tasks: a slot a wave, round robin, its MID columns loaded once for every task; a column's result -> LDS CONTRIB[task][pick], then each token's routed picks (top-k order) x their weights + the shared expert -> bf16 pairs pushed to every rank's FFN region.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def stage_down(c, tasks):
    """Rows 16 task .. of every task in ``tasks``: a slot a wave, round robin,
    its MID columns loaded once for every task; a column's result ->
    LDS ``CONTRIB[task][pick]``, then each token's routed picks (top-k order) x
    their weights + the shared expert -> bf16 pairs pushed to every rank's FFN
    region."""
    s, tid, lane, wave, a = c["S"], c["tid"], c["lane"], c["wave"], c["args"]
    nt = len(tasks)
    nu = n_slots(c)
    for u in range(wave, nu, WAVES):
        e = tab_ld(c, T_SLOT_E + u)
        rgs = [e * DOWN_TASKS + task for task in tasks]
        wvs = [weight_loads(a["w2"], rg, INTER, 0, K2_STEPS) for rg in rgs]
        svs = [scale_loads(a["s2"], rg, INTER, 0, K2_STEPS) for rg in rgs]
        cnt = slot_count(c, u)
        c["mb"].poll(
            [(c["upf"], u * UP_TILES + fx.min(lane, UP_TILES - 1), 1)], site=SITE_DOWN
        )
        c["mb"].acquire(Space.SCRATCH)
        # a lane's MID loads span tiles other lanes polled
        wave_sync(c["tab"] + (T_WAVE + wave), 64 * ((u - wave) // WAVES + 1))
        col = slot_col(c, cnt)
        xs = x_operands(c, c["mid"].value, u * 16 + col, INTER // 2, 0, K2_STEPS)
        pick = tab_ld(c, T_SLOT_PICK + u * 16 + col)
        for ti in range_constexpr(nt):
            acc = fx.Vector.filled(4, 0.0, fx.Float32)
            for i in range_constexpr(K2_STEPS):
                acc = step_mfma(acc, wvs[ti][i], scale_f32(svs[ti], rgs[ti], i), xs[i])
            if lane % 16 < cnt:
                for q in range_constexpr(4):
                    fx.ptr_store(
                        acc[q],
                        c["contrib"] + ((ti * 80 + pick) * 16 + 4 * (lane // 16) + q),
                    )
    gpu.barrier()
    if tid < nt * s * 8:
        ti = tid // (s * 8)
        t = tid % (s * 8) // 8
        r2 = tid % 8
        task = tasks[0]
        for q in range_constexpr(1, nt):
            task = (ti == q).select(tasks[q], task)
        # AITER's stage 2: bf16(weight x acc) a pick, added in bf16 (its
        # atomics' order is a race; here the shared expert, then top-k order)
        vals = []
        for h in range_constexpr(2):
            r = 2 * r2 + h
            tot = fx.Float32(0.0)
            if const_expr(c["experts"] != Experts.ROUTED):
                tot = bf16_round(
                    fx.ptr_load(c["contrib"] + ((ti * 80 + TOPK * s + t) * 16 + r))
                )
            if const_expr(c["experts"] != Experts.SHARED):
                for k in range_constexpr(TOPK):
                    w = tab_ld(c, T_PICK_W + t * TOPK + k).bitcast(fx.Float32)
                    v = fx.ptr_load(c["contrib"] + ((ti * 80 + t * TOPK + k) * 16 + r))
                    tot = bf16_round(tot + bf16_round(w * v))
            vals.append(row_valid(c, t).select(tot, fx.Float32(0.0)))
        word = bf16_pair(vals[0], vals[1]).bitcast(fx.Int32)
        cp = task * 8 + r2
        for p in range_constexpr(c["tp"]):
            put_word(c, c["peer"][p], ffn_pair(c, c["rank"], t, cp), word)
    gpu.barrier()

stage_norm(c, t)

Token t's Gemma RMSNorm of SRES -> XN, then XRDY.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def stage_norm(c, t):
    """Token t's Gemma RMSNorm of SRES -> XN, then XRDY."""
    tid, a = c["tid"], c["args"]
    part = fx.Float32(0.0)
    if tid < SEAM_TASKS:
        got = c["mb"].poll([(c["ssq"], t * SEAM_TASKS + tid, 1)], site=SITE_NORM)
        part = got[0][0].bitcast(fx.Float32)
        c["mb"].acquire(Space.SCRATCH)
    cta_sync()
    var = block_sum(c, part, 0) / fx.Float32(HIDDEN)
    rstd = fx.Float32(1.0) / fx.Float32(fmath.sqrt(var + fx.Float32(EPS)))
    per = HIDDEN // THREADS  # 12
    for q in range_constexpr(per // 2):
        col = (q * THREADS + tid) * 2
        sv = load_dev(c["sres"].value, t * HIDDEN + col, 2)
        ww = fx.Int32(
            bo.buffer_load(rsrc(a["norm_w"]), col // 2, vec_width=1, dtype=T.i32)
        )
        wl, wh = bf2_f32(ww)
        ol = fx.Int32(sv[0]).bitcast(fx.Float32) * rstd * (fx.Float32(1.0) + wl)
        oh = fx.Int32(sv[1]).bitcast(fx.Float32) * rstd * (fx.Float32(1.0) + wh)
        store_dev(
            c["xn"].value,
            t * (HIDDEN // 2) + col // 2,
            [bf16_pair(ol, oh).bitcast(fx.Int32)],
        )
    flag_after_stores(c, c["xrdy"], t, tid == 0)

stage_push(c, task)

Triples 16 task ..: this rank's partial as triple pairs to every rank's ATTN.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def stage_push(c, task):
    """Triples 16 task ..: this rank's partial as triple pairs to every rank's
    ATTN."""
    s, tid = c["S"], c["tid"]
    if tid < s * TASK_TRIPLES:
        t = tid // TASK_TRIPLES
        tr = task * TASK_TRIPLES + tid % TASK_TRIPLES
        bits = triple_bits(c["args"]["part"], t, tr)
        for p in range_constexpr(c["tp"]):
            c["mb"].put_triple(c["peer"][p], attn_pair(c, c["rank"], t, tr), bits)

stage_reduce(c, task)

Triples 16 task ..: the ranks' values summed in rank order (the all-reduce's bf16), + residual in f32 -> residual out, the f32 sum (SRES) and the task's sum of squares (SSQ, a pair a token).

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def stage_reduce(c, task):
    """Triples 16 task ..: the ranks' values summed in rank order (the
    all-reduce's bf16), + residual in f32 -> residual out, the f32 sum (SRES)
    and the task's sum of squares (SSQ, a pair a token)."""
    s, tid, a = c["S"], c["tid"], c["args"]
    sq = fx.Float32(0.0)
    if tid < s * TASK_TRIPLES:
        t = tid // TASK_TRIPLES
        tr = task * TASK_TRIPLES + tid % TASK_TRIPLES
        got = c["mb"].poll(
            [(c["own"], attn_pair(c, src, t, tr), 1) for src in range(c["tp"])],
            site=SITE_REDUCE,
            triples=True,
        )
        res = triple_bits(a["res_in"], t, tr)
        col = t * HIDDEN + 3 * tr
        for k in range_constexpr(3):
            word = [got[src][0][k // 2] for src in range(c["tp"])]
            acc = bf16_bits_f32(word[0] >> (16 * (k % 2)))
            for src in range_constexpr(1, c["tp"]):
                acc = acc + bf16_bits_f32(word[src] >> (16 * (k % 2)))
            v = bf16_round(acc) + bf16_bits_f32(fx.Int32(res >> (16 * k)))
            bo.buffer_store(v.to(fx.BFloat16), rsrc(a["res_out"]), col + k)
            store_dev(c["sres"].value, col + k, [v.bitcast(fx.Int32)])
            sq = sq + v * v
        sq = butterfly(sq, (8, 4, 2, 1))
    cta_sync(drain=True)
    if (tid < s * TASK_TRIPLES) & (tid % TASK_TRIPLES == 0):
        t = tid // TASK_TRIPLES
        c["mb"].publish(c["ssq"], t * SEAM_TASKS + task, sq)

stage_route(c, t)

Token t: a thread sums 4 of an expert's 16 logit parts, wave 0 the 4 sums in order (two experts a lane); top-4 of sigmoid + bias (ties to the lower expert), weights the selected sigmoids x (2.0 / their sum).

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def stage_route(c, t):
    """Token t: a thread sums 4 of an expert's 16 logit parts, wave 0 the 4
    sums in order (two experts a lane); top-4 of sigmoid + bias (ties to the
    lower expert), weights the selected sigmoids x (2.0 / their sum)."""
    tid, lane, wave, red, a = c["tid"], c["lane"], c["wave"], c["red"], c["args"]
    quarter = ROUTER_PARTS // 4
    e = tid % EXPERTS
    q = tid // EXPERTS
    got = c["mb"].poll(
        [
            (c["logit"], (t * EXPERTS + e) * ROUTER_PARTS + q * quarter + p, 2)
            for p in range(0, quarter, 2)
        ],
        site=SITE_ROUTE,
    )
    c["mb"].acquire(Space.SCRATCH)
    part = fx.Float32(0.0)
    for p in range_constexpr(quarter):
        part = part + got[p // 2][p % 2].bitcast(fx.Float32)
    fx.ptr_store(part, red + tid)
    gpu.barrier()
    if wave == 0:
        logits = []
        for i in range_constexpr(2):
            tot = fx.Float32(0.0)
            for qq in range_constexpr(4):
                tot = tot + fx.ptr_load(red + (qq * EXPERTS + lane + 64 * i))
            logits.append(tot)
        sc = [_sigmoid(logits[i]) for i in range(2)]
        bias = [
            fx.Float32(
                bo.buffer_load(rsrc(a["bias"]), lane + 64 * i, vec_width=1, dtype=T.f32)
            )
            for i in range(2)
        ]
        neg = fx.Float32(float("-inf"))
        val = [sc[i] + bias[i] for i in range(2)]
        val = [fx.isnan(v).select(neg, v) for v in val]
        sc = [fx.isnan(v).select(fx.Float32(0.0), v) for v in sc]
        taken = [fx.Int32(0), fx.Int32(0)]
        total = fx.Float32(0.0)
        my_id, my_w = fx.Int32(0), fx.Float32(0.0)
        for k in range_constexpr(TOPK):
            v0 = (taken[0] == 0).select(val[0], neg)
            v1 = (taken[1] == 0).select(val[1], neg)
            pick1 = v1 > v0
            v = pick1.select(v1, v0)
            idx = pick1.select(fx.Int32(lane + 64), fx.Int32(lane))
            best = butterfly(v, (32, 16, 8, 4, 2, 1), fx.max)
            cand = (v == best).select(idx, fx.Int32(1 << 20))
            win = butterfly(cand, (32, 16, 8, 4, 2, 1), fx.min)
            mine0 = win == lane
            mine1 = win == lane + 64
            w_mine = mine0.select(sc[0], mine1.select(sc[1], fx.Float32(0.0)))
            w = butterfly(w_mine, (32, 16, 8, 4, 2, 1), fx.max)
            taken[0] = mine0.select(fx.Int32(1), taken[0])
            taken[1] = mine1.select(fx.Int32(1), taken[1])
            total = total + w
            my_id = (lane == k).select(win, my_id)
            my_w = (lane == k).select(w, my_w)
        scale = fx.Float32(ROUTED_SCALE) / fx.max(total, fx.Float32(1e-20))
        live = row_valid(c, t)
        my_id = live.select(my_id, fx.Int32(-1))
        if lane < TOPK:
            c["mb"].publish(c["route"], t * 2 * TOPK + lane, my_id)
            c["mb"].put(c["route"], t * 2 * TOPK + TOPK + lane, my_w * scale)

stage_router(c, task)

Experts 16 tile .. over K part part (384 columns, 48 a wave) for every token: FP32 weight x f32(X) on f32 MFMAs, the waves summed in order -> LOGIT[token][expert][part]. A lane's float4 of weights and four bf16 of X feed four MFMAs (K slot = its lane group), the same K order on both.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def stage_router(c, task):
    """Experts 16 tile .. over K part ``part`` (384 columns, 48 a wave) for
    every token: FP32 weight x f32(X) on f32 MFMAs, the waves summed in order
    -> ``LOGIT[token][expert][part]``. A lane's float4 of weights and four bf16 of
    X feed four MFMAs (K slot = its lane group), the same K order on both."""
    s, tid, lane, wave, red, a = (
        c["S"],
        c["tid"],
        c["lane"],
        c["wave"],
        c["red"],
        c["args"],
    )
    tile, part = task // ROUTER_PARTS, task % ROUTER_PARTS
    k0 = part * ROUTER_K + wave * (ROUTER_K // WAVES) + (lane // 16) * 4
    chunks = ROUTER_K // WAVES // 16
    e = tile * 16 + lane % 16
    wv = [
        fx.Vector(
            bo.buffer_load(
                rsrc(a["router_w"]), e * HIDDEN + k0 + 16 * ch, vec_width=4, dtype=T.f32
            )
        )
        for ch in range(chunks)
    ]
    await_x(c, s)
    tok = fx.min(lane % 16, s - 1)
    xv = [
        load_dev(c["xn"].value, (tok * HIDDEN + k0 + 16 * ch) // 2, 2)
        for ch in range(chunks)
    ]
    acc = fx.Vector.filled(4, 0.0, fx.Float32)
    for ch in range_constexpr(chunks):
        for h in range_constexpr(2):
            lo, hi = bf2_f32(fx.Int32(xv[ch][h]))
            acc = mfma_f32(wv[ch][2 * h], lo, acc)
            acc = mfma_f32(wv[ch][2 * h + 1], hi, acc)
    fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
    gpu.barrier()
    if tid < 16 * s:
        r = tid % 16
        t = tid // 16
        at = (16 * (r // 4) + t) * 4 + r % 4
        tot = fx.Float32(0.0)
        for w in range_constexpr(WAVES):
            tot = tot + fx.ptr_load(red + (w * 64 * 4 + at))
        # released: a routed up CTA reads XN through logit -> route -> table
        c["mb"].publish(
            c["logit"], (t * EXPERTS + tile * 16 + r) * ROUTER_PARTS + part, tot
        )
    gpu.barrier()

stage_up(c, u, j, weights=None)

Slot u's tile pair j: gate rows 16 j .., up rows 768 + 16 j .. of its expert over this wave's K steps; waves summed in order; SwiGLU-OAI -> bf16 MID[u][column][16 j ..], then UPF.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

@traced
def stage_up(c, u, j, weights=None):
    """Slot u's tile pair j: gate rows 16 j .., up rows 768 + 16 j .. of its
    expert over this wave's K steps; waves summed in order; SwiGLU-OAI ->
    bf16 ``MID[u][column][16 j ..]``, then UPF."""
    tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"]
    e = tab_ld(c, T_SLOT_E + u)
    cnt = slot_count(c, u)
    st0 = wave * UP_WAVE_STEPS
    rgs = [e * (2 * UP_TILES) + j, e * (2 * UP_TILES) + UP_TILES + j]
    if const_expr(weights is None):
        weights = up_weights(c, rgs, st0)
    wvs, svs = weights
    tok = tab_ld(c, T_SLOT_TOK + u * 16 + slot_col(c, cnt))
    xs = x_operands(c, c["xn"].value, tok, HIDDEN // 2, st0, UP_WAVE_STEPS)
    accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)]
    for i in range_constexpr(UP_WAVE_STEPS):
        for h in range_constexpr(2):
            accs[h] = step_mfma(accs[h], wvs[h][i], scale_f32(svs[h], rgs[h], i), xs[i])
    for h in range_constexpr(2):
        fx.ptr_store(accs[h], red + ((wave * 2 + h) * 64 + lane) * 4)
    gpu.barrier()
    if tid < 256:
        col = tid // 16
        r = tid % 16
        if col < cnt:
            at = (16 * (r // 4) + col) * 4 + r % 4
            g = fx.Float32(0.0)
            up = fx.Float32(0.0)
            for w in range_constexpr(WAVES):
                g = g + fx.ptr_load(red + ((w * 2) * 64 * 4 + at))
                up = up + fx.ptr_load(red + ((w * 2 + 1) * 64 * 4 + at))
            v = _swiglu(g, up).to(fx.BFloat16)
            store16(c["mid"].value, (u * 16 + col) * INTER + j * 16 + r, v, "agent")
    flag_after_stores(c, c["upf"], u * UP_TILES + j, tid == 0)

triple_bits(buf, t, triple)

Row t's bf16 at columns 3 triple .. of a [*, HIDDEN] bf16 buffer: one 8-byte load from the dword holding the first, shifted to bits 0 .. 47.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

def triple_bits(buf, t, triple):
    """Row t's bf16 at columns 3 triple .. of a [*, HIDDEN] bf16 buffer: one
    8-byte load from the dword holding the first, shifted to bits 0 .. 47."""
    v = fx.Vector(
        bo.buffer_load(
            rsrc(buf),
            t * (HIDDEN // 2) + (3 * triple) // 2,
            vec_width=2,
            dtype=T.i32,
        )
    )
    w = (fx.Int64(fx.Uint32(fx.Int32(v[1]))) << 32) | fx.Int64(
        fx.Uint32(fx.Int32(v[0]))
    )
    return w >> fx.Int64(16 * (triple & 1))

valid_mask(c)

The step's real rows (a padded row's slot is -1), a bit a row: every wave's lane t < S reads row t's slot.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

def valid_mask(c):
    """The step's real rows (a padded row's slot is -1), a bit a row: every
    wave's lane t < S reads row t's slot."""
    s, lane = c["S"], c["lane"]
    hi = fx.Int32(
        bo.buffer_load(
            rsrc(c["args"]["slots"]),
            2 * fx.min(lane, s - 1) + 1,
            vec_width=1,
            dtype=T.i32,
        )
    )
    ball = ballot((lane < s) & (hi >= 0))
    return fx.Int32(ball & fx.Int64((1 << s) - 1))

x_operands(c, base, row, k_dwords, st0, steps)

The lane's B operands of steps st0 ..: row row of a bf16 [*, K] scratch (k_dwords a row), 8 K values an MFMA.

Source code in vllm/models/minimax_m3/amd/specialized_decode/mono/ffn.py

def x_operands(c, base, row, k_dwords, st0, steps):
    """The lane's B operands of steps st0 ..: row ``row`` of a bf16 [*, K]
    scratch (``k_dwords`` a row), 8 K values an MFMA."""
    lane = c["lane"]
    at = row * k_dwords + (lane // 16) * 16
    return [
        [load_dev(base, at + (st0 + i) * 64 + d * 4, 4) for d in range(4)]
        for i in range(steps)
    ]