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)

Functions:

• k2_launch –

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

• peer_layout –

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

• stage_oproj –

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

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)

Host launcher: core [S, 1536] (o_proj's input), prefix updated in place, out the MoE output [S, 7168]; MoE weights as moe.

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,
):
    """Host launcher: ``core`` [S, 1536] (o_proj's input), ``prefix`` updated in
    place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``."""
    f = build(key)
    f(
        core.data_ptr(),
        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,
        stream=torch.cuda.current_stream(),
    )

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

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.

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."""
    ops = [k2a.o_loads(c, fx.min(gi + O_N * j, OTASKS - 1)) for j in range(O_ROUNDS)]
    k2a.load_core(c)
    for j in range_constexpr(O_ROUNDS):
        t = gi + O_N * j
        if t < OTASKS:
            k2a.stage_o(c, t, ops[j])