File size: 9,310 Bytes
e7654c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | """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)
|