k1_dev/k1_dev.py

"""fableous / ultra-kernels K1 dev harness.

Standalone microbenchmark of the Gemma4 MTP drafter iteration loop on GPU.
Mirrors vllm/model_executor/models/gemma4_mtp.py math exactly (RMSNorm without
+1, q_norm per head, neox RoPE, scaling=1.0, gelu_tanh GEGLU, layer_scalar,
centroid head) with synthetic bf16 weights and a synthetic contiguous KV cache
(drafter is Q-only / KV-shared: K,V are FIXED across all 7 iterations).

Variants timed (7-iteration loop, batch 1):
  ref_eager  -- plain PyTorch ops
  ref_graph  -- same captured in a CUDA graph (mimics onegraph)
  mega       -- cooperative-launch CUDA megakernel (if MEGA_SRC present)

Prints per-variant: ms/loop, ms/iter, and max-abs-diff of drafted tokens +
final backbone hidden vs ref_eager.
"""

from __future__ import annotations

import math
import os
import time

import torch

torch.manual_seed(0)

DEV = "cuda"
DT = torch.bfloat16

# drafter dims (gemma-4-E4B assistant config)
HID = 256
INTER = 2048
NHEAD = 4
NKV = 2
HD_S = 256          # sliding-layer head_dim
HD_G = 512          # global-layer head_dim
BB = 2560           # backbone hidden size
VOCAB = 262144
NCENT = 2048
VPC = VOCAB // NCENT            # 128 tokens per centroid
KCENT = int(os.environ.get("KCENT", "64"))   # CENTROID_TOP_K in the stack
LAYER_TYPES = ["S", "S", "S", "F"]
WINDOW = 512
SEQ = int(os.environ.get("SEQ", "1024"))     # committed context length
K_SPEC = 7
EPS = 1e-6
ROPE_LOCAL = 10000.0
ROPE_GLOBAL = 1000000.0
NORMALIZER = BB ** 0.5


def make_weights() -> dict:
    g = lambda *s: (torch.randn(*s, device=DEV, dtype=torch.float32) * 0.02).to(DT)
    w = {
        "embed": g(VOCAB, BB),                  # shared backbone-dim table
        "pre_proj": g(HID, 2 * BB),
        "post_proj": g(BB, HID),
        "final_norm": g(HID),
        "cent_w": g(NCENT, HID),
        "lm_head": g(VOCAB, HID),               # draft-dim head table
        "token_ordering": torch.randperm(VOCAB, device=DEV),
        "layers": [],
    }
    for lt in LAYER_TYPES:
        hd = HD_G if lt == "F" else HD_S
        w["layers"].append(
            {
                "in_norm": g(HID),
                "q_proj": g(NHEAD * hd, HID),
                "q_norm": g(hd),
                "post_attn_norm": g(HID),
                "pre_ffn_norm": g(HID),
                "gate": g(INTER, HID),
                "up": g(INTER, HID),
                "down": g(HID, INTER),
                "post_ffn_norm": g(HID),
                "layer_scalar": torch.ones(1, device=DEV, dtype=torch.float32),
            }
        )
    return w


def make_kv() -> list:
    kv = []
    for lt in LAYER_TYPES:
        hd = HD_G if lt == "F" else HD_S
        k = (torch.randn(SEQ, NKV, hd, device=DEV, dtype=torch.float32) * 0.5).to(DT)
        v = (torch.randn(SEQ, NKV, hd, device=DEV, dtype=torch.float32) * 0.5).to(DT)
        kv.append((k, v))
    return kv


def rms(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    # production CUDA rms_norm: f32 internal, ONE rounding after weight mult
    xf = x.float()
    return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + EPS) * w.float()).to(x.dtype)


ROPE_PARTIAL_F = 0.25   # global layer: proportional partial rotary


