vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.oproj.prepared

Output projection fused with its TP all-reduce and post-attention AttnRes: rank partials go into Lamport slots of a multicast [TP, 64, 7168] mailbox (parity alternates by layer) and the consumer sums and re-arms them. TP16 B<=8 fuses producer and consumer (MLA layers below B5 keep them split); above B8 the producer runs 16-token groups. At twoshot_batches, given TwoShot buffers, the all-reduce is common.twoshot (chunks pushed to their column owner, owner sums + multicasts, all Lamport-polled), so each GPU receives ~2 x 14 KB per token instead of TP x 14 KB. bind_ahead: fused rows that also compute the next layer's AttnRes statistics (consumer ahead).

Classes:

• OprojMailbox –

  Sentinel-armed multicast mailbox; defaults to [parity, TP, 64, 7168] BF16.

• PreparedOproj –

  One decoder layer's o-projection + all-reduce + post-attention AttnRes.

Functions:

• compile_oproj –

  Producer, consumer or fused kernel over a [TP, cap, 7168] parity view;

• engine –

  Producer engine above B8 (GB300): grouped mma.sync for TP8 up to B16, else

• producer –

  B<=8: one mma.sync pass, TP8 stages 8 (MLA 4) of 12 K blocks per warp. Above

OprojMailbox

Sentinel-armed multicast mailbox; defaults to [parity, TP, 64, 7168] BF16. multicast=False: a peer-store mailbox (handle.buffer_ptrs only, mc 0), which spends no NVLink multicast object.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py

class OprojMailbox:
    """Sentinel-armed multicast mailbox; defaults to [parity, TP, 64, 7168] BF16.
    `multicast=False`: a peer-store mailbox (`handle.buffer_ptrs` only, `mc`
    0), which spends no NVLink multicast object."""

    def __init__(self, group, device, tp, rank, shape=None, multicast=True):
        self.tp, self.rank = tp, rank
        shape = (2, tp, 64, 7168) if shape is None else shape
        self.buffer = symm_mem.empty(shape, dtype=torch.bfloat16, device=device)
        if multicast:
            self.handle, self.mc = multicast_rendezvous(
                self.buffer, group, "Kimi K3 o-projection mailbox"
            )
        else:
            self.handle, self.mc = peer_rendezvous(self.buffer, group), 0
        if multicast and not self.mc:  # slot() offsets this address
            raise RuntimeError(
                "Kimi K3 specialized o-projection needs multicast memory"
            )
        self.buffer.view(torch.int32).fill_(-0x80000000)
        torch.accelerator.synchronize(device)
        torch.distributed.barrier(group=group)

    def slot(self, parity):
        return self.buffer[parity], self.mc + parity * self.buffer[0].numel() * 2

PreparedOproj

One decoder layer's o-projection + all-reduce + post-attention AttnRes.

Methods:

• bind_ahead –

  The fused rows also compute target's AttnRes statistics into

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py

