"""Runtime for the Ling-mini-2.0 megakernel: loads this repository's packed 4-bit model onto the GPU and runs the kernel. One kernel launch runs every layer of up to 256 tokens and chooses each next token itself (greedy decoding). from ling_kernel import LingKernel mk = LingKernel() preds = mk.forced(prompt_ids) # reads a prompt; preds[-1] is the next token new = mk.greedy(preds[-1], len(prompt_ids), 64) # 64 more tokens Needs an NVIDIA RTX 50-series GPU (sm_120), the kernel library (megakernel_ling.dll on Windows, built from megakernel_ling.cu) and PyTorch with CUDA. """ from __future__ import annotations import ctypes import json import sys import time from pathlib import Path import torch from safetensors import safe_open HERE = Path(__file__).resolve().parent H, NL, NQ, NKV, HD, NE, TOPK, FI, FD, VOCAB = 2048, 20, 16, 4, 128, 256, 8, 512, 5120, 157184 QD, KVD = NQ * HD, NKV * HD QKVD = QD + 2 * KVD QKVB = QKVD + 4 # one q/k/v accumulator: the 3072 sums, then the sum of x^2 NCHUNK = 5 MAX_STEPS = 256 # tokens per launch ROPE_DIM = 64 # rotated dims per head KERNEL_FORMAT = 4 POINTERS = ("layers", "dense0", "embed", "final_norm", "head_c", "head_s", "rope_cos", "rope_sin", "k_cache", "v_cache", "x", "qkv", "act", "att_m", "att_l", "att_o", "logits", "best", "in_tokens", "out_tokens", "bar", "stamps", "experts", "racc", "head_order") INTS = ("max_pos", "pos0", "n_steps", "forced", "write_logits") class Params(ctypes.Structure): _fields_ = [(n, ctypes.c_void_p) for n in POINTERS] + [(n, ctypes.c_int) for n in INTS] def library(path: Path | None = None): name = "megakernel_ling.dll" if sys.platform == "win32" else "libmegakernel_ling.so" lib = ctypes.CDLL(str(path or HERE / name)) lib.mk_launch.argtypes = [ctypes.POINTER(Params)] lib.mk_launch.restype = ctypes.c_int lib.mk_launch_async.argtypes = [ctypes.POINTER(Params), ctypes.c_int, ctypes.c_void_p] lib.mk_launch_async.restype = ctypes.c_int lib.mk_error.argtypes = [ctypes.c_int] lib.mk_error.restype = ctypes.c_char_p lib.mk_layer_size.restype = ctypes.c_longlong lib.mk_dense0_size.restype = ctypes.c_longlong if lib.mk_format() != KERNEL_FORMAT or lib.mk_params_size() != ctypes.sizeof(Params): raise RuntimeError("megakernel_ling library and ling_kernel.py do not match") return lib def rope_tables(theta: float, n: int): """cos/sin [n, 64] for positions 0..n-1 (RoPE on the first 64 dims of each head).""" inv = 1.0 / (theta ** (torch.arange(0, ROPE_DIM, 2, device="cuda", dtype=torch.float32) / ROPE_DIM)) f = torch.arange(n, device="cuda", dtype=torch.float32)[:, None] * inv[None] emb = torch.cat([f, f], -1) return emb.cos().contiguous(), emb.sin().contiguous() class State: """One conversation's GPU state: its KV cache and the kernel's work buffers.""" def __init__(self, max_pos: int): f32 = dict(device="cuda", dtype=torch.float32) self.max_pos = max_pos self.k_cache = torch.zeros(NL, max_pos, KVD, device="cuda", dtype=torch.float16) self.v_cache = torch.zeros(NL, max_pos, KVD, device="cuda", dtype=torch.float16) self.x, self.qkv, self.act = torch.zeros(H, **f32), torch.zeros(2 * QKVB, **f32), torch.zeros(FD, **f32) self.att_m, self.att_l = torch.zeros(NQ * NCHUNK, **f32), torch.zeros(NQ * NCHUNK, **f32) self.att_o = torch.zeros(NQ * NCHUNK * HD, **f32) self.best = torch.zeros(2, device="cuda", dtype=torch.int64) self.bar = torch.zeros(2, device="cuda", dtype=torch.int32) self.racc = torch.zeros(2 * (NE + 1), **f32) self.out = torch.zeros(MAX_STEPS, device="cuda", dtype=torch.int32) self.keep = None def zero(self) -> None: for t in (self.best, self.bar, self.racc, self.qkv, self.act): t.zero_() class LingKernel: def __init__(self, max_pos: int = 4096, log=print, dll: Path | None = None): t = time.time() self.lib = library(dll) meta = json.loads((HERE / "packed" / "meta.json").read_text()) LB, DB = self.lib.mk_layer_size(), self.lib.mk_dense0_size() if (meta["kernel_format"], meta["layer_bytes"], meta["dense0_bytes"]) != (KERNEL_FORMAT, LB, DB): raise RuntimeError("packed/ was written for another kernel layout") self.config = json.loads((HERE / "config.json").read_text()) self.layers = torch.empty(NL * LB, dtype=torch.uint8, device="cuda") with safe_open(str(HERE / "packed" / "layers.safetensors"), framework="pt") as f: part = f.get_slice("layers") n = self.layers.numel() for a in range(0, n, 1 << 28): # 256 MB at a time b = min(n, a + (1 << 28)) self.layers[a:b].copy_(part[a:b]) with safe_open(str(HERE / "packed" / "extras.safetensors"), framework="pt") as f: g = lambda k: f.get_tensor(k).cuda().contiguous() # noqa: E731 self.head_c, self.head_s, self.head_order = g("head_c"), g("head_s"), g("head_order") self.dense0, self.embed, self.final_norm = g("dense0"), g("embed"), g("final_norm") self.rope_cos, self.rope_sin = rope_tables(self.config["rope_theta"], max_pos) self.max_pos = max_pos self.sms = self.lib.mk_blocks() self.state = State(max_pos) self.logits = None if log: log(f" loaded the 4-bit model ({(self.layers.numel() + self.dense0.numel()) / 1e9:.2f} GB of layers) " f"in {time.time() - t:.0f} s") def params(self, st: State, tokens: torch.Tensor, pos0: int, n: int, forced: bool, logits: bool = False) -> Params: ptr = {"layers": self.layers, "dense0": self.dense0, "embed": self.embed, "final_norm": self.final_norm, "head_c": self.head_c, "head_s": self.head_s, "rope_cos": self.rope_cos, "rope_sin": self.rope_sin, "k_cache": st.k_cache, "v_cache": st.v_cache, "x": st.x, "qkv": st.qkv, "act": st.act, "att_m": st.att_m, "att_l": st.att_l, "att_o": st.att_o, "logits": self.logits if logits else None, "best": st.best, "in_tokens": tokens, "out_tokens": st.out, "bar": st.bar, "stamps": None, "experts": None, "racc": st.racc, "head_order": self.head_order} return Params(**{k: (v.data_ptr() if v is not None else None) for k, v in ptr.items()}, max_pos=st.max_pos, pos0=pos0, n_steps=n, forced=int(forced), write_logits=int(logits)) def _launch(self, ins: list[int], pos0: int, n: int, forced: bool, logits: bool = False) -> float: """One synchronous launch on the whole GPU with this object's own state; returns seconds.""" st = self.state assert 1 <= n <= MAX_STEPS and pos0 + n <= st.max_pos if logits and self.logits is None: self.logits = torch.empty(MAX_STEPS, VOCAB, device="cuda", dtype=torch.float32) tokens = torch.tensor(ins, device="cuda", dtype=torch.int32) st.zero() p = self.params(st, tokens, pos0, n, forced, logits) torch.cuda.synchronize() t = time.perf_counter() err = self.lib.mk_launch(ctypes.byref(p)) seconds = time.perf_counter() - t if err: raise RuntimeError(f"megakernel launch failed: {self.lib.mk_error(err).decode()}") return seconds def launch_async(self, st: State, ins: list[int], pos0: int, n: int, forced: bool, nblocks: int, stream) -> None: """A launch for another conversation on nblocks SMs and its own CUDA stream (several chats at once); returns at once, the caller synchronizes. Its output tokens land in st.out.""" with torch.cuda.stream(stream): tokens = torch.tensor(ins, device="cuda", dtype=torch.int32) st.zero() p = self.params(st, tokens, pos0, n, forced) err = self.lib.mk_launch_async(ctypes.byref(p), nblocks, ctypes.c_void_p(stream.cuda_stream)) if err: raise RuntimeError(f"launch failed ({err})") st.keep = tokens def forced(self, tokens: list[int], pos0: int = 0, logits: bool = False): """Teacher-forced: one step per given token from position pos0. Returns the predicted next token after each (and all logits when asked).""" preds, rows = [], [] for c in range(0, len(tokens), MAX_STEPS): chunk = tokens[c:c + MAX_STEPS] self._launch(chunk, pos0 + c, len(chunk), True, logits) preds += self.state.out[:len(chunk)].tolist() if logits: rows.append(self.logits[:len(chunk)].clone()) return (preds, torch.cat(rows)) if logits else preds def greedy(self, first: int, pos0: int, n: int) -> list[int]: """n greedy tokens after `first`, which is fed at position pos0.""" new = [] while len(new) < n: k = min(MAX_STEPS, n - len(new)) self._launch([new[-1] if new else first], pos0 + len(new), k, False) new += self.state.out[:k].tolist() return new def timed_greedy(self, first: int, pos0: int, n: int) -> float: return self._launch([first], pos0, n, False)