| """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 |
| NCHUNK = 5 |
| MAX_STEPS = 256 |
| ROPE_DIM = 64 |
| 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): |
| 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() |
| 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) |
|
|