class PreparedOproj:
    """One decoder layer's o-projection + all-reduce + post-attention AttnRes."""

    def __init__(self, decoder, mailbox, batches, mla, front=None, twoshot=None):
        weight = decoder.self_attn.o_proj.weight
        norm, out_norm = decoder.mlp_res_norm, decoder.post_attention_layernorm
        weights = (norm.weight, decoder.mlp_res_proj.weight.squeeze(0), out_norm.weight)
        self.views = tuple(to_cute(w.reshape(-1)) for w in weights)
        self.eps = tuple(cutlass.Float32(n.variance_epsilon) for n in (norm, out_norm))
        parity, rank = decoder.layer_idx % 2, cutlass.Int32(mailbox.rank)
        local_box, mc = mailbox.slot(parity)
        self.device, self.box = weight.device, to_cute(local_box)
        self.head = (self.box, cutlass.Int64(mc), rank)
        self.twoshot: Any = None
        if twoshot is not None:
            recv, peers, gather, gather_mc = twoshot.views(parity)
            recv = to_cute(recv)
            self.twoshot = (
                (recv, to_cute(peers), rank),
                to_cute(gather),
                (recv, cutlass.Int64(gather_mc), rank),
            )
        self.wview, tp, k = to_cute(weight), mailbox.tp, weight.shape[1]
        nb = decoder.prev_valid_blocks + decoder.is_block_write_layer
        delta = not decoder.is_block_write_layer
        self.front = None if front is None else (*map(to_cute, front[:2]), front[2])
        self.codes, self.twoshots = {}, twoshot_batches(tp) if self.twoshot else set()
        self.shapes: dict[int, tuple] = {}  # fused rows: compile_oproj arguments
        self.ahead: dict[int, Any] = {}
        self.ahead_views: tuple = ()
        for b in batches:
            publish = front is not None and b in front[2]
            fused = tp == 16 and (not mla or b >= 5) and b <= 8
            kinds = ("fused",) if fused else ("producer", "consumer")
            twoshot = b in self.twoshots and not fused
            shape = (k, b, tp, nb, delta, mla)
            self.codes[b] = tuple(
                compile_oproj(kind, *shape, publish and kind != "producer", 64, twoshot)
                for kind in kinds
            )
            if fused:
                self.shapes[b] = (*shape, publish)

    def bind_ahead(self, target, mixture, scal, rows):
        """The fused `rows` also compute `target`'s AttnRes statistics into
        `mixture` [cap, 7168] / `scal` [cap, 2] FP32; `target` must read the
        saved blocks of this kernel's AttnRes. Returns the rows bound."""
        norm, query = target.self_attention_res_norm, target.self_attention_res_proj
        weights = (to_cute(w.reshape(-1)) for w in (norm.weight, query.weight))
        stats = to_cute(mixture), to_cute(scal)
        eps = cutlass.Float32(norm.variance_epsilon)
        self.ahead_views = (*weights, *stats, eps)
        cap = mixture.shape[0]
        for b in sorted(rows):
            shape = self.shapes.get(b)
            if shape is None or b > cap or shape[3] == 0:
                continue
            self.ahead[b] = compile_oproj("fused", *shape, 64, False, cap)
        return tuple(sorted(self.ahead))

    def run(self, local, prefix, blocks, out):
        batch, s = local.shape[0], stream(self.device)
        codes, x = self.codes[batch], to_cute(local)
        front = (None, None)
        if self.front is not None and batch in self.front[2]:
            front = self.front[:2]
        rows = (dynamic_view(prefix), dynamic_view(blocks), *self.views)
        rows += (dynamic_view(out), *front, *self.eps, s)
        if batch in self.ahead:
            stats = (*rows[:-1], *self.ahead_views, s)
            self.ahead[batch](x, self.wview, *self.head, *stats)
        elif len(codes) == 1:
            codes[0](x, self.wview, *self.head, *rows)
        elif batch in self.twoshots:
            head, gather, tail = self.twoshot
            codes[0](x, self.wview, *head, s)
            codes[1](gather, *rows[:-1], *tail, s)
        else:
            codes[0](x, self.wview, *self.head, s)
            codes[1](self.box, *rows)
        return out, prefix, blocks

bind_ahead(target, mixture, scal, rows)

The fused rows also compute target's AttnRes statistics into mixture [cap, 7168] / scal [cap, 2] FP32; target must read the saved blocks of this kernel's AttnRes. Returns the rows bound.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py

def bind_ahead(self, target, mixture, scal, rows):
    """The fused `rows` also compute `target`'s AttnRes statistics into
    `mixture` [cap, 7168] / `scal` [cap, 2] FP32; `target` must read the
    saved blocks of this kernel's AttnRes. Returns the rows bound."""
    norm, query = target.self_attention_res_norm, target.self_attention_res_proj
    weights = (to_cute(w.reshape(-1)) for w in (norm.weight, query.weight))
    stats = to_cute(mixture), to_cute(scal)
    eps = cutlass.Float32(norm.variance_epsilon)
    self.ahead_views = (*weights, *stats, eps)
    cap = mixture.shape[0]
    for b in sorted(rows):
        shape = self.shapes.get(b)
        if shape is None or b > cap or shape[3] == 0:
            continue
        self.ahead[b] = compile_oproj("fused", *shape, 64, False, cap)
    return tuple(sorted(self.ahead))

compile_oproj(kind, k, batch, tp, nb, has_delta, mla, front=False, cap=64, twoshot=False, ahead=0) cached