def make_cs_caches() -> list:
    """Per-layer f32 cos_sin caches [SEQ, hd] mirroring vllm:
    sliding = default full rotation; global = proportional partial (0.25),
    identity (cos=1, sin=0) on non-rotated pairs."""
    caches = []
    posv = torch.arange(SEQ, device=DEV, dtype=torch.float32)
    for lt in LAYER_TYPES:
        hd = HD_G if lt == "F" else HD_S
        half = hd // 2
        if lt == "F":
            n_rot = int(hd * ROPE_PARTIAL_F) // 2          # 64 pairs
            exp = torch.arange(0, 2 * n_rot, 2, device=DEV, dtype=torch.float32) / hd
            inv = 1.0 / (ROPE_GLOBAL ** exp)
            inv = torch.cat([inv, torch.zeros(half - n_rot, device=DEV)])
        else:
            exp = torch.arange(0, hd, 2, device=DEV, dtype=torch.float32) / hd
            inv = 1.0 / (ROPE_LOCAL ** exp)
        ang = posv[:, None] * inv[None, :]
        caches.append(torch.cat([torch.cos(ang), torch.sin(ang)], dim=-1).contiguous())
    return caches


def rope(q: torch.Tensor, pos: int, cs: torch.Tensor) -> torch.Tensor:
    hd = q.shape[-1]
    half = hd // 2
    cos, sin = cs[pos, :half], cs[pos, half:]
    qf = q.float()
    q1, q2 = qf[..., :half], qf[..., half:]
    return torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1).to(q.dtype)


def ref_iteration(w: dict, kv: list, cs_caches: list, token: torch.Tensor,
                  hidden: torch.Tensor, pos: int):
    """One drafter iteration. token: int64[1]; hidden: bf16[1, BB] backbone-dim."""
    emb = (w["embed"][token].float() * NORMALIZER).to(DT)     # [1, BB] (round)
    x = torch.cat([emb, hidden], dim=-1)                       # [1, 2*BB]
    h = (x.float() @ w["pre_proj"].float().T).to(DT)           # [1, HID]

    for li, lt in enumerate(LAYER_TYPES):
        lw = w["layers"][li]
        hd = HD_G if lt == "F" else HD_S
        theta = ROPE_GLOBAL if lt == "F" else ROPE_LOCAL
        residual = h
        hn = rms(h, lw["in_norm"])
        q = (hn.float() @ lw["q_proj"].float().T).to(DT).view(NHEAD, hd)
        q = rms(q, lw["q_norm"])
        q = rope(q, pos, cs_caches[li])

        k, v = kv[li]
        if lt == "S":
            lo = max(0, pos - WINDOW + 1) if WINDOW else 0
        else:
            lo = 0
        ks, vs = k[lo:pos + 1], v[lo:pos + 1]                  # [T, NKV, hd]
        rep = NHEAD // NKV
        attn_out = torch.empty(NHEAD, hd, device=DEV, dtype=torch.float32)
        for hh in range(NHEAD):
            kvh = hh // rep
            scores = (ks[:, kvh].float() @ q[hh].float())       # [T]
            p = torch.softmax(scores, dim=0)
            attn_out[hh] = p @ vs[:, kvh].float()
        a = attn_out.reshape(1, NHEAD * hd).to(DT)

        # o_proj weight: reuse q_proj-transposed shape -- store separately
        o = (a.float() @ w["layers"][li]["o_proj"].float().T).to(DT)  # [1, HID]
        o = rms(o, lw["post_attn_norm"])
        h = o + residual
        residual = h
        hn = rms(h, lw["pre_ffn_norm"])
        gate = (hn.float() @ lw["gate"].float().T).to(DT)
        up = (hn.float() @ lw["up"].float().T).to(DT)
        gl = torch.nn.functional.gelu(gate.float(), approximate="tanh").to(DT)
        act = gl * up                                  # bf16 mult (rounds)
        d = (act.float() @ lw["down"].float().T).to(DT)
        d = rms(d, lw["post_ffn_norm"])
        h = (d + residual) * lw["layer_scalar"].to(DT)

    draft_h = rms(h, w["final_norm"])                          # [1, HID]
    backbone_h = (draft_h.float() @ w["post_proj"].float().T).to(DT)  # [1, BB]

    # centroid head: bf16 scores, deterministic ascending centroid order,
    # bf16 logits, tie-break LEFT (matches kernel + production fused argmax)
    cs = (draft_h.float() @ w["cent_w"].float().T).to(DT)      # [1, NCENT] bf16
    _, top_c = torch.topk(cs.float(), k=KCENT, dim=-1)         # [1, KCENT]
    top_c, _ = torch.sort(top_c, dim=-1)                       # ascending ids
    clusters = w["token_ordering"].view(NCENT, VPC)
    selected = clusters[top_c]                                 # [1, KCENT, VPC]
    sel = selected.reshape(-1)
    rows = w["lm_head"][sel].view(1, KCENT * VPC, HID)
    logits = torch.einsum(
        "td,tsd->ts", draft_h.float(), rows.float()
    ).to(DT).float()                                           # bf16-rounded
    nxt = sel[logits.argmax(-1)].view(1)                       # first max = tie LEFT
    return nxt, backbone_h


