vllm.models.kimi_k3.amd.specialized_decode.mono.attention.mla

M1a mla_pre: one Kimi-K3 MLA layer from the AttnRes seam to AITER's decode input, for a spec-verify step (S rows) at TP8 + DCP8 (the DCP group the TP group: rank r's 12 heads are heads 12 r .. of the gathered query) or TP8 + DCP1 (ranks 1: the query is this rank's 12 heads, no gather).

One launch a layer, BLOCKS CTAs x THREADS threads:

attn CTAs 0 .. 55   AttnRes + input_layernorm as K1 (``kda.stage_*``) ->
                    bf16 x rows, an XRDY flag a slice
gemv CTAs 56 ..     fused_qkv_a (132 x 16 rows) and g_proj (96 x 16), task
                    gi and gi + 200: q_c | kv_c | k_pe -> QA (bf16 pairs);
                    the gate bf16(sigmoid(bf16 g)) -> ``gate_out``
norm (token t on CTA t)
                    q_a / kv_a RMSNorm: q_c -> plain bf16 rows, a QCR flag
                    a token; kv_c (normed) | k_pe -> fp8 cache entry when
                    the token's slot is this rank's (>= 0)
q_b (16 rows, CTAs 64 ..)
                    q_c rows in LDS: q_nope -> QN (bf16 pairs); q_pe -> fp8
                    -> every rank's QG (DCP1: ``q_out``)
W_UK (a wave a (token tile, head, 16 latent rows), CTAs 208 ..)
                    q_nope -> MXFP4 (aiter's even scale, a lane a 32-block),
                    one 16x16x128 FP4 MFMA against W_K -> bf16 -> fp8 ->
                    every rank's QG (DCP1: ``q_out``)
                    (``qbf16``, DCP1: bf16 -> ``q_out``, M1c's)
unpack (every CTA, DCP8)
                    this rank's QG -> the dense gathered ``q_out``
                    [S, ranks 12, 576] fp8

Rounding follows the vLLM path where it is elementwise (bf16 hand-offs, the fp8 quants, the MXFP4 quant of the W_UK input); split-K sums differ only in their order.

Functions:

• bf_words –

  n dwords of a bf16 row at word i (device scope).

• block_sum2 –

  Waves 0 .. 5's sum and waves 6, 7's sum of v (q / kv threads).

• even_code –

  Aiter's _mxfp4_scale_from_amax code from |bf16| bits am: the

• fp8_word –

  Four f32 -> four E4M3 (vLLM's saturating quant, scale 1).

• gemv_tasks –

  Gemv CTA gi: task gi (weights issued before x lands), then gi + 200.

• mla_pre –

  Host launcher. w_k (12, 512, 64) u8 MXFP4 and w_k_scale (12, 512,

• peer_pairs –

  QG: every source rank's [S, 12, 144] fp8 words, a pair each.

• proj_loads –

  Task task's 16 weight rows (fused_qkv_a, then g_proj) for this wave;

• put_q –

  Query values 4 w .. (four f32, bf16-exact) of (token t, local head h):

• put_qg –

  fp8 word w of (token t, local head h) -> every rank's QG; at DCP1

• stage_norm –

  Token t: q_c (threads 0 .. 383, 4 columns each) and kv_c (384 ..) RMSNorm;

• stage_proj –

  Task task against the LDS x (tokens t0 .. t0 + n): q_c | kv_c | k_pe

• stage_qb –

  q_b rows 16 task ..: head h = task / 12, its nope rows (sub < 8) -> QN or

• stage_uk –

  Wave task (tile, head, 16 latent rows): W_K's rows against the 16 tokens'

• stage_unpack –

  This rank's QG -> q_out [S, QH, 576]: a thread 2 words a round (QG

bf_words(ptr, i, n)

n dwords of a bf16 row at word i (device scope).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

def bf_words(ptr, i, n):
    """``n`` dwords of a bf16 row at word ``i`` (device scope)."""
    return fx.Vector(
        bo.buffer_load(rsrc(ptr), i, vec_width=n, dtype=T.i32, cache_modifier=CM_DEV)
    )

block_sum2(c, v)

Waves 0 .. 5's sum and waves 6, 7's sum of v (q / kv threads).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

@traced
def block_sum2(c, v):
    """Waves 0 .. 5's sum and waves 6, 7's sum of ``v`` (q / kv threads)."""
    tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"]
    v = butterfly(v, (32, 16, 8, 4, 2, 1))
    gpu.barrier()
    if lane == 0:
        fx.ptr_store(v, red + wave)
    gpu.barrier()
    q = fx.ptr_load(red + 0)
    for w in range_constexpr(1, 6):
        q = q + fx.ptr_load(red + w)
    kv = fx.ptr_load(red + 6) + fx.ptr_load(red + 7)
    _ = tid
    return q, kv

even_code(am)

Aiter's _mxfp4_scale_from_amax code from |bf16| bits am: the exponent of amax rounded at mantissa bit 21, minus 2, clamped at 0.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

def even_code(am):
    """Aiter's ``_mxfp4_scale_from_amax`` code from |bf16| bits ``am``: the
    exponent of amax rounded at mantissa bit 21, minus 2, clamped at 0."""
    e = (((am << 16) + fx.Int32(0x200000)) >> 23) & 0xFF
    return fx.max(e - 2, fx.Int32(0))

fp8_word(v)

Four f32 -> four E4M3 (vLLM's saturating quant, scale 1).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

def fp8_word(v):
    """Four f32 -> four E4M3 (vLLM's saturating quant, scale 1)."""
    return fp8_pack4(*[clamp_fp8(x) for x in v])

gemv_tasks(c, gi)

Gemv CTA gi: task gi (weights issued before x lands), then gi + 200.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

@traced
def gemv_tasks(c, gi):
    """Gemv CTA ``gi``: task gi (weights issued before x lands), then gi + 200."""
    s, on, tls, tid = c["S"], c["on"], c["tls"], c["tid"]
    ops0 = proj_loads(c, gi)
    kda.wait_x(c)
    for t0, n in gemv.token_tiles(s, XT):
        kda.load_x(c, t0, n)
        stage_proj(c, gi, ops0, t0, n)
    stamp(on, tls, tid, 4)
    if gi + GEMV_CTAS < GEMV_TASKS:
        ops1 = proj_loads(c, gi + GEMV_CTAS)
        for t0, n in gemv.token_tiles(s, XT):
            if const_expr(s > XT):
                kda.load_x(c, t0, n)
            stage_proj(c, gi + GEMV_CTAS, ops1, t0, n)
    stamp(on, tls, tid, 5)

mla_pre(key, *, prefix, delta, blocks, ares_nw, ares_qk, in_nw, w_qkv, w_g, q_nw, kv_nw, w_qb, w_k, w_k_scale, kv_cache, slots, gate_out, q_out, scratch, peers, rank, layer, epoch, health, alarm, tl=None)

Host launcher. w_k (12, 512, 64) u8 MXFP4 and w_k_scale (12, 512, 4) e8m0, any strides but W_K's last; kv_cache [blocks, block_size, 576] one-byte, entries dense; slots [S] int64 (-1: another rank's token); gate_out [S, 1536] bf16; q_out [S, q_heads(key), 576] fp8 (dense; bf16 with qbf16); peers every rank's peer buffer base (i64 a rank; unread at ranks 1, may be None); layer a slot below kda.TAG_SLOTS.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

def mla_pre(
    key: MlaPreBuild,
    *,
    prefix,
    delta,
    blocks,
    ares_nw,
    ares_qk,
    in_nw,
    w_qkv,
    w_g,
    q_nw,
    kv_nw,
    w_qb,
    w_k,
    w_k_scale,
    kv_cache,
    slots,
    gate_out,
    q_out,
    scratch,
    peers,
    rank,
    layer,
    epoch,
    health,
    alarm,
    tl=None,
):
    """Host launcher. ``w_k`` (12, 512, 64) u8 MXFP4 and ``w_k_scale`` (12, 512,
    4) e8m0, any strides but W_K's last; ``kv_cache`` [blocks, block_size, 576]
    one-byte, entries dense; ``slots`` [S] int64 (-1: another rank's token);
    ``gate_out`` [S, 1536] bf16; ``q_out`` [S, q_heads(key), 576] fp8 (dense; bf16 with
    ``qbf16``);
    ``peers`` every rank's peer buffer base (i64 a rank; unread at ranks 1,
    may be None); ``layer`` a slot below
    ``kda.TAG_SLOTS``."""
    assert w_k.stride(-1) == 1 and kv_cache.stride(-1) == 1
    assert slots.dtype == torch.int64 and slots.size(0) == key.tokens
    assert q_out.is_contiguous() and gate_out.is_contiguous()
    assert q_out.shape[1:] == (q_heads(key), ENTRY), q_out.shape
    assert q_out.element_size() == (2 if key.qbf16 else 1), q_out.dtype
    f = build(key)
    f(
        prefix.data_ptr(),
        ptr(delta),
        blocks.data_ptr(),
        blocks.stride(0),
        blocks.stride(1),
        ares_nw.data_ptr(),
        ares_qk.data_ptr(),
        in_nw.data_ptr(),
        w_qkv.data_ptr(),
        w_g.data_ptr(),
        q_nw.data_ptr(),
        kv_nw.data_ptr(),
        w_qb.data_ptr(),
        w_k.data_ptr(),
        w_k.stride(0),
        w_k.stride(1),
        w_k_scale.data_ptr(),
        w_k_scale.stride(0),
        w_k_scale.stride(1),
        w_k_scale.stride(2),
        kv_cache.data_ptr(),
        kv_cache.stride(0),
        kv_cache.stride(1),
        kv_cache.size(1),
        slots.data_ptr(),
        gate_out.data_ptr(),
        q_out.data_ptr(),
        scratch.data_ptr(),
        ptr(peers),
        rank,
        layer,
        epoch.data_ptr(),
        health.data_ptr(),
        alarm,
        ptr(tl),
        stream=torch.cuda.current_stream(),
    )

peer_pairs(s, ranks)

QG: every source rank's [S, 12, 144] fp8 words, a pair each.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

def peer_pairs(s: int, ranks: int) -> int:
    """QG: every source rank's [S, 12, 144] fp8 words, a pair each."""
    return ranks * s * NH * QW

proj_loads(c, task)

Task task's 16 weight rows (fused_qkv_a, then g_proj) for this wave; a task past the last loads the last.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

def proj_loads(c, task):
    """Task ``task``'s 16 weight rows (fused_qkv_a, then g_proj) for this wave;
    a task past the last loads the last."""
    a = c["args"]
    task = fx.min(task, GEMV_TASKS - 1)
    is_g = uniform((task >= A_TASKS).select(fx.Int32(1), fx.Int32(0)))
    w = a["w_qkv"] + fx.Int64(is_g) * (a["w_g"] - a["w_qkv"])
    row0 = (task - is_g * A_TASKS) * ROWS
    return gemv.loads(c["lane"], c["wave"], w, row0, HIDDEN)

put_q(c, t, h, w, v)

Query values 4 w .. (four f32, bf16-exact) of (token t, local head h): fp8 -> QG / q_out, or with qbf16 bf16 pairs -> q_out.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

def put_q(c, t, h, w, v):
    """Query values ``4 w ..`` (four f32, bf16-exact) of (token t, local head h):
    fp8 -> QG / ``q_out``, or with ``qbf16`` bf16 pairs -> ``q_out``."""
    if const_expr(c["qbf16"]):
        pair = fx.Vector.from_elements(
            [bf16_pair(v[0], v[1]), bf16_pair(v[2], v[3])], fx.Float32
        )
        bo.buffer_store(pair, rsrc(c["args"]["q_out"]), ((t * NH + h) * QW + w) * 2)
    if const_expr(not c["qbf16"]):
        put_qg(c, t, h, w, fp8_word(v))

put_qg(c, t, h, w, word)

fp8 word w of (token t, local head h) -> every rank's QG; at DCP1 straight to q_out (nothing in this launch reads it).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

def put_qg(c, t, h, w, word):
    """fp8 word ``w`` of (token t, local head h) -> every rank's QG; at DCP1
    straight to ``q_out`` (nothing in this launch reads it)."""
    if const_expr(c["R"] == 1):
        bo.buffer_store(word, rsrc(c["args"]["q_out"]), (t * NH + h) * QW + w)
    if const_expr(c["R"] > 1):
        i = qg_pair(c, c["rank"], t, h, w)
        for p in range_constexpr(c["R"]):
            c["put_words"](c["qg"][p], i, [word])

stage_norm(c, t)

Token t: q_c (threads 0 .. 383, 4 columns each) and kv_c (384 ..) RMSNorm; q_c -> plain bf16 row + QCR; kv_c | k_pe (threads 0 .. 15) -> fp8 entry.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

@traced
def stage_norm(c, t):
    """Token t: q_c (threads 0 .. 383, 4 columns each) and kv_c (384 ..) RMSNorm;
    q_c -> plain bf16 row + QCR; kv_c | k_pe (threads 0 .. 15) -> fp8 entry."""
    tid, a = c["tid"], c["args"]
    is_q = tid < QL // 4
    col = is_q.select(tid * 4, QL + (tid - QL // 4) * 4)
    pe = QL + KVL + fx.min(tid, ROPE // 4 - 1) * 4
    got = c["poll"](
        [(c["qa"], (t * QKV_A + col) // 2, 2), (c["qa"], (t * QKV_A + pe) // 2, 2)]
    )
    w = got[0]
    v = [bf_lo(w[0]), bf_hi(w[0]), bf_lo(w[1]), bf_hi(w[1])]
    sq = v[0] * v[0] + v[1] * v[1] + v[2] * v[2] + v[3] * v[3]
    qs, kvs = block_sum2(c, sq)
    rq = hw_rsq(qs * fx.Float32(1.0 / QL) + fx.Float32(c["q_eps"]))
    rkv = hw_rsq(kvs * fx.Float32(1.0 / KVL) + fx.Float32(c["kv_eps"]))
    nw = ld_bf4(a["q_nw"], fx.min(tid, QL // 4 - 1) * 4)
    kw = ld_bf4(a["kv_nw"], fx.max(tid - QL // 4, 0) * 4)
    qn = [bf16r(v[j] * rq * nw[j]) for j in range(4)]
    kn = [bf16r(v[j] * rkv * kw[j]) for j in range(4)]
    if is_q:
        bo.buffer_store(
            kda.bf4_words(qn),
            rsrc(c["qc"]),
            (t * QL + tid * 4) // 2,
            cache_modifier=CM_DEV,
        )
    slot = uniform(ld_i32(a["slots"], t * 2))
    if slot >= 0:
        blk = slot // a["kv_bs"]
        off = slot % a["kv_bs"]
        base = (
            a["kv_cache"]
            + fx.Int64(blk) * fx.Int64(a["kv_sb"])
            + fx.Int64(off) * fx.Int64(a["kv_ss"])
        )
        if tid >= QL // 4:
            bo.buffer_store(fp8_word(kn), rsrc(base), tid - QL // 4)
        if tid < ROPE // 4:
            p = got[1]
            pv = [bf_lo(p[0]), bf_hi(p[0]), bf_lo(p[1]), bf_hi(p[1])]
            bo.buffer_store(fp8_word(pv), rsrc(base), KVL // 4 + tid)
    publish(c["put"], c["qcr"], t, fx.Int32(1), tid == 0)
    gpu.barrier()

stage_proj(c, task, ops, t0, n)

Task task against the LDS x (tokens t0 .. t0 + n): q_c | kv_c | k_pe rows -> QA, g rows -> the gate.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

@traced
def stage_proj(c, task, ops, t0, n):
    """Task ``task`` against the LDS x (tokens t0 .. t0 + n): q_c | kv_c | k_pe
    rows -> QA, g rows -> the gate."""
    tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"]
    acc = gemv.mfmas(lane, wave, c["xl"], HIDDEN + LDS_PAD, HIDDEN, n, ops)
    fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
    gpu.barrier()
    r2 = tid % 8
    tt = tid // 8
    live = (tid < 8 * n) & (task < GEMV_TASKS)
    v0 = bf16r(row_sum(red, 2 * r2, fx.min(tt, n - 1)))
    v1 = bf16r(row_sum(red, 2 * r2 + 1, fx.min(tt, n - 1)))
    t = t0 + tt
    if live & (task < A_TASKS):
        c["put_bf"](c["qa"], t * QKV_A + task * ROWS + 2 * r2, [v0, v1])
    if live & (task >= A_TASKS):
        col = (task - A_TASKS) * ROWS + 2 * r2
        g = bf16_pair(sigmoid(v0), sigmoid(v1))
        bo.buffer_store(g, rsrc(c["args"]["gate_out"]), (t * GROWS + col) // 2)
    gpu.barrier()

stage_qb(c, task)

q_b rows 16 task ..: head h = task / 12, its nope rows (sub < 8) -> QN or pe rows -> fp8 -> QG.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

@traced
def stage_qb(c, task):
    """q_b rows 16 task ..: head h = task / 12, its nope rows (sub < 8) -> QN or
    pe rows -> fp8 -> QG."""
    s, tid, lane, wave, red = c["S"], c["tid"], c["lane"], c["wave"], c["red"]
    a = c["args"]
    h = task // QB_SUB
    sub = task % QB_SUB
    ops = gemv.loads(lane, wave, a["w_qb"], task * ROWS, QL)
    if tid < s:
        c["poll"]([(c["qcr"], tid, 1)])
    gpu.barrier()
    for t0, n in gemv.token_tiles(s):
        gemv.dev_rows_to_lds(tid, c["qc"], QL, n, c["xl"], QCROW, t0=t0)
        acc = gemv.mfmas(lane, wave, c["xl"], QCROW, QL, n, ops)
        fx.ptr_store(acc, red + (wave * 64 + lane) * 4)
        gpu.barrier()
        r2 = tid % 8
        tt = fx.min(tid // 8, n - 1)
        if (tid < 8 * n) & (sub < NOPE // ROWS):
            v0 = row_sum(red, 2 * r2, tt)
            v1 = row_sum(red, 2 * r2 + 1, tt)
            at = ((t0 + tt) * NH + h) * NOPE + sub * ROWS + 2 * r2
            c["put_bf"](c["qn"], at, [v0, v1])
        r4 = tid % 4
        t4 = fx.min(tid // 4, n - 1)
        if (tid < 4 * n) & (sub >= NOPE // ROWS):
            v = [bf16r(row_sum(red, 4 * r4 + j, t4)) for j in range(4)]
            w = KVL // 4 + (sub - NOPE // ROWS) * 4 + r4
            put_q(c, t0 + t4, h, w, v)
        gpu.barrier()

stage_uk(c, task)

Wave task (tile, head, 16 latent rows): W_K's rows against the 16 tokens' MXFP4 q_nope -> bf16 -> fp8 -> QG (a lane 4 rows of one token).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

@traced
def stage_uk(c, task):
    """Wave task (tile, head, 16 latent rows): W_K's rows against the 16 tokens'
    MXFP4 q_nope -> bf16 -> fp8 -> QG (a lane 4 rows of one token)."""
    s, lane = c["S"], c["lane"]
    a = c["args"]
    ng = task % UK_GROUPS
    h = (task // UK_GROUPS) % NH
    t0 = task // (UK_GROUPS * NH) * 16
    n = ng * ROWS + lane % 16
    kb = lane // 16
    wo = h * a["wk_sh"] + n * a["wk_sn"] + kb * 16
    wv = fx.Vector(bo.buffer_load(rsrc(a["w_k"]), wo // 4, vec_width=4, dtype=T.i32))
    so = h * a["ks_sh"] + n * a["ks_sn"] + kb * a["ks_sk"]
    sw = fx.Int32(bo.buffer_load(rsrc(a["w_ks"]), so // 4, vec_width=1, dtype=T.i32))
    sa = (sw >> ((so % 4) * 8)) & 0xFF
    t = t0 + lane % 16
    tc = fx.min(t, s - 1)
    p0 = (((tc * NH + h) * NOPE) + kb * 32) // 2
    got = c["poll"]([(c["qn"], p0 + 2 * j, 2) for j in range(8)])
    ws = [got[j // 2][j % 2] for j in range(16)]
    am = fx.Int32(0)
    for w in ws:
        am = fx.max(am, w & 0x7FFF)
        am = fx.max(am, (w >> 16) & 0x7FFF)
    # code 0 (amax < 2^-125) quantizes by 2^-126: the tiny lane stays ~0
    code = fx.max(even_code(am), fx.Int32(1))
    scale = pow2(code)
    xs = []
    for q in range_constexpr(4):
        pk = fx.Int32(0)
        for j in range_constexpr(4):
            pk = pk_fp4_bf16(pk, ws[4 * q + j], scale, j)
        xs.append(pk)
    xv = fx.Vector.from_elements(xs, fx.Int32)
    acc = fx.Vector.filled(4, 0.0, fx.Float32)
    acc = mfma_scaled(wv, xv, acc, sa, code, FP4, FP4, 0, 0)
    v = [bf16r(acc[j]) for j in range(4)]
    if t < s:
        put_q(c, t, h, ng * 4 + kb, v)

stage_unpack(c)

This rank's QG -> q_out [S, QH, 576]: a thread 2 words a round (QG pair 2 u, u = global thread + round x grid). Gathered head g also lands at R x 12 + g while that is a pad head (g < QH - R x 12).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/attention/mla.py

@traced
def stage_unpack(c):
    """This rank's QG -> ``q_out`` [S, QH, 576]: a thread 2 words a round (QG
    pair 2 u, u = global thread + round x grid). Gathered head g also lands
    at R x 12 + g while that is a pad head (g < QH - R x 12)."""
    s, R, tid, bid, QH = c["S"], c["R"], c["tid"], c["bid"], c["QH"]
    pad = QH - R * NH
    nu = R * s * NH * QW // 2
    grid = BLOCKS * THREADS
    rounds = (nu + grid - 1) // grid
    g = bid * THREADS + tid
    us = [fx.min(g + grid * i, nu - 1) for i in range(rounds)]
    got = c["poll"]([(c["qg_own"], 2 * u, 2) for u in us])
    for i in range_constexpr(rounds):
        u = us[i]
        w2 = u % (QW // 2)
        h = (u // (QW // 2)) % NH
        t = (u // (NH * QW // 2)) % s
        src = u // (s * NH * QW // 2)
        gh = src * NH + h
        at = (t * QH + gh) * QW + 2 * w2
        if g + grid * i < nu:
            word = fx.Vector.from_elements(got[i], fx.Int32)
            bo.buffer_store(word, rsrc(c["args"]["q_out"]), at)
            if const_expr(pad > 0):  # noqa: SIM102 (compile-time, then traced)
                if gh < pad:
                    bo.buffer_store(word, rsrc(c["args"]["q_out"]), at + R * NH * QW)