IQ.Lvbs/IQ.Pilot
8
7
Fork 2
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

IQ.Pilot Prebuilt Release @ 27f668a

Browse Source
This commit is contained in:
IQ.Lvbs CI [bot]
2026-09-03 18:23:24 -05:00
commit b073c5182b
2554 changed files with 679696 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
artifacts/package_runtime/tinygrad/llm/__init__.py Normal file
View File

2
artifacts/package_runtime/tinygrad/llm/__main__.py Normal file
View File

@@ -0,0 +1,2 @@
from tinygrad.llm.cli import main
if __name__ == "__main__": main()

42
artifacts/package_runtime/tinygrad/llm/chat.html Normal file
View File

@@ -0,0 +1,42 @@
<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
* { margin: 0 }
body { background: #212121; color: #e3e3e3; font-family: system-ui;
height: 100vh; display: flex; flex-direction: column }
#chat { flex: 1; overflow-y: auto; padding: 20px }
.msg { padding: 10px 16px; margin: 8px 0; white-space: pre-wrap; border-radius: 18px }
.user { background: #2f2f2f; margin-left: auto; width: fit-content; max-width: 70% }
#input { max-width: 768px; width: 100%; margin: 20px auto; padding: 14px 20px;
background: #2f2f2f; color: inherit; font: inherit;
border: none; outline: none; resize: none; border-radius: 24px; field-sizing: content }
</style></head><body><div id="chat"></div>
<textarea id="input" rows="1" placeholder="Ask anything" autofocus></textarea>
<script>
input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) { e.preventDefault(); send() } }
const msgs = [];
async function send() {
if (!input.value.trim()) return;
msgs.push({role: 'user', content: input.value.trim()});
chat.innerHTML += '<div class="msg user">' + input.value.trim().replace(/</g, '&lt;') + '</div>';
input.value = '';
const d = document.createElement('div'); d.className = 'msg'; chat.appendChild(d);
const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'},
body: JSON.stringify({model: 'llama', messages: msgs, stream: true, temperature: 0.7})});
let buf = '', txt = '', rsn = '';
for (const rd = r.body.getReader(), dec = new TextDecoder();;) {
const {done, value} = await rd.read();
if (done) break;
buf += dec.decode(value, {stream: true});
const lines = buf.split('\n');
buf = lines.pop();
for (const ln of lines)
if (ln.startsWith('data: ') && !ln.includes('[DONE]'))
try { const dl = JSON.parse(ln.slice(6)).choices[0]?.delta;
if (dl?.reasoning_content) { const s = document.createElement('span'); s.style.color = '#888';
s.textContent = dl.reasoning_content; rsn += dl.reasoning_content; d.appendChild(s) }
if (dl?.content) { const s = document.createElement('span');
s.textContent = dl.content; txt += dl.content; d.appendChild(s) } } catch {}
chat.scrollTop = chat.scrollHeight;
}
const m = {role:'assistant', content:txt}; if (rsn) m.reasoning_content = rsn; msgs.push(m);
}
</script></body></html>

203
artifacts/package_runtime/tinygrad/llm/cli.py Normal file
View File

