Kimi-Linear Decode GPT-6.1 Sol

"""Single-launch, persistent Triton implementation of the hybrid decode motif.

All large matrices stay in the original packed int4 representation. The one
resident grid executes the entire motif and synchronizes through an acquire /
release device counter. No expanded weight matrices or expanded KV cache exist.
"""
import torch
from torch import nn
import triton
import triton.language as tl


class QuantLinear(nn.Module):
    def __init__(self, k, n, group=128):
        super().__init__()
        self.register_buffer('w_q', torch.empty(k // 2, n, dtype=torch.uint8))
        self.register_buffer('scales', torch.empty(k // group, n, dtype=torch.bfloat16))
        self.register_buffer('zeros', torch.empty(k // group, n, dtype=torch.bfloat16))


class QuantExperts(nn.Module):
    def __init__(self, e, k, n, group=128):
        super().__init__()
        self.register_buffer('w_q', torch.empty(e, k // 2, n, dtype=torch.uint8))
        self.register_buffer('scales', torch.empty(e, k // group, n, dtype=torch.bfloat16))
        self.register_buffer('zeros', torch.empty(e, k // group, n, dtype=torch.bfloat16))


class KDA(nn.Module):
    def __init__(self, c):
        super().__init__()
        d, a = c.hidden, c.kda_heads * c.kda_head_dim
        self.q_proj = QuantLinear(d, a)
        self.k_proj = QuantLinear(d, a)
        self.v_proj = QuantLinear(d, a)
        self.g_proj = QuantLinear(d, a)
        self.o_proj = QuantLinear(a, d)
        self.beta_proj = nn.Linear(d, c.kda_heads, bias=False, dtype=c.dtype)
        self.conv_w = nn.Parameter(torch.empty(3, a, c.short_conv, dtype=c.dtype))


class MLA(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.q_proj = QuantLinear(c.hidden, c.mla_heads * (c.qk_nope + c.qk_rope))
        self.kv_a = QuantLinear(c.hidden, c.kv_lora + c.qk_rope)
        self.kv_b = QuantLinear(c.kv_lora, c.mla_heads * (c.qk_nope + c.v_head))
        self.o_proj = QuantLinear(c.mla_heads * c.v_head, c.hidden)


class MoE(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.router = nn.Linear(c.hidden, c.n_experts, bias=False, dtype=c.dtype)
        self.gate = QuantExperts(c.n_experts, c.hidden, c.moe_inter)
        self.up = QuantExperts(c.n_experts, c.hidden, c.moe_inter)
        self.down = QuantExperts(c.n_experts, c.moe_inter, c.hidden)
        self.s_gate = QuantExperts(c.n_shared, c.hidden, c.moe_inter)
        self.s_up = QuantExperts(c.n_shared, c.hidden, c.moe_inter)
        self.s_down = QuantExperts(c.n_shared, c.moe_inter, c.hidden)


class Block(nn.Module):
    def __init__(self, c, kind):
        super().__init__()
        self.attn_norm = nn.Parameter(torch.ones(c.hidden, dtype=c.dtype))
        self.moe_norm = nn.Parameter(torch.ones(c.hidden, dtype=c.dtype))
        self.attn = KDA(c) if kind == 'K' else MLA(c)
        self.moe = MoE(c)


@triton.jit
def _barrier(Counter, GRID: tl.constexpr):
    # The grid is smaller than the SM count; every CTA can remain resident.
    tl.debug_barrier()
    ticket = tl.atomic_add(Counter, 1, sem='acq_rel', scope='gpu')
    target = (ticket // GRID + 1) * GRID
    while tl.atomic_add(Counter, 0, sem='acquire', scope='gpu') < target:
        pass
    tl.debug_barrier()


@triton.jit
def _pointer(W, layer, slot: tl.constexpr, DT: tl.constexpr):
    return tl.load(W + layer * 40 + slot).to(tl.pointer_type(DT))


@triton.jit
def _gemv(X, Q, Scale, Zero, n, K: tl.constexpr, N: tl.constexpr, BN: tl.constexpr):
    """Stream one group at a time and reduce without materializing any weights."""
    kk = tl.arange(0, 128)
    acc = tl.full((128, BN), 0, tl.float32)
    for g in range(K // 128):
        k = g * 128 + kk
        packed = tl.load(Q + (k[:, None] // 2) * N + n[None, :], n[None, :] < N, 0)
        iq = (packed >> ((k[:, None] & 1) * 4)) & 15
        s = tl.load(Scale + g * N + n, n < N, 0).to(tl.float32)
        z = tl.load(Zero + g * N + n, n < N, 0).to(tl.float32)
        wt = ((iq.to(tl.float32) - z[None, :]) * s[None, :]).to(tl.bfloat16).to(tl.float32)
        x = tl.load(X + k).to(tl.float32)
        acc += x[:, None] * wt
    return tl.sum(acc, 0)


@triton.jit
def _norm(Hidden, X, Weight):
    d = tl.arange(0, 4096)
    h = tl.load(Hidden + d, d < 2304, 0).to(tl.float32)
    w = tl.load(Weight + d, d < 2304, 0).to(tl.float32)
    inv = tl.rsqrt(tl.sum(h * h, 0) / 2304 + 1.e-6)
    tl.store(X + d, (h * inv * w).to(tl.bfloat16).to(tl.float32), d < 2304)


@triton.jit
def _mega(Hin, Hout, W, Counter, Work,
          S0, CQ0, CK0, CV0, S1, CQ1, CK1, CV1, S2, CQ2, CK2, CV2,
          Cold, Rold, Cnew, Rnew, L,
          E: tl.constexpr, GRID: tl.constexpr, BN: tl.constexpr):
    pid = tl.program_id(0)
    # Work is fp32 scratch, allocated once per model. Distinct phase lifetimes
    # allow these arrays to be reused across all four layers.
    X = Work
    A = Work + 4096
    B = Work + 20480
    O = Work + 20544
    Idx = Work + 24640
    RW = Work + 24672
    FF = Work + 24704
    FO = Work + 33920
    Qlatent = Work + 54656
    Parts = Work + 71040
    Latent = Work + 137088
    if pid == 0:
        d = tl.arange(0, 4096)
        tl.store(Hout + d, tl.load(Hin + d, d < 2304, 0), d < 2304)
    _barrier(Counter, GRID)
    for layer in range(4):
        if pid == 0:
            _norm(Hout, X, _pointer(W, layer, 0, tl.bfloat16))
        _barrier(Counter, GRID)
        if layer < 3:
            if layer == 0:
                S, CQ, CK, CV = S0, CQ0, CK0, CV0
            elif layer == 1:
                S, CQ, CK, CV = S1, CQ1, CK1, CV1
            else:
                S, CQ, CK, CV = S2, CQ2, CK2, CV2
            for task in range(pid, 4 * (4096 // BN), GRID):
                which = task // (4096 // BN)
                n = (task % (4096 // BN)) * BN + tl.arange(0, BN)
                # Projection slots 2,5,8,11. Runtime selection keeps one body.
                qp = tl.load(W + layer * 40 + 2 + which * 3).to(tl.pointer_type(tl.uint8))
                sp = tl.load(W + layer * 40 + 3 + which * 3).to(tl.pointer_type(tl.bfloat16))
                zp = tl.load(W + layer * 40 + 4 + which * 3).to(tl.pointer_type(tl.bfloat16))
                val = _gemv(X, qp, sp, zp, n, 2304, 4096, BN).to(tl.bfloat16).to(tl.float32)
                if which < 3:
                    if which == 0:
                        C = CQ
                    elif which == 1:
                        C = CK
                    else:
                        C = CV
                    cw = _pointer(W, layer, 18, tl.bfloat16) + which * 4096 * 4
                    c0 = tl.load(C + n).to(tl.float32)
                    c1 = tl.load(C + 4096 + n).to(tl.float32)
                    c2 = tl.load(C + 8192 + n).to(tl.float32)
                    y = c0 * tl.load(cw + n * 4).to(tl.float32)
                    y += c1 * tl.load(cw + n * 4 + 1).to(tl.float32)
                    y += c2 * tl.load(cw + n * 4 + 2).to(tl.float32)
                    y += val * tl.load(cw + n * 4 + 3).to(tl.float32)
                    tl.store(C + n, c1)
                    tl.store(C + 4096 + n, c2)
                    tl.store(C + 8192 + n, val)
                    val = (y / (1 + tl.exp(-y))).to(tl.bfloat16).to(tl.float32)
                else:
                    val = 1 / (1 + tl.exp(val))
                tl.store(A + which * 4096 + n, val)
            if pid < 32:
                d = tl.arange(0, 4096)
                bp = _pointer(W, layer, 17, tl.bfloat16)
                x = tl.load(X + d, d < 2304, 0)
                b = tl.load(bp + pid * 2304 + d, d < 2304, 0).to(tl.float32)
                bv = tl.sum(x * b, 0).to(tl.bfloat16).to(tl.float32)
                tl.store(B + pid, 1 / (1 + tl.exp(-bv)))
            _barrier(Counter, GRID)
            # Independent state columns: each worker owns 32 value channels.
            for task in range(pid, 128, GRID):
                head = task // 4
                ki = tl.arange(0, 128)
                vi = (task % 4) * 32 + tl.arange(0, 32)
                q = tl.load(A + head * 128 + ki) * (128 ** -0.5)
                k = tl.load(A + 4096 + head * 128 + ki)
                v = tl.load(A + 8192 + head * 128 + vi)
                decay = tl.load(A + 12288 + head * 128 + ki)
                beta = tl.load(B + head)
                address = S + head * 16384 + ki[:, None] * 128 + vi[None, :]
                sv = tl.load(address) * decay[:, None]
                pred = tl.sum(sv * k[:, None], 0)
                sv += beta * k[:, None] * (v - pred)[None, :]
                tl.store(address, sv)
                out = tl.sum(sv * q[:, None], 0).to(tl.bfloat16).to(tl.float32)
                tl.store(O + head * 128 + vi, out)
            _barrier(Counter, GRID)
        else:
            # Query and current compressed KV projections.
            for task in range(pid, triton.cdiv(6144, BN) + triton.cdiv(576, BN), GRID):
                if task < 6144 // BN:
                    n = task * BN + tl.arange(0, BN)
                    val = _gemv(X, _pointer(W, layer, 2, tl.uint8),
                                _pointer(W, layer, 3, tl.bfloat16), _pointer(W, layer, 4, tl.bfloat16),
                                n, 2304, 6144, BN).to(tl.bfloat16).to(tl.float32)
                    tl.store(A + n, val)
                else:
                    n = (task - 6144 // BN) * BN + tl.arange(0, BN)
                    val = _gemv(X, _pointer(W, layer, 5, tl.uint8),
                                _pointer(W, layer, 6, tl.bfloat16), _pointer(W, layer, 7, tl.bfloat16),
                                n, 2304, 576, BN).to(tl.bfloat16).to(tl.float32)
                    tl.store(B + n, val, n < 576)
            # Copy the fed caches, including when they came from another model.
            for t in range(pid * 1024, L * 512, GRID * 1024):
                ix = t + tl.arange(0, 1024)
                tl.store(Cnew + ix, tl.load(Cold + ix, ix < L * 512, 0), ix < L * 512)
            for t in range(pid * 1024, L * 64, GRID * 1024):
                ix = t + tl.arange(0, 1024)
                tl.store(Rnew + ix, tl.load(Rold + ix, ix < L * 64, 0), ix < L * 64)
            _barrier(Counter, GRID)
            if pid == 0:
                r = tl.arange(0, 512)
                tl.store(Cnew + L * 512 + r, tl.load(B + r))
            if pid < 32:
                r = tl.arange(0, 32)
                angle = L * tl.exp(-tl.log(10000.) * (2 * r / 64.))
                co, si = tl.cos(angle), tl.sin(angle)
                qe = tl.load(A + pid * 192 + 128 + r * 2)
                qo = tl.load(A + pid * 192 + 129 + r * 2)
                tl.store(A + pid * 192 + 128 + r * 2, (qe * co - qo * si).to(tl.bfloat16).to(tl.float32))
                tl.store(A + pid * 192 + 129 + r * 2, (qo * co + qe * si).to(tl.bfloat16).to(tl.float32))
                if pid == 0:
                    ke, ko = tl.load(B + 512 + 2 * r), tl.load(B + 513 + 2 * r)
                    tl.store(Rnew + L * 64 + 2 * r, ke * co - ko * si)
                    tl.store(Rnew + L * 64 + 2 * r + 1, ko * co + ke * si)
            # Absorb W_k into each head's query, without storing W_k.
            for task in range(pid, 128, GRID):
                head = task // 4
                k = (task % 4) * 128 + tl.arange(0, 128)
                n = head * 256 + tl.arange(0, 128)
                qp = _pointer(W, layer, 8, tl.uint8)
                sp = _pointer(W, layer, 9, tl.bfloat16)
                zp = _pointer(W, layer, 10, tl.bfloat16)
                packed = tl.load(qp + (k[:, None] // 2) * 8192 + n[None, :])
                iq = (packed >> ((k[:, None] & 1) * 4)) & 15
                sc = tl.load(sp + (k[:, None] // 128) * 8192 + n[None, :]).to(tl.float32)
                ze = tl.load(zp + (k[:, None] // 128) * 8192 + n[None, :]).to(tl.float32)
                wt = ((iq.to(tl.float32) - ze) * sc).to(tl.bfloat16).to(tl.float32)
                q = tl.load(A + head * 192 + tl.arange(0, 128))
                ql = tl.sum(wt * q[None, :], 1)
                tl.store(Qlatent + head * 512 + k, ql)
            _barrier(Counter, GRID)
            # Four independent ranges per head, online softmax in latent space.
            for task in range(pid, 128, GRID):
                head, split = task // 4, task % 4
                d = tl.arange(0, 512)
                rr = tl.arange(0, 64)
                ql = tl.load(Qlatent + head * 512 + d)
                qr = tl.load(A + head * 192 + 128 + rr)
                m = tl.full((), -float('inf'), tl.float32)
                den = tl.full((), 0, tl.float32)
                acc = tl.full((512,), 0, tl.float32)
                chunks = tl.cdiv(L + 1, 4 * 32)
                for block in range(chunks):
                    tok = (split * chunks + block) * 32 + tl.arange(0, 32)
                    c = tl.load(Cnew + tok[:, None] * 512 + d[None, :], tok[:, None] <= L, 0).to(tl.float32)
                    kr = tl.load(Rnew + tok[:, None] * 64 + rr[None, :], tok[:, None] <= L, 0).to(tl.float32)
                    score = (tl.sum(c * ql[None, :], 1) + tl.sum(kr * qr[None, :], 1)) * (192 ** -0.5)
                    score = tl.where(tok <= L, score, -float('inf'))
                    nm = tl.maximum(m, tl.max(score, 0))
                    alpha = tl.exp(m - nm)
                    p = tl.exp(score - nm)
                    acc = acc * alpha + tl.sum(c * p[:, None], 0)
                    den = den * alpha + tl.sum(p, 0)
                    m = nm
                tl.store(Parts + task * 516 + d, acc)
                tl.store(Parts + task * 516 + 512, den)
                tl.store(Parts + task * 516 + 513, m)
            _barrier(Counter, GRID)
            if pid < 32:
                splits = tl.arange(0, 4)
                d = tl.arange(0, 512)
                mm = tl.load(Parts + (pid * 4 + splits) * 516 + 513)
                dd = tl.load(Parts + (pid * 4 + splits) * 516 + 512)
                factor = tl.exp(mm - tl.max(mm, 0))
                av = tl.load(Parts + (pid * 4 + splits[:, None]) * 516 + d[None, :])
                out = tl.sum(av * factor[:, None], 0) / tl.sum(dd * factor, 0)
                tl.store(Latent + pid * 512 + d, out)
            _barrier(Counter, GRID)
            for task in range(pid, 4096 // BN, GRID):
                head = task // (128 // BN)
                n = head * 256 + 128 + (task % (128 // BN)) * BN + tl.arange(0, BN)
                val = _gemv(Latent + head * 512, _pointer(W, layer, 8, tl.uint8),
                            _pointer(W, layer, 9, tl.bfloat16), _pointer(W, layer, 10, tl.bfloat16),
                            n, 512, 8192, BN).to(tl.bfloat16).to(tl.float32)
                tl.store(O + task * BN + tl.arange(0, BN), val)
            _barrier(Counter, GRID)
        # Output projection, residual, and second RMS normalization.
        for task in range(pid, triton.cdiv(2304, BN), GRID):
            n = task * BN + tl.arange(0, BN)
            val = _gemv(O, _pointer(W, layer, 14, tl.uint8),
                        _pointer(W, layer, 15, tl.bfloat16), _pointer(W, layer, 16, tl.bfloat16),
                        n, 4096, 2304, BN).to(tl.bfloat16).to(tl.float32)
            h = tl.load(Hout + n, n < 2304, 0).to(tl.float32)
            tl.store(Hout + n, h + val, n < 2304)
        _barrier(Counter, GRID)
        if pid == 0:
            _norm(Hout, X, _pointer(W, layer, 1, tl.bfloat16))
        _barrier(Counter, GRID)
        # Tiny bf16 router, rounded like nn.Linear in the oracle.
        for expert in range(pid, E, GRID):
            d = tl.arange(0, 4096)
            x = tl.load(X + d, d < 2304, 0)
            r = tl.load(_pointer(W, layer, 19, tl.bfloat16) + expert * 2304 + d, d < 2304, 0).to(tl.float32)
            tl.store(B + expert, tl.sum(x * r, 0).to(tl.bfloat16).to(tl.float32))
        _barrier(Counter, GRID)
        if pid == 0:
            r = tl.arange(0, triton.next_power_of_2(E))
            logits = tl.load(B + r, r < E, -float('inf'))
            maximum = tl.max(logits, 0)
            chosen = tl.full((8,), 0, tl.float32)
            ar = tl.arange(0, 8)
            for j in tl.static_range(8):
                best = tl.max(logits, 0)
                ix = tl.min(tl.where(logits == best, r, 2147483647), 0)
                tl.store(Idx + j, ix)
                chosen = tl.where(ar == j, tl.exp(best - maximum), chosen)
                logits = tl.where(r == ix, -float('inf'), logits)
            weights = chosen / tl.sum(chosen, 0) * 2.446
            tl.store(RW + ar, weights)
        _barrier(Counter, GRID)
        for task in range(pid, 9 * (1024 // BN), GRID):
            selected = task // (1024 // BN)
            n = (task % (1024 // BN)) * BN + tl.arange(0, BN)
            if selected < 8:
                ix = tl.load(Idx + selected).to(tl.int32)
                qg = _pointer(W, layer, 20, tl.uint8) + ix * (2304 // 2) * 1024
                sg = _pointer(W, layer, 21, tl.bfloat16) + ix * 18 * 1024
                zg = _pointer(W, layer, 22, tl.bfloat16) + ix * 18 * 1024
                qu = _pointer(W, layer, 23, tl.uint8) + ix * (2304 // 2) * 1024
                su = _pointer(W, layer, 24, tl.bfloat16) + ix * 18 * 1024
                zu = _pointer(W, layer, 25, tl.bfloat16) + ix * 18 * 1024
            else:
                qg, sg, zg = _pointer(W, layer, 29, tl.uint8), _pointer(W, layer, 30, tl.bfloat16), _pointer(W, layer, 31, tl.bfloat16)
                qu, su, zu = _pointer(W, layer, 32, tl.uint8), _pointer(W, layer, 33, tl.bfloat16), _pointer(W, layer, 34, tl.bfloat16)
            gv = _gemv(X, qg, sg, zg, n, 2304, 1024, BN)
            uv = _gemv(X, qu, su, zu, n, 2304, 1024, BN)
            tl.store(FF + selected * 1024 + n, gv / (1 + tl.exp(-gv)) * uv)
        _barrier(Counter, GRID)
        for task in range(pid, 9 * triton.cdiv(2304, BN), GRID):
            selected = task // triton.cdiv(2304, BN)
            n = (task % triton.cdiv(2304, BN)) * BN + tl.arange(0, BN)
            if selected < 8:
                ix = tl.load(Idx + selected).to(tl.int32)
                qd = _pointer(W, layer, 26, tl.uint8) + ix * 512 * 2304
                sd = _pointer(W, layer, 27, tl.bfloat16) + ix * 8 * 2304
                zd = _pointer(W, layer, 28, tl.bfloat16) + ix * 8 * 2304
            else:
                qd, sd, zd = _pointer(W, layer, 35, tl.uint8), _pointer(W, layer, 36, tl.bfloat16), _pointer(W, layer, 37, tl.bfloat16)
            val = _gemv(FF + selected * 1024, qd, sd, zd, n, 1024, 2304, BN)
            tl.store(FO + selected * 2304 + n, val, n < 2304)
        _barrier(Counter, GRID)
        for task in range(pid, triton.cdiv(2304, 128), GRID):
            n = task * 128 + tl.arange(0, 128)
            total = tl.full((128,), 0, tl.float32)
            for j in tl.static_range(8):
                total += tl.load(FO + j * 2304 + n, n < 2304, 0) * tl.load(RW + j)
            total += tl.load(FO + 8 * 2304 + n, n < 2304, 0)
            total = total.to(tl.bfloat16).to(tl.float32)
            h = tl.load(Hout + n, n < 2304, 0).to(tl.float32)
            tl.store(Hout + n, h + total, n < 2304)
        _barrier(Counter, GRID)


class Model(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self.register_buffer('_counter', torch.zeros(1, dtype=torch.int32), persistent=False)
        self.register_buffer('_work', torch.empty(154000, dtype=torch.float32), persistent=False)
        self.register_buffer('_weights', torch.empty(4, 40, dtype=torch.int64), persistent=False)
        self._ready = False

    def _prepare(self):
        rows = []
        for block in self.blocks:
            a, m = block.attn, block.moe
            row = [0] * 40
            row[0], row[1] = block.attn_norm.data_ptr(), block.moe_norm.data_ptr()
            def qset(slot, mod):
                row[slot:slot + 3] = [mod.w_q.data_ptr(), mod.scales.data_ptr(), mod.zeros.data_ptr()]
            qset(2, a.q_proj)
            if isinstance(a, KDA):
                qset(5, a.k_proj)
                qset(8, a.v_proj)
                qset(11, a.g_proj)
                row[17], row[18] = a.beta_proj.weight.data_ptr(), a.conv_w.data_ptr()
            else:
                qset(5, a.kv_a)
                qset(8, a.kv_b)
            qset(14, a.o_proj)
            row[19] = m.router.weight.data_ptr()
            for slot, mod in zip((20, 23, 26, 29, 32, 35), (m.gate, m.up, m.down, m.s_gate, m.s_up, m.s_down)):
                qset(slot, mod)
            rows.append(row)
        self._weights.copy_(torch.tensor(rows, dtype=torch.int64))
        self._ready = True

    def load_state_dict(self, state_dict, strict=True, assign=False):
        result = super().load_state_dict(state_dict, strict=strict, assign=assign)
        self._prepare()
        return result

    def _apply(self, fn, recurse=True):
        result = super()._apply(fn, recurse)
        self._ready = False
        return result

    @torch.no_grad()
    def step(self, hidden, state):
        if not self._ready:
            self._prepare()
        h = torch.empty_like(hidden)
        oldc, oldr = state[3]['c_kv'], state[3]['k_rope']
        length = oldc.shape[0]
        newc = torch.empty((length + 1, 512), device=hidden.device, dtype=torch.bfloat16)
        newr = torch.empty((length + 1, 64), device=hidden.device, dtype=torch.bfloat16)
        _mega[(128,)](hidden, h, self._weights, self._counter, self._work,
            state[0]['S'], state[0]['cq'], state[0]['ck'], state[0]['cv'],
            state[1]['S'], state[1]['cq'], state[1]['ck'], state[1]['cv'],
            state[2]['S'], state[2]['cq'], state[2]['ck'], state[2]['cv'],
            oldc, oldr, newc, newr, length,
            E=self.cfg.n_experts, GRID=128, BN=32, num_warps=4, enable_fp_fusion=False)
        state[3]['c_kv'], state[3]['k_rope'] = newc, newr
        return h, state