def add_oproj(w: dict) -> None:
    g = lambda *s: (torch.randn(*s, device=DEV, dtype=torch.float32) * 0.02).to(DT)
    for li, lt in enumerate(LAYER_TYPES):
        hd = HD_G if lt == "F" else HD_S
        w["layers"][li]["o_proj"] = g(HID, NHEAD * hd)


def run_loop(w, kv, cs_caches, first_token, hidden0, pos):
    tokens = []
    hidden = hidden0
    tok = first_token
    for _ in range(K_SPEC):
        tok, hidden = ref_iteration(w, kv, cs_caches, tok, hidden, pos)
        tokens.append(tok)
    return torch.cat(tokens), hidden


def bench(fn, n_warm=10, n_iter=50) -> float:
    for _ in range(n_warm):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n_iter):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / n_iter * 1e3


def main() -> None:
    print(f"[k1] device={torch.cuda.get_device_name(0)} torch={torch.__version__}")
    print(f"[k1] SEQ={SEQ} KCENT={KCENT} K_SPEC={K_SPEC}")
    w = make_weights()
    add_oproj(w)
    kv = make_kv()
    cs_caches = make_cs_caches()
    first = torch.randint(0, VOCAB, (1,), device=DEV)
    hidden0 = (torch.randn(1, BB, device=DEV, dtype=torch.float32) * 0.5).to(DT)
    pos = SEQ - 1

    # reference (eager)
    ref_tokens, ref_hidden = run_loop(w, kv, cs_caches, first, hidden0, pos)
    print(f"[k1] ref tokens: {ref_tokens.tolist()}")

    t_eager = bench(lambda: run_loop(w, kv, cs_caches, first, hidden0, pos))
    print(f"[k1] ref_eager : {t_eager:8.3f} ms/loop  {t_eager/K_SPEC:7.3f} ms/iter")

    # CUDA-graphed loop (static buffers like onegraph)
    s_first = first.clone()
    s_hidden = hidden0.clone()
    s_out = torch.zeros(K_SPEC, dtype=torch.int64, device=DEV)

    def graph_body():
        hidden = s_hidden
        tok = s_first
        for i in range(K_SPEC):
            tok, hidden = ref_iteration(w, kv, cs_caches, tok, hidden, pos)
            s_out[i] = tok[0]

    for _ in range(3):
        graph_body()
    torch.cuda.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        graph_body()
    graph.replay()
    torch.cuda.synchronize()
    assert s_out.tolist() == ref_tokens.tolist(), (s_out.tolist(), ref_tokens.tolist())
    t_graph = bench(graph.replay)
    print(f"[k1] ref_graph : {t_graph:8.3f} ms/loop  {t_graph/K_SPEC:7.3f} ms/iter")

    # roofline estimate
    wb = sum(
        t.numel() * t.element_size()
        for lw in w["layers"]
        for t in lw.values()
    ) + w["pre_proj"].numel() * 2 + w["post_proj"].numel() * 2 + w["cent_w"].numel() * 2
    kvb = 0
    for li, lt in enumerate(LAYER_TYPES):
        hd = HD_G if lt == "F" else HD_S
        t = min(SEQ, WINDOW) if lt == "S" else SEQ
        kvb += t * NKV * hd * 2 * 2
    gather = KCENT * VPC * HID * 2
    per_iter = (wb + kvb + gather) / 1e6
    print(f"[k1] roofline  : ~{per_iter:.1f} MB/iter -> {per_iter/600*1000:7.3f} ms/iter @600GB/s")

    if os.path.exists("megakernel.cu"):
        import mega_runner
        mega_runner.run_mega(
            w, kv, cs_caches, first, hidden0, pos, SEQ, KCENT,
            ref_tokens, ref_hidden, k_spec=K_SPEC,
        )


if __name__ == "__main__":
    main()