vllm.models.kimi_k3.amd.specialized_decode.mono.layer

K2 layer_post: one Kimi-K3 layer from the attention's core output (the o_proj input) to the layer's MoE output, in one launch -- K2a's o_proj, its all-reduce and the MLP AttnRes seam, then K2b's latent MoE (moe).

As ATOM's layer_post, the two halves share the launch so that the MoE's static weights stream while the AttnRes chain (partials -> reduce -> slice -> gate -> mix) is still settling: the CTAs that will run the router, latent and shared gate/up GEMVs issue their weights before they first wait.

CTAs 0 .. 55     K2a's reduce + AttnRes (slice, gate, mix -> x rows + XRDY);
                 then route (0 .. S-1) and the shared SiTU (8 .. 8+S-1)
CTAs 56 .. 159   o_proj (448 x 16 rows) -> partials to every rank; then
                 the shared gate/up (96 tasks, weights right after o_proj)
CTAs 160 .. 243  router parts (28) and latent parts (56): weights first
every CTA        ug; down (0 .. 223) or the shared down (224 .. 255);
                 lnorm (56 .. 83), up (0 .. 55), reduce (200 .. 255)

An MLA layer's K2 runs its M1b (mla_back) first, in the same launch: every CTA pushes, CTAs 0 .. 55 combine and write core (a flag a head and 4 tokens), and the o_proj CTAs, their weights in flight, poll the flags of a core tile before they load it. That removes the M1b -> K2 launch boundary.

Functions:

