J3ff/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Prebuilt Release @ ab07000

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:42 -05:00
commit 9f9c9a70cc
3729 changed files with 778697 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tinygrad_repo/tinygrad/llm/__init__.py Normal file
View File

2
tinygrad_repo/tinygrad/llm/__main__.py Normal file
View File

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

38
tinygrad_repo/tinygrad/llm/chat.html Normal file
View File

@@ -0,0 +1,38 @@
<!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 = '';
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 { d.textContent += JSON.parse(ln.slice(6)).choices[0]?.delta?.content || '' } catch {}
chat.scrollTop = chat.scrollHeight;
}
msgs.push({role: 'assistant', content: d.textContent});
}
</script></body></html>

235
tinygrad_repo/tinygrad/llm/cli.py Normal file
View File

@@ -0,0 +1,235 @@
from __future__ import annotations
import sys, argparse, codecs, typing, re, unicodedata, json, uuid, time, pathlib
from tinygrad import nn
from tinygrad.uop.ops import UOp, Ops
from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored, Context, fetch, profile_marker
from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
from tinygrad.llm.model import Transformer
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)
def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
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)
return SimpleTokenizer(dict(normal_tokens), dict(special_tokens), 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'))
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 role(self, role:str):
if self.preset == 'olmo': return self.encode("<|" + role + "|>\n") # OLMoE Instruct format
if self.preset == 'kimi-k2': return self.encode("<|im_" + role + "|>" + role + "<|im_middle|>")
if self.preset == 'qwen2': return self.encode("<|im_start|>" + role + "\n")
if self.preset == 'glm4': return self.encode("<|" + role + "|>")
if self.preset == 'tekken':
if role == 'user': return self.encode("[INST]")
if role == 'assistant': return []
raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.preset}'")
return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
def end_turn(self):
if self.preset == 'olmo': return self.encode("\n")
if self.preset == 'kimi-k2': return [self.eos_id]
if self.preset == 'qwen2': return [self.eos_id] + self.encode("\n")
if self.preset == 'glm4': return []
if self.preset == 'tekken': return self.encode("[/INST]")
return [self.eos_id]
def prefix(self) -> list[int]:
return ([] if self.bos_id is None else [self.bos_id]) + (self.encode("<sop>") if self.preset == 'glm4' else [])
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.5:27b": "https://huggingface.co/unsloth/Qwen3.5-27B-GGUF/resolve/main/Qwen3.5-27B-Q4_K_M.gguf",
"qwen3.5:35b-a3b": "https://huggingface.co/unsloth/Qwen3.5-35B-A3B-GGUF/resolve/main/Qwen3.5-35B-A3B-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",
}
# *** simple OpenAI API compatible server with web interface on http://localhost:8000/ ***
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):
model, tok = self.server.model, self.server.tok
cache_start_pos = model.get_start_pos(ids)
stderr_log(f"{self.path} {colored('--', 'BLACK')} "
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}
yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
out: list[int] = []
finish_reason = "stop"
st = time.perf_counter()
dec = tok.stream_decoder()
for next_id in model.generate(ids, temperature=temperature):
if len(out) == 0: stderr_log(f"prefill:{(len(ids)-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)
yield {"choices": [{"index":0, "delta":{"content":dec(next_id)}, "finish_reason":None}], **tmpl}
if max_tokens is not None and len(out) >= max_tokens:
finish_reason = "length"
break
if (tail := dec()): yield {"choices": [{"index":0, "delta":{"content":tail}, "finish_reason":None}], **tmpl}
yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
if include_usage:
yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
et = time.perf_counter()
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')} total:{et-st:6.2f}s\n")
def do_POST(self):
tok = self.server.tok
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":
# extract tokens, last assistant message is treated as prefill
ids: list[int] = tok.prefix()
for i, msg in enumerate(body["messages"]):
ids += tok.role(msg["role"])
content = msg["content"]
if isinstance(content, str): ids += tok.encode(content)
elif isinstance(content, list):
for c in content:
if c["type"] == "text": ids += tok.encode(c["text"])
else: raise RuntimeError(f"unhandled type: {c['type']}")
else: raise RuntimeError(f"unknown content type: {type(content)}")
if msg["role"] == "assistant" and i == len(body["messages"]) - 1: break
ids += tok.end_turn()
else: ids += tok.role("assistant")
# 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)))
if body.get("stream"): self.stream_json(chunks)
else:
out, finish_reason = [], "stop"
for c in chunks:
if c["choices"] and c["choices"][0].get("delta", {}).get("content"): out.append(c["choices"][0]["delta"]["content"])
if c["choices"] and c["choices"][0].get("finish_reason"): finish_reason = c["choices"][0]["finish_reason"]
self.send_data(json.dumps({**c, "object":"chat.completion",
"choices":[{"index":0, "message":{"role":"assistant","content":"".join(out)}, "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):
self.model, self.model_name, self.tok = model, model_name, tok
super().__init__(server_address, Handler)
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")
# get tokenizer
tok = SimpleTokenizer.from_gguf_kv(kv)
# warmup the JIT
if args.warmup or args.serve:
# run 2 tokens through the model twice to capture the JIT before serving
with Context(DEBUG=max(DEBUG.value, 1)):
for _ in range(2): list(zip(range(2), model.generate([0])))
# start server
if args.serve: LLMServer(('', args.serve), model, model_name, tok).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()
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")): next(gen)
exit(0)
# interactive chat
ids: list[int] = tok.prefix()
while 1:
try:
ids += tok.role("user") + tok.encode(input('>>> ')) + tok.end_turn() + tok.role("assistant")
except EOFError:
break
dec = tok.stream_decoder()
for next_id in model.generate(ids):
sys.stdout.write(dec(next_id) if not tok.is_end(next_id) else dec() + "\n\n")
sys.stdout.flush()
if tok.is_end(next_id): break
if __name__ == "__main__": main()

177
tinygrad_repo/tinygrad/llm/gguf.py Normal file
View File

@@ -0,0 +1,177 @@
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.stack(*[ Tensor(2**(i*b), device=t.device, dtype=t.dtype) for i in range(8//b) ]), 0xff >> (8 - b)
return t.unsqueeze(-1).expand((*t.shape,8//b)).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).expand((-1, 256))
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([0, 7, 14, 21], device=t.device, dtype=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([0, 2, 4, 6, 8, 10, 12, 14], device=t.device, dtype=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).cast(dtypes.uint32)).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

416
tinygrad_repo/tinygrad/llm/model.py Normal file
View File

@@ -0,0 +1,416 @@
from __future__ import annotations
import functools, itertools, pathlib
from dataclasses import dataclass, replace
from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
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, device=device)[:(dim // 2)] / dim))
freqs = Tensor.arange(end, device=device).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
return freqs.cos().cat(freqs.sin(), dim=-1).contiguous()
class ExpertWeights:
"""Like nn.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, device=x.device).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, device=x.device).reshape(1,1,n,1) < Tensor.arange(n, device=x.device).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
@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
full_attention_interval: int = 0
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 = nn.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 = nn.Linear(config.dim, config.shared_expert_dim, bias=False)
self.ffn_up_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False)
self.ffn_down_shexp = nn.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 = nn.Linear(config.dim, config.hidden_dim, bias=False)
self.ffn_up = nn.Linear(config.dim, config.hidden_dim, bias=False)
self.ffn_down = nn.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 = nn.Linear(config.dim, q_proj_out, bias=config.qkv_bias)
self.attn_k = nn.Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
self.attn_v = nn.Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
self.attn_output = nn.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, device=x.device, 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 = nn.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 = nn.Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False)
else:
self.attn_q = nn.Linear(config.dim, config.n_heads * config.head_dim, bias=False)
self.attn_kv_a_mqa = nn.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 = nn.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:]
q = (q_nope @ self.attn_k_b["weight"].transpose(-1, -2)).cat(apply_rope(q_rope, self.freqs_cis[start_pos:start_pos+T]), 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 = apply_rope(
kv_a[..., self.config.kv_lora_rank:].reshape(B, T, 1, self.config.rope_dim).transpose(1, 2),
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, device=x.device, 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, self.attn_gate = nn.Linear(config.dim, self.conv_channels, bias=False), nn.Linear(config.dim, ssm.inner_size, bias=False)
self.ssm_alpha, self.ssm_beta = nn.Linear(config.dim, self.num_v_heads, bias=False), nn.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(self.num_v_heads)}
self.ssm_a = Tensor.zeros(self.num_v_heads)
self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), nn.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.attn_gate(x).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_alpha(x).float() + self.ssm_dt["bias"]).softplus() * self.ssm_a).reshape(B, self.num_v_heads, 1, 1).exp()
# 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))
return self.ssm_out((core_attn_out * out_gate.silu()).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_v_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(config, config.ssm) if config.ssm and (i+1) % config.full_attention_interval != 0 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 = nn.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
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')})
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 ('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')),
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,
full_attention_interval=kv.get(f'{arch}.full_attention_interval', 0),
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 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:
sp, nt = v_start_pos.bind(start_pos), v_toks.bind(min(chunk_size, len(tokens) - start_pos))
out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
start_pos += nt.val
# 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]