From 5331889b0676064c6eade4eec80471935662e4e1 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 30 Jul 2026 13:54:01 +0000 Subject: [PATCH] speed for testing --- extra/benchmark_llm.py | 113 ++++++++++++++++++++ test/external/external_test_llm_opencode.py | 90 ++++++++++++++++ test/unit/test_jit.py | 23 +++- tinygrad/codegen/__init__.py | 11 +- tinygrad/device.py | 1 + tinygrad/engine/jit.py | 3 +- tinygrad/engine/realize.py | 22 +++- tinygrad/helpers.py | 1 + tinygrad/llm/cli.py | 5 +- tinygrad/llm/cpu.py | 2 +- tinygrad/llm/gguf.py | 20 +++- tinygrad/llm/model.py | 9 -- tinygrad/runtime/ops_cpu.py | 2 + 13 files changed, 277 insertions(+), 25 deletions(-) create mode 100644 extra/benchmark_llm.py create mode 100644 test/external/external_test_llm_opencode.py diff --git a/extra/benchmark_llm.py b/extra/benchmark_llm.py new file mode 100644 index 0000000000..68c242185a --- /dev/null +++ b/extra/benchmark_llm.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""Benchmark tinygrad LLM prefill and decode independently.""" +from __future__ import annotations + +import argparse, gc, json, statistics, time +from dataclasses import asdict, dataclass + +from tinygrad import Context, Device, Tensor, UOp +from tinygrad.helpers import fetch, profile_marker +from tinygrad.llm.cli import models +from tinygrad.llm.model import Transformer + +@dataclass +class Result: + prompt_tokens: int + decode_tokens: int + time_to_first_token_s: float + prefill_tokens_per_s: float + decode_tokens_per_s: float + decode_p50_ms: float + decode_p95_ms: float + output_tokens: list[int] + +def percentile(values:list[float], percentile:float) -> float: + ordered = sorted(values) + return ordered[round((len(ordered) - 1) * percentile)] + +def synthetic_prompt(length:int, vocab_size:int, salt:int) -> list[int]: + assert length > 0 and vocab_size > 256 + return [256 + salt % (vocab_size - 256)] + [256 + (i * 7919) % (vocab_size - 256) for i in range(1, length)] + +def benchmark(model:Transformer, prompt:list[int], decode_tokens:int, chunk_size:int) -> Result: + gen = model.generate(prompt.copy(), chunk_size=chunk_size) + profile_marker(f"prefill {len(prompt)} start") + begin = time.perf_counter() + output_tokens = [next(gen)] + ttft = time.perf_counter() - begin + profile_marker(f"prefill {len(prompt)} end") + + decode_times: list[float] = [] + profile_marker(f"decode {len(prompt)} start") + for _ in range(decode_tokens): + begin = time.perf_counter() + output_tokens.append(next(gen)) + decode_times.append(time.perf_counter() - begin) + profile_marker(f"decode {len(prompt)} end") + + return Result(len(prompt), decode_tokens, ttft, len(prompt) / ttft, decode_tokens / sum(decode_times), + statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3, output_tokens) + +def benchmark_decode_position(model:Transformer, position:int, decode_tokens:int) -> Result: + token = Tensor([[0]], dtype="int32", device=Device.DEFAULT).realize() + temperature = Tensor([0.0], device=Device.DEFAULT).realize() + decode_times, output_tokens = [], [] + for pos in range(position, position + decode_tokens): + begin = time.perf_counter() + output_tokens.append(int(model(token, UOp.variable("start_pos", 0, model.max_context-1).bind(pos), temperature).realize().item())) + decode_times.append(time.perf_counter() - begin) + return Result(position, decode_tokens, 0.0, 0.0, decode_tokens / sum(decode_times), + statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3, output_tokens) + +def main() -> None: + parser = argparse.ArgumentParser(description="Measure LLM prefill and steady-state decode speed") + parser.add_argument("--model", default="qwen3:0.6b", help="Model preset or local GGUF path") + parser.add_argument("--max-context", type=int, default=32768) + parser.add_argument("--prompt-tokens", type=int, nargs="+", default=[128, 2048, 8192]) + parser.add_argument("--decode-tokens", type=int, default=32) + parser.add_argument("--decode-position", type=int, nargs="+") + parser.add_argument("--chunk-size", type=int, default=256) + parser.add_argument("--beam", type=int, default=2) + parser.add_argument("--jit-batch-size", type=int, default=448) + parser.add_argument("--parallel-compile", type=int, default=12) + parser.add_argument("--realize", action="store_true") + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + + if args.decode_tokens < 1: parser.error("--decode-tokens must be positive") + if args.chunk_size < 1: parser.error("--chunk-size must be positive") + if args.decode_position is None and max(args.prompt_tokens) + args.decode_tokens >= args.max_context: + parser.error("prompt plus decode tokens must fit within --max-context") + if args.decode_position is not None and max(args.decode_position) + args.decode_tokens >= args.max_context: + parser.error("decode position plus decode tokens must fit within --max-context") + + begin = time.perf_counter() + path = fetch(models.get(args.model, args.model)) + fetched = time.perf_counter() + model, kv = Transformer.from_gguf(path, args.max_context, realize=args.realize) + loaded = time.perf_counter() + vocab_size = len(kv["tokenizer.ggml.tokens"]) + print(f"startup: fetch={fetched-begin:.2f}s load={loaded-fetched:.2f}s", flush=True) + with Context(BEAM=args.beam, JIT_BATCH_SIZE=args.jit_batch_size, PARALLEL_COMPILE=args.parallel_compile): + model.warmup(args.chunk_size) + startup = time.perf_counter() - begin + print(f"startup: warmup={startup-(loaded-begin):.2f}s total={startup:.2f}s", flush=True) + gc.freeze() + + results = [benchmark_decode_position(model, pos, args.decode_tokens) for pos in args.decode_position] if args.decode_position is not None else \ + [benchmark(model, synthetic_prompt(n, vocab_size, salt=i+1), args.decode_tokens, args.chunk_size) + for i, n in enumerate(args.prompt_tokens)] + if args.json: + print(json.dumps({"model": args.model, "max_context": args.max_context, "chunk_size": args.chunk_size, + "beam": args.beam, "jit_batch_size": args.jit_batch_size, "parallel_compile": args.parallel_compile, + "realize": args.realize, "startup_s": startup, "results": [asdict(x) for x in results]}, indent=2)) + return + + print(f"model={args.model} max_context={args.max_context} chunk_size={args.chunk_size} beam={args.beam} " + f"jit_batch_size={args.jit_batch_size} parallel_compile={args.parallel_compile} realize={args.realize} startup={startup:.2f}s") + print(f"{'prompt':>8} {'TTFT':>10} {'prefill':>14} {'decode':>14} {'decode p50':>12} {'decode p95':>12}") + for result in results: + print(f"{result.prompt_tokens:8d} {result.time_to_first_token_s:9.3f}s {result.prefill_tokens_per_s:11.1f} t/s " + f"{result.decode_tokens_per_s:11.1f} t/s {result.decode_p50_ms:9.2f} ms {result.decode_p95_ms:9.2f} ms") + +if __name__ == "__main__": main() diff --git a/test/external/external_test_llm_opencode.py b/test/external/external_test_llm_opencode.py new file mode 100644 index 0000000000..31c422d820 --- /dev/null +++ b/test/external/external_test_llm_opencode.py @@ -0,0 +1,90 @@ +"""Real-model OpenCode regression. + +Run against an existing server: + RUN_LLM_OPENCODE_REGRESSION=1 LLM_BASE_URL=http://127.0.0.1:8000/v1 \ + python -m pytest test/external/external_test_llm_opencode.py -v + +Or set LLM_GGUF and let the test start the tinygrad server. +""" +from __future__ import annotations + +import json, os, pathlib, re, shutil, socket, subprocess, sys, tempfile, time, unittest, urllib.request + +RUN_REGRESSION = os.getenv("RUN_LLM_OPENCODE_REGRESSION") == "1" + +def _server_ready(base_url:str) -> bool: + try: + with urllib.request.urlopen(base_url.rstrip("/") + "/models", timeout=1) as response: return response.status == 200 + except OSError: return False + +@unittest.skipUnless(RUN_REGRESSION, "set RUN_LLM_OPENCODE_REGRESSION=1 to run the OpenCode regression") +class TestLLMOpenCode(unittest.TestCase): + server:subprocess.Popen|None = None + server_log:tempfile._TemporaryFileWrapper|None = None + + @classmethod + def setUpClass(cls): + if shutil.which("opencode") is None: raise unittest.SkipTest("opencode is not installed") + if base_url := os.getenv("LLM_BASE_URL"): + cls.base_url = base_url.rstrip("/") + if not cls.base_url.endswith("/v1"): cls.base_url += "/v1" + if not _server_ready(cls.base_url): raise RuntimeError(f"LLM server is not responding at {cls.base_url}") + return + + model = pathlib.Path(os.environ["LLM_GGUF"]) + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + cls.base_url = f"http://127.0.0.1:{port}/v1" + cls.server_log = tempfile.NamedTemporaryFile(mode="w+", prefix="tinygrad-llm-") + cls.server = subprocess.Popen( + [sys.executable, "-m", "tinygrad.llm", "--model", str(model), "--serve", str(port), "--max_context", "262144"], + stdout=cls.server_log, stderr=subprocess.STDOUT, start_new_session=True) + deadline = time.monotonic() + 180 + while time.monotonic() < deadline and cls.server.poll() is None: + if _server_ready(cls.base_url): return + time.sleep(0.25) + cls.server_log.seek(0) + raise RuntimeError(f"LLM server failed to start:\n{cls.server_log.read()[-8000:]}") + + @classmethod + def tearDownClass(cls): + if cls.server is not None: + cls.server.terminate() + try: cls.server.wait(timeout=10) + except subprocess.TimeoutExpired: + cls.server.kill() + cls.server.wait(timeout=10) + if cls.server_log is not None: cls.server_log.close() + + def run_opencode(self, prompt:str, cwd:pathlib.Path) -> str: + config = cwd / "opencode.json" + config.write_text(json.dumps({ + "$schema": "https://opencode.ai/config.json", "permission": {"*": "allow"}, "formatter": False, "lsp": False, + "provider": {"regression": {"npm": "@ai-sdk/openai-compatible", "options": {"baseURL": self.base_url}, + "models": {"tinygrad": {"name": "tinygrad"}}}}, + })) + env = os.environ | {"OPENCODE_CONFIG": str(config)} + result = subprocess.run( + ["opencode", "run", "--pure", "--auto", "--dir", str(cwd), "-m", "regression/tinygrad", prompt], + cwd=cwd, env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=120) + self.assertEqual(result.returncode, 0, result.stdout) + return re.sub(r"\x1b\[[0-?]*[ -/]*[@-~]", "", result.stdout) + + def test_read_tool(self): + with tempfile.TemporaryDirectory() as directory: + cwd, marker = pathlib.Path(directory), "tinygrad-opencode-regression-7f3a91c2" + (cwd / "exact.txt").write_text(marker + "\n") + output = self.run_opencode("Read exact.txt with a tool and reply with its exact contents, with no other text.", cwd) + self.assertRegex(output, r"(?im)^\s*(?:→|>)\s*Read\s+exact\.txt\s*$") + self.assertIn(marker, output) + + def test_shell_tool(self): + with tempfile.TemporaryDirectory() as directory: + cwd = pathlib.Path(directory) + output = self.run_opencode( + "Use the shell tool to run `printf tinygrad-shell-regression > shell-regression.txt`, then report completion.", cwd) + self.assertRegex(output, r"(?im)^\s*(?:\$|→|>)\s*.*printf\s+tinygrad-shell-regression") + self.assertEqual((cwd / "shell-regression.txt").read_text(), "tinygrad-shell-regression") + +if __name__ == "__main__": unittest.main() diff --git a/test/unit/test_jit.py b/test/unit/test_jit.py index 7735776fec..ff1d402172 100644 --- a/test/unit/test_jit.py +++ b/test/unit/test_jit.py @@ -1,7 +1,9 @@ import unittest, numpy as np +from unittest.mock import patch from test.helpers import assert_jit_cache_len from tinygrad import Tensor, TinyJit, Context, UOp, dtypes -from tinygrad.engine.jit import JitError +from tinygrad.engine.jit import JitError, graph_split_rewrite +from tinygrad.uop.ops import Ops def _simple_test(add, extract=lambda x: x, N=10): for _ in range(5): @@ -12,8 +14,25 @@ def _simple_test(add, extract=lambda x: x, N=10): assert_jit_cache_len(add, 1) class TestJit(unittest.TestCase): + def test_graph_batch_size_limit(self): + class FakeGraph: + @staticmethod + def supports_uop(_devs, _call): return True + + def graph_sizes(limit:int|None) -> list[int]: + class FakeDevice: + graph, graph_batch_size_limit = FakeGraph, limit + dev = FakeDevice() + buf = UOp.new_buffer("FAKE", 1, dtypes.float) + prg = UOp(Ops.PROGRAM, src=(UOp.sink(),)) + with patch("tinygrad.engine.jit.Device", {"FAKE":dev}): + linear = graph_split_rewrite(UOp(Ops.LINEAR, src=tuple(prg.call(buf) for _ in range(20))), max_batch_size=4) + return [len(call.src[0].src[0].src) for call in linear.src] + + self.assertEqual(graph_sizes(4), [4, 4, 4, 4, 4]) + self.assertEqual(graph_sizes(None), [4, 8, 8]) + def test_jitbeam_triggers_beam(self): - from unittest.mock import patch from tinygrad.helpers import getenv as _getenv @TinyJit def add(a, b): return (a+b).realize() diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 936fd99225..b10d267e71 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -482,8 +482,15 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: return prg to_program_cache: dict[tuple, UOp] = {} -def to_program(ast:UOp, renderer:Renderer) -> UOp: +def program_cache_key(ast:UOp, renderer:Renderer) -> tuple: config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT) - key = (ast.key, type(renderer), renderer.target, *[x.value for x in config]) + return (ast.key, type(renderer), renderer.target, *[x.value for x in config]) + +def parallel_to_program(args:tuple[UOp, Renderer, tuple]) -> tuple[tuple, UOp]: + ast, renderer, key = args + return key, do_to_program(ast, renderer) + +def to_program(ast:UOp, renderer:Renderer) -> UOp: + key = program_cache_key(ast, renderer) if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer) return prg diff --git a/tinygrad/device.py b/tinygrad/device.py index 9be9ec446a..91dbc37c49 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -333,6 +333,7 @@ class Program(Generic[DeviceType]): class Compiled: profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device. + graph_batch_size_limit:int|None = None pm_lower:Any = None pm_bufferize:Any = None diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 96f92a1b88..da763857c9 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -39,7 +39,8 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp: if len(current_batch) <= 1 and not getenv("GRAPH_ONE_KERNEL"): new_src.extend(current_batch) else: new_src.append(create_graph_call(current_batch)) - max_batch_size *= 2 + max_batch_size = min((max_batch_size * 2, *(dev.graph_batch_size_limit for dev in current_batch_devs + if dev.graph_batch_size_limit is not None))) if DEBUG >= 2: print(f"JIT GRAPHing batch with {len(current_batch)} kernels") current_batch, current_batch_devs = [], [] diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index f11c9688d4..e8af9dafe4 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -1,13 +1,15 @@ from __future__ import annotations from typing import cast, Iterator, Any, Sequence -import time, random, itertools, math, contextlib, weakref, array +import time, random, itertools, math, contextlib, weakref, array, os, multiprocessing +from concurrent.futures import ProcessPoolExecutor +from concurrent.futures.process import BrokenProcessPool from dataclasses import dataclass, replace, field from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple -from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events +from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, PARALLEL_COMPILE from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite from tinygrad.device import Device, Buffer, MultiBuffer from tinygrad.renderer import Estimates -from tinygrad.codegen import to_program +from tinygrad.codegen import to_program, to_program_cache, program_cache_key, parallel_to_program from tinygrad.codegen.opt.postrange import bufs_from_ast # **************** Helpers **************** @@ -268,6 +270,20 @@ if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_li def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> UOp: if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True) if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True) + if jit and PARALLEL_COMPILE.value: + pending:dict[tuple, tuple[UOp, Any, tuple]] = {} + for call in linear.toposort(): + if call.op is not Ops.CALL or call.src[0].op not in (Ops.SINK, Ops.PROGRAM): continue + renderer = Device[call.device if isinstance(call.device, str) else call.device[0]].renderer + key = program_cache_key(call.src[0], renderer) + if key not in to_program_cache: pending.setdefault(key, (call.src[0], renderer, key)) + if len(pending) >= 16: + workers = min(PARALLEL_COMPILE.value, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1), len(pending)) + try: + start_method = "fork" if all(renderer.target.device == "CPU" for _,renderer,_ in pending.values()) else "spawn" + with ProcessPoolExecutor(workers, mp_context=multiprocessing.get_context(start_method)) as pool: + for key,program in pool.map(parallel_to_program, pending.values()): to_program_cache[key] = program + except BrokenProcessPool: pass linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True) if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, jit=jit) return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 86d0136b18..758b153d26 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -232,6 +232,7 @@ class _DEV(ContextVar): DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0) IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0) JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32) +PARALLEL_COMPILE = ContextVar("PARALLEL_COMPILE", 0) CHUNK_SIZE = 2**20 # TinyFS content-addressed store: blob chunk + hash-tree node granularity WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0) TRAINING = ContextVar("TRAINING", 0) diff --git a/tinygrad/llm/cli.py b/tinygrad/llm/cli.py index e799faf66c..d647101300 100644 --- a/tinygrad/llm/cli.py +++ b/tinygrad/llm/cli.py @@ -163,9 +163,8 @@ def main(): # 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]))) + with Context(DEBUG=max(DEBUG.value, 1), PARALLEL_COMPILE=getenv("PARALLEL_COMPILE", 12)): + model.warmup() # start server if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever() diff --git a/tinygrad/llm/cpu.py b/tinygrad/llm/cpu.py index 5ec8627100..51b6e519b3 100644 --- a/tinygrad/llm/cpu.py +++ b/tinygrad/llm/cpu.py @@ -1679,7 +1679,7 @@ def _cpu_topk_uop(out:UOp, sel:UOp, x:UOp, k:int, bias:UOp|None=None, normalize: indices[worst_slot.valid(take)].store(index.cast(dtypes.int32))).end(index) sorted_values = selected - # Ascending score order, with larger indices first on ties, matches reversing the C fallback's descending list. + # Ascending score order, with larger indices first on ties, matches the reference implementation's reversed descending list. for end in range(k - 1, 0, -1): for slot in range(end): left_score, right_score = scores.after(sorted_values)[slot].load(), scores.after(sorted_values)[slot + 1].load() diff --git a/tinygrad/llm/gguf.py b/tinygrad/llm/gguf.py index 0a199048fb..5bcc9473b0 100644 --- a/tinygrad/llm/gguf.py +++ b/tinygrad/llm/gguf.py @@ -1,10 +1,11 @@ -import functools, io, pathlib, re, struct, weakref +import functools, io, pathlib, re, struct, weakref, mmap from typing import Any, Callable from tinygrad.tensor import Tensor from tinygrad.uop.ops import UOp from tinygrad.dtype import dtypes -from tinygrad.helpers import prod, round_up +from tinygrad.helpers import prod, round_up, mv_address +from tinygrad.device import Device 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. @@ -22,6 +23,17 @@ _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)} _quantized_tensors:weakref.WeakKeyDictionary[UOp, tuple[UOp, int]] = weakref.WeakKeyDictionary() +_cpu_mapped_ggufs:dict[tuple[pathlib.Path, int, int], tuple[mmap.mmap, Tensor]] = {} + +def _gguf_tensor(path:pathlib.Path) -> Tensor: + path = path.resolve() + if not Device.DEFAULT.startswith("CPU"): return Tensor(path).to(None) + stat = path.stat() + key = (path, stat.st_mtime_ns, stat.st_size) + if key not in _cpu_mapped_ggufs: + with path.open("rb") as f: mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_COPY) + _cpu_mapped_ggufs[key] = mm, Tensor.from_blob(mv_address(memoryview(mm)), (len(mm),), dtype=dtypes.uint8, device=Device.DEFAULT) + return _cpu_mapped_ggufs[key][1] def get_ggml_quantization(tensor:Tensor) -> tuple[Tensor, int]|None: if (meta:=_quantized_tensors.get(tensor.uop)) is None: return None @@ -184,8 +196,8 @@ def gguf_load(fn: Tensor|str|pathlib.Path) -> tuple[dict, dict[str, 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)).to(None)) + kv, sd = _gguf_parse(fn if isinstance(fn, Tensor) else _gguf_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).to(None))[1]) + for pp in _gguf_split_paths(pathlib.Path(fn), kv)[1:]: sd.update(_gguf_parse(_gguf_tensor(pp))[1]) return kv, sd diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 48e6d0a3d2..cb7242095b 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -1378,15 +1378,6 @@ class Transformer: prefill_batch = getenv("PREFILL_JIT_BATCH_SIZE", 16 if str(device).startswith("CPU") else 128) with Context(JIT_BATCH_SIZE=prefill_batch): next(warm) next(warm) - # AMD flash decode specializes on its attention partition. Fused CPU decode uses one graph for the full cache. - if self.max_context > short_decode_len and not str(device).startswith("CPU"): - self.rollout_jits[self.max_context] = TinyJit( - functools.partial(self.forward_recurrent_decode, decode_len=self.max_context, sample=False)) - self.rollout_jits[self.max_context].cnt = 1 - long_result = self(Tensor([[0]], dtype="int32", device=device), - UOp.variable("start_pos", 0, self.max_context-1).bind(short_decode_len), Tensor([0.0], device=device)) - assert isinstance(long_result, Tensor) - long_result.realize() self._warming_up = False else: for salt in range(2): next(self.generate([salt] + [0] * (warm_len - 1), chunk_size=chunk_size)) diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 40a1ac2021..fbb00fa134 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -281,6 +281,8 @@ class CPUAllocator(HCQAllocator): def _unmap(self, mb): pass # CPU _do_map returns a view wrapper, nothing to release class CPUDevice(HCQCompiled): + graph_batch_size_limit = 64 + pm_lower = PatternMatcher([ (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_host_queue)])