• extra_peer_end –

  The end (from the kernels' peer base) of every region the OPT_* bits

• k2_launch –

  Host launcher: core [S, 1536] (o_proj's input), prefix updated in

• k2_sent –

  (OPT_SENTS) Whether build key's all-reduces go through the sentinel

• m1b_ctok –

  The fused M1b's tokens a combine task (over its NA CTAs).

• peer_layout –

  K2a's AR region, then K2b's: one parity's (name -> (byte offset, bytes)).

• rlogp_half_bytes –

  (OPT_RLOG_PARTS) One epoch parity's RLOGP, tagged pairs

• sent_layout –

  (OPT_SENTS) One layer parity's sentinel regions, name -> (byte

• sent_offset –

  (OPT_SENTS) Byte offset (from the kernels' peer base) of the sentinel

• stage_oproj –

  o_proj tasks gi, gi + 200, gi + 400 (< 448): every task's weights in

• tag_layout –

  (OPT_TAG_X) The tagged rows' scratch region (name -> (byte offset,

• xl_words –

  The int32 words of the bf16 row tile LDS xl the GEMV stages share:

extra_peer_end()

The end (from the kernels' peer base) of every region the OPT_* bits add to the peer buffer.

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

def extra_peer_end() -> int:
    """The end (from the kernels' peer base) of every region the OPT_* bits
    add to the peer buffer."""
    return rlogp_offset() + 2 * rlogp_half_bytes()

k2_launch(key, *, core, w_o, prefix, blocks, ares_nw, ares_qk, in_nw, w_gate, bias, w_ld, w_sgu, w_sd, w13, w13s, w2, w2s, ln_w, w_up, out, scratch, peers, rank, epoch, layer, health, alarm, tl=None, m1b_in=None, m1c_in=None)

Host launcher: core [S, 1536] (o_proj's input), prefix updated in place, out the MoE output [S, 7168]; MoE weights as moe. With key.m1b_qlen, m1b_in is mla_back.mla_post's (o, lse, seq_lens, qsl, w_v, w_v_scale, gate) and this launch writes core first. With key.m1c, m1c_in is mla_attn.mla_attn's (q, kv, pages, indptr, scale, kv_scale), its partials M1b's o / lse (split 0's views).

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

def k2_launch(
    key: K2Build,
    *,
    core,
    w_o,
    prefix,
    blocks,
    ares_nw,
    ares_qk,
    in_nw,
    w_gate,
    bias,
    w_ld,
    w_sgu,
    w_sd,
    w13,
    w13s,
    w2,
    w2s,
    ln_w,
    w_up,
    out,
    scratch,
    peers,
    rank,
    epoch,
    layer,
    health,
    alarm,
    tl=None,
    m1b_in=None,
    m1c_in=None,
):
    """Host launcher: ``core`` [S, 1536] (o_proj's input), ``prefix`` updated in
    place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``. With
    ``key.m1b_qlen``, ``m1b_in`` is ``mla_back.mla_post``'s (o, lse, seq_lens,
    qsl, w_v, w_v_scale, gate) and this launch writes ``core`` first. With
    ``key.m1c``, ``m1c_in`` is ``mla_attn.mla_attn``'s (q, kv, pages, indptr,
    scale, kv_scale), its partials M1b's o / lse (split 0's views)."""
    assert (m1b_in is not None) == (key.m1b_qlen > 0)
    assert (m1c_in is not None) == key.m1c
    # the regions the OPT_* bits use exist
    assert SENT_READY or not (k2_sent(key) or k2_rlogp(key)), (
        "OPT_SENTS / OPT_RLOG_PARTS without the runner's peer regions"
    )
    assert TAG_READY or not (key.opt & OPT_TAG_X), (
        "OPT_TAG_X without the runner's region"
    )
    if m1c_in is not None and m1c.wide_kv(m1c_in["kv"]):
        key = replace(key, m1c_wide=True)
    f = build(key)
    _run_compiled(
        f,
        core.data_ptr(),
        *(m1b.host_args(key.m1b_ranks, **m1b_in) if m1b_in is not None else [0] * 16),
        *(m1c.host_args(**m1c_in) if m1c_in is not None else [0] * 6),
        w_o.data_ptr(),
        prefix.data_ptr(),
        blocks.data_ptr(),
        blocks.stride(0),
        blocks.stride(1),
        ares_nw.data_ptr(),
        ares_qk.data_ptr(),
        in_nw.data_ptr(),
        w_gate.data_ptr(),
        bias.data_ptr(),
        w_ld.data_ptr(),
        w_sgu.data_ptr(),
        w_sd.data_ptr(),
        w13.data_ptr(),
        w13s.data_ptr(),
        w2.data_ptr(),
        w2s.data_ptr(),
        ln_w.data_ptr(),
        w_up.data_ptr(),
        out.data_ptr(),
        scratch.data_ptr(),
        peers.data_ptr(),
        rank,
        epoch.data_ptr(),
        layer,
        health.data_ptr(),
        alarm,
        ptr(tl),
        torch.cuda.current_stream(),
    )

k2_sent(key)

(OPT_SENTS) Whether build key's all-reduces go through the sentinel regions. The one predicate the push and the reduce of every one of them use: a push the reduce does not take would leave words neither polled nor re-armed, and the next launch's reader would hang.

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

def k2_sent(key: K2Build) -> bool:
    """(OPT_SENTS) Whether build ``key``'s all-reduces go through the sentinel
    regions. The one predicate the push and the reduce of every one of them
    use: a push the reduce does not take would leave words neither polled nor
    re-armed, and the next launch's reader would hang."""
    return bool(key.opt & OPT_SENTS) and key.tokens <= SENT_ROWS

m1b_ctok(key)

The fused M1b's tokens a combine task (over its NA CTAs).

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

def m1b_ctok(key: K2Build) -> int:
    """The fused M1b's tokens a combine task (over its NA CTAs)."""
    if not key.m1b_qlen:
        return m1b.TT
    return m1b.ctok(key.tokens, key.m1b_qlen, key.m1b_splits, NA)

peer_layout(s)

K2a's AR region, then K2b's: one parity's (name -> (byte offset, bytes)).

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

def peer_layout(s):
    """K2a's AR region, then K2b's: one parity's (name -> (byte offset, bytes))."""
    pairs = {"ar": k2a.ar_half_pairs(s), **k2b.peer_pairs(s)}
    return pair_layout(pairs.items())

rlogp_half_bytes()

(OPT_RLOG_PARTS) One epoch parity's RLOGP, tagged pairs [src rank][part][token][router row], SENT_ROWS rows.

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

def rlogp_half_bytes() -> int:
    """(OPT_RLOG_PARTS) One epoch parity's RLOGP, tagged pairs
    `[src rank][part][token][router row]`, SENT_ROWS rows."""
    return TP_ * k2b.R_PARTS * SENT_ROWS * (k2b.E // TP_) * 8

sent_layout()

(OPT_SENTS) One layer parity's sentinel regions, name -> (byte offset, bytes): a word a (rank, bf16 pair, row), SENT_ROWS rows.

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

def sent_layout() -> dict:
    """(OPT_SENTS) One layer parity's sentinel regions, name -> (byte
    offset, bytes): a word a (rank, bf16 pair, row), SENT_ROWS rows."""
    n = {
        "ars": TP_ * SENT_ROWS * HIDDEN // 2 * 4,
        "latrs": TP_ * SENT_ROWS * k2b.LAT // 2 * 4,
        "finals": TP_ * SENT_ROWS * HIDDEN // 2 * 4,
    }
    out, off = {}, 0
    for name, b in n.items():
        out[name] = (off, b)
        off += -(-b // 256) * 256
    return out

sent_offset()

(OPT_SENTS) Byte offset (from the kernels' peer base) of the sentinel regions: past K2's regions and M1's at their largest (DCP8), so no other region moves. Two layer parities.

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

def sent_offset() -> int:
    """(OPT_SENTS) Byte offset (from the kernels' peer base) of the sentinel
    regions: past K2's regions and M1's at their largest (DCP8), so no other
    region moves. Two layer parities."""
    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import mla

    end = (
        peer_bytes(k2b.MAX_S)
        + mla.peer_bytes(k2b.MAX_S, TP_)
        + m1b.peer_bytes(k2b.MAX_S, TP_)
    )
    return -(-end // 4096) * 4096

stage_oproj(c, gi)

o_proj tasks gi, gi + 200, gi + 400 (< 448): every task's weights in flight at once (one memory latency, not three), then the MFMAs a core tile of 16 rows at a time.

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

@traced
def stage_oproj(c, gi):
    """o_proj tasks gi, gi + 200, gi + 400 (< 448): every task's weights in
    flight at once (one memory latency, not three), then the MFMAs a core tile
    of 16 rows at a time."""
    ops = [k2a.o_loads(c, fx.min(gi + O_N * j, OTASKS - 1)) for j in range(O_ROUNDS)]
    for t0, n in gemv.token_tiles(c["S"]):
        cm = 0
        if const_expr(c["m1b"]):
            # the fused M1b's tasks of these rows (a head, ``TC`` tokens), a
            # thread a flag
            tc = c["TC"]
            f0 = t0 // tc * m1b.NH
            nf = ((t0 + n + tc - 1) // tc) * m1b.NH - f0
            tid = c["tid"]
            if tid < nf:
                c["poll"]([(c["crdy"], f0 + tid, 1)])
            gpu.barrier()
            cm = CM_DEV
        k2a.load_core(c, t0, n, cm)
        for j in range_constexpr(O_ROUNDS):
            t = gi + O_N * j
            if t < OTASKS:
                k2a.stage_o(c, t, ops[j], t0, n)

tag_layout()

(OPT_TAG_X) The tagged rows' scratch region (name -> (byte offset, bytes)), past every launch's layout at its largest (K1's, M1a's, K2's: as MonoRunner sizes the scratch), so no other launch's plain word sits where these keep tags.

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

def tag_layout() -> dict:
    """(OPT_TAG_X) The tagged rows' scratch region (name -> (byte offset,
    bytes)), past every launch's layout at its largest (K1's, M1a's,
    K2's: as ``MonoRunner`` sizes the scratch), so no other launch's plain
    word sits where these keep tags."""
    from vllm.models.kimi_k3.amd.specialized_decode.mono.attention import kda, mla

    S = k2b.MAX_S
    end = max(
        kda.scratch_bytes(
            kda.KdaPreBuild(tokens=S, qlen=S, nblocks=8, delta=True, state_len=3 + S)
        ),
        scratch_bytes(K2Build(tokens=S, nblocks=8)),
        mla.scratch_bytes(mla.MlaPreBuild(tokens=S, nblocks=8, delta=True)),
    )
    # + 1 MiB: a fused K2's M1b flags (``crdy``) can outgrow the unfused key's
    off = -(-(end + (1 << 20)) // 4096) * 4096
    return {"xt": (off, S * HIDDEN // 2 * 8)}

xl_words(s)

The int32 words of the bf16 row tile LDS xl the GEMV stages share: the core tile (o_proj), the early stages' x tiles, every h row (the shared down) and the LN tile (up).

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

def xl_words(s):
    """The int32 words of the bf16 row tile LDS ``xl`` the GEMV stages share:
    the core tile (o_proj), the early stages' x tiles, every h row (the shared
    down) and the LN tile (up)."""
    tiled = [k2a.OK, HIDDEN // k2b.R_PARTS, HIDDEN // k2b.L_PARTS, HIDDEN, k2b.LAT]
    rows = [(k, min(s, gemv.tile_rows(k))) for k in tiled] + [(k2b.SI, s)]
    return max(n * (k + gemv.LDS_PAD) // 2 for k, n in rows)