Producer, consumer or fused kernel over a [TP, cap, 7168] parity view; front: the consumer also publishes x_moe [1, 4, 7168] and the epoch hint. twoshot: the two-shot pair (arguments as TwoShot.views). ahead: the fused kernel also stores the next layer's AttnRes statistics into [ahead, 7168] / [ahead, 2] FP32 rows (after the AttnRes rows: next norm, next query, the two rows, next eps).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py

@cache
def compile_oproj(
    kind, k, batch, tp, nb, has_delta, mla, front=False, cap=64, twoshot=False, ahead=0
):
    """Producer, consumer or fused kernel over a [TP, cap, 7168] parity view;
    `front`: the consumer also publishes x_moe [1, 4, 7168] and the epoch hint.
    `twoshot`: the two-shot pair (arguments as `TwoShot.views`). `ahead`: the
    fused kernel also stores the next layer's AttnRes statistics into [ahead, 7168] /
    [ahead, 2] FP32 rows (after the AttnRes rows: next norm, next query, the two
    rows, next eps)."""
    box = fake((1 if twoshot else tp, cap, 7168), (cap * 7168, 7168, 1))
    x, w = fake((batch, k), (k, 1)), fake((7168, k), (k, 1))
    mc, slot, eps = cutlass.Int64(0), cutlass.Int32(0), cutlass.Float32(1e-6)
    recv, peers = fake((tp, cap, 7168 // tp)), fake((tp,), None, torch.int64)
    if kind.startswith("producer"):  # producer_mma / producer_umma: engine A/B
        kernel = producer(k, batch, tp, mla, twoshot, kind[9:] or None)
        tail = (recv, peers, slot) if twoshot else (box, mc, slot)
        return cute.compile(kernel, x, w, *tail, make_fake_stream())
    row, weight = dynamic_fake(2), fake((7168,), (1,))
    xmoe = fake((1, 4, 7168), (4 * 7168, 7168, 1)) if front else None
    hints = fake((224,), (1,), torch.int32) if front else None
    rows: tuple = (
        row,
        dynamic_fake(3),
        weight,
        weight,
        weight,
        row,
        xmoe,
        hints,
        eps,
        eps,
    )
    if kind == "consumer":  # B64: one 128-thread group per CTA polls every rank
        options = dict(has_delta=has_delta)
        groups = 4 if batch <= 32 and not twoshot else 1
        kernel = OprojConsumer(nb, groups=groups, twoshot=twoshot, **options)
        extra = (recv, mc, slot) if twoshot else ()
        return cute.compile(kernel, box, *rows, *extra, make_fake_stream())
    kernel = OprojFused(nb, has_delta=has_delta, ahead=ahead > 0)
    if ahead:
        stats = (fake((ahead, n), None, torch.float32) for n in (7168, 2))
        rows += (weight, weight, *stats, eps)
    return cute.compile(kernel, x, w, box, mc, slot, *rows, make_fake_stream())

engine(batch, tp)

Producer engine above B8 (GB300): grouped mma.sync for TP8 up to B16, else tcgen05.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py

def engine(batch, tp):
    """Producer engine above B8 (GB300): grouped mma.sync for TP8 up to B16, else
    tcgen05."""
    default = "mma" if tp == 8 and batch <= 16 else "umma"
    return default

producer(k, batch, tp, mla, twoshot=False, kind=None)

B<=8: one mma.sync pass, TP8 stages 8 (MLA 4) of 12 K blocks per warp. Above B8 (kind overrides engine): the tcgen05 engine, or 2-tile warps over 8 K parts in 16-token groups, TP8 staging 2 of 6 K blocks per warp.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py

def producer(k, batch, tp, mla, twoshot=False, kind=None):
    """B<=8: one mma.sync pass, TP8 stages 8 (MLA 4) of 12 K blocks per warp. Above
    B8 (`kind` overrides `engine`): the tcgen05 engine, or 2-tile warps over 8 K
    parts in 16-token groups, TP8 staging 2 of 6 K blocks per warp."""
    if batch <= 8:
        staged = 0 if tp == 16 else 4 if mla else 8
        return OprojProducer(k, staged_blocks=staged, twoshot=twoshot)
    if (kind or engine(batch, tp)) == "umma":
        return OprojUmma(k, twoshot=twoshot)
    grouped = dict(rows=2, group_tiles=2, twoshot=twoshot)
    return OprojProducer(k, staged_blocks=0 if tp == 16 else 2, **grouped)