@@ -0,0 +1,203 @@
from __future__ import annotations
import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time
from typing import TYPE_CHECKING
from tinygrad import nn
from tinygrad.uop.ops import UOp, Ops
from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, Context, fetch, profile_marker, getenv
from tinygrad.llm.model import Transformer
if TYPE_CHECKING:
import jinja2
class SimpleTokenizer:
def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3",
bos_id:int|None=None, eos_id:int=0, eot_id:int|None=None):
preset = {"qwen35":"qwen2","qwen35moe":"qwen2"}.get(preset, preset)
if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo","kimi-k2","tekken","glm4"):
raise ValueError(f"Invalid tokenizer preset '{preset}'")
# https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
bs = [*range(33, 127), *range(161, 173), *range(174, 256)] # bytes that map to themselves
self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
# https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
# 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
# compact adjacent codepoints into ranges: listing them all makes re spend seconds on large prompts
def ucat_range(pre:str) -> str:
cps = enumerate(cp for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
runs = [list(g) for _, g in itertools.groupby(cps, lambda e: e[1]-e[0])]
return "".join(re.escape(chr(g[0][1])) + (f"-{re.escape(chr(g[-1][1]))}" if len(g) > 1 else "") for g in runs)
r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)")
self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}
self._special_tokens = special_tokens
self._tok2bytes = {tid: tok for tok, tid in self._normal_tokens.items()} | {tid: tok.encode() for tok, tid in self._special_tokens.items()}
self.preset = preset
self.bos_id, self.eos_id, self.eot_id = bos_id, eos_id, eot_id
@staticmethod
def from_gguf_kv(kv:dict):
# https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
normal_tokens, special_tokens = partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
special_tokens_dict = dict(special_tokens)
return SimpleTokenizer(dict(normal_tokens), special_tokens_dict, kv["tokenizer.ggml.pre"],
bos_id=kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None,
eos_id=kv.get('tokenizer.ggml.eos_token_id', 0), eot_id=kv.get('tokenizer.ggml.eot_token_id', special_tokens_dict.get('<|im_end|>')))
def _encode_word(self, word:bytes) -> list[int]:
if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
parts = [bytes([b]) for b in word]
# greedily merge any parts that we can
while True:
i = min([(sys.maxsize, -1)] + [(self._normal_tokens.get(parts[j]+parts[j+1], sys.maxsize), j) for j in range(len(parts)-1)])[1]
if i == -1: break
parts[i:i+2] = [parts[i] + parts[i+1]]
try: return [self._normal_tokens[p] for p in parts]
except KeyError: raise RuntimeError("token not found")
def _encode_sentence(self, chunk:str) -> list[int]:
return [tok for word in self._split_to_word.findall(chunk) for tok in self._encode_word(word.encode())]
def encode(self, text:str) -> list[int]:
tokens: list[int] = []
pos = 0
for match in self._split_to_sentence.finditer(text):
tokens.extend(self._encode_sentence(text[pos:match.start(0)]) + [self._special_tokens[text[match.start(0):match.end(0)]]])
pos = match.end(0)
return tokens + self._encode_sentence(text[pos:])
def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode(errors='replace')
def stream_decoder(self) -> typing.Callable[..., str]:
dec = codecs.getincrementaldecoder('utf-8')('replace')
def _decode(tid:int|None=None) -> str: return dec.decode(self._tok2bytes[tid]) if tid is not None else dec.decode(b'', final=True)
return _decode
def is_end(self, token_id:int) -> bool: return token_id in (self.eos_id, self.eot_id)
models = {
"llama3.2:1b": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
"llama3.2:1b-q4": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q4_K_M.gguf",
"llama3.2:3b": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q6_K.gguf",
"llama3.2:3b-f16": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-f16.gguf",
"llama3.1:8b": "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf",
"qwen3:0.6b": "https://huggingface.co/Qwen/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q8_0.gguf",
"qwen3:1.7b": "https://huggingface.co/unsloth/Qwen3-1.7B-GGUF/resolve/main/Qwen3-1.7B-Q4_K_M.gguf",
"qwen3:8b": "https://huggingface.co/Qwen/Qwen3-8B-GGUF/resolve/main/Qwen3-8B-Q4_K_M.gguf",
"qwen3:30b-a3b": "https://huggingface.co/Qwen/Qwen3-30B-A3B-GGUF/resolve/main/Qwen3-30B-A3B-Q4_K_M.gguf",
"qwen3.5:0.8b": "https://huggingface.co/unsloth/Qwen3.5-0.8B-GGUF/resolve/main/Qwen3.5-0.8B-Q8_0.gguf",
"qwen3.5:4b": "https://huggingface.co/unsloth/Qwen3.5-4B-GGUF/resolve/main/Qwen3.5-4B-Q4_K_M.gguf",
"qwen3.5:9b": "https://huggingface.co/unsloth/Qwen3.5-9B-GGUF/resolve/main/Qwen3.5-9B-Q4_K_M.gguf",
"qwen3.6:27b": "https://huggingface.co/unsloth/Qwen3.6-27B-GGUF/resolve/main/Qwen3.6-27B-Q4_K_M.gguf",
"qwen3.6:35b-a3b": "https://huggingface.co/unsloth/Qwen3.6-35B-A3B-GGUF/resolve/main/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf",
"olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
"moonlight": "https://huggingface.co/gabriellarson/Moonlight-16B-A3B-Instruct-GGUF/resolve/main/Moonlight-16B-A3B-Instruct-Q4_K_M.gguf",
"glm-4.7-flash": "https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/resolve/main/GLM-4.7-Flash-Q4_K_M.gguf",
}
class FallbackTemplate:
# minimal jinja2.Template-compatible chat template without jinja2, no tool calling support
def __init__(self, tok:SimpleTokenizer): self.tok = tok
def role(self, role:str) -> str:
if self.tok.preset == 'olmo': return "<|" + role + "|>\n" # OLMoE Instruct format
if self.tok.preset == 'kimi-k2': return "<|im_" + role + "|>" + role + "<|im_middle|>"
if self.tok.preset == 'qwen2': return "<|im_start|>" + role + "\n"
if self.tok.preset == 'glm4': return "<|" + role + "|>"
if self.tok.preset == 'tekken':
if role == 'user': return "[INST]"
if role == 'assistant': return ""
raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.tok.preset}'")
return "<|start_header_id|>" + role + "<|end_header_id|>\n\n"
def end_turn(self) -> str:
if self.tok.preset == 'olmo': return "\n"
if self.tok.preset == 'kimi-k2': return self.tok.decode([self.tok.eos_id])
if self.tok.preset == 'qwen2': return self.tok.decode([self.tok.eos_id]) + "\n"
if self.tok.preset == 'glm4': return ""
if self.tok.preset == 'tekken': return "[/INST]"
return self.tok.decode([self.tok.eos_id])
def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True, preserve_thinking:bool=False) -> str:
out = self.tok.decode([] if self.tok.bos_id is None else [self.tok.bos_id]) + ("<sop>" if self.tok.preset == 'glm4' else "")
for msg in messages:
out += self.role(msg["role"])
content = msg.get("content")
if isinstance(content, str): out += content
elif isinstance(content, list):
for c in content:
if c["type"] == "text": out += c["text"]
else: raise RuntimeError(f"unhandled type: {c['type']}")
elif content is not None: raise RuntimeError(f"unknown content type: {type(content)}")
out += self.end_turn()
return out + self.role("assistant") if add_generation_prompt else out
from tinygrad.llm.serve import LLMServer
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", "-m", default=list(models.keys())[0], help=f"Model choice ({', '.join(models.keys())}) or path to a local GGUF file")
parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
parser.add_argument("--serve", nargs='?', type=int, const=8000, metavar="PORT", help="Run OpenAI compatible API (optional port, default 8000)")
parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
args = parser.parse_args()
# load the model
model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
model_name = kv.get('general.name') or kv.get('general.basename') or args.model
file_sizes = [y.nbytes() for y in UOp.sink(*[x.uop for x in nn.state.get_parameters(model)]).toposort() if y.op is Ops.BUFFER]
print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params, "
f"max context {args.max_context} on {nn.state.get_parameters(model)[0].device}")
# get tokenizer
tok = SimpleTokenizer.from_gguf_kv(kv)
# use the model's chat template if jinja2 is available (enables model-specific formatting)
template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
if (ct := kv.get('tokenizer.chat_template')) is not None:
try:
import jinja2
env = jinja2.Environment()
env.filters['tojson'] = lambda obj, **kwargs: json.dumps(obj, **kwargs) # jinja2's tojson escapes <>& for HTML safety
env.globals['raise_exception'] = lambda msg: (_ for _ in ()).throw(RuntimeError(msg))
env.globals['strftime_now'] = lambda fmt: time.strftime(fmt)
env.globals['bos_token'] = tok.decode([tok.bos_id]) if tok.bos_id is not None else ""
env.globals['eos_token'] = tok.decode([tok.eos_id])
template = env.from_string(ct)
except ImportError: print("warning: jinja2 is not installed, the model's chat template is disabled")
# warmup the JIT
if args.warmup or args.serve:
with Context(DEBUG=max(DEBUG.value, 1)): model.warmup()
# start server
if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()
# do benchmark
if args.benchmark is not None:
gen = model.generate(toks:=[tok.bos_id or 0])
for i in range(args.benchmark):
profile_marker(f"decode @ {i}")
GlobalCounters.reset()
if (log:=getenv("BENCHMARK_LOG", "")): from extra.bench_log import WallTimeEvent, BenchEvent
with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s,"
f" {GlobalCounters.global_mem//1000000}/{GlobalCounters.mem_used//1000000} MB -- "+\
tok.decode(toks).replace("\n", "\\n")):
if log:
with WallTimeEvent(BenchEvent.STEP): next(gen)
else: next(gen)
exit(0)
# interactive chat
messages: list[dict] = []
while 1:
try: messages.append({"role":"user", "content":input('>>> ')})
except EOFError: break
ids = tok.encode(template.render(messages=messages, add_generation_prompt=True))
reply, dec = "", tok.stream_decoder()
for next_id in model.generate(ids):
if tok.is_end(next_id):
sys.stdout.write(dec() + "\n\n")
break
reply += (piece := dec(next_id))
sys.stdout.write(piece)
sys.stdout.flush()
messages.append({"role":"assistant", "content":reply})
if __name__ == "__main__": main()

176
artifacts/package_runtime/tinygrad/llm/gguf.py Normal file
View File

