speed for testing

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()

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()

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()

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

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

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 = [], []
 

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)

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)

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()

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()

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

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))

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)])