@@ -0,0 +1,176 @@
import functools, io, pathlib, re, struct
from typing import Any, Callable
from tinygrad.tensor import Tensor
from tinygrad.dtype import dtypes
from tinygrad.helpers import prod, round_up
from tinygrad.nn.state import TensorIO
# ggml packs each iq grid entry as N bytes (N=4 for uint32 grids, N=8 for uint64 grids) in a single word. See ggml-common.h.
@functools.lru_cache(None)
def _ggml_iq_grid(device: str, grid: tuple[int, ...], grid_shape: tuple[int, int]) -> Tensor:
values = [float((w >> (8*i)) & 0xFF) for w in grid for i in range(grid_shape[1])]
return Tensor(values, dtype=dtypes.float32, device=device).reshape(grid_shape)
# native types {ggml_type: dtype}
_GGML_NATIVE = {0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8, 25: dtypes.int16,
26: dtypes.int32, 27: dtypes.int64, 28: dtypes.float64, 30: dtypes.bfloat16}
# quant types {ggml_type: (number of elements, number of bytes)}
_GGML_QUANT = {2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17), 41:(128,18)}
def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
"""
Converts ggml tensor data to a tinygrad tensor.
Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 24),
int16 (id: 25), int32 (id: 26), int64 (id: 27), float64 (id: 28), bfloat16 (id: 30)
Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q5_0 (id: 6),
Q5_1 (id: 7), Q8_0 (id: 8), Q4_K (id: 12), Q5_K (id: 13),
Q6_K (id: 14), IQ3_XXS (id: 18), IQ3_S (id: 21), IQ2_S (id: 22), IQ4_XS (id: 23), MXFP4 (id: 39), Q1_0 (id: 41)
"""
# https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356
if (dtype := _GGML_NATIVE.get(ggml_type)) is not None:
return t[:dtype.itemsize * n].contiguous().bitcast(dtype)
def q_to_uint8(t: Tensor, b: int) -> Tensor:
# TODO: rewrite with arange?
shift_tensor, bitmask = Tensor.const(tuple(2**(i*b) for i in range(8//b)), t.dtype), 0xff >> (8 - b)
return t.unsqueeze(-1).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None:
from tinygrad.runtime.autogen import ggml_common as _ggml
blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
if ggml_type == 3:
d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
return q_to_uint8(blocks[:,4:], 4).bitcast(dtypes.int8) * d + m
if ggml_type in (6, 7):
d = blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
qh_off = 2 if ggml_type == 6 else 4
qh = q_to_uint8(blocks[:,qh_off:qh_off+4], 1).reshape((-1, 8, 4)).transpose(-1, -2).flatten(-2).bitcast(dtypes.int8)
q = q_to_uint8(blocks[:,qh_off+4:], 4).bitcast(dtypes.int8) + qh * 16
return q * d + (blocks[:,2:4].bitcast(dtypes.float16).cast(dtypes.float32) if ggml_type == 7 else -16 * d)
if ggml_type == 8: return blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32) * blocks[:,2:].bitcast(dtypes.int8)
# Q4_K: 256 elements per 144-byte block (d:2, dmin:2, scales:12, qs:128)
# Q5_K: 256 elements per 176-byte block (d:2, dmin:2, scales:12, qh:32, qs:128)
if ggml_type in (12, 13):
d, dmin = (blocks[:,i:i+2].bitcast(dtypes.float16).cast(dtypes.float32).unsqueeze(-1) for i in [0, 2])
s = blocks[:,4:16] # 12 bytes: 6-bit scales[0-3], 6-bit mins[0-3], high bits[4-7]
sc = s[:,0:4].bitwise_and(63).cat(s[:,8:12].bitwise_and(0xF).bitwise_or(s[:,0:4].rshift(6).lshift(4)), dim=-1)
mn = s[:,4:8].bitwise_and(63).cat(s[:,8:12].rshift(4).bitwise_or(s[:,4:8].rshift(6).lshift(4)), dim=-1)
qs_off = 48 if ggml_type == 13 else 16
q = Tensor.stack((qs:=blocks[:,qs_off:qs_off+128].reshape(-1,4,32)).bitwise_and(0xF), qs.rshift(4), dim=2).reshape(-1,8,32)
if ggml_type == 13: q = q + q_to_uint8(blocks[:,16:48], 1).reshape(-1, 8, 32) * 16
return (d * sc.unsqueeze(-1) * q - dmin * mn.unsqueeze(-1)).flatten(-2)
if ggml_type == 14:
xl, xh = q_to_uint8(blocks[:,:128].reshape((-1, 2, 64)), 4), q_to_uint8(blocks[:,128:192].reshape((-1, 2, 32)), 2).lshift(4)
scales = blocks[:,192:208].bitcast(dtypes.int8).unsqueeze(-1).expand((-1, 16, 16)).reshape((-1, 256))
d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32)
return d * (xl.bitwise_or(xh).bitcast(dtypes.int8) - 32).flatten(-2) * scales
if ggml_type == 18:
d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
scale_words = blocks[:, 66:98].bitcast(dtypes.uint32)
db = d * (scale_words.rshift(28).cast(dtypes.float32) + 0.5).reshape((-1, 8, 1, 1)) * 0.5
sign_idx = scale_words.unsqueeze(-1).rshift(Tensor.const((0, 7, 14, 21), dtypes.uint32)).bitwise_and(0x7F).reshape((-1, 32)).cast(dtypes.int32)
even_signs = Tensor([i | (0x80 if i.bit_count() % 2 else 0) for i in range(128)], dtype=dtypes.uint8, device=t.device)
signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8))
grid = _ggml_iq_grid(t.device, _ggml.iq3xxs_grid, (256, 4))[blocks[:, 2:66]].reshape((-1, 8, 4, 8))
return (db * grid * signs).flatten(-3)
if ggml_type == 21:
d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
scales = (1 + 2 * q_to_uint8(blocks[:, 106:110].reshape((-1, 4, 1)), 4).reshape((-1, 8))).cast(dtypes.float32).reshape((-1, 8, 1, 1))
qh = q_to_uint8(blocks[:, 66:74].reshape((-1, 8, 1)), 1).reshape((-1, 64)).cast(dtypes.uint16)
signs = (q_to_uint8(blocks[:, 74:106].reshape((-1, 32, 1)), 1).reshape((-1, 256)) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8))
q = blocks[:, 2:66].cast(dtypes.uint16) + qh.lshift(8)
return (d * scales * _ggml_iq_grid(t.device, _ggml.iq3s_grid, (512, 4))[q].reshape((-1, 8, 4, 8)) * signs).flatten(-3)
if ggml_type == 22:
d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
db = d * (q_to_uint8(blocks[:, 74:82].reshape((-1, 8, 1)), 4).reshape((-1, 16)).cast(dtypes.float32) + 0.5).reshape((-1, 16, 1, 1)) * 0.25
signs = (q_to_uint8(blocks[:, 34:66].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 16, 2, 8))
qh = q_to_uint8(blocks[:, 66:74].reshape((-1, 8, 1)), 2).reshape((-1, 32)).cast(dtypes.uint16)
q = blocks[:, 2:34].cast(dtypes.uint16) + qh.lshift(8)
return (db * _ggml_iq_grid(t.device, _ggml.iq2s_grid, (1024, 8))[q].reshape((-1, 16, 2, 8)) * signs).flatten(-3)
if ggml_type == 23:
d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1))
scale_shifts = Tensor.const((0, 2, 4, 6, 8, 10, 12, 14), dtypes.uint16)
iq4_xs_lut = Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)
scales_l = Tensor.stack((sl:=blocks[:, 4:8]).bitwise_and(0xF), sl.rshift(4), dim=2).reshape((-1, 8))
scales_h = blocks[:, 2:4].bitcast(dtypes.uint16).unsqueeze(-1).rshift(scale_shifts).bitwise_and(0x03).reshape((-1, 8)).cast(dtypes.uint8)
scales = (scales_l.bitwise_or(scales_h.lshift(4)).bitcast(dtypes.int8) - 32).cast(dtypes.float32).reshape((-1, 8, 1))
q = (qs:=blocks[:, 8:].reshape((-1, 8, 16))).bitwise_and(0xF).cat(qs.rshift(4), dim=2)
return (d * scales * iq4_xs_lut[q]).flatten(-2)
if ggml_type == 39:
e = blocks[:, 0].cast(dtypes.uint32)
small_bits = Tensor([0x00200000, 0x00400000], dtype=dtypes.uint32, device=t.device)[e.clip(0, 1).cast(dtypes.int32)] # e = 0 or e = 1 case
d = (e < 2).where(small_bits, (e - 1) * 0x00800000).bitcast(dtypes.float32).unsqueeze(-1)
codes = q_to_uint8(blocks[:, 1:17], 4)
fp4_lut = Tensor([0.0, 1.0, 2.0, 3.0, 4.0, 6.0, 8.0, 12.0,
-0.0,-1.0,-2.0,-3.0,-4.0,-6.0,-8.0,-12.0],
dtype=dtypes.float32, device=t.device)
fp4_val = fp4_lut[codes]
return (fp4_val * d).flatten(-2)[:n]
if ggml_type == 41:
d = blocks[:,:2].bitcast(dtypes.float16)
bits = q_to_uint8(blocks[:,2:], 1).reshape(-1, 8, 16).transpose(-1, -2).flatten(-2).bitcast(dtypes.int8)
return d * (bits * 2 - 1)
raise ValueError(f"GGML type '{ggml_type}' is not supported!")
def _read_unpack(fmt: str, n: int, r:io.BufferedIOBase): return struct.unpack(fmt, r.read(n))[0]
def read_str(r:io.BufferedIOBase): return str(r.read(read_uint64(r)), "utf-8")
def read_arr(r:io.BufferedIOBase):
item_reader, n = readers[read_int32(r)], read_uint64(r)
return [item_reader(r) for _ in range(n)]
readers: dict[int, Callable[[io.BufferedIOBase], Any]] = { 8: read_str, 9: read_arr,
**{ t: functools.partial(_read_unpack, "<"+f, nb) for t,f,nb in \
[ (0,"c",1), (1,"b",1), (2,"H",2), (3,"h",2), (4,"I",4), (5,"i",4), (6,"f",4), (7,"?",1), (10,"Q",8), (11,"q",8), (12,"d",8) ] } }
read_uint32, read_int32, read_uint64, read_int64 = readers[4], readers[5], readers[10], readers[11]
def _gguf_parse(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
# TODO: remove the need for copy to default device
tensor = tensor.to(None).realize()
r = io.BufferedReader(TensorIO(tensor), 1_000_000)
magic, version, n_tensors, n_kv = r.read(4), read_int32(r), read_int64(r), read_int64(r)
if magic != b"GGUF" or version not in [2, 3]: raise ValueError("Invalid GGUF format!")
kv_data = {}
for _ in range(n_kv):
k, typ = read_str(r), read_int32(r)
kv_data[k] = readers[typ](r)
t_infos = [ (read_str(r), tuple(read_uint64(r) for _ in range(read_uint32(r))), read_int32(r), read_uint64(r)) for _ in range(n_tensors) ]
alignment, pos = kv_data.get("general.alignment", 32), r.tell()
data_start = round_up(pos, alignment)
state_dict = {name: ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims)) for name, dims, typ, off in t_infos}
return kv_data, state_dict
def _gguf_split_paths(path: pathlib.Path, kv: dict) -> list[pathlib.Path]:
if (total := kv.get('split.count', 1)) <= 1: return [path]
if kv.get('split.no', 0) != 0: raise ValueError(f"multi-part GGUF must be loaded from the first split, got split.no={kv['split.no']}")
if not (m := re.match(r"^(.*)-00001-of-\d{5}\.gguf$", str(path))): raise ValueError(f"first split path must end with -00001-of-NNNNN.gguf: {path}")
return [pathlib.Path(f"{m.group(1)}-{i:05d}-of-{total:05d}.gguf") for i in range(1, total+1)]
def gguf_load(fn: Tensor|str|pathlib.Path) -> tuple[dict, dict[str, Tensor]]:
"""
Loads a .gguf file, returning the `kv_data` and `state_dict`. Multi-part splits are auto-merged when loaded by path.
```python
import pathlib
from tinygrad import Device, Tensor
from tinygrad.llm.gguf import gguf_load
gguf_tensor = Tensor(pathlib.Path("Meta-Llama-3-8B-Instruct.Q4_0.gguf")).to(Device.DEFAULT)
kv_data, state_dict = gguf_load(gguf_tensor)
```
NOTE: The provided tensor must be on a device that supports execution.
"""
kv, sd = _gguf_parse(fn if isinstance(fn, Tensor) else Tensor(pathlib.Path(fn)))
if kv.get('split.count', 1) <= 1: return kv, sd
if isinstance(fn, Tensor): raise ValueError("multi-part GGUF requires a path argument (got Tensor)")
for pp in _gguf_split_paths(pathlib.Path(fn), kv)[1:]: sd.update(_gguf_parse(Tensor(pp))[1])
return kv, sd

447
artifacts/package_runtime/tinygrad/llm/model.py Normal file
View File

@@ -0,0 +1,447 @@
from __future__ import annotations
import functools, itertools, pathlib
from dataclasses import dataclass, replace
from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
from tinygrad.nn import Linear
from tinygrad.llm.gguf import gguf_load
from tinygrad.uop.ops import resolve
@functools.cache
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor:
freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
return freqs.cos().cat(freqs.sin(), dim=-1).clone(device)
class ExpertWeights:
"""Like Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features)."""
def __init__(self, num_experts:int, in_features:int, out_features:int):
self.weight = Tensor.zeros(num_experts, out_features, in_features)
def __call__(self, sel:Tensor, x:Tensor) -> Tensor:
# sel: (B, T, k), x: (B, T, 1, in) or (B, T, k, in) -> output: (B, T, k, out)
return (x.unsqueeze(-2) @ self.weight[sel].transpose(-1, -2)).contiguous().squeeze(-2)
def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
assert x.shape[-1] % 2 == 0
cos, sin = freqs_cis.reshape(1, 1, x.shape[2], -1).chunk(2, dim=-1)
x1, x2 = x.chunk(2, dim=-1)
return (x1 * cos - x2 * sin).cat(x2 * cos + x1 * sin, dim=-1)
def pairwise_topk(x: Tensor, k: int) -> tuple[Tensor, Tensor]:
n = x.shape[-1]
vals = Tensor.arange(n).reshape(1,1,n).cast(x.dtype).expand(x.shape)
cmp = (x.unsqueeze(-1) > x.unsqueeze(-2)) | ((x.unsqueeze(-1) == x.unsqueeze(-2)) & \
(Tensor.arange(n).reshape(1,1,n,1) < Tensor.arange(n).reshape(1,1,1,n)))
sel = x.const_like(0).scatter(-1, cmp.sum(axis=-1).cast('int32'), vals)[:,:,n-k:].cast('int32')
return x.gather(-1, sel), sel
@dataclass(frozen=True)
class SSMConfig:
conv_kernel: int
state_size: int
group_count: int
time_step_rank: int
inner_size: int
kda: bool = False
@dataclass(frozen=True)
class TransformerConfig:
num_blocks: int
dim: int
hidden_dim: int
n_heads: int
n_kv_heads: int
norm_eps: float
vocab_size: int
head_dim: int
rope_theta: float
rope_dim: int
v_head_dim: int
max_context: int = 0
qk_norm: int = 0
num_experts: int = 0
num_experts_per_tok: int = 0
norm_topk_prob: bool = False
q_lora_rank: int = 0
kv_lora_rank: int = 0
shared_expert_dim: int = 0
ssm_layers: tuple[bool, ...] = ()
attn_output_gate: bool = False
ssm: SSMConfig|None = None
shared_expert_gate: bool = True
leading_dense_blocks: int = 0
dense_hidden_dim: int = 0
routed_scaling_factor: float = 1.0
qkv_bias: bool = False
expert_bias: bool = False
class FFNBlock:
def __init__(self, config:TransformerConfig):
self.config = config
# --- RMSNorms --------------------------------------------------------
self.attn_norm = nn.RMSNorm(config.dim, config.norm_eps)
self.ffn_norm = nn.RMSNorm(config.dim, config.norm_eps)
# --- feed-forward (MoE or dense) -------------------------------------
if config.num_experts > 0:
self.ffn_gate_inp = Linear(config.dim, config.num_experts, bias=False) # router
if config.expert_bias: self.exp_probs_b = {"bias": Tensor.zeros(config.num_experts)}
self.ffn_gate_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim)
self.ffn_up_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim)
self.ffn_down_exps = ExpertWeights(config.num_experts, config.hidden_dim, config.dim)
if config.shared_expert_dim > 0:
self.ffn_gate_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
self.ffn_up_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
self.ffn_down_shexp = Linear(config.shared_expert_dim, config.dim, bias=False)
if config.shared_expert_gate: self.ffn_gate_inp_shexp = {"weight": Tensor.zeros(config.dim)}
else:
self.ffn_gate = Linear(config.dim, config.hidden_dim, bias=False)
self.ffn_up = Linear(config.dim, config.hidden_dim, bias=False)
self.ffn_down = Linear(config.hidden_dim, config.dim, bias=False)
def _feed_forward(self, x:Tensor) -> Tensor:
if hasattr(self, 'ffn_gate_exps'):
h = x.unsqueeze(2) # (B, T, 1, D) - add expert dim for broadcasting
logits = self.ffn_gate_inp(x)
if hasattr(self, 'exp_probs_b'):
probs = logits.sigmoid()
_, sel = pairwise_topk(probs + self.exp_probs_b["bias"], self.config.num_experts_per_tok)
probs = probs.gather(-1, sel)
if self.config.norm_topk_prob: probs = probs / probs.sum(axis=-1, keepdim=True)
else:
vals, sel = pairwise_topk(logits, self.config.num_experts_per_tok)
probs = vals.softmax(-1) if self.config.norm_topk_prob else logits.softmax(-1).gather(-1, sel)
probs = probs * self.config.routed_scaling_factor
x_down = self.ffn_down_exps(sel, (self.ffn_gate_exps(sel, h).silu() * self.ffn_up_exps(sel, h)).contiguous()) # (B, T, k, D)
out = (x_down * probs.unsqueeze(-1)).sum(axis=2) # (B, T, D)
if hasattr(self, 'ffn_gate_shexp'):
shexp = self.ffn_down_shexp(self.ffn_gate_shexp(x).silu().contiguous() * self.ffn_up_shexp(x))
if hasattr(self, 'ffn_gate_inp_shexp'): shexp = shexp * (x * self.ffn_gate_inp_shexp["weight"]).sum(axis=-1, keepdim=True).sigmoid()
out = out + shexp
return out
# TODO: remove the need for this contiguous
return self.ffn_down(self.ffn_gate(x).silu().contiguous() * self.ffn_up(x))
# given the token-prefix match, return how much cached state this block can still reuse
def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return prefix_len
# return writes that reset this block's state after a cache mismatch
def _state_reset_ops(self) -> list[Tensor]: return []
def _init_state(self, x:Tensor): raise NotImplementedError
def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: raise NotImplementedError
def __call__(self, x: Tensor, start_pos: int|UOp):
self._init_state(x)
# we pass in the weights implicitly so we unpack the GGUF on the fly
@function(precompile=True, allow_implicit=True)
def _run(x:Tensor, start_pos:int|UOp):
h = x + self._attention(self.attn_norm(x), start_pos)
return (h + self._feed_forward(self.ffn_norm(h))).contiguous()
return _run(x, start_pos)
class TransformerBlock(FFNBlock):
def __init__(self, config:TransformerConfig):
super().__init__(config)
assert config.v_head_dim == config.head_dim, "TransformerBlock requires v_head_dim == head_dim"
# --- attention projections (all linear, bias-free) ------------------
q_proj_out = config.head_dim * config.n_heads * (2 if config.attn_output_gate else 1)
kv_proj_out = config.head_dim * config.n_kv_heads
self.attn_q = Linear(config.dim, q_proj_out, bias=config.qkv_bias)
self.attn_k = Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
self.attn_v = Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
self.attn_output = Linear(config.head_dim * config.n_heads, config.dim, bias=False)
if config.qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(config.qk_norm, config.norm_eps), nn.RMSNorm(config.qk_norm, config.norm_eps)
def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
q, k, v = self.attn_q(x), self.attn_k(x), self.attn_v(x)
if self.config.qk_norm and self.config.qk_norm != self.config.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
B, T, _ = x.shape
if self.config.attn_output_gate:
qg = q.reshape(B, T, self.config.n_heads, 2, self.config.head_dim)
q, gate = qg[:, :, :, 0, :], qg[:, :, :, 1, :].reshape(B, T, self.config.n_heads * self.config.head_dim)
q = q.reshape(B, T, self.config.n_heads, self.config.head_dim).transpose(1, 2) # (B,H,T,Hd)
k = k.reshape(B, T, self.config.n_kv_heads, self.config.head_dim).transpose(1, 2) # (B,KvH,T,Hd)
v = v.reshape(B, T, self.config.n_kv_heads, self.config.head_dim).transpose(1, 2) # (B,KvH,T,Hd)
if self.config.qk_norm == self.config.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
q = apply_rope(q[..., :self.config.rope_dim], self.freqs_cis[start_pos:start_pos+T]).cat(q[..., self.config.rope_dim:], dim=-1)
k = apply_rope(k[..., :self.config.rope_dim], self.freqs_cis[start_pos:start_pos+T]).cat(k[..., self.config.rope_dim:], dim=-1)
# NOTE: we don't want to change self.cache_kv, the function API doesn't support this well
assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop)))
k = assigned_kv[0, :, :, 0:start_pos+T, :]
v = assigned_kv[1, :, :, 0:start_pos+T, :]
#self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
#k = self.cache_kv[0, :, :, 0:start_pos+T, :]
#v = self.cache_kv[1, :, :, 0:start_pos+T, :]
# NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
# TODO: this if statement should be removed and it shouldn't generate extra kernels
mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
if resolve(T != 1) else None
attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True) # (B,H,T,Hd)
attn = attn.transpose(1, 2).reshape(B, T, -1) # back to (B,T,D)
return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid()))
def _init_state(self, x:Tensor):
if not hasattr(self, "cache_kv"):
# TODO: how is the dtype of this determined?
self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, device=x.device)
self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
class MLATransformerBlock(FFNBlock):
def __init__(self, config:TransformerConfig):
super().__init__(config)
qk_nope_head_dim = config.head_dim - config.rope_dim
if config.q_lora_rank > 0:
self.attn_q_a = Linear(config.dim, config.q_lora_rank, bias=False)
self.attn_q_a_norm = nn.RMSNorm(config.q_lora_rank, config.norm_eps)
self.attn_q_b = Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False)
else:
self.attn_q = Linear(config.dim, config.n_heads * config.head_dim, bias=False)
self.attn_kv_a_mqa = Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False)
self.attn_kv_a_norm = nn.RMSNorm(config.kv_lora_rank, config.norm_eps)
self.attn_k_b = {"weight": Tensor.zeros(config.n_heads, config.kv_lora_rank, qk_nope_head_dim)}
self.attn_v_b = {"weight": Tensor.zeros(config.n_heads, config.v_head_dim, config.kv_lora_rank)}
self.attn_output = Linear(config.n_heads * config.v_head_dim, config.dim, bias=False)
def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
B, T, _ = x.shape
q_nope_head_dim = self.config.head_dim - self.config.rope_dim
q_proj = self.attn_q_b(self.attn_q_a_norm(self.attn_q_a(x))) if self.config.q_lora_rank > 0 else self.attn_q(x)
q = q_proj.reshape(B, T, self.config.n_heads, self.config.head_dim).transpose(1, 2)
q_nope, q_rope = q[..., :q_nope_head_dim], q[..., q_nope_head_dim:]
if not self.config.ssm or not self.config.ssm.kda: q_rope = apply_rope(q_rope, self.freqs_cis[start_pos:start_pos+T])
q = (q_nope @ self.attn_k_b["weight"].transpose(-1, -2)).cat(q_rope, dim=-1)
kv_a = self.attn_kv_a_mqa(x)
c_kv = self.attn_kv_a_norm(kv_a[..., :self.config.kv_lora_rank])
k_rope = kv_a[..., self.config.kv_lora_rank:].reshape(B, T, 1, self.config.rope_dim).transpose(1, 2)
if not self.config.ssm or not self.config.ssm.kda: k_rope = apply_rope(k_rope, self.freqs_cis[start_pos:start_pos+T])
k_store = c_kv.reshape(B, 1, T, self.config.kv_lora_rank).cat(k_rope.reshape(B, 1, T, self.config.rope_dim), dim=-1)
k = Tensor(self.cache_k.uop.after(self.cache_k[:, :, start_pos:start_pos+T, :].uop.store(k_store.uop)))[:, :, 0:start_pos+T, :]
v = k[..., :self.config.kv_lora_rank]
mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
if resolve(T != 1) else None
attn = q @ k.transpose(-1, -2) * (1.0 / self.config.head_dim ** 0.5)
if mask is not None: attn = attn + mask
attn = attn.softmax(-1)
attn = ((attn @ v) @ self.attn_v_b["weight"].transpose(-1, -2)).transpose(1, 2).reshape(B, T, -1)
return self.attn_output(attn)
def _init_state(self, x:Tensor):
if not hasattr(self, "cache_k"):
self.cache_k = Tensor.empty(x.shape[0], 1, self.config.max_context, self.config.kv_lora_rank + self.config.rope_dim, device=x.device)
self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
class GatedDeltaNetBlock(FFNBlock):
def __init__(self, config:TransformerConfig, ssm:SSMConfig):
super().__init__(config)
self.head_k_dim, self.num_k_heads, self.num_v_heads = ssm.state_size, ssm.group_count, ssm.time_step_rank
assert self.num_v_heads % self.num_k_heads == 0
self.head_v_dim, self.ssm_conv_kernel = ssm.inner_size // ssm.time_step_rank, ssm.conv_kernel
self.conv_channels, self.q_dim = ssm.inner_size + 2*ssm.group_count*ssm.state_size, ssm.state_size*ssm.group_count
self.attn_qkv = Linear(config.dim, self.conv_channels, bias=False)
if ssm.kda:
self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False)
self.ssm_f_a, self.ssm_f_b = Linear(config.dim, self.head_k_dim, bias=False), Linear(self.head_k_dim, ssm.inner_size, bias=False)
else:
self.attn_gate = Linear(config.dim, ssm.inner_size, bias=False)
self.ssm_alpha = Linear(config.dim, self.num_v_heads, bias=False)
self.ssm_beta = Linear(config.dim, self.num_v_heads, bias=False)
self.ssm_conv1d = {"weight": Tensor.zeros(self.conv_channels, self.ssm_conv_kernel)}
self.ssm_dt = {"bias": Tensor.zeros(ssm.inner_size if ssm.kda else self.num_v_heads)}
self.ssm_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads)
self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), Linear(ssm.inner_size, config.dim, bias=False)
def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
B, T, _ = x.shape
assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
# input processing
x = x.half()
out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
# qkv conv
conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
conv_out = (conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1).silu()
q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)
q = q.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
k = k.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
v = v.reshape(B, self.num_v_heads, self.head_v_dim)
q, k, v = q.mul(self.head_k_dim**-0.5).unsqueeze(-1), k.unsqueeze(-1), v.unsqueeze(-1)
# recurrent
recurrent_state = self.recurrent_state * alpha
recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2)
# store the updated state
conv_state_store = self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop)
recurrent_state_store = self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop)
recurrent_state = Tensor(self.recurrent_state.uop.after(recurrent_state_store, conv_state_store))
# output
core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
# recurrent state can't be partially reused after divergence, force a full rebuild
def _state_reset_ops(self):
return [self.conv_state.assign(self.conv_state.const_like(0)),
self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return 0 if prefix_len != cached_len else prefix_len
def _init_state(self, x):
if not hasattr(self, "conv_state"):
self.conv_state = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device).clone()
self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).clone()
class Transformer:
def __init__(self, config:TransformerConfig):
dense_config = replace(config, num_experts=0, num_experts_per_tok=0, shared_expert_dim=0, hidden_dim=config.dense_hidden_dim or config.hidden_dim)
if config.ssm: config = replace(config, qk_norm=config.head_dim)
block_cls = MLATransformerBlock if config.kv_lora_rank > 0 else TransformerBlock
self.blk:list[FFNBlock] = [GatedDeltaNetBlock(dense_config if i < config.leading_dense_blocks else config, config.ssm)
if config.ssm and config.ssm_layers[i] else
block_cls(dense_config if i < config.leading_dense_blocks else config) for i in range(config.num_blocks)]
self.token_embd = nn.Embedding(config.vocab_size, config.dim)
self.output_norm = nn.RMSNorm(config.dim, config.norm_eps)
self.output = Linear(config.dim, config.vocab_size, bias=False)
self.max_context = config.max_context
self.has_recurrent_block = any(isinstance(b, GatedDeltaNetBlock) for b in self.blk)
self._cached_tokens: list[int] = []
# we specialize the JIT for prefill and rollout
self.prefill_jit = TinyJit(self.forward)
self.rollout_jit = TinyJit(self.forward)
def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
x = self.token_embd(tokens).float() # (B, T, D)
for block in self.blk: x = block(x, start_pos)
logits = self.output(self.output_norm(x))[:, -1, :]
# Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)
def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
return (self.prefill_jit if resolve(tokens.shape[1] != 1) else self.rollout_jit)(tokens.contiguous(), start_pos, temperature)
@staticmethod
def from_gguf(gguf:Tensor|str|pathlib.Path, max_context:int|None=None,
realize=bool(getenv("REALIZE", 0))) -> tuple[Transformer, dict]:
# TODO: remove the need for copy to default device
kv, state_dict = gguf_load(gguf.to(None).realize() if isinstance(gguf, Tensor) else gguf)
# all state items should be float16, not float32
state_dict = {k:v.cast('float16') if getenv("HALF", 1) else v for k,v in state_dict.items()}
# some models like Llama 3.2 don't have an output.weight, they just tie to the token_embd.weight
if 'output.weight' not in state_dict: state_dict['output.weight'] = state_dict['token_embd.weight']
arch = kv['general.architecture']
max_context = min(max_context, kv[f'{arch}.context_length']) if max_context is not None else kv[f'{arch}.context_length']
n_heads, n_kv_heads = kv[f'{arch}.attention.head_count'], kv[f'{arch}.attention.head_count_kv']
ssm = None
ssm_layers: tuple[bool, ...] = ()
if arch in ('qwen35', 'qwen35moe'):
ssm = SSMConfig(**{k: kv[f'{arch}.ssm.{k}'] for k in ('conv_kernel','state_size','group_count','time_step_rank','inner_size')})
ssm_layers = tuple((i+1) % kv[f'{arch}.full_attention_interval'] != 0 for i in range(kv[f'{arch}.block_count']))
elif arch == 'kimi-linear':
ssm_layers = tuple(x == 0 for x in n_kv_heads)
n_kv_heads = max(n_kv_heads)
ssm = SSMConfig(kv[f'{arch}.ssm.conv_kernel'], kv[f'{arch}.kda.head_dim'], n_heads, n_heads, n_heads*kv[f'{arch}.kda.head_dim'], kda=True)
for i, is_ssm in enumerate(ssm_layers):
if not is_ssm: continue
state_dict[f"blk.{i}.attn_qkv.weight"] = state_dict.pop(f"blk.{i}.attn_q.weight").cat(
state_dict.pop(f"blk.{i}.attn_k.weight"), state_dict.pop(f"blk.{i}.attn_v.weight"), dim=0).contiguous()
state_dict[f"blk.{i}.ssm_conv1d.weight"] = state_dict.pop(f"blk.{i}.ssm_conv1d_q.weight").cat(
state_dict.pop(f"blk.{i}.ssm_conv1d_k.weight"), state_dict.pop(f"blk.{i}.ssm_conv1d_v.weight"), dim=0).squeeze(1).contiguous()
state_dict[f"blk.{i}.ssm_out.weight"] = state_dict.pop(f"blk.{i}.attn_output.weight")
if arch in ('qwen35', 'qwen35moe', 'glm4moe'):
state_dict = {k.replace('post_attention_norm', 'ffn_norm'):v for k,v in state_dict.items()}
kv_lora_rank = kv.get(f'{arch}.attention.kv_lora_rank', 0)
head_dim = kv.get(f'{arch}.attention.key_length_mla', kv.get(f'{arch}.attention.key_length', kv[f'{arch}.embedding_length'] // n_heads))
rope_dim = kv.get(f'{arch}.rope.dimension_count', head_dim)
# Permute RoPE weights from interleaved to half-split layout.
for name in state_dict:
if arch == 'kimi-linear': continue
if ('attn_q.weight' in name or 'attn_q_b.weight' in name) and (arch == 'llama' or kv_lora_rank):
w = state_dict[name].reshape(n_heads, state_dict[name].shape[0]//n_heads, -1)
prefix = head_dim-rope_dim
state_dict[name] = w[:, :prefix].cat(w[:, prefix:].rearrange("n (h two) d -> n (two h) d", two=2), dim=1).reshape(-1, w.shape[-1])
elif arch == 'llama' and 'attn_k.weight' in name:
w = state_dict[name].reshape(n_kv_heads, state_dict[name].shape[0]//n_kv_heads, -1)
state_dict[name] = w.rearrange("n (h two) d -> n (two h) d", two=2).reshape(-1, w.shape[-1])
elif kv_lora_rank and 'attn_kv_a_mqa.weight' in name:
state_dict[name] = state_dict[name][:kv_lora_rank].cat(state_dict[name][kv_lora_rank:].rearrange("(h two) d -> (two h) d", two=2), dim=0)
config = TransformerConfig(
num_blocks=kv[f'{arch}.block_count'] - kv.get(f'{arch}.nextn_predict_layers', 0), dim=kv[f'{arch}.embedding_length'],
hidden_dim=kv.get(f'{arch}.expert_feed_forward_length', kv.get(f'{arch}.feed_forward_length', 0)),
n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'],
vocab_size=len(kv['tokenizer.ggml.tokens']),
head_dim=head_dim,
rope_theta=kv[f'{arch}.rope.freq_base'],
rope_dim=rope_dim,
v_head_dim=kv.get(f'{arch}.attention.value_length_mla', kv.get(f'{arch}.attention.value_length', head_dim)),
max_context=max_context,
qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0,
num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0),
norm_topk_prob=kv.get(f'{arch}.expert_weights_norm', arch in ('qwen3moe', 'qwen35moe', 'kimi-linear')),
kv_lora_rank=kv_lora_rank, q_lora_rank=kv.get(f'{arch}.attention.q_lora_rank', 0),
leading_dense_blocks=kv.get(f'{arch}.leading_dense_block_count', 0),
shared_expert_dim=kv.get(
f'{arch}.expert_shared_feed_forward_length',
kv.get(f'{arch}.expert_shared_count', 0) * kv.get(f'{arch}.expert_feed_forward_length', 0)),
shared_expert_gate=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.ffn_gate_inp_shexp.weight" in state_dict,
dense_hidden_dim=kv.get(f'{arch}.feed_forward_length', 0) if kv.get(f'{arch}.leading_dense_block_count', 0) else 0,
routed_scaling_factor=kv.get(f'{arch}.expert_weights_scale', 1.0), attn_output_gate=arch in ('qwen35', 'qwen35moe'), ssm=ssm,
ssm_layers=ssm_layers,
qkv_bias='blk.0.attn_q.bias' in state_dict,
expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict)
model = Transformer(config)
nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False) # NOTE: rope_freqs.weight (32,) is unused
# NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
if realize:
for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
Tensor.realize(*params)
return model, kv
def warmup(self):
for _ in range(2): list(zip(range(2), self.generate([0])))
def get_start_pos(self, tokens:list[int]) -> int:
prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
if self.has_recurrent_block: chunk_size = 1
v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
v_toks = UOp.variable("toks", 1, chunk_size)
# TODO: use UOp.variable for temperature once float variables are supported
temp = Tensor([temperature])
# assign all input tokens once, then slice from start_pos for the model call
t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32").reshape(1, self.max_context)
# recompute start_pos from what's currently valid in the caches
start_pos = self.get_start_pos(tokens)
if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets)
out, prompt_len = None, len(tokens)
while len(tokens) < self.max_context:
n_toks = min(chunk_size, len(tokens) - start_pos)
sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
start_pos += n_toks
# chunked prefill: keep processing until all prompt tokens are consumed
if start_pos < len(tokens): continue
tokens.append(int(out.item()))
self._cached_tokens = tokens[:-1]
yield tokens[-1]

167
artifacts/package_runtime/tinygrad/llm/serve.py Normal file
View File

@@ -0,0 +1,167 @@
from __future__ import annotations
import json, pathlib, re, time, typing, uuid
from typing import TYPE_CHECKING
from tinygrad.helpers import DEBUG, colored, stderr_log
from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
if TYPE_CHECKING:
from tinygrad.llm.cli import SimpleTokenizer
from tinygrad.llm.model import Transformer
def parse_tool_call(s:str) -> tuple[str, typing.Any]|None:
s = s.strip()
if s.startswith("{"): # hermes JSON format: {"name": ..., "arguments": {...}}
try:
call = json.loads(s)
return call["name"], call.get("arguments", call.get("parameters", {}))
except (json.JSONDecodeError, KeyError): return None
# XML format: <function=name>\n<parameter=key>\nvalue\n</parameter>...</function>
if (fm := re.match(r"<function=([^>]+)>\s*(.*?)\s*(?:</function>)?$", s, re.DOTALL)):
args = {}
for pm in re.finditer(r"<parameter=([^>]+)>(.*?)</parameter>", fm.group(2), re.DOTALL):
value = re.sub(r"^\r?\n|\r?\n\Z", "", pm.group(2))
try: args[pm.group(1)] = json.loads(value)
except json.JSONDecodeError: args[pm.group(1)] = value
return fm.group(1), args
return None
def normalize_messages(messages:list[dict]) -> None:
# chat templates expect tool_call arguments as dicts (OpenAI clients send JSON strings)
for m in messages:
for tc in m.get("tool_calls") or []:
if "function" in tc and isinstance(args := tc["function"].get("arguments"), str):
try: tc["function"]["arguments"] = json.loads(args)
except json.JSONDecodeError: pass
class StreamRouter:
# routes streamed output text to (field, text) deltas, keeping tool_call regions in .buf for the final parse
def __init__(self, reasoning:bool=False):
self.buf = ""
self.mode = "reasoning" if reasoning else "undecided" # output inside a think block is sent as reasoning_content
def split(self, tag:str, final:bool) -> tuple[str, bool]:
# split buf on the first full tag, holding back a partial tag at the end unless final
if tag in self.buf:
before, self.buf = self.buf.split(tag, 1)
return before, True
hold = max((i for i in range(1, min(len(self.buf), len(tag))+1) if tag.startswith(self.buf[-i:])), default=0) if not final else 0
emit, self.buf = self.buf[:len(self.buf)-hold], self.buf[len(self.buf)-hold:]
return emit, False
def route(self, piece:str, final:bool=False) -> typing.Iterator[tuple[str, str]]:
self.buf += piece
if self.mode == "undecided": # decide whether the output starts with a think block
if not final and len(self.buf) < len("<think>") and "<think>".startswith(self.buf): return
self.mode, self.buf = ("reasoning", self.buf[len("<think>"):]) if self.buf.startswith("<think>") else ("content", self.buf)
if self.mode == "reasoning":
emit, done = self.split("</think>", final)
if emit: yield "reasoning_content", emit
if not done: return
self.mode = "content"
if self.mode == "tool": return
emit, found = self.split("<tool_call>", final)
if emit: yield "content", emit
if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
class Handler(HTTPRequestHandler):
server: LLMServer
def log_request(self, code='-', size='-'): pass
def do_GET(self):
if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":self.server.model_name,"object":"model"}]}).encode())
else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html")
def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0,
reasoning:bool=False):
model, tok = self.server.model, self.server.tok
prompt_tokens = len(ids)
cache_start_pos = model.get_start_pos(ids)
stderr_log(f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d} {colored('--', 'BLACK')} ")
tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
def chunk(d:dict): return {"choices": [{"index":0, "delta":d, "finish_reason":None}], **tmpl}
out: list[int] = []
finish_reason = "stop"
st = pt = time.perf_counter()
dec = tok.stream_decoder()
router = StreamRouter(reasoning)
def log_stats(interrupted:bool=False):
et = time.perf_counter()
total = f"total:{et-st:6.2f}s"
stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s {colored('--', 'BLACK')} "
f"out:{len(out):5d} {colored('--', 'BLACK')} {colored(total, 'red') if interrupted else total}\n")
completed = False
try:
yield chunk({"role":"assistant", "content":""})
for next_id in model.generate(ids, temperature=temperature):
if len(out) == 0:
stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s {colored('--', 'BLACK')} ")
if tok.is_end(next_id): break
out.append(next_id)
for field, delta in router.route(dec(next_id)): yield chunk({field:delta})
if max_tokens is not None and len(out) >= max_tokens:
finish_reason = "length"
break
for field, delta in router.route(dec(), final=True): yield chunk({field:delta})
tool_calls: list[dict] = []
for m in re.finditer(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", router.buf, re.DOTALL):
if (parsed := parse_tool_call(m.group(1))) is None:
stderr_log(f"failed to parse tool call: {m.group(1)[:200]}")
yield chunk({"content":m.group(0)}) # don't silently drop output the client can't use
else:
name, args = parsed
tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
"function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}})
if tool_calls:
yield chunk({"tool_calls":tool_calls})
if finish_reason == "stop": finish_reason = "tool_calls"
completed = True
yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
if include_usage:
yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out),
"total_tokens": prompt_tokens + len(out)}, **tmpl}
log_stats()
except GeneratorExit:
if not completed: log_stats(interrupted=True)
raise
def do_POST(self):
request_st = time.perf_counter()
stderr_log(f"{self.path} {colored('--', 'BLACK')} ")
raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
if DEBUG >= 1: print(json.dumps(body, indent=2))
if self.path == "/v1/chat/completions":
# render and tokenize
normalize_messages(body["messages"])
rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True, preserve_thinking=True)
ids: list[int] = self.server.tok.encode(rendered)
stderr_log(f"prep:{(time.perf_counter()-request_st)*1e3:5.0f} ms {colored('--', 'BLACK')} ")
if len(ids) >= self.server.model.max_context:
stderr_log(f"{colored('context length exceeded', 'red')} in:{len(ids):5d} max:{self.server.model.max_context:5d}\n")
return self.send_data(json.dumps({"error":{"message":f"prompt has {len(ids)} tokens, but the model context is "
f"{self.server.model.max_context}", "type":"invalid_request_error", "param":"messages", "code":"context_length_exceeded"}}).encode(),
status_code=400)
# reply
max_tokens = body.get("max_completion_tokens") or body.get("max_tokens")
chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False),
max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)),
reasoning=rendered.rstrip().endswith("<think>"))
if body.get("stream"): self.stream_json(chunks)
else:
out, reasoning, tool_calls, finish_reason = [], [], [], "stop"
for c in chunks:
if not c["choices"]: continue
choice = c["choices"][0]
if (delta := choice.get("delta", {})):
if delta.get("content"): out.append(delta["content"])
if delta.get("reasoning_content"): reasoning.append(delta["reasoning_content"])
tool_calls += [{k:v for k, v in tc.items() if k != "index"} for tc in delta.get("tool_calls", [])]
if choice.get("finish_reason"): finish_reason = choice["finish_reason"]
message: dict[str, typing.Any] = {"role":"assistant", "content":"".join(out) or None}
if reasoning: message["reasoning_content"] = "".join(reasoning)
if tool_calls: message["tool_calls"] = tool_calls
self.send_data(json.dumps({**c, "object":"chat.completion",
"choices":[{"index":0, "message":message, "finish_reason":finish_reason}]}).encode())
else:
raise RuntimeError(f"unhandled path {self.path}")
class LLMServer(TCPServerWithReuse):
def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer, template:typing.Any):
self.model, self.model_name, self.tok, self.template = model, model_name, tok, template
super().__init__(server_address, Handler)